mirror of
https://github.com/coder/coder.git
synced 2026-06-02 20:48:20 +00:00
feat(coderd/database): add boundary correlation columns to aibridge_interceptions
Add boundary_session_id (UUID NULL) and boundary_sequence_number (BIGINT NULL) columns to aibridge_interceptions with a partial index on boundary_session_id. No FK to boundary_sessions (soft reference, resolved at query time). Update the InsertAIBridgeInterception query and dbgen to accept the new fields. Still no readers/writers using them. Also renumber boundary_log_scopes migration from 000481 to 000482 to resolve a duplicate migration number conflict with boundary_sessions_and_logs.
This commit is contained in:
Sas Swart
@@ -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)
|
||||
|
||||
@@ -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{
|
||||
|
||||
Generated
+9
-1
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
@@ -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;
|
||||
Generated
+4
@@ -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
|
||||
|
||||
Generated
+24
-8
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 *;
|
||||
|
||||
|
||||
Reference in New Issue
Block a user