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
7 changes: 7 additions & 0 deletions coderd/apidoc/docs.go

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

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

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

64 changes: 49 additions & 15 deletions coderd/database/db2sdk/db2sdk.go
Original file line number Diff line number Diff line change
Expand Up @@ -1142,34 +1142,43 @@ func AIBridgeSession(row database.ListAIBridgeSessionsRow) codersdk.AIBridgeSess
return session
}

// AIBridgeSessionThreadsParams groups the session row and its subresources for
// AIBridgeSessionThreads. Named fields avoid transposing the several adjacent
// slice arguments (notably TopDomains and NetworkCalls) at the call site.
type AIBridgeSessionThreadsParams struct {
Session database.ListAIBridgeSessionsRow
Interceptions []database.ListAIBridgeSessionThreadsRow
TokenUsages []database.AIBridgeTokenUsage
ToolUsages []database.AIBridgeToolUsage
UserPrompts []database.AIBridgeUserPrompt
ModelThoughts []database.AIBridgeModelThought
TopDomains []database.GetAIBridgeSessionTopDomainsRow
NetworkCalls []database.BoundaryLog
}

// AIBridgeSessionThreads converts session metadata and thread interceptions
// into the threads response. It groups interceptions into threads, builds
// agentic actions from tool usages and model thoughts, and aggregates
// token usage with metadata.
func AIBridgeSessionThreads(
session database.ListAIBridgeSessionsRow,
interceptions []database.ListAIBridgeSessionThreadsRow,
tokenUsages []database.AIBridgeTokenUsage,
toolUsages []database.AIBridgeToolUsage,
userPrompts []database.AIBridgeUserPrompt,
modelThoughts []database.AIBridgeModelThought,
topDomains []database.GetAIBridgeSessionTopDomainsRow,
) codersdk.AIBridgeSessionThreadsResponse {
func AIBridgeSessionThreads(p AIBridgeSessionThreadsParams) codersdk.AIBridgeSessionThreadsResponse {
session := p.Session
interceptions := p.Interceptions

// Index subresources by interception ID.
tokensByInterception := make(map[uuid.UUID][]database.AIBridgeTokenUsage, len(interceptions))
for _, tu := range tokenUsages {
for _, tu := range p.TokenUsages {
tokensByInterception[tu.InterceptionID] = append(tokensByInterception[tu.InterceptionID], tu)
}
toolsByInterception := make(map[uuid.UUID][]database.AIBridgeToolUsage, len(interceptions))
for _, tu := range toolUsages {
for _, tu := range p.ToolUsages {
toolsByInterception[tu.InterceptionID] = append(toolsByInterception[tu.InterceptionID], tu)
}
promptsByInterception := make(map[uuid.UUID][]database.AIBridgeUserPrompt, len(interceptions))
for _, up := range userPrompts {
for _, up := range p.UserPrompts {
promptsByInterception[up.InterceptionID] = append(promptsByInterception[up.InterceptionID], up)
}
thoughtsByInterception := make(map[uuid.UUID][]database.AIBridgeModelThought, len(interceptions))
for _, mt := range modelThoughts {
for _, mt := range p.ModelThoughts {
thoughtsByInterception[mt.InterceptionID] = append(thoughtsByInterception[mt.InterceptionID], mt)
}

Expand Down Expand Up @@ -1207,7 +1216,7 @@ func AIBridgeSessionThreads(

// Aggregate session-level token usage metadata from all token
// usages in the session (not just the page).
sessionTokenMeta := aggregateTokenMetadata(tokenUsages)
sessionTokenMeta := aggregateTokenMetadata(p.TokenUsages)

resp := codersdk.AIBridgeSessionThreadsResponse{
ID: session.SessionID,
Expand Down Expand Up @@ -1253,7 +1262,7 @@ func AIBridgeSessionThreads(
Blocked: session.NetworkCallsBlocked,
}
}
for _, d := range topDomains {
for _, d := range p.TopDomains {
resp.NetworkTopDomains = append(resp.NetworkTopDomains, codersdk.AIBridgeSessionNetworkDomain{
Domain: d.Domain,
Count: d.Count,
Expand All @@ -1262,9 +1271,34 @@ func AIBridgeSessionThreads(
// from the last row processed.
resp.NetworkDomainCount = d.TotalDomains
}
resp.NetworkCallLogs = AgentFirewallLogs(p.NetworkCalls)
return resp
}

// AgentFirewallLogs converts boundary logs to their SDK representation.
// Allowed is derived from MatchedRule being non-NULL.
func AgentFirewallLogs(logs []database.BoundaryLog) []codersdk.AgentFirewallLog {
results := make([]codersdk.AgentFirewallLog, 0, len(logs))
for _, l := range logs {
bl := codersdk.AgentFirewallLog{
ID: l.ID,
SessionID: l.SessionID,
SequenceNumber: l.SequenceNumber,
Allowed: l.MatchedRule.Valid,
CreatedAt: l.CreatedAt,
Proto: l.Proto,
Method: l.Method,
Detail: l.Detail,
CapturedAt: &l.CapturedAt,
}
if l.MatchedRule.Valid {
bl.MatchedRule = &l.MatchedRule.String
}
results = append(results, bl)
}
return results
}

func buildAIBridgeThread(
threadID uuid.UUID,
interceptions []database.AIBridgeInterception,
Expand Down
7 changes: 7 additions & 0 deletions coderd/database/dbauthz/dbauthz.go
Original file line number Diff line number Diff line change
Expand Up @@ -6836,6 +6836,13 @@ func (q *querier) ListAIBridgeModels(ctx context.Context, arg database.ListAIBri
return q.db.ListAuthorizedAIBridgeModels(ctx, arg, prep)
}

func (q *querier) ListAIBridgeSessionNetworkCalls(ctx context.Context, arg database.ListAIBridgeSessionNetworkCallsParams) ([]database.BoundaryLog, error) {
if err := q.authorizeContext(ctx, policy.ActionRead, rbac.ResourceAibridgeInterception); err != nil {
return nil, err
}
return q.db.ListAIBridgeSessionNetworkCalls(ctx, arg)
}

func (q *querier) ListAIBridgeSessionThreads(ctx context.Context, arg database.ListAIBridgeSessionThreadsParams) ([]database.ListAIBridgeSessionThreadsRow, error) {
prep, err := prepareSQLFilter(ctx, q.auth, policy.ActionRead, rbac.ResourceAibridgeInterception.Type)
if err != nil {
Expand Down
6 changes: 6 additions & 0 deletions coderd/database/dbauthz/dbauthz_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6975,6 +6975,12 @@ func (s *MethodTestSuite) TestAIBridge() {
check.Args(params).Asserts(rbac.ResourceAibridgeInterception, policy.ActionRead).Returns([]database.GetAIBridgeSessionTopDomainsRow{})
}))

s.Run("ListAIBridgeSessionNetworkCalls", s.Mocked(func(db *dbmock.MockStore, faker *gofakeit.Faker, check *expects) {
params := database.ListAIBridgeSessionNetworkCallsParams{SessionID: "sess", Limit: 1000}
db.EXPECT().ListAIBridgeSessionNetworkCalls(gomock.Any(), params).Return([]database.BoundaryLog{}, nil).AnyTimes()
check.Args(params).Asserts(rbac.ResourceAibridgeInterception, policy.ActionRead).Returns([]database.BoundaryLog{})
}))

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()
Expand Down
8 changes: 8 additions & 0 deletions coderd/database/dbmetrics/querymetrics.go

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

15 changes: 15 additions & 0 deletions coderd/database/dbmock/dbmock.go

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

20 changes: 20 additions & 0 deletions coderd/database/querier.go

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

79 changes: 79 additions & 0 deletions coderd/database/queries.sql.go

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

39 changes: 39 additions & 0 deletions coderd/database/queries/aibridge.sql
Original file line number Diff line number Diff line change
Expand Up @@ -618,6 +618,45 @@ FROM domains
ORDER BY count DESC, domain ASC
LIMIT COALESCE(NULLIF(@limit_::integer, 0), 5);

-- name: ListAIBridgeSessionNetworkCalls :many
-- Returns the individual Agent Firewall network calls made during an AI
-- session, ordered chronologically. All protocols are included, unlike
-- GetAIBridgeSessionTopDomains which considers only HTTP egress, so the list
-- covers the same events the network_calls summary in ListAIBridgeSessions
-- counts. The list is capped at @limit_ rows, so its length equals the summary
-- total only for sessions at or below the cap. The summary stays authoritative
-- for whole-session totals.
--
-- Windowing mirrors that summary and GetAIBridgeSessionTopDomains: 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, and falls back to the maximum
-- sequence_number for the last interception so the window stays an
-- index-satisfiable range.
SELECT bl.*
FROM aibridge_interceptions afi
LEFT JOIN LATERAL (
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
) 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 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
-- created_at leads because a session can span several firewall sessions, whose
-- sequence numbers are independent streams. id breaks remaining ties so the row
-- that lands on the limit boundary is stable across identical requests.
ORDER BY bl.created_at ASC, bl.sequence_number ASC, bl.id ASC
LIMIT COALESCE(NULLIF(@limit_::integer, 0), 1000);

-- name: ListAIBridgeSessionThreads :many
-- Returns all interceptions belonging to paginated threads within a session.
-- Threads are paginated by (started_at, thread_id) cursor.
Expand Down
8 changes: 7 additions & 1 deletion codersdk/aibridge.go
Original file line number Diff line number Diff line change
Expand Up @@ -202,7 +202,13 @@ type AIBridgeSessionThreadsResponse struct {
// 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"`
// NetworkCallLogs is the chronological list of individual network calls made
// during the session, holding the earliest calls up to a server-side cap.
// NetworkCalls remains authoritative for whole-session totals, so a shorter
// list than NetworkCalls.Total means the list was truncated. Empty when the
// session did not pass through Agent Firewall.
NetworkCallLogs []AgentFirewallLog `json:"network_call_logs,omitempty"`
Threads []AIBridgeThread `json:"threads"`

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.

Threads will be always present, right? no need for omitempty ?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

There is no omitempty on threads? Only on NetworkCallLogs, which is absent in the case where agent firewall wasn't used.

}

// AIBridgeSessionThreadsTokenUsage represents aggregated token usage
Expand Down
14 changes: 14 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