diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 8dbaa10bfa..1618d9d3bf 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -194,6 +194,8 @@ func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterce ThreadRootInterceptionID: uuid.NullUUID{UUID: rootID, Valid: rootID != uuid.Nil}, CredentialKind: credentialKindOrDefault(in.CredentialKind), CredentialHint: in.CredentialHint, + BoundarySessionID: uuid.NullUUID{}, + BoundarySequenceNumber: sql.NullInt64{}, }) if err != nil { return nil, xerrors.Errorf("start interception: %w", err) diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 416a2b7257..e007357d76 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -1999,6 +1999,8 @@ func AIBridgeInterception(t testing.TB, db database.Store, seed database.InsertA ClientSessionID: seed.ClientSessionID, CredentialKind: takeFirst(seed.CredentialKind, database.CredentialKindCentralized), CredentialHint: takeFirst(seed.CredentialHint, ""), + BoundarySessionID: seed.BoundarySessionID, + BoundarySequenceNumber: seed.BoundarySequenceNumber, }) if endedAt != nil { interception, err = db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{ diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 82aa376d34..2e3ec9982d 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1369,7 +1369,9 @@ CREATE TABLE aibridge_interceptions ( session_id text GENERATED ALWAYS AS (COALESCE(client_session_id, ((thread_root_id)::text)::character varying, ((id)::text)::character varying)) STORED NOT NULL, provider_name text DEFAULT ''::text NOT NULL, credential_kind credential_kind DEFAULT 'centralized'::credential_kind NOT NULL, - credential_hint character varying(15) DEFAULT ''::character varying NOT NULL + credential_hint character varying(15) DEFAULT ''::character varying NOT NULL, + boundary_session_id uuid, + boundary_sequence_number bigint ); COMMENT ON TABLE aibridge_interceptions IS 'Audit log of requests intercepted by AI Bridge'; @@ -1390,6 +1392,10 @@ COMMENT ON COLUMN aibridge_interceptions.credential_kind IS 'How the request was COMMENT ON COLUMN aibridge_interceptions.credential_hint IS 'Masked credential identifier for audit (e.g. sk-a***efgh).'; +COMMENT ON COLUMN aibridge_interceptions.boundary_session_id IS 'The Boundary session ID, linking this Bridge interception to a Boundary confinement session.'; + +COMMENT ON COLUMN aibridge_interceptions.boundary_sequence_number IS 'The Boundary sequence number from the request header. Used to determine exact ordering of network requests relative to Boundary audit events. NULL when the request did not pass through Boundary.'; + CREATE TABLE aibridge_model_thoughts ( interception_id uuid NOT NULL, content text NOT NULL, @@ -4161,6 +4167,8 @@ CREATE INDEX idx_ai_provider_keys_provider_id ON ai_provider_keys USING btree (p CREATE INDEX idx_ai_providers_enabled ON ai_providers USING btree (enabled) WHERE (deleted = false); +CREATE INDEX idx_aibridge_interceptions_boundary_session_id ON aibridge_interceptions USING btree (boundary_session_id) WHERE (boundary_session_id IS NOT NULL); + 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); diff --git a/coderd/database/migrations/000483_aibridge_boundary_session.down.sql b/coderd/database/migrations/000483_aibridge_boundary_session.down.sql new file mode 100644 index 0000000000..93eb42d67e --- /dev/null +++ b/coderd/database/migrations/000483_aibridge_boundary_session.down.sql @@ -0,0 +1,5 @@ +DROP INDEX IF EXISTS idx_aibridge_interceptions_boundary_session_id; + +ALTER TABLE aibridge_interceptions + DROP COLUMN IF EXISTS boundary_sequence_number, + DROP COLUMN IF EXISTS boundary_session_id; diff --git a/coderd/database/migrations/000483_aibridge_boundary_session.up.sql b/coderd/database/migrations/000483_aibridge_boundary_session.up.sql new file mode 100644 index 0000000000..51de547e48 --- /dev/null +++ b/coderd/database/migrations/000483_aibridge_boundary_session.up.sql @@ -0,0 +1,15 @@ +-- No FK to boundary_sessions: Bridge interceptions may be recorded before +-- the boundary_sessions row exists, since boundary log delivery is async. +-- boundary_session_id is a soft reference resolved at query time. +ALTER TABLE aibridge_interceptions + ADD COLUMN boundary_session_id UUID NULL, + ADD COLUMN boundary_sequence_number BIGINT NULL; + +COMMENT ON COLUMN aibridge_interceptions.boundary_session_id IS + 'The Boundary session ID, linking this Bridge interception to a Boundary confinement session.'; +COMMENT ON COLUMN aibridge_interceptions.boundary_sequence_number IS + 'The Boundary sequence number from the request header. Used to determine exact ordering of network requests relative to Boundary audit events. NULL when the request did not pass through Boundary.'; + +CREATE INDEX idx_aibridge_interceptions_boundary_session_id + ON aibridge_interceptions (boundary_session_id) + WHERE boundary_session_id IS NOT NULL; diff --git a/coderd/database/models.go b/coderd/database/models.go index ebfaa7a051..490ceac248 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4381,6 +4381,10 @@ type AIBridgeInterception struct { CredentialKind CredentialKind `db:"credential_kind" json:"credential_kind"` // Masked credential identifier for audit (e.g. sk-a***efgh). CredentialHint string `db:"credential_hint" json:"credential_hint"` + // The Boundary session ID, linking this Bridge interception to a Boundary confinement session. + BoundarySessionID uuid.NullUUID `db:"boundary_session_id" json:"boundary_session_id"` + // The Boundary sequence number from the request header. Used to determine exact ordering of network requests relative to Boundary audit events. NULL when the request did not pass through Boundary. + BoundarySequenceNumber sql.NullInt64 `db:"boundary_sequence_number" json:"boundary_sequence_number"` } // Audit log of model thinking in intercepted requests in AI Bridge diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index dc646121dc..351ee49f8d 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1066,7 +1066,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, client_session_id, session_id, provider_name, credential_kind, credential_hint + id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id, session_id, provider_name, credential_kind, credential_hint, boundary_session_id, boundary_sequence_number FROM aibridge_interceptions WHERE @@ -1093,6 +1093,8 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.BoundarySessionID, + &i.BoundarySequenceNumber, ) return i, err } @@ -1127,7 +1129,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, client_session_id, session_id, provider_name, credential_kind, credential_hint + id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id, session_id, provider_name, credential_kind, credential_hint, boundary_session_id, boundary_sequence_number FROM aibridge_interceptions ` @@ -1158,6 +1160,8 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.BoundarySessionID, + &i.BoundarySequenceNumber, ); err != nil { return nil, err } @@ -1306,11 +1310,11 @@ func (q *sqlQuerier) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, const insertAIBridgeInterception = `-- name: InsertAIBridgeInterception :one INSERT INTO aibridge_interceptions ( - id, api_key_id, initiator_id, provider, provider_name, model, metadata, started_at, client, client_session_id, thread_parent_id, thread_root_id, credential_kind, credential_hint + id, api_key_id, initiator_id, provider, provider_name, model, metadata, started_at, client, client_session_id, thread_parent_id, thread_root_id, credential_kind, credential_hint, boundary_session_id, boundary_sequence_number ) VALUES ( - $1, $2, $3, $4, $5, $6, COALESCE($7::jsonb, '{}'::jsonb), $8, $9, $10, $11::uuid, $12::uuid, $13, $14 + $1, $2, $3, $4, $5, $6, COALESCE($7::jsonb, '{}'::jsonb), $8, $9, $10, $11::uuid, $12::uuid, $13, $14, $15::uuid, $16 ) -RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id, session_id, provider_name, credential_kind, credential_hint +RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id, session_id, provider_name, credential_kind, credential_hint, boundary_session_id, boundary_sequence_number ` type InsertAIBridgeInterceptionParams struct { @@ -1328,6 +1332,8 @@ type InsertAIBridgeInterceptionParams struct { ThreadRootInterceptionID uuid.NullUUID `db:"thread_root_interception_id" json:"thread_root_interception_id"` CredentialKind CredentialKind `db:"credential_kind" json:"credential_kind"` CredentialHint string `db:"credential_hint" json:"credential_hint"` + BoundarySessionID uuid.NullUUID `db:"boundary_session_id" json:"boundary_session_id"` + BoundarySequenceNumber sql.NullInt64 `db:"boundary_sequence_number" json:"boundary_sequence_number"` } func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) { @@ -1346,6 +1352,8 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA arg.ThreadRootInterceptionID, arg.CredentialKind, arg.CredentialHint, + arg.BoundarySessionID, + arg.BoundarySequenceNumber, ) var i AIBridgeInterception err := row.Scan( @@ -1365,6 +1373,8 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.BoundarySessionID, + &i.BoundarySequenceNumber, ) return i, err } @@ -1597,7 +1607,7 @@ func (q *sqlQuerier) ListAIBridgeClients(ctx context.Context, arg ListAIBridgeCl 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.client_session_id, aibridge_interceptions.session_id, aibridge_interceptions.provider_name, aibridge_interceptions.credential_kind, aibridge_interceptions.credential_hint, + 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, aibridge_interceptions.session_id, aibridge_interceptions.provider_name, aibridge_interceptions.credential_kind, aibridge_interceptions.credential_hint, aibridge_interceptions.boundary_session_id, aibridge_interceptions.boundary_sequence_number, visible_users.id, visible_users.username, visible_users.name, visible_users.avatar_url FROM aibridge_interceptions @@ -1720,6 +1730,8 @@ func (q *sqlQuerier) ListAIBridgeInterceptions(ctx context.Context, arg ListAIBr &i.AIBridgeInterception.ProviderName, &i.AIBridgeInterception.CredentialKind, &i.AIBridgeInterception.CredentialHint, + &i.AIBridgeInterception.BoundarySessionID, + &i.AIBridgeInterception.BoundarySequenceNumber, &i.VisibleUser.ID, &i.VisibleUser.Username, &i.VisibleUser.Name, @@ -1915,7 +1927,7 @@ WITH paginated_threads AS ( ) SELECT COALESCE(aibridge_interceptions.thread_root_id, aibridge_interceptions.id) AS thread_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, aibridge_interceptions.session_id, aibridge_interceptions.provider_name, aibridge_interceptions.credential_kind, aibridge_interceptions.credential_hint + 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, aibridge_interceptions.session_id, aibridge_interceptions.provider_name, aibridge_interceptions.credential_kind, aibridge_interceptions.credential_hint, aibridge_interceptions.boundary_session_id, aibridge_interceptions.boundary_sequence_number FROM aibridge_interceptions JOIN @@ -1979,6 +1991,8 @@ func (q *sqlQuerier) ListAIBridgeSessionThreads(ctx context.Context, arg ListAIB &i.AIBridgeInterception.ProviderName, &i.AIBridgeInterception.CredentialKind, &i.AIBridgeInterception.CredentialHint, + &i.AIBridgeInterception.BoundarySessionID, + &i.AIBridgeInterception.BoundarySequenceNumber, ); err != nil { return nil, err } @@ -2400,7 +2414,7 @@ UPDATE aibridge_interceptions WHERE id = $3::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, client_session_id, session_id, provider_name, credential_kind, credential_hint +RETURNING id, initiator_id, provider, model, started_at, metadata, ended_at, api_key_id, client, thread_parent_id, thread_root_id, client_session_id, session_id, provider_name, credential_kind, credential_hint, boundary_session_id, boundary_sequence_number ` type UpdateAIBridgeInterceptionEndedParams struct { @@ -2429,6 +2443,8 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.BoundarySessionID, + &i.BoundarySequenceNumber, ) return i, err } diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index a1b49d25cd..557ba2cb9f 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -1,8 +1,8 @@ -- name: InsertAIBridgeInterception :one INSERT INTO aibridge_interceptions ( - id, api_key_id, initiator_id, provider, provider_name, model, metadata, started_at, client, client_session_id, thread_parent_id, thread_root_id, credential_kind, credential_hint + id, api_key_id, initiator_id, provider, provider_name, model, metadata, started_at, client, client_session_id, thread_parent_id, thread_root_id, credential_kind, credential_hint, boundary_session_id, boundary_sequence_number ) VALUES ( - @id, @api_key_id, @initiator_id, @provider, @provider_name, @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, @credential_kind, @credential_hint + @id, @api_key_id, @initiator_id, @provider, @provider_name, @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, @credential_kind, @credential_hint, sqlc.narg('boundary_session_id')::uuid, sqlc.narg('boundary_sequence_number') ) RETURNING *;