fix(server): track assistant segments per item

Keep segment lifecycle state for every provider item so interleaved deltas, approval pauses, and item completions cannot affect sibling assistant messages.

t3-code[bot] committed 2 months ago

Commit 2917cba1 parent 6bdc9322 files · +280 −119Browse files

Unified · Split

2 files

apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts +174 −9
apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts +106 −110

apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts+174−9

@@ -2220,10 +2220,7 @@ describe("ProviderRuntimeIngestion", () => { · 2219 unmodified lines
2220	threadId: asThreadId("thread-1"),	2220	threadId: asThreadId("thread-1"),
2221	turnId,	2221	turnId,
2222	});	2222	});
2223	await waitForThread(	2223	await waitForThread(harness.readModel, (thread) => thread.session?.activeTurnId === turnId);
2224	harness.readModel,
2225	(thread) => thread.session?.activeTurnId === turnId,
2226	);
2227		2224
2228	for (const [itemId, delta] of [	2225	for (const [itemId, delta] of [
2229	["item-buffered-first", "first item"],	2226	["item-buffered-first", "first item"],
@@ -2251,12 +2248,11 @@ describe("ProviderRuntimeIngestion", () => { · 21 unmodified lines
2251	payload: { state: "completed" },	2248	payload: { state: "completed" },
2252	});	2249	});
2253		2250
2254	const thread = await waitForThread(	2251	const thread = await waitForThread(harness.readModel, (entry) =>
2255	harness.readModel,	2252	entry.messages.some(
2256	(entry) =>	2253	(message: ProviderRuntimeTestMessage) =>
2257	entry.messages.some((message: ProviderRuntimeTestMessage) =>
2258	message.id === "assistant:item-buffered-second" && !message.streaming,	2254	message.id === "assistant:item-buffered-second" && !message.streaming,
2259	),	2255	),
2260	);	2256	);
2261	expect(thread.messages).toEqual(	2257	expect(thread.messages).toEqual(
2262	expect.arrayContaining([	2258	expect.arrayContaining([
@@ -2266,6 +2262,175 @@ describe("ProviderRuntimeIngestion", () => { · 3 unmodified lines
2266	);	2262	);
2267	});	2263	});
2268		2264
		2265	it("keeps interleaved buffered assistant deltas with their provider items", async () => {
		2266	const harness = await createHarness();
		2267	const base = {
		2268	provider: ProviderDriverKind.make("codex"),
		2269	createdAt: "2026-01-01T00:00:00.000Z",
		2270	threadId: asThreadId("thread-1"),
		2271	turnId: asTurnId("turn-buffered-interleaved-items"),
		2272	};
		2273
		2274	await harness.emitAndDrain([
		2275	{ ...base, type: "turn.started", eventId: asEventId("interleaved-start") },
		2276	{
		2277	...base,
		2278	type: "content.delta",
		2279	eventId: asEventId("interleaved-a-1"),
		2280	itemId: asItemId("item-interleaved-a"),
		2281	payload: { streamKind: "assistant_text", delta: "first " },
		2282	},
		2283	{
		2284	...base,
		2285	type: "content.delta",
		2286	eventId: asEventId("interleaved-b"),
		2287	itemId: asItemId("item-interleaved-b"),
		2288	payload: { streamKind: "assistant_text", delta: "second" },
		2289	},
		2290	{
		2291	...base,
		2292	type: "content.delta",
		2293	eventId: asEventId("interleaved-a-2"),
		2294	itemId: asItemId("item-interleaved-a"),
		2295	payload: { streamKind: "assistant_text", delta: "third" },
		2296	},
		2297	{
		2298	...base,
		2299	type: "turn.completed",
		2300	eventId: asEventId("interleaved-completed"),
		2301	payload: { state: "completed" },
		2302	},
		2303	]);
		2304
		2305	const messages = (await harness.readModel()).threads[0]?.messages ?? [];
		2306	expect(messages).toEqual(
		2307	expect.arrayContaining([
		2308	expect.objectContaining({ id: "assistant:item-interleaved-a", text: "first third" }),
		2309	expect.objectContaining({ id: "assistant:item-interleaved-b", text: "second" }),
		2310	]),
		2311	);
		2312	expect(
		2313	messages.some((message) => message.id === "assistant:item-interleaved-a:segment:1"),
		2314	).toBe(false);
		2315	});
		2316
		2317	it("finalizes every open assistant item at an approval boundary", async () => {
		2318	const harness = await createHarness();
		2319	const base = {
		2320	provider: ProviderDriverKind.make("codex"),
		2321	createdAt: "2026-01-01T00:00:00.000Z",
		2322	threadId: asThreadId("thread-1"),
		2323	turnId: asTurnId("turn-buffered-multiple-items-approval"),
		2324	};
		2325
		2326	await harness.emitAndDrain([
		2327	{ ...base, type: "turn.started", eventId: asEventId("multiple-approval-start") },
		2328	{
		2329	...base,
		2330	type: "content.delta",
		2331	eventId: asEventId("multiple-approval-a"),
		2332	itemId: asItemId("item-multiple-approval-a"),
		2333	payload: { streamKind: "assistant_text", delta: "first item" },
		2334	},
		2335	{
		2336	...base,
		2337	type: "content.delta",
		2338	eventId: asEventId("multiple-approval-b"),
		2339	itemId: asItemId("item-multiple-approval-b"),
		2340	payload: { streamKind: "assistant_text", delta: "second item" },
		2341	},
		2342	{
		2343	...base,
		2344	type: "request.opened",
		2345	eventId: asEventId("multiple-approval-opened"),
		2346	requestId: ApprovalRequestId.make("req-multiple-approval"),
		2347	payload: { requestType: "command_execution_approval", detail: "pwd" },
		2348	},
		2349	]);
		2350
		2351	const messages = (await harness.readModel()).threads[0]?.messages ?? [];
		2352	expect(messages).toEqual(
		2353	expect.arrayContaining([
		2354	expect.objectContaining({
		2355	id: "assistant:item-multiple-approval-a",
		2356	text: "first item",
		2357	streaming: false,
		2358	}),
		2359	expect.objectContaining({
		2360	id: "assistant:item-multiple-approval-b",
		2361	text: "second item",
		2362	streaming: false,
		2363	}),
		2364	]),
		2365	);
		2366	});
		2367
		2368	it("completes only the matching assistant item when items overlap", async () => {
		2369	const harness = await createHarness({ serverSettings: { enableLegacyTokenStreaming: true } });
		2370	const base = {
		2371	provider: ProviderDriverKind.make("codex"),
		2372	createdAt: "2026-01-01T00:00:00.000Z",
		2373	threadId: asThreadId("thread-1"),
		2374	turnId: asTurnId("turn-streaming-overlapping-items"),
		2375	};
		2376
		2377	await harness.emitAndDrain([
		2378	{ ...base, type: "turn.started", eventId: asEventId("overlap-start") },
		2379	{
		2380	...base,
		2381	type: "content.delta",
		2382	eventId: asEventId("overlap-a"),
		2383	itemId: asItemId("item-overlap-a"),
		2384	payload: { streamKind: "assistant_text", delta: "first" },
		2385	},
		2386	{
		2387	...base,
		2388	type: "content.delta",
		2389	eventId: asEventId("overlap-b-1"),
		2390	itemId: asItemId("item-overlap-b"),
		2391	payload: { streamKind: "assistant_text", delta: "second" },
		2392	},
		2393	{
		2394	...base,
		2395	type: "item.completed",
		2396	eventId: asEventId("overlap-a-completed"),
		2397	itemId: asItemId("item-overlap-a"),
		2398	payload: { itemType: "assistant_message", status: "completed" },
		2399	},
		2400	]);
		2401
		2402	let messages = (await harness.readModel()).threads[0]?.messages ?? [];
		2403	expect(messages.find((message) => message.id === "assistant:item-overlap-a")?.streaming).toBe(
		2404	false,
		2405	);
		2406	expect(messages.find((message) => message.id === "assistant:item-overlap-b")?.streaming).toBe(
		2407	true,
		2408	);
		2409
		2410	await harness.emitAndDrain([
		2411	{
		2412	...base,
		2413	type: "content.delta",
		2414	eventId: asEventId("overlap-b-2"),
		2415	itemId: asItemId("item-overlap-b"),
		2416	payload: { streamKind: "assistant_text", delta: " item" },
		2417	},
		2418	{
		2419	...base,
		2420	type: "item.completed",
		2421	eventId: asEventId("overlap-b-completed"),
		2422	itemId: asItemId("item-overlap-b"),
		2423	payload: { itemType: "assistant_message", status: "completed" },
		2424	},
		2425	]);
		2426
		2427	messages = (await harness.readModel()).threads[0]?.messages ?? [];
		2428	expect(messages.find((message) => message.id === "assistant:item-overlap-b")).toMatchObject({
		2429	text: "second item",
		2430	streaming: false,
		2431	});
		2432	});
		2433
2269	it("flushes and completes buffered assistant text when an approval request opens", async () => {	2434	it("flushes and completes buffered assistant text when an approval request opens", async () => {
2270	const harness = await createHarness();	2435	const harness = await createHarness();
2271	const now = "2026-01-01T00:00:00.000Z";	2436	const now = "2026-01-01T00:00:00.000Z";

apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts+106−110

 }
 interface AssistantSegmentState {
-  baseKey: string;
   nextSegmentIndex: number;
   activeMessageId: MessageId | null;
 }
   return left === right;
 }
-function hasAssistantMessageForTurn(
-  messages: ReadonlyArray<OrchestrationMessage>,
-  turnId: TurnId,
-  options?: { readonly streamingOnly?: boolean },
-): boolean {
-  for (let index = 0; index < messages.length; index += 1) {
-    const message = messages[index];
-    if (!message) {
-      continue;
-    }
-    if (message.role !== "assistant" || message.turnId !== turnId) {
-      continue;
-    }
-    if (options?.streamingOnly === true && !message.streaming) {
-      continue;
-    }
-    return true;
-  }
-  return false;
-}
 function findMessageById(
   messages: ReadonlyArray<OrchestrationMessage>,
   messageId: MessageId,
     lookup: () => Effect.succeed(""),
   });
-  const assistantSegmentStateByTurnKey = yield* Cache.make<string, AssistantSegmentState>({
+  const assistantSegmentStateByTurnKey = yield* Cache.make<
+    string,
+    Map<string, AssistantSegmentState>
+  >({
     capacity: TURN_MESSAGE_IDS_BY_TURN_CACHE_CAPACITY,
     timeToLive: TURN_MESSAGE_IDS_BY_TURN_TTL,
-    lookup: () =>
+    lookup: () => Effect.succeed(new Map()),
-      Effect.die(
-        new Error("assistant segment state should be read through getOption before initialization"),
-      ),
   });
   const bufferedProposedPlanById = yield* Cache.make<string, { text: string; createdAt: string }>({
   const clearAssistantMessageIdsForTurn = (threadId: ThreadId, turnId: TurnId) =>
     Cache.invalidate(turnMessageIdsByTurnKey, providerTurnKey(threadId, turnId));
-  const getAssistantSegmentStateForTurn = (threadId: ThreadId, turnId: TurnId) =>
+  const getAssistantSegmentStatesForTurn = (threadId: ThreadId, turnId: TurnId) =>
     Cache.getOption(assistantSegmentStateByTurnKey, providerTurnKey(threadId, turnId));
+  const getAssistantSegmentStateForTurn = (threadId: ThreadId, turnId: TurnId, baseKey: string) =>
+    getAssistantSegmentStatesForTurn(threadId, turnId).pipe(
+      Effect.map((states) =>
+        Option.flatMap(states, (entries) => Option.fromNullishOr(entries.get(baseKey))),
+      ),
+    );
   const setAssistantSegmentStateForTurn = (
     threadId: ThreadId,
     turnId: TurnId,
+    baseKey: string,
     state: AssistantSegmentState,
-  ) => Cache.set(assistantSegmentStateByTurnKey, providerTurnKey(threadId, turnId), state);
+  ) =>
+    getAssistantSegmentStatesForTurn(threadId, turnId).pipe(
+      Effect.flatMap((existingStates) => {
+        const nextStates = new Map(Option.getOrElse(existingStates, () => new Map()));
+        nextStates.set(baseKey, state);
+        return Cache.set(
+          assistantSegmentStateByTurnKey,
+          providerTurnKey(threadId, turnId),
+          nextStates,
+        );
+      }),
+    );
   const clearAssistantSegmentStateForTurn = (threadId: ThreadId, turnId: TurnId) =>
     Cache.invalidate(assistantSegmentStateByTurnKey, providerTurnKey(threadId, turnId));
-  const getActiveAssistantMessageIdForTurn = (threadId: ThreadId, turnId: TurnId) =>
-    getAssistantSegmentStateForTurn(threadId, turnId).pipe(
-      Effect.map((state) =>
-        Option.flatMap(state, (entry) =>
-          entry.activeMessageId ? Option.some(entry.activeMessageId) : Option.none(),
-        ),
-      ),
-    );
   const startAssistantSegmentForTurn = (input: {
     threadId: ThreadId;
     turnId: TurnId;
     baseKey: string;
   }) =>
-    getAssistantSegmentStateForTurn(input.threadId, input.turnId).pipe(
+    getAssistantSegmentStateForTurn(input.threadId, input.turnId, input.baseKey).pipe(
       Effect.flatMap((existingState) =>
         Effect.gen(function* () {
-          const nextState = Option.match(existingState, {
+          const segmentIndex = Option.match(existingState, {
-            onNone: () => ({
+            onNone: () => 0,
-              baseKey: input.baseKey,
+            onSome: (state) => state.nextSegmentIndex,
-              nextSegmentIndex: 1,
-              activeMessageId: assistantSegmentMessageId(input.baseKey, 0),
-            }),
-            onSome: (state) => {
-              const segmentIndex = state.baseKey === input.baseKey ? state.nextSegmentIndex : 0;
-              const messageId = assistantSegmentMessageId(input.baseKey, segmentIndex);
-              return {
-                baseKey: input.baseKey,
-                nextSegmentIndex: state.baseKey === input.baseKey ? state.nextSegmentIndex + 1 : 1,
-                activeMessageId: messageId,
-              } satisfies AssistantSegmentState;
-            },
           });
-          yield* setAssistantSegmentStateForTurn(input.threadId, input.turnId, nextState);
+          const nextState: AssistantSegmentState = {
+            nextSegmentIndex: segmentIndex + 1,
+            activeMessageId: assistantSegmentMessageId(input.baseKey, segmentIndex),
+          };
+          yield* setAssistantSegmentStateForTurn(
+            input.threadId,
+            input.turnId,
+            input.baseKey,
+            nextState,
+          );
           return nextState.activeMessageId!;
         }),
       ),
       }
       const baseKey = assistantSegmentBaseKeyFromEvent(input.event);
-      const segmentState = yield* getAssistantSegmentStateForTurn(input.threadId, input.turnId);
+      const segmentState = yield* getAssistantSegmentStateForTurn(
-      if (
+        input.threadId,
-        Option.isSome(segmentState) &&
+        input.turnId,
-        segmentState.value.baseKey === baseKey &&
+        baseKey,
-        segmentState.value.activeMessageId !== null
+      );
-      ) {
+      if (Option.isSome(segmentState) && segmentState.value.activeMessageId !== null) {
         return segmentState.value.activeMessageId;
       }
       yield* clearAssistantMessageState(input.messageId);
     });
-  const finalizeActiveAssistantSegmentForTurn = (input: {
+  const finalizeActiveAssistantSegmentsForTurn = (input: {
     event: ProviderRuntimeEvent;
     threadId: ThreadId;
     turnId: TurnId;
     createdAt: string;
     commandTag: string;
     finalDeltaCommandTag: string;
-    hasProjectedMessage: boolean;
+    projectedMessageIds: ReadonlySet<MessageId>;
     flushedMessageIds?: ReadonlySet<MessageId>;
   }) =>
     Effect.gen(function* () {
-      const activeMessageId = yield* getActiveAssistantMessageIdForTurn(
+      const segmentStates = yield* getAssistantSegmentStatesForTurn(input.threadId, input.turnId);
-        input.threadId,
+      if (Option.isNone(segmentStates)) {
-        input.turnId,
-      );
-      if (Option.isNone(activeMessageId)) {
         return;
       }
-      yield* finalizeAssistantMessage({
+      yield* Effect.forEach(
-        event: input.event,
+        segmentStates.value,
-        threadId: input.threadId,
+        ([baseKey, state]) =>
-        messageId: activeMessageId.value,
+          Effect.gen(function* () {
-        turnId: input.turnId,
+            if (state.activeMessageId === null) {
-        createdAt: input.createdAt,
+              return;
-        commandTag: input.commandTag,
+            }
-        finalDeltaCommandTag: input.finalDeltaCommandTag,
-        hasProjectedMessage:
-          input.hasProjectedMessage ||
-          (input.flushedMessageIds?.has(activeMessageId.value) ?? false),
-      });
-      yield* forgetAssistantMessageId(input.threadId, input.turnId, activeMessageId.value);
-      const state = yield* getAssistantSegmentStateForTurn(input.threadId, input.turnId);
+            const messageId = state.activeMessageId;
-      if (Option.isSome(state)) {
+            yield* finalizeAssistantMessage({
-        yield* setAssistantSegmentStateForTurn(input.threadId, input.turnId, {
+              event: input.event,
-          ...state.value,
+              threadId: input.threadId,
-          activeMessageId: null,
+              messageId,
-        });
+              turnId: input.turnId,
-      }
+              createdAt: input.createdAt,
+              commandTag: input.commandTag,
+              finalDeltaCommandTag: input.finalDeltaCommandTag,
+              hasProjectedMessage:
+                input.projectedMessageIds.has(messageId) ||
+                (input.flushedMessageIds?.has(messageId) ?? false),
+            });
+            yield* forgetAssistantMessageId(input.threadId, input.turnId, messageId);
+            yield* setAssistantSegmentStateForTurn(input.threadId, input.turnId, baseKey, {
+              ...state,
+              activeMessageId: null,
+            });
+          }),
+        { concurrency: 1 },
+      ).pipe(Effect.asVoid);
     });
   const upsertProposedPlan = (input: {
                     : "assistant-delta-flush-on-user-input-requested",
               })
             : new Set<MessageId>();
-        yield* finalizeActiveAssistantSegmentForTurn({
+        yield* finalizeActiveAssistantSegmentsForTurn({
           event,
           threadId: thread.id,
           turnId: pauseForUserTurnId,
             event.type === "request.opened"
               ? "assistant-delta-finalize-on-request-opened"
               : "assistant-delta-finalize-on-user-input-requested",
-          hasProjectedMessage:
+          projectedMessageIds: new Set(
-            detailedThread !== null &&
+            (detailedThread?.messages ?? [])
-            hasAssistantMessageForTurn(detailedThread.messages, pauseForUserTurnId, {
+              .filter(
-              streamingOnly: true,
+                (message) => message.role === "assistant" && message.turnId === pauseForUserTurnId,
-            }),
+              )
+              .map((message) => message.id),
+          ),
           flushedMessageIds,
         });
       }
       const assistantCompletion =
         event.type === "item.completed" && event.payload.itemType === "assistant_message"
           ? {
-              messageId: MessageId.make(
+              baseKey: assistantSegmentBaseKeyFromEvent(event),
-                `assistant:${event.itemId ?? event.turnId ?? event.eventId}`,
+              messageId: assistantSegmentMessageId(assistantSegmentBaseKeyFromEvent(event), 0),
-              ),
               fallbackText: event.payload.detail,
             }
           : undefined;
         const detailedThread = yield* getLoadedThreadDetail();
         const messages = detailedThread?.messages ?? [];
         const turnId = toTurnId(event.turnId);
-        const activeAssistantMessageId = turnId
+        const segmentState = turnId
-          ? yield* getActiveAssistantMessageIdForTurn(thread.id, turnId)
+          ? yield* getAssistantSegmentStateForTurn(thread.id, turnId, assistantCompletion.baseKey)
-          : Option.none<MessageId>();
+          : Option.none<AssistantSegmentState>();
-        const hasAssistantMessagesForTurn =
+        const activeAssistantMessageId = Option.flatMap(segmentState, (state) =>
-          turnId !== undefined ? hasAssistantMessageForTurn(messages, turnId) : false;
+          Option.fromNullishOr(state.activeMessageId),
+        );
         const assistantMessageId = Option.getOrElse(
           activeAssistantMessageId,
           () => assistantCompletion.messageId,
           !existingAssistantMessage || existingAssistantMessage.text.length === 0;
         const shouldSkipRedundantCompletion =
-          Option.isNone(activeAssistantMessageId) &&
+          Option.isSome(segmentState) && segmentState.value.activeMessageId === null;
-          turnId !== undefined &&
-          hasAssistantMessagesForTurn &&
-          (assistantCompletion.fallbackText?.trim().length ?? 0) === 0;
         if (!shouldSkipRedundantCompletion) {
           if (turnId && Option.isNone(activeAssistantMessageId)) {
           if (turnId) {
             yield* forgetAssistantMessageId(thread.id, turnId, assistantMessageId);
+            if (Option.isSome(segmentState)) {
+              yield* setAssistantSegmentStateForTurn(
+                thread.id,
+                turnId,
+                assistantCompletion.baseKey,
+                {
+                  ...segmentState.value,
+                  activeMessageId: null,
+                },
+              );
+            }
           }
         }
-        if (turnId) {
-          yield* clearAssistantSegmentStateForTurn(thread.id, turnId);
-        }
       }
       if (proposedPlanCompletion) {