diff --git a/coderd/database/dump.sql b/coderd/database/dump.sql index 79e987cc084..816b010b092 100644 --- a/coderd/database/dump.sql +++ b/coderd/database/dump.sql @@ -769,6 +769,65 @@ BEGIN END; $$; +CREATE FUNCTION aibridge_session_merge_value(arr text[], value text) RETURNS text[] + LANGUAGE sql IMMUTABLE + AS $$ + SELECT CASE + WHEN value IS NULL THEN arr + WHEN arr @> ARRAY[value] THEN arr + ELSE arr || value + END; +$$; + +CREATE FUNCTION aibridge_session_track_interception() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + INSERT INTO aibridge_sessions ( + session_id, initiator_id, started_at, last_active_at, + providers, provider_names, models, client + ) + VALUES ( + NEW.session_id, NEW.initiator_id, NEW.started_at, + -- Prompts are recorded before the interception ends, so the prompt + -- trigger finds no row for a session's first interception. Pick them up + -- here instead; GREATEST skips the NULL when there are none. + GREATEST(NEW.started_at, ( + SELECT MAX(created_at) FROM aibridge_user_prompts + WHERE interception_id = NEW.id + )), + ARRAY[NEW.provider], ARRAY[NEW.provider_name], ARRAY[NEW.model], + COALESCE(NEW.client, 'Unknown') + ) + ON CONFLICT (session_id, initiator_id) DO UPDATE SET + started_at = LEAST(aibridge_sessions.started_at, EXCLUDED.started_at), + last_active_at = GREATEST(aibridge_sessions.last_active_at, EXCLUDED.last_active_at), + providers = aibridge_session_merge_value(aibridge_sessions.providers, NEW.provider), + provider_names = aibridge_session_merge_value(aibridge_sessions.provider_names, NEW.provider_name), + models = aibridge_session_merge_value(aibridge_sessions.models, NEW.model); + -- client is deliberately absent: the first interception to complete + -- sets it and later ones leave it alone, since a session_id comes from + -- a single client. + RETURN NULL; +END; +$$; + +CREATE FUNCTION aibridge_session_track_prompt() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + UPDATE aibridge_sessions s + SET last_active_at = GREATEST(s.last_active_at, NEW.created_at) + -- Join aibridge_user_prompts with aibridge_interceptions to enrich the + -- prompt with session_id and initiator_id, then filter the session by them. + FROM aibridge_interceptions ai + WHERE ai.id = NEW.interception_id + AND s.session_id = ai.session_id + AND s.initiator_id = ai.initiator_id; + RETURN NULL; +END; +$$; + CREATE FUNCTION bump_chat_queue_version_on_queued_message_change() RETURNS trigger LANGUAGE plpgsql AS $$ @@ -1644,6 +1703,25 @@ CREATE TABLE aibridge_model_thoughts ( COMMENT ON TABLE aibridge_model_thoughts IS 'Audit log of model thinking in intercepted requests in AI Bridge'; +CREATE TABLE aibridge_sessions ( + session_id text NOT NULL, + initiator_id uuid NOT NULL, + started_at timestamp with time zone NOT NULL, + last_active_at timestamp with time zone NOT NULL, + providers text[] DEFAULT '{}'::text[] NOT NULL, + provider_names text[] DEFAULT '{}'::text[] NOT NULL, + models text[] DEFAULT '{}'::text[] NOT NULL, + client text DEFAULT 'Unknown'::text NOT NULL +); + +COMMENT ON TABLE aibridge_sessions IS 'Materialized view of AI Bridge sessions, maintained by triggers on aibridge_interceptions and aibridge_user_prompts. Each row summarizes the interceptions sharing same session_id and initiator.'; + +COMMENT ON COLUMN aibridge_sessions.started_at IS 'Earliest started_at across the session''s interceptions. Paired with last_active_at so time-range filters can test whether the session overlaps the requested window.'; + +COMMENT ON COLUMN aibridge_sessions.last_active_at IS 'Timestamp of the latest event in the session: the most recent user prompt or interception start, whichever is later. Sort key for the sessions list, and the upper bound for time-range filters.'; + +COMMENT ON COLUMN aibridge_sessions.client IS 'The client that issued the session. Scalar rather than an array because a session_id originates from one client.'; + CREATE TABLE aibridge_token_usages ( id uuid NOT NULL, interception_id uuid NOT NULL, @@ -4275,6 +4353,9 @@ ALTER TABLE ONLY ai_user_daily_spend ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_pkey PRIMARY KEY (id); +ALTER TABLE ONLY aibridge_sessions + ADD CONSTRAINT aibridge_sessions_pkey PRIMARY KEY (session_id, initiator_id); + ALTER TABLE ONLY aibridge_token_usages ADD CONSTRAINT aibridge_token_usages_pkey PRIMARY KEY (id); @@ -4719,6 +4800,18 @@ CREATE INDEX idx_aibridge_interceptions_thread_root_id ON aibridge_interceptions CREATE INDEX idx_aibridge_model_thoughts_interception_id ON aibridge_model_thoughts USING btree (interception_id); +CREATE INDEX idx_aibridge_sessions_client ON aibridge_sessions USING btree (client); + +CREATE INDEX idx_aibridge_sessions_initiator ON aibridge_sessions USING btree (initiator_id); + +CREATE INDEX idx_aibridge_sessions_last_active ON aibridge_sessions USING btree (last_active_at DESC, session_id DESC); + +CREATE INDEX idx_aibridge_sessions_models ON aibridge_sessions USING gin (models); + +CREATE INDEX idx_aibridge_sessions_provider_names ON aibridge_sessions USING gin (provider_names); + +CREATE INDEX idx_aibridge_sessions_providers ON aibridge_sessions USING gin (providers); + CREATE INDEX idx_aibridge_token_usages_effective_group_id_created_at ON aibridge_token_usages USING btree (effective_group_id, created_at) WHERE (effective_group_id IS NOT NULL); CREATE INDEX idx_aibridge_token_usages_interception_id ON aibridge_token_usages USING btree (interception_id); @@ -5055,6 +5148,10 @@ CREATE OR REPLACE VIEW provisioner_job_stats AS LEFT JOIN provisioner_job_timings pjt ON ((pjt.job_id = pj.id))) GROUP BY pj.id, wb.workspace_id; +CREATE TRIGGER aibridge_interceptions_track_session AFTER INSERT OR UPDATE ON aibridge_interceptions FOR EACH ROW WHEN ((new.ended_at IS NOT NULL)) EXECUTE FUNCTION aibridge_session_track_interception(); + +CREATE TRIGGER aibridge_user_prompts_track_session AFTER INSERT ON aibridge_user_prompts FOR EACH ROW EXECUTE FUNCTION aibridge_session_track_prompt(); + CREATE TRIGGER inhibit_enqueue_if_disabled BEFORE INSERT ON notification_messages FOR EACH ROW EXECUTE FUNCTION inhibit_enqueue_if_disabled(); CREATE TRIGGER protect_deleting_organizations BEFORE UPDATE ON organizations FOR EACH ROW WHEN (((new.deleted = true) AND (old.deleted = false))) EXECUTE FUNCTION protect_deleting_organizations(); @@ -5138,6 +5235,9 @@ ALTER TABLE ONLY ai_seat_state ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id); +ALTER TABLE ONLY aibridge_sessions + ADD CONSTRAINT aibridge_sessions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id) ON DELETE CASCADE; + ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_user_id_uuid_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; diff --git a/coderd/database/foreign_key_constraint.go b/coderd/database/foreign_key_constraint.go index a24e3d73b5c..41009b4ed4e 100644 --- a/coderd/database/foreign_key_constraint.go +++ b/coderd/database/foreign_key_constraint.go @@ -11,6 +11,7 @@ const ( ForeignKeyAIProvidersSettingsKeyID ForeignKeyConstraint = "ai_providers_settings_key_id_fkey" // ALTER TABLE ONLY ai_providers ADD CONSTRAINT ai_providers_settings_key_id_fkey FOREIGN KEY (settings_key_id) REFERENCES dbcrypt_keys(active_key_digest); ForeignKeyAISeatStateUserID ForeignKeyConstraint = "ai_seat_state_user_id_fkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_user_id_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyAibridgeInterceptionsInitiatorID ForeignKeyConstraint = "aibridge_interceptions_initiator_id_fkey" // ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id); + ForeignKeyAibridgeSessionsInitiatorID ForeignKeyConstraint = "aibridge_sessions_initiator_id_fkey" // ALTER TABLE ONLY aibridge_sessions ADD CONSTRAINT aibridge_sessions_initiator_id_fkey FOREIGN KEY (initiator_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyAPIKeysUserIDUUID ForeignKeyConstraint = "api_keys_user_id_uuid_fkey" // ALTER TABLE ONLY api_keys ADD CONSTRAINT api_keys_user_id_uuid_fkey FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE; ForeignKeyBoundaryLogsOwnerID ForeignKeyConstraint = "boundary_logs_owner_id_fkey" // ALTER TABLE ONLY boundary_logs ADD CONSTRAINT boundary_logs_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE SET NULL; ForeignKeyBoundarySessionsOwnerID ForeignKeyConstraint = "boundary_sessions_owner_id_fkey" // ALTER TABLE ONLY boundary_sessions ADD CONSTRAINT boundary_sessions_owner_id_fkey FOREIGN KEY (owner_id) REFERENCES users(id) ON DELETE SET NULL; diff --git a/coderd/database/migrations/000569_aibridge_sessions.down.sql b/coderd/database/migrations/000569_aibridge_sessions.down.sql new file mode 100644 index 00000000000..93cb2ac45d2 --- /dev/null +++ b/coderd/database/migrations/000569_aibridge_sessions.down.sql @@ -0,0 +1,8 @@ +DROP TRIGGER IF EXISTS aibridge_user_prompts_track_session ON aibridge_user_prompts; +DROP TRIGGER IF EXISTS aibridge_interceptions_track_session ON aibridge_interceptions; + +DROP FUNCTION IF EXISTS aibridge_session_track_prompt(); +DROP FUNCTION IF EXISTS aibridge_session_track_interception(); +DROP FUNCTION IF EXISTS aibridge_session_merge_value(text[], text); + +DROP TABLE IF EXISTS aibridge_sessions; diff --git a/coderd/database/migrations/000569_aibridge_sessions.up.sql b/coderd/database/migrations/000569_aibridge_sessions.up.sql new file mode 100644 index 00000000000..8fde89c0faf --- /dev/null +++ b/coderd/database/migrations/000569_aibridge_sessions.up.sql @@ -0,0 +1,162 @@ +-- Materializes AI Bridge sessions, the logical grouping of interceptions +-- sharing a session_id. Without a table of their own, ordering and filtering +-- the sessions list meant aggregating every interception on each page load, and +-- no index could serve the ordering. Storing the ordering key and the +-- filterable attributes here makes a page an index scan that stops after LIMIT +-- rows. +-- +-- The two timestamps are used for time-range filtering and sorting: +-- +-- column | definition | used for +-- ---------------+---------------------------------------+------------------------ +-- started_at | MIN(interception.started_at) | filtering +-- last_active_at | MAX(prompt times, interception starts)| ordering, filtering +CREATE TABLE aibridge_sessions ( + session_id text NOT NULL, + initiator_id uuid NOT NULL REFERENCES users(id) ON DELETE CASCADE, + started_at timestamptz NOT NULL, + -- Ordering key: the latest event in the session, meaning the most recent + -- user prompt or interception start, whichever is later. + last_active_at timestamptz NOT NULL, + -- Filter attributes belong to interceptions, so they are denormalized here + -- to keep filtered pages on the index. Sessions genuinely may span several + -- providers and models, so those are arrays; client is a scalar because a + -- session_id is issued by a single client. + providers text[] NOT NULL DEFAULT '{}', + provider_names text[] NOT NULL DEFAULT '{}', + models text[] NOT NULL DEFAULT '{}', + client text NOT NULL DEFAULT 'Unknown', + -- session_id alone is not unique: it derives from the client-supplied + -- client_session_id, so two users can present the same value. Keying on + -- both columns keeps their sessions separate, matching the + -- GROUP BY session_id, initiator_id the query used before. + PRIMARY KEY (session_id, initiator_id) +); + +COMMENT ON TABLE aibridge_sessions IS 'Materialized view of AI Bridge sessions, maintained by triggers on aibridge_interceptions and aibridge_user_prompts. Each row summarizes the interceptions sharing same session_id and initiator.'; +COMMENT ON COLUMN aibridge_sessions.started_at IS 'Earliest started_at across the session''s interceptions. Paired with last_active_at so time-range filters can test whether the session overlaps the requested window.'; +COMMENT ON COLUMN aibridge_sessions.last_active_at IS 'Timestamp of the latest event in the session: the most recent user prompt or interception start, whichever is later. Sort key for the sessions list, and the upper bound for time-range filters.'; +COMMENT ON COLUMN aibridge_sessions.client IS 'The client that issued the session. Scalar rather than an array because a session_id originates from one client.'; + +-- Serves ORDER BY last_active_at DESC, session_id DESC LIMIT n for the ListAIBridgeSessions query. +CREATE INDEX idx_aibridge_sessions_last_active + ON aibridge_sessions (last_active_at DESC, session_id DESC); +-- `started_at` is deliberately left unindexed for now, as indexing it doesn't seem to provide much benefit. +-- Revisit later if necessary. + +-- Answers the initiator and client filters. +CREATE INDEX idx_aibridge_sessions_initiator + ON aibridge_sessions (initiator_id); +CREATE INDEX idx_aibridge_sessions_client + ON aibridge_sessions (client); + +-- Answers the array membership filters. +CREATE INDEX idx_aibridge_sessions_providers + ON aibridge_sessions USING gin (providers); +CREATE INDEX idx_aibridge_sessions_provider_names + ON aibridge_sessions USING gin (provider_names); +CREATE INDEX idx_aibridge_sessions_models + ON aibridge_sessions USING gin (models); + +-- Adds value to arr only when absent, so repeated interceptions with the same +-- provider or model do not grow the arrays without bound. +CREATE FUNCTION aibridge_session_merge_value(arr text[], value text) RETURNS text[] + LANGUAGE sql + IMMUTABLE + AS $$ + SELECT CASE + WHEN value IS NULL THEN arr + WHEN arr @> ARRAY[value] THEN arr + ELSE arr || value + END; +$$; + +-- Upserts the session row when an interception completes. +-- +-- Every accumulator is monotonic: last_active_at only moves forward, started_at +-- only moves back, and the arrays only grow. Interceptions can therefore arrive +-- in any order and out of order relative to prompts, and the row converges on +-- the same values. +CREATE FUNCTION aibridge_session_track_interception() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + INSERT INTO aibridge_sessions ( + session_id, initiator_id, started_at, last_active_at, + providers, provider_names, models, client + ) + VALUES ( + NEW.session_id, NEW.initiator_id, NEW.started_at, + -- Prompts are recorded before the interception ends, so the prompt + -- trigger finds no row for a session's first interception. Pick them up + -- here instead; GREATEST skips the NULL when there are none. + GREATEST(NEW.started_at, ( + SELECT MAX(created_at) FROM aibridge_user_prompts + WHERE interception_id = NEW.id + )), + ARRAY[NEW.provider], ARRAY[NEW.provider_name], ARRAY[NEW.model], + COALESCE(NEW.client, 'Unknown') + ) + ON CONFLICT (session_id, initiator_id) DO UPDATE SET + started_at = LEAST(aibridge_sessions.started_at, EXCLUDED.started_at), + last_active_at = GREATEST(aibridge_sessions.last_active_at, EXCLUDED.last_active_at), + providers = aibridge_session_merge_value(aibridge_sessions.providers, NEW.provider), + provider_names = aibridge_session_merge_value(aibridge_sessions.provider_names, NEW.provider_name), + models = aibridge_session_merge_value(aibridge_sessions.models, NEW.model); + -- client is deliberately absent: the first interception to complete + -- sets it and later ones leave it alone, since a session_id comes from + -- a single client. + RETURN NULL; +END; +$$; + +-- Creates or updates the session row for each completed interception, keeping +-- aibridge_sessions in sync with aibridge_interceptions. +CREATE TRIGGER aibridge_interceptions_track_session + AFTER INSERT OR UPDATE ON aibridge_interceptions + FOR EACH ROW + WHEN (NEW.ended_at IS NOT NULL) + EXECUTE FUNCTION aibridge_session_track_interception(); + +-- Advances the session's last_active_at when a prompt arrives, keeping +-- aibridge_sessions in sync with aibridge_user_prompts. +CREATE FUNCTION aibridge_session_track_prompt() RETURNS trigger + LANGUAGE plpgsql + AS $$ +BEGIN + UPDATE aibridge_sessions s + SET last_active_at = GREATEST(s.last_active_at, NEW.created_at) + -- Join aibridge_user_prompts with aibridge_interceptions to enrich the + -- prompt with session_id and initiator_id, then filter the session by them. + FROM aibridge_interceptions ai + WHERE ai.id = NEW.interception_id + AND s.session_id = ai.session_id + AND s.initiator_id = ai.initiator_id; + RETURN NULL; +END; +$$; + +CREATE TRIGGER aibridge_user_prompts_track_session + AFTER INSERT ON aibridge_user_prompts + FOR EACH ROW + EXECUTE FUNCTION aibridge_session_track_prompt(); + +-- Backfills sessions from existing interceptions. +INSERT INTO aibridge_sessions ( + session_id, initiator_id, started_at, last_active_at, + providers, provider_names, models, client +) +SELECT + ai.session_id, + ai.initiator_id, + MIN(ai.started_at), + GREATEST(MAX(up.created_at), MAX(ai.started_at)), + ARRAY_AGG(DISTINCT ai.provider ORDER BY ai.provider), + ARRAY_AGG(DISTINCT ai.provider_name ORDER BY ai.provider_name), + ARRAY_AGG(DISTINCT ai.model ORDER BY ai.model), + COALESCE((ARRAY_AGG(ai.client ORDER BY ai.started_at, ai.id))[1], 'Unknown') +FROM aibridge_interceptions ai +LEFT JOIN aibridge_user_prompts up ON up.interception_id = ai.id +WHERE ai.ended_at IS NOT NULL +GROUP BY ai.session_id, ai.initiator_id +ON CONFLICT (session_id, initiator_id) DO NOTHING; diff --git a/coderd/database/migrations/testdata/fixtures/000569_aibridge_sessions.up.sql b/coderd/database/migrations/testdata/fixtures/000569_aibridge_sessions.up.sql new file mode 100644 index 00000000000..42da42bb8eb --- /dev/null +++ b/coderd/database/migrations/testdata/fixtures/000569_aibridge_sessions.up.sql @@ -0,0 +1,44 @@ +-- Inserts a completed interception and a prompt so the triggers added in this +-- migration populate aibridge_sessions. +-- +-- The interception carries an ended_at because only completed interceptions are +-- tracked. +INSERT INTO + aibridge_interceptions ( + id, + initiator_id, + provider, + provider_name, + model, + client, + client_session_id, + started_at, + ended_at + ) +VALUES ( + '1f7c4a5e-6b2d-4c8f-9a1e-2d3b4c5e6f70', + '30095c71-380b-457a-8995-97b8ee6e5307', -- admin@coder.com, from 000022_initial_v0.6.6.up.sql + 'anthropic', + 'anthropic-prod', + 'claude-sonnet-4-6', + 'claude-code', + 'fixture-session-1', + '2025-09-15 12:45:13.921148+00', + '2025-09-15 12:45:21.674413+00' + ); + +INSERT INTO + aibridge_user_prompts ( + id, + interception_id, + provider_response_id, + prompt, + created_at + ) +VALUES ( + '2a8d5b6f-7c3e-4d9a-8b2f-3e4c5d6f7a81', + '1f7c4a5e-6b2d-4c8f-9a1e-2d3b4c5e6f70', + 'msg_fixture_1', + 'hello from a migration fixture', + '2025-09-15 12:45:18.000000+00' + ); diff --git a/coderd/database/models.go b/coderd/database/models.go index a68b7e54bc9..060a2f9d6b9 100644 --- a/coderd/database/models.go +++ b/coderd/database/models.go @@ -4872,6 +4872,21 @@ type APIKey struct { AllowList AllowList `db:"allow_list" json:"allow_list"` } +// Materialized view of AI Bridge sessions, maintained by triggers on aibridge_interceptions and aibridge_user_prompts. Each row summarizes the interceptions sharing same session_id and initiator. +type AibridgeSession struct { + SessionID string `db:"session_id" json:"session_id"` + InitiatorID uuid.UUID `db:"initiator_id" json:"initiator_id"` + // Earliest started_at across the session's interceptions. Paired with last_active_at so time-range filters can test whether the session overlaps the requested window. + StartedAt time.Time `db:"started_at" json:"started_at"` + // Timestamp of the latest event in the session: the most recent user prompt or interception start, whichever is later. Sort key for the sessions list, and the upper bound for time-range filters. + LastActiveAt time.Time `db:"last_active_at" json:"last_active_at"` + Providers []string `db:"providers" json:"providers"` + ProviderNames []string `db:"provider_names" json:"provider_names"` + Models []string `db:"models" json:"models"` + // The client that issued the session. Scalar rather than an array because a session_id originates from one client. + Client string `db:"client" json:"client"` +} + type AuditLog struct { ID uuid.UUID `db:"id" json:"id"` Time time.Time `db:"time" json:"time"` diff --git a/coderd/database/querier.go b/coderd/database/querier.go index 7e052926c40..8435d6a1bfe 100644 --- a/coderd/database/querier.go +++ b/coderd/database/querier.go @@ -94,6 +94,9 @@ type sqlcQuerier interface { CleanTailnetLostPeers(ctx context.Context) error CleanTailnetTunnels(ctx context.Context) error CleanupDeletedMCPServerIDsFromChats(ctx context.Context) error + // Counts the sessions ListAIBridgeSessions would return for the same filters. + // Reads aibridge_sessions so the filters are answered by indexes, rather than + // counting distinct groups across every interception. 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 @@ -1236,9 +1239,14 @@ 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 from the + // aibridge_sessions index, then do expensive lateral joins (tokens, prompts, + // per-session aggregates) only for the ~page-size result set. + // + // aibridge_sessions carries the ordering key and the filterable attributes, + // so both ORDER BY and the filters are answered by indexes. Deriving them + // from aibridge_interceptions instead meant grouping every interception on + // every request, which grew linearly with the table. // 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 4697e3499a1..ad3b2b4ba95 100644 --- a/coderd/database/queries.sql.go +++ b/coderd/database/queries.sql.go @@ -1030,49 +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 + aibridge_sessions s WHERE - -- Remove inflight interceptions (ones which lack an ended_at value). - aibridge_interceptions.ended_at IS NOT NULL - -- Filter by time frame - AND CASE - WHEN $1::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at >= $1::timestamptz + -- Filter by time frame. + CASE + WHEN $1::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.last_active_at >= $1::timestamptz ELSE true END AND CASE - WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at <= $2::timestamptz + WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.started_at <= $2::timestamptz ELSE true END -- Filter initiator_id AND CASE - WHEN $3::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN aibridge_interceptions.initiator_id = $3::uuid + WHEN $3::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN s.initiator_id = $3::uuid ELSE true END - -- Filter provider + -- Filter provider. A session can span several providers and models, so + -- these are array membership tests answered by GIN indexes. AND CASE - WHEN $4::text != '' THEN aibridge_interceptions.provider = $4::text + WHEN $4::text != '' THEN s.providers @> ARRAY[$4::text] ELSE true END -- Filter provider_name AND CASE - WHEN $5::text != '' THEN aibridge_interceptions.provider_name = $5::text + WHEN $5::text != '' THEN s.provider_names @> ARRAY[$5::text] ELSE true END -- Filter model AND CASE - WHEN $6::text != '' THEN aibridge_interceptions.model = $6::text + WHEN $6::text != '' THEN s.models @> ARRAY[$6::text] ELSE true END -- Filter client AND CASE - WHEN $7::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = $7::text + WHEN $7::text != '' THEN s.client = $7::text ELSE true END -- Filter session_id AND CASE - WHEN $8::text != '' THEN aibridge_interceptions.session_id = $8::text + WHEN $8::text != '' THEN s.session_id = $8::text ELSE true END -- Authorize Filter clause will be injected below in CountAuthorizedAIBridgeSessions @@ -1090,6 +1089,9 @@ type CountAIBridgeSessionsParams struct { SessionID string `db:"session_id" json:"session_id"` } +// Counts the sessions ListAIBridgeSessions would return for the same filters. +// Reads aibridge_sessions so the filters are answered by indexes, rather than +// counting distinct groups across every interception. func (q *sqlQuerier) CountAIBridgeSessions(ctx context.Context, arg CountAIBridgeSessionsParams) (int64, error) { row := q.db.QueryRowContext(ctx, countAIBridgeSessions, arg.StartedAfter, @@ -1110,7 +1112,7 @@ const deleteOldAIBridgeRecords = `-- name: DeleteOldAIBridgeRecords :one WITH -- We don't have FK relationships between the dependent tables and aibridge_interceptions, so we can't rely on DELETE CASCADE. to_delete AS ( - SELECT id FROM aibridge_interceptions + SELECT id, session_id, initiator_id FROM aibridge_interceptions WHERE started_at < $1::timestamp with time zone ), -- CTEs are executed in order. @@ -1138,13 +1140,30 @@ WITH DELETE FROM aibridge_interceptions WHERE id IN (SELECT id FROM to_delete) RETURNING 1 + ), + -- Delete the session if: + -- There is a deleted interception linked to this session. + -- There are no interceptions linked to this session with a timestamp after the deletion timestamp. + -- + -- Sibling CTEs share the same snapshot, so deleted interceptions are still visible when we delete the session. + sessions AS ( + DELETE FROM aibridge_sessions s + WHERE (s.session_id, s.initiator_id) IN (SELECT session_id, initiator_id FROM to_delete) + AND NOT EXISTS ( + SELECT 1 FROM aibridge_interceptions ai + WHERE ai.session_id = s.session_id + AND ai.initiator_id = s.initiator_id + AND ai.started_at >= $1::timestamp with time zone + ) + RETURNING 1 ) SELECT ( (SELECT COUNT(*) FROM model_thoughts) + (SELECT COUNT(*) FROM tool_usages) + (SELECT COUNT(*) FROM token_usages) + (SELECT COUNT(*) FROM user_prompts) + - (SELECT COUNT(*) FROM interceptions) + (SELECT COUNT(*) FROM interceptions) + + (SELECT COUNT(*) FROM sessions) )::bigint as total_deleted ` @@ -2221,106 +2240,81 @@ 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 - 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 + -- Resolve the cursor's last_active_at once so the planner cannot + -- re-evaluate it per row. MAX collapses the rare case of one session_id + -- belonging to two initiators into a single value, and guarantees exactly + -- one row so the scalar subquery below is always valid. + SELECT MAX(s.last_active_at) AS last_active_at + FROM aibridge_sessions s + WHERE s.session_id = $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. 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 + s.session_id, + s.initiator_id, + s.started_at, + s.last_active_at 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 + aibridge_sessions s WHERE - -- Remove inflight interceptions (ones which lack an ended_at value). - ai.ended_at IS NOT NULL - -- Filter by time frame - AND CASE - WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN ai.started_at >= $2::timestamptz + -- Filter by time frame. + CASE + WHEN $2::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.last_active_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 + WHEN $3::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.started_at <= $3::timestamptz ELSE true END -- Filter initiator_id AND CASE - WHEN $4::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ai.initiator_id = $4::uuid + WHEN $4::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN s.initiator_id = $4::uuid ELSE true END - -- Filter provider + -- Filter provider. A session can span several providers and models, + -- so these are array membership tests answered by GIN indexes. AND CASE - WHEN $5::text != '' THEN ai.provider = $5::text + WHEN $5::text != '' THEN s.providers @> ARRAY[$5::text] ELSE true END -- Filter provider_name AND CASE - WHEN $6::text != '' THEN ai.provider_name = $6::text + WHEN $6::text != '' THEN s.provider_names @> ARRAY[$6::text] ELSE true END -- Filter model AND CASE - WHEN $7::text != '' THEN ai.model = $7::text + WHEN $7::text != '' THEN s.models @> ARRAY[$7::text] ELSE true END -- Filter client AND CASE - WHEN $8::text != '' THEN COALESCE(ai.client, 'Unknown') = $8::text + WHEN $8::text != '' THEN s.client = $8::text ELSE true END -- Filter session_id AND CASE - WHEN $9::text != '' THEN ai.session_id = $9::text + WHEN $9::text != '' THEN s.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 + -- 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. + AND CASE WHEN $1::text != '' THEN ( - (COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at)), ai.session_id) < ( + (s.last_active_at, s.session_id) < ( (SELECT last_active_at FROM cursor_pos), $1::text ) ) ELSE true END + -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeSessions + -- @authorize_filter ORDER BY - last_active_at DESC, - ai.session_id DESC + s.last_active_at DESC, + s.session_id DESC LIMIT COALESCE(NULLIF($11::integer, 0), 100) OFFSET $10 ) @@ -2334,9 +2328,11 @@ SELECT sr.models::text[] AS models, COALESCE(sr.client, '')::varchar(64) AS client, sr.metadata::jsonb AS metadata, - sp.started_at::timestamptz AS started_at, - sp.ended_at::timestamptz AS ended_at, - sp.threads, + -- Fall back to the summary row, which is NOT NULL, so a session that has + -- somehow lost all interceptions still does not return NULL. + COALESCE(sr.started_at, sp.started_at)::timestamptz AS started_at, + COALESCE(sr.ended_at, sp.last_active_at)::timestamptz AS ended_at, + sr.threads, COALESCE(st.input_tokens, 0)::bigint AS input_tokens, COALESCE(st.output_tokens, 0)::bigint AS output_tokens, COALESCE(st.cache_read_input_tokens, 0)::bigint AS cache_read_input_tokens, @@ -2351,13 +2347,21 @@ FROM JOIN visible_users ON visible_users.id = sp.initiator_id LEFT JOIN LATERAL ( + -- Per-session aggregates over the page's sessions only. started_at, + -- ended_at and threads are computed here rather than stored on + -- aibridge_sessions because this scan already reads exactly the rows they + -- summarize, so they cost nothing extra and stay consistent with the + -- provider and model arrays alongside them. SELECT (ARRAY_AGG(ai.client ORDER BY ai.started_at, ai.id))[1] AS client, (ARRAY_AGG(ai.metadata ORDER BY ai.started_at, ai.id))[1] AS metadata, ARRAY_AGG(DISTINCT ai.provider ORDER BY ai.provider) AS providers, ARRAY_AGG(DISTINCT ai.model ORDER BY ai.model) AS models, ARRAY_AGG(ai.id) AS interception_ids, - BOOL_OR(ai.agent_firewall_session_id IS NOT NULL) AS firewall_active + BOOL_OR(ai.agent_firewall_session_id IS NOT NULL) AS firewall_active, + 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 FROM aibridge_interceptions ai WHERE ai.session_id = sp.session_id AND ai.initiator_id = sp.initiator_id @@ -2457,9 +2461,14 @@ 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 from the +// aibridge_sessions index, then do expensive lateral joins (tokens, prompts, +// per-session aggregates) only for the ~page-size result set. +// +// aibridge_sessions carries the ordering key and the filterable attributes, +// so both ORDER BY and the filters are answered by indexes. Deriving them +// from aibridge_interceptions instead meant grouping every interception on +// every request, which grew linearly with the table. // 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/aibridge.sql b/coderd/database/queries/aibridge.sql index 546fef5d84d..35753a7c3e6 100644 --- a/coderd/database/queries/aibridge.sql +++ b/coderd/database/queries/aibridge.sql @@ -260,7 +260,7 @@ FROM WITH -- We don't have FK relationships between the dependent tables and aibridge_interceptions, so we can't rely on DELETE CASCADE. to_delete AS ( - SELECT id FROM aibridge_interceptions + SELECT id, session_id, initiator_id FROM aibridge_interceptions WHERE started_at < @before_time::timestamp with time zone ), -- CTEs are executed in order. @@ -288,6 +288,22 @@ WITH DELETE FROM aibridge_interceptions WHERE id IN (SELECT id FROM to_delete) RETURNING 1 + ), + -- Delete the session if: + -- There is a deleted interception linked to this session. + -- There are no interceptions linked to this session with a timestamp after the deletion timestamp. + -- + -- Sibling CTEs share the same snapshot, so deleted interceptions are still visible when we delete the session. + sessions AS ( + DELETE FROM aibridge_sessions s + WHERE (s.session_id, s.initiator_id) IN (SELECT session_id, initiator_id FROM to_delete) + AND NOT EXISTS ( + SELECT 1 FROM aibridge_interceptions ai + WHERE ai.session_id = s.session_id + AND ai.initiator_id = s.initiator_id + AND ai.started_at >= @before_time::timestamp with time zone + ) + RETURNING 1 ) -- Cumulative count. SELECT ( @@ -295,54 +311,57 @@ SELECT ( (SELECT COUNT(*) FROM tool_usages) + (SELECT COUNT(*) FROM token_usages) + (SELECT COUNT(*) FROM user_prompts) + - (SELECT COUNT(*) FROM interceptions) + (SELECT COUNT(*) FROM interceptions) + + (SELECT COUNT(*) FROM sessions) )::bigint as total_deleted; -- name: CountAIBridgeSessions :one +-- Counts the sessions ListAIBridgeSessions would return for the same filters. +-- Reads aibridge_sessions so the filters are answered by indexes, rather than +-- counting distinct groups across every interception. SELECT - COUNT(DISTINCT (aibridge_interceptions.session_id, aibridge_interceptions.initiator_id)) + COUNT(*) FROM - aibridge_interceptions + aibridge_sessions s WHERE - -- Remove inflight interceptions (ones which lack an ended_at value). - aibridge_interceptions.ended_at IS NOT NULL - -- 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 + -- Filter by time frame. + CASE + WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.last_active_at >= @started_after::timestamptz ELSE true END AND CASE - WHEN @started_before::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN aibridge_interceptions.started_at <= @started_before::timestamptz + WHEN @started_before::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.started_at <= @started_before::timestamptz ELSE true END -- Filter initiator_id AND CASE - WHEN @initiator_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN aibridge_interceptions.initiator_id = @initiator_id::uuid + WHEN @initiator_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN s.initiator_id = @initiator_id::uuid ELSE true END - -- Filter provider + -- Filter provider. A session can span several providers and models, so + -- these are array membership tests answered by GIN indexes. AND CASE - WHEN @provider::text != '' THEN aibridge_interceptions.provider = @provider::text + WHEN @provider::text != '' THEN s.providers @> ARRAY[@provider::text] ELSE true END -- Filter provider_name AND CASE - WHEN @provider_name::text != '' THEN aibridge_interceptions.provider_name = @provider_name::text + WHEN @provider_name::text != '' THEN s.provider_names @> ARRAY[@provider_name::text] ELSE true END -- Filter model AND CASE - WHEN @model::text != '' THEN aibridge_interceptions.model = @model::text + WHEN @model::text != '' THEN s.models @> ARRAY[@model::text] ELSE true END -- Filter client AND CASE - WHEN @client::text != '' THEN COALESCE(aibridge_interceptions.client, 'Unknown') = @client::text + WHEN @client::text != '' THEN s.client = @client::text ELSE true END -- Filter session_id AND CASE - WHEN @session_id::text != '' THEN aibridge_interceptions.session_id = @session_id::text + WHEN @session_id::text != '' THEN s.session_id = @session_id::text ELSE true END -- Authorize Filter clause will be injected below in CountAuthorizedAIBridgeSessions @@ -354,110 +373,90 @@ 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 from the +-- aibridge_sessions index, then do expensive lateral joins (tokens, prompts, +-- per-session aggregates) only for the ~page-size result set. +-- +-- aibridge_sessions carries the ordering key and the filterable attributes, +-- so both ORDER BY and the filters are answered by indexes. Deriving them +-- from aibridge_interceptions instead meant grouping every interception on +-- every request, which grew linearly with the table. 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 - 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 + -- Resolve the cursor's last_active_at once so the planner cannot + -- re-evaluate it per row. MAX collapses the rare case of one session_id + -- belonging to two initiators into a single value, and guarantees exactly + -- one row so the scalar subquery below is always valid. + SELECT MAX(s.last_active_at) AS last_active_at + FROM aibridge_sessions s + WHERE s.session_id = @after_session_id ), 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. 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 + s.session_id, + s.initiator_id, + s.started_at, + s.last_active_at 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 + aibridge_sessions s WHERE - -- Remove inflight interceptions (ones which lack an ended_at value). - ai.ended_at IS NOT NULL - -- 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 + -- Filter by time frame. + CASE + WHEN @started_after::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.last_active_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 + WHEN @started_before::timestamptz != '0001-01-01 00:00:00+00'::timestamptz THEN s.started_at <= @started_before::timestamptz ELSE true END -- Filter initiator_id AND CASE - WHEN @initiator_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN ai.initiator_id = @initiator_id::uuid + WHEN @initiator_id::uuid != '00000000-0000-0000-0000-000000000000'::uuid THEN s.initiator_id = @initiator_id::uuid ELSE true END - -- Filter provider + -- Filter provider. A session can span several providers and models, + -- so these are array membership tests answered by GIN indexes. AND CASE - WHEN @provider::text != '' THEN ai.provider = @provider::text + WHEN @provider::text != '' THEN s.providers @> ARRAY[@provider::text] ELSE true END -- Filter provider_name AND CASE - WHEN @provider_name::text != '' THEN ai.provider_name = @provider_name::text + WHEN @provider_name::text != '' THEN s.provider_names @> ARRAY[@provider_name::text] ELSE true END -- Filter model AND CASE - WHEN @model::text != '' THEN ai.model = @model::text + WHEN @model::text != '' THEN s.models @> ARRAY[@model::text] ELSE true END -- Filter client AND CASE - WHEN @client::text != '' THEN COALESCE(ai.client, 'Unknown') = @client::text + WHEN @client::text != '' THEN s.client = @client::text ELSE true END -- Filter session_id AND CASE - WHEN @session_id::text != '' THEN ai.session_id = @session_id::text + WHEN @session_id::text != '' THEN s.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 + -- 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. + AND CASE WHEN @after_session_id::text != '' THEN ( - (COALESCE(MAX(latest_prompt.latest_prompt_at), MIN(ai.started_at)), ai.session_id) < ( + (s.last_active_at, s.session_id) < ( (SELECT last_active_at FROM cursor_pos), @after_session_id::text ) ) ELSE true END + -- Authorize Filter clause will be injected below in ListAuthorizedAIBridgeSessions + -- @authorize_filter ORDER BY - last_active_at DESC, - ai.session_id DESC + s.last_active_at DESC, + s.session_id DESC LIMIT COALESCE(NULLIF(@limit_::integer, 0), 100) OFFSET @offset_ ) @@ -471,9 +470,11 @@ SELECT sr.models::text[] AS models, COALESCE(sr.client, '')::varchar(64) AS client, sr.metadata::jsonb AS metadata, - sp.started_at::timestamptz AS started_at, - sp.ended_at::timestamptz AS ended_at, - sp.threads, + -- Fall back to the summary row, which is NOT NULL, so a session that has + -- somehow lost all interceptions still does not return NULL. + COALESCE(sr.started_at, sp.started_at)::timestamptz AS started_at, + COALESCE(sr.ended_at, sp.last_active_at)::timestamptz AS ended_at, + sr.threads, COALESCE(st.input_tokens, 0)::bigint AS input_tokens, COALESCE(st.output_tokens, 0)::bigint AS output_tokens, COALESCE(st.cache_read_input_tokens, 0)::bigint AS cache_read_input_tokens, @@ -488,13 +489,21 @@ FROM JOIN visible_users ON visible_users.id = sp.initiator_id LEFT JOIN LATERAL ( + -- Per-session aggregates over the page's sessions only. started_at, + -- ended_at and threads are computed here rather than stored on + -- aibridge_sessions because this scan already reads exactly the rows they + -- summarize, so they cost nothing extra and stay consistent with the + -- provider and model arrays alongside them. SELECT (ARRAY_AGG(ai.client ORDER BY ai.started_at, ai.id))[1] AS client, (ARRAY_AGG(ai.metadata ORDER BY ai.started_at, ai.id))[1] AS metadata, ARRAY_AGG(DISTINCT ai.provider ORDER BY ai.provider) AS providers, ARRAY_AGG(DISTINCT ai.model ORDER BY ai.model) AS models, ARRAY_AGG(ai.id) AS interception_ids, - BOOL_OR(ai.agent_firewall_session_id IS NOT NULL) AS firewall_active + BOOL_OR(ai.agent_firewall_session_id IS NOT NULL) AS firewall_active, + 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 FROM aibridge_interceptions ai WHERE ai.session_id = sp.session_id AND ai.initiator_id = sp.initiator_id diff --git a/coderd/database/unique_constraint.go b/coderd/database/unique_constraint.go index 4b1a4376f2d..5b6274e128e 100644 --- a/coderd/database/unique_constraint.go +++ b/coderd/database/unique_constraint.go @@ -14,6 +14,7 @@ const ( UniqueAISeatStatePkey UniqueConstraint = "ai_seat_state_pkey" // ALTER TABLE ONLY ai_seat_state ADD CONSTRAINT ai_seat_state_pkey PRIMARY KEY (user_id); UniqueAIUserDailySpendPkey UniqueConstraint = "ai_user_daily_spend_pkey" // ALTER TABLE ONLY ai_user_daily_spend ADD CONSTRAINT ai_user_daily_spend_pkey PRIMARY KEY (user_id, effective_group_id, day); UniqueAibridgeInterceptionsPkey UniqueConstraint = "aibridge_interceptions_pkey" // ALTER TABLE ONLY aibridge_interceptions ADD CONSTRAINT aibridge_interceptions_pkey PRIMARY KEY (id); + UniqueAibridgeSessionsPkey UniqueConstraint = "aibridge_sessions_pkey" // ALTER TABLE ONLY aibridge_sessions ADD CONSTRAINT aibridge_sessions_pkey PRIMARY KEY (session_id, initiator_id); UniqueAibridgeTokenUsagesPkey UniqueConstraint = "aibridge_token_usages_pkey" // ALTER TABLE ONLY aibridge_token_usages ADD CONSTRAINT aibridge_token_usages_pkey PRIMARY KEY (id); UniqueAibridgeToolUsagesPkey UniqueConstraint = "aibridge_tool_usages_pkey" // ALTER TABLE ONLY aibridge_tool_usages ADD CONSTRAINT aibridge_tool_usages_pkey PRIMARY KEY (id); UniqueAibridgeUserPromptsPkey UniqueConstraint = "aibridge_user_prompts_pkey" // ALTER TABLE ONLY aibridge_user_prompts ADD CONSTRAINT aibridge_user_prompts_pkey PRIMARY KEY (id);