Keep active-turn prompts visible in budgeted timelines

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 }),

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,

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<string>();
   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<string, string>();
+  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,13 +686,34 @@ 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,
       itemStart: 95,
     });
     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,
     );

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,

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,

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, {