diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 79e987cc084..3e443638d53 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 00000000000..9e9ce32c3aa --- /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 00000000000..26ef4aff3ca --- /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 70940611cce..9ddea849d3f 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 bfb40a67a65..780f2ba1c1a 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 352142a2316..f7a2f10bd3c 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 546fef5d84d..2d4f1e9c440 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 cba71e26178..f976a3bf61e 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) }) }