opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add RecordInterceptionEnded rpc (#20494)

Browse Source 
Adds RPC that marks interception as completed.
Added to aibridge in https://github.com/coder/aibridge/pull/43

fixes https://github.com/coder/internal/issues/1051
This commit is contained in:
Paweł Banaszewski
2025-10-27 16:38:00 +01:00
committed by GitHub
parent ffe22a0ffc
commit d0fb4599f0
10 changed files with 584 additions and 299 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/x/aibridged/aibridged_test.go
+1
View File
@@ -305,6 +305,7 @@ func TestRouting(t *testing.T) {
interceptionID = in.GetId()
return &proto.RecordInterceptionResponse{}, nil
})
client.EXPECT().RecordInterceptionEnded(gomock.Any(), gomock.Any()).Times(tc.expectedHits)
// Given: aibridged is started.
srv, err := aibridged.New(t.Context(), pool, func(ctx context.Context) (aibridged.DRPCClient, error) {
enterprise/x/aibridged/aibridgedmock/clientmock.go
+15
View File
@@ -116,6 +116,21 @@ func (mr *MockDRPCClientMockRecorder) RecordInterception(ctx, in any) *gomock.Ca
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordInterception", reflect.TypeOf((*MockDRPCClient)(nil).RecordInterception), ctx, in)
}
// RecordInterceptionEnded mocks base method.
func (m *MockDRPCClient) RecordInterceptionEnded(ctx context.Context, in *proto.RecordInterceptionEndedRequest) (*proto.RecordInterceptionEndedResponse, error) {
m.ctrl.T.Helper()
ret := m.ctrl.Call(m, "RecordInterceptionEnded", ctx, in)
ret0, _ := ret[0].(*proto.RecordInterceptionEndedResponse)
ret1, _ := ret[1].(error)
return ret0, ret1
}
// RecordInterceptionEnded indicates an expected call of RecordInterceptionEnded.
func (mr *MockDRPCClientMockRecorder) RecordInterceptionEnded(ctx, in any) *gomock.Call {
mr.mock.ctrl.T.Helper()
return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "RecordInterceptionEnded", reflect.TypeOf((*MockDRPCClient)(nil).RecordInterceptionEnded), ctx, in)
}
// RecordPromptUsage mocks base method.
func (m *MockDRPCClient) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUsageRequest) (*proto.RecordPromptUsageResponse, error) {
m.ctrl.T.Helper()
enterprise/x/aibridged/proto/aibridged.pb.go
+431 -293
View File
File diff suppressed because it is too large Load Diff
enterprise/x/aibridged/proto/aibridged.proto
+8
View File
@@ -11,6 +11,7 @@ service Recorder {
// RecordInterception creates a new interception record to which all other sub-resources
// (token, prompt, tool uses) will be related.
rpc RecordInterception(RecordInterceptionRequest) returns (RecordInterceptionResponse);
rpc RecordInterceptionEnded(RecordInterceptionEndedRequest) returns (RecordInterceptionEndedResponse);
rpc RecordTokenUsage(RecordTokenUsageRequest) returns (RecordTokenUsageResponse);
rpc RecordPromptUsage(RecordPromptUsageRequest) returns (RecordPromptUsageResponse);
rpc RecordToolUsage(RecordToolUsageRequest) returns (RecordToolUsageResponse);
@@ -45,6 +46,13 @@ message RecordInterceptionRequest {
message RecordInterceptionResponse {}
message RecordInterceptionEndedRequest {
string id = 1; // UUID.
google.protobuf.Timestamp ended_at = 2;
}
message RecordInterceptionEndedResponse {}
message RecordTokenUsageRequest {
string interception_id = 1; // UUID.
string msg_id = 2; // ID provided by provider.
enterprise/x/aibridged/proto/aibridged_drpc.pb.go
+43 -3
View File
@@ -39,6 +39,7 @@ type DRPCRecorderClient interface {
DRPCConn() drpc.Conn
RecordInterception(ctx context.Context, in *RecordInterceptionRequest) (*RecordInterceptionResponse, error)
RecordInterceptionEnded(ctx context.Context, in *RecordInterceptionEndedRequest) (*RecordInterceptionEndedResponse, error)
RecordTokenUsage(ctx context.Context, in *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error)
RecordPromptUsage(ctx context.Context, in *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error)
RecordToolUsage(ctx context.Context, in *RecordToolUsageRequest) (*RecordToolUsageResponse, error)
@@ -63,6 +64,15 @@ func (c *drpcRecorderClient) RecordInterception(ctx context.Context, in *RecordI
return out, nil
}
func (c *drpcRecorderClient) RecordInterceptionEnded(ctx context.Context, in *RecordInterceptionEndedRequest) (*RecordInterceptionEndedResponse, error) {
out := new(RecordInterceptionEndedResponse)
err := c.cc.Invoke(ctx, "/proto.Recorder/RecordInterceptionEnded", drpcEncoding_File_enterprise_x_aibridged_proto_aibridged_proto{}, in, out)
if err != nil {
return nil, err
}
return out, nil
}
func (c *drpcRecorderClient) RecordTokenUsage(ctx context.Context, in *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) {
out := new(RecordTokenUsageResponse)
err := c.cc.Invoke(ctx, "/proto.Recorder/RecordTokenUsage", drpcEncoding_File_enterprise_x_aibridged_proto_aibridged_proto{}, in, out)
@@ -92,6 +102,7 @@ func (c *drpcRecorderClient) RecordToolUsage(ctx context.Context, in *RecordTool
type DRPCRecorderServer interface {
RecordInterception(context.Context, *RecordInterceptionRequest) (*RecordInterceptionResponse, error)
RecordInterceptionEnded(context.Context, *RecordInterceptionEndedRequest) (*RecordInterceptionEndedResponse, error)
RecordTokenUsage(context.Context, *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error)
RecordPromptUsage(context.Context, *RecordPromptUsageRequest) (*RecordPromptUsageResponse, error)
RecordToolUsage(context.Context, *RecordToolUsageRequest) (*RecordToolUsageResponse, error)
@@ -103,6 +114,10 @@ func (s *DRPCRecorderUnimplementedServer) RecordInterception(context.Context, *R
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
func (s *DRPCRecorderUnimplementedServer) RecordInterceptionEnded(context.Context, *RecordInterceptionEndedRequest) (*RecordInterceptionEndedResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
func (s *DRPCRecorderUnimplementedServer) RecordTokenUsage(context.Context, *RecordTokenUsageRequest) (*RecordTokenUsageResponse, error) {
return nil, drpcerr.WithCode(errors.New("Unimplemented"), drpcerr.Unimplemented)
}
@@ -117,7 +132,7 @@ func (s *DRPCRecorderUnimplementedServer) RecordToolUsage(context.Context, *Reco
type DRPCRecorderDescription struct{}
func (DRPCRecorderDescription) NumMethods() int { return 4 }
func (DRPCRecorderDescription) NumMethods() int { return 5 }
func (DRPCRecorderDescription) Method(n int) (string, drpc.Encoding, drpc.Receiver, interface{}, bool) {
switch n {
@@ -131,6 +146,15 @@ func (DRPCRecorderDescription) Method(n int) (string, drpc.Encoding, drpc.Receiv
)
}, DRPCRecorderServer.RecordInterception, true
case 1:
return "/proto.Recorder/RecordInterceptionEnded", drpcEncoding_File_enterprise_x_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
RecordInterceptionEnded(
ctx,
in1.(*RecordInterceptionEndedRequest),
)
}, DRPCRecorderServer.RecordInterceptionEnded, true
case 2:
return "/proto.Recorder/RecordTokenUsage", drpcEncoding_File_enterprise_x_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
@@ -139,7 +163,7 @@ func (DRPCRecorderDescription) Method(n int) (string, drpc.Encoding, drpc.Receiv
in1.(*RecordTokenUsageRequest),
)
}, DRPCRecorderServer.RecordTokenUsage, true
case 2:
case 3:
return "/proto.Recorder/RecordPromptUsage", drpcEncoding_File_enterprise_x_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
@@ -148,7 +172,7 @@ func (DRPCRecorderDescription) Method(n int) (string, drpc.Encoding, drpc.Receiv
in1.(*RecordPromptUsageRequest),
)
}, DRPCRecorderServer.RecordPromptUsage, true
case 3:
case 4:
return "/proto.Recorder/RecordToolUsage", drpcEncoding_File_enterprise_x_aibridged_proto_aibridged_proto{},
func(srv interface{}, ctx context.Context, in1, in2 interface{}) (drpc.Message, error) {
return srv.(DRPCRecorderServer).
@@ -182,6 +206,22 @@ func (x *drpcRecorder_RecordInterceptionStream) SendAndClose(m *RecordIntercepti
return x.CloseSend()
}
type DRPCRecorder_RecordInterceptionEndedStream interface {
drpc.Stream
SendAndClose(*RecordInterceptionEndedResponse) error
}
type drpcRecorder_RecordInterceptionEndedStream struct {
drpc.Stream
}
func (x *drpcRecorder_RecordInterceptionEndedStream) SendAndClose(m *RecordInterceptionEndedResponse) error {
if err := x.MsgSend(m, drpcEncoding_File_enterprise_x_aibridged_proto_aibridged_proto{}); err != nil {
return err
}
return x.CloseSend()
}
type DRPCRecorder_RecordTokenUsageStream interface {
drpc.Stream
SendAndClose(*RecordTokenUsageResponse) error
enterprise/x/aibridged/translator.go
+8
View File
@@ -35,6 +35,14 @@ func (t *recorderTranslation) RecordInterception(ctx context.Context, req *aibri
return err
}
func (t *recorderTranslation) RecordInterceptionEnded(ctx context.Context, req *aibridge.InterceptionRecordEnded) error {
_, err := t.client.RecordInterceptionEnded(ctx, &proto.RecordInterceptionEndedRequest{
Id: req.ID,
EndedAt: timestamppb.New(req.EndedAt),
})
return err
}
func (t *recorderTranslation) RecordPromptUsage(ctx context.Context, req *aibridge.PromptUsageRecord) error {
_, err := t.client.RecordPromptUsage(ctx, &proto.RecordPromptUsageRequest{
InterceptionId: req.InterceptionID,
enterprise/x/aibridgedserver/aibridgedserver.go
+21
View File
@@ -55,6 +55,7 @@ type store interface {
InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) (database.AIBridgeTokenUsage, error)
InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) (database.AIBridgeUserPrompt, error)
InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) (database.AIBridgeToolUsage, error)
UpdateAIBridgeInterceptionEnded(ctx context.Context, intcID database.UpdateAIBridgeInterceptionEndedParams) (database.AIBridgeInterception, error)
// MCPConfigurator-related queries.
GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error)
@@ -129,6 +130,26 @@ func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterce
return &proto.RecordInterceptionResponse{}, nil
}
func (s *Server) RecordInterceptionEnded(ctx context.Context, in *proto.RecordInterceptionEndedRequest) (*proto.RecordInterceptionEndedResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
intcID, err := uuid.Parse(in.GetId())
if err != nil {
return nil, xerrors.Errorf("invalid interception ID %q: %w", in.GetId(), err)
}
_, err = s.store.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{
ID: intcID,
EndedAt: in.EndedAt.AsTime(),
})
if err != nil {
return nil, xerrors.Errorf("end interception: %w", err)
}
return &proto.RecordInterceptionEndedResponse{}, nil
}
func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsageRequest) (*proto.RecordTokenUsageResponse, error) {
//nolint:gocritic // AIBridged has specific authz rules.
ctx = dbauthz.AsAIBridged(ctx)
enterprise/x/aibridgedserver/aibridgedserver_test.go
+54
View File
@@ -423,6 +423,60 @@ func TestRecordInterception(t *testing.T) {
)
}
func TestRecordInterceptionEnded(t *testing.T) {
t.Parallel()
testRecordMethod(t,
func(srv *aibridgedserver.Server, ctx context.Context, req *proto.RecordInterceptionEndedRequest) (*proto.RecordInterceptionEndedResponse, error) {
return srv.RecordInterceptionEnded(ctx, req)
},
[]testRecordMethodCase[*proto.RecordInterceptionEndedRequest]{
{
name: "ok",
request: &proto.RecordInterceptionEndedRequest{
Id: uuid.UUID{1}.String(),
EndedAt: timestamppb.Now(),
},
setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {
interceptionID, err := uuid.Parse(req.GetId())
assert.NoError(t, err, "parse interception UUID")
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), database.UpdateAIBridgeInterceptionEndedParams{
ID: interceptionID,
EndedAt: req.EndedAt.AsTime(),
}).Return(database.AIBridgeInterception{
ID: interceptionID,
InitiatorID: uuid.UUID{2},
Provider: "prov",
Model: "mod",
StartedAt: time.Now(),
EndedAt: sql.NullTime{Time: req.EndedAt.AsTime(), Valid: true},
}, nil)
},
},
{
name: "bad_uuid_error",
request: &proto.RecordInterceptionEndedRequest{
Id: "this-is-not-uuid",
},
setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {},
expectedErr: "invalid interception ID",
},
{
name: "database_error",
request: &proto.RecordInterceptionEndedRequest{
Id: uuid.UUID{1}.String(),
EndedAt: timestamppb.Now(),
},
setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionEndedRequest) {
db.EXPECT().UpdateAIBridgeInterceptionEnded(gomock.Any(), gomock.Any()).Return(database.AIBridgeInterception{}, sql.ErrConnDone)
},
expectedErr: "end interception: " + sql.ErrConnDone.Error(),
},
},
)
}
func TestRecordTokenUsage(t *testing.T) {
t.Parallel()
go.mod
+1 -1
View File
@@ -476,7 +476,7 @@ require (
github.com/anthropics/anthropic-sdk-go v1.13.0
github.com/brianvoe/gofakeit/v7 v7.8.0
github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225
github.com/coder/aibridge v0.1.4
github.com/coder/aibridge v0.1.5
github.com/coder/aisdk-go v0.0.9
github.com/coder/boundary v1.0.1-0.20250925154134-55a44f2a7945
github.com/coder/preview v1.0.4
go.sum
+2 -2
View File
@@ -915,8 +915,8 @@ github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 h1:aQ3y1lwWyqYPiWZThqv
github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8=
github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225 h1:tRIViZ5JRmzdOEo5wUWngaGEFBG8OaE1o2GIHN5ujJ8=
github.com/coder/agentapi-sdk-go v0.0.0-20250505131810-560d1d88d225/go.mod h1:rNLVpYgEVeu1Zk29K64z6Od8RBP9DwqCu9OfCzh8MR4=
github.com/coder/aibridge v0.1.4 h1:MCbrq33RCrk6v16ZbQnabfUVaCAOmJR4mPwc+UvagQs=
github.com/coder/aibridge v0.1.4/go.mod h1:Q5MCfKMcKYmYl4qH1Zd0rltmPaUBPKFvIPs2k9q6qeY=
github.com/coder/aibridge v0.1.5 h1:uSrltfLZWF2qOaq9RDzJW/26Ow1wMFwcwObBM0WikME=
github.com/coder/aibridge v0.1.5/go.mod h1:Q5MCfKMcKYmYl4qH1Zd0rltmPaUBPKFvIPs2k9q6qeY=
github.com/coder/aisdk-go v0.0.9 h1:Vzo/k2qwVGLTR10ESDeP2Ecek1SdPfZlEjtTfMveiVo=
github.com/coder/aisdk-go v0.0.9/go.mod h1:KF6/Vkono0FJJOtWtveh5j7yfNrSctVTpwgweYWSp5M=
github.com/coder/boundary v1.0.1-0.20250925154134-55a44f2a7945 h1:hDUf02kTX8EGR3+5B+v5KdYvORs4YNfDPci0zCs+pC0=