perf(agent-sessions): rank the page on the index, then aggregate only that page

apps/api/src/routes/internal/ai-sessions.http.test.ts

@@ -358,6 +358,179 @@ describe("POST /internal/ai-sessions/spans", () => {
 	})
 })
 
+/**
+ * The list is two reads, and the handler is what holds them together: it derives
+ * the fan-out's window from the page, and it re-imposes the page's order on an
+ * aggregation that cannot know it. Both are invisible in either query alone.
+ */
+describe("POST /internal/ai-sessions/list", () => {
+	const LIST_BODY = { ...WINDOW, limit: 3 }
+
+	/** A stage-one row: a session id and the extent of its agent spans. */
+	const pageRow = (sessionId: string, agentStart: string, agentEnd: string) => ({
+		sessionId,
+		agentStart,
+		agentEnd,
+	})
+
+	/** A stage-two row, in the wire shape the aggregation's SELECT decodes. */
+	const listRow = (sessionId: string, startTime: string) => ({
+		sessionId,
+		vendorId: "eve",
+		vendorVersion: "1",
+		traceCount: "1",
+		spanCount: "12",
+		errorSpanCount: "0",
+		serviceNames: ["agent-runner"],
+		startTime,
+		endTime: "2026-08-19 10:45:00.000000000",
+		durationMs: "1000",
+	})
+
+	// Deliberately not in start order: the page ranks on the first AGENT span and
+	// the aggregation orders by the first span of any kind, so the handler must
+	// not be able to reconstruct one from the other.
+	const PAGE = [
+		pageRow("wrun_beta", "2026-08-19 10:20:00.000000000", "2026-08-19 10:30:00.000000000"),
+		pageRow("wrun_alpha", "2026-08-19 10:05:00.000000000", "2026-08-19 10:40:00.000000000"),
+		pageRow(`trace:${TRACE_ID}`, "2026-08-19 09:50:00.000000000", "2026-08-19 10:00:00.000000000"),
+	]
+
+	it("bounds the aggregation by the page's own agent spans, not the caller's window", async () => {
+		let pageSql: string | undefined
+		let listSql: string | undefined
+		const harness = makeHarness({
+			compiledQuery: (_tenant, compiled, options) => {
+				if (options?.context === "aiSessionsPage") {
+					pageSql = compiledQueryOf(compiled).sql
+					return compiledQueryOf(compiled).decodeRows(PAGE).pipe(Effect.orDie)
+				}
+				listSql = compiledQueryOf(compiled).sql
+				return compiledQueryOf(compiled)
+					.decodeRows(PAGE.map((row) => listRow(row.sessionId, row.agentStart)))
+					.pipe(Effect.orDie)
+			},
+		})
+
+		try {
+			const response = await harness.post("/internal/ai-sessions/list", LIST_BODY)
+			expect(response.status).toBe(200)
+			// Stage one is the only read that sees the caller's window.
+			expect(pageSql).toContain("FROM ai_trace_index")
+			expect(pageSql).not.toContain("trace_detail_spans")
+			expect(pageSql).toContain(`Timestamp <= '${WINDOW.endTime}'`)
+			expect(pageSql).toContain("LIMIT 3")
+			// Stage two reads the fan-out table over the page's extent — the min
+			// agentStart and the max agentEnd of the rows stage one returned, padded.
+			// The caller's window would be 30 days of partitions on the page the UI
+			// actually offers.
+			expect(listSql).toContain("FROM trace_detail_spans")
+			expect(listSql).toContain("Timestamp >= '2026-08-19 09:50:00.000000000' - INTERVAL 3600 SECOND")
+			expect(listSql).toContain("Timestamp <= '2026-08-19 10:40:00.000000000' + INTERVAL 3600 SECOND")
+			// The caller's range reaches NO level of stage two — not the fan-out and
+			// not either of its two `ai_trace_index` reads. The handler hands it
+			// `orgId` and the page's two bounds, and nothing else.
+			expect(listSql).not.toContain(WINDOW.startTime)
+			expect(listSql).not.toContain(WINDOW.endTime)
+			// Three levels take that lower bound — the fan-out padded, and each of
+			// the two `ai_trace_index` reads exactly.
+			expect(listSql?.split("Timestamp >= '2026-08-19 09:50:00.000000000'").length).toBe(4)
+			expect(listSql).not.toContain("__PARAM_")
+		} finally {
+			await harness.dispose()
+		}
+	})
+
+	it("seeks stage two by exactly the session ids stage one ranked", async () => {
+		let listSql: string | undefined
+		const harness = makeHarness({
+			compiledQuery: (_tenant, compiled, options) => {
+				if (options?.context === "aiSessionsPage") {
+					return compiledQueryOf(compiled).decodeRows(PAGE).pipe(Effect.orDie)
+				}
+				listSql = compiledQueryOf(compiled).sql
+				return compiledQueryOf(compiled)
+					.decodeRows(PAGE.map((row) => listRow(row.sessionId, row.agentStart)))
+					.pipe(Effect.orDie)
+			},
+		})
+
+		try {
+			await harness.post("/internal/ai-sessions/list", LIST_BODY)
+			for (const row of PAGE) {
+				expect(listSql).toContain(`'${row.sessionId}'`)
+			}
+		} finally {
+			await harness.dispose()
+		}
+	})
+
+	it("answers in the page's order, dropping a session the aggregation lost", async () => {
+		const harness = makeHarness({
+			compiledQuery: (_tenant, compiled, options) =>
+				options?.context === "aiSessionsPage"
+					? compiledQueryOf(compiled).decodeRows(PAGE).pipe(Effect.orDie)
+					: compiledQueryOf(compiled)
+							// Reversed, and one short: the aggregation's own ORDER BY is
+							// meaningless to the client, and a trace whose spans fell outside
+							// the padded window returns nothing at all.
+							.decodeRows([
+								listRow(`trace:${TRACE_ID}`, "2026-08-19 09:49:00.000000000"),
+								listRow("wrun_beta", "2026-08-19 10:19:00.000000000"),
+							])
+							.pipe(Effect.orDie),
+		})
+
+		try {
+			const response = await harness.post("/internal/ai-sessions/list", LIST_BODY)
+			expect(response.status).toBe(200)
+			// The page's order is the order that was paged; re-sorting here would
+			// let a row jump between pages on a scroll. A session with no row is
+			// dropped rather than shown with blank counts.
+			expect((response.body.data as ReadonlyArray<{ sessionId: string }>).map((r) => r.sessionId)).toEqual([
+				"wrun_beta",
+				`trace:${TRACE_ID}`,
+			])
+			// Three ranked, two returned. `ranked` is what the client pages on: on
+			// `data.length` this short page reads as the end of the list, and the
+			// next offset would be one too low and re-show a session. The gap is
+			// real — the two MVs are written one after the other from the same
+			// insert, so the newest session can be ranked before its spans land.
+			expect(response.body.ranked).toBe(PAGE.length)
+			expect(response.body.ranked).toBe(3)
+		} finally {
+			await harness.dispose()
+		}
+	})
+
+	it("answers an empty page without touching trace_detail_spans", async () => {
+		const contexts: Array<string | undefined> = []
+		const harness = makeHarness({
+			compiledQuery: (_tenant, compiled, options) => {
+				contexts.push(options?.context)
+				return compiledQueryOf(compiled).decodeRows([]).pipe(Effect.orDie)
+			},
+		})
+
+		try {
+			const response = await harness.post("/internal/ai-sessions/list", LIST_BODY)
+			expect(response.status).toBe(200)
+			// `ranked` is omitted entirely, not sent as 0: the handler short-circuits
+			// with `new ListAiSessionsResponse({ data: [] })`, and the field is an
+			// `optionalKey`. The client's `?? data.length` fallback reads it as 0
+			// either way, which is what ends the scroll.
+			expect(response.body).toEqual({ data: [] })
+			expect("ranked" in response.body).toBe(false)
+			// One read, not two. With no ids to seek by, the fan-out's `IN ()` is a
+			// builder defect, and the shape it would have compiled to reads the whole
+			// padded window for nothing.
+			expect(contexts).toEqual(["aiSessionsPage"])
+		} finally {
+			await harness.dispose()
+		}
+	})
+})
+
 describe("POST /internal/ai-sessions/facets", () => {
 	// `pick("vendor")` in the handler and `facet("vendor", …)` in the query are
 	// two independent string literals in two packages. If either drifts both

apps/api/src/routes/internal/ai-sessions.http.ts

@@ -34,26 +34,65 @@ export const HttpAiSessionsInternalLive = HttpApiBuilder.group(
 					Effect.gen(function* () {
 						const tenant = yield* CurrentTenant.Context
 						yield* Effect.annotateCurrentSpan({ orgId: tenant.orgId })
-						const compiled = CH.compile(
-							Integrations.aiSessionListQuery({
-								limit: payload.limit,
-								offset: payload.offset,
-								vendorIds: payload.vendorIds,
-								serviceNames: payload.serviceNames,
-							}),
-							{
-								orgId: tenant.orgId,
-								startTime: payload.startTime,
-								endTime: payload.endTime,
-							},
+						const filters = { vendorIds: payload.vendorIds, serviceNames: payload.serviceNames }
+						const window = { orgId: tenant.orgId, startTime: payload.startTime, endTime: payload.endTime }
+						// Two reads, not one: the page is ranked on `ai_trace_index` over the
+						// caller's whole window, and only then is that page aggregated over
+						// `trace_detail_spans` — inside the hours its own agent spans cover,
+						// never the caller's window. See `aiSessionListQuery` for what the
+						// single-read shape cost.
+						const page = yield* warehouse.compiledQuery(
+							tenant,
+							CH.compile(
+								Integrations.aiSessionPageQuery({
+									...filters,
+									limit: payload.limit,
+									offset: payload.offset,
+								}),
+								window,
+							),
+							{ profile: "list", context: "aiSessionsPage" },
 						)
+						if (page.length === 0) {
+							return new ListAiSessionsResponse({ data: [] })
+						}
+						// Fixed-width warehouse literals, so they sort as the instants do.
+						const fanOutStart = page.map((row) => row.agentStart).reduce((a, b) => (a < b ? a : b))
+						const fanOutEnd = page.map((row) => row.agentEnd).reduce((a, b) => (a < b ? b : a))
 						// The row schema already coerces the UInt64 aggregates and decodes
 						// exactly the response's fields, so rows pass through unmapped.
-						const rows = yield* warehouse.compiledQuery(tenant, compiled, {
-							profile: "list",
-							context: "listAiSessions",
+						const rows = yield* warehouse.compiledQuery(
+							tenant,
+							CH.compile(
+								Integrations.aiSessionListQuery({
+									...filters,
+									sessionIds: page.map((row) => row.sessionId),
+								}),
+								{ orgId: tenant.orgId, fanOutStart, fanOutEnd },
+							),
+							{ profile: "list", context: "listAiSessions" },
+						)
+						// The page's order is the order that was paged, so it is the order
+						// shown: the aggregation sorts by the true first span, which leads
+						// the first agent span the page ranked on by under a second.
+						//
+						// A session the aggregation did not return is dropped rather than
+						// shown blank, and `ranked` tells the client the page was still a
+						// full one. It happens: `ai_trace_index` and `trace_detail_spans`
+						// are two materialized views written one after the other from the
+						// same `traces` insert, so the newest session — ranked first — can
+						// have index rows a moment before it has span rows, and at the far
+						// end of retention the two tables' TTL merges run on their own
+						// clocks. The two counts on the span are how often.
+						yield* Effect.annotateCurrentSpan({
+							"maple.ai.page_size": page.length,
+							"maple.ai.aggregated": rows.length,
+						})
+						const byId = new Map(rows.map((row) => [row.sessionId, row]))
+						return new ListAiSessionsResponse({
+							data: page.flatMap((row) => byId.get(row.sessionId) ?? []),
+							ranked: page.length,
 						})
-						return new ListAiSessionsResponse({ data: rows })
 					}),
 				)
 				.handle("facets", ({ payload }) =>