History

perf(client-runtime): keep turn and checkpoint refs stable while streaming (#9145)

Theo Browne committed 2 months ago

Commit 6866fd61 parent 7e460f42 files · +183 −9Browse files

Unified · Split

2 files
  1. packages/client-runtime/src/state/threadReducer.test.ts +133 −0
  2. packages/client-runtime/src/state/threadReducer.ts +50 −9
packages/client-runtime/src/state/threadReducer.test.ts+133−0
@@ -497,6 +497,139 @@ describe("applyThreadDetailEvent", () => { · 496 unmodified lines
497497 expect(result.thread.latestTurn?.completedAt).toBeNull();
498498 }
499499 });
500+
501+ it("keeps latestTurn and checkpoints references across a streaming delta", () => {
502+ const streamingThread: OrchestrationThread = {
503+ ...baseThread,
504+ session: {
505+ threadId: ThreadId.make("thread-1"),
506+ status: "running",
507+ providerName: "claude",
508+ runtimeMode: "full-access",
509+ activeTurnId: TurnId.make("turn-1"),
510+ lastError: null,
511+ updatedAt: "2026-04-01T06:59:00.000Z",
512+ },
513+ latestTurn: {
514+ turnId: TurnId.make("turn-1"),
515+ state: "running",
516+ requestedAt: "2026-04-01T06:59:00.000Z",
517+ startedAt: "2026-04-01T06:59:00.000Z",
518+ completedAt: null,
519+ assistantMessageId: MessageId.make("msg-2"),
520+ },
521+ messages: [
522+ {
523+ id: MessageId.make("msg-2"),
524+ role: "assistant",
525+ text: "Hello",
526+ turnId: TurnId.make("turn-1"),
527+ streaming: true,
528+ createdAt: "2026-04-01T06:00:00.000Z",
529+ updatedAt: "2026-04-01T06:00:00.000Z",
530+ },
531+ ],
532+ checkpoints: [
533+ {
534+ turnId: TurnId.make("turn-1"),
535+ checkpointTurnCount: 1,
536+ checkpointRef: CheckpointRef.make("ref-1"),
537+ status: "ready",
538+ files: [],
539+ assistantMessageId: MessageId.make("msg-2"),
540+ completedAt: "2026-04-01T06:00:30.000Z",
541+ },
542+ ],
543+ };
544+
545+ const result = applyThreadDetailEvent(streamingThread, {
546+ ...baseEventFields,
547+ sequence: 9,
548+ occurredAt: "2026-04-01T07:00:00.000Z",
549+ aggregateKind: "thread",
550+ aggregateId: ThreadId.make("thread-1"),
551+ type: "thread.message-sent",
552+ payload: {
553+ threadId: ThreadId.make("thread-1"),
554+ messageId: MessageId.make("msg-2"),
555+ role: "assistant",
556+ text: ", world",
557+ turnId: TurnId.make("turn-1"),
558+ streaming: true,
559+ createdAt: "2026-04-01T06:00:00.000Z",
560+ updatedAt: "2026-04-01T07:00:00.000Z",
561+ },
562+ });
563+
564+ expect(result.kind).toBe("updated");
565+ if (result.kind === "updated") {
566+ expect(result.thread.messages).not.toBe(streamingThread.messages);
567+ expect(result.thread.messages[0]?.text).toBe("Hello, world");
568+ expect(result.thread.latestTurn).toBe(streamingThread.latestTurn);
569+ expect(result.thread.checkpoints).toBe(streamingThread.checkpoints);
570+ }
571+ });
572+
573+ it("replaces latestTurn and checkpoints when the first assistant message binds the turn", () => {
574+ const unboundThread: OrchestrationThread = {
575+ ...baseThread,
576+ session: {
577+ threadId: ThreadId.make("thread-1"),
578+ status: "running",
579+ providerName: "claude",
580+ runtimeMode: "full-access",
581+ activeTurnId: TurnId.make("turn-1"),
582+ lastError: null,
583+ updatedAt: "2026-04-01T06:59:00.000Z",
584+ },
585+ latestTurn: {
586+ turnId: TurnId.make("turn-1"),
587+ state: "running",
588+ requestedAt: "2026-04-01T06:59:00.000Z",
589+ startedAt: "2026-04-01T06:59:00.000Z",
590+ completedAt: null,
591+ assistantMessageId: null,
592+ },
593+ checkpoints: [
594+ {
595+ turnId: TurnId.make("turn-1"),
596+ checkpointTurnCount: 1,
597+ checkpointRef: CheckpointRef.make("ref-1"),
598+ status: "ready",
599+ files: [],
600+ assistantMessageId: null,
601+ completedAt: "2026-04-01T06:59:30.000Z",
602+ },
603+ ],
604+ };
605+
606+ const result = applyThreadDetailEvent(unboundThread, {
607+ ...baseEventFields,
608+ sequence: 9,
609+ occurredAt: "2026-04-01T07:00:00.000Z",
610+ aggregateKind: "thread",
611+ aggregateId: ThreadId.make("thread-1"),
612+ type: "thread.message-sent",
613+ payload: {
614+ threadId: ThreadId.make("thread-1"),
615+ messageId: MessageId.make("msg-2"),
616+ role: "assistant",
617+ text: "Hello",
618+ turnId: TurnId.make("turn-1"),
619+ streaming: true,
620+ createdAt: "2026-04-01T07:00:00.000Z",
621+ updatedAt: "2026-04-01T07:00:00.000Z",
622+ },
623+ });
624+
625+ expect(result.kind).toBe("updated");
626+ if (result.kind === "updated") {
627+ expect(result.thread.latestTurn).not.toBe(unboundThread.latestTurn);
628+ expect(result.thread.latestTurn?.assistantMessageId).toBe("msg-2");
629+ expect(result.thread.checkpoints).not.toBe(unboundThread.checkpoints);
630+ expect(result.thread.checkpoints[0]?.assistantMessageId).toBe("msg-2");
631+ }
632+ });
500633 });
501634
502635 describe("thread.session-set", () => {
packages/client-runtime/src/state/threadReducer.ts+50−9
@@ -340,16 +340,18 @@ export function applyThreadDetailEvent( · 339 unmodified lines
340340 // assistant message only settles the turn once the session is no longer
341341 // running it — providers may emit several assistant messages per turn
342342 // (commentary between tool calls), and the turn must stay unsettled
343− // until the provider reports turn end.
343+ // until the provider reports turn end. Streaming deltas recompute the
344+ // same record, so the previous reference is kept when nothing changed.
344345 const turnStillRunning =
345346 event.payload.turnId !== null &&
346347 thread.session?.status === "running" &&
347348 thread.session.activeTurnId === event.payload.turnId;
348349 const settlesTurn = !event.payload.streaming && !turnStillRunning;
349− const latestTurn: OrchestrationThread["latestTurn"] =
350+ const latestTurn = reuseLatestTurn(
351+ thread.latestTurn,
350352 event.payload.role === "assistant" &&
351− event.payload.turnId !== null &&
352− (thread.latestTurn === null || thread.latestTurn.turnId === event.payload.turnId)
353+ event.payload.turnId !== null &&
354+ (thread.latestTurn === null || thread.latestTurn.turnId === event.payload.turnId)
353355 ? {
354356 turnId: event.payload.turnId,
355357 state: settlesTurn
@@ -374,9 +376,11 @@ export function applyThreadDetailEvent( · 18 unmodified lines
374376 : null,
375377 assistantMessageId: event.payload.messageId,
376378 }
377− : thread.latestTurn;
379+ : thread.latestTurn,
380+ );
378381
379− // Rebind checkpoint assistant message IDs for assistant messages.
382+ // Rebind checkpoint assistant message IDs for assistant messages. The
383+ // helper hands back the same array when the entry is already bound.
380384 const checkpoints =
381385 event.payload.role === "assistant" && event.payload.turnId !== null
382386 ? rebindCheckpointAssistantMessage(
@@ -403,7 +407,8 @@ export function applyThreadDetailEvent( · 20 unmodified lines
403407 // Leaving the "running" session status is the turn-end signal: settle a
404408 // still-running latest turn so its duration reflects the whole turn.
405409 const settledTurnState = settledTurnStateForSessionStatus(event.payload.session.status);
406− const latestTurn: OrchestrationLatestTurn | null =
410+ const latestTurn = reuseLatestTurn(
411+ thread.latestTurn,
407412 event.payload.session.status === "running" && event.payload.session.activeTurnId !== null
408413 ? {
409414 turnId: event.payload.session.activeTurnId,
@@ -433,7 +438,8 @@ export function applyThreadDetailEvent( · 23 unmodified lines
433438 // "running" is the authoritative turn end.
434439 completedAt: event.payload.session.updatedAt,
435440 }
436− : thread.latestTurn;
441+ : thread.latestTurn,
442+ );
437443
438444 return {
439445 kind: "updated",
@@ -687,11 +693,46 @@ function checkpointStatusToTurnState( · 247 unmodified lines
687693 }
688694 }
689695
696+/**
697+ * Returns `previous` when `next` matches it field for field, otherwise `next`.
698+ * Streaming cases recompute the latest turn on every delta, and keeping the
699+ * old reference lets selectors and memos keyed on `latestTurn` skip work.
700+ */
701+function reuseLatestTurn(
702+ previous: OrchestrationLatestTurn | null,
703+ next: OrchestrationLatestTurn | null,
704+): OrchestrationLatestTurn | null {
705+ if (previous === null || next === null) {
706+ return next;
707+ }
708+ return previous.turnId === next.turnId &&
709+ previous.state === next.state &&
710+ previous.requestedAt === next.requestedAt &&
711+ previous.startedAt === next.startedAt &&
712+ previous.completedAt === next.completedAt &&
713+ previous.assistantMessageId === next.assistantMessageId &&
714+ previous.sourceProposedPlan?.threadId === next.sourceProposedPlan?.threadId &&
715+ previous.sourceProposedPlan?.planId === next.sourceProposedPlan?.planId
716+ ? previous
717+ : next;
718+}
719+
720+/**
721+ * Points the checkpoint for `turnId` at `messageId`. Returns the input array
722+ * untouched when no checkpoint needs rebinding, so streaming deltas for an
723+ * already-bound message do not allocate a new `checkpoints` reference.
724+ */
690725 function rebindCheckpointAssistantMessage(
691726 checkpoints: ReadonlyArray<OrchestrationCheckpointSummary>,
692727 turnId: TurnId,
693728 messageId: MessageId,
694−): OrchestrationCheckpointSummary[] {
729+): ReadonlyArray<OrchestrationCheckpointSummary> {
730+ const needsRebind = checkpoints.some(
731+ (entry) => entry.turnId === turnId && entry.assistantMessageId !== messageId,
732+ );
733+ if (!needsRebind) {
734+ return checkpoints;
735+ }
695736 return Arr.map(checkpoints, (entry) =>
696737 entry.turnId === turnId ? { ...entry, assistantMessageId: messageId } : entry,
697738 );