From 7c19c7cdb20819f039678d30f7d72e51ff44e606 Mon Sep 17 00:00:00 2001 From: Yevhenii Shcherbina Date: Thu, 13 Aug 2026 15:16:08 +0000 Subject: [PATCH] perf: serve the AI Bridge sessions list from an indexed anti-join ListAIBridgeSessions and CountAIBridgeSessions aggregated every completed interception before applying LIMIT, so a page load scaled with the size of aibridge_interceptions. On a 2,000,000 interception dataset the list took 5.0 to 5.6 s and the count 5.7 s. Represent a session by its latest interception that matches the filters, found with a NOT EXISTS probe for a newer interception in the same session. A partial index on (session_id, initiator_id, started_at DESC, id DESC) answers the probe, so the ordered scan stops after LIMIT sessions and the expensive aggregation runs only for that page. On the same dataset the list takes 126 ms and the count 490 ms. Pagination is now keyset on the (started_at, id) of that row, and cursor_pos resolves the cursor with the same filters, which fixes filtered pagination repeating sessions across pages. The list is now ordered by the session's latest completed interception rather than by last_active_at. last_active_at is still returned for display, and the two can disagree when interceptions in a session overlap in time. --- coderd/database/dump.sql | 2 + ...569_aibridge_session_latest_index.down.sql | 1 + ...00569_aibridge_session_latest_index.up.sql | 11 + coderd/database/modelqueries.go | 8 +- coderd/database/querier.go | 16 +- coderd/database/queries.sql.go | 261 ++++++--- coderd/database/queries/aibridge.sql | 233 ++++++-- enterprise/coderd/aibridge_test.go | 519 ++++++++++++++++-- 8 files changed, 880 insertions(+), 171 deletions(-) create mode 100644 coderd/database/migrations/000569_aibridge_session_latest_index.down.sql create mode 100644 coderd/database/migrations/000569_aibridge_session_latest_index.up.sql diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 79e987cc084e9..3e443638d5309 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -4709,6 +4709,8 @@ CREATE INDEX idx_aibridge_interceptions_provider ON aibridge_interceptions USING CREATE INDEX idx_aibridge_interceptions_session_id ON aibridge_interceptions USING btree (session_id) WHERE (ended_at IS NOT NULL); +CREATE INDEX idx_aibridge_interceptions_session_latest ON aibridge_interceptions USING btree (session_id, initiator_id, started_at DESC, id DESC) WHERE (ended_at IS NOT NULL); + CREATE INDEX idx_aibridge_interceptions_sessions_filter ON aibridge_interceptions USING btree (initiator_id, started_at DESC, id DESC) WHERE (ended_at IS NOT NULL); CREATE INDEX idx_aibridge_interceptions_started_id_desc ON aibridge_interceptions USING btree (started_at DESC, id DESC); diff --git a/coderd/database/migrations/000569_aibridge_session_latest_index.down.sql b/coderd/database/migrations/000569_aibridge_session_latest_index.down.sql new file mode 100644 index 0000000000000..9e9ce32c3aa2f --- /dev/null +++ b/coderd/database/migrations/000569_aibridge_session_latest_index.down.sql @@ -0,0 +1 @@ +DROP INDEX IF EXISTS idx_aibridge_interceptions_session_latest; diff --git a/coderd/database/migrations/000569_aibridge_session_latest_index.up.sql b/coderd/database/migrations/000569_aibridge_session_latest_index.up.sql new file mode 100644 index 0000000000000..26ef4aff3cabf --- /dev/null +++ b/coderd/database/migrations/000569_aibridge_session_latest_index.up.sql @@ -0,0 +1,11 @@ +-- Serves the AI Bridge sessions list. The list needs one row per session, +-- ordered by the session's latest interception, which an anti-join finds by +-- probing for a newer interception in the same session. Ordering the index by +-- (session_id, initiator_id, started_at DESC, id DESC) makes that probe an +-- index-only lookup, so the ordered scan of the list can stop after LIMIT +-- sessions instead of aggregating every interception. +-- +-- Partial on ended_at because the list excludes in-flight interceptions. +CREATE INDEX idx_aibridge_interceptions_session_latest + ON aibridge_interceptions (session_id, initiator_id, started_at DESC, id DESC) + WHERE ended_at IS NOT NULL; diff --git a/coderd/database/modelqueries.go b/coderd/database/modelqueries.go index 70940611ccee4..9ddea849d3fae 100644 --- a/coderd/database/modelqueries.go +++ b/coderd/database/modelqueries.go @@ -1026,15 +1026,17 @@ func (q *sqlQuerier) ListAuthorizedAIBridgeSessions(ctx context.Context, arg Lis } query := fmt.Sprintf("-- name: ListAuthorizedAIBridgeSessions :many\n%s", filtered) + // Argument order matches the placeholders in the generated + // listAIBridgeSessions query. rows, err := q.db.QueryContext(ctx, query, arg.AfterSessionID, arg.StartedAfter, arg.StartedBefore, - arg.InitiatorID, arg.Provider, arg.ProviderName, arg.Model, arg.Client, + arg.InitiatorID, arg.SessionID, arg.Offset, arg.Limit, @@ -1095,14 +1097,16 @@ func (q *sqlQuerier) CountAuthorizedAIBridgeSessions(ctx context.Context, arg Co } query := fmt.Sprintf("-- name: CountAuthorizedAIBridgeSessions :one\n%s", filtered) + // Argument order matches the placeholders in the generated + // countAIBridgeSessions query. rows, err := q.db.QueryContext(ctx, query, arg.StartedAfter, arg.StartedBefore, - arg.InitiatorID, arg.Provider, arg.ProviderName, arg.Model, arg.Client, + arg.InitiatorID, arg.SessionID, ) if err != nil { diff --git a/coderd/database/querier.go b/coderd/database/querier.go index bfb40a67a651c..780f2ba1c1a26 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -91,6 +91,10 @@ type sqlcQuerier interface { CleanTailnetLostPeers(ctx context.Context) error CleanTailnetTunnels(ctx context.Context) error CleanupDeletedMCPServerIDsFromChats(ctx context.Context) error + // Counts one row per session by keeping only each session's latest matching + // interception, the same anti-join ListAIBridgeSessions uses. Counting the + // anti-join lets idx_aibridge_interceptions_session_latest answer the probe, + // instead of hashing every interception to deduplicate session keys. CountAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams) (int64, error) CountAuditLogs(ctx context.Context, arg CountAuditLogsParams) (int64, error) // Cheap queue-length check used by ChatMachine.Update when deciding @@ -1233,9 +1237,15 @@ type sqlcQuerier interface { // the most recent user prompt. A "session" is a logical grouping of // interceptions that share the same session_id (set by the client). // - // 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. + // Pagination-first strategy: identify the page of sessions with an ordered + // index scan that stops after LIMIT rows, then do the expensive aggregation + // (tokens, prompts, first-interception metadata) only for that page. + // + // Sessions are represented by their latest matching interception, found by an + // anti-join probing for a newer interception in the same session. That makes + // the session key a single indexed row, so ordering the list is an index scan + // over idx_aibridge_interceptions_session_latest rather than an aggregate over + // every interception. // 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. diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go index 352142a23168d..f7a2f10bd3cd8 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1030,12 +1030,48 @@ func (q *sqlQuerier) CalculateAIBridgeInterceptionsTelemetrySummary(ctx context. const countAIBridgeSessions = `-- name: CountAIBridgeSessions :one SELECT - COUNT(DISTINCT (aibridge_interceptions.session_id, aibridge_interceptions.initiator_id)) + COUNT(*) FROM aibridge_interceptions WHERE -- Remove inflight interceptions (ones which lack an ended_at value). aibridge_interceptions.ended_at IS NOT NULL + -- Keep one interception per session: the latest one that matches the + -- filters. The filters below are repeated inside the probe so a session + -- is represented by its latest matching interception, not by an interception + -- that the filters exclude. + AND NOT EXISTS ( + SELECT 1 + FROM aibridge_interceptions newer + WHERE newer.session_id = aibridge_interceptions.session_id + AND newer.initiator_id = aibridge_interceptions.initiator_id + AND newer.ended_at IS NOT NULL + AND CASE + WHEN $1::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at >= $1::timestamptz + ELSE true + END + AND CASE + WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at <= $2::timestamptz + ELSE true + END + AND CASE + WHEN $3::text != '' THEN newer.provider = $3::text + ELSE true + END + AND CASE + WHEN $4::text != '' THEN newer.provider_name = $4::text + ELSE true + END + AND CASE + WHEN $5::text != '' THEN newer.model = $5::text + ELSE true + END + AND CASE + WHEN $6::text != '' THEN COALESCE(newer.client, 'Unknown') = $6::text + ELSE true + END + AND (newer.started_at, newer.id) > (aibridge_interceptions.started_at, aibridge_interceptions.id) + ) -- Filter by time frame AND CASE WHEN $1::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at >= $1::timestamptz @@ -1047,27 +1083,27 @@ WHERE END -- Filter initiator_id AND CASE - WHEN $3::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN aibridge_interceptions.initiator_id = $3::uuid + WHEN $7::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN aibridge_interceptions.initiator_id = $7::uuid ELSE true END -- Filter provider AND CASE - WHEN $4::text != '' THEN aibridge_interceptions.provider = $4::text + WHEN $3::text != '' THEN aibridge_interceptions.provider = $3::text ELSE true END -- Filter provider_name AND CASE - WHEN $5::text != '' THEN aibridge_interceptions.provider_name = $5::text + WHEN $4::text != '' THEN aibridge_interceptions.provider_name = $4::text ELSE true END -- Filter model AND CASE - WHEN $6::text != '' THEN aibridge_interceptions.model = $6::text + WHEN $5::text != '' THEN aibridge_interceptions.model = $5::text ELSE true END -- Filter client AND CASE - WHEN $7::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = $7::text + WHEN $6::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = $6::text ELSE true END -- Filter session_id @@ -1082,23 +1118,27 @@ WHERE type CountAIBridgeSessionsParams struct { StartedAfter time.Time `db:"started_after" json:"started_after"` StartedBefore time.Time `db:"started_before" json:"started_before"` - InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` Provider string `db:"provider" json:"provider"` ProviderName string `db:"provider_name" json:"provider_name"` Model string `db:"model" json:"model"` Client string `db:"client" json:"client"` + InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` SessionID string `db:"session_id" json:"session_id"` } +// Counts one row per session by keeping only each session's latest matching +// interception, the same anti-join ListAIBridgeSessions uses. Counting the +// anti-join lets idx_aibridge_interceptions_session_latest answer the probe, +// instead of hashing every interception to deduplicate session keys. func (q *sqlQuerier) CountAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams) (int64, error) { row := q.db.QueryRowContext(ctx, countAIBridgeSessions, arg.StartedAfter, arg.StartedBefore, - arg.InitiatorID, arg.Provider, arg.ProviderName, arg.Model, arg.Client, + arg.InitiatorID, arg.SessionID, ) var count int64 @@ -2221,46 +2261,97 @@ func (q *sqlQuerier) ListAIBridgeSessionThreads(ctx context.Context, arg ListAIB const listAIBridgeSessions = `-- name: ListAIBridgeSessions :many WITH cursor_pos AS ( - -- Resolve the cursor's last_active_at once, outside the HAVING clause, - -- so the planner cannot accidentally re-evaluate it per group. Direct - -- LEFT JOIN is safe here since we only use MAX/MIN aggregates (no COUNT - -- affected by fan-out from multiple prompts per interception). - -- COALESCE falls back to MIN(ai.started_at) so the cursor value is - -- never NULL, which would silently drop rows from the HAVING comparison. - SELECT COALESCE(MAX(up.created_at), MIN(ai.started_at)) AS last_active_at + -- Resolve the cursor session's latest matching interception once. The same + -- filters apply here as in session_keys, so the cursor sits at the same + -- position in the ordering the page uses. + SELECT ai.started_at, ai.id FROM aibridge_interceptions ai - LEFT JOIN aibridge_user_prompts up ON up.interception_id = ai.id - WHERE ai.session_id = $1 AND ai.ended_at IS NOT NULL + WHERE ai.session_id = $1 + AND ai.ended_at IS NOT NULL + AND CASE + WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN ai.started_at >= $2::timestamptz + ELSE true + END + AND CASE + WHEN $3::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN ai.started_at <= $3::timestamptz + ELSE true + END + AND CASE + WHEN $4::text != '' THEN ai.provider = $4::text + ELSE true + END + AND CASE + WHEN $5::text != '' THEN ai.provider_name = $5::text + ELSE true + END + AND CASE + WHEN $6::text != '' THEN ai.model = $6::text + ELSE true + END + AND CASE + WHEN $7::text != '' THEN COALESCE(ai.client, 'Unknown') = $7::text + ELSE true + END + -- session_id is not unique on its own: it derives from a client-supplied + -- value, so two users can present the same one. Restricting the cursor to + -- the filtered initiator keeps paging on the same user's session. + AND CASE + WHEN $8::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ai.initiator_id = $8::uuid + ELSE true + END + ORDER BY ai.started_at DESC, ai.id DESC + LIMIT 1 ), -session_page AS ( - -- Paginate at the session level first; only cheap aggregates here. - -- A lateral correlated subquery for prompts keeps the join one-to-one - -- with aibridge_interceptions so COUNT(*) for thread tallies is not - -- inflated. LIMIT 1 combined with the (interception_id, created_at DESC) - -- index makes this an index-only lookup per interception row rather than - -- a full-table-scan GROUP BY over all prompts. - -- last_active_at is the latest prompt timestamp, falling back to - -- MIN(started_at) for sessions with no prompts. The COALESCE ensures - -- it is never NULL so the HAVING row-value cursor comparison is safe. +session_keys AS ( + -- One row per session: its latest interception that matches the filters. + -- The ORDER BY plus LIMIT means the scan stops as soon as the page is full. SELECT ai.session_id, ai.initiator_id, - MIN(ai.started_at) AS started_at, - MAX(ai.ended_at) AS ended_at, - COUNT(*) FILTER (WHERE ai.thread_root_id IS NULL) AS threads, - COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at))::timestamptz AS last_active_at + ai.started_at AS last_interception_at, + ai.id AS last_interception_id FROM aibridge_interceptions ai - LEFT JOIN LATERAL ( - SELECT created_at AS latest_prompt_at - FROM aibridge_user_prompts - WHERE interception_id = ai.id - ORDER BY created_at DESC - LIMIT 1 - ) latest_prompt ON true WHERE -- Remove inflight interceptions (ones which lack an ended_at value). ai.ended_at IS NOT NULL + -- Keep only the session's latest matching interception. The filters are + -- repeated inside the probe so that a filtered list orders sessions by + -- their latest matching interception. initiator_id and session_id need no + -- repetition: the probe already joins on both, so their filters hold for + -- the probed rows by construction. + AND NOT EXISTS ( + SELECT 1 + FROM aibridge_interceptions newer + WHERE newer.session_id = ai.session_id + AND newer.initiator_id = ai.initiator_id + AND newer.ended_at IS NOT NULL + AND CASE + WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at >= $2::timestamptz + ELSE true + END + AND CASE + WHEN $3::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at <= $3::timestamptz + ELSE true + END + AND CASE + WHEN $4::text != '' THEN newer.provider = $4::text + ELSE true + END + AND CASE + WHEN $5::text != '' THEN newer.provider_name = $5::text + ELSE true + END + AND CASE + WHEN $6::text != '' THEN newer.model = $6::text + ELSE true + END + AND CASE + WHEN $7::text != '' THEN COALESCE(newer.client, 'Unknown') = $7::text + ELSE true + END + AND (newer.started_at, newer.id) > (ai.started_at, ai.id) + ) -- Filter by time frame AND CASE WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN ai.started_at >= $2::timestamptz @@ -2272,27 +2363,27 @@ session_page AS ( END -- Filter initiator_id AND CASE - WHEN $4::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ai.initiator_id = $4::uuid + WHEN $8::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ai.initiator_id = $8::uuid ELSE true END -- Filter provider AND CASE - WHEN $5::text != '' THEN ai.provider = $5::text + WHEN $4::text != '' THEN ai.provider = $4::text ELSE true END -- Filter provider_name AND CASE - WHEN $6::text != '' THEN ai.provider_name = $6::text + WHEN $5::text != '' THEN ai.provider_name = $5::text ELSE true END -- Filter model AND CASE - WHEN $7::text != '' THEN ai.model = $7::text + WHEN $6::text != '' THEN ai.model = $6::text ELSE true END -- Filter client AND CASE - WHEN $8::text != '' THEN COALESCE(ai.client, 'Unknown') = $8::text + WHEN $7::text != '' THEN COALESCE(ai.client, 'Unknown') = $7::text ELSE true END -- Filter session_id @@ -2300,29 +2391,57 @@ session_page AS ( WHEN $9::text != '' THEN ai.session_id = $9::text ELSE true END - -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeSessions - -- @authorize_filter - GROUP BY - ai.session_id, ai.initiator_id - HAVING - -- Cursor pagination: uses a composite (last_active_at, session_id) cursor to - -- support keyset pagination. The less-than comparison matches the DESC - -- sort order so rows after the cursor come later in results. The cursor - -- value comes from cursor_pos to guarantee single evaluation. - CASE - WHEN $1::text != '' THEN ( - (COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at)), ai.session_id) < ( - (SELECT last_active_at FROM cursor_pos), - $1::text - ) - ) + -- Cursor pagination: (started_at, id) of the session's latest matching + -- interception is unique and matches the ORDER BY below, so pages are + -- disjoint and complete. Sessions are keyed by a single row, so a session + -- cannot reappear on a later page through an older interception. + AND CASE + WHEN $1::text != '' THEN + (ai.started_at, ai.id) < (SELECT started_at, id FROM cursor_pos) ELSE true END + -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeSessions. + -- The probe above needs no filter of its own: it joins on initiator_id, and + -- authorization for an interception depends only on that column. + -- @authorize_filter ORDER BY - last_active_at DESC, - ai.session_id DESC + ai.started_at DESC, + ai.id DESC LIMIT COALESCE(NULLIF($11::integer, 0), 100) OFFSET $10 +), +session_page AS ( + -- Aggregate only the page's sessions. A lateral correlated subquery for + -- prompts keeps the join one-to-one with aibridge_interceptions so COUNT(*) + -- for thread tallies is not inflated. LIMIT 1 combined with the + -- (interception_id, created_at DESC) index makes this an index-only lookup + -- per interception row. + -- last_active_at is the latest prompt timestamp, falling back to + -- MIN(started_at) for sessions with no prompts. + SELECT + k.session_id, + k.initiator_id, + k.last_interception_at, + k.last_interception_id, + MIN(ai.started_at) AS started_at, + MAX(ai.ended_at) AS ended_at, + COUNT(*) FILTER (WHERE ai.thread_root_id IS NULL) AS threads, + COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at))::timestamptz AS last_active_at + FROM + session_keys k + JOIN aibridge_interceptions ai + ON ai.session_id = k.session_id + AND ai.initiator_id = k.initiator_id + AND ai.ended_at IS NOT NULL + LEFT JOIN LATERAL ( + SELECT created_at AS latest_prompt_at + FROM aibridge_user_prompts + WHERE interception_id = ai.id + ORDER BY created_at DESC + LIMIT 1 + ) latest_prompt ON true + GROUP BY + k.session_id, k.initiator_id, k.last_interception_at, k.last_interception_id ) SELECT sp.session_id, @@ -2411,19 +2530,19 @@ LEFT JOIN LATERAL ( AND afi.agent_firewall_sequence_number IS NOT NULL ) bnc ON true ORDER BY - sp.last_active_at DESC, - sp.session_id DESC + sp.last_interception_at DESC, + sp.last_interception_id DESC ` type ListAIBridgeSessionsParams struct { AfterSessionID string `db:"after_session_id" json:"after_session_id"` StartedAfter time.Time `db:"started_after" json:"started_after"` StartedBefore time.Time `db:"started_before" json:"started_before"` - InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` Provider string `db:"provider" json:"provider"` ProviderName string `db:"provider_name" json:"provider_name"` Model string `db:"model" json:"model"` Client string `db:"client" json:"client"` + InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` SessionID string `db:"session_id" json:"session_id"` Offset int32 `db:"offset_" json:"offset_"` Limit int32 `db:"limit_" json:"limit_"` @@ -2457,9 +2576,15 @@ type ListAIBridgeSessionsRow struct { // the most recent user prompt. A "session" is a logical grouping of // interceptions that share the same session_id (set by the client). // -// 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. +// Pagination-first strategy: identify the page of sessions with an ordered +// index scan that stops after LIMIT rows, then do the expensive aggregation +// (tokens, prompts, first-interception metadata) only for that page. +// +// Sessions are represented by their latest matching interception, found by an +// anti-join probing for a newer interception in the same session. That makes +// the session key a single indexed row, so ordering the list is an index scan +// over idx_aibridge_interceptions_session_latest rather than an aggregate over +// every interception. // 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. @@ -2471,11 +2596,11 @@ func (q *sqlQuerier) ListAIBridgeSessions(ctx context.Context, arg ListAIBridgeS arg.AfterSessionID, arg.StartedAfter, arg.StartedBefore, - arg.InitiatorID, arg.Provider, arg.ProviderName, arg.Model, arg.Client, + arg.InitiatorID, arg.SessionID, arg.Offset, arg.Limit, diff --git a/coderd/database/queries/aibridge.sql b/coderd/database/queries/aibridge.sql index 546fef5d84dc4..2d4f1e9c44031 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -299,13 +299,53 @@ SELECT ( )::bigint as total_deleted; -- name: CountAIBridgeSessions :one +-- Counts one row per session by keeping only each session's latest matching +-- interception, the same anti-join ListAIBridgeSessions uses. Counting the +-- anti-join lets idx_aibridge_interceptions_session_latest answer the probe, +-- instead of hashing every interception to deduplicate session keys. SELECT - COUNT(DISTINCT (aibridge_interceptions.session_id, aibridge_interceptions.initiator_id)) + COUNT(*) FROM aibridge_interceptions WHERE -- Remove inflight interceptions (ones which lack an ended_at value). aibridge_interceptions.ended_at IS NOT NULL + -- Keep one interception per session: the latest one that matches the + -- filters. The filters below are repeated inside the probe so a session + -- is represented by its latest matching interception, not by an interception + -- that the filters exclude. + AND NOT EXISTS ( + SELECT 1 + FROM aibridge_interceptions newer + WHERE newer.session_id = aibridge_interceptions.session_id + AND newer.initiator_id = aibridge_interceptions.initiator_id + AND newer.ended_at IS NOT NULL + AND CASE + WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at >= @started_after::timestamptz + ELSE true + END + AND CASE + WHEN @started_before::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at <= @started_before::timestamptz + ELSE true + END + AND CASE + WHEN @provider::text != '' THEN newer.provider = @provider::text + ELSE true + END + AND CASE + WHEN @provider_name::text != '' THEN newer.provider_name = @provider_name::text + ELSE true + END + AND CASE + WHEN @model::text != '' THEN newer.model = @model::text + ELSE true + END + AND CASE + WHEN @client::text != '' THEN COALESCE(newer.client, 'Unknown') = @client::text + ELSE true + END + AND (newer.started_at, newer.id) > (aibridge_interceptions.started_at, aibridge_interceptions.id) + ) -- Filter by time frame AND CASE WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at >= @started_after::timestamptz @@ -354,50 +394,107 @@ WHERE -- the most recent user prompt. A "session" is a logical grouping of -- interceptions that share the same session_id (set by the client). -- --- 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. +-- Pagination-first strategy: identify the page of sessions with an ordered +-- index scan that stops after LIMIT rows, then do the expensive aggregation +-- (tokens, prompts, first-interception metadata) only for that page. +-- +-- Sessions are represented by their latest matching interception, found by an +-- anti-join probing for a newer interception in the same session. That makes +-- the session key a single indexed row, so ordering the list is an index scan +-- over idx_aibridge_interceptions_session_latest rather than an aggregate over +-- every interception. WITH cursor_pos AS ( - -- Resolve the cursor's last_active_at once, outside the HAVING clause, - -- so the planner cannot accidentally re-evaluate it per group. Direct - -- LEFT JOIN is safe here since we only use MAX/MIN aggregates (no COUNT - -- affected by fan-out from multiple prompts per interception). - -- COALESCE falls back to MIN(ai.started_at) so the cursor value is - -- never NULL, which would silently drop rows from the HAVING comparison. - SELECT COALESCE(MAX(up.created_at), MIN(ai.started_at)) AS last_active_at + -- Resolve the cursor session's latest matching interception once. The same + -- filters apply here as in session_keys, so the cursor sits at the same + -- position in the ordering the page uses. + SELECT ai.started_at, ai.id FROM aibridge_interceptions ai - LEFT JOIN aibridge_user_prompts up ON up.interception_id = ai.id - WHERE ai.session_id = @after_session_id AND ai.ended_at IS NOT NULL + WHERE ai.session_id = @after_session_id + AND ai.ended_at IS NOT NULL + AND CASE + WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN ai.started_at >= @started_after::timestamptz + ELSE true + END + AND CASE + WHEN @started_before::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN ai.started_at <= @started_before::timestamptz + ELSE true + END + AND CASE + WHEN @provider::text != '' THEN ai.provider = @provider::text + ELSE true + END + AND CASE + WHEN @provider_name::text != '' THEN ai.provider_name = @provider_name::text + ELSE true + END + AND CASE + WHEN @model::text != '' THEN ai.model = @model::text + ELSE true + END + AND CASE + WHEN @client::text != '' THEN COALESCE(ai.client, 'Unknown') = @client::text + ELSE true + END + -- session_id is not unique on its own: it derives from a client-supplied + -- value, so two users can present the same one. Restricting the cursor to + -- the filtered initiator keeps paging on the same user's session. + AND CASE + WHEN @initiator_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ai.initiator_id = @initiator_id::uuid + ELSE true + END + ORDER BY ai.started_at DESC, ai.id DESC + LIMIT 1 ), -session_page AS ( - -- Paginate at the session level first; only cheap aggregates here. - -- A lateral correlated subquery for prompts keeps the join one-to-one - -- with aibridge_interceptions so COUNT(*) for thread tallies is not - -- inflated. LIMIT 1 combined with the (interception_id, created_at DESC) - -- index makes this an index-only lookup per interception row rather than - -- a full-table-scan GROUP BY over all prompts. - -- last_active_at is the latest prompt timestamp, falling back to - -- MIN(started_at) for sessions with no prompts. The COALESCE ensures - -- it is never NULL so the HAVING row-value cursor comparison is safe. +session_keys AS ( + -- One row per session: its latest interception that matches the filters. + -- The ORDER BY plus LIMIT means the scan stops as soon as the page is full. SELECT ai.session_id, ai.initiator_id, - MIN(ai.started_at) AS started_at, - MAX(ai.ended_at) AS ended_at, - COUNT(*) FILTER (WHERE ai.thread_root_id IS NULL) AS threads, - COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at))::timestamptz AS last_active_at + ai.started_at AS last_interception_at, + ai.id AS last_interception_id FROM aibridge_interceptions ai - LEFT JOIN LATERAL ( - SELECT created_at AS latest_prompt_at - FROM aibridge_user_prompts - WHERE interception_id = ai.id - ORDER BY created_at DESC - LIMIT 1 - ) latest_prompt ON true WHERE -- Remove inflight interceptions (ones which lack an ended_at value). ai.ended_at IS NOT NULL + -- Keep only the session's latest matching interception. The filters are + -- repeated inside the probe so that a filtered list orders sessions by + -- their latest matching interception. initiator_id and session_id need no + -- repetition: the probe already joins on both, so their filters hold for + -- the probed rows by construction. + AND NOT EXISTS ( + SELECT 1 + FROM aibridge_interceptions newer + WHERE newer.session_id = ai.session_id + AND newer.initiator_id = ai.initiator_id + AND newer.ended_at IS NOT NULL + AND CASE + WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at >= @started_after::timestamptz + ELSE true + END + AND CASE + WHEN @started_before::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN newer.started_at <= @started_before::timestamptz + ELSE true + END + AND CASE + WHEN @provider::text != '' THEN newer.provider = @provider::text + ELSE true + END + AND CASE + WHEN @provider_name::text != '' THEN newer.provider_name = @provider_name::text + ELSE true + END + AND CASE + WHEN @model::text != '' THEN newer.model = @model::text + ELSE true + END + AND CASE + WHEN @client::text != '' THEN COALESCE(newer.client, 'Unknown') = @client::text + ELSE true + END + AND (newer.started_at, newer.id) > (ai.started_at, ai.id) + ) -- Filter by time frame AND CASE WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN ai.started_at >= @started_after::timestamptz @@ -437,29 +534,57 @@ session_page AS ( WHEN @session_id::text != '' THEN ai.session_id = @session_id::text ELSE true END - -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeSessions - -- @authorize_filter - GROUP BY - ai.session_id, ai.initiator_id - HAVING - -- Cursor pagination: uses a composite (last_active_at, session_id) cursor to - -- support keyset pagination. The less-than comparison matches the DESC - -- sort order so rows after the cursor come later in results. The cursor - -- value comes from cursor_pos to guarantee single evaluation. - CASE - WHEN @after_session_id::text != '' THEN ( - (COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at)), ai.session_id) < ( - (SELECT last_active_at FROM cursor_pos), - @after_session_id::text - ) - ) + -- Cursor pagination: (started_at, id) of the session's latest matching + -- interception is unique and matches the ORDER BY below, so pages are + -- disjoint and complete. Sessions are keyed by a single row, so a session + -- cannot reappear on a later page through an older interception. + AND CASE + WHEN @after_session_id::text != '' THEN + (ai.started_at, ai.id) < (SELECT started_at, id FROM cursor_pos) ELSE true END + -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeSessions. + -- The probe above needs no filter of its own: it joins on initiator_id, and + -- authorization for an interception depends only on that column. + -- @authorize_filter ORDER BY - last_active_at DESC, - ai.session_id DESC + ai.started_at DESC, + ai.id DESC LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100) OFFSET @offset_ +), +session_page AS ( + -- Aggregate only the page's sessions. A lateral correlated subquery for + -- prompts keeps the join one-to-one with aibridge_interceptions so COUNT(*) + -- for thread tallies is not inflated. LIMIT 1 combined with the + -- (interception_id, created_at DESC) index makes this an index-only lookup + -- per interception row. + -- last_active_at is the latest prompt timestamp, falling back to + -- MIN(started_at) for sessions with no prompts. + SELECT + k.session_id, + k.initiator_id, + k.last_interception_at, + k.last_interception_id, + MIN(ai.started_at) AS started_at, + MAX(ai.ended_at) AS ended_at, + COUNT(*) FILTER (WHERE ai.thread_root_id IS NULL) AS threads, + COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at))::timestamptz AS last_active_at + FROM + session_keys k + JOIN aibridge_interceptions ai + ON ai.session_id = k.session_id + AND ai.initiator_id = k.initiator_id + AND ai.ended_at IS NOT NULL + LEFT JOIN LATERAL ( + SELECT created_at AS latest_prompt_at + FROM aibridge_user_prompts + WHERE interception_id = ai.id + ORDER BY created_at DESC + LIMIT 1 + ) latest_prompt ON true + GROUP BY + k.session_id, k.initiator_id, k.last_interception_at, k.last_interception_id ) SELECT sp.session_id, @@ -554,8 +679,8 @@ LEFT JOIN LATERAL ( AND afi.agent_firewall_sequence_number IS NOT NULL ) bnc ON true ORDER BY - sp.last_active_at DESC, - sp.session_id DESC + sp.last_interception_at DESC, + sp.last_interception_id DESC ; -- name: GetAIBridgeSessionTopDomains :many diff --git a/enterprise/coderd/aibridge_test.go b/enterprise/coderd/aibridge_test.go index cba71e261786c..f976a3bf61ebc 100644 --- a/enterprise/coderd/aibridge_test.go +++ b/enterprise/coderd/aibridge_test.go @@ -87,6 +87,40 @@ func auditLogByNewSpendLimit(t *testing.T, rows []database.GetAuditLogsOffsetRow return matches[0] } +// pageAIBridgeSessionIDs walks the whole session list with keyset pagination, +// asserting that no session id is returned on more than one page, and returns +// the session ids in the order they were paged through. +func pageAIBridgeSessionIDs(ctx context.Context, t *testing.T, client *codersdk.Client, filter codersdk.AIBridgeListSessionsFilter, pageSize int) []string { + t.Helper() + + var ( + ids []string + seen = make(map[string]struct{}) + after string + ) + // Bound the loop so a broken cursor cannot page forever. + for range 100 { + f := filter + f.Pagination = codersdk.Pagination{Limit: pageSize} + f.AfterSessionID = after + res, err := client.AIBridgeListSessions(ctx, f) + require.NoError(t, err) + if len(res.Sessions) == 0 { + return ids + } + require.LessOrEqual(t, len(res.Sessions), pageSize, "page larger than requested limit") + for _, s := range res.Sessions { + _, dup := seen[s.ID] + require.Falsef(t, dup, "session %q returned on more than one page (pages so far: %v)", s.ID, ids) + seen[s.ID] = struct{}{} + ids = append(ids, s.ID) + } + after = res.Sessions[len(res.Sessions)-1].ID + } + t.Fatalf("keyset pagination did not terminate, collected %d ids", len(ids)) + return nil +} + func TestAIBridgeListSessions(t *testing.T) { t.Parallel() @@ -1132,23 +1166,28 @@ func TestAIBridgeListSessions(t *testing.T) { require.Equal(t, "session-c", res.Sessions[2].ID) }) - // PromptlessSessionSortsByStartedAt verifies that a session whose root - // interception has no associated user prompts still appears in results and - // sorts by MIN(started_at) as a fallback. Without the COALESCE fallback a - // NULL last_active_at would cause the HAVING row-value comparison to - // evaluate to NULL (not false), silently dropping the session from all - // result pages. + // PromptlessSessionOrdersByLatestInterception verifies that a session whose + // interception has no associated user prompts still appears in results, and + // that it takes its place in the ordering from its latest completed + // interception. // - // Three sessions are arranged so that the promptless session sits between - // two prompted sessions in sort order: + // The list orders sessions by the started_at of the session's latest + // completed interception. last_active_at is still returned for display: it is + // the latest prompt timestamp, falling back to MIN(started_at) for sessions + // with no prompts. The two can disagree, for example when interceptions in a + // session overlap in time, or when a prompt lands after a later interception + // has already started. // - // A: started=now, prompt=now → last_active_at=now - // B: started=now-1h, NO prompt → last_active_at=now-1h (fallback) - // C: started=now-2h, prompt=now-30m → last_active_at=now-30m + // Three sessions, each with a single interception: // - // Sort order by last_active_at DESC: C (now-30m) > B (now-1h), so: A, C, B. - // B disappearing would indicate the fallback is broken. - t.Run("PromptlessSessionSortsByStartedAt", func(t *testing.T) { + // A: started=now, prompt=now -> last_active_at=now + // B: started=now-1h, NO prompt -> last_active_at=now-1h (fallback) + // C: started=now-2h, prompt=now-30m -> last_active_at=now-30m + // + // Order by latest interception DESC: A, B, C. C sorts last even though its + // last_active_at (now-30m) is the second highest. B disappearing would + // indicate the promptless session is being dropped. + t.Run("PromptlessSessionOrdersByLatestInterception", func(t *testing.T) { t.Parallel() client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) ctx := testutil.Context(t, testutil.WaitLong) @@ -1168,7 +1207,8 @@ func TestAIBridgeListSessions(t *testing.T) { CreatedAt: now, }) - // Session B: no prompt at all, exercises the MIN(started_at) fallback. + // Session B: no prompt at all, exercises the MIN(started_at) fallback + // for last_active_at. bEndedAt := now.Add(time.Minute) bInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ InitiatorID: firstUser.UserID, @@ -1176,8 +1216,8 @@ func TestAIBridgeListSessions(t *testing.T) { ClientSessionID: sql.NullString{String: "session-b", Valid: true}, }, &bEndedAt) - // Session C: has a prompt more recent than B's started_at, so C sorts - // above B even though C started earlier. + // Session C: has a prompt more recent than B's started_at. It still sorts + // below B, because ordering follows the latest completed interception. cEndedAt := now.Add(time.Minute) cInterception := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ InitiatorID: firstUser.UserID, @@ -1190,36 +1230,44 @@ func TestAIBridgeListSessions(t *testing.T) { CreatedAt: now.Add(-30 * time.Minute), }) - //nolint:gocritic // Owner role is irrelevant; testing sort fallback. + //nolint:gocritic // Owner role is irrelevant; testing sort order. res, err := client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{}) require.NoError(t, err) require.Len(t, res.Sessions, 3, "promptless session B must appear in results") - // Expected order: A (last_active_at=now), C (last_active_at=now-30m), B (last_active_at=now-1h via fallback). - require.Equal(t, aInterception.SessionID, res.Sessions[0].ID, "session A should be first") - require.Equal(t, cInterception.SessionID, res.Sessions[1].ID, "session C should be second (prompt=now-30m beats B's started_at=now-1h)") - require.Equal(t, bInterception.SessionID, res.Sessions[2].ID, "session B should be last (no prompt, falls back to started_at=now-1h)") + // Expected order by latest completed interception: A (now), B (now-1h), + // C (now-2h). + require.Equal(t, aInterception.SessionID, res.Sessions[0].ID, "session A should be first (interception at now)") + require.Equal(t, bInterception.SessionID, res.Sessions[1].ID, "session B should be second (interception at now-1h)") + require.Equal(t, cInterception.SessionID, res.Sessions[2].ID, "session C should be last (interception at now-2h)") // All sessions have last_active_at; session B falls back to started_at. require.NotZero(t, res.Sessions[0].LastActiveAt, "session A should have last_active_at set") - require.NotZero(t, res.Sessions[1].LastActiveAt, "session C should have last_active_at set") - require.WithinDuration(t, bInterception.StartedAt, res.Sessions[2].LastActiveAt, time.Millisecond, "session B has no prompts, last_active_at should equal started_at") + require.WithinDuration(t, bInterception.StartedAt, res.Sessions[1].LastActiveAt, time.Millisecond, "session B has no prompts, last_active_at should equal started_at") + // C's last_active_at is higher than B's, yet C sorts below B: the display + // value and the ordering key are allowed to disagree. + require.WithinDuration(t, now.Add(-30*time.Minute), res.Sessions[2].LastActiveAt, time.Millisecond, "session C last_active_at should be its prompt timestamp") + require.True(t, res.Sessions[2].LastActiveAt.After(res.Sessions[1].LastActiveAt), "session C is more recently active than B but still sorts below it") }) - // SortsByLastActive verifies that sessions are ordered by last_active_at. - // Every session here has at least one prompt, so last_active_at equals - // the latest prompt timestamp rather than the started_at fallback. + // SortsByLatestInterception verifies that sessions are ordered by the + // started_at of their latest completed interception. + // + // last_active_at is still returned for display, and is the latest prompt + // timestamp (falling back to MIN(started_at) for promptless sessions). The + // ordering key and last_active_at can disagree, for example when + // interceptions in a session overlap in time, or when a prompt lands after a + // later interception has already started, which is what this fixture does. // - // Three sessions are created with intentionally crossing timestamps so that - // the "prompt time" order differs from the "started_at" order: + // Three sessions with intentionally crossing timestamps: // - // X: started=now, prompt=now → last_active_at = now - // Y: started=now-2h, prompt=now-30m → last_active_at = now-30m - // Z: started=now-1h, prompt=now-1h → last_active_at = now-1h + // X: started=now, prompt=now -> last_active_at = now + // Y: started=now-2h, prompt=now-30m -> last_active_at = now-30m + // Z: started=now-1h, prompt=now-1h -> last_active_at = now-1h // - // Order by started_at DESC: X, Z, Y - // Order by last_active_at DESC: X, Y, Z - t.Run("SortsByLastActive", func(t *testing.T) { + // Order by latest interception DESC: X, Z, Y + // Order by last_active_at DESC: X, Y, Z + t.Run("SortsByLatestInterception", func(t *testing.T) { t.Parallel() client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) ctx := testutil.Context(t, testutil.WaitLong) @@ -1270,16 +1318,399 @@ func TestAIBridgeListSessions(t *testing.T) { require.NoError(t, err) require.Len(t, res.Sessions, 3) - // Expected order: X (now), Y (now-30m), Z (now-1h). - // If sorted by started_at the order would be X, Z, Y. - require.Equal(t, xInterception.SessionID, res.Sessions[0].ID, "session X should be first (prompt=now)") - require.Equal(t, yInterception.SessionID, res.Sessions[1].ID, "session Y should be second (prompt=now-30m beats Z's now-1h)") - require.Equal(t, zInterception.SessionID, res.Sessions[2].ID, "session Z should be last (prompt=now-1h)") + // Expected order: X (now), Z (now-1h), Y (now-2h). + // If sorted by last_active_at the order would be X, Y, Z. + require.Equal(t, xInterception.SessionID, res.Sessions[0].ID, "session X should be first (interception at now)") + require.Equal(t, zInterception.SessionID, res.Sessions[1].ID, "session Z should be second (interception at now-1h)") + require.Equal(t, yInterception.SessionID, res.Sessions[2].ID, "session Y should be last (interception at now-2h)") + + // last_active_at still reports the latest prompt per session, so Y + // reports more recent activity than Z while sorting below it. + require.WithinDuration(t, now, res.Sessions[0].LastActiveAt, time.Millisecond, "session X last_active_at should be its prompt timestamp") + require.WithinDuration(t, now.Add(-time.Hour), res.Sessions[1].LastActiveAt, time.Millisecond, "session Z last_active_at should be its prompt timestamp") + require.WithinDuration(t, now.Add(-30*time.Minute), res.Sessions[2].LastActiveAt, time.Millisecond, "session Y last_active_at should be its prompt timestamp") + require.True(t, res.Sessions[2].LastActiveAt.After(res.Sessions[1].LastActiveAt), "session Y is more recently active than Z but still sorts below it") + }) + + // KeysetPaginationWithModelFilter pages through a filtered list in small + // pages and asserts the pages are disjoint and their union equals the + // unpaginated filtered list. This is the regression test for the cursor + // ignoring the filters, which made filtered pages repeat sessions. + // + // Every matching session also contains a newer non-matching interception + // carrying a newer prompt, so the session's position in the filtered + // ordering must come from its latest *matching* interception. The old query + // resolved the cursor from all of the session's interceptions, so the cursor + // sat later in the ordering than the cursor session itself and the session + // was returned again on the next page. + t.Run("KeysetPaginationWithModelFilter", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + unique := time.Now().UnixNano() + wantModel := fmt.Sprintf("model-want-%d", unique) + otherModel := fmt.Sprintf("model-other-%d", unique) + + // mkInterception creates a completed interception plus one prompt at its + // started_at, so last_active_at tracks the interception timestamps. + mkInterception := func(sessionID, model string, startedAt time.Time) { + endedAt := startedAt.Add(time.Minute) + intc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Model: model, + StartedAt: startedAt, + ClientSessionID: sql.NullString{String: sessionID, Valid: true}, + }, &endedAt) + dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + InterceptionID: intc.ID, + Prompt: fmt.Sprintf("prompt for %s", sessionID), + CreatedAt: startedAt, + }) + } + + const sessionCount = 8 + var wantSessionIDs []string + for i := range sessionCount { + sessionID := fmt.Sprintf("session-%d-%d", unique, i) + base := now.Add(-time.Duration(i) * time.Hour) + + // Sessions 2 and 5 only ever use the other model, so the filtered + // list must skip them. + matches := i != 2 && i != 5 + if matches { + mkInterception(sessionID, wantModel, base) + wantSessionIDs = append(wantSessionIDs, sessionID) + } + + // A newer interception on the other model, with a newer prompt. It + // must not determine the session's position in the filtered + // ordering. + mkInterception(sessionID, otherModel, base.Add(30*time.Minute)) + } + + filter := codersdk.AIBridgeListSessionsFilter{Model: wantModel} - // All sessions have LastActiveAt populated. - require.NotNil(t, res.Sessions[0].LastActiveAt, "session X should have last_active_at set") - require.NotNil(t, res.Sessions[1].LastActiveAt, "session Y should have last_active_at set") - require.NotNil(t, res.Sessions[2].LastActiveAt, "session Z should have last_active_at set") + // Unpaginated filtered list. + //nolint:gocritic // Owner role is irrelevant; testing pagination. + all, err := client.AIBridgeListSessions(ctx, filter) + require.NoError(t, err) + allIDs := make([]string, 0, len(all.Sessions)) + for _, s := range all.Sessions { + allIDs = append(allIDs, s.ID) + } + require.Equal(t, wantSessionIDs, allIDs, "filtered list must contain exactly the matching sessions, newest first") + + // Paging in twos must yield the same sessions, once each, in the same + // order. + pagedIDs := pageAIBridgeSessionIDs(ctx, t, client, filter, 2) + require.Equal(t, allIDs, pagedIDs, "filtered keyset pages must be disjoint and complete") + + // A page size that does not divide the result set exercises the final + // short page. + pagedIDs = pageAIBridgeSessionIDs(ctx, t, client, filter, 4) + require.Equal(t, allIDs, pagedIDs, "filtered keyset pages must be disjoint and complete") + }) + + // KeysetPaginationUnfiltered pages through an unfiltered list and asserts + // the pages are disjoint and their union equals the unpaginated list. + t.Run("KeysetPaginationUnfiltered", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + unique := time.Now().UnixNano() + + // Six sessions, each with two completed interceptions, so a session + // keyed by the wrong interception would show up twice. + const sessionCount = 6 + var wantSessionIDs []string + for i := range sessionCount { + sessionID := fmt.Sprintf("session-%d-%d", unique, i) + base := now.Add(-time.Duration(i) * time.Hour) + for _, offset := range []time.Duration{0, 20 * time.Minute} { + startedAt := base.Add(offset) + endedAt := startedAt.Add(time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + StartedAt: startedAt, + ClientSessionID: sql.NullString{String: sessionID, Valid: true}, + }, &endedAt) + } + wantSessionIDs = append(wantSessionIDs, sessionID) + } + + //nolint:gocritic // Owner role is irrelevant; testing pagination. + all, err := client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{}) + require.NoError(t, err) + allIDs := make([]string, 0, len(all.Sessions)) + for _, s := range all.Sessions { + allIDs = append(allIDs, s.ID) + } + require.Equal(t, wantSessionIDs, allIDs) + + for _, pageSize := range []int{1, 2, 4} { + pagedIDs := pageAIBridgeSessionIDs(ctx, t, client, codersdk.AIBridgeListSessionsFilter{}, pageSize) + require.Equalf(t, allIDs, pagedIDs, "unfiltered keyset pages (size %d) must be disjoint and complete", pageSize) + } + }) + + // OrderedByLatestMatchingInterception verifies that with a per-interception + // filter, sessions are ordered by their latest interception that matches the + // filter, not by their latest interception overall. + // + // S1: wantModel at now-30m, otherModel at now-1m + // S2: wantModel at now-10m + // + // Unfiltered order: S1 (now-1m), S2 (now-10m). + // Filtered by wantModel: S2 (now-10m), S1 (now-30m). + t.Run("OrderedByLatestMatchingInterception", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + unique := time.Now().UnixNano() + wantModel := fmt.Sprintf("model-want-%d", unique) + otherModel := fmt.Sprintf("model-other-%d", unique) + s1 := fmt.Sprintf("session-1-%d", unique) + s2 := fmt.Sprintf("session-2-%d", unique) + + mkInterception := func(sessionID, model string, startedAt time.Time) { + endedAt := startedAt.Add(time.Second) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Model: model, + StartedAt: startedAt, + ClientSessionID: sql.NullString{String: sessionID, Valid: true}, + }, &endedAt) + } + + mkInterception(s1, wantModel, now.Add(-30*time.Minute)) + mkInterception(s1, otherModel, now.Add(-time.Minute)) + mkInterception(s2, wantModel, now.Add(-10*time.Minute)) + + // Unfiltered: S1 leads on its newest interception overall. + //nolint:gocritic // Owner role is irrelevant; testing sort order. + res, err := client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{}) + require.NoError(t, err) + require.Len(t, res.Sessions, 2) + require.Equal(t, s1, res.Sessions[0].ID, "unfiltered: S1 first (newest interception now-1m)") + require.Equal(t, s2, res.Sessions[1].ID, "unfiltered: S2 second (newest interception now-10m)") + + // Filtered: S1's newest matching interception is now-30m, so it sorts + // below S2. + res, err = client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{Model: wantModel}) + require.NoError(t, err) + require.Len(t, res.Sessions, 2) + require.Equal(t, s2, res.Sessions[0].ID, "filtered: S2 first (matching interception now-10m)") + require.Equal(t, s1, res.Sessions[1].ID, "filtered: S1 second (matching interception now-30m)") + }) + + // CountMatchesListedRows verifies the count query agrees with the number of + // rows an unpaginated list returns, with and without a filter. + t.Run("CountMatchesListedRows", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + unique := time.Now().UnixNano() + wantModel := fmt.Sprintf("model-want-%d", unique) + otherModel := fmt.Sprintf("model-other-%d", unique) + + // Five sessions. Sessions 0, 1 and 3 have at least one interception on + // wantModel; sessions 2 and 4 do not. Sessions with multiple + // interceptions must still count once. + const sessionCount = 5 + for i := range sessionCount { + sessionID := fmt.Sprintf("session-%d-%d", unique, i) + base := now.Add(-time.Duration(i) * time.Hour) + model := wantModel + if i == 2 || i == 4 { + model = otherModel + } + for _, offset := range []time.Duration{0, 10 * time.Minute} { + startedAt := base.Add(offset) + endedAt := startedAt.Add(time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Model: model, + StartedAt: startedAt, + ClientSessionID: sql.NullString{String: sessionID, Valid: true}, + }, &endedAt) + } + } + + //nolint:gocritic // Owner role is irrelevant; testing count. + res, err := client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{}) + require.NoError(t, err) + require.Len(t, res.Sessions, sessionCount) + require.EqualValues(t, len(res.Sessions), res.Count, "unfiltered count must match listed rows") + + res, err = client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{Model: wantModel}) + require.NoError(t, err) + require.Len(t, res.Sessions, 3) + require.EqualValues(t, len(res.Sessions), res.Count, "filtered count must match listed rows") + + res, err = client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{Model: otherModel}) + require.NoError(t, err) + require.Len(t, res.Sessions, 2) + require.EqualValues(t, len(res.Sessions), res.Count, "filtered count must match listed rows") + }) + + // NewerInflightInterception verifies a session is still listed, and keyed by + // its latest *completed* interception, when a newer interception is still in + // flight. Sessions whose interceptions are all in flight are excluded, which + // the InflightSessions subtest above covers. + t.Run("NewerInflightInterception", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + unique := time.Now().UnixNano() + mixed := fmt.Sprintf("session-mixed-%d", unique) + inflightOnly := fmt.Sprintf("session-inflight-%d", unique) + other := fmt.Sprintf("session-other-%d", unique) + + // Mixed session: completed interception at now-2h, in-flight one at now. + // Its key is the completed interception, so it sorts below `other`. + mixedEndedAt := now.Add(-2*time.Hour + time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + StartedAt: now.Add(-2 * time.Hour), + ClientSessionID: sql.NullString{String: mixed, Valid: true}, + }, &mixedEndedAt) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + StartedAt: now, + ClientSessionID: sql.NullString{String: mixed, Valid: true}, + }, nil) + + // Fully in-flight session: excluded entirely. + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + StartedAt: now.Add(-time.Minute), + ClientSessionID: sql.NullString{String: inflightOnly, Valid: true}, + }, nil) + + otherEndedAt := now.Add(-time.Hour + time.Minute) + dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + StartedAt: now.Add(-time.Hour), + ClientSessionID: sql.NullString{String: other, Valid: true}, + }, &otherEndedAt) + + //nolint:gocritic // Owner role is irrelevant; testing inflight handling. + res, err := client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{}) + require.NoError(t, err) + require.Len(t, res.Sessions, 2) + require.EqualValues(t, 2, res.Count) + require.Equal(t, other, res.Sessions[0].ID, "session keyed by its completed interception at now-1h") + require.Equal(t, mixed, res.Sessions[1].ID, "mixed session keyed by its completed interception at now-2h, not the in-flight one at now") + }) + + // AggregatesSpanUnfilteredInterceptions pins the aggregate scope: with a + // per-interception filter, started_at, ended_at, threads and last_active_at + // (and the providers/models/client/token laterals) are computed over all + // completed interceptions of the session, not only the ones matching the + // filter. Only the session's position in the list is decided by its latest + // matching interception. + // + // One session with two standalone (thread-root) interceptions that overlap in + // time: + // + // otherModel: started=now-2h, ended=now-10m, prompt=now-30m + // wantModel: started=now-1h, ended=now-59m, prompt=now-45m + // + // Filtering on wantModel still reports started_at=now-2h, ended_at=now-10m, + // threads=2 and last_active_at=now-30m. Aggregating only the matching + // interception would report now-1h, now-59m, 1 and now-45m. + t.Run("AggregatesSpanUnfilteredInterceptions", func(t *testing.T) { + t.Parallel() + client, db, firstUser := coderdenttest.NewWithDatabase(t, aibridgeOpts(t)) + ctx := testutil.Context(t, testutil.WaitLong) + + now := dbtime.Now() + unique := time.Now().UnixNano() + wantModel := fmt.Sprintf("model-want-%d", unique) + otherModel := fmt.Sprintf("model-other-%d", unique) + sessionID := fmt.Sprintf("session-%d", unique) + + // Non-matching interception: starts first and ends last, so it decides + // both reported timestamps. Its prompt is the most recent in the session. + otherStartedAt := now.Add(-2 * time.Hour) + otherEndedAt := now.Add(-10 * time.Minute) + otherIntc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "openai", + Model: otherModel, + StartedAt: otherStartedAt, + ClientSessionID: sql.NullString{String: sessionID, Valid: true}, + }, &otherEndedAt) + dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + InterceptionID: otherIntc.ID, + Prompt: "prompt on the non-matching interception", + CreatedAt: now.Add(-30 * time.Minute), + }) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: otherIntc.ID, + InputTokens: 200, + OutputTokens: 20, + CreatedAt: otherStartedAt, + }) + + // Matching interception: a second thread root, fully inside the other + // interception's window. + wantStartedAt := now.Add(-time.Hour) + wantEndedAt := now.Add(-59 * time.Minute) + wantIntc := dbgen.AIBridgeInterception(t, db, database.InsertAIBridgeInterceptionParams{ + InitiatorID: firstUser.UserID, + Provider: "anthropic", + Model: wantModel, + StartedAt: wantStartedAt, + ClientSessionID: sql.NullString{String: sessionID, Valid: true}, + }, &wantEndedAt) + dbgen.AIBridgeUserPrompt(t, db, database.InsertAIBridgeUserPromptParams{ + InterceptionID: wantIntc.ID, + Prompt: "prompt on the matching interception", + CreatedAt: now.Add(-45 * time.Minute), + }) + dbgen.AIBridgeTokenUsage(t, db, database.InsertAIBridgeTokenUsageParams{ + InterceptionID: wantIntc.ID, + InputTokens: 100, + OutputTokens: 10, + CreatedAt: wantStartedAt, + }) + + //nolint:gocritic // Owner role is irrelevant; testing aggregate scope. + res, err := client.AIBridgeListSessions(ctx, codersdk.AIBridgeListSessionsFilter{Model: wantModel}) + require.NoError(t, err) + require.Len(t, res.Sessions, 1) + require.EqualValues(t, 1, res.Count) + + s := res.Sessions[0] + require.Equal(t, sessionID, s.ID) + require.WithinDuration(t, otherStartedAt, s.StartedAt, time.Millisecond, + "started_at is MIN over all completed interceptions, including the filtered-out one") + require.NotNil(t, s.EndedAt) + require.WithinDuration(t, otherEndedAt, *s.EndedAt, time.Millisecond, + "ended_at is MAX over all completed interceptions, including the filtered-out one") + require.EqualValues(t, 2, s.Threads, + "threads counts thread roots across all completed interceptions, including the filtered-out one") + require.WithinDuration(t, now.Add(-30*time.Minute), s.LastActiveAt, time.Millisecond, + "last_active_at is the latest prompt across all completed interceptions, including the filtered-out one") + + // The remaining per-session fields are unfiltered for the same reason. + require.ElementsMatch(t, []string{wantModel, otherModel}, s.Models) + require.ElementsMatch(t, []string{"anthropic", "openai"}, s.Providers) + require.NotNil(t, s.LastPrompt) + require.Equal(t, "prompt on the non-matching interception", *s.LastPrompt) + require.EqualValues(t, 300, s.TokenUsageSummary.InputTokens) + require.EqualValues(t, 30, s.TokenUsageSummary.OutputTokens) }) }