From adad5bdd493d2f1354898b934feec4ce8a0133c1 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Mon, 22 Jun 2026 15:17:37 +0200 Subject: [PATCH] feat: surface agent firewall correlation in AI Bridge sessions API (#26416) Add `agent_firewall_session_id` and `agent_firewall_sequence_number` fields to `AIBridgeThread` in the `GET /api/v2/aibridge/sessions/{session_id}` response. These fields link each thread to its agent firewall confinement session so the frontend can discover the boundary session and compute sequence ranges for interleaving firewall events within the thread timeline. The database columns already exist on `aibridge_interceptions` (migration 000520) and are already selected by `ListAIBridgeSessionThreads`. This PR surfaces them through the SDK type and the `db2sdk` conversion. Depends on #24814 **Naming note:** The RFC uses `boundary_session_id` / `boundary_sequence_number`, but the codebase standardized on `agent_firewall_*` naming in the DB migration. The API fields follow the existing convention. > [!NOTE] > This PR was authored by Coder Agents. --- coderd/apidoc/docs.go | 9 ++++++ coderd/apidoc/swagger.json | 9 ++++++ coderd/database/db2sdk/db2sdk.go | 8 +++++ codersdk/aibridge.go | 9 ++++++ docs/reference/api/aibridge.md | 2 ++ docs/reference/api/schemas.md | 30 +++++++++++-------- enterprise/coderd/aibridge_test.go | 47 ++++++++++++++++++++++++++++++ site/src/api/typesGenerated.ts | 13 +++++++++ 8 files changed, 115 insertions(+), 12 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index a56092c1b4..c4ca12fa33 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -15141,6 +15141,15 @@ const docTemplate = `{ "codersdk.AIBridgeThread": { "type": "object", "properties": { + "agent_firewall_sequence_number": { + "description": "AgentFirewallSequenceNumber is the firewall sequence number from\nthe root interception. Used to determine the position of this\nLLM request in the firewall event stream. Nil when the request\ndid not pass through the agent firewall.", + "type": "integer" + }, + "agent_firewall_session_id": { + "description": "AgentFirewallSessionID links this thread to an agent firewall\nconfinement session. Nil when the request did not pass through\nthe agent firewall.", + "type": "string", + "format": "uuid" + }, "agentic_actions": { "type": "array", "items": { diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 4bde3bc0c4..04112e4cfa 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -13489,6 +13489,15 @@ "codersdk.AIBridgeThread": { "type": "object", "properties": { + "agent_firewall_sequence_number": { + "description": "AgentFirewallSequenceNumber is the firewall sequence number from\nthe root interception. Used to determine the position of this\nLLM request in the firewall event stream. Nil when the request\ndid not pass through the agent firewall.", + "type": "integer" + }, + "agent_firewall_session_id": { + "description": "AgentFirewallSessionID links this thread to an agent firewall\nconfinement session. Nil when the request did not pass through\nthe agent firewall.", + "type": "string", + "format": "uuid" + }, "agentic_actions": { "type": "array", "items": { diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index b2ccca09d6..24b4d60404 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1274,6 +1274,14 @@ func buildAIBridgeThread( if prompts := promptsByInterception[rootIntc.ID]; len(prompts) > 0 { thread.Prompt = &prompts[0].Prompt } + if rootIntc.AgentFirewallSessionID.Valid { + id := rootIntc.AgentFirewallSessionID.UUID + thread.AgentFirewallSessionID = &id + } + if rootIntc.AgentFirewallSequenceNumber.Valid { + n := rootIntc.AgentFirewallSequenceNumber.Int32 + thread.AgentFirewallSequenceNumber = &n + } } // Compute thread time bounds from interceptions. diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index d8356a559f..6c32af0ecd 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -80,6 +80,15 @@ type AIBridgeThread struct { EndedAt *time.Time `json:"ended_at,omitempty" format:"date-time"` TokenUsage AIBridgeSessionThreadsTokenUsage `json:"token_usage"` AgenticActions []AIBridgeAgenticAction `json:"agentic_actions"` + // AgentFirewallSessionID links this thread to an agent firewall + // confinement session. Nil when the request did not pass through + // the agent firewall. + AgentFirewallSessionID *uuid.UUID `json:"agent_firewall_session_id,omitempty" format:"uuid"` + // AgentFirewallSequenceNumber is the firewall sequence number from + // the root interception. Used to determine the position of this + // LLM request in the firewall event stream. Nil when the request + // did not pass through the agent firewall. + AgentFirewallSequenceNumber *int32 `json:"agent_firewall_sequence_number,omitempty"` } // AIBridgeAgenticAction represents a tool call with associated diff --git a/docs/reference/api/aibridge.md b/docs/reference/api/aibridge.md index 4a5757fe01..ca86148e1a 100644 --- a/docs/reference/api/aibridge.md +++ b/docs/reference/api/aibridge.md @@ -191,6 +191,8 @@ curl -X GET http://coder-server:8080/api/v2/aibridge/sessions/{session_id} \ "started_at": "2019-08-24T14:15:22Z", "threads": [ { + "agent_firewall_sequence_number": 0, + "agent_firewall_session_id": "3735294f-18b1-4e7a-a269-99c30f0b30e7", "agentic_actions": [ { "model": "string", diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index 044fdf99a7..acecb1f60c 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -654,6 +654,8 @@ "started_at": "2019-08-24T14:15:22Z", "threads": [ { + "agent_firewall_sequence_number": 0, + "agent_firewall_session_id": "3735294f-18b1-4e7a-a269-99c30f0b30e7", "agentic_actions": [ { "model": "string", @@ -791,6 +793,8 @@ ```json { + "agent_firewall_sequence_number": 0, + "agent_firewall_session_id": "3735294f-18b1-4e7a-a269-99c30f0b30e7", "agentic_actions": [ { "model": "string", @@ -850,18 +854,20 @@ ### Properties -| Name | Type | Required | Restrictions | Description | -|-------------------|----------------------------------------------------------------------------------------|----------|--------------|-------------| -| `agentic_actions` | array of [codersdk.AIBridgeAgenticAction](#codersdkaibridgeagenticaction) | false | | | -| `credential_hint` | string | false | | | -| `credential_kind` | string | false | | | -| `ended_at` | string | false | | | -| `id` | string | false | | | -| `model` | string | false | | | -| `prompt` | string | false | | | -| `provider` | string | false | | | -| `started_at` | string | false | | | -| `token_usage` | [codersdk.AIBridgeSessionThreadsTokenUsage](#codersdkaibridgesessionthreadstokenusage) | false | | | +| Name | Type | Required | Restrictions | Description | +|----------------------------------|----------------------------------------------------------------------------------------|----------|--------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `agent_firewall_sequence_number` | integer | false | | Agent firewall sequence number is the firewall sequence number from the root interception. Used to determine the position of this LLM request in the firewall event stream. Nil when the request did not pass through the agent firewall. | +| `agent_firewall_session_id` | string | false | | Agent firewall session ID links this thread to an agent firewall confinement session. Nil when the request did not pass through the agent firewall. | +| `agentic_actions` | array of [codersdk.AIBridgeAgenticAction](#codersdkaibridgeagenticaction) | false | | | +| `credential_hint` | string | false | | | +| `credential_kind` | string | false | | | +| `ended_at` | string | false | | | +| `id` | string | false | | | +| `model` | string | false | | | +| `prompt` | string | false | | | +| `provider` | string | false | | | +| `started_at` | string | false | | | +| `token_usage` | [codersdk.AIBridgeSessionThreadsTokenUsage](#codersdkaibridgesessionthreadstokenusage) | false | | | ## codersdk.AIBridgeToolCall diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 86cb319047..fcd7a27611 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -1457,6 +1457,53 @@ func TestAIBridgeGetSessionThreads(t *testing.T) { require.Equal(t, "sk-a...efgh", res.Threads[0].CredentialHint) }) + t.Run("ThreadsWithAgentFirewallCorrelation", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + fwSessionID := uuid.New() + + // Thread with firewall correlation on the root interception. + rootEndedAt := now.Add(time.Minute) + root := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + StartedAt: now, + ClientSessionID: sql.NullString{String: "fw-session", Valid: true}, + AgentFirewallSessionID: uuid.NullUUID{UUID: fwSessionID, Valid: true}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: 5, Valid: true}, + }, &rootEndedAt) + + // Thread without firewall correlation in the same session. + noFWEndedAt := now.Add(2 * time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "openai", + Model: "gpt-4", + StartedAt: now.Add(time.Minute), + ClientSessionID: sql.NullString{String: "fw-session", Valid: true}, + }, &noFWEndedAt) + + res, err := client.AIBridgeGetSessionThreads(ctx, "fw-session", uuid.Nil, uuid.Nil, 0) + require.NoError(t, err) + require.Equal(t, "fw-session", res.ID) + require.Len(t, res.Threads, 2) + + // First thread has firewall correlation. + require.Equal(t, root.ID, res.Threads[0].ID) + require.NotNil(t, res.Threads[0].AgentFirewallSessionID) + require.Equal(t, fwSessionID, *res.Threads[0].AgentFirewallSessionID) + require.NotNil(t, res.Threads[0].AgentFirewallSequenceNumber) + require.Equal(t, int32(5), *res.Threads[0].AgentFirewallSequenceNumber) + + // Second thread has no firewall correlation. + require.Nil(t, res.Threads[1].AgentFirewallSessionID) + require.Nil(t, res.Threads[1].AgentFirewallSequenceNumber) + }) + t.Run("ThreadsWithAgenticActions", func(t *testing.T) { t.Parallel() client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) diff --git a/site/src/api/typesGenerated.ts b/site/src/api/typesGenerated.ts index 78769ed7d2..9ed9d80e91 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -203,6 +203,19 @@ export interface AIBridgeThread { readonly ended_at?: string; readonly token_usage: AIBridgeSessionThreadsTokenUsage; readonly agentic_actions: readonly AIBridgeAgenticAction[]; + /** + * AgentFirewallSessionID links this thread to an agent firewall + * confinement session. Nil when the request did not pass through + * the agent firewall. + */ + readonly agent_firewall_session_id?: string; + /** + * AgentFirewallSequenceNumber is the firewall sequence number from + * the root interception. Used to determine the position of this + * LLM request in the firewall event stream. Nil when the request + * did not pass through the agent firewall. + */ + readonly agent_firewall_sequence_number?: number; } // From codersdk/aibridge.go