diff --git a/apps/app/src/components/thread/timeline/useThreadTimelineController.merge.test.ts b/apps/app/src/components/thread/timeline/useThreadTimelineController.merge.test.ts index 8a84e12c26..4217c6b5a4 100644 --- a/apps/app/src/components/thread/timeline/useThreadTimelineController.merge.test.ts +++ b/apps/app/src/components/thread/timeline/useThreadTimelineController.merge.test.ts @@ -167,6 +167,20 @@ describe("timeline page row merging", () => { ]); }); + it("keeps a pinned active-turn prompt ahead of prepended in-turn work", () => { + const prompt = userRow({ id: "active-turn-prompt", sequence: 1 }); + const rows = prependOlderTimelineRows({ + olderRows: [prompt, commandRow({ id: "older-work", sequence: 10 })], + loadedRows: [prompt, commandRow({ id: "latest-work", sequence: 20 })], + }); + + expect(rows.map((row) => row.id)).toEqual([ + "active-turn-prompt", + "older-work", + "latest-work", + ]); + }); + it("keeps distinct byte-budget slices of one finished turn", () => { const olderCommands = [ commandRow({ id: "command-1", sequence: 10 }), diff --git a/apps/server/src/services/threads/timeline.ts b/apps/server/src/services/threads/timeline.ts index 26fd49fe5f..cb4d378c5b 100644 --- a/apps/server/src/services/threads/timeline.ts +++ b/apps/server/src/services/threads/timeline.ts @@ -19,6 +19,7 @@ import type { } from "@bb/server-contract"; import { findStoredTimelineWindowByteBudgetFloor, + findFirstStoredTurnInputAcceptedRow, findTimelineWindowBudgetFloorSequence, getStoredEventRowsByParentToolCallIdsDataBytes, getEnvironment, @@ -37,6 +38,7 @@ import { listLatestBackgroundTaskStateRowsByItemIds, listLatestGoalEventRowsByThreadIds, listLatestOpenBackgroundTaskStateRowsForThread, + listStoredClientTurnRequestRowsByKeys, listStoredTimelineWindowEventRows, listTodoSnapshotEventRowsForThread, listStoredToolCallRowsByItemIds, @@ -1021,6 +1023,8 @@ interface ResolveTimelineSegmentWindowArgs { } interface ResolvedTimelineSegmentWindow { + /** Unfinished turn whose work is cut by this sequence window. */ + activeTurnId: string | null; beforeSequence: number | undefined; byteWindowSequenceStart: number | null; /** @@ -1061,8 +1065,15 @@ function applyTimelineWindowByteBudget( if (floor.kind === "single-event-too-large") { const hasOlderRows = floor.hasOlderRows || args.window.knownHasOlderSegments === true; + const activeTurnId = + args.window.activeTurnId ?? + findUnfinishedTurnCoveringSequence(db, { + sequence: floor.sequenceStart, + threadId: args.threadId, + }); return { ...args.window, + activeTurnId, byteWindowSequenceStart: floor.sequenceStart, knownHasOlderSegments: hasOlderRows, oversizedEventPlaceholder: { @@ -1093,8 +1104,16 @@ function applyTimelineWindowByteBudget( return args.window; } + const activeTurnId = + args.window.activeTurnId ?? + findUnfinishedTurnCoveringSequence(db, { + sequence: floor.sequenceStart, + threadId: args.threadId, + }); + return { ...args.window, + activeTurnId, byteWindowSequenceStart: floor.sequenceStart, requiresWholeItemClosure: true, sequenceWindowStart: { @@ -1169,7 +1188,10 @@ function resolveTimelineWindowBounds( args: ResolveTimelineWindowBoundsArgs, ): Pick< ResolvedTimelineSegmentWindow, - "effectiveSegmentLimit" | "sequenceStart" | "sequenceWindowStart" + | "activeTurnId" + | "effectiveSegmentLimit" + | "sequenceStart" + | "sequenceWindowStart" > & { affordableAnchorCount: number } { const { anchors, budgetFloorSequence, segmentLimit, threadId } = args; const affordable = countAffordableAnchors( @@ -1194,6 +1216,7 @@ function resolveTimelineWindowBounds( }) ) { return { + activeTurnId: unfinishedTurnId, affordableAnchorCount: 0, effectiveSegmentLimit: segmentLimit, sequenceWindowStart: { @@ -1209,6 +1232,7 @@ function resolveTimelineWindowBounds( // empty thread rather than a slow one. const segmentCount = Math.max(1, affordable); return { + activeTurnId: null, affordableAnchorCount: segmentCount, effectiveSegmentLimit: segmentCount, sequenceWindowStart: null, @@ -1235,6 +1259,7 @@ function resolveTimelineSegmentWindow( ): ResolvedTimelineSegmentWindow { const { eventBudget, page, threadId } = args; const noAnchors: ResolvedTimelineSegmentWindow = { + activeTurnId: null, beforeSequence: undefined, byteWindowSequenceStart: null, requiresWholeItemClosure: false, @@ -1314,6 +1339,7 @@ function resolveTimelineSegmentWindow( threadId, }); return { + activeTurnId: bounds.activeTurnId, // Every cursor names the first sequence the page that issued it covered, // so this page ends exactly there. Reading up to the *next anchor* past // the cursor instead — and trimming that segment off after projecting it @@ -1353,6 +1379,7 @@ function resolveTimelineSegmentWindow( threadId, }); return { + activeTurnId: bounds.activeTurnId, beforeSequence: undefined, byteWindowSequenceStart: null, requiresWholeItemClosure: bounds.sequenceWindowStart !== null, @@ -1367,6 +1394,33 @@ function resolveTimelineSegmentWindow( }; } +function ensureTimelineWindowActiveTurnInputRows( + db: DbConnection, + args: TimelineWindowRowsArgs & { activeTurnId: string | null }, +): StoredEventRow[] { + if (args.activeTurnId === null) { + return [...args.rows]; + } + + const acceptedInputRow = findFirstStoredTurnInputAcceptedRow(db, { + threadId: args.threadId, + turnId: args.activeTurnId, + }); + if (!acceptedInputRow) { + return [...args.rows]; + } + + const clientRequestId = parseAcceptedInputClientRequestId(acceptedInputRow); + const requestRows = listStoredClientTurnRequestRowsByKeys(db, { + keys: [{ requestId: clientRequestId, threadId: args.threadId }], + }); + return mergeStoredEventRowsById([ + ...requestRows, + acceptedInputRow, + ...args.rows, + ]); +} + function selectStandardTimelineEventRows( db: DbConnection, thread: Thread, @@ -1402,15 +1456,23 @@ function selectStandardTimelineEventRows( threadId: thread.id, }; const windowRows = listStoredTimelineWindowEventRows(db, windowArgs); + const windowRowsWithActiveTurnInput = ensureTimelineWindowActiveTurnInputRows( + db, + { + activeTurnId: window.activeTurnId, + rows: windowRows, + threadId: thread.id, + }, + ); const wholeItemWindowRows = window.requiresWholeItemClosure ? ensureSequenceWindowWholeItemRows(db, { beforeSequence, maxInlineOutputChars, - rows: windowRows, + rows: windowRowsWithActiveTurnInput, sequenceStart, threadId: thread.id, }) - : windowRows; + : windowRowsWithActiveTurnInput; const selectedRowsWithTurnStarts = ensureTimelineWindowTurnStartedRows(db, { threadId: thread.id, rows: wholeItemWindowRows, diff --git a/apps/server/test/services/threads/timeline-in-turn-window.test.ts b/apps/server/test/services/threads/timeline-in-turn-window.test.ts index 9ce0850579..7295b9c4e6 100644 --- a/apps/server/test/services/threads/timeline-in-turn-window.test.ts +++ b/apps/server/test/services/threads/timeline-in-turn-window.test.ts @@ -473,7 +473,7 @@ function walkAllPages( thread: Thread, eventBudget: number, ): WalkResult { - const rowsByPage: string[][] = []; + const rowsByPage: TimelineRow[][] = []; const seenCursors = new Set(); let cursor: TimelinePaginationCursor | null = null; let maxEventRowCount = 0; @@ -483,7 +483,7 @@ function walkAllPages( const { profile, response } = buildPage(db, thread, eventBudget, cursor); pages += 1; maxEventRowCount = Math.max(maxEventRowCount, profile.eventRowCount); - rowsByPage.push(response.rows.map((row) => JSON.stringify(row))); + rowsByPage.push(response.rows); if (!response.timelinePage.hasOlderRows) { break; } @@ -499,11 +499,25 @@ function walkAllPages( expect(pages).toBeLessThan(100); } - return { maxEventRowCount, pages, rows: rowsByPage.reverse().flat() }; + const serializedRowsById = new Map(); + for (const row of rowsByPage.reverse().flat()) { + const serialized = JSON.stringify(row); + const existing = serializedRowsById.get(row.id); + if (existing !== undefined) { + expect(existing).toBe(serialized); + continue; + } + serializedRowsById.set(row.id, serialized); + } + return { + maxEventRowCount, + pages, + rows: [...serializedRowsById.values()], + }; } describe("in-turn timeline windows", () => { - it("bounds a running turn that is larger than the whole budget", () => { + it("bounds a running turn while keeping its initiating user message", () => { const { db, thread } = setup(); // One 300-item turn still running: no user message inside it to cut on. seedTurns(db, thread, { completeLastTurn: false, itemsPerTurn: [300] }); @@ -516,6 +530,14 @@ describe("in-turn timeline windows", () => { const budgeted = buildPage(db, thread, 100, null); expect(budgeted.profile.eventRowCount).toBeLessThanOrEqual(120); + expect(budgeted.response.rows).toContainEqual( + expect.objectContaining({ + kind: "conversation", + role: "user", + text: "User message 1", + turnRequest: expect.objectContaining({ status: "accepted" }), + }), + ); expect(budgeted.response.rows.length).toBeLessThan( unbudgeted.response.rows.length, ); @@ -664,6 +686,16 @@ describe("in-turn timeline windows", () => { itemStart: 75, }); const second = buildPage(db, thread, LARGE_BUDGET, null).response; + const secondOlderCursor = second.timelinePage.olderCursor; + if (secondOlderCursor === null) { + throw new Error("expected the byte-budgeted turn to have an older page"); + } + const secondOlder = buildPage( + db, + thread, + LARGE_BUDGET, + secondOlderCursor, + ).response; appendCommandItems(db, thread, { commandChars: 25_000, count: 10, @@ -671,6 +703,17 @@ describe("in-turn timeline windows", () => { }); const third = buildPage(db, thread, LARGE_BUDGET, null).response; + for (const response of [first, second, secondOlder, third]) { + expect(response.rows).toContainEqual( + expect.objectContaining({ + kind: "conversation", + role: "user", + text: "User message 1", + turnRequest: expect.objectContaining({ status: "accepted" }), + }), + ); + } + expect(first.timelinePage.olderCursor?.anchorSeq).not.toBe( second.timelinePage.olderCursor?.anchorSeq, ); diff --git a/packages/db/src/data/events.ts b/packages/db/src/data/events.ts index 2195baafc3..6b70881505 100644 --- a/packages/db/src/data/events.ts +++ b/packages/db/src/data/events.ts @@ -888,6 +888,11 @@ export interface ListOpenTurnInputAcceptedRowsByThreadIdsArgs { threadIds: readonly string[]; } +export interface FindFirstStoredTurnInputAcceptedRowArgs { + threadId: string; + turnId: string; +} + export interface ThreadClientTurnRequestKey { requestId: ClientTurnRequestId; threadId: string; @@ -1229,6 +1234,27 @@ export function listOpenTurnInputAcceptedRowsByThreadIds( ); } +export function findFirstStoredTurnInputAcceptedRow( + db: DbQueryConnection, + args: FindFirstStoredTurnInputAcceptedRowArgs, +): StoredEventRow | null { + return ( + db + .select(storedEventRowFields) + .from(events) + .where( + and( + eq(events.threadId, args.threadId), + eq(events.turnId, args.turnId), + eq(events.type, "turn/input/accepted"), + ), + ) + .orderBy(events.sequence) + .limit(1) + .get() ?? null + ); +} + export function listStoredClientTurnRequestRowsByKeys( db: DbQueryConnection, args: ListStoredClientTurnRequestRowsByKeysArgs, diff --git a/packages/db/src/data/index.ts b/packages/db/src/data/index.ts index 246bdad43a..f9f96efa14 100644 --- a/packages/db/src/data/index.ts +++ b/packages/db/src/data/index.ts @@ -328,6 +328,7 @@ export { deleteThreadEventSuffixInTransaction, getHighWaterMarks, findStoredClientTurnRequestSequenceByRequestId, + findFirstStoredTurnInputAcceptedRow, findStoredEventRow, getActiveStoredTurnId, hasRootStoredTurnStarted, @@ -402,6 +403,7 @@ export type { DeleteThreadEventSuffixArgs, DeleteThreadEventSuffixResult, FindStoredClientTurnRequestSequenceByRequestIdArgs, + FindFirstStoredTurnInputAcceptedRowArgs, GetStoredTurnRequestEventForTurnArgs, GetRootStoredTurnStartedSequenceArgs, GetLatestThreadInterruptedReasonArgs, diff --git a/packages/db/test/data/events.test.ts b/packages/db/test/data/events.test.ts index 1f059fd6b1..ce3f441d12 100644 --- a/packages/db/test/data/events.test.ts +++ b/packages/db/test/data/events.test.ts @@ -18,6 +18,7 @@ import { appendStoredThreadEventInTransaction, appendStoredThreadEventsInTransaction, findStoredEventRow, + findFirstStoredTurnInputAcceptedRow, findStoredTimelineWindowByteBudgetFloor, getActiveStoredTurnId, getHighWaterMarks, @@ -1433,6 +1434,47 @@ describe("events", () => { expect(rowsByThreadId.get(otherThread.id)?.sequence).toBe(1); }); + it("finds the first accepted input for one exact turn", () => { + const { db, thread } = setup(); + + insertEvents(db, noopNotifier, [ + { + threadId: thread.id, + sequence: 1, + type: "turn/input/accepted", + ...createTurnEventFields({ turnId: "turn-target" }), + data: JSON.stringify({ clientRequestId: "creq_23456789aa" }), + }, + { + threadId: thread.id, + sequence: 2, + type: "turn/input/accepted", + ...createTurnEventFields({ turnId: "turn-other" }), + data: JSON.stringify({ clientRequestId: "creq_23456789ab" }), + }, + { + threadId: thread.id, + sequence: 3, + type: "turn/input/accepted", + ...createTurnEventFields({ turnId: "turn-target" }), + data: JSON.stringify({ clientRequestId: "creq_23456789ac" }), + }, + ]); + + expect( + findFirstStoredTurnInputAcceptedRow(db, { + threadId: thread.id, + turnId: "turn-target", + })?.sequence, + ).toBe(1); + expect( + findFirstStoredTurnInputAcceptedRow(db, { + threadId: thread.id, + turnId: "turn-missing", + }), + ).toBeNull(); + }); + it("lists client turn request rows by thread/request keys", () => { const { db, project, thread } = setup(); const otherThread = createThread(db, noopNotifier, {