Thanks to visit codestin.com
Credit goes to github.com

Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 19 additions & 0 deletions coderd/apidoc/docs.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

19 changes: 19 additions & 0 deletions coderd/apidoc/swagger.json

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

9 changes: 9 additions & 0 deletions coderd/database/db2sdk/db2sdk.go
Original file line number Diff line number Diff line change
Expand Up @@ -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{}
Expand Down
4 changes: 2 additions & 2 deletions coderd/database/dump.sql

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Original file line number Diff line number Diff line change
@@ -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;
Original file line number Diff line number Diff line change
@@ -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);
3 changes: 3 additions & 0 deletions coderd/database/modelqueries.go
Original file line number Diff line number Diff line change
Expand Up @@ -1052,6 +1052,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
}
Expand Down
41 changes: 39 additions & 2 deletions coderd/database/queries.sql.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

35 changes: 33 additions & 2 deletions coderd/database/queries/aibridge.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
17 changes: 15 additions & 2 deletions codersdk/aibridge.go

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be simpler to inline two nil-able int64s instead of making an embedded struct?

NetworkCallsTotal   *int64 `json:"network_calls_count,omitempty"`
NetworkCallsBlocked *int64 `json:"network_calls_blocked,omitempty"`

Original file line number Diff line number Diff line change
Expand Up @@ -66,8 +66,13 @@ type AIBridgeSession struct {
EndedAt *time.Time `json:"ended_at,omitempty" format:"date-time"`
Threads int64 `json:"threads"`
TokenUsageSummary AIBridgeSessionTokenUsageSummary `json:"token_usage_summary"`
LastPrompt *string `json:"last_prompt,omitempty"`
LastActiveAt time.Time `json:"last_active_at" format:"date-time"`
// 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"`
LastPrompt *string `json:"last_prompt,omitempty"`
LastActiveAt time.Time `json:"last_active_at" format:"date-time"`
}

type AIBridgeSessionTokenUsageSummary struct {
Expand All @@ -77,6 +82,14 @@ type AIBridgeSessionTokenUsageSummary struct {
CacheWriteInputTokens int64 `json:"cache_write_input_tokens"`
}

// AIBridgeSessionNetworkCallSummary aggregates the Agent Firewall network
// calls made during a session. Blocked counts calls denied by the firewall
// allow-list.
type AIBridgeSessionNetworkCallSummary struct {
Total int64 `json:"total"`
Blocked int64 `json:"blocked"`
}

type AIBridgeListSessionsResponse struct {
Count int64 `json:"count"`
Sessions []AIBridgeSession `json:"sessions"`
Expand Down
4 changes: 4 additions & 0 deletions docs/reference/api/aigateway.md

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Loading
Loading