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

feat: add RecordInterceptionEnded rpc#20494

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to ourterms of service andprivacy statement. We’ll occasionally send you account related emails.

Already on GitHub?Sign in to your account

Merged
pawbana merged 3 commits intomainfrompb/aibridge-intc-ended-at-rpc
Oct 27, 2025
Merged
Show file tree
Hide file tree
Changes fromall commits
Commits
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletionsenterprise/x/aibridged/aibridged_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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) {
Expand Down
15 changes: 15 additions & 0 deletionsenterprise/x/aibridged/aibridgedmock/clientmock.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

724 changes: 431 additions & 293 deletionsenterprise/x/aibridged/proto/aibridged.pb.go
View file
Open in desktop

Large diffs are not rendered by default.

8 changes: 8 additions & 0 deletionsenterprise/x/aibridged/proto/aibridged.proto
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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);
Expand DownExpand Up@@ -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.
Expand Down
46 changes: 43 additions & 3 deletionsenterprise/x/aibridged/proto/aibridged_drpc.pb.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more abouthow customized files appear on GitHub.

8 changes: 8 additions & 0 deletionsenterprise/x/aibridged/translator.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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,
Expand Down
21 changes: 21 additions & 0 deletionsenterprise/x/aibridgedserver/aibridgedserver.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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)
Expand DownExpand Up@@ -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)
Expand Down
54 changes: 54 additions & 0 deletionsenterprise/x/aibridgedserver/aibridgedserver_test.go
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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()

Expand Down
2 changes: 1 addition & 1 deletiongo.mod
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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
Expand Down
4 changes: 2 additions & 2 deletionsgo.sum
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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=
Expand Down
Loading
[8]ページ先頭
©2009-2025 Movatter.jp