feat: track client sessions (#22470)

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
+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
}