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

Commitd0fb459

Browse files
authored
feat: add RecordInterceptionEnded rpc (#20494)
Adds RPC that marks interception as completed.Added to aibridge incoder/aibridge#43fixescoder/internal#1051
1 parentffe22a0 commitd0fb459

File tree

10 files changed

+584
-299
lines changed

10 files changed

+584
-299
lines changed

‎enterprise/x/aibridged/aibridged_test.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,7 @@ func TestRouting(t *testing.T) {
305305
interceptionID=in.GetId()
306306
return&proto.RecordInterceptionResponse{},nil
307307
})
308+
client.EXPECT().RecordInterceptionEnded(gomock.Any(),gomock.Any()).Times(tc.expectedHits)
308309

309310
// Given: aibridged is started.
310311
srv,err:=aibridged.New(t.Context(),pool,func(ctx context.Context) (aibridged.DRPCClient,error) {

‎enterprise/x/aibridged/aibridgedmock/clientmock.go‎

Lines changed: 15 additions & 0 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎enterprise/x/aibridged/proto/aibridged.pb.go‎

Lines changed: 431 additions & 293 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎enterprise/x/aibridged/proto/aibridged.proto‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ service Recorder {
1111
// RecordInterception creates a new interception record to which all other sub-resources
1212
// (token, prompt, tool uses) will be related.
1313
rpcRecordInterception(RecordInterceptionRequest)returns (RecordInterceptionResponse);
14+
rpcRecordInterceptionEnded(RecordInterceptionEndedRequest)returns (RecordInterceptionEndedResponse);
1415
rpcRecordTokenUsage(RecordTokenUsageRequest)returns (RecordTokenUsageResponse);
1516
rpcRecordPromptUsage(RecordPromptUsageRequest)returns (RecordPromptUsageResponse);
1617
rpcRecordToolUsage(RecordToolUsageRequest)returns (RecordToolUsageResponse);
@@ -45,6 +46,13 @@ message RecordInterceptionRequest {
4546

4647
messageRecordInterceptionResponse {}
4748

49+
messageRecordInterceptionEndedRequest {
50+
stringid=1;// UUID.
51+
google.protobuf.Timestampended_at=2;
52+
}
53+
54+
messageRecordInterceptionEndedResponse {}
55+
4856
messageRecordTokenUsageRequest {
4957
stringinterception_id=1;// UUID.
5058
stringmsg_id=2;// ID provided by provider.

‎enterprise/x/aibridged/proto/aibridged_drpc.pb.go‎

Lines changed: 43 additions & 3 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎enterprise/x/aibridged/translator.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,14 @@ func (t *recorderTranslation) RecordInterception(ctx context.Context, req *aibri
3535
returnerr
3636
}
3737

38+
func (t*recorderTranslation)RecordInterceptionEnded(ctx context.Context,req*aibridge.InterceptionRecordEnded)error {
39+
_,err:=t.client.RecordInterceptionEnded(ctx,&proto.RecordInterceptionEndedRequest{
40+
Id:req.ID,
41+
EndedAt:timestamppb.New(req.EndedAt),
42+
})
43+
returnerr
44+
}
45+
3846
func (t*recorderTranslation)RecordPromptUsage(ctx context.Context,req*aibridge.PromptUsageRecord)error {
3947
_,err:=t.client.RecordPromptUsage(ctx,&proto.RecordPromptUsageRequest{
4048
InterceptionId:req.InterceptionID,

‎enterprise/x/aibridgedserver/aibridgedserver.go‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ type store interface {
5555
InsertAIBridgeTokenUsage(ctx context.Context,arg database.InsertAIBridgeTokenUsageParams) (database.AIBridgeTokenUsage,error)
5656
InsertAIBridgeUserPrompt(ctx context.Context,arg database.InsertAIBridgeUserPromptParams) (database.AIBridgeUserPrompt,error)
5757
InsertAIBridgeToolUsage(ctx context.Context,arg database.InsertAIBridgeToolUsageParams) (database.AIBridgeToolUsage,error)
58+
UpdateAIBridgeInterceptionEnded(ctx context.Context,intcID database.UpdateAIBridgeInterceptionEndedParams) (database.AIBridgeInterception,error)
5859

5960
// MCPConfigurator-related queries.
6061
GetExternalAuthLinksByUserID(ctx context.Context,userID uuid.UUID) ([]database.ExternalAuthLink,error)
@@ -129,6 +130,26 @@ func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterce
129130
return&proto.RecordInterceptionResponse{},nil
130131
}
131132

133+
func (s*Server)RecordInterceptionEnded(ctx context.Context,in*proto.RecordInterceptionEndedRequest) (*proto.RecordInterceptionEndedResponse,error) {
134+
//nolint:gocritic // AIBridged has specific authz rules.
135+
ctx=dbauthz.AsAIBridged(ctx)
136+
137+
intcID,err:=uuid.Parse(in.GetId())
138+
iferr!=nil {
139+
returnnil,xerrors.Errorf("invalid interception ID %q: %w",in.GetId(),err)
140+
}
141+
142+
_,err=s.store.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
143+
ID:intcID,
144+
EndedAt:in.EndedAt.AsTime(),
145+
})
146+
iferr!=nil {
147+
returnnil,xerrors.Errorf("end interception: %w",err)
148+
}
149+
150+
return&proto.RecordInterceptionEndedResponse{},nil
151+
}
152+
132153
func (s*Server)RecordTokenUsage(ctx context.Context,in*proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse,error) {
133154
//nolint:gocritic // AIBridged has specific authz rules.
134155
ctx=dbauthz.AsAIBridged(ctx)

‎enterprise/x/aibridgedserver/aibridgedserver_test.go‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -423,6 +423,60 @@ func TestRecordInterception(t *testing.T) {
423423
)
424424
}
425425

426+
funcTestRecordInterceptionEnded(t*testing.T) {
427+
t.Parallel()
428+
429+
testRecordMethod(t,
430+
func(srv*aibridgedserver.Server,ctx context.Context,req*proto.RecordInterceptionEndedRequest) (*proto.RecordInterceptionEndedResponse,error) {
431+
returnsrv.RecordInterceptionEnded(ctx,req)
432+
},
433+
[]testRecordMethodCase[*proto.RecordInterceptionEndedRequest]{
434+
{
435+
name:"ok",
436+
request:&proto.RecordInterceptionEndedRequest{
437+
Id: uuid.UUID{1}.String(),
438+
EndedAt:timestamppb.Now(),
439+
},
440+
setupMocks:func(t*testing.T,db*dbmock.MockStore,req*proto.RecordInterceptionEndedRequest) {
441+
interceptionID,err:=uuid.Parse(req.GetId())
442+
assert.NoError(t,err,"parse interception UUID")
443+
444+
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), database.UpdateAIBridgeInterceptionEndedParams{
445+
ID:interceptionID,
446+
EndedAt:req.EndedAt.AsTime(),
447+
}).Return(database.AIBridgeInterception{
448+
ID:interceptionID,
449+
InitiatorID: uuid.UUID{2},
450+
Provider:"prov",
451+
Model:"mod",
452+
StartedAt:time.Now(),
453+
EndedAt: sql.NullTime{Time:req.EndedAt.AsTime(),Valid:true},
454+
},nil)
455+
},
456+
},
457+
{
458+
name:"bad_uuid_error",
459+
request:&proto.RecordInterceptionEndedRequest{
460+
Id:"this-is-not-uuid",
461+
},
462+
setupMocks:func(t*testing.T,db*dbmock.MockStore,req*proto.RecordInterceptionEndedRequest) {},
463+
expectedErr:"invalid interception ID",
464+
},
465+
{
466+
name:"database_error",
467+
request:&proto.RecordInterceptionEndedRequest{
468+
Id: uuid.UUID{1}.String(),
469+
EndedAt:timestamppb.Now(),
470+
},
471+
setupMocks:func(t*testing.T,db*dbmock.MockStore,req*proto.RecordInterceptionEndedRequest) {
472+
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(),gomock.Any()).Return(database.AIBridgeInterception{},sql.ErrConnDone)
473+
},
474+
expectedErr:"end interception: "+sql.ErrConnDone.Error(),
475+
},
476+
},
477+
)
478+
}
479+
426480
funcTestRecordTokenUsage(t*testing.T) {
427481
t.Parallel()
428482

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -476,7 +476,7 @@ require (
476476
github.com/anthropics/anthropic-sdk-gov1.13.0
477477
github.com/brianvoe/gofakeit/v7v7.8.0
478478
github.com/coder/agentapi-sdk-gov0.0.0-20250505131810-560d1d88d225
479-
github.com/coder/aibridgev0.1.4
479+
github.com/coder/aibridgev0.1.5
480480
github.com/coder/aisdk-gov0.0.9
481481
github.com/coder/boundaryv1.0.1-0.20250925154134-55a44f2a7945
482482
github.com/coder/previewv1.0.4

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -915,8 +915,8 @@ github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 h1:aQ3y1lwWyqYPiWZThqv
915915
github.com/cncf/xds/gov0.0.0-20250501225837-2ac532fd4443/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8=
916916
github.com/coder/agentapi-sdk-gov0.0.0-20250505131810-560d1d88d225 h1:tRIViZ5JRmzdOEo5wUWngaGEFBG8OaE1o2GIHN5ujJ8=
917917
github.com/coder/agentapi-sdk-gov0.0.0-20250505131810-560d1d88d225/go.mod h1:rNLVpYgEVeu1Zk29K64z6Od8RBP9DwqCu9OfCzh8MR4=
918-
github.com/coder/aibridgev0.1.4 h1:MCbrq33RCrk6v16ZbQnabfUVaCAOmJR4mPwc+UvagQs=
919-
github.com/coder/aibridgev0.1.4/go.mod h1:Q5MCfKMcKYmYl4qH1Zd0rltmPaUBPKFvIPs2k9q6qeY=
918+
github.com/coder/aibridgev0.1.5 h1:uSrltfLZWF2qOaq9RDzJW/26Ow1wMFwcwObBM0WikME=
919+
github.com/coder/aibridgev0.1.5/go.mod h1:Q5MCfKMcKYmYl4qH1Zd0rltmPaUBPKFvIPs2k9q6qeY=
920920
github.com/coder/aisdk-gov0.0.9 h1:Vzo/k2qwVGLTR10ESDeP2Ecek1SdPfZlEjtTfMveiVo=
921921
github.com/coder/aisdk-gov0.0.9/go.mod h1:KF6/Vkono0FJJOtWtveh5j7yfNrSctVTpwgweYWSp5M=
922922
github.com/coder/boundaryv1.0.1-0.20250925154134-55a44f2a7945 h1:hDUf02kTX8EGR3+5B+v5KdYvORs4YNfDPci0zCs+pC0=

0 commit comments

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