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

Commitc94e2cc

Browse files
committed
chore: record user agent in interception metadata
1 parent3e33a1c commitc94e2cc

File tree

7 files changed

+194
-6
lines changed

7 files changed

+194
-6
lines changed

‎enterprise/aibridged/aibridged_integration_test.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package aibridged_test
33
import (
44
"bytes"
55
"context"
6+
"encoding/json"
67
"fmt"
78
"net/http"
89
"net/http/httptest"
@@ -208,10 +209,13 @@ func TestIntegration(t *testing.T) {
208209
}
209210
]
210211
}`))
212+
userAgent:="userAgent123"
211213
require.NoError(t,err,"make request to test server")
212214
req.Header.Add("Authorization","Bearer "+apiKey.Key)
213215
req.Header.Add("Accept","application/json")
216+
req.Header.Add("User-Agent",userAgent)
214217

218+
require.Equal(t,userAgent,req.UserAgent())
215219
// When: aibridged handles the request.
216220
rec:=httptest.NewRecorder()
217221
srv.ServeHTTP(rec,req)
@@ -234,6 +238,13 @@ func TestIntegration(t *testing.T) {
234238
require.True(t,intc0.StartedAt.Before(intc0.EndedAt.Time))
235239
require.Less(t,intc0.EndedAt.Time.Sub(intc0.StartedAt),5*time.Second)
236240

241+
require.True(t,intc0.Metadata.Valid)
242+
meta:=map[string]any{}
243+
err=json.Unmarshal(intc0.Metadata.RawMessage,&meta)
244+
require.NoError(t,err)
245+
require.Contains(t,meta,"user-agent")
246+
require.Equal(t,userAgent,meta["user-agent"])
247+
237248
prompts,err:=db.GetAIBridgeUserPromptsByInterceptionID(ctx,interceptions[0].ID)
238249
require.NoError(t,err)
239250
require.Len(t,prompts,1)

‎enterprise/aibridged/http.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,9 +67,10 @@ func (s *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
6767
}
6868

6969
handler,err:=s.GetRequestHandler(ctx,Request{
70-
SessionKey:key,
7170
APIKeyID:resp.ApiKeyId,
7271
InitiatorID:id,
72+
SessionKey:key,
73+
UserAgent:r.UserAgent(),
7374
})
7475
iferr!=nil {
7576
logger.Warn(ctx,"failed to acquire request handler",slog.Error(err))

‎enterprise/aibridged/pool.go‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -129,7 +129,11 @@ func (p *CachedBridgePool) Acquire(ctx context.Context, req Request, clientFn Cl
129129
returnnil,xerrors.Errorf("acquire client: %w",err)
130130
}
131131

132-
return&recorderTranslation{apiKeyID:req.APIKeyID,client:client},nil
132+
return&recorderTranslation{
133+
apiKeyID:req.APIKeyID,
134+
client:client,
135+
userAgent:req.UserAgent,
136+
},nil
133137
})
134138

135139
// Slow path.

‎enterprise/aibridged/pool_test.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,9 +56,10 @@ func TestPool(t *testing.T) {
5656

5757
// ...and it will return it when acquired again.
5858
instB,err:=pool.Acquire(t.Context(), aibridged.Request{
59-
SessionKey:"key",
59+
SessionKey:"differentkey",
6060
InitiatorID:id,
6161
APIKeyID:apiKeyID1.String(),
62+
UserAgent:"some user-agent",
6263
},clientFn,newMockMCPFactory(mcpProxy))
6364
require.NoError(t,err,"acquire pool instance")
6465
require.Same(t,inst,instB)

‎enterprise/aibridged/request.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,8 @@ package aibridged
33
import"github.com/google/uuid"
44

55
typeRequeststruct {
6-
SessionKeystring
76
APIKeyIDstring
87
InitiatorID uuid.UUID
8+
SessionKeystring
9+
UserAgentstring
910
}

‎enterprise/aibridged/translator.go‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,15 +16,27 @@ import (
1616
"github.com/coder/aibridge"
1717
)
1818

19+
const (
20+
META_KEY_USERAGENT="user-agent"
21+
)
22+
1923
var_ aibridge.Recorder=&recorderTranslation{}
2024

2125
// recorderTranslation satisfies the aibridge.Recorder interface and translates calls into dRPC calls to aibridgedserver.
2226
typerecorderTranslationstruct {
23-
apiKeyIDstring
24-
client proto.DRPCRecorderClient
27+
apiKeyIDstring
28+
client proto.DRPCRecorderClient
29+
userAgentstring
2530
}
2631

2732
func (t*recorderTranslation)RecordInterception(ctx context.Context,req*aibridge.InterceptionRecord)error {
33+
ift.userAgent!="" {
34+
ifreq.Metadata==nil {
35+
req.Metadata=map[string]any{}
36+
}
37+
req.Metadata[META_KEY_USERAGENT]=t.userAgent
38+
}
39+
2840
_,err:=t.client.RecordInterception(ctx,&proto.RecordInterceptionRequest{
2941
Id:req.ID,
3042
ApiKeyId:t.apiKeyID,
Lines changed: 158 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,158 @@
1+
package aibridged
2+
3+
import (
4+
"context"
5+
"testing"
6+
"time"
7+
8+
"github.com/google/uuid"
9+
"github.com/stretchr/testify/require"
10+
"google.golang.org/protobuf/proto"
11+
"google.golang.org/protobuf/types/known/anypb"
12+
"google.golang.org/protobuf/types/known/structpb"
13+
"google.golang.org/protobuf/types/known/timestamppb"
14+
15+
"github.com/coder/aibridge"
16+
abpb"github.com/coder/coder/v2/enterprise/aibridged/proto"
17+
"github.com/coder/coder/v2/testutil"
18+
)
19+
20+
typemockClientstruct {
21+
abpb.DRPCRecorderClient
22+
got*abpb.RecordInterceptionRequest
23+
}
24+
25+
func (mc*mockClient)RecordInterception(ctx context.Context,in*abpb.RecordInterceptionRequest) (*abpb.RecordInterceptionResponse,error) {
26+
mc.got=in
27+
return&abpb.RecordInterceptionResponse{},nil
28+
}
29+
30+
funcmustAnypbNew(t*testing.T,src proto.Message)*anypb.Any {
31+
ret,err:=anypb.New(src)
32+
require.NoError(t,err)
33+
returnret
34+
}
35+
36+
funcTestRecordInterception(t*testing.T) {
37+
t.Parallel()
38+
39+
tests:= []struct {
40+
namestring
41+
apiKeyIDstring
42+
userAgentstring
43+
in aibridge.InterceptionRecord
44+
expect*abpb.RecordInterceptionRequest
45+
}{
46+
{
47+
name:"simple",
48+
apiKeyID:"key",
49+
userAgent:"user-agent",
50+
in: aibridge.InterceptionRecord{
51+
ID: uuid.UUID{1}.String(),
52+
InitiatorID: uuid.UUID{2}.String(),
53+
Provider:"prov",
54+
Model:"model",
55+
Metadata:map[string]any{"some":"data"},
56+
StartedAt:time.UnixMicro(123),
57+
},
58+
expect:&abpb.RecordInterceptionRequest{
59+
Id: uuid.UUID{1}.String(),
60+
ApiKeyId:"key",
61+
InitiatorId: uuid.UUID{2}.String(),
62+
Provider:"prov",
63+
Model:"model",
64+
Metadata:map[string]*anypb.Any{
65+
"some":mustAnypbNew(t,structpb.NewStringValue("data")),
66+
META_KEY_USERAGENT:mustAnypbNew(t,structpb.NewStringValue("user-agent")),
67+
},
68+
StartedAt:timestamppb.New(time.UnixMicro(123)),
69+
},
70+
},
71+
{
72+
name:"empty-user-agent",
73+
apiKeyID:"key",
74+
in: aibridge.InterceptionRecord{
75+
ID: uuid.UUID{1}.String(),
76+
InitiatorID: uuid.UUID{2}.String(),
77+
Provider:"prov",
78+
Model:"model",
79+
Metadata:map[string]any{"some":"data"},
80+
StartedAt:time.UnixMicro(123),
81+
},
82+
expect:&abpb.RecordInterceptionRequest{
83+
Id: uuid.UUID{1}.String(),
84+
ApiKeyId:"key",
85+
InitiatorId: uuid.UUID{2}.String(),
86+
Provider:"prov",
87+
Model:"model",
88+
Metadata:map[string]*anypb.Any{
89+
"some":mustAnypbNew(t,structpb.NewStringValue("data")),
90+
},
91+
StartedAt:timestamppb.New(time.UnixMicro(123)),
92+
},
93+
},
94+
{
95+
name:"overrides-user-agent",
96+
apiKeyID:"key",
97+
userAgent:"user-agent",
98+
in: aibridge.InterceptionRecord{
99+
ID: uuid.UUID{1}.String(),
100+
InitiatorID: uuid.UUID{2}.String(),
101+
Provider:"prov",
102+
Model:"model",
103+
Metadata:map[string]any{META_KEY_USERAGENT:"key-already-set"},
104+
StartedAt:time.UnixMicro(123),
105+
},
106+
expect:&abpb.RecordInterceptionRequest{
107+
Id: uuid.UUID{1}.String(),
108+
ApiKeyId:"key",
109+
InitiatorId: uuid.UUID{2}.String(),
110+
Provider:"prov",
111+
Model:"model",
112+
Metadata:map[string]*anypb.Any{
113+
META_KEY_USERAGENT:mustAnypbNew(t,structpb.NewStringValue("user-agent")),
114+
},
115+
StartedAt:timestamppb.New(time.UnixMicro(123)),
116+
},
117+
},
118+
{
119+
name:"user-agent-empty-metadata",
120+
apiKeyID:"key",
121+
userAgent:"user-agent",
122+
in: aibridge.InterceptionRecord{
123+
ID: uuid.UUID{1}.String(),
124+
InitiatorID: uuid.UUID{2}.String(),
125+
Provider:"prov",
126+
Model:"model",
127+
StartedAt:time.UnixMicro(123),
128+
},
129+
expect:&abpb.RecordInterceptionRequest{
130+
Id: uuid.UUID{1}.String(),
131+
ApiKeyId:"key",
132+
InitiatorId: uuid.UUID{2}.String(),
133+
Provider:"prov",
134+
Model:"model",
135+
Metadata:map[string]*anypb.Any{
136+
META_KEY_USERAGENT:mustAnypbNew(t,structpb.NewStringValue("user-agent")),
137+
},
138+
StartedAt:timestamppb.New(time.UnixMicro(123)),
139+
},
140+
},
141+
}
142+
143+
for_,tc:=rangetests {
144+
t.Run(tc.name,func(t*testing.T) {
145+
t.Parallel()
146+
ctx:=testutil.Context(t,testutil.WaitShort)
147+
148+
mc:=&mockClient{}
149+
rt:=&recorderTranslation{
150+
apiKeyID:tc.apiKeyID,
151+
client:mc,
152+
userAgent:tc.userAgent,
153+
}
154+
rt.RecordInterception(ctx,&tc.in)
155+
require.Equal(t,tc.expect,mc.got)
156+
})
157+
}
158+
}

0 commit comments

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