Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit87f12f6

Browse files
committed
added api key id to recorded instead of pushing it through context/actor
1 parent8de5657 commit87f12f6

File tree

5 files changed

+38
-14
lines changed

5 files changed

+38
-14
lines changed

‎enterprise/aibridged/http.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
5757
}
5858

5959
// Rewire request context to include actor.
60-
r=r.WithContext(aibridge.AsActor(ctx,resp.GetApiKeyId(),resp.GetOwnerId(),nil))
60+
r=r.WithContext(aibridge.AsActor(ctx,resp.GetOwnerId(),nil))
6161

6262
id,err:=uuid.Parse(resp.GetOwnerId())
6363
iferr!=nil {
@@ -68,6 +68,7 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
6868

6969
handler,err:=s.GetRequestHandler(ctx,Request{
7070
SessionKey:key,
71+
APIKeyID:resp.ApiKeyId,
7172
InitiatorID:id,
7273
})
7374
iferr!=nil {

‎enterprise/aibridged/pool.go‎

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -111,17 +111,9 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
111111
// may visit the slow path unnecessarily.
112112
deferp.cache.Wait()
113113

114-
recorder:=aibridge.NewRecorder(p.logger.Named("recorder"),func() (aibridge.Recorder,error) {
115-
client,err:=clientFn()
116-
iferr!=nil {
117-
returnnil,xerrors.Errorf("acquire client: %w",err)
118-
}
119-
120-
return&recorderTranslation{client:client},nil
121-
})
122-
123114
// Fast path.
124-
bridge,ok:=p.cache.Get(req.InitiatorID.String())
115+
cacheKey:=req.InitiatorID.String()+req.APIKeyID
116+
bridge,ok:=p.cache.Get(cacheKey)
125117
ifok&&bridge!=nil {
126118
// TODO: future improvement:
127119
// Once we can detect token expiry against an MCP server, we no longer need to let these instances
@@ -131,6 +123,15 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
131123
returnbridge,nil
132124
}
133125

126+
recorder:=aibridge.NewRecorder(p.logger.Named("recorder"),func() (aibridge.Recorder,error) {
127+
client,err:=clientFn()
128+
iferr!=nil {
129+
returnnil,xerrors.Errorf("acquire client: %w",err)
130+
}
131+
132+
return&recorderTranslation{apiKeyID:req.APIKeyID,client:client},nil
133+
})
134+
134135
// Slow path.
135136
// Creating an *aibridge.RequestBridge may take some time, so gate all subsequent callers behind the initial request and return the resulting value.
136137
// TODO: track startup time since it adds latency to first request (histogram count will also help us see how often this occurs).
@@ -158,7 +159,7 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
158159
returnnil,xerrors.Errorf("create new request bridge: %w",err)
159160
}
160161

161-
p.cache.SetWithTTL(req.InitiatorID.String(),bridge,cacheCost,p.options.TTL)
162+
p.cache.SetWithTTL(cacheKey,bridge,cacheCost,p.options.TTL)
162163

163164
returnbridge,nil
164165
})

‎enterprise/aibridged/pool_test.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,15 @@ func TestPool(t *testing.T) {
5050
inst,err:=pool.Acquire(t.Context(), aibridged.Request{
5151
SessionKey:"key",
5252
InitiatorID:id,
53+
APIKeyID:"apiKeyId",
5354
},clientFn,newMockMCPFactory(mcpProxy))
5455
require.NoError(t,err,"acquire pool instance")
5556

5657
// ...and it will return it when acquired again.
5758
instB,err:=pool.Acquire(t.Context(), aibridged.Request{
5859
SessionKey:"key",
5960
InitiatorID:id,
61+
APIKeyID:"apiKeyId",
6062
},clientFn,newMockMCPFactory(mcpProxy))
6163
require.NoError(t,err,"acquire pool instance")
6264
require.Same(t,inst,instB)
@@ -74,6 +76,7 @@ func TestPool(t *testing.T) {
7476
inst2,err:=pool.Acquire(t.Context(), aibridged.Request{
7577
SessionKey:"key",
7678
InitiatorID:id2,
79+
APIKeyID:"apiKeyId",
7780
},clientFn,newMockMCPFactory(mcpProxy))
7881
require.NoError(t,err,"acquire pool instance")
7982
require.NotSame(t,inst,inst2)
@@ -84,6 +87,23 @@ func TestPool(t *testing.T) {
8487
require.EqualValues(t,1,metrics.Hits())
8588
require.EqualValues(t,2,metrics.Misses())
8689

90+
// Different instance is returned for different api key id
91+
mcpProxy.EXPECT().Init(gomock.Any()).Times(1).Return(nil)
92+
93+
inst2B,err:=pool.Acquire(t.Context(), aibridged.Request{
94+
SessionKey:"key",
95+
InitiatorID:id2,
96+
APIKeyID:"newApiKeyId",
97+
},clientFn,newMockMCPFactory(mcpProxy))
98+
require.NoError(t,err,"acquire pool instance 2B")
99+
require.NotSame(t,inst2,inst2B)
100+
101+
metrics=pool.Metrics()
102+
require.EqualValues(t,3,metrics.KeysAdded())
103+
require.EqualValues(t,2,metrics.KeysEvicted())
104+
require.EqualValues(t,1,metrics.Hits())
105+
require.EqualValues(t,3,metrics.Misses())
106+
87107
// TODO: add test for expiry.
88108
// This requires Go 1.25's [synctest](https://pkg.go.dev/testing/synctest) since the
89109
// internal cache lib cannot be tested using coder/quartz.

‎enterprise/aibridged/request.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,5 +4,6 @@ import "github.com/google/uuid"
44

55
typeRequeststruct {
66
SessionKeystring
7+
APIKeyIDstring
78
InitiatorID uuid.UUID
89
}

‎enterprise/aibridged/translator.go‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,14 @@ var _ aibridge.Recorder = &recorderTranslation{}
2020

2121
// recorderTranslation satisfies the aibridge.Recorder interface and translates calls into dRPC calls to aibridgedserver.
2222
typerecorderTranslationstruct {
23-
client proto.DRPCRecorderClient
23+
apiKeyIDstring
24+
client proto.DRPCRecorderClient
2425
}
2526

2627
func (t*recorderTranslation)RecordInterception(ctx context.Context,req*aibridge.InterceptionRecord)error {
2728
_,err:=t.client.RecordInterception(ctx,&proto.RecordInterceptionRequest{
2829
Id:req.ID,
29-
ApiKeyId:req.APIKeyID,
30+
ApiKeyId:t.apiKeyID,
3031
InitiatorId:req.InitiatorID,
3132
Provider:req.Provider,
3233
Model:req.Model,

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp