From aaa5bce3522e676a5c36c947d34d41fa30a6ab8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Banaszewski?= Date: Wed, 17 Jun 2026 16:30:59 +0000 Subject: [PATCH] feat(coderd/database): add AI Gateway key auth lookup and last-used queries Add GetAIGatewayKeyByHashedSecret for authenticating standalone AI Gateway replicas, and UpdateAIGatewayKeyLastUsedAt to record liveness for active DRPC sessions. The last-used write is an internal system operation authorized against ResourceSystem, keeping AI Gateway keys immutable from a user's perspective. Part of AIGOV-308. Generated with Coder Agents. --- coderd/database/dbauthz/dbauthz.go | 18 +++++++++++ coderd/database/dbauthz/dbauthz_test.go | 11 +++++++ coderd/database/dbmetrics/querymetrics.go | 16 +++++++++ coderd/database/dbmock/dbmock.go | 29 +++++++++++++++++ coderd/database/querier.go | 2 ++ coderd/database/queries.sql.go | 36 +++++++++++++++++++++ coderd/database/queries/ai_gateway_keys.sql | 10 ++++++ 7 files changed, 122 insertions(+) diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 47f38d4d0a74f..8c6ef430d535f 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2753,6 +2753,13 @@ func (q *querier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, in return q.db.GetAIBridgeUserPromptsByInterceptionID(ctx, interceptionID) } +func (q *querier) GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (database.AIGatewayKey, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAIGatewayKey); err != nil { + return database.AIGatewayKey{}, err + } + return q.db.GetAIGatewayKeyByHashedSecret(ctx, hashedSecret) +} + func (q *querier) GetAIModelPriceByProviderModel(ctx context.Context, arg database.GetAIModelPriceByProviderModelParams) (database.AIModelPrice, error) { if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAiModelPrice); err != nil { return database.AIModelPrice{}, err @@ -6972,6 +6979,17 @@ func (q *querier) UpdateAIBridgeInterceptionEnded(ctx context.Context, params da return q.db.UpdateAIBridgeInterceptionEnded(ctx, params) } +// UpdateAIGatewayKeyLastUsedAt records liveness for an active Gateway DRPC +// session. It is an internal system write, not a user-facing mutation, so it +// authorizes against ResourceSystem to keep AI Gateway keys immutable from a +// user's perspective. +func (q *querier) UpdateAIGatewayKeyLastUsedAt(ctx context.Context, arg database.UpdateAIGatewayKeyLastUsedAtParams) error { + if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceSystem); err != nil { + return err + } + return q.db.UpdateAIGatewayKeyLastUsedAt(ctx, arg) +} + func (q *querier) UpdateAIProvider(ctx context.Context, arg database.UpdateAIProviderParams) (database.AIProvider, error) { if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIProvider); err != nil { return database.AIProvider{}, err diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 9565455ffb173..c48ee076762c0 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6920,6 +6920,17 @@ func (s *MethodTestSuite) TestAIBridge() { dbm.EXPECT().DeleteAIGatewayKey(gomock.Any(), id).Return(database.DeleteAIGatewayKeyRow{}, nil).AnyTimes() check.Args(id).Asserts(rbac.ResourceAIGatewayKey, policy.ActionDelete).Returns(database.DeleteAIGatewayKeyRow{}) })) + s.Run("GetAIGatewayKeyByHashedSecret", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + hashedSecret := []byte("hashed-secret") + row := database.AIGatewayKey{} + dbm.EXPECT().GetAIGatewayKeyByHashedSecret(gomock.Any(), hashedSecret).Return(row, nil).AnyTimes() + check.Args(hashedSecret).Asserts(rbac.ResourceAIGatewayKey, policy.ActionRead).Returns(row) + })) + s.Run("UpdateAIGatewayKeyLastUsedAt", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) { + params := database.UpdateAIGatewayKeyLastUsedAtParams{ID: uuid.New()} + dbm.EXPECT().UpdateAIGatewayKeyLastUsedAt(gomock.Any(), params).Return(nil).AnyTimes() + check.Args(params).Asserts(rbac.ResourceSystem, policy.ActionUpdate).Returns() + })) } func (s *MethodTestSuite) TestTelemetry() { diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index ee9e80d781566..61dab877e63c0 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1130,6 +1130,14 @@ func (m queryMetricsStore) GetAIBridgeUserPromptsByInterceptionID(ctx context.Co return r0, r1 } +func (m queryMetricsStore) GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (database.AIGatewayKey, error) { + start := time.Now() + r0, r1 := m.s.GetAIGatewayKeyByHashedSecret(ctx, hashedSecret) + m.queryLatencies.WithLabelValues("GetAIGatewayKeyByHashedSecret").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIGatewayKeyByHashedSecret").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetAIModelPriceByProviderModel(ctx context.Context, arg database.GetAIModelPriceByProviderModelParams) (database.AIModelPrice, error) { start := time.Now() r0, r1 := m.s.GetAIModelPriceByProviderModel(ctx, arg) @@ -5026,6 +5034,14 @@ func (m queryMetricsStore) UpdateAIBridgeInterceptionEnded(ctx context.Context, return r0, r1 } +func (m queryMetricsStore) UpdateAIGatewayKeyLastUsedAt(ctx context.Context, arg database.UpdateAIGatewayKeyLastUsedAtParams) error { + start := time.Now() + r0 := m.s.UpdateAIGatewayKeyLastUsedAt(ctx, arg) + m.queryLatencies.WithLabelValues("UpdateAIGatewayKeyLastUsedAt").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "UpdateAIGatewayKeyLastUsedAt").Inc() + return r0 +} + func (m queryMetricsStore) UpdateAIProvider(ctx context.Context, arg database.UpdateAIProviderParams) (database.AIProvider, error) { start := time.Now() r0, r1 := m.s.UpdateAIProvider(ctx, arg) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index 7c7c3aec7eb8d..98b949d73f548 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1947,6 +1947,21 @@ func (mr *MockStoreMockRecorder) GetAIBridgeUserPromptsByInterceptionID(ctx, int return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBridgeUserPromptsByInterceptionID", reflect.TypeOf((*MockStore)(nil).GetAIBridgeUserPromptsByInterceptionID), ctx, interceptionID) } +// GetAIGatewayKeyByHashedSecret mocks base method. +func (m *MockStore) GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (database.AIGatewayKey, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAIGatewayKeyByHashedSecret", ctx, hashedSecret) + ret0, _ := ret[0].(database.AIGatewayKey) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAIGatewayKeyByHashedSecret indicates an expected call of GetAIGatewayKeyByHashedSecret. +func (mr *MockStoreMockRecorder) GetAIGatewayKeyByHashedSecret(ctx, hashedSecret any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIGatewayKeyByHashedSecret", reflect.TypeOf((*MockStore)(nil).GetAIGatewayKeyByHashedSecret), ctx, hashedSecret) +} + // GetAIModelPriceByProviderModel mocks base method. func (m *MockStore) GetAIModelPriceByProviderModel(ctx context.Context, arg database.GetAIModelPriceByProviderModelParams) (database.AIModelPrice, error) { m.ctrl.T.Helper() @@ -9473,6 +9488,20 @@ func (mr *MockStoreMockRecorder) UpdateAIBridgeInterceptionEnded(ctx, arg any) * return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateAIBridgeInterceptionEnded", reflect.TypeOf((*MockStore)(nil).UpdateAIBridgeInterceptionEnded), ctx, arg) } +// UpdateAIGatewayKeyLastUsedAt mocks base method. +func (m *MockStore) UpdateAIGatewayKeyLastUsedAt(ctx context.Context, arg database.UpdateAIGatewayKeyLastUsedAtParams) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "UpdateAIGatewayKeyLastUsedAt", ctx, arg) + ret0, _ := ret[0].(error) + return ret0 +} + +// UpdateAIGatewayKeyLastUsedAt indicates an expected call of UpdateAIGatewayKeyLastUsedAt. +func (mr *MockStoreMockRecorder) UpdateAIGatewayKeyLastUsedAt(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateAIGatewayKeyLastUsedAt", reflect.TypeOf((*MockStore)(nil).UpdateAIGatewayKeyLastUsedAt), ctx, arg) +} + // UpdateAIProvider mocks base method. func (m *MockStore) UpdateAIProvider(ctx context.Context, arg database.UpdateAIProviderParams) (database.AIProvider, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 46452555302fd..96978f532dd27 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -301,6 +301,7 @@ type sqlcQuerier interface { GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeTokenUsage, error) GetAIBridgeToolUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeToolUsage, error) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeUserPrompt, error) + GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (AIGatewayKey, error) GetAIModelPriceByProviderModel(ctx context.Context, arg GetAIModelPriceByProviderModelParams) (AIModelPrice, error) GetAIProviderByID(ctx context.Context, id uuid.UUID) (AIProvider, error) // Lock the provider row until the model-config write completes. The @@ -1292,6 +1293,7 @@ type sqlcQuerier interface { UnpinChatByID(ctx context.Context, id uuid.UUID) error UnsetDefaultChatModelConfigs(ctx context.Context) error UpdateAIBridgeInterceptionEnded(ctx context.Context, arg UpdateAIBridgeInterceptionEndedParams) (AIBridgeInterception, error) + UpdateAIGatewayKeyLastUsedAt(ctx context.Context, arg UpdateAIGatewayKeyLastUsedAtParams) error UpdateAIProvider(ctx context.Context, arg UpdateAIProviderParams) (AIProvider, error) UpdateAPIKeyByID(ctx context.Context, arg UpdateAPIKeyByIDParams) error UpdateChatACLByID(ctx context.Context, arg UpdateChatACLByIDParams) error diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index cc5301a857a6a..9f1d8c82cf098 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -137,6 +137,26 @@ func (q *sqlQuerier) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) (Dele return i, err } +const getAIGatewayKeyByHashedSecret = `-- name: GetAIGatewayKeyByHashedSecret :one +SELECT id, created_at, name, secret_prefix, hashed_secret, last_used_at +FROM ai_gateway_keys +WHERE hashed_secret = $1 +` + +func (q *sqlQuerier) GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (AIGatewayKey, error) { + row := q.db.QueryRowContext(ctx, getAIGatewayKeyByHashedSecret, hashedSecret) + var i AIGatewayKey + err := row.Scan( + &i.ID, + &i.CreatedAt, + &i.Name, + &i.SecretPrefix, + &i.HashedSecret, + &i.LastUsedAt, + ) + return i, err +} + const insertAIGatewayKey = `-- name: InsertAIGatewayKey :one INSERT INTO ai_gateway_keys (id, name, secret_prefix, hashed_secret, created_at) VALUES ($1, $4, $2, $3, NOW()) @@ -217,6 +237,22 @@ func (q *sqlQuerier) ListAIGatewayKeys(ctx context.Context) ([]ListAIGatewayKeys return items, nil } +const updateAIGatewayKeyLastUsedAt = `-- name: UpdateAIGatewayKeyLastUsedAt :exec +UPDATE ai_gateway_keys +SET last_used_at = $1 +WHERE id = $2 +` + +type UpdateAIGatewayKeyLastUsedAtParams struct { + LastUsedAt sql.NullTime `db:"last_used_at" json:"last_used_at"` + ID uuid.UUID `db:"id" json:"id"` +} + +func (q *sqlQuerier) UpdateAIGatewayKeyLastUsedAt(ctx context.Context, arg UpdateAIGatewayKeyLastUsedAtParams) error { + _, err := q.db.ExecContext(ctx, updateAIGatewayKeyLastUsedAt, arg.LastUsedAt, arg.ID) + return err +} + const deleteAIProviderKey = `-- name: DeleteAIProviderKey :exec DELETE FROM ai_provider_keys diff --git a/coderd/database/queries/ai_gateway_keys.sql b/coderd/database/queries/ai_gateway_keys.sql index 308d0cb89d1aa..227fee13240aa 100644 --- a/coderd/database/queries/ai_gateway_keys.sql +++ b/coderd/database/queries/ai_gateway_keys.sql @@ -11,3 +11,13 @@ ORDER BY created_at ASC; -- name: DeleteAIGatewayKey :one DELETE FROM ai_gateway_keys WHERE id = $1 RETURNING id, name, secret_prefix, created_at, last_used_at; + +-- name: GetAIGatewayKeyByHashedSecret :one +SELECT * +FROM ai_gateway_keys +WHERE hashed_secret = $1; + +-- name: UpdateAIGatewayKeyLastUsedAt :exec +UPDATE ai_gateway_keys +SET last_used_at = @last_used_at +WHERE id = @id;