feat(coderd/database): add AI Gateway key auth lookup and last-used queries (#26505)

Adds DB methods`GetAIGatewayKeyIDByHashedSecret` and `UpdateAIGatewayKeyLastUsedAt`.
`GetAIGatewayKeyIDByHashedSecret` - returns AI Gateway key ID by hashed secret value.
`UpdateAIGatewayKeyLastUsedAt` - updates last used timestamp for given AI Gateway key. 
Used by standalone AI Gateway for authentication and keeping track of currently used keys.
This commit is contained in:
Paweł Banaszewski
2026-06-26 18:16:01 +02:00
committed by GitHub
parent e71d4ca69b
commit 0f1e792f3f
24 changed files with 271 additions and 7 deletions
+21
View File
@@ -486,6 +486,7 @@ var (
rbac.ResourceOauth2AppSecret.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
rbac.ResourceChat.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
rbac.ResourceAIProvider.Type: {policy.ActionCreate, policy.ActionRead, policy.ActionUpdate, policy.ActionDelete},
rbac.ResourceAIGatewayKey.Type: {policy.ActionRead, policy.ActionUpdate},
}),
User: []rbac.Permission{},
ByOrgID: map[string]rbac.OrgPermissions{},
@@ -2753,6 +2754,16 @@ func (q *querier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, in
return q.db.GetAIBridgeUserPromptsByInterceptionID(ctx, interceptionID)
}
// Authenticates a standalone AI Gateway replica by its hashed key secret, returning the matched key.
func (q *querier) GetAIGatewayKeyByHashedSecret(ctx context.Context, hashedSecret []byte) (database.AIGatewayKey, error) {
// Standalone AI Gateway has no Coder identity, so this runs under the
// system actor reading the AI Gateway key it authenticates against.
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
@@ -7022,6 +7033,16 @@ func (q *querier) UpdateAIBridgeInterceptionEnded(ctx context.Context, params da
return q.db.UpdateAIBridgeInterceptionEnded(ctx, params)
}
// Records liveness for a key used in active DRPC session between coderd and standalone AI Gateway.
func (q *querier) UpdateAIGatewayKeyLastUsedAt(ctx context.Context, id uuid.UUID) (int64, error) {
// Standalone AI Gateway has no Coder identity, so this runs under the
// system actor recording connection liveness on the AI Gateway key.
if err := q.authorizeContext(ctx, policy.ActionUpdate, rbac.ResourceAIGatewayKey); err != nil {
return 0, err
}
return q.db.UpdateAIGatewayKeyLastUsedAt(ctx, id)
}
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
+11
View File
@@ -6968,6 +6968,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")
key := database.AIGatewayKey{ID: uuid.New(), HashedSecret: hashedSecret}
dbm.EXPECT().GetAIGatewayKeyByHashedSecret(gomock.Any(), hashedSecret).Return(key, nil).AnyTimes()
check.Args(hashedSecret).Asserts(rbac.ResourceAIGatewayKey, policy.ActionRead).Returns(key)
}))
s.Run("UpdateAIGatewayKeyLastUsedAt", s.Mocked(func(dbm *dbmock.MockStore, _ *gofakeit.Faker, check *expects) {
id := uuid.New()
dbm.EXPECT().UpdateAIGatewayKeyLastUsedAt(gomock.Any(), id).Return(int64(1), nil).AnyTimes()
check.Args(id).Asserts(rbac.ResourceAIGatewayKey, policy.ActionUpdate).Returns(int64(1))
}))
}
func (s *MethodTestSuite) TestTelemetry() {