diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 7ef90f7e8a..78481a71f4 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -84,6 +84,34 @@ const docTemplate = `{ } } }, + "/aibridge/clients": { + "get": { + "produces": [ + "application/json" + ], + "tags": [ + "AI Bridge" + ], + "summary": "List AI Bridge clients", + "operationId": "list-ai-bridge-clients", + "responses": { + "200": { + "description": "OK", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/aibridge/interceptions": { "get": { "produces": [ diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 0d220380ed..b6c8ae7490 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -65,6 +65,30 @@ } } }, + "/aibridge/clients": { + "get": { + "produces": ["application/json"], + "tags": ["AI Bridge"], + "summary": "List AI Bridge clients", + "operationId": "list-ai-bridge-clients", + "responses": { + "200": { + "description": "OK", + "schema": { + "type": "array", + "items": { + "type": "string" + } + } + } + }, + "security": [ + { + "CoderSessionToken": [] + } + ] + } + }, "/aibridge/interceptions": { "get": { "produces": ["application/json"], diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index ea3284664b..bf75a2595a 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -5356,6 +5356,14 @@ func (q *querier) InsertWorkspaceResourceMetadata(ctx context.Context, arg datab return q.db.InsertWorkspaceResourceMetadata(ctx, arg) } +func (q *querier) ListAIBridgeClients(ctx context.Context, arg database.ListAIBridgeClientsParams) ([]string, 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.ListAuthorizedAIBridgeClients(ctx, arg, prep) +} + func (q *querier) ListAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams) ([]database.ListAIBridgeInterceptionsRow, error) { prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceAibridgeInterception.Type) if err != nil { @@ -7302,6 +7310,14 @@ func (q *querier) ListAuthorizedAIBridgeModels(ctx context.Context, arg database return q.ListAIBridgeModels(ctx, arg) } +func (q *querier) ListAuthorizedAIBridgeClients(ctx context.Context, arg database.ListAIBridgeClientsParams, _ rbac.PreparedAuthorized) ([]string, error) { + // TODO: Delete this function, all ListAIBridgeClients should be + // authorized. For now just call ListAIBridgeClients 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.ListAIBridgeClients(ctx, arg) +} + func (q *querier) ListAuthorizedAIBridgeSessions(ctx context.Context, arg database.ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]database.ListAIBridgeSessionsRow, error) { return q.db.ListAuthorizedAIBridgeSessions(ctx, arg, prepared) } diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index a76e892b44..4e67955f3d 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -5542,6 +5542,20 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(params, emptyPreparedAuthorized{}).Asserts() })) + s.Run("ListAIBridgeClients", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.ListAIBridgeClientsParams{} + db.EXPECT().ListAuthorizedAIBridgeClients(gomock.Any(), params, gomock.Any()).Return([]string{}, nil).AnyTimes() + // No asserts here because SQLFilter. + check.Args(params).Asserts() + })) + + s.Run("ListAuthorizedAIBridgeClients", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.ListAIBridgeClientsParams{} + db.EXPECT().ListAuthorizedAIBridgeClients(gomock.Any(), params, gomock.Any()).Return([]string{}, nil).AnyTimes() + // No asserts here because SQLFilter. + check.Args(params, emptyPreparedAuthorized{}).Asserts() + })) + s.Run("ListAIBridgeSessions", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { params := database.ListAIBridgeSessionsParams{} db.EXPECT().ListAuthorizedAIBridgeSessions(gomock.Any(), params, gomock.Any()).Return([]database.ListAIBridgeSessionsRow{}, nil).AnyTimes() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index 1b8707a08a..1ed0237ef5 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -3752,6 +3752,14 @@ func (m queryMetricsStore) InsertWorkspaceResourceMetadata(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) ListAIBridgeClients(ctx context.Context, arg database.ListAIBridgeClientsParams) ([]string, error) { + start := time.Now() + r0, r1 := m.s.ListAIBridgeClients(ctx, arg) + m.queryLatencies.WithLabelValues("ListAIBridgeClients").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListAIBridgeClients").Inc() + return r0, r1 +} + func (m queryMetricsStore) ListAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams) ([]database.ListAIBridgeInterceptionsRow, error) { start := time.Now() r0, r1 := m.s.ListAIBridgeInterceptions(ctx, arg) @@ -5296,6 +5304,14 @@ func (m queryMetricsStore) ListAuthorizedAIBridgeModels(ctx context.Context, arg return r0, r1 } +func (m queryMetricsStore) ListAuthorizedAIBridgeClients(ctx context.Context, arg database.ListAIBridgeClientsParams, prepared rbac.PreparedAuthorized) ([]string, error) { + start := time.Now() + r0, r1 := m.s.ListAuthorizedAIBridgeClients(ctx, arg, prepared) + m.queryLatencies.WithLabelValues("ListAuthorizedAIBridgeClients").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "ListAuthorizedAIBridgeClients").Inc() + return r0, r1 +} + func (m queryMetricsStore) ListAuthorizedAIBridgeSessions(ctx context.Context, arg database.ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]database.ListAIBridgeSessionsRow, error) { start := time.Now() r0, r1 := m.s.ListAuthorizedAIBridgeSessions(ctx, arg, prepared) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 3ea7b17ce7..d789e2b1f5 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -7022,6 +7022,21 @@ func (mr *MockStoreMockRecorder) InsertWorkspaceResourceMetadata(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "InsertWorkspaceResourceMetadata", reflect.TypeOf((*MockStore)(nil).InsertWorkspaceResourceMetadata), ctx, arg) } +// ListAIBridgeClients mocks base method. +func (m *MockStore) ListAIBridgeClients(ctx context.Context, arg database.ListAIBridgeClientsParams) ([]string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAIBridgeClients", ctx, arg) + ret0, _ := ret[0].([]string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAIBridgeClients indicates an expected call of ListAIBridgeClients. +func (mr *MockStoreMockRecorder) ListAIBridgeClients(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeClients", reflect.TypeOf((*MockStore)(nil).ListAIBridgeClients), ctx, arg) +} + // ListAIBridgeInterceptions mocks base method. func (m *MockStore) ListAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams) ([]database.ListAIBridgeInterceptionsRow, error) { m.ctrl.T.Helper() @@ -7157,6 +7172,21 @@ func (mr *MockStoreMockRecorder) ListAIBridgeUserPromptsByInterceptionIDs(ctx, i return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAIBridgeUserPromptsByInterceptionIDs", reflect.TypeOf((*MockStore)(nil).ListAIBridgeUserPromptsByInterceptionIDs), ctx, interceptionIds) } +// ListAuthorizedAIBridgeClients mocks base method. +func (m *MockStore) ListAuthorizedAIBridgeClients(ctx context.Context, arg database.ListAIBridgeClientsParams, prepared rbac.PreparedAuthorized) ([]string, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "ListAuthorizedAIBridgeClients", ctx, arg, prepared) + ret0, _ := ret[0].([]string) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// ListAuthorizedAIBridgeClients indicates an expected call of ListAuthorizedAIBridgeClients. +func (mr *MockStoreMockRecorder) ListAuthorizedAIBridgeClients(ctx, arg, prepared any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ListAuthorizedAIBridgeClients", reflect.TypeOf((*MockStore)(nil).ListAuthorizedAIBridgeClients), ctx, arg, prepared) +} + // ListAuthorizedAIBridgeInterceptions mocks base method. func (m *MockStore) ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg database.ListAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) ([]database.ListAIBridgeInterceptionsRow, error) { m.ctrl.T.Helper() diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index 9e2f00fc4f..4409ddc9d0 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -813,6 +813,7 @@ type aibridgeQuerier interface { ListAuthorizedAIBridgeInterceptions(ctx context.Context, arg ListAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeInterceptionsRow, error) CountAuthorizedAIBridgeInterceptions(ctx context.Context, arg CountAIBridgeInterceptionsParams, prepared rbac.PreparedAuthorized) (int64, error) ListAuthorizedAIBridgeModels(ctx context.Context, arg ListAIBridgeModelsParams, prepared rbac.PreparedAuthorized) ([]string, error) + ListAuthorizedAIBridgeClients(ctx context.Context, arg ListAIBridgeClientsParams, prepared rbac.PreparedAuthorized) ([]string, error) ListAuthorizedAIBridgeSessions(ctx context.Context, arg ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionsRow, error) CountAuthorizedAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) (int64, error) ListAuthorizedAIBridgeSessionThreads(ctx context.Context, arg ListAIBridgeSessionThreadsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionThreadsRow, error) @@ -950,6 +951,35 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeModels(ctx context.Context, arg ListA return items, nil } +func (q *sqlQuerier) ListAuthorizedAIBridgeClients(ctx context.Context, arg ListAIBridgeClientsParams, prepared rbac.PreparedAuthorized) ([]string, 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(listAIBridgeClients, fmt.Sprintf(" AND %s", authorizedFilter)) + if err != nil { + return nil, xerrors.Errorf("insert authorized filter: %w", err) + } + + query := fmt.Sprintf("-- name: ListAIBridgeClients :many\n%s", filtered) + rows, err := q.db.QueryContext(ctx, query, arg.Client, arg.Offset, arg.Limit) + if err != nil { + return nil, err + } + defer rows.Close() + var items []string + for rows.Next() { + var client string + if err := rows.Scan(&client); err != nil { + return nil, err + } + items = append(items, client) + } + return items, nil +} + func (q *sqlQuerier) ListAuthorizedAIBridgeSessions(ctx context.Context, arg ListAIBridgeSessionsParams, prepared rbac.PreparedAuthorized) ([]ListAIBridgeSessionsRow, error) { authorizedFilter, err := prepared.CompileToSQL(ctx, regosql.ConvertConfig{ VariableConverter: regosql.AIBridgeInterceptionConverter(), diff --git a/coderd/database/querier.go b/coderd/database/querier.go index d45b564173..b6a7fc07ce 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -775,6 +775,7 @@ 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) + ListAIBridgeClients(ctx context.Context, arg ListAIBridgeClientsParams) ([]string, error) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBridgeInterceptionsParams) ([]ListAIBridgeInterceptionsRow, error) // Finds all unique AI Bridge interception telemetry summaries combinations // (provider, model, client) in the given timeframe for telemetry reporting. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 4280a44e68..1f78c9e86b 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -909,6 +909,58 @@ func (q *sqlQuerier) InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIB return i, err } +const listAIBridgeClients = `-- name: ListAIBridgeClients :many +SELECT + COALESCE(client, 'Unknown') AS client +FROM + aibridge_interceptions +WHERE + ended_at IS NOT NULL + -- Filter client (prefix match to allow B-tree index usage). + AND CASE + WHEN $1::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') LIKE $1::text || '%' + ELSE true + END + -- We use an ` + "`" + `@authorize_filter` + "`" + ` as we are attempting to list clients + -- that are relevant to the user and what they are allowed to see. + -- Authorize Filter clause will be injected below in + -- ListAIBridgeClientsAuthorized. + -- @authorize_filter +GROUP BY + client +LIMIT COALESCE(NULLIF($3::integer, 0), 100) +OFFSET $2 +` + +type ListAIBridgeClientsParams struct { + Client string `db:"client" json:"client"` + Offset int32 `db:"offset_" json:"offset_"` + Limit int32 `db:"limit_" json:"limit_"` +} + +func (q *sqlQuerier) ListAIBridgeClients(ctx context.Context, arg ListAIBridgeClientsParams) ([]string, error) { + rows, err := q.db.QueryContext(ctx, listAIBridgeClients, arg.Client, arg.Offset, arg.Limit) + if err != nil { + return nil, err + } + defer rows.Close() + var items []string + for rows.Next() { + var client string + if err := rows.Scan(&client); err != nil { + return nil, err + } + items = append(items, client) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const listAIBridgeInterceptions = `-- name: ListAIBridgeInterceptions :many SELECT aibridge_interceptions.id, aibridge_interceptions.initiator_id, aibridge_interceptions.provider, aibridge_interceptions.model, aibridge_interceptions.started_at, aibridge_interceptions.metadata, aibridge_interceptions.ended_at, aibridge_interceptions.api_key_id, aibridge_interceptions.client, aibridge_interceptions.thread_parent_id, aibridge_interceptions.thread_root_id, aibridge_interceptions.client_session_id, aibridge_interceptions.session_id, diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 0d62ddb301..f4f03ff1cc 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -680,3 +680,27 @@ ORDER BY LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100) OFFSET @offset_ ; + + +-- name: ListAIBridgeClients :many +SELECT + COALESCE(client, 'Unknown') AS client +FROM + aibridge_interceptions +WHERE + ended_at IS NOT NULL + -- Filter client (prefix match to allow B-tree index usage). + AND CASE + WHEN @client::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') LIKE @client::text || '%' + ELSE true + END + -- We use an `@authorize_filter` as we are attempting to list clients + -- that are relevant to the user and what they are allowed to see. + -- Authorize Filter clause will be injected below in + -- ListAIBridgeClientsAuthorized. + -- @authorize_filter +GROUP BY + client +LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100) +OFFSET @offset_ +; diff --git a/coderd/searchquery/search.go b/coderd/searchquery/search.go index 814ba985ab..330c2e6eb4 100644 --- a/coderd/searchquery/search.go +++ b/coderd/searchquery/search.go @@ -474,6 +474,34 @@ func AIBridgeModels(query string, page codersdk.Pagination) (database.ListAIBrid return filter, parser.Errors } +func AIBridgeClients(query string, page codersdk.Pagination) (database.ListAIBridgeClientsParams, []codersdk.ValidationError) { + // nolint:exhaustruct // Empty values just means "don't filter by that field". + filter := database.ListAIBridgeClientsParams{ + // #nosec G115 - Safe conversion for pagination offset which is expected to be within int32 range + Offset: int32(page.Offset), + // #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 { + values.Add("client", term) + return nil + }) + if len(errors) > 0 { + return filter, errors + } + + parser := httpapi.NewQueryParamParser() + filter.Client = parser.String(values, "", "client") + + parser.ErrorExcessParams(values) + return filter, parser.Errors +} + // Tasks parses a search query for tasks. // // Supported query parameters: diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 1b80b73488..4c382940af 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -324,3 +324,17 @@ func (c *Client) AIBridgeGetSessionThreads(ctx context.Context, sessionID string var resp AIBridgeSessionThreadsResponse return resp, json.NewDecoder(res.Body).Decode(&resp) } + +// AIBridgeListClients returns the distinct AI clients visible to the caller. +func (c *Client) AIBridgeListClients(ctx context.Context) ([]string, error) { + res, err := c.Request(ctx, http.MethodGet, "/api/v2/aibridge/clients", nil) + if err != nil { + return nil, err + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + return nil, ReadBodyAsError(res) + } + var clients []string + return clients, json.NewDecoder(res.Body).Decode(&clients) +} diff --git a/docs/reference/api/aibridge.md b/docs/reference/api/aibridge.md index 6852e42d91..8d69db7c72 100644 --- a/docs/reference/api/aibridge.md +++ b/docs/reference/api/aibridge.md @@ -1,5 +1,38 @@ # AI Bridge +## List AI Bridge clients + +### Code samples + +```shell +# Example request using curl +curl -X GET http://coder-server:8080/api/v2/aibridge/clients \ + -H 'Accept: application/json' \ + -H 'Coder-Session-Token: API_KEY' +``` + +`GET /aibridge/clients` + +### Example responses + +> 200 Response + +```json +[ + "string" +] +``` + +### Responses + +| Status | Meaning | Description | Schema | +|--------|---------------------------------------------------------|-------------|-----------------| +| 200 | [OK](https://tools.ietf.org/html/rfc7231#section-6.3.1) | OK | array of string | + +