opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore: add API key ID to interceptions (#20513)

Browse Source 
Adds APIKeyID to interceptions.
Needed for tracking API key usage with bridge.
fixes https://github.com/coder/coder/issues/20001
This commit is contained in:
Paweł Banaszewski
2025-11-10 13:46:41 +01:00
committed by GitHub
parent 81c3375670
commit 991831b1dd
26 changed files with 349 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/db2sdk/db2sdk.go
+3
View File
@@ -974,6 +974,9 @@ func AIBridgeInterception(interception database.AIBridgeInterception, initiator
UserPrompts: sdkUserPrompts,
ToolUsages: sdkToolUsages,
}
if interception.APIKeyID.Valid {
intc.APIKeyID = &interception.APIKeyID.String
}
if interception.EndedAt.Valid {
intc.EndedAt = &interception.EndedAt.Time
}
coderd/database/dbgen/dbgen.go
+1
View File
@@ -1496,6 +1496,7 @@ func ClaimPrebuild(
func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertAIBridgeInterceptionParams, endedAt *time.Time) database.AIBridgeInterception {
interception, err := db.InsertAIBridgeInterception(genCtx, database.InsertAIBridgeInterceptionParams{
ID: takeFirst(seed.ID, uuid.New()),
APIKeyID: seed.APIKeyID,
InitiatorID: takeFirst(seed.InitiatorID, uuid.New()),
Provider: takeFirst(seed.Provider, "provider"),
Model: takeFirst(seed.Model, "model"),
coderd/database/dump.sql
Generated
+2 -1
View File
@@ -1056,7 +1056,8 @@ CREATE TABLE aibridge_interceptions (
model text NOT NULL,
started_at timestamp with time zone NOT NULL,
metadata jsonb,
ended_at timestamp with time zone
ended_at timestamp with time zone,
api_key_id text
);
COMMENT ON TABLE aibridge_interceptions IS 'Audit log of requests intercepted by AI Bridge';
coderd/database/migrations/000396_add_aibridge_interceptions_api_key_id.down.sql
+1
View File
@@ -0,0 +1 @@
ALTER TABLE aibridge_interceptions DROP COLUMN api_key_id;
coderd/database/migrations/000396_add_aibridge_interceptions_api_key_id.up.sql
+2
View File
@@ -0,0 +1,2 @@
-- column is nullable to not break interceptions recorded before this column was added
ALTER TABLE aibridge_interceptions ADD COLUMN api_key_id text;
coderd/database/modelqueries.go
+1
View File
@@ -805,6 +805,7 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, ar
&i.AIBridgeInterception.StartedAt,
&i.AIBridgeInterception.Metadata,
&i.AIBridgeInterception.EndedAt,
&i.AIBridgeInterception.APIKeyID,
&i.VisibleUser.ID,
&i.VisibleUser.Username,
&i.VisibleUser.Name,
coderd/database/models.go
+1
View File
@@ -3614,6 +3614,7 @@ type AIBridgeInterception struct {
StartedAt time.Time `db:"started_at" json:"started_at"`
Metadata pqtype.NullRawMessage `db:"metadata" json:"metadata"`
EndedAt sql.NullTime `db:"ended_at" json:"ended_at"`
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
}
// Audit log of tokens used by intercepted requests in AI Bridge
coderd/database/queries.sql.go
+14 -7
View File
@@ -326,7 +326,7 @@ func (q *sqlQuerier) CountAIBridgeInterceptions(ctx context.Context, arg CountAI
const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
SELECT
id, initiator_id, provider, model, started_at, metadata, ended_at
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
FROM
aibridge_interceptions
WHERE
@@ -344,13 +344,14 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU
&i.StartedAt,
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
)
return i, err
}
const getAIBridgeInterceptions = `-- name: GetAIBridgeInterceptions :many
SELECT
id, initiator_id, provider, model, started_at, metadata, ended_at
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
FROM
aibridge_interceptions
`
@@ -372,6 +373,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn
&i.StartedAt,
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
); err != nil {
return nil, err
}
@@ -517,15 +519,16 @@ func (q *sqlQuerier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context,
const insertAIBridgeInterception = `-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (
id, initiator_id, provider, model, metadata, started_at
id, api_key_id, initiator_id, provider, model, metadata, started_at
) VALUES (
$1, $2, $3, $4, COALESCE($5::jsonb, '{}'::jsonb), $6
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7
)
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
`
type InsertAIBridgeInterceptionParams struct {
ID uuid.UUID `db:"id" json:"id"`
APIKeyID sql.NullString `db:"api_key_id" json:"api_key_id"`
InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"`
Provider string `db:"provider" json:"provider"`
Model string `db:"model" json:"model"`
@@ -536,6 +539,7 @@ type InsertAIBridgeInterceptionParams struct {
func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) {
row := q.db.QueryRowContext(ctx, insertAIBridgeInterception,
arg.ID,
arg.APIKeyID,
arg.InitiatorID,
arg.Provider,
arg.Model,
@@ -551,6 +555,7 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
&i.StartedAt,
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
)
return i, err
}
@@ -689,7 +694,7 @@ func (q *sqlQuerier) InsertAIBridgeUserPrompt(ctx context.Context, arg InsertAIB
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.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,
visible_users.id, visible_users.username, visible_users.name, visible_users.avatar_url
FROM
aibridge_interceptions
@@ -787,6 +792,7 @@ func (q *sqlQuerier) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBr
&i.AIBridgeInterception.StartedAt,
&i.AIBridgeInterception.Metadata,
&i.AIBridgeInterception.EndedAt,
&i.AIBridgeInterception.APIKeyID,
&i.VisibleUser.ID,
&i.VisibleUser.Username,
&i.VisibleUser.Name,
@@ -993,7 +999,7 @@ UPDATE aibridge_interceptions
WHERE
id = $2::uuid
AND ended_at IS NULL
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id
`
type UpdateAIBridgeInterceptionEndedParams struct {
@@ -1012,6 +1018,7 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up
&i.StartedAt,
&i.Metadata,
&i.EndedAt,
&i.APIKeyID,
)
return i, err
}
coderd/database/queries/aibridge.sql
+2 -2
View File
@@ -1,8 +1,8 @@
-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (
id, initiator_id, provider, model, metadata, started_at
id, api_key_id, initiator_id, provider, model, metadata, started_at
) VALUES (
@id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at
@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at
)
RETURNING *;