From bf4f9ca01018ba08d06c34b89622c9a98f88bb40 Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Wed, 22 Jul 2026 14:02:31 +0000 Subject: [PATCH 1/5] feat: add network calls summary to AI session threads API Expose total and blocked network calls plus top destination domains on the AI session threads endpoint. Total and blocked reuse the existing Agent Firewall aggregation from the sessions list query; top domains are a new GetAIBridgeSessionTopDomains aggregation over boundary logs using the same interception-window correlation. Refs AIGOV-463 Co-Authored-By: Claude Opus 4.8 (1M context) --- coderd/apidoc/docs.go | 29 ++++++ coderd/apidoc/swagger.json | 29 ++++++ coderd/database/db2sdk/db2sdk.go | 19 ++++ coderd/database/dbauthz/dbauthz.go | 7 ++ coderd/database/dbauthz/dbauthz_test.go | 6 ++ coderd/database/dbmetrics/querymetrics.go | 8 ++ coderd/database/dbmock/dbmock.go | 15 ++++ coderd/database/querier.go | 13 +++ coderd/database/queries.sql.go | 89 ++++++++++++++++++ coderd/database/queries/aibridge.sql | 54 +++++++++++ codersdk/aibridge.go | 19 +++- docs/reference/api/aigateway.md | 11 +++ docs/reference/api/schemas.md | 60 +++++++++---- enterprise/coderd/aibridge.go | 13 ++- enterprise/coderd/aibridge_test.go | 104 ++++++++++++++++++++++ site/src/api/typesGenerated.ts | 24 +++++ 16 files changed, 483 insertions(+), 17 deletions(-) diff --git a/coderd/apidoc/docs.go b/coderd/apidoc/docs.go index 388c35f114445..473155543e14e 100644 --- a/coderd/apidoc/docs.go +++ b/coderd/apidoc/docs.go @@ -15597,6 +15597,17 @@ const docTemplate = `{ } } }, + "codersdk.AIBridgeSessionNetworkDomain": { + "type": "object", + "properties": { + "count": { + "type": "integer" + }, + "domain": { + "type": "string" + } + } + }, "codersdk.AIBridgeSessionThreadsResponse": { "type": "object", "properties": { @@ -15623,6 +15634,24 @@ const docTemplate = `{ "type": "string" } }, + "network_calls": { + "description": "NetworkCalls summarizes the Agent Firewall network calls made during the\nsession. A nil value means the session did not pass through Agent\nFirewall, so network call monitoring was not active, which the UI\nsurfaces as \"Disabled\".", + "allOf": [ + { + "$ref": "#/definitions/codersdk.AIBridgeSessionNetworkCallSummary" + } + ] + }, + "network_domain_count": { + "type": "integer" + }, + "network_top_domains": { + "description": "NetworkTopDomains lists the most contacted destination hosts, ordered by\ncall count descending. NetworkDomainCount is the total number of distinct\ndomains, used to render a \"+N more\" overflow beyond the listed domains.", + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeSessionNetworkDomain" + } + }, "page_ended_at": { "type": "string", "format": "date-time" diff --git a/coderd/apidoc/swagger.json b/coderd/apidoc/swagger.json index 4553c63fc92f4..0076f613e4183 100644 --- a/coderd/apidoc/swagger.json +++ b/coderd/apidoc/swagger.json @@ -13891,6 +13891,17 @@ } } }, + "codersdk.AIBridgeSessionNetworkDomain": { + "type": "object", + "properties": { + "count": { + "type": "integer" + }, + "domain": { + "type": "string" + } + } + }, "codersdk.AIBridgeSessionThreadsResponse": { "type": "object", "properties": { @@ -13917,6 +13928,24 @@ "type": "string" } }, + "network_calls": { + "description": "NetworkCalls summarizes the Agent Firewall network calls made during the\nsession. A nil value means the session did not pass through Agent\nFirewall, so network call monitoring was not active, which the UI\nsurfaces as \"Disabled\".", + "allOf": [ + { + "$ref": "#/definitions/codersdk.AIBridgeSessionNetworkCallSummary" + } + ] + }, + "network_domain_count": { + "type": "integer" + }, + "network_top_domains": { + "description": "NetworkTopDomains lists the most contacted destination hosts, ordered by\ncall count descending. NetworkDomainCount is the total number of distinct\ndomains, used to render a \"+N more\" overflow beyond the listed domains.", + "type": "array", + "items": { + "$ref": "#/definitions/codersdk.AIBridgeSessionNetworkDomain" + } + }, "page_ended_at": { "type": "string", "format": "date-time" diff --git a/coderd/database/db2sdk/db2sdk.go b/coderd/database/db2sdk/db2sdk.go index 234b56739ea84..75deb23053d31 100644 --- a/coderd/database/db2sdk/db2sdk.go +++ b/coderd/database/db2sdk/db2sdk.go @@ -1153,6 +1153,7 @@ func AIBridgeSessionThreads( toolUsages []database.AIBridgeToolUsage, userPrompts []database.AIBridgeUserPrompt, modelThoughts []database.AIBridgeModelThought, + topDomains []database.GetAIBridgeSessionTopDomainsRow, ) codersdk.AIBridgeSessionThreadsResponse { // Index subresources by interception ID. tokensByInterception := make(map[uuid.UUID][]database.AIBridgeTokenUsage, len(interceptions)) @@ -1243,6 +1244,24 @@ func AIBridgeSessionThreads( if !session.EndedAt.IsZero() { resp.EndedAt = &session.EndedAt } + // NetworkCalls is only meaningful when the session passed through Agent + // Firewall. When it did not, leave it nil so the UI renders "Disabled" + // rather than a misleading zero count. + if session.FirewallActive { + resp.NetworkCalls = &codersdk.AIBridgeSessionNetworkCallSummary{ + Total: session.NetworkCallsTotal, + Blocked: session.NetworkCallsBlocked, + } + } + for _, d := range topDomains { + resp.NetworkTopDomains = append(resp.NetworkTopDomains, codersdk.AIBridgeSessionNetworkDomain{ + Domain: d.Domain, + Count: d.Count, + }) + // TotalDomains is the same on every row (a window aggregate); take it + // from the last row processed. + resp.NetworkDomainCount = d.TotalDomains + } return resp } diff --git a/coderd/database/dbauthz/dbauthz.go b/coderd/database/dbauthz/dbauthz.go index 3b51ccceb65ec..30a6389a0b73e 100644 --- a/coderd/database/dbauthz/dbauthz.go +++ b/coderd/database/dbauthz/dbauthz.go @@ -2811,6 +2811,13 @@ func (q *querier) GetAIBridgeInterceptions(ctx context.Context) ([]database.AIBr return fetchWithPostFilter(q.auth, policy.ActionRead, fetch)(ctx, nil) } +func (q *querier) GetAIBridgeSessionTopDomains(ctx context.Context, arg database.GetAIBridgeSessionTopDomainsParams) ([]database.GetAIBridgeSessionTopDomainsRow, error) { + if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAibridgeInterception); err != nil { + return nil, err + } + return q.db.GetAIBridgeSessionTopDomains(ctx, arg) +} + func (q *querier) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]database.AIBridgeTokenUsage, error) { // All aibridge_token_usages records belong to the initiator of their associated interception. if err := q.authorizeAIBridgeInterceptionAction(ctx, policy.ActionRead, interceptionID); err != nil { diff --git a/coderd/database/dbauthz/dbauthz_test.go b/coderd/database/dbauthz/dbauthz_test.go index 40bf232df43a6..503c80465ae93 100644 --- a/coderd/database/dbauthz/dbauthz_test.go +++ b/coderd/database/dbauthz/dbauthz_test.go @@ -6962,6 +6962,12 @@ func (s *MethodTestSuite) TestAIBridge() { check.Args(params, emptyPreparedAuthorized{}).Asserts() })) + s.Run("GetAIBridgeSessionTopDomains", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { + params := database.GetAIBridgeSessionTopDomainsParams{SessionID: "sess", Limit: 5} + db.EXPECT().GetAIBridgeSessionTopDomains(gomock.Any(), params).Return([]database.GetAIBridgeSessionTopDomainsRow{}, nil).AnyTimes() + check.Args(params).Asserts(rbac.ResourceAibridgeInterception, policy.ActionRead).Returns([]database.GetAIBridgeSessionTopDomainsRow{}) + })) + s.Run("ListAIBridgeTokenUsagesByInterceptionIDs", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) { ids := []uuid.UUID{{1}} db.EXPECT().ListAIBridgeTokenUsagesByInterceptionIDs(gomock.Any(), ids).Return([]database.AIBridgeTokenUsage{}, nil).AnyTimes() diff --git a/coderd/database/dbmetrics/querymetrics.go b/coderd/database/dbmetrics/querymetrics.go index a47adc4aae79a..33d5364aa8e0a 100644 --- a/coderd/database/dbmetrics/querymetrics.go +++ b/coderd/database/dbmetrics/querymetrics.go @@ -1105,6 +1105,14 @@ func (m queryMetricsStore) GetAIBridgeInterceptions(ctx context.Context) ([]data return r0, r1 } +func (m queryMetricsStore) GetAIBridgeSessionTopDomains(ctx context.Context, arg database.GetAIBridgeSessionTopDomainsParams) ([]database.GetAIBridgeSessionTopDomainsRow, error) { + start := time.Now() + r0, r1 := m.s.GetAIBridgeSessionTopDomains(ctx, arg) + m.queryLatencies.WithLabelValues("GetAIBridgeSessionTopDomains").Observe(time.Since(start).Seconds()) + m.queryCounts.WithLabelValues(httpmw.ExtractHTTPRoute(ctx), httpmw.ExtractHTTPMethod(ctx), "GetAIBridgeSessionTopDomains").Inc() + return r0, r1 +} + func (m queryMetricsStore) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]database.AIBridgeTokenUsage, error) { start := time.Now() r0, r1 := m.s.GetAIBridgeTokenUsagesByInterceptionID(ctx, interceptionID) diff --git a/coderd/database/dbmock/dbmock.go b/coderd/database/dbmock/dbmock.go index c63d297527b45..c7b824b24d98f 100644 --- a/coderd/database/dbmock/dbmock.go +++ b/coderd/database/dbmock/dbmock.go @@ -1903,6 +1903,21 @@ func (mr *MockStoreMockRecorder) GetAIBridgeInterceptions(ctx any) *gomock.Call return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBridgeInterceptions", reflect.TypeOf((*MockStore)(nil).GetAIBridgeInterceptions), ctx) } +// GetAIBridgeSessionTopDomains mocks base method. +func (m *MockStore) GetAIBridgeSessionTopDomains(ctx context.Context, arg database.GetAIBridgeSessionTopDomainsParams) ([]database.GetAIBridgeSessionTopDomainsRow, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetAIBridgeSessionTopDomains", ctx, arg) + ret0, _ := ret[0].([]database.GetAIBridgeSessionTopDomainsRow) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// GetAIBridgeSessionTopDomains indicates an expected call of GetAIBridgeSessionTopDomains. +func (mr *MockStoreMockRecorder) GetAIBridgeSessionTopDomains(ctx, arg any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetAIBridgeSessionTopDomains", reflect.TypeOf((*MockStore)(nil).GetAIBridgeSessionTopDomains), ctx, arg) +} + // GetAIBridgeTokenUsagesByInterceptionID mocks base method. func (m *MockStore) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]database.AIBridgeTokenUsage, error) { m.ctrl.T.Helper() diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 7ed0f7bc02b4a..8f074c7edc84f 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -302,6 +302,19 @@ type sqlcQuerier interface { // the root), we return its own ID as the root. GetAIBridgeInterceptionLineageByToolCallID(ctx context.Context, toolCallID string) (GetAIBridgeInterceptionLineageByToolCallIDRow, error) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeInterception, error) + // Returns the most contacted destination hosts for an AI session, ordered by + // call count descending and limited to the top @limit_ rows. total_domains is + // the number of distinct domains across the whole session, used to render a + // "+N more" overflow beyond the returned rows. Only HTTP egress is considered; + // dns/git/fs boundary logs do not carry a domain in the same shape. + // + // Windowing mirrors the network_calls aggregation in ListAIBridgeSessions: + // each interception's boundary logs fall in the open interval (this seq, next + // interception's seq) within the same firewall session. The exclusive lower + // bound drops the interception's own LLM-provider call. next_seq considers all + // interceptions in the firewall session so windows never bleed across AI + // sessions that share one firewall session. + GetAIBridgeSessionTopDomains(ctx context.Context, arg GetAIBridgeSessionTopDomainsParams) ([]GetAIBridgeSessionTopDomainsRow, error) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeTokenUsage, error) GetAIBridgeToolUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeToolUsage, error) GetAIBridgeUserPromptsByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeUserPrompt, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index e9df5bb7977f0..cbaac37c5be64 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1272,6 +1272,95 @@ func (q *sqlQuerier) GetAIBridgeInterceptions(ctx context.Context) ([]AIBridgeIn return items, nil } +const getAIBridgeSessionTopDomains = `-- name: GetAIBridgeSessionTopDomains :many +WITH session_boundary_logs AS ( + SELECT bl.detail + FROM aibridge_interceptions afi + LEFT JOIN LATERAL ( + SELECT MIN(nxt.agent_firewall_sequence_number) AS next_seq + FROM aibridge_interceptions nxt + WHERE nxt.agent_firewall_session_id = afi.agent_firewall_session_id + AND nxt.agent_firewall_sequence_number > afi.agent_firewall_sequence_number + ) w ON true + JOIN boundary_logs bl + ON bl.session_id = afi.agent_firewall_session_id + AND bl.sequence_number > afi.agent_firewall_sequence_number + AND (w.next_seq IS NULL OR bl.sequence_number < w.next_seq) + WHERE afi.session_id = $2::text + AND afi.ended_at IS NOT NULL + AND afi.agent_firewall_session_id IS NOT NULL + AND afi.agent_firewall_sequence_number IS NOT NULL + AND bl.proto = 'http' +), +extracted AS ( + -- Strip an optional scheme, then keep the host up to the first port, path, + -- query, or fragment delimiter. + SELECT substring(detail from '^(?:[A-Za-z][A-Za-z0-9+.-]*://)?([^/:?#]+)') AS domain + FROM session_boundary_logs +), +domains AS ( + SELECT domain, COUNT(*)::bigint AS count + FROM extracted + WHERE domain IS NOT NULL AND domain != '' + GROUP BY domain +) +SELECT + -- COALESCE keeps sqlc from typing the grouped column as nullable; the + -- domains CTE already filters out NULL/empty hosts. + COALESCE(domain, '')::text AS domain, + count, + COUNT(*) OVER ()::bigint AS total_domains +FROM domains +ORDER BY count DESC, domain ASC +LIMIT COALESCE(NULLIF($1::integer, 0), 5) +` + +type GetAIBridgeSessionTopDomainsParams struct { + Limit int32 `db:"limit_" json:"limit_"` + SessionID string `db:"session_id" json:"session_id"` +} + +type GetAIBridgeSessionTopDomainsRow struct { + Domain string `db:"domain" json:"domain"` + Count int64 `db:"count" json:"count"` + TotalDomains int64 `db:"total_domains" json:"total_domains"` +} + +// Returns the most contacted destination hosts for an AI session, ordered by +// call count descending and limited to the top @limit_ rows. total_domains is +// the number of distinct domains across the whole session, used to render a +// "+N more" overflow beyond the returned rows. Only HTTP egress is considered; +// dns/git/fs boundary logs do not carry a domain in the same shape. +// +// Windowing mirrors the network_calls aggregation in ListAIBridgeSessions: +// each interception's boundary logs fall in the open interval (this seq, next +// interception's seq) within the same firewall session. The exclusive lower +// bound drops the interception's own LLM-provider call. next_seq considers all +// interceptions in the firewall session so windows never bleed across AI +// sessions that share one firewall session. +func (q *sqlQuerier) GetAIBridgeSessionTopDomains(ctx context.Context, arg GetAIBridgeSessionTopDomainsParams) ([]GetAIBridgeSessionTopDomainsRow, error) { + rows, err := q.db.QueryContext(ctx, getAIBridgeSessionTopDomains, arg.Limit, arg.SessionID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []GetAIBridgeSessionTopDomainsRow + for rows.Next() { + var i GetAIBridgeSessionTopDomainsRow + if err := rows.Scan(&i.Domain, &i.Count, &i.TotalDomains); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Close(); err != nil { + return nil, err + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + const getAIBridgeTokenUsagesByInterceptionID = `-- name: GetAIBridgeTokenUsagesByInterceptionID :many SELECT id, interception_id, provider_response_id, input_tokens, output_tokens, metadata, created_at, cache_read_input_tokens, cache_write_input_tokens, effective_group_id, input_price_micros, output_price_micros, cache_read_price_micros, cache_write_price_micros, cost_micros diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 63635b6ae22ef..79b17a6b787db 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -551,6 +551,60 @@ ORDER BY sp.session_id DESC ; +-- name: GetAIBridgeSessionTopDomains :many +-- Returns the most contacted destination hosts for an AI session, ordered by +-- call count descending and limited to the top @limit_ rows. total_domains is +-- the number of distinct domains across the whole session, used to render a +-- "+N more" overflow beyond the returned rows. Only HTTP egress is considered; +-- dns/git/fs boundary logs do not carry a domain in the same shape. +-- +-- Windowing mirrors the network_calls aggregation in ListAIBridgeSessions: +-- each interception's boundary logs fall in the open interval (this seq, next +-- interception's seq) within the same firewall session. The exclusive lower +-- bound drops the interception's own LLM-provider call. next_seq considers all +-- interceptions in the firewall session so windows never bleed across AI +-- sessions that share one firewall session. +WITH session_boundary_logs AS ( + SELECT bl.detail + FROM aibridge_interceptions afi + LEFT JOIN LATERAL ( + SELECT MIN(nxt.agent_firewall_sequence_number) AS next_seq + FROM aibridge_interceptions nxt + WHERE nxt.agent_firewall_session_id = afi.agent_firewall_session_id + AND nxt.agent_firewall_sequence_number > afi.agent_firewall_sequence_number + ) w ON true + JOIN boundary_logs bl + ON bl.session_id = afi.agent_firewall_session_id + AND bl.sequence_number > afi.agent_firewall_sequence_number + AND (w.next_seq IS NULL OR bl.sequence_number < w.next_seq) + WHERE afi.session_id = @session_id::text + AND afi.ended_at IS NOT NULL + AND afi.agent_firewall_session_id IS NOT NULL + AND afi.agent_firewall_sequence_number IS NOT NULL + AND bl.proto = 'http' +), +extracted AS ( + -- Strip an optional scheme, then keep the host up to the first port, path, + -- query, or fragment delimiter. + SELECT substring(detail from '^(?:[A-Za-z][A-Za-z0-9+.-]*://)?([^/:?#]+)') AS domain + FROM session_boundary_logs +), +domains AS ( + SELECT domain, COUNT(*)::bigint AS count + FROM extracted + WHERE domain IS NOT NULL AND domain != '' + GROUP BY domain +) +SELECT + -- COALESCE keeps sqlc from typing the grouped column as nullable; the + -- domains CTE already filters out NULL/empty hosts. + COALESCE(domain, '')::text AS domain, + count, + COUNT(*) OVER ()::bigint AS total_domains +FROM domains +ORDER BY count DESC, domain ASC +LIMIT COALESCE(NULLIF(@limit_::integer, 0), 5); + -- name: ListAIBridgeSessionThreads :many -- Returns all interceptions belonging to paginated threads within a session. -- Threads are paginated by (started_at, thread_id) cursor. diff --git a/codersdk/aibridge.go b/codersdk/aibridge.go index 66f754bae6118..31d6ae52cafd6 100644 --- a/codersdk/aibridge.go +++ b/codersdk/aibridge.go @@ -165,6 +165,13 @@ type AIBridgeSessionNetworkCallSummary struct { Blocked int64 `json:"blocked"` } +// AIBridgeSessionNetworkDomain is one destination host contacted during a +// session, with the number of network calls made to it. +type AIBridgeSessionNetworkDomain struct { + Domain string `json:"domain"` + Count int64 `json:"count"` +} + type AIBridgeListSessionsResponse struct { Count int64 `json:"count"` Sessions []AIBridgeSession `json:"sessions"` @@ -185,7 +192,17 @@ type AIBridgeSessionThreadsResponse struct { StartedAt time.Time `json:"started_at" format:"date-time"` EndedAt *time.Time `json:"ended_at,omitempty" format:"date-time"` TokenUsageSummary AIBridgeSessionThreadsTokenUsage `json:"token_usage_summary"` - Threads []AIBridgeThread `json:"threads"` + // NetworkCalls summarizes the Agent Firewall network calls made during the + // session. A nil value means the session did not pass through Agent + // Firewall, so network call monitoring was not active, which the UI + // surfaces as "Disabled". + NetworkCalls *AIBridgeSessionNetworkCallSummary `json:"network_calls,omitempty"` + // NetworkTopDomains lists the most contacted destination hosts, ordered by + // call count descending. NetworkDomainCount is the total number of distinct + // domains, used to render a "+N more" overflow beyond the listed domains. + NetworkTopDomains []AIBridgeSessionNetworkDomain `json:"network_top_domains,omitempty"` + NetworkDomainCount int64 `json:"network_domain_count,omitempty"` + Threads []AIBridgeThread `json:"threads"` } // AIBridgeSessionThreadsTokenUsage represents aggregated token usage diff --git a/docs/reference/api/aigateway.md b/docs/reference/api/aigateway.md index 3d8ed690729b8..935309c481289 100644 --- a/docs/reference/api/aigateway.md +++ b/docs/reference/api/aigateway.md @@ -195,6 +195,17 @@ Alias: also available at /api/v2/aibridge/sessions/{session_id} for backward com "models": [ "string" ], + "network_calls": { + "blocked": 0, + "total": 0 + }, + "network_domain_count": 0, + "network_top_domains": [ + { + "count": 0, + "domain": "string" + } + ], "page_ended_at": "2019-08-24T14:15:22Z", "page_started_at": "2019-08-24T14:15:22Z", "providers": [ diff --git a/docs/reference/api/schemas.md b/docs/reference/api/schemas.md index f8295c7d1ba4c..0f460573c3860 100644 --- a/docs/reference/api/schemas.md +++ b/docs/reference/api/schemas.md @@ -655,6 +655,22 @@ | `blocked` | integer | false | | | | `total` | integer | false | | | +## codersdk.AIBridgeSessionNetworkDomain + +```json +{ + "count": 0, + "domain": "string" +} +``` + +### Properties + +| Name | Type | Required | Restrictions | Description | +|----------|---------|----------|--------------|-------------| +| `count` | integer | false | | | +| `domain` | string | false | | | + ## codersdk.AIBridgeSessionThreadsResponse ```json @@ -675,6 +691,17 @@ "models": [ "string" ], + "network_calls": { + "blocked": 0, + "total": 0 + }, + "network_domain_count": 0, + "network_top_domains": [ + { + "count": 0, + "domain": "string" + } + ], "page_ended_at": "2019-08-24T14:15:22Z", "page_started_at": "2019-08-24T14:15:22Z", "providers": [ @@ -758,21 +785,24 @@ ### Properties -| Name | Type | Required | Restrictions | Description | -|-----------------------|----------------------------------------------------------------------------------------|----------|--------------|-------------| -| `client` | string | false | | | -| `ended_at` | string | false | | | -| `id` | string | false | | | -| `initiator` | [codersdk.MinimalUser](#codersdkminimaluser) | false | | | -| `metadata` | object | false | | | -| » `[any property]` | any | false | | | -| `models` | array of string | false | | | -| `page_ended_at` | string | false | | | -| `page_started_at` | string | false | | | -| `providers` | array of string | false | | | -| `started_at` | string | false | | | -| `threads` | array of [codersdk.AIBridgeThread](#codersdkaibridgethread) | false | | | -| `token_usage_summary` | [codersdk.AIBridgeSessionThreadsTokenUsage](#codersdkaibridgesessionthreadstokenusage) | false | | | +| Name | Type | Required | Restrictions | Description | +|------------------------|------------------------------------------------------------------------------------------|----------|--------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------| +| `client` | string | false | | | +| `ended_at` | string | false | | | +| `id` | string | false | | | +| `initiator` | [codersdk.MinimalUser](#codersdkminimaluser) | false | | | +| `metadata` | object | false | | | +| » `[any property]` | any | false | | | +| `models` | array of string | false | | | +| `network_calls` | [codersdk.AIBridgeSessionNetworkCallSummary](#codersdkaibridgesessionnetworkcallsummary) | false | | Network calls summarizes the Agent Firewall network calls made during the session. A nil value means the session did not pass through Agent Firewall, so network call monitoring was not active, which the UI surfaces as "Disabled". | +| `network_domain_count` | integer | false | | | +| `network_top_domains` | array of [codersdk.AIBridgeSessionNetworkDomain](#codersdkaibridgesessionnetworkdomain) | false | | Network top domains lists the most contacted destination hosts, ordered by call count descending. NetworkDomainCount is the total number of distinct domains, used to render a "+N more" overflow beyond the listed domains. | +| `page_ended_at` | string | false | | | +| `page_started_at` | string | false | | | +| `providers` | array of string | false | | | +| `started_at` | string | false | | | +| `threads` | array of [codersdk.AIBridgeThread](#codersdkaibridgethread) | false | | | +| `token_usage_summary` | [codersdk.AIBridgeSessionThreadsTokenUsage](#codersdkaibridgesessionthreadstokenusage) | false | | | ## codersdk.AIBridgeSessionThreadsTokenUsage diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index 884bd23afdbc8..ee2aa6bbdb84a 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -369,6 +369,7 @@ func (api *API) aiBridgeGetSessionThreads(rw http.ResponseWriter, r *http.Reques toolUsages []database.AIBridgeToolUsage userPrompts []database.AIBridgeUserPrompt modelThoughts []database.AIBridgeModelThought + topDomains []database.GetAIBridgeSessionTopDomainsRow ) err = api.Database.InTx(func(db database.Store) error { // Validate cursor IDs before querying threads. The SQL @@ -435,6 +436,16 @@ func (api *API) aiBridgeGetSessionThreads(rw http.ResponseWriter, r *http.Reques return xerrors.Errorf("list model thoughts: %w", err) } + // Aggregate the session's top network destinations. Scoped by + // session ID (not the page) so the summary reflects the whole session. + topDomains, err = db.GetAIBridgeSessionTopDomains(ctx, database.GetAIBridgeSessionTopDomainsParams{ + SessionID: sessionIDParam, + Limit: 5, + }) + if err != nil { + return xerrors.Errorf("get session top domains: %w", err) + } + return nil }, &database.TxOptions{ Isolation: sql.LevelRepeatableRead, @@ -456,7 +467,7 @@ func (api *API) aiBridgeGetSessionThreads(rw http.ResponseWriter, r *http.Reques return } - resp := db2sdk.AIBridgeSessionThreads(session, threadRows, tokenUsages, toolUsages, userPrompts, modelThoughts) + resp := db2sdk.AIBridgeSessionThreads(session, threadRows, tokenUsages, toolUsages, userPrompts, modelThoughts, topDomains) httpapi.Write(ctx, rw, http.StatusOK, resp) } diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index 00c7b4dde7a81..a065d7c425a0c 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -1672,6 +1672,110 @@ func TestAIBridgeGetSessionThreads(t *testing.T) { require.Nil(t, res.Threads[1].AgentFirewallSequenceNumber) }) + t.Run("NetworkSummary", func(t *testing.T) { + t.Parallel() + // Use the raw store so boundary logs can be seeded directly. No RBAC + // role grants boundary_log:create; they are written by the agent path. + db, ps := dbtestutil.NewDB(t) + opts := aibridgeOpts(t) + opts.Options.Database = db + opts.Options.Pubsub = ps + client, _, firstUser := coderdenttest.NewWithDatabase(t, opts) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + fw := uuid.New() + + // One interception marked at firewall seq 0, so its window is (0, +inf) + // and the LLM-provider call logged at seq 0 is excluded. + endedAt := now.Add(time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + StartedAt: now, + ClientSessionID: sql.NullString{String: "net-session", Valid: true}, + AgentFirewallSessionID: uuid.NullUUID{UUID: fw, Valid: true}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: 0, Valid: true}, + }, &endedAt) + + type logSeed struct { + seq int32 + proto string + detail string + allowed bool + } + seeds := []logSeed{ + {0, "http", "https://api.github.com/llm", true}, // LLM call, excluded + {1, "http", "https://api.github.com/repos/coder", true}, // github egress + {2, "http", "https://api.github.com/repos/other", true}, // github egress + {3, "http", "https://registry.npmjs.org/lodash", false}, // npm egress, blocked + {4, "http", "https://api.github.com/repos/more", true}, // github egress + {5, "dns", "example.com", true}, // non-http, ignored by top domains + } + logs := make([]database.BoundaryLog, 0, len(seeds)) + for _, s := range seeds { + // A non-empty matched_rule marks the call as allowed; an empty rule + // is stored as NULL, which counts as blocked. + rule := "" + if s.allowed { + rule = "allow " + s.detail + } + logs = append(logs, database.BoundaryLog{ + SessionID: fw, + OwnerID: uuid.NullUUID{UUID: firstUser.UserID, Valid: true}, + SequenceNumber: s.seq, + CapturedAt: now, + CreatedAt: now, + Proto: s.proto, + Method: "GET", + Detail: s.detail, + MatchedRule: sql.NullString{String: rule, Valid: rule != ""}, + }) + } + dbgen.BoundaryLogs(t, db, logs) + + res, err := client.AIBridgeGetSessionThreads(ctx, "net-session", uuid.Nil, uuid.Nil, 0) + require.NoError(t, err) + + // total counts seq 1-5 (LLM call at seq 0 excluded); one blocked. + require.NotNil(t, res.NetworkCalls) + require.EqualValues(t, 5, res.NetworkCalls.Total) + require.EqualValues(t, 1, res.NetworkCalls.Blocked) + + // Top domains covers HTTP egress only: github x3, npm x1. The dns log is + // excluded. Two distinct domains, ordered by count descending. + require.Len(t, res.NetworkTopDomains, 2) + require.Equal(t, "api.github.com", res.NetworkTopDomains[0].Domain) + require.EqualValues(t, 3, res.NetworkTopDomains[0].Count) + require.Equal(t, "registry.npmjs.org", res.NetworkTopDomains[1].Domain) + require.EqualValues(t, 1, res.NetworkTopDomains[1].Count) + require.EqualValues(t, 2, res.NetworkDomainCount) + }) + + t.Run("NetworkSummaryDisabled", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + endedAt := now.Add(time.Minute) + // No firewall correlation: network monitoring was not active. + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + StartedAt: now, + ClientSessionID: sql.NullString{String: "no-fw-session", Valid: true}, + }, &endedAt) + + res, err := client.AIBridgeGetSessionThreads(ctx, "no-fw-session", uuid.Nil, uuid.Nil, 0) + require.NoError(t, err) + require.Nil(t, res.NetworkCalls) + require.Empty(t, res.NetworkTopDomains) + require.EqualValues(t, 0, res.NetworkDomainCount) + }) + 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 5dadf27b2b235..93c4081b4cde3 100644 --- a/site/src/api/typesGenerated.ts +++ b/site/src/api/typesGenerated.ts @@ -162,6 +162,16 @@ export interface AIBridgeSessionNetworkCallSummary { readonly blocked: number; } +// From codersdk/aibridge.go +/** + * AIBridgeSessionNetworkDomain is one destination host contacted during a + * session, with the number of network calls made to it. + */ +export interface AIBridgeSessionNetworkDomain { + readonly domain: string; + readonly count: number; +} + // From codersdk/aibridge.go /** * AIBridgeSessionThreadsResponse is the response for GET @@ -181,6 +191,20 @@ export interface AIBridgeSessionThreadsResponse { readonly started_at: string; readonly ended_at?: string; readonly token_usage_summary: AIBridgeSessionThreadsTokenUsage; + /** + * NetworkCalls summarizes the Agent Firewall network calls made during the + * session. A nil value means the session did not pass through Agent + * Firewall, so network call monitoring was not active, which the UI + * surfaces as "Disabled". + */ + readonly network_calls?: AIBridgeSessionNetworkCallSummary; + /** + * NetworkTopDomains lists the most contacted destination hosts, ordered by + * call count descending. NetworkDomainCount is the total number of distinct + * domains, used to render a "+N more" overflow beyond the listed domains. + */ + readonly network_top_domains?: readonly AIBridgeSessionNetworkDomain[]; + readonly network_domain_count?: number; readonly threads: readonly AIBridgeThread[]; } From 0e5efd95f8cfefd5d3923cfdaed37303f96f8d8d Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Thu, 23 Jul 2026 09:53:53 +0000 Subject: [PATCH 2/5] test: cover network-call windowing invariants and reduce top-domains fetch Add multi-interception and shared-firewall-session tests asserting the network_calls summary total/blocked and top-domains counts partition correctly across consecutive windows and do not bleed across AI sessions that share one firewall session. Reduce the top-domains fetch to a single row, since the summary card renders only the most-contacted domain plus a "+N more" count derived from NetworkDomainCount (a window aggregate independent of the row cap). Document that the domain-extraction regex assumes scheme+host(+port) detail without userinfo or IPv6 literal hosts, and add a port-suffixed test row that pins host stripping. Refs AIGOV-463 Co-Authored-By: Claude Opus 4.8 (1M context) --- coderd/database/queries.sql.go | 6 +- coderd/database/queries/aibridge.sql | 6 +- enterprise/coderd/aibridge.go | 8 +- enterprise/coderd/aibridge_test.go | 193 ++++++++++++++++++++++----- 4 files changed, 172 insertions(+), 41 deletions(-) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index cbaac37c5be64..bdf7c2565eafb 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1294,7 +1294,11 @@ WITH session_boundary_logs AS ( ), extracted AS ( -- Strip an optional scheme, then keep the host up to the first port, path, - -- query, or fragment delimiter. + -- query, or fragment delimiter. This assumes HTTP egress detail is a plain + -- scheme+host(+port) URL: it does not handle userinfo (user@host, which + -- would be captured into the host) or IPv6 literal hosts ([::1], where the + -- leading '[' is captured and the ':' terminates early). Boundary HTTP logs + -- do not currently emit those forms; revisit this extraction if they do. SELECT substring(detail from '^(?:[A-Za-z][A-Za-z0-9+.-]*://)?([^/:?#]+)') AS domain FROM session_boundary_logs ), diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 79b17a6b787db..9521263c46084 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -585,7 +585,11 @@ WITH session_boundary_logs AS ( ), extracted AS ( -- Strip an optional scheme, then keep the host up to the first port, path, - -- query, or fragment delimiter. + -- query, or fragment delimiter. This assumes HTTP egress detail is a plain + -- scheme+host(+port) URL: it does not handle userinfo (user@host, which + -- would be captured into the host) or IPv6 literal hosts ([::1], where the + -- leading '[' is captured and the ':' terminates early). Boundary HTTP logs + -- do not currently emit those forms; revisit this extraction if they do. SELECT substring(detail from '^(?:[A-Za-z][A-Za-z0-9+.-]*://)?([^/:?#]+)') AS domain FROM session_boundary_logs ), diff --git a/enterprise/coderd/aibridge.go b/enterprise/coderd/aibridge.go index ee2aa6bbdb84a..a428c31b7a092 100644 --- a/enterprise/coderd/aibridge.go +++ b/enterprise/coderd/aibridge.go @@ -436,11 +436,13 @@ func (api *API) aiBridgeGetSessionThreads(rw http.ResponseWriter, r *http.Reques return xerrors.Errorf("list model thoughts: %w", err) } - // Aggregate the session's top network destinations. Scoped by - // session ID (not the page) so the summary reflects the whole session. + // Aggregate the session's top network destination. Scoped by session + // ID (not the page) so the summary reflects the whole session. The + // summary card renders only the single most-contacted domain plus a + // "+N more" count derived from NetworkDomainCount, so we fetch one row. topDomains, err = db.GetAIBridgeSessionTopDomains(ctx, database.GetAIBridgeSessionTopDomainsParams{ SessionID: sessionIDParam, - Limit: 5, + Limit: 1, }) if err != nil { return xerrors.Errorf("get session top domains: %w", err) diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index a065d7c425a0c..1541174af8d6c 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -1562,6 +1562,39 @@ func TestAIBridgeConcurrencyLimiting(t *testing.T) { } } +type boundaryLogSeed struct { + seq int32 + proto string + detail string + allowed bool +} + +// seedBoundaryLogs writes boundary logs for a firewall session via the raw +// store. A non-empty matched_rule marks a call allowed; a blocked call stores a +// NULL rule. No RBAC role grants boundary_log:create, so tests seed directly. +func seedBoundaryLogs(t *testing.T, db database.Store, fw, ownerID uuid.UUID, at time.Time, seeds []boundaryLogSeed) { + t.Helper() + logs := make([]database.BoundaryLog, 0, len(seeds)) + for _, s := range seeds { + rule := "" + if s.allowed { + rule = "allow " + s.detail + } + logs = append(logs, database.BoundaryLog{ + SessionID: fw, + OwnerID: uuid.NullUUID{UUID: ownerID, Valid: true}, + SequenceNumber: s.seq, + CapturedAt: at, + CreatedAt: at, + Proto: s.proto, + Method: "GET", + Detail: s.detail, + MatchedRule: sql.NullString{String: rule, Valid: rule != ""}, + }) + } + dbgen.BoundaryLogs(t, db, logs) +} + func TestAIBridgeGetSessionThreads(t *testing.T) { t.Parallel() @@ -1699,60 +1732,148 @@ func TestAIBridgeGetSessionThreads(t *testing.T) { AgentFirewallSequenceNumber: sql.NullInt32{Int32: 0, Valid: true}, }, &endedAt) - type logSeed struct { - seq int32 - proto string - detail string - allowed bool - } - seeds := []logSeed{ + seedBoundaryLogs(t, db, fw, firstUser.UserID, now, []boundaryLogSeed{ {0, "http", "https://api.github.com/llm", true}, // LLM call, excluded {1, "http", "https://api.github.com/repos/coder", true}, // github egress {2, "http", "https://api.github.com/repos/other", true}, // github egress {3, "http", "https://registry.npmjs.org/lodash", false}, // npm egress, blocked {4, "http", "https://api.github.com/repos/more", true}, // github egress {5, "dns", "example.com", true}, // non-http, ignored by top domains - } - logs := make([]database.BoundaryLog, 0, len(seeds)) - for _, s := range seeds { - // A non-empty matched_rule marks the call as allowed; an empty rule - // is stored as NULL, which counts as blocked. - rule := "" - if s.allowed { - rule = "allow " + s.detail - } - logs = append(logs, database.BoundaryLog{ - SessionID: fw, - OwnerID: uuid.NullUUID{UUID: firstUser.UserID, Valid: true}, - SequenceNumber: s.seq, - CapturedAt: now, - CreatedAt: now, - Proto: s.proto, - Method: "GET", - Detail: s.detail, - MatchedRule: sql.NullString{String: rule, Valid: rule != ""}, - }) - } - dbgen.BoundaryLogs(t, db, logs) + {6, "http", "https://api.github.com:8080/repos", true}, // port-suffixed; host stripped to api.github.com + }) res, err := client.AIBridgeGetSessionThreads(ctx, "net-session", uuid.Nil, uuid.Nil, 0) require.NoError(t, err) - // total counts seq 1-5 (LLM call at seq 0 excluded); one blocked. + // total counts seq 1-6 (LLM call at seq 0 excluded); one blocked. require.NotNil(t, res.NetworkCalls) - require.EqualValues(t, 5, res.NetworkCalls.Total) + require.EqualValues(t, 6, res.NetworkCalls.Total) require.EqualValues(t, 1, res.NetworkCalls.Blocked) - // Top domains covers HTTP egress only: github x3, npm x1. The dns log is - // excluded. Two distinct domains, ordered by count descending. - require.Len(t, res.NetworkTopDomains, 2) + // Top domains covers HTTP egress only and is capped at one row (the + // summary card renders a single domain). github wins with 4 HTTP calls: + // seqs 1, 2, 4, and the port-suffixed seq 6 whose host strips to + // api.github.com (proving the port is not treated as a separate host). + // The dns log (seq 5) is excluded from domains. NetworkDomainCount is a + // window aggregate independent of the row cap, so it still reports the + // two distinct HTTP domains (github, npm). + require.Len(t, res.NetworkTopDomains, 1) + require.Equal(t, "api.github.com", res.NetworkTopDomains[0].Domain) + require.EqualValues(t, 4, res.NetworkTopDomains[0].Count) + require.EqualValues(t, 2, res.NetworkDomainCount) + }) + + t.Run("NetworkMultipleInterceptions", func(t *testing.T) { + t.Parallel() + // Two interceptions in the same firewall session must produce two + // consecutive, non-overlapping windows: (0, 5) for the first and + // (5, +inf) for the second. Each interception's own LLM call (logged at + // its own sequence) is excluded by the exclusive lower bound. + db, ps := dbtestutil.NewDB(t) + opts := aibridgeOpts(t) + opts.Options.Database = db + opts.Options.Pubsub = ps + client, _, firstUser := coderdenttest.NewWithDatabase(t, opts) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + fw := uuid.New() + + for _, seq := range []int32{0, 5} { + endedAt := now.Add(time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + StartedAt: now, + ClientSessionID: sql.NullString{String: "multi-net", Valid: true}, + AgentFirewallSessionID: uuid.NullUUID{UUID: fw, Valid: true}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: seq, Valid: true}, + }, &endedAt) + } + + seedBoundaryLogs(t, db, fw, firstUser.UserID, now, []boundaryLogSeed{ + {0, "http", "https://api.github.com/llm", true}, // interception 1 LLM call, excluded + {1, "http", "https://api.github.com/a", true}, // window (0,5) + {2, "http", "https://api.github.com/b", true}, // window (0,5) + {3, "http", "https://registry.npmjs.org/x", false}, // window (0,5), blocked + {5, "http", "https://api.github.com/llm2", true}, // interception 2 LLM call, excluded + {6, "http", "https://api.github.com/c", true}, // window (5,+inf) + {7, "http", "https://registry.npmjs.org/y", false}, // window (5,+inf), blocked + }) + + res, err := client.AIBridgeGetSessionThreads(ctx, "multi-net", uuid.Nil, uuid.Nil, 0) + require.NoError(t, err) + + // Both windows contribute: seqs 1,2,3,6,7. The two LLM calls (0, 5) are + // excluded. Blocked = seqs 3 and 7. + require.NotNil(t, res.NetworkCalls) + require.EqualValues(t, 5, res.NetworkCalls.Total) + require.EqualValues(t, 2, res.NetworkCalls.Blocked) + // github: seqs 1,2,6 = 3; npm: seqs 3,7 = 2. Two distinct domains. + require.Len(t, res.NetworkTopDomains, 1) require.Equal(t, "api.github.com", res.NetworkTopDomains[0].Domain) require.EqualValues(t, 3, res.NetworkTopDomains[0].Count) - require.Equal(t, "registry.npmjs.org", res.NetworkTopDomains[1].Domain) - require.EqualValues(t, 1, res.NetworkTopDomains[1].Count) require.EqualValues(t, 2, res.NetworkDomainCount) }) + t.Run("NetworkSharedFirewallSessionNoBleed", func(t *testing.T) { + t.Parallel() + // Two AI sessions share one firewall session. next_seq considers every + // interception in the firewall session, so session A's window is bounded + // by session B's interception and B's calls never bleed into A's counts. + db, ps := dbtestutil.NewDB(t) + opts := aibridgeOpts(t) + opts.Options.Database = db + opts.Options.Pubsub = ps + client, _, firstUser := coderdenttest.NewWithDatabase(t, opts) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + fw := uuid.New() + + // Session A anchored at firewall seq 0; session B at seq 10. + for _, s := range []struct { + session string + seq int32 + }{{"sess-a", 0}, {"sess-b", 10}} { + endedAt := now.Add(time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "anthropic", + Model: "claude-sonnet-4-20250514", + StartedAt: now, + ClientSessionID: sql.NullString{String: s.session, Valid: true}, + AgentFirewallSessionID: uuid.NullUUID{UUID: fw, Valid: true}, + AgentFirewallSequenceNumber: sql.NullInt32{Int32: s.seq, Valid: true}, + }, &endedAt) + } + + seedBoundaryLogs(t, db, fw, firstUser.UserID, now, []boundaryLogSeed{ + {0, "http", "https://api.github.com/llm-a", true}, // A's LLM call, excluded + {1, "http", "https://api.github.com/a1", true}, // A's window (0,10) + {2, "http", "https://registry.npmjs.org/a2", false}, // A's window (0,10), blocked + {10, "http", "https://api.github.com/llm-b", true}, // B's LLM call, excluded + {11, "http", "https://api.github.com/b1", true}, // B's window (10,+inf) + {12, "http", "https://api.github.com/b2", true}, // B's window (10,+inf) + {13, "http", "https://registry.npmjs.org/b3", false}, // B's window (10,+inf), blocked + }) + + // Session A sees only its own two calls (seqs 1, 2), not B's. + resA, err := client.AIBridgeGetSessionThreads(ctx, "sess-a", uuid.Nil, uuid.Nil, 0) + require.NoError(t, err) + require.NotNil(t, resA.NetworkCalls) + require.EqualValues(t, 2, resA.NetworkCalls.Total) + require.EqualValues(t, 1, resA.NetworkCalls.Blocked) + + // Session B sees only its own three calls (seqs 11, 12, 13), not A's. + resB, err := client.AIBridgeGetSessionThreads(ctx, "sess-b", uuid.Nil, uuid.Nil, 0) + require.NoError(t, err) + require.NotNil(t, resB.NetworkCalls) + require.EqualValues(t, 3, resB.NetworkCalls.Total) + require.EqualValues(t, 1, resB.NetworkCalls.Blocked) + }) + t.Run("NetworkSummaryDisabled", func(t *testing.T) { t.Parallel() client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) From 56621bd41128ecbf09c4dba7313a07606387b8df Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Tue, 28 Jul 2026 12:28:16 +0000 Subject: [PATCH 3/5] perf(coderd/database): bound network-call windows with a sentinel sequence number --- coderd/database/querier.go | 4 +++- coderd/database/queries.sql.go | 19 ++++++++++++++----- coderd/database/queries/aibridge.sql | 19 ++++++++++++++----- 3 files changed, 31 insertions(+), 11 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 8f074c7edc84f..c9ffd35e50c62 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -313,7 +313,9 @@ type sqlcQuerier interface { // interception's seq) within the same firewall session. The exclusive lower // bound drops the interception's own LLM-provider call. next_seq considers all // interceptions in the firewall session so windows never bleed across AI - // sessions that share one firewall session. + // sessions that share one firewall session, and falls back to the maximum + // sequence_number for the last interception so the window stays an + // index-satisfiable range. GetAIBridgeSessionTopDomains(ctx context.Context, arg GetAIBridgeSessionTopDomainsParams) ([]GetAIBridgeSessionTopDomainsRow, error) GetAIBridgeTokenUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeTokenUsage, error) GetAIBridgeToolUsagesByInterceptionID(ctx context.Context, interceptionID uuid.UUID) ([]AIBridgeToolUsage, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index bdf7c2565eafb..6b1cf378d5bad 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1277,7 +1277,7 @@ WITH session_boundary_logs AS ( SELECT bl.detail FROM aibridge_interceptions afi LEFT JOIN LATERAL ( - SELECT MIN(nxt.agent_firewall_sequence_number) AS next_seq + SELECT COALESCE(MIN(nxt.agent_firewall_sequence_number), 2147483647) AS next_seq FROM aibridge_interceptions nxt WHERE nxt.agent_firewall_session_id = afi.agent_firewall_session_id AND nxt.agent_firewall_sequence_number > afi.agent_firewall_sequence_number @@ -1285,7 +1285,7 @@ WITH session_boundary_logs AS ( JOIN boundary_logs bl ON bl.session_id = afi.agent_firewall_session_id AND bl.sequence_number > afi.agent_firewall_sequence_number - AND (w.next_seq IS NULL OR bl.sequence_number < w.next_seq) + AND bl.sequence_number < w.next_seq WHERE afi.session_id = $2::text AND afi.ended_at IS NOT NULL AND afi.agent_firewall_session_id IS NOT NULL @@ -1341,7 +1341,9 @@ type GetAIBridgeSessionTopDomainsRow struct { // interception's seq) within the same firewall session. The exclusive lower // bound drops the interception's own LLM-provider call. next_seq considers all // interceptions in the firewall session so windows never bleed across AI -// sessions that share one firewall session. +// sessions that share one firewall session, and falls back to the maximum +// sequence_number for the last interception so the window stays an +// index-satisfiable range. func (q *sqlQuerier) GetAIBridgeSessionTopDomains(ctx context.Context, arg GetAIBridgeSessionTopDomainsParams) ([]GetAIBridgeSessionTopDomainsRow, error) { rows, err := q.db.QueryContext(ctx, getAIBridgeSessionTopDomains, arg.Limit, arg.SessionID) if err != nil { @@ -2257,12 +2259,19 @@ LEFT JOIN LATERAL ( -- (logged at exactly its sequence number), leaving the agent's other -- egress. next_seq considers all interceptions in the firewall session so -- windows never bleed across AI sessions that share one firewall session. + -- + -- The last interception in a firewall session has no successor, so next_seq + -- falls back to the maximum sequence_number instead of NULL. That keeps the + -- window a plain range that the (session_id, sequence_number) index can + -- satisfy end to end: an OR'd NULL check cannot be an index bound, which + -- made every interception scan the firewall session's logs from its own + -- sequence number to the end and discard the overshoot. SELECT COUNT(*)::bigint AS total, COUNT(*) FILTER (WHERE bl.matched_rule IS NULL)::bigint AS blocked FROM aibridge_interceptions afi LEFT JOIN LATERAL ( - SELECT MIN(nxt.agent_firewall_sequence_number) AS next_seq + SELECT COALESCE(MIN(nxt.agent_firewall_sequence_number), 2147483647) AS next_seq FROM aibridge_interceptions nxt WHERE nxt.agent_firewall_session_id = afi.agent_firewall_session_id AND nxt.agent_firewall_sequence_number > afi.agent_firewall_sequence_number @@ -2270,7 +2279,7 @@ LEFT JOIN LATERAL ( JOIN boundary_logs bl ON bl.session_id = afi.agent_firewall_session_id AND bl.sequence_number > afi.agent_firewall_sequence_number - AND (w.next_seq IS NULL OR bl.sequence_number < w.next_seq) + AND bl.sequence_number < w.next_seq WHERE afi.id = ANY(sr.interception_ids) AND afi.agent_firewall_session_id IS NOT NULL AND afi.agent_firewall_sequence_number IS NOT NULL diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 9521263c46084..767dfa0979b7c 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -528,12 +528,19 @@ LEFT JOIN LATERAL ( -- (logged at exactly its sequence number), leaving the agent's other -- egress. next_seq considers all interceptions in the firewall session so -- windows never bleed across AI sessions that share one firewall session. + -- + -- The last interception in a firewall session has no successor, so next_seq + -- falls back to the maximum sequence_number instead of NULL. That keeps the + -- window a plain range that the (session_id, sequence_number) index can + -- satisfy end to end: an OR'd NULL check cannot be an index bound, which + -- made every interception scan the firewall session's logs from its own + -- sequence number to the end and discard the overshoot. SELECT COUNT(*)::bigint AS total, COUNT(*) FILTER (WHERE bl.matched_rule IS NULL)::bigint AS blocked FROM aibridge_interceptions afi LEFT JOIN LATERAL ( - SELECT MIN(nxt.agent_firewall_sequence_number) AS next_seq + SELECT COALESCE(MIN(nxt.agent_firewall_sequence_number), 2147483647) AS next_seq FROM aibridge_interceptions nxt WHERE nxt.agent_firewall_session_id = afi.agent_firewall_session_id AND nxt.agent_firewall_sequence_number > afi.agent_firewall_sequence_number @@ -541,7 +548,7 @@ LEFT JOIN LATERAL ( JOIN boundary_logs bl ON bl.session_id = afi.agent_firewall_session_id AND bl.sequence_number > afi.agent_firewall_sequence_number - AND (w.next_seq IS NULL OR bl.sequence_number < w.next_seq) + AND bl.sequence_number < w.next_seq WHERE afi.id = ANY(sr.interception_ids) AND afi.agent_firewall_session_id IS NOT NULL AND afi.agent_firewall_sequence_number IS NOT NULL @@ -563,12 +570,14 @@ ORDER BY -- interception's seq) within the same firewall session. The exclusive lower -- bound drops the interception's own LLM-provider call. next_seq considers all -- interceptions in the firewall session so windows never bleed across AI --- sessions that share one firewall session. +-- sessions that share one firewall session, and falls back to the maximum +-- sequence_number for the last interception so the window stays an +-- index-satisfiable range. WITH session_boundary_logs AS ( SELECT bl.detail FROM aibridge_interceptions afi LEFT JOIN LATERAL ( - SELECT MIN(nxt.agent_firewall_sequence_number) AS next_seq + SELECT COALESCE(MIN(nxt.agent_firewall_sequence_number), 2147483647) AS next_seq FROM aibridge_interceptions nxt WHERE nxt.agent_firewall_session_id = afi.agent_firewall_session_id AND nxt.agent_firewall_sequence_number > afi.agent_firewall_sequence_number @@ -576,7 +585,7 @@ WITH session_boundary_logs AS ( JOIN boundary_logs bl ON bl.session_id = afi.agent_firewall_session_id AND bl.sequence_number > afi.agent_firewall_sequence_number - AND (w.next_seq IS NULL OR bl.sequence_number < w.next_seq) + AND bl.sequence_number < w.next_seq WHERE afi.session_id = @session_id::text AND afi.ended_at IS NOT NULL AND afi.agent_firewall_session_id IS NOT NULL From d4fb3200b4d56306826ff51aebcbe16786a0980b Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Thu, 30 Jul 2026 11:53:05 +0200 Subject: [PATCH 4/5] Update coderd/database/queries/aibridge.sql Co-authored-by: Cian Johnston --- coderd/database/queries/aibridge.sql | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 767dfa0979b7c..977d200151b6e 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -529,12 +529,12 @@ LEFT JOIN LATERAL ( -- egress. next_seq considers all interceptions in the firewall session so -- windows never bleed across AI sessions that share one firewall session. -- - -- The last interception in a firewall session has no successor, so next_seq - -- falls back to the maximum sequence_number instead of NULL. That keeps the - -- window a plain range that the (session_id, sequence_number) index can - -- satisfy end to end: an OR'd NULL check cannot be an index bound, which - -- made every interception scan the firewall session's logs from its own - -- sequence number to the end and discard the overshoot. +-- The last interception in a session has no next row, so next_seq uses +-- the largest sequence_number instead of NULL. The lookup stays a plain +-- range, so the (session_id, sequence_number) index answers it alone. +-- With NULL and an OR check, the index cannot bound the range: each +-- interception reads every log to the end of the session and throws +-- most of them away. SELECT COUNT(*)::bigint AS total, COUNT(*) FILTER (WHERE bl.matched_rule IS NULL)::bigint AS blocked From 16d4f5d16cd931e549ba39f55d7e1e32ddbdff5d Mon Sep 17 00:00:00 2001 From: Sas Swart Date: Thu, 30 Jul 2026 10:16:08 +0000 Subject: [PATCH 5/5] make gen --- coderd/database/querier.go | 6 ++++++ coderd/database/queries.sql.go | 12 ++++++------ 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/coderd/database/querier.go b/coderd/database/querier.go index c9ffd35e50c62..0ab566b7d1b98 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -1258,6 +1258,12 @@ type sqlcQuerier interface { // Pagination-first strategy: identify the page of sessions cheaply via a // single GROUP BY scan, then do expensive lateral joins (tokens, prompts, // first-interception metadata) only for the ~page-size result set. + // The last interception in a session has no next row, so next_seq uses + // the largest sequence_number instead of NULL. The lookup stays a plain + // range, so the (session_id, sequence_number) index answers it alone. + // With NULL and an OR check, the index cannot bound the range: each + // interception reads every log to the end of the session and throws + // most of them away. ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeSessionsParams) ([]ListAIBridgeSessionsRow, error) ListAIBridgeTokenUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeTokenUsage, error) ListAIBridgeToolUsagesByInterceptionIDs(ctx context.Context, interceptionIds []uuid.UUID) ([]AIBridgeToolUsage, error) diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 6b1cf378d5bad..7ba20c7fc061d 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -2260,12 +2260,6 @@ LEFT JOIN LATERAL ( -- egress. next_seq considers all interceptions in the firewall session so -- windows never bleed across AI sessions that share one firewall session. -- - -- The last interception in a firewall session has no successor, so next_seq - -- falls back to the maximum sequence_number instead of NULL. That keeps the - -- window a plain range that the (session_id, sequence_number) index can - -- satisfy end to end: an OR'd NULL check cannot be an index bound, which - -- made every interception scan the firewall session's logs from its own - -- sequence number to the end and discard the overshoot. SELECT COUNT(*)::bigint AS total, COUNT(*) FILTER (WHERE bl.matched_rule IS NULL)::bigint AS blocked @@ -2334,6 +2328,12 @@ type ListAIBridgeSessionsRow struct { // Pagination-first strategy: identify the page of sessions cheaply via a // single GROUP BY scan, then do expensive lateral joins (tokens, prompts, // first-interception metadata) only for the ~page-size result set. +// The last interception in a session has no next row, so next_seq uses +// the largest sequence_number instead of NULL. The lookup stays a plain +// range, so the (session_id, sequence_number) index answers it alone. +// With NULL and an OR check, the index cannot bound the range: each +// interception reads every log to the end of the session and throws +// most of them away. func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeSessionsParams) ([]ListAIBridgeSessionsRow, error) { rows, err := q.db.QueryContext(ctx, listAIBridgeSessions, arg.AfterSessionID,