From 0a6ba5d51a02603b5d1ccf9ea3c47f9555ee2a20 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Fri, 26 Sep 2025 16:20:33 +0200 Subject: [PATCH] feat: add endpoint to list aibridge interceptions (#19929) Co-authored-by: Dean Sheather --- coderd/apidoc/docs.go | 197 ++++++++ coderd/apidoc/swagger.json | 193 +++++++ coderd/coderd.go | 57 ++- coderd/database/db2sdk/db2sdk.go | 82 +++ coderd/database/dbauthz/dbauthz.go | 57 ++- coderd/database/dbauthz/dbauthz_test.go | 92 +++- coderd/database/dbgen/dbgen.go | 65 +++ coderd/database/dbmetrics/querymetrics.go | 53 +- coderd/database/dbmock/dbmock.go | 96 +++- coderd/database/dump.sql | 6 + ...374_aibridge_interception_indices.down.sql | 5 + ...00374_aibridge_interception_indices.up.sql | 9 + coderd/database/modelqueries.go | 55 ++ coderd/database/querier.go | 10 +- coderd/database/queries.sql.go | 339 ++++++++++++- coderd/database/queries/aibridge.sql | 149 +++++- coderd/rbac/regosql/configs.go | 17 +- coderd/rbac/roles.go | 2 + coderd/searchquery/search.go | 48 +- codersdk/aibridge.go | 125 +++++ docs/reference/api/aibridge.md | 96 ++++ docs/reference/api/schemas.md | 239 +++++++++ enterprise/coderd/aibridge.go | 159 ++++++ enterprise/coderd/aibridge_test.go | 478 ++++++++++++++++++ enterprise/coderd/aibridged.go | 19 +- enterprise/coderd/coderd.go | 26 + .../x/aibridgedserver/aibridgedserver.go | 12 +- .../x/aibridgedserver/aibridgedserver_test.go | 47 +- site/src/api/typesGenerated.ts | 57 +++ 29 files changed, 2646 insertions(+), 144 deletions(-) create mode 100644 coderd/database/migrations/000374_aibridge_interception_indices.down.sql create mode 100644 coderd/database/migrations/000374_aibridge_interception_indices.up.sql create mode 100644 codersdk/aibridge.go create mode 100644 docs/reference/api/aibridge.md create mode 100644 enterprise/coderd/aibridge.go create mode 100644 enterprise/coderd/aibridge_test.go diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 21681cd19b..df881137bc 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -85,6 +85,51 @@ const docTemplate = `{ } } }, + "/api/experimental/aibridge/interceptions": { + "get": { + "security": [ + { + "CoderSessionToken": [] + } + ], + "produces": [ + "application/json" + ], + "tags": [ + "AIBridge" + ], + "summary": "List AIBridge interceptions", + "operationId": "list-aibridge-interceptions", + "parameters": [ + { + "type": "string", + "description": "Search query in the format ` + "`" + `key:value` + "`" + `. Available keys are: initiator, provider, model, started_after, started_before.", + "name": "q", + "in": "query" + }, + { + "type": "integer", + "description": "Page limit", + "name": "limit", + "in": "query" + }, + { + "type": "string", + "description": "Cursor pagination after ID", + "name": "after_id", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.AIBridgeListInterceptionsResponse" + } + } + } + } + }, "/appearance": { "get": { "security": [ @@ -11226,6 +11271,62 @@ const docTemplate = `{ } } }, + "codersdk.AIBridgeInterception": { + "type": "object", + "properties": { + "id": { + "type": "string", + "format": "uuid" + }, + "initiator_id": { + "type": "string", + "format": "uuid" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "model": { + "type": "string" + }, + "provider": { + "type": "string" + }, + "started_at": { + "type": "string", + "format": "date-time" + }, + "token_usages": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeTokenUsage" + } + }, + "tool_usages": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeToolUsage" + } + }, + "user_prompts": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeUserPrompt" + } + } + } + }, + "codersdk.AIBridgeListInterceptionsResponse": { + "type": "object", + "properties": { + "results": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeInterception" + } + } + } + }, "codersdk.AIBridgeOpenAIConfig": { "type": "object", "properties": { @@ -11237,6 +11338,102 @@ const docTemplate = `{ } } }, + "codersdk.AIBridgeTokenUsage": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "id": { + "type": "string", + "format": "uuid" + }, + "input_tokens": { + "type": "integer" + }, + "interception_id": { + "type": "string", + "format": "uuid" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "output_tokens": { + "type": "integer" + }, + "provider_response_id": { + "type": "string" + } + } + }, + "codersdk.AIBridgeToolUsage": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "id": { + "type": "string", + "format": "uuid" + }, + "injected": { + "type": "boolean" + }, + "input": { + "type": "string" + }, + "interception_id": { + "type": "string", + "format": "uuid" + }, + "invocation_error": { + "type": "string" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "provider_response_id": { + "type": "string" + }, + "server_url": { + "type": "string" + }, + "tool": { + "type": "string" + } + } + }, + "codersdk.AIBridgeUserPrompt": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "id": { + "type": "string", + "format": "uuid" + }, + "interception_id": { + "type": "string", + "format": "uuid" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "prompt": { + "type": "string" + }, + "provider_response_id": { + "type": "string" + } + } + }, "codersdk.AIConfig": { "type": "object", "properties": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 203d7533bc..5fa20c5867 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -65,6 +65,47 @@ } } }, + "/api/experimental/aibridge/interceptions": { + "get": { + "security": [ + { + "CoderSessionToken": [] + } + ], + "produces": ["application/json"], + "tags": ["AIBridge"], + "summary": "List AIBridge interceptions", + "operationId": "list-aibridge-interceptions", + "parameters": [ + { + "type": "string", + "description": "Search query in the format `key:value`. Available keys are: initiator, provider, model, started_after, started_before.", + "name": "q", + "in": "query" + }, + { + "type": "integer", + "description": "Page limit", + "name": "limit", + "in": "query" + }, + { + "type": "string", + "description": "Cursor pagination after ID", + "name": "after_id", + "in": "query" + } + ], + "responses": { + "200": { + "description": "OK", + "schema": { + "$ref": "#/definitions/codersdk.AIBridgeListInterceptionsResponse" + } + } + } + } + }, "/appearance": { "get": { "security": [ @@ -9954,6 +9995,62 @@ } } }, + "codersdk.AIBridgeInterception": { + "type": "object", + "properties": { + "id": { + "type": "string", + "format": "uuid" + }, + "initiator_id": { + "type": "string", + "format": "uuid" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "model": { + "type": "string" + }, + "provider": { + "type": "string" + }, + "started_at": { + "type": "string", + "format": "date-time" + }, + "token_usages": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeTokenUsage" + } + }, + "tool_usages": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeToolUsage" + } + }, + "user_prompts": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeUserPrompt" + } + } + } + }, + "codersdk.AIBridgeListInterceptionsResponse": { + "type": "object", + "properties": { + "results": { + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeInterception" + } + } + } + }, "codersdk.AIBridgeOpenAIConfig": { "type": "object", "properties": { @@ -9965,6 +10062,102 @@ } } }, + "codersdk.AIBridgeTokenUsage": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "id": { + "type": "string", + "format": "uuid" + }, + "input_tokens": { + "type": "integer" + }, + "interception_id": { + "type": "string", + "format": "uuid" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "output_tokens": { + "type": "integer" + }, + "provider_response_id": { + "type": "string" + } + } + }, + "codersdk.AIBridgeToolUsage": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "id": { + "type": "string", + "format": "uuid" + }, + "injected": { + "type": "boolean" + }, + "input": { + "type": "string" + }, + "interception_id": { + "type": "string", + "format": "uuid" + }, + "invocation_error": { + "type": "string" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "provider_response_id": { + "type": "string" + }, + "server_url": { + "type": "string" + }, + "tool": { + "type": "string" + } + } + }, + "codersdk.AIBridgeUserPrompt": { + "type": "object", + "properties": { + "created_at": { + "type": "string", + "format": "date-time" + }, + "id": { + "type": "string", + "format": "uuid" + }, + "interception_id": { + "type": "string", + "format": "uuid" + }, + "metadata": { + "type": "object", + "additionalProperties": {} + }, + "prompt": { + "type": "string" + }, + "provider_response_id": { + "type": "string" + } + } + }, "codersdk.AIConfig": { "type": "object", "properties": { diff --git a/coderd/coderd.go b/coderd/coderd.go index 1724c369ec..15f2767c22 100644 --- a/coderd/coderd.go +++ b/coderd/coderd.go @@ -1000,38 +1000,41 @@ func New(options *Options) *API { // Experimental routes are not guaranteed to be stable and may change at any time. r.Route("/api/experimental", func(r chi.Router) { + api.ExperimentalHandler = r + r.NotFound(func(rw http.ResponseWriter, _ *http.Request) { httpapi.RouteNotFound(rw) }) - - // Only this group should be subject to apiKeyMiddleware; aibridged will mount its own - // router and handles key validation in a different fashion. - // See enterprise/x/aibridged/http.go. - r.Group(func(r chi.Router) { + r.Use( + // Specific routes can specify different limits, but every rate + // limit must be configurable by the admin. + apiRateLimiter, + httpmw.ReportCLITelemetry(api.Logger, options.Telemetry), + ) + r.Route("/aitasks", func(r chi.Router) { + r.Use(apiKeyMiddleware) + r.Get("/prompts", api.aiTasksPrompts) + }) + r.Route("/tasks", func(r chi.Router) { r.Use(apiKeyMiddleware) - r.Route("/aitasks", func(r chi.Router) { - r.Get("/prompts", api.aiTasksPrompts) - }) - r.Route("/tasks", func(r chi.Router) { - r.Use(apiRateLimiter) - r.Get("/", api.tasksList) + r.Get("/", api.tasksList) - r.Route("/{user}", func(r chi.Router) { - r.Use(httpmw.ExtractOrganizationMembersParam(options.Database, api.HTTPAuth.Authorize)) - r.Get("/{id}", api.taskGet) - r.Delete("/{id}", api.taskDelete) - r.Post("/{id}/send", api.taskSend) - r.Get("/{id}/logs", api.taskLogs) - r.Post("/", api.tasksCreate) - }) - }) - r.Route("/mcp", func(r chi.Router) { - r.Use( - httpmw.RequireExperimentWithDevBypass(api.Experiments, codersdk.ExperimentOAuth2, codersdk.ExperimentMCPServerHTTP), - ) - // MCP HTTP transport endpoint with mandatory authentication - r.Mount("/http", api.mcpHTTPHandler()) + r.Route("/{user}", func(r chi.Router) { + r.Use(httpmw.ExtractOrganizationMembersParam(options.Database, api.HTTPAuth.Authorize)) + r.Get("/{id}", api.taskGet) + r.Delete("/{id}", api.taskDelete) + r.Post("/{id}/send", api.taskSend) + r.Get("/{id}/logs", api.taskLogs) + r.Post("/", api.tasksCreate) }) }) + r.Route("/mcp", func(r chi.Router) { + r.Use( + apiKeyMiddleware, + httpmw.RequireExperimentWithDevBypass(api.Experiments, codersdk.ExperimentOAuth2, codersdk.ExperimentMCPServerHTTP), + ) + // MCP HTTP transport endpoint with mandatory authentication + r.Mount("/http", api.mcpHTTPHandler()) + }) }) r.Route("/api/v2", func(r chi.Router) { @@ -1727,6 +1730,8 @@ type API struct { // APIHandler serves "/api/v2" APIHandler chi.Router + // ExperimentalHandler serves "/api/experimental" + ExperimentalHandler chi.Router // RootHandler serves "/" RootHandler chi.Router diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 97a7376a2e..eaa79c152f 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -13,6 +13,7 @@ import ( "github.com/google/uuid" "github.com/hashicorp/hcl/v2" + "github.com/sqlc-dev/pqtype" "golang.org/x/xerrors" "tailscale.com/tailcfg" @@ -918,3 +919,84 @@ func PreviewParameterValidation(v *previewtypes.ParameterValidation) codersdk.Pr Monotonic: v.Monotonic, } } + +func AIBridgeInterception(interception database.AIBridgeInterception, tokenUsages []database.AIBridgeTokenUsage, userPrompts []database.AIBridgeUserPrompt, toolUsages []database.AIBridgeToolUsage) codersdk.AIBridgeInterception { + sdkTokenUsages := List(tokenUsages, AIBridgeTokenUsage) + sort.Slice(sdkTokenUsages, func(i, j int) bool { + // created_at ASC + return sdkTokenUsages[i].CreatedAt.Before(sdkTokenUsages[j].CreatedAt) + }) + sdkUserPrompts := List(userPrompts, AIBridgeUserPrompt) + sort.Slice(sdkUserPrompts, func(i, j int) bool { + // created_at ASC + return sdkUserPrompts[i].CreatedAt.Before(sdkUserPrompts[j].CreatedAt) + }) + sdkToolUsages := List(toolUsages, AIBridgeToolUsage) + sort.Slice(sdkToolUsages, func(i, j int) bool { + // created_at ASC + return sdkToolUsages[i].CreatedAt.Before(sdkToolUsages[j].CreatedAt) + }) + return codersdk.AIBridgeInterception{ + ID: interception.ID, + InitiatorID: interception.InitiatorID, + Provider: interception.Provider, + Model: interception.Model, + Metadata: jsonOrEmptyMap(interception.Metadata), + StartedAt: interception.StartedAt, + TokenUsages: sdkTokenUsages, + UserPrompts: sdkUserPrompts, + ToolUsages: sdkToolUsages, + } +} + +func AIBridgeTokenUsage(usage database.AIBridgeTokenUsage) codersdk.AIBridgeTokenUsage { + return codersdk.AIBridgeTokenUsage{ + ID: usage.ID, + InterceptionID: usage.InterceptionID, + ProviderResponseID: usage.ProviderResponseID, + InputTokens: usage.InputTokens, + OutputTokens: usage.OutputTokens, + Metadata: jsonOrEmptyMap(usage.Metadata), + CreatedAt: usage.CreatedAt, + } +} + +func AIBridgeUserPrompt(prompt database.AIBridgeUserPrompt) codersdk.AIBridgeUserPrompt { + return codersdk.AIBridgeUserPrompt{ + ID: prompt.ID, + InterceptionID: prompt.InterceptionID, + ProviderResponseID: prompt.ProviderResponseID, + Prompt: prompt.Prompt, + Metadata: jsonOrEmptyMap(prompt.Metadata), + CreatedAt: prompt.CreatedAt, + } +} + +func AIBridgeToolUsage(usage database.AIBridgeToolUsage) codersdk.AIBridgeToolUsage { + return codersdk.AIBridgeToolUsage{ + ID: usage.ID, + InterceptionID: usage.InterceptionID, + ProviderResponseID: usage.ProviderResponseID, + ServerURL: usage.ServerUrl.String, + Tool: usage.Tool, + Input: usage.Input, + Injected: usage.Injected, + InvocationError: usage.InvocationError.String, + Metadata: jsonOrEmptyMap(usage.Metadata), + CreatedAt: usage.CreatedAt, + } +} + +func jsonOrEmptyMap(rawMessage pqtype.NullRawMessage) map[string]any { + var m map[string]any + if !rawMessage.Valid { + return m + } + + err := json.Unmarshal(rawMessage.RawMessage, &m) + if err != nil { + // Don't reuse m + return map[string]any{} + } + return m +} diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 6704c82118..2490ea2323 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -3842,26 +3842,26 @@ func (q *querier) InsertAIBridgeInterception(ctx context.Context, arg database.I return insert(q.log, q.auth, rbac.ResourceAibridgeInterception.WithOwner(arg.InitiatorID.String()), q.db.InsertAIBridgeInterception)(ctx, arg) } -func (q *querier) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error { +func (q *querier) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) (database.AIBridgeTokenUsage, error) { // All aibridge_token_usages records belong to the initiator of their associated interception. if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil { - return err + return database.AIBridgeTokenUsage{}, err } return q.db.InsertAIBridgeTokenUsage(ctx, arg) } -func (q *querier) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error { +func (q *querier) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) (database.AIBridgeToolUsage, error) { // All aibridge_tool_usages records belong to the initiator of their associated interception. if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil { - return err + return database.AIBridgeToolUsage{}, err } return q.db.InsertAIBridgeToolUsage(ctx, arg) } -func (q *querier) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error { +func (q *querier) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) (database.AIBridgeUserPrompt, error) { // All aibridge_user_prompts records belong to the initiator of their associated interception. if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionUpdate, arg.InterceptionID); err != nil { - return err + return database.AIBridgeUserPrompt{}, err } return q.db.InsertAIBridgeUserPrompt(ctx, arg) } @@ -4409,6 +4409,44 @@ func (q *querier) InsertWorkspaceResourceMetadata(ctx context.Context, arg datab return q.db.InsertWorkspaceResourceMetadata(ctx, arg) } +func (q *querier) ListAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams) ([]database.AIBridgeInterception, error) { + prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceAibridgeInterception.Type) + if err != nil { + return nil, xerrors.Errorf("(dev error) prepare sql filter: %w", err) + } + return q.db.ListAuthorizedAIBridgeInterceptions(ctx, arg, prep) +} + +func (q *querier) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Context, interceptionIDs []uuid.UUID) ([]database.AIBridgeTokenUsage, error) { + // This function is a system function until we implement a join for aibridge interceptions. + // Matches the behavior of the workspaces listing endpoint. + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil { + return nil, err + } + + return q.db.ListAIBridgeTokenUsagesByInterceptionIDs(ctx, interceptionIDs) +} + +func (q *querier) ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context, interceptionIDs []uuid.UUID) ([]database.AIBridgeToolUsage, error) { + // This function is a system function until we implement a join for aibridge interceptions. + // Matches the behavior of the workspaces listing endpoint. + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil { + return nil, err + } + + return q.db.ListAIBridgeToolUsagesByInterceptionIDs(ctx, interceptionIDs) +} + +func (q *querier) ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context, interceptionIDs []uuid.UUID) ([]database.AIBridgeUserPrompt, error) { + // This function is a system function until we implement a join for aibridge interceptions. + // Matches the behavior of the workspaces listing endpoint. + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceSystem); err != nil { + return nil, err + } + + return q.db.ListAIBridgeUserPromptsByInterceptionIDs(ctx, interceptionIDs) +} + func (q *querier) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) { return fetchWithPostFilter(q.auth, policy.ActionRead, q.db.ListProvisionerKeysByOrganization)(ctx, organizationID) } @@ -5761,3 +5799,10 @@ func (q *querier) GetAuthorizedConnectionLogsOffset(ctx context.Context, arg dat func (q *querier) CountAuthorizedConnectionLogs(ctx context.Context, arg database.CountConnectionLogsParams, _ rbac.PreparedAuthorized) (int64, error) { return q.CountConnectionLogs(ctx, arg) } + +func (q *querier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams, _ rbac.PreparedAuthorized) ([]database.AIBridgeInterception, error) { + // TODO: Delete this function, all ListAIBridgeInterceptions should be authorized. For now just call ListAIBridgeInterceptions on the authz querier. + // This cannot be deleted for now because it's included in the + // database.Store interface, so dbauthz needs to implement it. + return q.ListAIBridgeInterceptions(ctx, arg) +} diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index e94e7275d9..4b96c4e3eb 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -4334,13 +4334,6 @@ func TestInsertAPIKey_AsPrebuildsUser(t *testing.T) { } func (s *MethodTestSuite) TestAIBridge() { - s.Run("GetAIBridgeInterceptionByID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - intID := uuid.UUID{2} - intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) - db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() - check.Args(intID).Asserts(intc, policy.ActionRead).Returns(intc) - })) - s.Run("InsertAIBridgeInterception", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { initID := uuid.UUID{3} user := testutil.Fake(s.T(), faker, database.User{ID: initID}) @@ -4360,30 +4353,43 @@ func (s *MethodTestSuite) TestAIBridge() { s.Run("InsertAIBridgeTokenUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { intID := uuid.UUID{2} intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) - params := database.InsertAIBridgeTokenUsageParams{InterceptionID: intc.ID} db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation. - db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), params).Return(nil).AnyTimes() - check.Args(params).Asserts(intc, policy.ActionUpdate) - })) - s.Run("InsertAIBridgeToolUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - intID := uuid.UUID{2} - intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) - params := database.InsertAIBridgeToolUsageParams{InterceptionID: intc.ID} - db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation. - db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), params).Return(nil).AnyTimes() + params := database.InsertAIBridgeTokenUsageParams{InterceptionID: intc.ID} + expected := testutil.Fake(s.T(), faker, database.AIBridgeTokenUsage{InterceptionID: intc.ID}) + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), params).Return(expected, nil).AnyTimes() check.Args(params).Asserts(intc, policy.ActionUpdate) })) s.Run("InsertAIBridgeUserPrompt", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { intID := uuid.UUID{2} intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) - params := database.InsertAIBridgeUserPromptParams{InterceptionID: intc.ID} db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation. - db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), params).Return(nil).AnyTimes() + + params := database.InsertAIBridgeUserPromptParams{InterceptionID: intc.ID} + expected := testutil.Fake(s.T(), faker, database.AIBridgeUserPrompt{InterceptionID: intc.ID}) + db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), params).Return(expected, nil).AnyTimes() check.Args(params).Asserts(intc, policy.ActionUpdate) })) + s.Run("InsertAIBridgeToolUsage", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + intID := uuid.UUID{2} + intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) + db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation. + + params := database.InsertAIBridgeToolUsageParams{InterceptionID: intc.ID} + expected := testutil.Fake(s.T(), faker, database.AIBridgeToolUsage{InterceptionID: intc.ID}) + db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), params).Return(expected, nil).AnyTimes() + check.Args(params).Asserts(intc, policy.ActionUpdate) + })) + + s.Run("GetAIBridgeInterceptionByID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + intID := uuid.UUID{2} + intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) + db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() + check.Args(intID).Asserts(intc, policy.ActionRead).Returns(intc) + })) + s.Run("GetAIBridgeInterceptions", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { a := testutil.Fake(s.T(), faker, database.AIBridgeInterception{}) b := testutil.Fake(s.T(), faker, database.AIBridgeInterception{}) @@ -4401,6 +4407,16 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(intID).Asserts(intc, policy.ActionRead).Returns(toks) })) + s.Run("GetAIBridgeUserPromptsByInterceptionID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + intID := uuid.UUID{2} + intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) + pr := testutil.Fake(s.T(), faker, database.AIBridgeUserPrompt{InterceptionID: intID}) + prs := []database.AIBridgeUserPrompt{pr} + db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation. + db.EXPECT().GetAIBridgeUserPromptsByInterceptionID(gomock.Any(), intID).Return(prs, nil).AnyTimes() + check.Args(intID).Asserts(intc, policy.ActionRead).Returns(prs) + })) + s.Run("GetAIBridgeToolUsagesByInterceptionID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { intID := uuid.UUID{2} intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) @@ -4411,13 +4427,35 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(intID).Asserts(intc, policy.ActionRead).Returns(tools) })) - s.Run("GetAIBridgeUserPromptsByInterceptionID", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { - intID := uuid.UUID{2} - intc := testutil.Fake(s.T(), faker, database.AIBridgeInterception{ID: intID}) - pr := testutil.Fake(s.T(), faker, database.AIBridgeUserPrompt{InterceptionID: intID}) - prs := []database.AIBridgeUserPrompt{pr} - db.EXPECT().GetAIBridgeInterceptionByID(gomock.Any(), intID).Return(intc, nil).AnyTimes() // Validation. - db.EXPECT().GetAIBridgeUserPromptsByInterceptionID(gomock.Any(), intID).Return(prs, nil).AnyTimes() - check.Args(intID).Asserts(intc, policy.ActionRead).Returns(prs) + s.Run("ListAIBridgeInterceptions", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.ListAIBridgeInterceptionsParams{} + db.EXPECT().ListAuthorizedAIBridgeInterceptions(gomock.Any(), params, gomock.Any()).Return([]database.AIBridgeInterception{}, nil).AnyTimes() + // No asserts here because SQLFilter. + check.Args(params).Asserts() + })) + + s.Run("ListAuthorizedAIBridgeInterceptions", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.ListAIBridgeInterceptionsParams{} + db.EXPECT().ListAuthorizedAIBridgeInterceptions(gomock.Any(), params, gomock.Any()).Return([]database.AIBridgeInterception{}, nil).AnyTimes() + // No asserts here because SQLFilter. + check.Args(params, emptyPreparedAuthorized{}).Asserts() + })) + + s.Run("ListAIBridgeTokenUsagesByInterceptionIDs", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + ids := []uuid.UUID{{1}} + db.EXPECT().ListAIBridgeTokenUsagesByInterceptionIDs(gomock.Any(), ids).Return([]database.AIBridgeTokenUsage{}, nil).AnyTimes() + check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead).Returns([]database.AIBridgeTokenUsage{}) + })) + + s.Run("ListAIBridgeUserPromptsByInterceptionIDs", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + ids := []uuid.UUID{{1}} + db.EXPECT().ListAIBridgeUserPromptsByInterceptionIDs(gomock.Any(), ids).Return([]database.AIBridgeUserPrompt{}, nil).AnyTimes() + check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead).Returns([]database.AIBridgeUserPrompt{}) + })) + + s.Run("ListAIBridgeToolUsagesByInterceptionIDs", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + ids := []uuid.UUID{{1}} + db.EXPECT().ListAIBridgeToolUsagesByInterceptionIDs(gomock.Any(), ids).Return([]database.AIBridgeToolUsage{}, nil).AnyTimes() + check.Args(ids).Asserts(rbac.ResourceSystem, policy.ActionRead).Returns([]database.AIBridgeToolUsage{}) })) } diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index ef4248ce1c..9b5a3818fd 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -1471,6 +1471,71 @@ func ClaimPrebuild( return claimedWorkspace } +func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertAIBridgeInterceptionParams) database.AIBridgeInterception { + interception, err := db.InsertAIBridgeInterception(genCtx, database.InsertAIBridgeInterceptionParams{ + ID: takeFirst(seed.ID, uuid.New()), + InitiatorID: takeFirst(seed.InitiatorID, uuid.New()), + Provider: takeFirst(seed.Provider, "provider"), + Model: takeFirst(seed.Model, "model"), + Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")), + StartedAt: takeFirst(seed.StartedAt, dbtime.Now()), + }) + require.NoError(t, err, "insert aibridge interception") + return interception +} + +func AIBridgeTokenUsage(t testing.TB, db database.Store, seed database.InsertAIBridgeTokenUsageParams) database.AIBridgeTokenUsage { + usage, err := db.InsertAIBridgeTokenUsage(genCtx, database.InsertAIBridgeTokenUsageParams{ + ID: takeFirst(seed.ID, uuid.New()), + InterceptionID: takeFirst(seed.InterceptionID, uuid.New()), + ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"), + InputTokens: takeFirst(seed.InputTokens, 100), + OutputTokens: takeFirst(seed.OutputTokens, 100), + Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")), + CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()), + }) + require.NoError(t, err, "insert aibridge token usage") + return usage +} + +func AIBridgeUserPrompt(t testing.TB, db database.Store, seed database.InsertAIBridgeUserPromptParams) database.AIBridgeUserPrompt { + prompt, err := db.InsertAIBridgeUserPrompt(genCtx, database.InsertAIBridgeUserPromptParams{ + ID: takeFirst(seed.ID, uuid.New()), + InterceptionID: takeFirst(seed.InterceptionID, uuid.New()), + ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"), + Prompt: takeFirst(seed.Prompt, "prompt"), + Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")), + CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()), + }) + require.NoError(t, err, "insert aibridge user prompt") + return prompt +} + +func AIBridgeToolUsage(t testing.TB, db database.Store, seed database.InsertAIBridgeToolUsageParams) database.AIBridgeToolUsage { + serverURL := sql.NullString{} + if seed.ServerUrl.Valid { + serverURL = seed.ServerUrl + } + invocationError := sql.NullString{} + if seed.InvocationError.Valid { + invocationError = seed.InvocationError + } + toolUsage, err := db.InsertAIBridgeToolUsage(genCtx, database.InsertAIBridgeToolUsageParams{ + ID: takeFirst(seed.ID, uuid.New()), + InterceptionID: takeFirst(seed.InterceptionID, uuid.New()), + ProviderResponseID: takeFirst(seed.ProviderResponseID, "provider_response_id"), + Tool: takeFirst(seed.Tool, "tool"), + ServerUrl: serverURL, + Input: takeFirst(seed.Input, "input"), + Injected: takeFirst(seed.Injected, false), + InvocationError: invocationError, + Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")), + CreatedAt: takeFirst(seed.CreatedAt, dbtime.Now()), + }) + require.NoError(t, err, "insert aibridge tool usage") + return toolUsage +} + func provisionerJobTiming(t testing.TB, db database.Store, seed database.ProvisionerJobTiming) database.ProvisionerJobTiming { timing, err := db.InsertProvisionerJobTimings(genCtx, database.InsertProvisionerJobTimingsParams{ JobID: takeFirst(seed.JobID, uuid.New()), diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index d7d20f2355..f90c59cb5a 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -2210,25 +2210,25 @@ func (m queryMetricsStore) InsertAIBridgeInterception(ctx context.Context, arg d return r0, r1 } -func (m queryMetricsStore) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error { +func (m queryMetricsStore) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) (database.AIBridgeTokenUsage, error) { start := time.Now() - r0 := m.s.InsertAIBridgeTokenUsage(ctx, arg) + r0, r1 := m.s.InsertAIBridgeTokenUsage(ctx, arg) m.queryLatencies.WithLabelValues("InsertAIBridgeTokenUsage").Observe(time.Since(start).Seconds()) - return r0 + return r0, r1 } -func (m queryMetricsStore) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error { +func (m queryMetricsStore) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) (database.AIBridgeToolUsage, error) { start := time.Now() - r0 := m.s.InsertAIBridgeToolUsage(ctx, arg) + r0, r1 := m.s.InsertAIBridgeToolUsage(ctx, arg) m.queryLatencies.WithLabelValues("InsertAIBridgeToolUsage").Observe(time.Since(start).Seconds()) - return r0 + return r0, r1 } -func (m queryMetricsStore) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error { +func (m queryMetricsStore) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) (database.AIBridgeUserPrompt, error) { start := time.Now() - r0 := m.s.InsertAIBridgeUserPrompt(ctx, arg) + r0, r1 := m.s.InsertAIBridgeUserPrompt(ctx, arg) m.queryLatencies.WithLabelValues("InsertAIBridgeUserPrompt").Observe(time.Since(start).Seconds()) - return r0 + return r0, r1 } func (m queryMetricsStore) InsertAPIKey(ctx context.Context, arg database.InsertAPIKeyParams) (database.APIKey, error) { @@ -2665,6 +2665,34 @@ func (m queryMetricsStore) InsertWorkspaceResourceMetadata(ctx context.Context, return metadata, err } +func (m queryMetricsStore) ListAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams) ([]database.AIBridgeInterception, error) { + start := time.Now() + r0, r1 := m.s.ListAIBridgeInterceptions(ctx, arg) + m.queryLatencies.WithLabelValues("ListAIBridgeInterceptions").Observe(time.Since(start).Seconds()) + return r0, r1 +} + +func (m queryMetricsStore) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeTokenUsage, error) { + start := time.Now() + r0, r1 := m.s.ListAIBridgeTokenUsagesByInterceptionIDs(ctx, interceptionIds) + m.queryLatencies.WithLabelValues("ListAIBridgeTokenUsagesByInterceptionIDs").Observe(time.Since(start).Seconds()) + return r0, r1 +} + +func (m queryMetricsStore) ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeToolUsage, error) { + start := time.Now() + r0, r1 := m.s.ListAIBridgeToolUsagesByInterceptionIDs(ctx, interceptionIds) + m.queryLatencies.WithLabelValues("ListAIBridgeToolUsagesByInterceptionIDs").Observe(time.Since(start).Seconds()) + return r0, r1 +} + +func (m queryMetricsStore) ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeUserPrompt, error) { + start := time.Now() + r0, r1 := m.s.ListAIBridgeUserPromptsByInterceptionIDs(ctx, interceptionIds) + m.queryLatencies.WithLabelValues("ListAIBridgeUserPromptsByInterceptionIDs").Observe(time.Since(start).Seconds()) + return r0, r1 +} + func (m queryMetricsStore) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) { start := time.Now() r0, r1 := m.s.ListProvisionerKeysByOrganization(ctx, organizationID) @@ -3630,3 +3658,10 @@ func (m queryMetricsStore) CountAuthorizedConnectionLogs(ctx context.Context, ar m.queryLatencies.WithLabelValues("CountAuthorizedConnectionLogs").Observe(time.Since(start).Seconds()) return r0, r1 } + +func (m queryMetricsStore) ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) ([]database.AIBridgeInterception, error) { + start := time.Now() + r0, r1 := m.s.ListAuthorizedAIBridgeInterceptions(ctx, arg, prepared) + m.queryLatencies.WithLabelValues("ListAuthorizedAIBridgeInterceptions").Observe(time.Since(start).Seconds()) + return r0, r1 +} diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 66de8ec6e5..080f0390a6 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -4724,11 +4724,12 @@ func (mr *MockStoreMockRecorder) InsertAIBridgeInterception(ctx, arg any) *gomoc } // InsertAIBridgeTokenUsage mocks base method. -func (m *MockStore) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error { +func (m *MockStore) InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) (database.AIBridgeTokenUsage, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "InsertAIBridgeTokenUsage", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(database.AIBridgeTokenUsage) + ret1, _ := ret[1].(error) + return ret0, ret1 } // InsertAIBridgeTokenUsage indicates an expected call of InsertAIBridgeTokenUsage. @@ -4738,11 +4739,12 @@ func (mr *MockStoreMockRecorder) InsertAIBridgeTokenUsage(ctx, arg any) *gomock. } // InsertAIBridgeToolUsage mocks base method. -func (m *MockStore) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error { +func (m *MockStore) InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) (database.AIBridgeToolUsage, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "InsertAIBridgeToolUsage", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(database.AIBridgeToolUsage) + ret1, _ := ret[1].(error) + return ret0, ret1 } // InsertAIBridgeToolUsage indicates an expected call of InsertAIBridgeToolUsage. @@ -4752,11 +4754,12 @@ func (mr *MockStoreMockRecorder) InsertAIBridgeToolUsage(ctx, arg any) *gomock.C } // InsertAIBridgeUserPrompt mocks base method. -func (m *MockStore) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error { +func (m *MockStore) InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) (database.AIBridgeUserPrompt, error) { m.ctrl.T.Helper() ret := m.ctrl.Call(m, "InsertAIBridgeUserPrompt", ctx, arg) - ret0, _ := ret[0].(error) - return ret0 + ret0, _ := ret[0].(database.AIBridgeUserPrompt) + ret1, _ := ret[1].(error) + return ret0, ret1 } // InsertAIBridgeUserPrompt indicates an expected call of InsertAIBridgeUserPrompt. @@ -5680,6 +5683,81 @@ func (mr *MockStoreMockRecorder) InsertWorkspaceResourceMetadata(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertWorkspaceResourceMetadata", reflect.TypeOf((*MockStore)(nil).InsertWorkspaceResourceMetadata), ctx, arg) } +// ListAIBridgeInterceptions mocks base method. +func (m *MockStore) ListAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams) ([]database.AIBridgeInterception, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAIBridgeInterceptions", ctx, arg) + ret0, _ := ret[0].([]database.AIBridgeInterception) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAIBridgeInterceptions indicates an expected call of ListAIBridgeInterceptions. +func (mr *MockStoreMockRecorder) ListAIBridgeInterceptions(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeInterceptions", reflect.TypeOf((*MockStore)(nil).ListAIBridgeInterceptions), ctx, arg) +} + +// ListAIBridgeTokenUsagesByInterceptionIDs mocks base method. +func (m *MockStore) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeTokenUsage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAIBridgeTokenUsagesByInterceptionIDs", ctx, interceptionIds) + ret0, _ := ret[0].([]database.AIBridgeTokenUsage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAIBridgeTokenUsagesByInterceptionIDs indicates an expected call of ListAIBridgeTokenUsagesByInterceptionIDs. +func (mr *MockStoreMockRecorder) ListAIBridgeTokenUsagesByInterceptionIDs(ctx, interceptionIds any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeTokenUsagesByInterceptionIDs", reflect.TypeOf((*MockStore)(nil).ListAIBridgeTokenUsagesByInterceptionIDs), ctx, interceptionIds) +} + +// ListAIBridgeToolUsagesByInterceptionIDs mocks base method. +func (m *MockStore) ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeToolUsage, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAIBridgeToolUsagesByInterceptionIDs", ctx, interceptionIds) + ret0, _ := ret[0].([]database.AIBridgeToolUsage) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAIBridgeToolUsagesByInterceptionIDs indicates an expected call of ListAIBridgeToolUsagesByInterceptionIDs. +func (mr *MockStoreMockRecorder) ListAIBridgeToolUsagesByInterceptionIDs(ctx, interceptionIds any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeToolUsagesByInterceptionIDs", reflect.TypeOf((*MockStore)(nil).ListAIBridgeToolUsagesByInterceptionIDs), ctx, interceptionIds) +} + +// ListAIBridgeUserPromptsByInterceptionIDs mocks base method. +func (m *MockStore) ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]database.AIBridgeUserPrompt, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAIBridgeUserPromptsByInterceptionIDs", ctx, interceptionIds) + ret0, _ := ret[0].([]database.AIBridgeUserPrompt) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAIBridgeUserPromptsByInterceptionIDs indicates an expected call of ListAIBridgeUserPromptsByInterceptionIDs. +func (mr *MockStoreMockRecorder) ListAIBridgeUserPromptsByInterceptionIDs(ctx, interceptionIds any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeUserPromptsByInterceptionIDs", reflect.TypeOf((*MockStore)(nil).ListAIBridgeUserPromptsByInterceptionIDs), ctx, interceptionIds) +} + +// ListAuthorizedAIBridgeInterceptions mocks base method. +func (m *MockStore) ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) ([]database.AIBridgeInterception, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAuthorizedAIBridgeInterceptions", ctx, arg, prepared) + ret0, _ := ret[0].([]database.AIBridgeInterception) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAuthorizedAIBridgeInterceptions indicates an expected call of ListAuthorizedAIBridgeInterceptions. +func (mr *MockStoreMockRecorder) ListAuthorizedAIBridgeInterceptions(ctx, arg, prepared any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAuthorizedAIBridgeInterceptions", reflect.TypeOf((*MockStore)(nil).ListAuthorizedAIBridgeInterceptions), ctx, arg, prepared) +} + // ListProvisionerKeysByOrganization mocks base method. func (m *MockStore) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]database.ProvisionerKey, error) { m.ctrl.T.Helper() diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index b800962266..b517099522 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -3112,6 +3112,12 @@ CREATE INDEX idx_agent_stats_user_id ON workspace_agent_stats USING btree (user_ CREATE INDEX idx_aibridge_interceptions_initiator_id ON aibridge_interceptions USING btree (initiator_id); +CREATE INDEX idx_aibridge_interceptions_model ON aibridge_interceptions USING btree (model); + +CREATE INDEX idx_aibridge_interceptions_provider ON aibridge_interceptions USING btree (provider); + +CREATE INDEX idx_aibridge_interceptions_started_id_desc ON aibridge_interceptions USING btree (started_at DESC, id DESC); + CREATE INDEX idx_aibridge_token_usages_interception_id ON aibridge_token_usages USING btree (interception_id); CREATE INDEX idx_aibridge_token_usages_provider_response_id ON aibridge_token_usages USING btree (provider_response_id); diff --git a/coderd/database/migrations/000374_aibridge_interception_indices.down.sql b/coderd/database/migrations/000374_aibridge_interception_indices.down.sql new file mode 100644 index 0000000000..748d2a752d --- /dev/null +++ b/coderd/database/migrations/000374_aibridge_interception_indices.down.sql @@ -0,0 +1,5 @@ +DROP INDEX IF EXISTS idx_aibridge_interceptions_started_id_desc; + +DROP INDEX IF EXISTS idx_aibridge_interceptions_provider; + +DROP INDEX IF EXISTS idx_aibridge_interceptions_model; diff --git a/coderd/database/migrations/000374_aibridge_interception_indices.up.sql b/coderd/database/migrations/000374_aibridge_interception_indices.up.sql new file mode 100644 index 0000000000..14634b52c0 --- /dev/null +++ b/coderd/database/migrations/000374_aibridge_interception_indices.up.sql @@ -0,0 +1,9 @@ +-- This is used for consistent cursor pagination. +CREATE INDEX IF NOT EXISTS idx_aibridge_interceptions_started_id_desc + ON aibridge_interceptions (started_at DESC, id DESC); + +CREATE INDEX IF NOT EXISTS idx_aibridge_interceptions_provider + ON aibridge_interceptions (provider); + +CREATE INDEX IF NOT EXISTS idx_aibridge_interceptions_model + ON aibridge_interceptions (model); diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index a96acd0cf9..8e2d74fb8f 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -51,6 +51,7 @@ type customQuerier interface { userQuerier auditLogQuerier connectionLogQuerier + aibridgeQuerier } type templateQuerier interface { @@ -761,6 +762,60 @@ func (q *sqlQuerier) CountAuthorizedConnectionLogs(ctx context.Context, arg Coun return count, nil } +type aibridgeQuerier interface { + ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg ListAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) ([]AIBridgeInterception, error) +} + +func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg ListAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) ([]AIBridgeInterception, error) { + authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{ + VariableConverter: regosql.AIBridgeInterceptionConverter(), + }) + if err != nil { + return nil, xerrors.Errorf("compile authorized filter: %w", err) + } + filtered, err := insertAuthorizedFilter(listAIBridgeInterceptions, fmt.Sprintf(" AND %s", authorizedFilter)) + if err != nil { + return nil, xerrors.Errorf("insert authorized filter: %w", err) + } + + query := fmt.Sprintf("-- name: ListAuthorizedAIBridgeInterceptions :many\n%s", filtered) + rows, err := q.db.QueryContext(ctx, query, + arg.StartedAfter, + arg.StartedBefore, + arg.InitiatorID, + arg.Provider, + arg.Model, + arg.AfterID, + arg.Limit, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []AIBridgeInterception + for rows.Next() { + var i AIBridgeInterception + if err := rows.Scan( + &i.ID, + &i.InitiatorID, + &i.Provider, + &i.Model, + &i.StartedAt, + &i.Metadata, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + func insertAuthorizedFilter(query string, replaceWith string) (string, error) { if !strings.Contains(query, authorizedQueryPlaceholder) { return "", xerrors.Errorf("query does not contain authorized replace string, this is not an authorized query") diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 7d77f9aad5..a43042114c 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -508,9 +508,9 @@ type sqlcQuerier interface { GetWorkspacesByTemplateID(ctx context.Context, templateID uuid.UUID) ([]WorkspaceTable, error) GetWorkspacesEligibleForTransition(ctx context.Context, now time.Time) ([]GetWorkspacesEligibleForTransitionRow, error) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) - InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) error - InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBridgeToolUsageParams) error - InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIBridgeUserPromptParams) error + InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error) + InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBridgeToolUsageParams) (AIBridgeToolUsage, error) + InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIBridgeUserPromptParams) (AIBridgeUserPrompt, error) InsertAPIKey(ctx context.Context, arg InsertAPIKeyParams) (APIKey, error) // We use the organization_id as the id // for simplicity since all users is @@ -585,6 +585,10 @@ type sqlcQuerier interface { InsertWorkspaceProxy(ctx context.Context, arg InsertWorkspaceProxyParams) (WorkspaceProxy, error) InsertWorkspaceResource(ctx context.Context, arg InsertWorkspaceResourceParams) (WorkspaceResource, error) InsertWorkspaceResourceMetadata(ctx context.Context, arg InsertWorkspaceResourceMetadataParams) ([]WorkspaceResourceMetadatum, error) + ListAIBridgeInterceptions(ctx context.Context, arg ListAIBridgeInterceptionsParams) ([]AIBridgeInterception, error) + ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeTokenUsage, error) + ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeToolUsage, error) + ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeUserPrompt, error) ListProvisionerKeysByOrganization(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListProvisionerKeysByOrganizationExcludeReserved(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerKey, error) ListUserSecrets(ctx context.Context, userID uuid.UUID) ([]UserSecret, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index c367852e5d..e4c0300e82 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -112,7 +112,12 @@ func (q *sqlQuerier) ActivityBumpWorkspace(ctx context.Context, arg ActivityBump } const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one -SELECT id, initiator_id, provider, model, started_at, metadata FROM aibridge_interceptions WHERE id = $1::uuid +SELECT + id, initiator_id, provider, model, started_at, metadata +FROM + aibridge_interceptions +WHERE + id = $1::uuid ` func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UUID) (AIBridgeInterception, error) { @@ -130,7 +135,10 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU } const getAIBridgeInterceptions = `-- name: GetAIBridgeInterceptions :many -SELECT id, initiator_id, provider, model, started_at, metadata FROM aibridge_interceptions +SELECT + id, initiator_id, provider, model, started_at, metadata +FROM + aibridge_interceptions ` func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeInterception, error) { @@ -164,7 +172,13 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn } const getAIBridgeTokenUsagesByInterceptionID = `-- name: GetAIBridgeTokenUsagesByInterceptionID :many -SELECT id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at FROM aibridge_token_usages WHERE interception_id = $1::uuid +SELECT + id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at +FROM + aibridge_token_usages WHERE interception_id = $1::uuid +ORDER BY + created_at ASC, + id ASC ` func (q *sqlQuerier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeTokenUsage, error) { @@ -199,7 +213,15 @@ func (q *sqlQuerier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, } const getAIBridgeToolUsagesByInterceptionID = `-- name: GetAIBridgeToolUsagesByInterceptionID :many -SELECT id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at FROM aibridge_tool_usages WHERE interception_id = $1::uuid +SELECT + id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at +FROM + aibridge_tool_usages +WHERE + interception_id = $1::uuid +ORDER BY + created_at ASC, + id ASC ` func (q *sqlQuerier) GetAIBridgeToolUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeToolUsage, error) { @@ -237,7 +259,15 @@ func (q *sqlQuerier) GetAIBridgeToolUsagesByInterceptionID(ctx context.Context, } const getAIBridgeUserPromptsByInterceptionID = `-- name: GetAIBridgeUserPromptsByInterceptionID :many -SELECT id, interception_id, provider_response_id, prompt, metadata, created_at FROM aibridge_user_prompts WHERE interception_id = $1::uuid +SELECT + id, interception_id, provider_response_id, prompt, metadata, created_at +FROM + aibridge_user_prompts +WHERE + interception_id = $1::uuid +ORDER BY + created_at ASC, + id ASC ` func (q *sqlQuerier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeUserPrompt, error) { @@ -271,8 +301,11 @@ func (q *sqlQuerier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, } const insertAIBridgeInterception = `-- name: InsertAIBridgeInterception :one -INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, metadata, started_at) -VALUES ($1::uuid, $2::uuid, $3, $4, COALESCE($5::jsonb, '{}'::jsonb), $6) +INSERT INTO aibridge_interceptions ( + id, initiator_id, provider, model, metadata, started_at +) VALUES ( + $1, $2, $3, $4, COALESCE($5::jsonb, '{}'::jsonb), $6 +) RETURNING id, initiator_id, provider, model, started_at, metadata ` @@ -306,12 +339,13 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA return i, err } -const insertAIBridgeTokenUsage = `-- name: InsertAIBridgeTokenUsage :exec +const insertAIBridgeTokenUsage = `-- name: InsertAIBridgeTokenUsage :one INSERT INTO aibridge_token_usages ( id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at ) VALUES ( $1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7 ) +RETURNING id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at ` type InsertAIBridgeTokenUsageParams struct { @@ -324,8 +358,8 @@ type InsertAIBridgeTokenUsageParams struct { CreatedAt time.Time `db:"created_at" json:"created_at"` } -func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) error { - _, err := q.db.ExecContext(ctx, insertAIBridgeTokenUsage, +func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIBridgeTokenUsageParams) (AIBridgeTokenUsage, error) { + row := q.db.QueryRowContext(ctx, insertAIBridgeTokenUsage, arg.ID, arg.InterceptionID, arg.ProviderResponseID, @@ -334,15 +368,26 @@ func (q *sqlQuerier) InsertAIBridgeTokenUsage(ctx context.Context, arg InsertAIB arg.Metadata, arg.CreatedAt, ) - return err + var i AIBridgeTokenUsage + err := row.Scan( + &i.ID, + &i.InterceptionID, + &i.ProviderResponseID, + &i.InputTokens, + &i.OutputTokens, + &i.Metadata, + &i.CreatedAt, + ) + return i, err } -const insertAIBridgeToolUsage = `-- name: InsertAIBridgeToolUsage :exec +const insertAIBridgeToolUsage = `-- name: InsertAIBridgeToolUsage :one INSERT INTO aibridge_tool_usages ( id, interception_id, provider_response_id, tool, server_url, input, injected, invocation_error, metadata, created_at ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, COALESCE($9::jsonb, '{}'::jsonb), $10 ) +RETURNING id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at ` type InsertAIBridgeToolUsageParams struct { @@ -358,8 +403,8 @@ type InsertAIBridgeToolUsageParams struct { CreatedAt time.Time `db:"created_at" json:"created_at"` } -func (q *sqlQuerier) InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBridgeToolUsageParams) error { - _, err := q.db.ExecContext(ctx, insertAIBridgeToolUsage, +func (q *sqlQuerier) InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBridgeToolUsageParams) (AIBridgeToolUsage, error) { + row := q.db.QueryRowContext(ctx, insertAIBridgeToolUsage, arg.ID, arg.InterceptionID, arg.ProviderResponseID, @@ -371,15 +416,29 @@ func (q *sqlQuerier) InsertAIBridgeToolUsage(ctx context.Context, arg InsertAIBr arg.Metadata, arg.CreatedAt, ) - return err + var i AIBridgeToolUsage + err := row.Scan( + &i.ID, + &i.InterceptionID, + &i.ProviderResponseID, + &i.ServerUrl, + &i.Tool, + &i.Input, + &i.Injected, + &i.InvocationError, + &i.Metadata, + &i.CreatedAt, + ) + return i, err } -const insertAIBridgeUserPrompt = `-- name: InsertAIBridgeUserPrompt :exec +const insertAIBridgeUserPrompt = `-- name: InsertAIBridgeUserPrompt :one INSERT INTO aibridge_user_prompts ( id, interception_id, provider_response_id, prompt, metadata, created_at ) VALUES ( $1, $2, $3, $4, COALESCE($5::jsonb, '{}'::jsonb), $6 ) +RETURNING id, interception_id, provider_response_id, prompt, metadata, created_at ` type InsertAIBridgeUserPromptParams struct { @@ -391,8 +450,8 @@ type InsertAIBridgeUserPromptParams struct { CreatedAt time.Time `db:"created_at" json:"created_at"` } -func (q *sqlQuerier) InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIBridgeUserPromptParams) error { - _, err := q.db.ExecContext(ctx, insertAIBridgeUserPrompt, +func (q *sqlQuerier) InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIBridgeUserPromptParams) (AIBridgeUserPrompt, error) { + row := q.db.QueryRowContext(ctx, insertAIBridgeUserPrompt, arg.ID, arg.InterceptionID, arg.ProviderResponseID, @@ -400,7 +459,249 @@ func (q *sqlQuerier) InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIB arg.Metadata, arg.CreatedAt, ) - return err + var i AIBridgeUserPrompt + err := row.Scan( + &i.ID, + &i.InterceptionID, + &i.ProviderResponseID, + &i.Prompt, + &i.Metadata, + &i.CreatedAt, + ) + return i, err +} + +const listAIBridgeInterceptions = `-- name: ListAIBridgeInterceptions :many +SELECT + id, initiator_id, provider, model, started_at, metadata +FROM + aibridge_interceptions +WHERE + -- Filter by time frame + CASE + WHEN $1::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at >= $1::timestamptz + ELSE true + END + AND CASE + WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at <= $2::timestamptz + ELSE true + END + -- Filter initiator_id + AND CASE + WHEN $3::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN aibridge_interceptions.initiator_id = $3::uuid + ELSE true + END + -- Filter provider + AND CASE + WHEN $4::text != '' THEN aibridge_interceptions.provider = $4::text + ELSE true + END + -- Filter model + AND CASE + WHEN $5::text != '' THEN aibridge_interceptions.model = $5::text + ELSE true + END + -- Cursor pagination + AND CASE + WHEN $6::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ( + -- The pagination cursor is the last ID of the previous page. + -- The query is ordered by the started_at field, so select all + -- rows before the cursor and before the after_id UUID. + -- This uses a less than operator because we're sorting DESC. The + -- "after_id" terminology comes from our pagination parser in + -- coderd. + (aibridge_interceptions.started_at, aibridge_interceptions.id) < ( + (SELECT started_at FROM aibridge_interceptions WHERE id = $6), + $6::uuid + ) + ) + ELSE true + END + -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeInterceptions + -- @authorize_filter +ORDER BY + aibridge_interceptions.started_at DESC, + aibridge_interceptions.id DESC +LIMIT COALESCE(NULLIF($7::integer, 0), 100) +` + +type ListAIBridgeInterceptionsParams struct { + StartedAfter time.Time `db:"started_after" json:"started_after"` + StartedBefore time.Time `db:"started_before" json:"started_before"` + InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` + Provider string `db:"provider" json:"provider"` + Model string `db:"model" json:"model"` + AfterID uuid.UUID `db:"after_id" json:"after_id"` + Limit int32 `db:"limit_" json:"limit_"` +} + +func (q *sqlQuerier) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBridgeInterceptionsParams) ([]AIBridgeInterception, error) { + rows, err := q.db.QueryContext(ctx, listAIBridgeInterceptions, + arg.StartedAfter, + arg.StartedBefore, + arg.InitiatorID, + arg.Provider, + arg.Model, + arg.AfterID, + arg.Limit, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []AIBridgeInterception + for rows.Next() { + var i AIBridgeInterception + if err := rows.Scan( + &i.ID, + &i.InitiatorID, + &i.Provider, + &i.Model, + &i.StartedAt, + &i.Metadata, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listAIBridgeTokenUsagesByInterceptionIDs = `-- name: ListAIBridgeTokenUsagesByInterceptionIDs :many +SELECT + id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at +FROM + aibridge_token_usages +WHERE + interception_id = ANY($1::uuid[]) +ORDER BY + created_at ASC, + id ASC +` + +func (q *sqlQuerier) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeTokenUsage, error) { + rows, err := q.db.QueryContext(ctx, listAIBridgeTokenUsagesByInterceptionIDs, pq.Array(interceptionIds)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []AIBridgeTokenUsage + for rows.Next() { + var i AIBridgeTokenUsage + if err := rows.Scan( + &i.ID, + &i.InterceptionID, + &i.ProviderResponseID, + &i.InputTokens, + &i.OutputTokens, + &i.Metadata, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listAIBridgeToolUsagesByInterceptionIDs = `-- name: ListAIBridgeToolUsagesByInterceptionIDs :many +SELECT + id, interception_id, provider_response_id, server_url, tool, input, injected, invocation_error, metadata, created_at +FROM + aibridge_tool_usages +WHERE + interception_id = ANY($1::uuid[]) +ORDER BY + created_at ASC, + id ASC +` + +func (q *sqlQuerier) ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeToolUsage, error) { + rows, err := q.db.QueryContext(ctx, listAIBridgeToolUsagesByInterceptionIDs, pq.Array(interceptionIds)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []AIBridgeToolUsage + for rows.Next() { + var i AIBridgeToolUsage + if err := rows.Scan( + &i.ID, + &i.InterceptionID, + &i.ProviderResponseID, + &i.ServerUrl, + &i.Tool, + &i.Input, + &i.Injected, + &i.InvocationError, + &i.Metadata, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const listAIBridgeUserPromptsByInterceptionIDs = `-- name: ListAIBridgeUserPromptsByInterceptionIDs :many +SELECT + id, interception_id, provider_response_id, prompt, metadata, created_at +FROM + aibridge_user_prompts +WHERE + interception_id = ANY($1::uuid[]) +ORDER BY + created_at ASC, + id ASC +` + +func (q *sqlQuerier) ListAIBridgeUserPromptsByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeUserPrompt, error) { + rows, err := q.db.QueryContext(ctx, listAIBridgeUserPromptsByInterceptionIDs, pq.Array(interceptionIds)) + if err != nil { + return nil, err + } + defer rows.Close() + var items []AIBridgeUserPrompt + for rows.Next() { + var i AIBridgeUserPrompt + if err := rows.Scan( + &i.ID, + &i.InterceptionID, + &i.ProviderResponseID, + &i.Prompt, + &i.Metadata, + &i.CreatedAt, + ); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil } const deleteAPIKeyByID = `-- name: DeleteAPIKeyByID :exec diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index dab2073cea..79d41defd5 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -1,40 +1,163 @@ -- name: InsertAIBridgeInterception :one -INSERT INTO aibridge_interceptions (id, initiator_id, provider, model, metadata, started_at) -VALUES (@id::uuid, @initiator_id::uuid, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at) +INSERT INTO aibridge_interceptions ( + id, initiator_id, provider, model, metadata, started_at +) VALUES ( + @id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at +) RETURNING *; --- name: InsertAIBridgeTokenUsage :exec +-- name: InsertAIBridgeTokenUsage :one INSERT INTO aibridge_token_usages ( id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at ) VALUES ( @id, @interception_id, @provider_response_id, @input_tokens, @output_tokens, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at -); +) +RETURNING *; --- name: InsertAIBridgeUserPrompt :exec +-- name: InsertAIBridgeUserPrompt :one INSERT INTO aibridge_user_prompts ( id, interception_id, provider_response_id, prompt, metadata, created_at ) VALUES ( @id, @interception_id, @provider_response_id, @prompt, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at -); +) +RETURNING *; --- name: InsertAIBridgeToolUsage :exec +-- name: InsertAIBridgeToolUsage :one INSERT INTO aibridge_tool_usages ( id, interception_id, provider_response_id, tool, server_url, input, injected, invocation_error, metadata, created_at ) VALUES ( @id, @interception_id, @provider_response_id, @tool, @server_url, @input, @injected, @invocation_error, COALESCE(@metadata::jsonb, '{}'::jsonb), @created_at -); +) +RETURNING *; -- name: GetAIBridgeInterceptionByID :one -SELECT * FROM aibridge_interceptions WHERE id = @id::uuid; +SELECT + * +FROM + aibridge_interceptions +WHERE + id = @id::uuid; -- name: GetAIBridgeInterceptions :many -SELECT * FROM aibridge_interceptions; +SELECT + * +FROM + aibridge_interceptions; -- name: GetAIBridgeTokenUsagesByInterceptionID :many -SELECT * FROM aibridge_token_usages WHERE interception_id = @interception_id::uuid; +SELECT + * +FROM + aibridge_token_usages WHERE interception_id = @interception_id::uuid +ORDER BY + created_at ASC, + id ASC; -- name: GetAIBridgeUserPromptsByInterceptionID :many -SELECT * FROM aibridge_user_prompts WHERE interception_id = @interception_id::uuid; +SELECT + * +FROM + aibridge_user_prompts +WHERE + interception_id = @interception_id::uuid +ORDER BY + created_at ASC, + id ASC; -- name: GetAIBridgeToolUsagesByInterceptionID :many -SELECT * FROM aibridge_tool_usages WHERE interception_id = @interception_id::uuid; +SELECT + * +FROM + aibridge_tool_usages +WHERE + interception_id = @interception_id::uuid +ORDER BY + created_at ASC, + id ASC; + +-- name: ListAIBridgeInterceptions :many +SELECT + * +FROM + aibridge_interceptions +WHERE + -- Filter by time frame + CASE + WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at >= @started_after::timestamptz + ELSE true + END + AND CASE + WHEN @started_before::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at <= @started_before::timestamptz + ELSE true + END + -- Filter initiator_id + AND CASE + WHEN @initiator_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN aibridge_interceptions.initiator_id = @initiator_id::uuid + ELSE true + END + -- Filter provider + AND CASE + WHEN @provider::text != '' THEN aibridge_interceptions.provider = @provider::text + ELSE true + END + -- Filter model + AND CASE + WHEN @model::text != '' THEN aibridge_interceptions.model = @model::text + ELSE true + END + -- Cursor pagination + AND CASE + WHEN @after_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ( + -- The pagination cursor is the last ID of the previous page. + -- The query is ordered by the started_at field, so select all + -- rows before the cursor and before the after_id UUID. + -- This uses a less than operator because we're sorting DESC. The + -- "after_id" terminology comes from our pagination parser in + -- coderd. + (aibridge_interceptions.started_at, aibridge_interceptions.id) < ( + (SELECT started_at FROM aibridge_interceptions WHERE id = @after_id), + @after_id::uuid + ) + ) + ELSE true + END + -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeInterceptions + -- @authorize_filter +ORDER BY + aibridge_interceptions.started_at DESC, + aibridge_interceptions.id DESC +LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100) +; + +-- name: ListAIBridgeTokenUsagesByInterceptionIDs :many +SELECT + * +FROM + aibridge_token_usages +WHERE + interception_id = ANY(@interception_ids::uuid[]) +ORDER BY + created_at ASC, + id ASC; + +-- name: ListAIBridgeUserPromptsByInterceptionIDs :many +SELECT + * +FROM + aibridge_user_prompts +WHERE + interception_id = ANY(@interception_ids::uuid[]) +ORDER BY + created_at ASC, + id ASC; + +-- name: ListAIBridgeToolUsagesByInterceptionIDs :many +SELECT + * +FROM + aibridge_tool_usages +WHERE + interception_id = ANY(@interception_ids::uuid[]) +ORDER BY + created_at ASC, + id ASC; diff --git a/coderd/rbac/regosql/configs.go b/coderd/rbac/regosql/configs.go index 1c1e126ff6..b06d4d0583 100644 --- a/coderd/rbac/regosql/configs.go +++ b/coderd/rbac/regosql/configs.go @@ -54,7 +54,7 @@ func AuditLogConverter() *sqltypes.VariableConverter { matcher := sqltypes.NewVariableConverter().RegisterMatcher( resourceIDMatcher(), sqltypes.StringVarMatcher("COALESCE(audit_logs.organization_id :: text, '')", []string{"input", "object", "org_owner"}), - // Aduit logs have no user owner, only owner by an organization. + // Audit logs have no user owner, only owner by an organization. sqltypes.AlwaysFalse(userOwnerMatcher()), ) matcher.RegisterMatcher( @@ -78,6 +78,21 @@ func ConnectionLogConverter() *sqltypes.VariableConverter { return matcher } +func AIBridgeInterceptionConverter() *sqltypes.VariableConverter { + matcher := sqltypes.NewVariableConverter().RegisterMatcher( + resourceIDMatcher(), + // AIBridge interceptions are not tied to any organization. + sqltypes.StringVarMatcher("''", []string{"input", "object", "org_owner"}), + sqltypes.StringVarMatcher("initiator_id :: text", []string{"input", "object", "owner"}), + ) + matcher.RegisterMatcher( + // No ACLs on the aibridge interception type + sqltypes.AlwaysFalse(groupACLMatcher(matcher)), + sqltypes.AlwaysFalse(userACLMatcher(matcher)), + ) + return matcher +} + func UserConverter() *sqltypes.VariableConverter { matcher := sqltypes.NewVariableConverter().RegisterMatcher( resourceIDMatcher(), diff --git a/coderd/rbac/roles.go b/coderd/rbac/roles.go index c6770f31b0..bc75991d01 100644 --- a/coderd/rbac/roles.go +++ b/coderd/rbac/roles.go @@ -327,6 +327,8 @@ func ReloadBuiltinRoles(opts *RoleOptions) { // Allow auditors to query deployment stats and insights. ResourceDeploymentStats.Type: {policy.ActionRead}, ResourceDeploymentConfig.Type: {policy.ActionRead}, + // Allow auditors to query aibridge interceptions. + ResourceAibridgeInterception.Type: {policy.ActionRead}, }), Org: map[string][]Permission{}, User: []Permission{}, diff --git a/coderd/searchquery/search.go b/coderd/searchquery/search.go index c9a418ff7c..5a135584df 100644 --- a/coderd/searchquery/search.go +++ b/coderd/searchquery/search.go @@ -226,7 +226,8 @@ func Workspaces(ctx context.Context, db database.Store, query string, page coder filter.HasExternalAgent = parser.NullableBoolean(values, sql.NullBool{}, "has_external_agent") filter.OrganizationID = parseOrganization(ctx, db, parser, values, "organization") filter.Shared = parser.NullableBoolean(values, sql.NullBool{}, "shared") - filter.SharedWithUserID = parseUser(ctx, db, parser, values, "shared_with_user") + // TODO: support "me" by passing in the actorID + filter.SharedWithUserID = parseUser(ctx, db, parser, values, "shared_with_user", uuid.Nil) filter.SharedWithGroupID = parseGroup(ctx, db, parser, values, "shared_with_group") type paramMatch struct { @@ -304,6 +305,46 @@ func Templates(ctx context.Context, db database.Store, actorID uuid.UUID, query return filter, parser.Errors } +func AIBridgeInterceptions(ctx context.Context, db database.Store, query string, page codersdk.Pagination, actorID uuid.UUID) (database.ListAIBridgeInterceptionsParams, []codersdk.ValidationError) { + // nolint:exhaustruct // Empty values just means "don't filter by that field". + filter := database.ListAIBridgeInterceptionsParams{ + AfterID: page.AfterID, + // #nosec G115 - Safe conversion for pagination limit which is expected to be within int32 range + Limit: int32(page.Limit), + } + + if query == "" { + return filter, nil + } + + values, errors := searchTerms(query, func(term string, values url.Values) error { + // Default to the initiating user + values.Add("user", term) + return nil + }) + if len(errors) > 0 { + return filter, errors + } + + parser := httpapi.NewQueryParamParser() + filter.InitiatorID = parseUser(ctx, db, parser, values, "initiator", actorID) + filter.Provider = parser.String(values, "", "provider") + filter.Model = parser.String(values, "", "model") + + // Time must be between started_after and started_before. + filter.StartedAfter = parser.Time3339Nano(values, time.Time{}, "started_after") + filter.StartedBefore = parser.Time3339Nano(values, time.Time{}, "started_before") + if !filter.StartedBefore.IsZero() && !filter.StartedAfter.IsZero() && !filter.StartedBefore.After(filter.StartedAfter) { + parser.Errors = append(parser.Errors, codersdk.ValidationError{ + Field: "started_before", + Detail: `Query param "started_before" has invalid value: "started_before" must be after "started_after" if set`, + }) + } + + parser.ErrorExcessParams(values) + return filter, parser.Errors +} + func searchTerms(query string, defaultKey func(term string, values url.Values) error) (url.Values, []codersdk.ValidationError) { searchValues := make(url.Values) @@ -365,11 +406,14 @@ func parseOrganization(ctx context.Context, db database.Store, parser *httpapi.Q }) } -func parseUser(ctx context.Context, db database.Store, parser *httpapi.QueryParamParser, vals url.Values, queryParam string) uuid.UUID { +func parseUser(ctx context.Context, db database.Store, parser *httpapi.QueryParamParser, vals url.Values, queryParam string, actorID uuid.UUID) uuid.UUID { return httpapi.ParseCustom(parser, vals, uuid.Nil, queryParam, func(v string) (uuid.UUID, error) { if v == "" { return uuid.Nil, nil } + if v == codersdk.Me && actorID != uuid.Nil { + return actorID, nil + } userID, err := uuid.Parse(v) if err == nil { return userID, nil diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go new file mode 100644 index 0000000000..3101dab383 --- /dev/null +++ b/codersdk/aibridge.go @@ -0,0 +1,125 @@ +package codersdk + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strings" + "time" + + "github.com/google/uuid" +) + +type AIBridgeInterception struct { + ID uuid.UUID `json:"id" format:"uuid"` + InitiatorID uuid.UUID `json:"initiator_id" format:"uuid"` + Provider string `json:"provider"` + Model string `json:"model"` + Metadata map[string]any `json:"metadata"` + StartedAt time.Time `json:"started_at" format:"date-time"` + TokenUsages []AIBridgeTokenUsage `json:"token_usages"` + UserPrompts []AIBridgeUserPrompt `json:"user_prompts"` + ToolUsages []AIBridgeToolUsage `json:"tool_usages"` +} + +type AIBridgeTokenUsage struct { + ID uuid.UUID `json:"id" format:"uuid"` + InterceptionID uuid.UUID `json:"interception_id" format:"uuid"` + ProviderResponseID string `json:"provider_response_id"` + InputTokens int64 `json:"input_tokens"` + OutputTokens int64 `json:"output_tokens"` + Metadata map[string]any `json:"metadata"` + CreatedAt time.Time `json:"created_at" format:"date-time"` +} + +type AIBridgeUserPrompt struct { + ID uuid.UUID `json:"id" format:"uuid"` + InterceptionID uuid.UUID `json:"interception_id" format:"uuid"` + ProviderResponseID string `json:"provider_response_id"` + Prompt string `json:"prompt"` + Metadata map[string]any `json:"metadata"` + CreatedAt time.Time `json:"created_at" format:"date-time"` +} + +type AIBridgeToolUsage struct { + ID uuid.UUID `json:"id" format:"uuid"` + InterceptionID uuid.UUID `json:"interception_id" format:"uuid"` + ProviderResponseID string `json:"provider_response_id"` + ServerURL string `json:"server_url"` + Tool string `json:"tool"` + Input string `json:"input"` + Injected bool `json:"injected"` + InvocationError string `json:"invocation_error"` + Metadata map[string]any `json:"metadata"` + CreatedAt time.Time `json:"created_at" format:"date-time"` +} + +type AIBridgeListInterceptionsResponse struct { + Results []AIBridgeInterception `json:"results"` +} + +// @typescript-ignore AIBridgeListInterceptionsFilter +type AIBridgeListInterceptionsFilter struct { + // Limit defaults to 100, max is 1000. + // Offset based pagination is not supported for AIBridge interceptions. Use + // cursor pagination instead with after_id. + Pagination Pagination `json:"pagination,omitempty"` + + // Initiator is a user ID, username, or "me". + Initiator string `json:"initiator,omitempty"` + StartedBefore time.Time `json:"started_before,omitempty" format:"date-time"` + StartedAfter time.Time `json:"started_after,omitempty" format:"date-time"` + Provider string `json:"provider,omitempty"` + Model string `json:"model,omitempty"` + + FilterQuery string `json:"q,omitempty"` +} + +// asRequestOption returns a function that can be used in (*Client).Request. +// It modifies the request query parameters. +func (f AIBridgeListInterceptionsFilter) asRequestOption() RequestOption { + return func(r *http.Request) { + var params []string + // Make sure all user input is quoted to ensure it's parsed as a single + // string. + if f.Initiator != "" { + params = append(params, fmt.Sprintf("initiator:%q", f.Initiator)) + } + if !f.StartedBefore.IsZero() { + params = append(params, fmt.Sprintf("started_before:%q", f.StartedBefore.Format(time.RFC3339Nano))) + } + if !f.StartedAfter.IsZero() { + params = append(params, fmt.Sprintf("started_after:%q", f.StartedAfter.Format(time.RFC3339Nano))) + } + if f.Provider != "" { + params = append(params, fmt.Sprintf("provider:%q", f.Provider)) + } + if f.Model != "" { + params = append(params, fmt.Sprintf("model:%q", f.Model)) + } + if f.FilterQuery != "" { + // If custom stuff is added, just add it on here. + params = append(params, f.FilterQuery) + } + + q := r.URL.Query() + q.Set("q", strings.Join(params, " ")) + r.URL.RawQuery = q.Encode() + } +} + +// AIBridgeListInterceptions returns AIBridge interceptions with the given +// filter. +func (c *ExperimentalClient) AIBridgeListInterceptions(ctx context.Context, filter AIBridgeListInterceptionsFilter) (AIBridgeListInterceptionsResponse, error) { + res, err := c.Request(ctx, http.MethodGet, "/api/experimental/aibridge/interceptions", nil, filter.asRequestOption(), filter.Pagination.asRequestOption(), filter.Pagination.asRequestOption()) + if err != nil { + return AIBridgeListInterceptionsResponse{}, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return AIBridgeListInterceptionsResponse{}, ReadBodyAsError(res) + } + var resp AIBridgeListInterceptionsResponse + return resp, json.NewDecoder(res.Body).Decode(&resp) +} diff --git a/docs/reference/api/aibridge.md b/docs/reference/api/aibridge.md new file mode 100644 index 0000000000..6c929a3fa9 --- /dev/null +++ b/docs/reference/api/aibridge.md @@ -0,0 +1,96 @@ +# AIBridge + +## List AIBridge interceptions + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/api/experimental/aibridge/interceptions \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /api/experimental/aibridge/interceptions` + +### Parameters + +| Name | In | Type | Required | Description | +|------------|-------|---------|----------|------------------------------------------------------------------------------------------------------------------------| +| `q` | query | string | false | Search query in the format `key:value`. Available keys are: initiator, provider, model, started_after, started_before. | +| `limit` | query | integer | false | Page limit | +| `after_id` | query | string | false | Cursor pagination after ID | + +### Example responses + +> 200 Response + +```json +{ + "results": [ + { + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "initiator_id": "06588898-9a84-4b35-ba8f-f9cbd64946f3", + "metadata": { + "property1": null, + "property2": null + }, + "model": "string", + "provider": "string", + "started_at": "2019-08-24T14:15:22Z", + "token_usages": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "input_tokens": 0, + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "output_tokens": 0, + "provider_response_id": "string" + } + ], + "tool_usages": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "injected": true, + "input": "string", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "invocation_error": "string", + "metadata": { + "property1": null, + "property2": null + }, + "provider_response_id": "string", + "server_url": "string", + "tool": "string" + } + ], + "user_prompts": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "prompt": "string", + "provider_response_id": "string" + } + ] + } + ] +} +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|----------------------------------------------------------------------------------------------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | [codersdk.AIBridgeListInterceptionsResponse](schemas.md#codersdkaibridgelistinterceptionsresponse) | + +To perform this operation, you must be authenticated. [Learn more](authentication.md). diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index a1ac1dafb5..0459093732 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -375,6 +375,151 @@ | `enabled` | boolean | false | | | | `openai` | [codersdk.AIBridgeOpenAIConfig](#codersdkaibridgeopenaiconfig) | false | | | +## codersdk.AIBridgeInterception + +```json +{ + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "initiator_id": "06588898-9a84-4b35-ba8f-f9cbd64946f3", + "metadata": { + "property1": null, + "property2": null + }, + "model": "string", + "provider": "string", + "started_at": "2019-08-24T14:15:22Z", + "token_usages": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "input_tokens": 0, + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "output_tokens": 0, + "provider_response_id": "string" + } + ], + "tool_usages": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "injected": true, + "input": "string", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "invocation_error": "string", + "metadata": { + "property1": null, + "property2": null + }, + "provider_response_id": "string", + "server_url": "string", + "tool": "string" + } + ], + "user_prompts": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "prompt": "string", + "provider_response_id": "string" + } + ] +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|--------------------|---------------------------------------------------------------------|----------|--------------|-------------| +| `id` | string | false | | | +| `initiator_id` | string | false | | | +| `metadata` | object | false | | | +| » `[any property]` | any | false | | | +| `model` | string | false | | | +| `provider` | string | false | | | +| `started_at` | string | false | | | +| `token_usages` | array of [codersdk.AIBridgeTokenUsage](#codersdkaibridgetokenusage) | false | | | +| `tool_usages` | array of [codersdk.AIBridgeToolUsage](#codersdkaibridgetoolusage) | false | | | +| `user_prompts` | array of [codersdk.AIBridgeUserPrompt](#codersdkaibridgeuserprompt) | false | | | + +## codersdk.AIBridgeListInterceptionsResponse + +```json +{ + "results": [ + { + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "initiator_id": "06588898-9a84-4b35-ba8f-f9cbd64946f3", + "metadata": { + "property1": null, + "property2": null + }, + "model": "string", + "provider": "string", + "started_at": "2019-08-24T14:15:22Z", + "token_usages": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "input_tokens": 0, + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "output_tokens": 0, + "provider_response_id": "string" + } + ], + "tool_usages": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "injected": true, + "input": "string", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "invocation_error": "string", + "metadata": { + "property1": null, + "property2": null + }, + "provider_response_id": "string", + "server_url": "string", + "tool": "string" + } + ], + "user_prompts": [ + { + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "prompt": "string", + "provider_response_id": "string" + } + ] + } + ] +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|-----------|-------------------------------------------------------------------------|----------|--------------|-------------| +| `results` | array of [codersdk.AIBridgeInterception](#codersdkaibridgeinterception) | false | | | + ## codersdk.AIBridgeOpenAIConfig ```json @@ -391,6 +536,100 @@ | `base_url` | string | false | | | | `key` | string | false | | | +## codersdk.AIBridgeTokenUsage + +```json +{ + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "input_tokens": 0, + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "output_tokens": 0, + "provider_response_id": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|------------------------|---------|----------|--------------|-------------| +| `created_at` | string | false | | | +| `id` | string | false | | | +| `input_tokens` | integer | false | | | +| `interception_id` | string | false | | | +| `metadata` | object | false | | | +| » `[any property]` | any | false | | | +| `output_tokens` | integer | false | | | +| `provider_response_id` | string | false | | | + +## codersdk.AIBridgeToolUsage + +```json +{ + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "injected": true, + "input": "string", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "invocation_error": "string", + "metadata": { + "property1": null, + "property2": null + }, + "provider_response_id": "string", + "server_url": "string", + "tool": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|------------------------|---------|----------|--------------|-------------| +| `created_at` | string | false | | | +| `id` | string | false | | | +| `injected` | boolean | false | | | +| `input` | string | false | | | +| `interception_id` | string | false | | | +| `invocation_error` | string | false | | | +| `metadata` | object | false | | | +| » `[any property]` | any | false | | | +| `provider_response_id` | string | false | | | +| `server_url` | string | false | | | +| `tool` | string | false | | | + +## codersdk.AIBridgeUserPrompt + +```json +{ + "created_at": "2019-08-24T14:15:22Z", + "id": "497f6eca-6276-4993-bfeb-53cbbbba6f08", + "interception_id": "34d9b688-63ad-46f4-88b5-665c1e7f7824", + "metadata": { + "property1": null, + "property2": null + }, + "prompt": "string", + "provider_response_id": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|------------------------|--------|----------|--------------|-------------| +| `created_at` | string | false | | | +| `id` | string | false | | | +| `interception_id` | string | false | | | +| `metadata` | object | false | | | +| » `[any property]` | any | false | | | +| `prompt` | string | false | | | +| `provider_response_id` | string | false | | | + ## codersdk.AIConfig ```json diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go new file mode 100644 index 0000000000..0b51aecfa1 --- /dev/null +++ b/enterprise/coderd/aibridge.go @@ -0,0 +1,159 @@ +package coderd + +import ( + "context" + "fmt" + "net/http" + + "github.com/google/uuid" + "golang.org/x/xerrors" + + "github.com/coder/coder/v2/coderd" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/db2sdk" + "github.com/coder/coder/v2/coderd/database/dbauthz" + "github.com/coder/coder/v2/coderd/httpapi" + "github.com/coder/coder/v2/coderd/httpmw" + "github.com/coder/coder/v2/coderd/searchquery" + "github.com/coder/coder/v2/codersdk" +) + +const ( + maxListInterceptionsLimit = 1000 + defaultListInterceptionsLimit = 100 +) + +// aiBridgeListInterceptions returns all AIBridge interceptions a user can read. +// Optional filters with query params +// +// @Summary List AIBridge interceptions +// @ID list-aibridge-interceptions +// @Security CoderSessionToken +// @Produce json +// @Tags AIBridge +// @Param q query string false "Search query in the format `key:value`. Available keys are: initiator, provider, model, started_after, started_before." +// @Param limit query int false "Page limit" +// @Param after_id query string false "Cursor pagination after ID" +// @Success 200 {object} codersdk.AIBridgeListInterceptionsResponse +// @Router /api/experimental/aibridge/interceptions [get] +func (api *API) aiBridgeListInterceptions(rw http.ResponseWriter, r *http.Request) { + ctx := r.Context() + apiKey := httpmw.APIKey(r) + + page, ok := coderd.ParsePagination(rw, r) + if !ok { + return + } + if page.Offset != 0 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Offset pagination is not supported.", + Detail: "Offset pagination is not supported for AIBridge interceptions. Use cursor pagination instead with after_id..", + }) + return + } + if page.Limit == 0 { + page.Limit = defaultListInterceptionsLimit + } + if page.Limit > maxListInterceptionsLimit || page.Limit < 1 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Invalid pagination limit value.", + Detail: fmt.Sprintf("Pagination limit must be in range (0, %d]", maxListInterceptionsLimit), + }) + return + } + + queryStr := r.URL.Query().Get("q") + filter, errs := searchquery.AIBridgeInterceptions(ctx, api.Database, queryStr, page, apiKey.UserID) + if len(errs) > 0 { + httpapi.Write(ctx, rw, http.StatusBadRequest, codersdk.Response{ + Message: "Invalid workspace search query.", + Validations: errs, + }) + return + } + + var rows []database.AIBridgeInterception + err := api.Database.InTx(func(db database.Store) error { + // Ensure the after_id interception exists and is visible to the user. + if page.AfterID != uuid.Nil { + _, err := db.GetAIBridgeInterceptionByID(ctx, page.AfterID) + if err != nil { + return xerrors.Errorf("get aibridge interception by id %s for cursor pagination: %w", page.AfterID, err) + } + } + + var err error + // This only returns authorized interceptions (when using dbauthz). + rows, err = db.ListAIBridgeInterceptions(ctx, filter) + if err != nil { + return xerrors.Errorf("list aibridge interceptions: %w", err) + } + + return nil + }, nil) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Internal error getting AIBridge interceptions.", + Detail: err.Error(), + }) + return + } + + // This fetches the other rows associated with the interceptions. + items, err := populatedAndConvertAIBridgeInterceptions(ctx, api.Database, rows) + if err != nil { + httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{ + Message: "Internal error converting database rows to API response.", + Detail: err.Error(), + }) + return + } + + httpapi.Write(ctx, rw, http.StatusOK, codersdk.AIBridgeListInterceptionsResponse{ + Results: items, + }) +} + +func populatedAndConvertAIBridgeInterceptions(ctx context.Context, db database.Store, dbInterceptions []database.AIBridgeInterception) ([]codersdk.AIBridgeInterception, error) { + ids := make([]uuid.UUID, len(dbInterceptions)) + for i, row := range dbInterceptions { + ids[i] = row.ID + } + + //nolint:gocritic // This is a system function until we implement a join for aibridge interceptions. AIBridge interception subresources use the same authorization call as their parent. + tokenUsagesRows, err := db.ListAIBridgeTokenUsagesByInterceptionIDs(dbauthz.AsSystemRestricted(ctx), ids) + if err != nil { + return nil, xerrors.Errorf("get linked aibridge token usages from database: %w", err) + } + tokenUsagesMap := make(map[uuid.UUID][]database.AIBridgeTokenUsage, len(dbInterceptions)) + for _, row := range tokenUsagesRows { + tokenUsagesMap[row.InterceptionID] = append(tokenUsagesMap[row.InterceptionID], row) + } + + //nolint:gocritic // This is a system function until we implement a join for aibridge interceptions. AIBridge interception subresources use the same authorization call as their parent. + userPromptRows, err := db.ListAIBridgeUserPromptsByInterceptionIDs(dbauthz.AsSystemRestricted(ctx), ids) + if err != nil { + return nil, xerrors.Errorf("get linked aibridge user prompts from database: %w", err) + } + userPromptsMap := make(map[uuid.UUID][]database.AIBridgeUserPrompt, len(dbInterceptions)) + for _, row := range userPromptRows { + userPromptsMap[row.InterceptionID] = append(userPromptsMap[row.InterceptionID], row) + } + + //nolint:gocritic // This is a system function until we implement a join for aibridge interceptions. AIBridge interception subresources use the same authorization call as their parent. + toolUsagesRows, err := db.ListAIBridgeToolUsagesByInterceptionIDs(dbauthz.AsSystemRestricted(ctx), ids) + if err != nil { + return nil, xerrors.Errorf("get linked aibridge tool usages from database: %w", err) + } + toolUsagesMap := make(map[uuid.UUID][]database.AIBridgeToolUsage, len(dbInterceptions)) + for _, row := range toolUsagesRows { + toolUsagesMap[row.InterceptionID] = append(toolUsagesMap[row.InterceptionID], row) + } + + items := make([]codersdk.AIBridgeInterception, len(dbInterceptions)) + for i, row := range dbInterceptions { + items[i] = db2sdk.AIBridgeInterception(row, tokenUsagesMap[row.ID], userPromptsMap[row.ID], toolUsagesMap[row.ID]) + } + + return items, nil +} diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go new file mode 100644 index 0000000000..e323c0e2b8 --- /dev/null +++ b/enterprise/coderd/aibridge_test.go @@ -0,0 +1,478 @@ +package coderd_test + +import ( + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/require" + + "github.com/coder/coder/v2/coderd/coderdtest" + "github.com/coder/coder/v2/coderd/database" + "github.com/coder/coder/v2/coderd/database/db2sdk" + "github.com/coder/coder/v2/coderd/database/dbgen" + "github.com/coder/coder/v2/coderd/database/dbtime" + "github.com/coder/coder/v2/codersdk" + "github.com/coder/coder/v2/cryptorand" + "github.com/coder/coder/v2/enterprise/coderd/coderdenttest" + "github.com/coder/coder/v2/testutil" +) + +func TestAIBridgeListInterceptions(t *testing.T) { + t.Parallel() + + t.Run("EmptyDB", func(t *testing.T) { + t.Parallel() + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, _ := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + experimentalClient := codersdk.NewExperimentalClient(client) + ctx := testutil.Context(t, testutil.WaitLong) + res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{}) + require.NoError(t, err) + require.Empty(t, res.Results) + }) + + t.Run("OK", func(t *testing.T) { + t.Parallel() + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, db, _ := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + experimentalClient := codersdk.NewExperimentalClient(client) + ctx := testutil.Context(t, testutil.WaitLong) + + // Insert a bunch of test data. + now := dbtime.Now() + i1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + StartedAt: now.Add(-time.Hour), + }) + i1tok1 := dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: i1.ID, + CreatedAt: now, + }) + i1tok2 := dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: i1.ID, + CreatedAt: now.Add(-time.Minute), + }) + i1up1 := dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + InterceptionID: i1.ID, + CreatedAt: now, + }) + i1up2 := dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + InterceptionID: i1.ID, + CreatedAt: now.Add(-time.Minute), + }) + i1tool1 := dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ + InterceptionID: i1.ID, + CreatedAt: now, + }) + i1tool2 := dbgen.AIBridgeToolUsage(t, db, database.InsertAIBridgeToolUsageParams{ + InterceptionID: i1.ID, + CreatedAt: now.Add(-time.Minute), + }) + i2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + StartedAt: now, + }) + + // Convert to SDK types for response comparison. + // You may notice that the ordering of the inner arrays are ASC, this is + // intentional. + i1SDK := db2sdk.AIBridgeInterception(i1, []database.AIBridgeTokenUsage{i1tok2, i1tok1}, []database.AIBridgeUserPrompt{i1up2, i1up1}, []database.AIBridgeToolUsage{i1tool2, i1tool1}) + i2SDK := db2sdk.AIBridgeInterception(i2, nil, nil, nil) + + res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{}) + require.NoError(t, err) + require.Len(t, res.Results, 2) + require.Equal(t, i2SDK.ID, res.Results[0].ID) + require.Equal(t, i1SDK.ID, res.Results[1].ID) + + // Normalize timestamps in the response so we can compare the whole + // thing easily. + res.Results[0].StartedAt = i2SDK.StartedAt + res.Results[1].StartedAt = i1SDK.StartedAt + require.Len(t, res.Results[1].TokenUsages, 2) + require.Equal(t, i1SDK.TokenUsages[0].ID, res.Results[1].TokenUsages[0].ID) + require.Equal(t, i1SDK.TokenUsages[1].ID, res.Results[1].TokenUsages[1].ID) + res.Results[1].TokenUsages[0].CreatedAt = i1SDK.TokenUsages[0].CreatedAt + res.Results[1].TokenUsages[1].CreatedAt = i1SDK.TokenUsages[1].CreatedAt + require.Len(t, res.Results[1].UserPrompts, 2) + require.Equal(t, i1SDK.UserPrompts[0].ID, res.Results[1].UserPrompts[0].ID) + require.Equal(t, i1SDK.UserPrompts[1].ID, res.Results[1].UserPrompts[1].ID) + res.Results[1].UserPrompts[0].CreatedAt = i1SDK.UserPrompts[0].CreatedAt + res.Results[1].UserPrompts[1].CreatedAt = i1SDK.UserPrompts[1].CreatedAt + require.Len(t, res.Results[1].ToolUsages, 2) + require.Equal(t, i1SDK.ToolUsages[0].ID, res.Results[1].ToolUsages[0].ID) + require.Equal(t, i1SDK.ToolUsages[1].ID, res.Results[1].ToolUsages[1].ID) + res.Results[1].ToolUsages[0].CreatedAt = i1SDK.ToolUsages[0].CreatedAt + res.Results[1].ToolUsages[1].CreatedAt = i1SDK.ToolUsages[1].CreatedAt + + require.Equal(t, []codersdk.AIBridgeInterception{i2SDK, i1SDK}, res.Results) + }) + + t.Run("Pagination", func(t *testing.T) { + t.Parallel() + + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, db, _ := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + experimentalClient := codersdk.NewExperimentalClient(client) + ctx := testutil.Context(t, testutil.WaitLong) + + allInterceptionIDs := make([]uuid.UUID, 0, 20) + + // Create 10 interceptions with the same started_at time. The returned + // order for these should still be deterministic. + now := dbtime.Now() + for i := range 10 { + interception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.UUID{byte(i)}, + StartedAt: now, + }) + allInterceptionIDs = append(allInterceptionIDs, interception.ID) + } + + // Create 10 interceptions with a random started_at time. + for i := range 10 { + randomOffset, err := cryptorand.Intn(10000) + require.NoError(t, err) + randomOffsetDur := time.Duration(randomOffset) * time.Second + interception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.UUID{byte(i + 10)}, + StartedAt: now.Add(randomOffsetDur), + }) + allInterceptionIDs = append(allInterceptionIDs, interception.ID) + } + + // Get all interceptions one by one from the API using cursor + // pagination. + getAllInterceptionsOneByOne := func() []uuid.UUID { + interceptionIDs := []uuid.UUID{} + for { + afterID := uuid.Nil + if len(interceptionIDs) > 0 { + afterID = interceptionIDs[len(interceptionIDs)-1] + } + res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{ + Pagination: codersdk.Pagination{ + AfterID: afterID, + Limit: 1, + }, + }) + require.NoError(t, err) + if len(res.Results) == 0 { + break + } + require.Len(t, res.Results, 1) + interceptionIDs = append(interceptionIDs, res.Results[0].ID) + } + return interceptionIDs + } + + // First attempt: get all interceptions one by one. + gotInterceptionIDs1 := getAllInterceptionsOneByOne() + // We should have all of the interceptions returned: + require.ElementsMatch(t, allInterceptionIDs, gotInterceptionIDs1) + + // Second attempt: get all interceptions one by one again. + gotInterceptionIDs2 := getAllInterceptionsOneByOne() + // They should be returned in the exact same order. + require.Equal(t, gotInterceptionIDs1, gotInterceptionIDs2) + + // Try to get an invalid limit. + res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{ + Pagination: codersdk.Pagination{ + Limit: 1001, + }, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Contains(t, sdkErr.Message, "Invalid pagination limit value.") + require.Empty(t, res.Results) + }) + + t.Run("Authorized", func(t *testing.T) { + t.Parallel() + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + adminClient, db, firstUser := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + adminExperimentalClient := codersdk.NewExperimentalClient(adminClient) + ctx := testutil.Context(t, testutil.WaitLong) + + secondUserClient, secondUser := coderdtest.CreateAnotherUser(t, adminClient, firstUser.OrganizationID) + secondUserExperimentalClient := codersdk.NewExperimentalClient(secondUserClient) + + now := dbtime.Now() + i1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + StartedAt: now, + }) + i2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: secondUser.ID, + StartedAt: now.Add(-time.Hour), + }) + + // Admin can see all interceptions. + res, err := adminExperimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{}) + require.NoError(t, err) + require.Len(t, res.Results, 2) + require.Equal(t, i1.ID, res.Results[0].ID) + require.Equal(t, i2.ID, res.Results[1].ID) + + // Second user can only see their own interceptions. + res, err = secondUserExperimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{}) + require.NoError(t, err) + require.Len(t, res.Results, 1) + require.Equal(t, i2.ID, res.Results[0].ID) + }) + + t.Run("Filter", func(t *testing.T) { + t.Parallel() + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, db, firstUser := coderdenttest.NewWithDatabase(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + experimentalClient := codersdk.NewExperimentalClient(client) + _, secondUser := coderdtest.CreateAnotherUser(t, client, firstUser.OrganizationID) + + // Insert a bunch of test data with varying filterable fields. + now := dbtime.Now() + i1 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.MustParse("00000000-0000-0000-0000-000000000001"), + InitiatorID: firstUser.UserID, + Provider: "one", + Model: "one", + StartedAt: now, + }) + i2 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.MustParse("00000000-0000-0000-0000-000000000002"), + InitiatorID: firstUser.UserID, + Provider: "two", + Model: "two", + StartedAt: now.Add(-time.Hour), + }) + i3 := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + ID: uuid.MustParse("00000000-0000-0000-0000-000000000003"), + InitiatorID: secondUser.ID, + Provider: "three", + Model: "three", + StartedAt: now.Add(-2 * time.Hour), + }) + + // Convert to SDK types for response comparison. We don't care about the + // inner arrays for this test. + i1SDK := db2sdk.AIBridgeInterception(i1, nil, nil, nil) + i2SDK := db2sdk.AIBridgeInterception(i2, nil, nil, nil) + i3SDK := db2sdk.AIBridgeInterception(i3, nil, nil, nil) + + cases := []struct { + name string + filter codersdk.AIBridgeListInterceptionsFilter + want []codersdk.AIBridgeInterception + }{ + { + name: "NoFilter", + filter: codersdk.AIBridgeListInterceptionsFilter{}, + want: []codersdk.AIBridgeInterception{i1SDK, i2SDK, i3SDK}, + }, + { + name: "Initiator/NoMatch", + filter: codersdk.AIBridgeListInterceptionsFilter{Initiator: uuid.New().String()}, + want: []codersdk.AIBridgeInterception{}, + }, + { + name: "Initiator/Me", + filter: codersdk.AIBridgeListInterceptionsFilter{Initiator: codersdk.Me}, + want: []codersdk.AIBridgeInterception{i1SDK, i2SDK}, + }, + { + name: "Initiator/UserID", + filter: codersdk.AIBridgeListInterceptionsFilter{Initiator: secondUser.ID.String()}, + want: []codersdk.AIBridgeInterception{i3SDK}, + }, + { + name: "Initiator/Username", + filter: codersdk.AIBridgeListInterceptionsFilter{Initiator: secondUser.Username}, + want: []codersdk.AIBridgeInterception{i3SDK}, + }, + { + name: "Provider/NoMatch", + filter: codersdk.AIBridgeListInterceptionsFilter{Provider: "nonsense"}, + want: []codersdk.AIBridgeInterception{}, + }, + { + name: "Provider/OK", + filter: codersdk.AIBridgeListInterceptionsFilter{Provider: "two"}, + want: []codersdk.AIBridgeInterception{i2SDK}, + }, + { + name: "Model/NoMatch", + filter: codersdk.AIBridgeListInterceptionsFilter{Model: "nonsense"}, + want: []codersdk.AIBridgeInterception{}, + }, + { + name: "Model/OK", + filter: codersdk.AIBridgeListInterceptionsFilter{Model: "three"}, + want: []codersdk.AIBridgeInterception{i3SDK}, + }, + { + name: "StartedAfter/NoMatch", + filter: codersdk.AIBridgeListInterceptionsFilter{ + StartedAfter: i1.StartedAt.Add(10 * time.Minute), + }, + want: []codersdk.AIBridgeInterception{}, + }, + { + name: "StartedAfter/OK", + filter: codersdk.AIBridgeListInterceptionsFilter{ + StartedAfter: i2.StartedAt.Add(-10 * time.Minute), + }, + want: []codersdk.AIBridgeInterception{i1SDK, i2SDK}, + }, + { + name: "StartedBefore/NoMatch", + filter: codersdk.AIBridgeListInterceptionsFilter{ + StartedBefore: i3.StartedAt.Add(-10 * time.Minute), + }, + want: []codersdk.AIBridgeInterception{}, + }, + { + name: "StartedBefore/OK", + filter: codersdk.AIBridgeListInterceptionsFilter{ + StartedBefore: i3.StartedAt.Add(10 * time.Minute), + }, + want: []codersdk.AIBridgeInterception{i3SDK}, + }, + { + name: "BothBeforeAndAfter/NoMatch", + filter: codersdk.AIBridgeListInterceptionsFilter{ + StartedAfter: i1.StartedAt.Add(10 * time.Minute), + StartedBefore: i1.StartedAt.Add(20 * time.Minute), + }, + want: []codersdk.AIBridgeInterception{}, + }, + { + name: "BothBeforeAndAfter/OK", + filter: codersdk.AIBridgeListInterceptionsFilter{ + StartedAfter: i2.StartedAt.Add(-10 * time.Minute), + StartedBefore: i2.StartedAt.Add(10 * time.Minute), + }, + want: []codersdk.AIBridgeInterception{i2SDK}, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + res, err := experimentalClient.AIBridgeListInterceptions(ctx, tc.filter) + require.NoError(t, err) + // We just compare UUID strings for the sake of this test. + wantIDs := make([]string, len(tc.want)) + for i, r := range tc.want { + wantIDs[i] = r.ID.String() + } + gotIDs := make([]string, len(res.Results)) + for i, r := range res.Results { + gotIDs[i] = r.ID.String() + } + require.Equal(t, wantIDs, gotIDs) + }) + } + }) + + t.Run("FilterErrors", func(t *testing.T) { + t.Parallel() + dv := coderdtest.DeploymentValues(t) + dv.Experiments = []string{string(codersdk.ExperimentAIBridge)} + client, _ := coderdenttest.New(t, &coderdenttest.Options{ + Options: &coderdtest.Options{ + DeploymentValues: dv, + }, + }) + experimentalClient := codersdk.NewExperimentalClient(client) + + // No need to insert any test data, we're just testing the filter + // errors. + + cases := []struct { + name string + q string + want []codersdk.ValidationError + }{ + { + name: "UnknownUsername", + q: "initiator:unknown", + want: []codersdk.ValidationError{ + { + Field: "initiator", + Detail: `Query param "initiator" has invalid value: user "unknown" either does not exist, or you are unauthorized to view them`, + }, + }, + }, + { + name: "InvalidStartedAfter", + q: "started_after:invalid", + want: []codersdk.ValidationError{ + { + Field: "started_after", + Detail: `Query param "started_after" must be a valid date format (2006-01-02T15:04:05.999999999Z07:00): parsing time "INVALID" as "2006-01-02T15:04:05.999999999Z07:00": cannot parse "INVALID" as "2006"`, + }, + }, + }, + { + name: "InvalidStartedBefore", + q: "started_before:invalid", + want: []codersdk.ValidationError{ + { + Field: "started_before", + Detail: `Query param "started_before" must be a valid date format (2006-01-02T15:04:05.999999999Z07:00): parsing time "INVALID" as "2006-01-02T15:04:05.999999999Z07:00": cannot parse "INVALID" as "2006"`, + }, + }, + }, + { + name: "InvalidBeforeAfterRange", + // Before MUST be after After if both are set + q: `started_after:"2025-01-01T00:00:00Z" started_before:"2024-01-01T00:00:00Z"`, + want: []codersdk.ValidationError{ + { + Field: "started_before", + Detail: `Query param "started_before" has invalid value: "started_before" must be after "started_after" if set`, + }, + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + res, err := experimentalClient.AIBridgeListInterceptions(ctx, codersdk.AIBridgeListInterceptionsFilter{ + FilterQuery: tc.q, + }) + var sdkErr *codersdk.Error + require.ErrorAs(t, err, &sdkErr) + require.Equal(t, tc.want, sdkErr.Validations) + require.Empty(t, res.Results) + }) + } + }) +} diff --git a/enterprise/coderd/aibridged.go b/enterprise/coderd/aibridged.go index 51afb24e6a..bf991103b1 100644 --- a/enterprise/coderd/aibridged.go +++ b/enterprise/coderd/aibridged.go @@ -6,16 +6,13 @@ import ( "io" "net/http" - "github.com/go-chi/chi/v5" "golang.org/x/xerrors" "storj.io/drpc/drpcmux" "storj.io/drpc/drpcserver" "cdr.dev/slog" - "github.com/coder/coder/v2/coderd/httpmw" "github.com/coder/coder/v2/coderd/tracing" - "github.com/coder/coder/v2/codersdk" "github.com/coder/coder/v2/codersdk/drpcsdk" "github.com/coder/coder/v2/enterprise/x/aibridged" aibridgedproto "github.com/coder/coder/v2/enterprise/x/aibridged/proto" @@ -25,24 +22,12 @@ import ( // RegisterInMemoryAIBridgedHTTPHandler mounts [aibridged.Server]'s HTTP router onto // [API]'s router, so that requests to aibridged will be relayed from Coder's API server // to the in-memory aibridged. -func (api *API) RegisterInMemoryAIBridgedHTTPHandler(srv *aibridged.Server) { +func (api *API) RegisterInMemoryAIBridgedHTTPHandler(srv http.Handler) { if srv == nil { panic("aibridged cannot be nil") } - if api.AGPL.RootHandler == nil { - panic("api.RootHandler cannot be nil") - } - - aibridgeEndpoint := "/api/experimental/aibridge" - - r := chi.NewRouter() - r.Group(func(r chi.Router) { - r.Use(httpmw.RequireExperiment(api.AGPL.Experiments, codersdk.ExperimentAIBridge)) - r.HandleFunc("/*", http.StripPrefix(aibridgeEndpoint, srv).ServeHTTP) - }) - - api.AGPL.RootHandler.Mount(aibridgeEndpoint, r) + api.aibridgedHandler = srv } // CreateInMemoryAIBridgeServer creates a [aibridged.DRPCServer] and returns a diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go index 0d276eef86..73585a8e6b 100644 --- a/enterprise/coderd/coderd.go +++ b/enterprise/coderd/coderd.go @@ -226,6 +226,30 @@ func New(ctx context.Context, options *Options) (_ *API, err error) { return api.refreshEntitlements(ctx) } + api.AGPL.ExperimentalHandler.Group(func(r chi.Router) { + r.Route("/aibridge", func(r chi.Router) { + r.Use( + httpmw.RequireExperimentWithDevBypass(api.AGPL.Experiments, codersdk.ExperimentAIBridge), + ) + r.Group(func(r chi.Router) { + r.Use(apiKeyMiddleware) + r.Get("/interceptions", api.aiBridgeListInterceptions) + }) + + // This is a bit funky but since aibridge only exposes a HTTP + // handler, this is how it has to be. + r.HandleFunc("/*", func(rw http.ResponseWriter, r *http.Request) { + if api.aibridgedHandler == nil { + httpapi.Write(r.Context(), rw, http.StatusNotFound, codersdk.Response{ + Message: "aibridged handler not mounted", + }) + return + } + api.aibridgedHandler.ServeHTTP(rw, r) + }) + }) + }) + api.AGPL.APIHandler.Group(func(r chi.Router) { r.Get("/entitlements", api.serveEntitlements) // /regions overrides the AGPL /regions endpoint @@ -677,6 +701,8 @@ type API struct { licenseMetricsCollector *license.MetricsCollector tailnetService *tailnet.ClientService + + aibridgedHandler http.Handler } // writeEntitlementWarningsHeader writes the entitlement warnings to the response header diff --git a/enterprise/x/aibridgedserver/aibridgedserver.go b/enterprise/x/aibridgedserver/aibridgedserver.go index 9f36d90517..ee79cd15f0 100644 --- a/enterprise/x/aibridgedserver/aibridgedserver.go +++ b/enterprise/x/aibridgedserver/aibridgedserver.go @@ -53,9 +53,9 @@ var _ aibridged.DRPCServer = &Server{} type store interface { // Recorder-related queries. InsertAIBridgeInterception(ctx context.Context, arg database.InsertAIBridgeInterceptionParams) (database.AIBridgeInterception, error) - InsertAIBridgeTokenUsage(ctx context.Context, arg database.InsertAIBridgeTokenUsageParams) error - InsertAIBridgeUserPrompt(ctx context.Context, arg database.InsertAIBridgeUserPromptParams) error - InsertAIBridgeToolUsage(ctx context.Context, arg database.InsertAIBridgeToolUsageParams) error + 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) // MCPConfigurator-related queries. GetExternalAuthLinksByUserID(ctx context.Context, userID uuid.UUID) ([]database.ExternalAuthLink, error) @@ -139,7 +139,7 @@ func (s *Server) RecordTokenUsage(ctx context.Context, in *proto.RecordTokenUsag return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err) } - err = s.store.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{ + _, err = s.store.InsertAIBridgeTokenUsage(ctx, database.InsertAIBridgeTokenUsageParams{ ID: uuid.New(), InterceptionID: intcID, ProviderResponseID: in.GetMsgId(), @@ -163,7 +163,7 @@ func (s *Server) RecordPromptUsage(ctx context.Context, in *proto.RecordPromptUs return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err) } - err = s.store.InsertAIBridgeUserPrompt(ctx, database.InsertAIBridgeUserPromptParams{ + _, err = s.store.InsertAIBridgeUserPrompt(ctx, database.InsertAIBridgeUserPromptParams{ ID: uuid.New(), InterceptionID: intcID, ProviderResponseID: in.GetMsgId(), @@ -186,7 +186,7 @@ func (s *Server) RecordToolUsage(ctx context.Context, in *proto.RecordToolUsageR return nil, xerrors.Errorf("failed to parse interception_id %q: %w", in.GetInterceptionId(), err) } - err = s.store.InsertAIBridgeToolUsage(ctx, database.InsertAIBridgeToolUsageParams{ + _, err = s.store.InsertAIBridgeToolUsage(ctx, database.InsertAIBridgeToolUsageParams{ ID: uuid.New(), InterceptionID: intcID, ProviderResponseID: in.GetMsgId(), diff --git a/enterprise/x/aibridgedserver/aibridgedserver_test.go b/enterprise/x/aibridgedserver/aibridgedserver_test.go index fcf8deea47..03fec9398b 100644 --- a/enterprise/x/aibridgedserver/aibridgedserver_test.go +++ b/enterprise/x/aibridgedserver/aibridgedserver_test.go @@ -464,7 +464,18 @@ func TestRecordTokenUsage(t *testing.T) { return false } return true - })).Return(nil) + })).Return(database.AIBridgeTokenUsage{ + ID: uuid.New(), + InterceptionID: interceptionID, + ProviderResponseID: req.GetMsgId(), + InputTokens: req.GetInputTokens(), + OutputTokens: req.GetOutputTokens(), + Metadata: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(metadataJSON), + Valid: true, + }, + CreatedAt: req.GetCreatedAt().AsTime(), + }, nil) }, }, { @@ -488,7 +499,7 @@ func TestRecordTokenUsage(t *testing.T) { CreatedAt: timestamppb.Now(), }, setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordTokenUsageRequest) { - db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).Return(sql.ErrConnDone) + db.EXPECT().InsertAIBridgeTokenUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeTokenUsage{}, sql.ErrConnDone) }, expectedErr: "insert token usage", }, @@ -534,7 +545,17 @@ func TestRecordPromptUsage(t *testing.T) { return false } return true - })).Return(nil) + })).Return(database.AIBridgeUserPrompt{ + ID: uuid.New(), + InterceptionID: interceptionID, + ProviderResponseID: req.GetMsgId(), + Prompt: req.GetPrompt(), + Metadata: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(metadataJSON), + Valid: true, + }, + CreatedAt: req.GetCreatedAt().AsTime(), + }, nil) }, }, { @@ -556,7 +577,7 @@ func TestRecordPromptUsage(t *testing.T) { CreatedAt: timestamppb.Now(), }, setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordPromptUsageRequest) { - db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), gomock.Any()).Return(sql.ErrConnDone) + db.EXPECT().InsertAIBridgeUserPrompt(gomock.Any(), gomock.Any()).Return(database.AIBridgeUserPrompt{}, sql.ErrConnDone) }, expectedErr: "insert user prompt", }, @@ -622,7 +643,21 @@ func TestRecordToolUsage(t *testing.T) { return false } return true - })).Return(nil) + })).Return(database.AIBridgeToolUsage{ + ID: uuid.New(), + InterceptionID: interceptionID, + ProviderResponseID: req.GetMsgId(), + Tool: req.GetTool(), + ServerUrl: dbServerURL, + Input: req.GetInput(), + Injected: req.GetInjected(), + InvocationError: dbInvocationError, + Metadata: pqtype.NullRawMessage{ + RawMessage: json.RawMessage(metadataJSON), + Valid: true, + }, + CreatedAt: req.GetCreatedAt().AsTime(), + }, nil) }, }, { @@ -646,7 +681,7 @@ func TestRecordToolUsage(t *testing.T) { CreatedAt: timestamppb.Now(), }, setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordToolUsageRequest) { - db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), gomock.Any()).Return(sql.ErrConnDone) + db.EXPECT().InsertAIBridgeToolUsage(gomock.Any(), gomock.Any()).Return(database.AIBridgeToolUsage{}, sql.ErrConnDone) }, expectedErr: "insert tool usage", }, diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 651cf1f325..b480408fd4 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -19,12 +19,69 @@ export interface AIBridgeConfig { readonly anthropic: AIBridgeAnthropicConfig; } +// From codersdk/aibridge.go +export interface AIBridgeInterception { + readonly id: string; + readonly initiator_id: string; + readonly provider: string; + readonly model: string; + // empty interface{} type, falling back to unknown + readonly metadata: Record; + readonly started_at: string; + readonly token_usages: readonly AIBridgeTokenUsage[]; + readonly user_prompts: readonly AIBridgeUserPrompt[]; + readonly tool_usages: readonly AIBridgeToolUsage[]; +} + +// From codersdk/aibridge.go +export interface AIBridgeListInterceptionsResponse { + readonly results: readonly AIBridgeInterception[]; +} + // From codersdk/deployment.go export interface AIBridgeOpenAIConfig { readonly base_url: string; readonly key: string; } +// From codersdk/aibridge.go +export interface AIBridgeTokenUsage { + readonly id: string; + readonly interception_id: string; + readonly provider_response_id: string; + readonly input_tokens: number; + readonly output_tokens: number; + // empty interface{} type, falling back to unknown + readonly metadata: Record; + readonly created_at: string; +} + +// From codersdk/aibridge.go +export interface AIBridgeToolUsage { + readonly id: string; + readonly interception_id: string; + readonly provider_response_id: string; + readonly server_url: string; + readonly tool: string; + readonly input: string; + readonly injected: boolean; + readonly invocation_error: string; + // empty interface{} type, falling back to unknown + readonly metadata: Record; + readonly created_at: string; +} + +// From codersdk/aibridge.go +export interface AIBridgeUserPrompt { + readonly id: string; + readonly interception_id: string; + readonly provider_response_id: string; + readonly prompt: string; + // empty interface{} type, falling back to unknown + readonly metadata: Record; + readonly created_at: string; +} + // From codersdk/deployment.go export interface AIConfig { readonly bridge?: AIBridgeConfig;