From e1c7e61eb9c588aa4fd755e3c96b13985feefa0e Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Mon, 15 Jun 2026 12:34:16 +0200 Subject: [PATCH] feat(coderd): add Agent Firewall correlation columns to aibridge_interceptions (#24817) Add `agent_firewall_session_id` (UUID NULL) and `agent_firewall_sequence_number` (INT NULL) to `aibridge_interceptions` with a partial index on `agent_firewall_session_id`. No FK to `boundary_sessions` (soft reference, resolved at query time). `RecordInterception` reads the new fields from the proto request (merged in #25884) via `parseOptionalUUID` / `parseOptionalInt32` helpers. > This PR was authored by Coder Agents. --- coderd/aibridgedserver/aibridgedserver.go | 59 +++++++-- .../aibridgedserver/aibridgedserver_test.go | 122 ++++++++++++++++++ coderd/database/dbgen/dbgen.go | 30 +++-- coderd/database/dump.sql | 10 +- ...0_aibridge_agent_firewall_session.down.sql | 5 + ...520_aibridge_agent_firewall_session.up.sql | 15 +++ coderd/database/modelqueries.go | 2 + coderd/database/models.go | 4 + coderd/database/querier_test.go | 81 ++++++++++++ coderd/database/queries.sql.go | 56 +++++--- coderd/database/queries/aibridge.sql | 4 +- 11 files changed, 336 insertions(+), 52 deletions(-) create mode 100644 coderd/database/migrations/000520_aibridge_agent_firewall_session.down.sql create mode 100644 coderd/database/migrations/000520_aibridge_agent_firewall_session.up.sql diff --git a/coderd/aibridgedserver/aibridgedserver.go b/coderd/aibridgedserver/aibridgedserver.go index 94411de16d..c9fb665f71 100644 --- a/coderd/aibridgedserver/aibridgedserver.go +++ b/coderd/aibridgedserver/aibridgedserver.go @@ -180,21 +180,29 @@ func (s *Server) RecordInterception(ctx context.Context, in *proto.RecordInterce providerName = in.Provider } + agentFirewallSessionID, err := parseOptionalUUID(in.AgentFirewallSessionId) + if err != nil { + s.logger.Warn(ctx, "invalid agent firewall session ID in interception request", + slog.F("agent_firewall_session_id", in.GetAgentFirewallSessionId()), slog.Error(err)) + } + _, err = s.store.InsertAIBridgeInterception(ctx, database.InsertAIBridgeInterceptionParams{ - ID: intcID, - APIKeyID: sql.NullString{String: in.ApiKeyId, Valid: true}, - Client: sql.NullString{String: in.Client, Valid: in.Client != ""}, - ClientSessionID: sql.NullString{String: in.GetClientSessionId(), Valid: in.GetClientSessionId() != ""}, - InitiatorID: initID, - Provider: in.Provider, - ProviderName: providerName, - Model: in.Model, - Metadata: out, - StartedAt: in.StartedAt.AsTime(), - ThreadParentInterceptionID: uuid.NullUUID{UUID: parentID, Valid: parentID != uuid.Nil}, - ThreadRootInterceptionID: uuid.NullUUID{UUID: rootID, Valid: rootID != uuid.Nil}, - CredentialKind: credentialKindOrDefault(in.CredentialKind), - CredentialHint: in.CredentialHint, + ID: intcID, + APIKeyID: sql.NullString{String: in.ApiKeyId, Valid: true}, + Client: sql.NullString{String: in.Client, Valid: in.Client != ""}, + ClientSessionID: sql.NullString{String: in.GetClientSessionId(), Valid: in.GetClientSessionId() != ""}, + InitiatorID: initID, + Provider: in.Provider, + ProviderName: providerName, + Model: in.Model, + Metadata: out, + StartedAt: in.StartedAt.AsTime(), + ThreadParentInterceptionID: uuid.NullUUID{UUID: parentID, Valid: parentID != uuid.Nil}, + ThreadRootInterceptionID: uuid.NullUUID{UUID: rootID, Valid: rootID != uuid.Nil}, + CredentialKind: credentialKindOrDefault(in.CredentialKind), + CredentialHint: in.CredentialHint, + AgentFirewallSessionID: agentFirewallSessionID, + AgentFirewallSequenceNumber: parseOptionalInt32(in.AgentFirewallSequenceNumber), }) if err != nil { return nil, xerrors.Errorf("start interception: %w", err) @@ -688,3 +696,26 @@ func metadataToMap(in map[string]*anypb.Any) map[string]any { } return meta } + +// parseOptionalUUID converts an optional proto string to uuid.NullUUID. +// Returns a zero NullUUID if s is nil. If s is non-nil but not a valid UUID, it +// returns a zero NullUUID along with the parse error so the caller can decide +// how to surface it. +func parseOptionalUUID(s *string) (uuid.NullUUID, error) { + if s == nil { + return uuid.NullUUID{}, nil + } + id, err := uuid.Parse(*s) + if err != nil { + return uuid.NullUUID{}, err + } + return uuid.NullUUID{UUID: id, Valid: true}, nil +} + +// parseOptionalInt32 converts an optional proto int32 to sql.NullInt32. +func parseOptionalInt32(n *int32) sql.NullInt32 { + if n == nil { + return sql.NullInt32{} + } + return sql.NullInt32{Int32: *n, Valid: true} +} diff --git a/coderd/aibridgedserver/aibridgedserver_test.go b/coderd/aibridgedserver/aibridgedserver_test.go index 5930260e05..2d5ce57e97 100644 --- a/coderd/aibridgedserver/aibridgedserver_test.go +++ b/coderd/aibridgedserver/aibridgedserver_test.go @@ -688,6 +688,128 @@ func TestRecordInterception(t *testing.T) { }, nil) }, }, + { + name: "valid interception with agent firewall correlation", + request: &proto.RecordInterceptionRequest{ + Id: uuid.NewString(), + ApiKeyId: uuid.NewString(), + InitiatorId: uuid.NewString(), + Provider: "anthropic", + Model: "claude-4-opus", + Metadata: metadataProto, + StartedAt: timestamppb.Now(), + AgentFirewallSessionId: ptr.Ref(uuid.NewString()), + AgentFirewallSequenceNumber: ptr.Ref(int32(42)), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) { + interceptionID, err := uuid.Parse(req.GetId()) + assert.NoError(t, err, "parse interception UUID") + initiatorID, err := uuid.Parse(req.GetInitiatorId()) + assert.NoError(t, err, "parse interception initiator UUID") + agentFirewallSessionID, err := uuid.Parse(req.GetAgentFirewallSessionId()) + assert.NoError(t, err, "parse agent firewall session UUID") + + db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{ + ID: interceptionID, + APIKeyID: sql.NullString{String: req.ApiKeyId, Valid: true}, + InitiatorID: initiatorID, + Provider: req.GetProvider(), + ProviderName: req.GetProvider(), + Model: req.GetModel(), + Metadata: json.RawMessage(metadataJSON), + StartedAt: req.StartedAt.AsTime().UTC(), + CredentialKind: database.CredentialKindCentralized, + AgentFirewallSessionID: uuid.NullUUID{UUID: agentFirewallSessionID, Valid: true}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: 42, Valid: true}, + }).Return(database.AIBridgeInterception{ + ID: interceptionID, + InitiatorID: initiatorID, + Provider: req.GetProvider(), + Model: req.GetModel(), + StartedAt: req.StartedAt.AsTime().UTC(), + }, nil) + }, + }, + { + name: "absent agent firewall fields treated as null", + request: &proto.RecordInterceptionRequest{ + Id: uuid.NewString(), + ApiKeyId: uuid.NewString(), + InitiatorId: uuid.NewString(), + Provider: "anthropic", + Model: "claude-4-opus", + Metadata: metadataProto, + StartedAt: timestamppb.Now(), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) { + interceptionID, err := uuid.Parse(req.GetId()) + assert.NoError(t, err, "parse interception UUID") + initiatorID, err := uuid.Parse(req.GetInitiatorId()) + assert.NoError(t, err, "parse interception initiator UUID") + + db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{ + ID: interceptionID, + APIKeyID: sql.NullString{String: req.ApiKeyId, Valid: true}, + InitiatorID: initiatorID, + Provider: req.GetProvider(), + ProviderName: req.GetProvider(), + Model: req.GetModel(), + Metadata: json.RawMessage(metadataJSON), + StartedAt: req.StartedAt.AsTime().UTC(), + CredentialKind: database.CredentialKindCentralized, + AgentFirewallSessionID: uuid.NullUUID{}, + AgentFirewallSequenceNumber: sql.NullInt32{}, + }).Return(database.AIBridgeInterception{ + ID: interceptionID, + InitiatorID: initiatorID, + Provider: req.GetProvider(), + Model: req.GetModel(), + StartedAt: req.StartedAt.AsTime().UTC(), + }, nil) + }, + }, + { + name: "invalid agent firewall session ID treated as null", + request: &proto.RecordInterceptionRequest{ + Id: uuid.NewString(), + ApiKeyId: uuid.NewString(), + InitiatorId: uuid.NewString(), + Provider: "anthropic", + Model: "claude-4-opus", + Metadata: metadataProto, + StartedAt: timestamppb.Now(), + AgentFirewallSessionId: ptr.Ref("not-a-uuid"), + AgentFirewallSequenceNumber: ptr.Ref(int32(7)), + }, + setupMocks: func(t *testing.T, db *dbmock.MockStore, req *proto.RecordInterceptionRequest) { + interceptionID, err := uuid.Parse(req.GetId()) + assert.NoError(t, err, "parse interception UUID") + initiatorID, err := uuid.Parse(req.GetInitiatorId()) + assert.NoError(t, err, "parse interception initiator UUID") + + // Malformed agent firewall session ID is stored as null + // (and logged) rather than failing the interception. + db.EXPECT().InsertAIBridgeInterception(gomock.Any(), database.InsertAIBridgeInterceptionParams{ + ID: interceptionID, + APIKeyID: sql.NullString{String: req.ApiKeyId, Valid: true}, + InitiatorID: initiatorID, + Provider: req.GetProvider(), + ProviderName: req.GetProvider(), + Model: req.GetModel(), + Metadata: json.RawMessage(metadataJSON), + StartedAt: req.StartedAt.AsTime().UTC(), + CredentialKind: database.CredentialKindCentralized, + AgentFirewallSessionID: uuid.NullUUID{}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: 7, Valid: true}, + }).Return(database.AIBridgeInterception{ + ID: interceptionID, + InitiatorID: initiatorID, + Provider: req.GetProvider(), + Model: req.GetModel(), + StartedAt: req.StartedAt.AsTime().UTC(), + }, nil) + }, + }, { name: "invalid interception ID", request: &proto.RecordInterceptionRequest{ diff --git a/coderd/database/dbgen/dbgen.go b/coderd/database/dbgen/dbgen.go index 3955220efe..249038ffc0 100644 --- a/coderd/database/dbgen/dbgen.go +++ b/coderd/database/dbgen/dbgen.go @@ -1997,20 +1997,22 @@ 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"), - ProviderName: takeFirst(seed.ProviderName, "provider-name"), - Model: takeFirst(seed.Model, "model"), - Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")), - StartedAt: takeFirst(seed.StartedAt, dbtime.Now()), - Client: seed.Client, - ThreadParentInterceptionID: seed.ThreadParentInterceptionID, - ThreadRootInterceptionID: seed.ThreadRootInterceptionID, - ClientSessionID: seed.ClientSessionID, - CredentialKind: takeFirst(seed.CredentialKind, database.CredentialKindCentralized), - CredentialHint: takeFirst(seed.CredentialHint, ""), + ID: takeFirst(seed.ID, uuid.New()), + APIKeyID: seed.APIKeyID, + InitiatorID: takeFirst(seed.InitiatorID, uuid.New()), + Provider: takeFirst(seed.Provider, "provider"), + ProviderName: takeFirst(seed.ProviderName, "provider-name"), + Model: takeFirst(seed.Model, "model"), + Metadata: takeFirstSlice(seed.Metadata, json.RawMessage("{}")), + StartedAt: takeFirst(seed.StartedAt, dbtime.Now()), + Client: seed.Client, + ThreadParentInterceptionID: seed.ThreadParentInterceptionID, + ThreadRootInterceptionID: seed.ThreadRootInterceptionID, + ClientSessionID: seed.ClientSessionID, + CredentialKind: takeFirst(seed.CredentialKind, database.CredentialKindCentralized), + CredentialHint: takeFirst(seed.CredentialHint, ""), + AgentFirewallSessionID: seed.AgentFirewallSessionID, + AgentFirewallSequenceNumber: seed.AgentFirewallSequenceNumber, }) if endedAt != nil { interception, err = db.UpdateAIBridgeInterceptionEnded(genCtx, database.UpdateAIBridgeInterceptionEndedParams{ diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 6f4d588d67..05d293bf4b 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -1512,7 +1512,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, + agent_firewall_session_id uuid, + agent_firewall_sequence_number integer ); COMMENT ON TABLE aibridge_interceptions IS 'Audit log of requests intercepted by AI Bridge'; @@ -1533,6 +1535,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.agent_firewall_session_id IS 'The Agent Firewall session ID, linking this Bridge interception to an Agent Firewall confinement session.'; + +COMMENT ON COLUMN aibridge_interceptions.agent_firewall_sequence_number IS 'The Agent Firewall sequence number from the request header. Used to determine exact ordering of network requests relative to Agent Firewall audit events. NULL when the request did not pass through Agent Firewall.'; + CREATE TABLE aibridge_model_thoughts ( interception_id uuid NOT NULL, content text NOT NULL, @@ -4364,6 +4370,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_agent_firewall_session_id ON aibridge_interceptions USING btree (agent_firewall_session_id) WHERE (agent_firewall_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/000520_aibridge_agent_firewall_session.down.sql b/coderd/database/migrations/000520_aibridge_agent_firewall_session.down.sql new file mode 100644 index 0000000000..8d9476e1a1 --- /dev/null +++ b/coderd/database/migrations/000520_aibridge_agent_firewall_session.down.sql @@ -0,0 +1,5 @@ +DROP INDEX IF EXISTS idx_aibridge_interceptions_agent_firewall_session_id; + +ALTER TABLE aibridge_interceptions + DROP COLUMN IF EXISTS agent_firewall_sequence_number, + DROP COLUMN IF EXISTS agent_firewall_session_id; diff --git a/coderd/database/migrations/000520_aibridge_agent_firewall_session.up.sql b/coderd/database/migrations/000520_aibridge_agent_firewall_session.up.sql new file mode 100644 index 0000000000..3594b7a055 --- /dev/null +++ b/coderd/database/migrations/000520_aibridge_agent_firewall_session.up.sql @@ -0,0 +1,15 @@ +-- No FK to agent firewall sessions: Bridge interceptions may be recorded +-- before the session row exists, since Agent Firewall log delivery is async. +-- agent_firewall_session_id is a soft reference resolved at query time. +ALTER TABLE aibridge_interceptions + ADD COLUMN agent_firewall_session_id UUID NULL, + ADD COLUMN agent_firewall_sequence_number INT NULL; + +COMMENT ON COLUMN aibridge_interceptions.agent_firewall_session_id IS + 'The Agent Firewall session ID, linking this Bridge interception to an Agent Firewall confinement session.'; +COMMENT ON COLUMN aibridge_interceptions.agent_firewall_sequence_number IS + 'The Agent Firewall sequence number from the request header. Used to determine exact ordering of network requests relative to Agent Firewall audit events. NULL when the request did not pass through Agent Firewall.'; + +CREATE INDEX idx_aibridge_interceptions_agent_firewall_session_id + ON aibridge_interceptions (agent_firewall_session_id) + WHERE agent_firewall_session_id IS NOT NULL; diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index d89156d654..dd25d88739 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -1142,6 +1142,8 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeSessionThreads(ctx context.Context, a &i.AIBridgeInterception.ProviderName, &i.AIBridgeInterception.CredentialKind, &i.AIBridgeInterception.CredentialHint, + &i.AIBridgeInterception.AgentFirewallSessionID, + &i.AIBridgeInterception.AgentFirewallSequenceNumber, ); err != nil { return nil, err } diff --git a/coderd/database/models.go b/coderd/database/models.go index e52d5f6c4a..b18645c065 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4402,6 +4402,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 Agent Firewall session ID, linking this Bridge interception to an Agent Firewall confinement session. + AgentFirewallSessionID uuid.NullUUID `db:"agent_firewall_session_id" json:"agent_firewall_session_id"` + // The Agent Firewall sequence number from the request header. Used to determine exact ordering of network requests relative to Agent Firewall audit events. NULL when the request did not pass through Agent Firewall. + AgentFirewallSequenceNumber sql.NullInt32 `db:"agent_firewall_sequence_number" json:"agent_firewall_sequence_number"` } // Audit log of model thinking in intercepted requests in AI Bridge diff --git a/coderd/database/querier_test.go b/coderd/database/querier_test.go index 23f170067d..1e097c2c3c 100644 --- a/coderd/database/querier_test.go +++ b/coderd/database/querier_test.go @@ -10297,6 +10297,87 @@ func TestUpdateAIBridgeInterceptionEnded(t *testing.T) { }) } +func TestAIBridgeInterceptionAgentFirewallColumns(t *testing.T) { + t.Parallel() + db, _ := dbtestutil.NewDB(t) + + afwSessionID := uuid.New() + + t.Run("InsertAndReadWithFirewallFieldsSet", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + user := dbgen.User(t, db, database.User{}) + + inserted, err := db.InsertAIBridgeInterception(ctx, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + InitiatorID: user.ID, + Metadata: json.RawMessage("{}"), + CredentialKind: database.CredentialKindCentralized, + AgentFirewallSessionID: uuid.NullUUID{UUID: afwSessionID, Valid: true}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: 5, Valid: true}, + }) + require.NoError(t, err) + require.Equal(t, uuid.NullUUID{UUID: afwSessionID, Valid: true}, inserted.AgentFirewallSessionID) + require.Equal(t, sql.NullInt32{Int32: 5, Valid: true}, inserted.AgentFirewallSequenceNumber) + + got, err := db.GetAIBridgeInterceptionByID(ctx, inserted.ID) + require.NoError(t, err) + require.Equal(t, uuid.NullUUID{UUID: afwSessionID, Valid: true}, got.AgentFirewallSessionID) + require.Equal(t, sql.NullInt32{Int32: 5, Valid: true}, got.AgentFirewallSequenceNumber) + }) + + t.Run("InsertAndReadWithFirewallFieldsNull", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + user := dbgen.User(t, db, database.User{}) + + inserted, err := db.InsertAIBridgeInterception(ctx, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + InitiatorID: user.ID, + Metadata: json.RawMessage("{}"), + CredentialKind: database.CredentialKindCentralized, + // AgentFirewallSessionID and AgentFirewallSequenceNumber omitted (zero → NULL). + }) + require.NoError(t, err) + require.False(t, inserted.AgentFirewallSessionID.Valid) + require.False(t, inserted.AgentFirewallSequenceNumber.Valid) + + got, err := db.GetAIBridgeInterceptionByID(ctx, inserted.ID) + require.NoError(t, err) + require.False(t, got.AgentFirewallSessionID.Valid) + require.False(t, got.AgentFirewallSequenceNumber.Valid) + }) + + t.Run("UpdatePreservesFields", func(t *testing.T) { + t.Parallel() + ctx := testutil.Context(t, testutil.WaitLong) + + user := dbgen.User(t, db, database.User{}) + + inserted, err := db.InsertAIBridgeInterception(ctx, database.InsertAIBridgeInterceptionParams{ + ID: uuid.New(), + InitiatorID: user.ID, + Metadata: json.RawMessage("{}"), + CredentialKind: database.CredentialKindCentralized, + AgentFirewallSessionID: uuid.NullUUID{UUID: afwSessionID, Valid: true}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: 5, Valid: true}, + }) + require.NoError(t, err) + + updated, err := db.UpdateAIBridgeInterceptionEnded(ctx, database.UpdateAIBridgeInterceptionEndedParams{ + ID: inserted.ID, + EndedAt: time.Now(), + }) + require.NoError(t, err) + require.True(t, updated.EndedAt.Valid) + // UpdateAIBridgeInterceptionEnded must not clobber the agent firewall fields. + require.Equal(t, uuid.NullUUID{UUID: afwSessionID, Valid: true}, updated.AgentFirewallSessionID) + require.Equal(t, sql.NullInt32{Int32: 5, Valid: true}, updated.AgentFirewallSequenceNumber) + }) +} + func TestDeleteExpiredAPIKeys(t *testing.T) { t.Parallel() db, _ := dbtestutil.NewDB(t) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index cb2dfd3ce6..48088d112b 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1104,7 +1104,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, agent_firewall_session_id, agent_firewall_sequence_number FROM aibridge_interceptions WHERE @@ -1131,6 +1131,8 @@ func (q *sqlQuerier) GetAIBridgeInterceptionByID(ctx context.Context, id uuid.UU &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.AgentFirewallSessionID, + &i.AgentFirewallSequenceNumber, ) return i, err } @@ -1165,7 +1167,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, agent_firewall_session_id, agent_firewall_sequence_number FROM aibridge_interceptions ` @@ -1196,6 +1198,8 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.AgentFirewallSessionID, + &i.AgentFirewallSequenceNumber, ); err != nil { return nil, err } @@ -1344,28 +1348,30 @@ 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, agent_firewall_session_id, agent_firewall_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, agent_firewall_session_id, agent_firewall_sequence_number ` 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"` - ProviderName string `db:"provider_name" json:"provider_name"` - Model string `db:"model" json:"model"` - 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"` - CredentialKind CredentialKind `db:"credential_kind" json:"credential_kind"` - CredentialHint string `db:"credential_hint" json:"credential_hint"` + 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"` + ProviderName string `db:"provider_name" json:"provider_name"` + Model string `db:"model" json:"model"` + 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"` + CredentialKind CredentialKind `db:"credential_kind" json:"credential_kind"` + CredentialHint string `db:"credential_hint" json:"credential_hint"` + AgentFirewallSessionID uuid.NullUUID `db:"agent_firewall_session_id" json:"agent_firewall_session_id"` + AgentFirewallSequenceNumber sql.NullInt32 `db:"agent_firewall_sequence_number" json:"agent_firewall_sequence_number"` } func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertAIBridgeInterceptionParams) (AIBridgeInterception, error) { @@ -1384,6 +1390,8 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA arg.ThreadRootInterceptionID, arg.CredentialKind, arg.CredentialHint, + arg.AgentFirewallSessionID, + arg.AgentFirewallSequenceNumber, ) var i AIBridgeInterception err := row.Scan( @@ -1403,6 +1411,8 @@ func (q *sqlQuerier) InsertAIBridgeInterception(ctx context.Context, arg InsertA &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.AgentFirewallSessionID, + &i.AgentFirewallSequenceNumber, ) return i, err } @@ -1810,7 +1820,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.agent_firewall_session_id, aibridge_interceptions.agent_firewall_sequence_number FROM aibridge_interceptions JOIN @@ -1874,6 +1884,8 @@ func (q *sqlQuerier) ListAIBridgeSessionThreads(ctx context.Context, arg ListAIB &i.AIBridgeInterception.ProviderName, &i.AIBridgeInterception.CredentialKind, &i.AIBridgeInterception.CredentialHint, + &i.AIBridgeInterception.AgentFirewallSessionID, + &i.AIBridgeInterception.AgentFirewallSequenceNumber, ); err != nil { return nil, err } @@ -2295,7 +2307,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, agent_firewall_session_id, agent_firewall_sequence_number ` type UpdateAIBridgeInterceptionEndedParams struct { @@ -2324,6 +2336,8 @@ func (q *sqlQuerier) UpdateAIBridgeInterceptionEnded(ctx context.Context, arg Up &i.ProviderName, &i.CredentialKind, &i.CredentialHint, + &i.AgentFirewallSessionID, + &i.AgentFirewallSequenceNumber, ) return i, err } diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 6fbc83b2f5..73996dfa21 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, agent_firewall_session_id, agent_firewall_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('agent_firewall_session_id')::uuid, sqlc.narg('agent_firewall_sequence_number') ) RETURNING *;