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

feat: track client sessions (#22470)

Browse Source 
This change adds support for tracking client session IDs in AI Bridge interceptions to enable better session-based auditing.

Depends on https://github.com/coder/aibridge/pull/198  
Fixes https://github.com/coder/internal/issues/1337

The session ID field is optional and not universally supported by all clients.
This commit is contained in:
Danny Kopping
2026-03-06 14:43:53 +02:00
committed by GitHub
parent f9891416c0
commit 13e3df67d6
17 changed files with 373 additions and 235 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/dbgen/dbgen.go
+1
View File
@@ -1595,6 +1595,7 @@ func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertA
Client: seed.Client,
ThreadParentInterceptionID: seed.ThreadParentInterceptionID,
ThreadRootInterceptionID: seed.ThreadRootInterceptionID,
ClientSessionID: seed.ClientSessionID,
})
if endedAt != nil {
interception, err = db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{
coderd/database/dump.sql
Generated
+6 -1
View File
@@ -1046,7 +1046,8 @@ CREATE TABLE aibridge_interceptions (
api_key_id text,
client character varying(64) DEFAULT 'Unknown'::character varying,
thread_parent_id uuid,
thread_root_id uuid
thread_root_id uuid,
client_session_id character varying(256)
);
COMMENT ON TABLE aibridge_interceptions IS 'Audit log of requests intercepted by AI Bridge';
@@ -1057,6 +1058,8 @@ COMMENT ON COLUMN aibridge_interceptions.thread_parent_id IS 'The interception w
COMMENT ON COLUMN aibridge_interceptions.thread_root_id IS 'The root interception of the thread that this interception belongs to.';
COMMENT ON COLUMN aibridge_interceptions.client_session_id IS 'The session ID supplied by the client (optional and not universally supported).';
CREATE TABLE aibridge_token_usages (
id uuid NOT NULL,
interception_id uuid NOT NULL,
@@ -3450,6 +3453,8 @@ CREATE INDEX idx_agent_stats_user_id ON workspace_agent_stats USING btree (user_
CREATE INDEX idx_aibridge_interceptions_client ON aibridge_interceptions USING btree (client);
CREATE INDEX idx_aibridge_interceptions_client_session_id ON aibridge_interceptions USING btree (client_session_id) WHERE (client_session_id IS NOT NULL);
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);
coderd/database/migrations/000428_aibridge_sessions.down.sql
+4
View File
@@ -0,0 +1,4 @@
DROP INDEX IF EXISTS idx_aibridge_interceptions_client_session_id;
ALTER TABLE aibridge_interceptions
DROP COLUMN client_session_id;
coderd/database/migrations/000428_aibridge_sessions.up.sql
+7
View File
@@ -0,0 +1,7 @@
ALTER TABLE aibridge_interceptions
ADD COLUMN client_session_id VARCHAR(256) NULL;
COMMENT ON COLUMN aibridge_interceptions.client_session_id IS 'The session ID supplied by the client (optional and not universally supported).';
CREATE INDEX idx_aibridge_interceptions_client_session_id ON aibridge_interceptions (client_session_id)
WHERE client_session_id IS NOT NULL;
coderd/database/modelqueries.go
+1
View File
@@ -815,6 +815,7 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeInterceptions(ctx context.Context, ar
&i.AIBridgeInterception.Client,
&i.AIBridgeInterception.ThreadParentID,
&i.AIBridgeInterception.ThreadRootID,
&i.AIBridgeInterception.ClientSessionID,
&i.VisibleUser.ID,
&i.VisibleUser.Username,
&i.VisibleUser.Name,
coderd/database/models.go
+2
View File
@@ -3793,6 +3793,8 @@ type AIBridgeInterception struct {
ThreadParentID uuid.NullUUID `db:"thread_parent_id" json:"thread_parent_id"`
// The root interception of the thread that this interception belongs to.
ThreadRootID uuid.NullUUID `db:"thread_root_id" json:"thread_root_id"`
// The session ID supplied by the client (optional and not universally supported).
ClientSessionID sql.NullString `db:"client_session_id" json:"client_session_id"`
}
// Audit log of tokens used by intercepted requests in AI Bridge
coderd/database/queries.sql.go
+14 -7
View File
@@ -378,7 +378,7 @@ func (q *sqlQuerier) DeleteOldAIBridgeRecords(ctx context.Context, beforeTime ti
const getAIBridgeInterceptionByID = `-- name: GetAIBridgeInterceptionByID :one
SELECT
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id
FROM
aibridge_interceptions
WHERE
@@ -400,6 +400,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU
&i.Client,
&i.ThreadParentID,
&i.ThreadRootID,
&i.ClientSessionID,
)
return i, err
}
@@ -434,7 +435,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptionLineageByToolCallID(ctx context.Cont
const getAIBridgeInterceptions = `-- name: GetAIBridgeInterceptions :many
SELECT
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id
FROM
aibridge_interceptions
`
@@ -460,6 +461,7 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn
&i.Client,
&i.ThreadParentID,
&i.ThreadRootID,
&i.ClientSessionID,
); err != nil {
return nil, err
}
@@ -606,11 +608,11 @@ func (q *sqlQuerier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context,
const insertAIBridgeInterception = `-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (
id, api_key_id, initiator_id, provider, model, metadata, started_at, client, thread_parent_id, thread_root_id
id, api_key_id, initiator_id, provider, model, metadata, started_at, client, client_session_id, thread_parent_id, thread_root_id
) VALUES (
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7, $8, $9::uuid, $10::uuid
$1, $2, $3, $4, $5, COALESCE($6::jsonb, '{}'::jsonb), $7, $8, $9, $10::uuid, $11::uuid
)
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id
`
type InsertAIBridgeInterceptionParams struct {
@@ -622,6 +624,7 @@ type InsertAIBridgeInterceptionParams struct {
Metadata json.RawMessage `db:"metadata" json:"metadata"`
StartedAt time.Time `db:"started_at" json:"started_at"`
Client sql.NullString `db:"client" json:"client"`
ClientSessionID sql.NullString `db:"client_session_id" json:"client_session_id"`
ThreadParentInterceptionID uuid.NullUUID `db:"thread_parent_interception_id" json:"thread_parent_interception_id"`
ThreadRootInterceptionID uuid.NullUUID `db:"thread_root_interception_id" json:"thread_root_interception_id"`
}
@@ -636,6 +639,7 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
arg.Metadata,
arg.StartedAt,
arg.Client,
arg.ClientSessionID,
arg.ThreadParentInterceptionID,
arg.ThreadRootInterceptionID,
)
@@ -652,6 +656,7 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA
&i.Client,
&i.ThreadParentID,
&i.ThreadRootID,
&i.ClientSessionID,
)
return i, err
}
@@ -793,7 +798,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.api_key_id, aibridge_interceptions.client, aibridge_interceptions.thread_parent_id, aibridge_interceptions.thread_root_id,
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,
visible_users.id, visible_users.username, visible_users.name, visible_users.avatar_url
FROM
aibridge_interceptions
@@ -904,6 +909,7 @@ func (q *sqlQuerier) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBr
&i.AIBridgeInterception.Client,
&i.AIBridgeInterception.ThreadParentID,
&i.AIBridgeInterception.ThreadRootID,
&i.AIBridgeInterception.ClientSessionID,
&i.VisibleUser.ID,
&i.VisibleUser.Username,
&i.VisibleUser.Name,
@@ -1164,7 +1170,7 @@ UPDATE aibridge_interceptions
WHERE
id = $2::uuid
AND ended_at IS NULL
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id
RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id
`
type UpdateAIBridgeInterceptionEndedParams struct {
@@ -1187,6 +1193,7 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up
&i.Client,
&i.ThreadParentID,
&i.ThreadRootID,
&i.ClientSessionID,
)
return i, err
}
coderd/database/queries/aibridge.sql
+2 -2
View File
@@ -1,8 +1,8 @@
-- name: InsertAIBridgeInterception :one
INSERT INTO aibridge_interceptions (
id, api_key_id, initiator_id, provider, model, metadata, started_at, client, thread_parent_id, thread_root_id
id, api_key_id, initiator_id, provider, model, metadata, started_at, client, client_session_id, thread_parent_id, thread_root_id
) VALUES (
@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at, @client, sqlc.narg('thread_parent_interception_id')::uuid, sqlc.narg('thread_root_interception_id')::uuid
@id, @api_key_id, @initiator_id, @provider, @model, COALESCE(@metadata::jsonb, '{}'::jsonb), @started_at, @client, sqlc.narg('client_session_id'), sqlc.narg('thread_parent_interception_id')::uuid, sqlc.narg('thread_root_interception_id')::uuid
)
RETURNING *;