History

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
  1. apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts +174 −9
  2. 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
22202220 threadId: asThreadId("thread-1"),
22212221 turnId,
22222222 });
2223− await waitForThread(
2224− harness.readModel,
2225− (thread) => thread.session?.activeTurnId === turnId,
2226− );
2223+ await waitForThread(harness.readModel, (thread) => thread.session?.activeTurnId === turnId);
22272224
22282225 for (const [itemId, delta] of [
22292226 ["item-buffered-first", "first item"],
@@ -2251,12 +2248,11 @@ describe("ProviderRuntimeIngestion", () => { · 21 unmodified lines
22512248 payload: { state: "completed" },
22522249 });
22532250
2254− const thread = await waitForThread(
2255− harness.readModel,
2256− (entry) =>
2257− entry.messages.some((message: ProviderRuntimeTestMessage) =>
2251+ const thread = await waitForThread(harness.readModel, (entry) =>
2252+ entry.messages.some(
2253+ (message: ProviderRuntimeTestMessage) =>
22582254 message.id === "assistant:item-buffered-second" && !message.streaming,
2259− ),
2255+ ),
22602256 );
22612257 expect(thread.messages).toEqual(
22622258 expect.arrayContaining([
@@ -2266,6 +2262,175 @@ describe("ProviderRuntimeIngestion", () => { · 3 unmodified lines
22662262 );
22672263 });
22682264
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+
22692434 it("flushes and completes buffered assistant text when an approval request opens", async () => {
22702435 const harness = await createHarness();
22712436 const now = "2026-01-01T00:00:00.000Z";
apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts+106−110
@@ -90,7 +90,6 @@ function findTaskTitleInActivities( · 89 unmodified lines
9090 }
9191
9292 interface AssistantSegmentState {
93− baseKey: string;
9493 nextSegmentIndex: number;
9594 activeMessageId: MessageId | null;
9695 }
@@ -136,27 +135,6 @@ function sameId(left: string | null | undefined, right: string | null | undefine · 39 unmodified lines
136135 return left === right;
137136 }
138137
139−function hasAssistantMessageForTurn(
140− messages: ReadonlyArray<OrchestrationMessage>,
141− turnId: TurnId,
142− options?: { readonly streamingOnly?: boolean },
143−): boolean {
144− for (let index = 0; index < messages.length; index += 1) {
145− const message = messages[index];
146− if (!message) {
147− continue;
148− }
149− if (message.role !== "assistant" || message.turnId !== turnId) {
150− continue;
151− }
152− if (options?.streamingOnly === true && !message.streaming) {
153− continue;
154− }
155− return true;
156− }
157− return false;
158−}
159−
160138 function findMessageById(
161139 messages: ReadonlyArray<OrchestrationMessage>,
162140 messageId: MessageId,
@@ -973,13 +951,13 @@ const make = Effect.gen(function* () { · 810 unmodified lines
973951 lookup: () => Effect.succeed(""),
974952 });
975953
976− const assistantSegmentStateByTurnKey = yield* Cache.make<string, AssistantSegmentState>({
954+ const assistantSegmentStateByTurnKey = yield* Cache.make<
955+ string,
956+ Map<string, AssistantSegmentState>
957+ >({
977958 capacity: TURN_MESSAGE_IDS_BY_TURN_CACHE_CAPACITY,
978959 timeToLive: TURN_MESSAGE_IDS_BY_TURN_TTL,
979− lookup: () =>
980− Effect.die(
981− new Error("assistant segment state should be read through getOption before initialization"),
982− ),
960+ lookup: () => Effect.succeed(new Map()),
983961 });
984962
985963 const bufferedProposedPlanById = yield* Cache.make<string, { text: string; createdAt: string }>({
@@ -1071,52 +1049,59 @@ const make = Effect.gen(function* () { · 85 unmodified lines
10711049 const clearAssistantMessageIdsForTurn = (threadId: ThreadId, turnId: TurnId) =>
10721050 Cache.invalidate(turnMessageIdsByTurnKey, providerTurnKey(threadId, turnId));
10731051
1074− const getAssistantSegmentStateForTurn = (threadId: ThreadId, turnId: TurnId) =>
1052+ const getAssistantSegmentStatesForTurn = (threadId: ThreadId, turnId: TurnId) =>
10751053 Cache.getOption(assistantSegmentStateByTurnKey, providerTurnKey(threadId, turnId));
10761054
1055+ const getAssistantSegmentStateForTurn = (threadId: ThreadId, turnId: TurnId, baseKey: string) =>
1056+ getAssistantSegmentStatesForTurn(threadId, turnId).pipe(
1057+ Effect.map((states) =>
1058+ Option.flatMap(states, (entries) => Option.fromNullishOr(entries.get(baseKey))),
1059+ ),
1060+ );
1061+
10771062 const setAssistantSegmentStateForTurn = (
10781063 threadId: ThreadId,
10791064 turnId: TurnId,
1065+ baseKey: string,
10801066 state: AssistantSegmentState,
1081− ) => Cache.set(assistantSegmentStateByTurnKey, providerTurnKey(threadId, turnId), state);
1067+ ) =>
1068+ getAssistantSegmentStatesForTurn(threadId, turnId).pipe(
1069+ Effect.flatMap((existingStates) => {
1070+ const nextStates = new Map(Option.getOrElse(existingStates, () => new Map()));
1071+ nextStates.set(baseKey, state);
1072+ return Cache.set(
1073+ assistantSegmentStateByTurnKey,
1074+ providerTurnKey(threadId, turnId),
1075+ nextStates,
1076+ );
1077+ }),
1078+ );
10821079
10831080 const clearAssistantSegmentStateForTurn = (threadId: ThreadId, turnId: TurnId) =>
10841081 Cache.invalidate(assistantSegmentStateByTurnKey, providerTurnKey(threadId, turnId));
10851082
1086− const getActiveAssistantMessageIdForTurn = (threadId: ThreadId, turnId: TurnId) =>
1087− getAssistantSegmentStateForTurn(threadId, turnId).pipe(
1088− Effect.map((state) =>
1089− Option.flatMap(state, (entry) =>
1090− entry.activeMessageId ? Option.some(entry.activeMessageId) : Option.none(),
1091− ),
1092− ),
1093− );
1094−
10951083 const startAssistantSegmentForTurn = (input: {
10961084 threadId: ThreadId;
10971085 turnId: TurnId;
10981086 baseKey: string;
10991087 }) =>
1100− getAssistantSegmentStateForTurn(input.threadId, input.turnId).pipe(
1088+ getAssistantSegmentStateForTurn(input.threadId, input.turnId, input.baseKey).pipe(
11011089 Effect.flatMap((existingState) =>
11021090 Effect.gen(function* () {
1103− const nextState = Option.match(existingState, {
1104− onNone: () => ({
1105− baseKey: input.baseKey,
1106− nextSegmentIndex: 1,
1107− activeMessageId: assistantSegmentMessageId(input.baseKey, 0),
1108− }),
1109− onSome: (state) => {
1110− const segmentIndex = state.baseKey === input.baseKey ? state.nextSegmentIndex : 0;
1111− const messageId = assistantSegmentMessageId(input.baseKey, segmentIndex);
1112− return {
1113− baseKey: input.baseKey,
1114− nextSegmentIndex: state.baseKey === input.baseKey ? state.nextSegmentIndex + 1 : 1,
1115− activeMessageId: messageId,
1116− } satisfies AssistantSegmentState;
1117− },
1091+ const segmentIndex = Option.match(existingState, {
1092+ onNone: () => 0,
1093+ onSome: (state) => state.nextSegmentIndex,
11181094 });
1119− yield* setAssistantSegmentStateForTurn(input.threadId, input.turnId, nextState);
1095+ const nextState: AssistantSegmentState = {
1096+ nextSegmentIndex: segmentIndex + 1,
1097+ activeMessageId: assistantSegmentMessageId(input.baseKey, segmentIndex),
1098+ };
1099+ yield* setAssistantSegmentStateForTurn(
1100+ input.threadId,
1101+ input.turnId,
1102+ input.baseKey,
1103+ nextState,
1104+ );
11201105 return nextState.activeMessageId!;
11211106 }),
11221107 ),
@@ -1133,12 +1118,12 @@ const make = Effect.gen(function* () { · 10 unmodified lines
11331118 }
11341119
11351120 const baseKey = assistantSegmentBaseKeyFromEvent(input.event);
1136− const segmentState = yield* getAssistantSegmentStateForTurn(input.threadId, input.turnId);
1137− if (
1138− Option.isSome(segmentState) &&
1139− segmentState.value.baseKey === baseKey &&
1140− segmentState.value.activeMessageId !== null
1141− ) {
1121+ const segmentState = yield* getAssistantSegmentStateForTurn(
1122+ input.threadId,
1123+ input.turnId,
1124+ baseKey,
1125+ );
1126+ if (Option.isSome(segmentState) && segmentState.value.activeMessageId !== null) {
11421127 return segmentState.value.activeMessageId;
11431128 }
11441129
@@ -1313,46 +1298,51 @@ const make = Effect.gen(function* () { · 168 unmodified lines
13131298 yield* clearAssistantMessageState(input.messageId);
13141299 });
13151300
1316− const finalizeActiveAssistantSegmentForTurn = (input: {
1301+ const finalizeActiveAssistantSegmentsForTurn = (input: {
13171302 event: ProviderRuntimeEvent;
13181303 threadId: ThreadId;
13191304 turnId: TurnId;
13201305 createdAt: string;
13211306 commandTag: string;
13221307 finalDeltaCommandTag: string;
1323− hasProjectedMessage: boolean;
1308+ projectedMessageIds: ReadonlySet<MessageId>;
13241309 flushedMessageIds?: ReadonlySet<MessageId>;
13251310 }) =>
13261311 Effect.gen(function* () {
1327− const activeMessageId = yield* getActiveAssistantMessageIdForTurn(
1328− input.threadId,
1329− input.turnId,
1330− );
1331− if (Option.isNone(activeMessageId)) {
1312+ const segmentStates = yield* getAssistantSegmentStatesForTurn(input.threadId, input.turnId);
1313+ if (Option.isNone(segmentStates)) {
13321314 return;
13331315 }
13341316
1335− yield* finalizeAssistantMessage({
1336− event: input.event,
1337− threadId: input.threadId,
1338− messageId: activeMessageId.value,
1339− turnId: input.turnId,
1340− createdAt: input.createdAt,
1341− commandTag: input.commandTag,
1342− finalDeltaCommandTag: input.finalDeltaCommandTag,
1343− hasProjectedMessage:
1344− input.hasProjectedMessage ||
1345− (input.flushedMessageIds?.has(activeMessageId.value) ?? false),
1346− });
1347− yield* forgetAssistantMessageId(input.threadId, input.turnId, activeMessageId.value);
1317+ yield* Effect.forEach(
1318+ segmentStates.value,
1319+ ([baseKey, state]) =>
1320+ Effect.gen(function* () {
1321+ if (state.activeMessageId === null) {
1322+ return;
1323+ }
13481324
1349− const state = yield* getAssistantSegmentStateForTurn(input.threadId, input.turnId);
1350− if (Option.isSome(state)) {
1351− yield* setAssistantSegmentStateForTurn(input.threadId, input.turnId, {
1352− ...state.value,
1353− activeMessageId: null,
1354− });
1355− }
1325+ const messageId = state.activeMessageId;
1326+ yield* finalizeAssistantMessage({
1327+ event: input.event,
1328+ threadId: input.threadId,
1329+ messageId,
1330+ turnId: input.turnId,
1331+ createdAt: input.createdAt,
1332+ commandTag: input.commandTag,
1333+ finalDeltaCommandTag: input.finalDeltaCommandTag,
1334+ hasProjectedMessage:
1335+ input.projectedMessageIds.has(messageId) ||
1336+ (input.flushedMessageIds?.has(messageId) ?? false),
1337+ });
1338+ yield* forgetAssistantMessageId(input.threadId, input.turnId, messageId);
1339+ yield* setAssistantSegmentStateForTurn(input.threadId, input.turnId, baseKey, {
1340+ ...state,
1341+ activeMessageId: null,
1342+ });
1343+ }),
1344+ { concurrency: 1 },
1345+ ).pipe(Effect.asVoid);
13561346 });
13571347
13581348 const upsertProposedPlan = (input: {
@@ -1816,7 +1806,7 @@ const make = Effect.gen(function* () { · 457 unmodified lines
18161806 : "assistant-delta-flush-on-user-input-requested",
18171807 })
18181808 : new Set<MessageId>();
1819− yield* finalizeActiveAssistantSegmentForTurn({
1809+ yield* finalizeActiveAssistantSegmentsForTurn({
18201810 event,
18211811 threadId: thread.id,
18221812 turnId: pauseForUserTurnId,
@@ -1829,11 +1819,13 @@ const make = Effect.gen(function* () { · 6 unmodified lines
18291819 event.type === "request.opened"
18301820 ? "assistant-delta-finalize-on-request-opened"
18311821 : "assistant-delta-finalize-on-user-input-requested",
1832− hasProjectedMessage:
1833− detailedThread !== null &&
1834− hasAssistantMessageForTurn(detailedThread.messages, pauseForUserTurnId, {
1835− streamingOnly: true,
1836− }),
1822+ projectedMessageIds: new Set(
1823+ (detailedThread?.messages ?? [])
1824+ .filter(
1825+ (message) => message.role === "assistant" && message.turnId === pauseForUserTurnId,
1826+ )
1827+ .map((message) => message.id),
1828+ ),
18371829 flushedMessageIds,
18381830 });
18391831 }
@@ -1846,9 +1838,8 @@ const make = Effect.gen(function* () { · 6 unmodified lines
18461838 const assistantCompletion =
18471839 event.type === "item.completed" && event.payload.itemType === "assistant_message"
18481840 ? {
1849− messageId: MessageId.make(
1850− `assistant:${event.itemId ?? event.turnId ?? event.eventId}`,
1851− ),
1841+ baseKey: assistantSegmentBaseKeyFromEvent(event),
1842+ messageId: assistantSegmentMessageId(assistantSegmentBaseKeyFromEvent(event), 0),
18521843 fallbackText: event.payload.detail,
18531844 }
18541845 : undefined;
@@ -1865,11 +1856,12 @@ const make = Effect.gen(function* () { · 10 unmodified lines
18651856 const detailedThread = yield* getLoadedThreadDetail();
18661857 const messages = detailedThread?.messages ?? [];
18671858 const turnId = toTurnId(event.turnId);
1868− const activeAssistantMessageId = turnId
1869− ? yield* getActiveAssistantMessageIdForTurn(thread.id, turnId)
1870− : Option.none<MessageId>();
1871− const hasAssistantMessagesForTurn =
1872− turnId !== undefined ? hasAssistantMessageForTurn(messages, turnId) : false;
1859+ const segmentState = turnId
1860+ ? yield* getAssistantSegmentStateForTurn(thread.id, turnId, assistantCompletion.baseKey)
1861+ : Option.none<AssistantSegmentState>();
1862+ const activeAssistantMessageId = Option.flatMap(segmentState, (state) =>
1863+ Option.fromNullishOr(state.activeMessageId),
1864+ );
18731865 const assistantMessageId = Option.getOrElse(
18741866 activeAssistantMessageId,
18751867 () => assistantCompletion.messageId,
@@ -1879,10 +1871,7 @@ const make = Effect.gen(function* () { · 3 unmodified lines
18791871 !existingAssistantMessage || existingAssistantMessage.text.length === 0;
18801872
18811873 const shouldSkipRedundantCompletion =
1882− Option.isNone(activeAssistantMessageId) &&
1883− turnId !== undefined &&
1884− hasAssistantMessagesForTurn &&
1885− (assistantCompletion.fallbackText?.trim().length ?? 0) === 0;
1874+ Option.isSome(segmentState) && segmentState.value.activeMessageId === null;
18861875
18871876 if (!shouldSkipRedundantCompletion) {
18881877 if (turnId && Option.isNone(activeAssistantMessageId)) {
@@ -1905,12 +1894,19 @@ const make = Effect.gen(function* () { · 16 unmodified lines
19051894
19061895 if (turnId) {
19071896 yield* forgetAssistantMessageId(thread.id, turnId, assistantMessageId);
1897+ if (Option.isSome(segmentState)) {
1898+ yield* setAssistantSegmentStateForTurn(
1899+ thread.id,
1900+ turnId,
1901+ assistantCompletion.baseKey,
1902+ {
1903+ ...segmentState.value,
1904+ activeMessageId: null,
1905+ },
1906+ );
1907+ }
19081908 }
19091909 }
1910−
1911− if (turnId) {
1912− yield* clearAssistantSegmentStateForTurn(thread.id, turnId);
1913− }
19141910 }
19151911
19161912 if (proposedPlanCompletion) {