opensource/coder
3
0
Fork 0
mirror of https://github.com/coder/coder.git synced 2026-09-24 15:04:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add network calls column to AI sessions table (#27269)

Browse Source 
Add a "Total/blocked network calls" column to the AIBridge sessions
table.

Update `ListAIBridgeSessions` query to calculate network called made and
blocked per session. See query plan
[here](https://explain.dalibo.com/plan/54355c90b165ggb4).

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Sas Swart
2026-07-21 14:34:23 +02:00
committed by GitHub
co-authored by Claude Opus 4.8
parent bce9ba3356
commit a9a1dcc65d
14 changed files with 353 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
coderd/database/db2sdk/db2sdk.go
+9
View File
@@ -1114,6 +1114,15 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
CacheWriteInputTokens: row.CacheWriteInputTokens,
},
}
// 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 row.FirewallActive {
session.NetworkCalls = &codersdk.AIBridgeSessionNetworkCallSummary{
Total: row.NetworkCallsTotal,
Blocked: row.NetworkCallsBlocked,
}
}
// Ensure non-nil slices for JSON serialization.
if session.Providers == nil {
session.Providers = []string{}
coderd/database/dump.sql
Generated
+2 -2
View File
@@ -4670,7 +4670,7 @@ CREATE INDEX idx_ai_providers_enabled ON ai_providers USING btree (enabled) WHER
CREATE INDEX idx_ai_user_daily_spend_effective_group_id_day ON ai_user_daily_spend USING btree (effective_group_id, day);
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_agent_firewall_session_seq ON aibridge_interceptions USING btree (agent_firewall_session_id, agent_firewall_sequence_number) WHERE (agent_firewall_session_id IS NOT NULL);
CREATE INDEX idx_aibridge_interceptions_client ON aibridge_interceptions USING btree (client);
@@ -4724,7 +4724,7 @@ CREATE INDEX idx_audit_logs_time_desc ON audit_logs USING btree ("time" DESC);
CREATE INDEX idx_boundary_logs_captured_at ON boundary_logs USING btree (captured_at);
CREATE INDEX idx_boundary_logs_session_seq ON boundary_logs USING btree (session_id, sequence_number);
CREATE INDEX idx_boundary_logs_session_seq ON boundary_logs USING btree (session_id, sequence_number) INCLUDE (matched_rule);
CREATE INDEX idx_chat_debug_runs_chat_started ON chat_debug_runs USING btree (chat_id, started_at DESC);
coderd/database/migrations/000550_aibridge_firewall_seq_index.down.sql
+10
View File
@@ -0,0 +1,10 @@
DROP INDEX IF EXISTS idx_boundary_logs_session_seq;
CREATE INDEX idx_boundary_logs_session_seq
ON boundary_logs (session_id, sequence_number);
DROP INDEX IF EXISTS idx_aibridge_interceptions_agent_firewall_session_seq;
CREATE INDEX idx_aibridge_interceptions_agent_firewall_session_id
ON aibridge_interceptions (agent_firewall_session_id)
WHERE agent_firewall_session_id IS NOT NULL;
coderd/database/migrations/000550_aibridge_firewall_seq_index.up.sql
+15
View File
@@ -0,0 +1,15 @@
-- Replace the session-only index with a composite index on
-- (agent_firewall_session_id, agent_firewall_sequence_number). The sessions
-- list computes each interception's next firewall sequence number to bound the
-- boundary_logs it triggered; the composite index serves that lookup index-only
-- and still covers session-only lookups.
DROP INDEX IF EXISTS idx_aibridge_interceptions_agent_firewall_session_id;
CREATE INDEX idx_aibridge_interceptions_agent_firewall_session_seq
ON aibridge_interceptions (agent_firewall_session_id, agent_firewall_sequence_number)
WHERE agent_firewall_session_id IS NOT NULL;
DROP INDEX IF EXISTS idx_boundary_logs_session_seq;
CREATE INDEX idx_boundary_logs_session_seq
ON boundary_logs (session_id, sequence_number) INCLUDE (matched_rule);
coderd/database/modelqueries.go
+3
View File
@@ -1054,6 +1054,9 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeSessions(ctx context.Context, arg Lis
&i.CacheWriteInputTokens,
&i.LastPrompt,
&i.LastActiveAt,
&i.NetworkCallsTotal,
&i.NetworkCallsBlocked,
&i.FirewallActive,
); err != nil {
return nil, err
}
coderd/database/queries.sql.go
Generated
+39 -2
View File
@@ -2115,7 +2115,10 @@ SELECT
COALESCE(st.cache_read_input_tokens, 0)::bigint AS cache_read_input_tokens,
COALESCE(st.cache_write_input_tokens, 0)::bigint AS cache_write_input_tokens,
COALESCE(slp.prompt, '') AS last_prompt,
sp.last_active_at AS last_active_at
sp.last_active_at AS last_active_at,
COALESCE(bnc.total, 0)::bigint AS network_calls_total,
COALESCE(bnc.blocked, 0)::bigint AS network_calls_blocked,
COALESCE(sr.firewall_active, false) AS firewall_active
FROM
session_page sp
JOIN
@@ -2126,7 +2129,8 @@ LEFT JOIN LATERAL (
(ARRAY_AGG(ai.metadata ORDER BY ai.started_at, ai.id))[1] AS metadata,
ARRAY_AGG(DISTINCT ai.provider ORDER BY ai.provider) AS providers,
ARRAY_AGG(DISTINCT ai.model ORDER BY ai.model) AS models,
ARRAY_AGG(ai.id) AS interception_ids
ARRAY_AGG(ai.id) AS interception_ids,
BOOL_OR(ai.agent_firewall_session_id IS NOT NULL) AS firewall_active
FROM aibridge_interceptions ai
WHERE ai.session_id = sp.session_id
AND ai.initiator_id = sp.initiator_id
@@ -2151,6 +2155,33 @@ LEFT JOIN LATERAL (
ORDER BY up.created_at DESC, up.id DESC
LIMIT 1
) slp ON true
LEFT JOIN LATERAL (
-- Count Agent Firewall network calls attributed to this session. Each
-- interception marks a point in its firewall session's monotonic sequence
-- stream; the boundary logs it triggered 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
-- (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.
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
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.id = ANY(sr.interception_ids)
AND afi.agent_firewall_session_id IS NOT NULL
AND afi.agent_firewall_sequence_number IS NOT NULL
) bnc ON true
ORDER BY
sp.last_active_at DESC,
sp.session_id DESC
@@ -2189,6 +2220,9 @@ type ListAIBridgeSessionsRow struct {
CacheWriteInputTokens int64 `db:"cache_write_input_tokens" json:"cache_write_input_tokens"`
LastPrompt string `db:"last_prompt" json:"last_prompt"`
LastActiveAt time.Time `db:"last_active_at" json:"last_active_at"`
NetworkCallsTotal int64 `db:"network_calls_total" json:"network_calls_total"`
NetworkCallsBlocked int64 `db:"network_calls_blocked" json:"network_calls_blocked"`
FirewallActive bool `db:"firewall_active" json:"firewall_active"`
}
// Returns paginated sessions with aggregated metadata, token counts, and
@@ -2238,6 +2272,9 @@ func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeS
&i.CacheWriteInputTokens,
&i.LastPrompt,
&i.LastActiveAt,
&i.NetworkCallsTotal,
&i.NetworkCallsBlocked,
&i.FirewallActive,
); err != nil {
return nil, err
}
coderd/database/queries/aibridge.sql
+33 -2
View File
@@ -479,7 +479,10 @@ SELECT
COALESCE(st.cache_read_input_tokens, 0)::bigint AS cache_read_input_tokens,
COALESCE(st.cache_write_input_tokens, 0)::bigint AS cache_write_input_tokens,
COALESCE(slp.prompt, '') AS last_prompt,
sp.last_active_at AS last_active_at
sp.last_active_at AS last_active_at,
COALESCE(bnc.total, 0)::bigint AS network_calls_total,
COALESCE(bnc.blocked, 0)::bigint AS network_calls_blocked,
COALESCE(sr.firewall_active, false) AS firewall_active
FROM
session_page sp
JOIN
@@ -490,7 +493,8 @@ LEFT JOIN LATERAL (
(ARRAY_AGG(ai.metadata ORDER BY ai.started_at, ai.id))[1] AS metadata,
ARRAY_AGG(DISTINCT ai.provider ORDER BY ai.provider) AS providers,
ARRAY_AGG(DISTINCT ai.model ORDER BY ai.model) AS models,
ARRAY_AGG(ai.id) AS interception_ids
ARRAY_AGG(ai.id) AS interception_ids,
BOOL_OR(ai.agent_firewall_session_id IS NOT NULL) AS firewall_active
FROM aibridge_interceptions ai
WHERE ai.session_id = sp.session_id
AND ai.initiator_id = sp.initiator_id
@@ -515,6 +519,33 @@ LEFT JOIN LATERAL (
ORDER BY up.created_at DESC, up.id DESC
LIMIT 1
) slp ON true
LEFT JOIN LATERAL (
-- Count Agent Firewall network calls attributed to this session. Each
-- interception marks a point in its firewall session's monotonic sequence
-- stream; the boundary logs it triggered 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
-- (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.
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
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.id = ANY(sr.interception_ids)
AND afi.agent_firewall_session_id IS NOT NULL
AND afi.agent_firewall_sequence_number IS NOT NULL
) bnc ON true
ORDER BY
sp.last_active_at DESC,
sp.session_id DESC