diff --git a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts index 923e8c16f540..ebb876b88803 100644 --- a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts +++ b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts @@ -2207,6 +2207,230 @@ describe("ProviderRuntimeIngestion", () => { expect(message?.streaming).toBe(false); }); + it("keeps buffered assistant items separate within one turn", async () => { + const harness = await createHarness(); + const now = "2026-01-01T00:00:00.000Z"; + const turnId = asTurnId("turn-buffered-multiple-items"); + + harness.emit({ + type: "turn.started", + eventId: asEventId("evt-turn-started-buffered-multiple-items"), + provider: ProviderDriverKind.make("codex"), + createdAt: now, + threadId: asThreadId("thread-1"), + turnId, + }); + await waitForThread(harness.readModel, (thread) => thread.session?.activeTurnId === turnId); + + for (const [itemId, delta] of [ + ["item-buffered-first", "first item"], + ["item-buffered-second", "second item"], + ] as const) { + harness.emit({ + type: "content.delta", + eventId: asEventId(`evt-${itemId}`), + provider: ProviderDriverKind.make("codex"), + createdAt: now, + threadId: asThreadId("thread-1"), + turnId, + itemId: asItemId(itemId), + payload: { streamKind: "assistant_text", delta }, + }); + } + + harness.emit({ + type: "turn.completed", + eventId: asEventId("evt-turn-completed-buffered-multiple-items"), + provider: ProviderDriverKind.make("codex"), + createdAt: now, + threadId: asThreadId("thread-1"), + turnId, + payload: { state: "completed" }, + }); + + const thread = await waitForThread(harness.readModel, (entry) => + entry.messages.some( + (message: ProviderRuntimeTestMessage) => + message.id === "assistant:item-buffered-second" && !message.streaming, + ), + ); + expect(thread.messages).toEqual( + expect.arrayContaining([ + expect.objectContaining({ id: "assistant:item-buffered-first", text: "first item" }), + expect.objectContaining({ id: "assistant:item-buffered-second", text: "second item" }), + ]), + ); + }); + + it("keeps interleaved buffered assistant deltas with their provider items", async () => { + const harness = await createHarness(); + const base = { + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-01-01T00:00:00.000Z", + threadId: asThreadId("thread-1"), + turnId: asTurnId("turn-buffered-interleaved-items"), + }; + + await harness.emitAndDrain([ + { ...base, type: "turn.started", eventId: asEventId("interleaved-start") }, + { + ...base, + type: "content.delta", + eventId: asEventId("interleaved-a-1"), + itemId: asItemId("item-interleaved-a"), + payload: { streamKind: "assistant_text", delta: "first " }, + }, + { + ...base, + type: "content.delta", + eventId: asEventId("interleaved-b"), + itemId: asItemId("item-interleaved-b"), + payload: { streamKind: "assistant_text", delta: "second" }, + }, + { + ...base, + type: "content.delta", + eventId: asEventId("interleaved-a-2"), + itemId: asItemId("item-interleaved-a"), + payload: { streamKind: "assistant_text", delta: "third" }, + }, + { + ...base, + type: "turn.completed", + eventId: asEventId("interleaved-completed"), + payload: { state: "completed" }, + }, + ]); + + const messages = (await harness.readModel()).threads[0]?.messages ?? []; + expect(messages).toEqual( + expect.arrayContaining([ + expect.objectContaining({ id: "assistant:item-interleaved-a", text: "first third" }), + expect.objectContaining({ id: "assistant:item-interleaved-b", text: "second" }), + ]), + ); + expect( + messages.some((message) => message.id === "assistant:item-interleaved-a:segment:1"), + ).toBe(false); + }); + + it("finalizes every open assistant item at an approval boundary", async () => { + const harness = await createHarness(); + const base = { + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-01-01T00:00:00.000Z", + threadId: asThreadId("thread-1"), + turnId: asTurnId("turn-buffered-multiple-items-approval"), + }; + + await harness.emitAndDrain([ + { ...base, type: "turn.started", eventId: asEventId("multiple-approval-start") }, + { + ...base, + type: "content.delta", + eventId: asEventId("multiple-approval-a"), + itemId: asItemId("item-multiple-approval-a"), + payload: { streamKind: "assistant_text", delta: "first item" }, + }, + { + ...base, + type: "content.delta", + eventId: asEventId("multiple-approval-b"), + itemId: asItemId("item-multiple-approval-b"), + payload: { streamKind: "assistant_text", delta: "second item" }, + }, + { + ...base, + type: "request.opened", + eventId: asEventId("multiple-approval-opened"), + requestId: ApprovalRequestId.make("req-multiple-approval"), + payload: { requestType: "command_execution_approval", detail: "pwd" }, + }, + ]); + + const messages = (await harness.readModel()).threads[0]?.messages ?? []; + expect(messages).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: "assistant:item-multiple-approval-a", + text: "first item", + streaming: false, + }), + expect.objectContaining({ + id: "assistant:item-multiple-approval-b", + text: "second item", + streaming: false, + }), + ]), + ); + }); + + it("completes only the matching assistant item when items overlap", async () => { + const harness = await createHarness({ serverSettings: { enableLegacyTokenStreaming: true } }); + const base = { + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-01-01T00:00:00.000Z", + threadId: asThreadId("thread-1"), + turnId: asTurnId("turn-streaming-overlapping-items"), + }; + + await harness.emitAndDrain([ + { ...base, type: "turn.started", eventId: asEventId("overlap-start") }, + { + ...base, + type: "content.delta", + eventId: asEventId("overlap-a"), + itemId: asItemId("item-overlap-a"), + payload: { streamKind: "assistant_text", delta: "first" }, + }, + { + ...base, + type: "content.delta", + eventId: asEventId("overlap-b-1"), + itemId: asItemId("item-overlap-b"), + payload: { streamKind: "assistant_text", delta: "second" }, + }, + { + ...base, + type: "item.completed", + eventId: asEventId("overlap-a-completed"), + itemId: asItemId("item-overlap-a"), + payload: { itemType: "assistant_message", status: "completed" }, + }, + ]); + + let messages = (await harness.readModel()).threads[0]?.messages ?? []; + expect(messages.find((message) => message.id === "assistant:item-overlap-a")?.streaming).toBe( + false, + ); + expect(messages.find((message) => message.id === "assistant:item-overlap-b")?.streaming).toBe( + true, + ); + + await harness.emitAndDrain([ + { + ...base, + type: "content.delta", + eventId: asEventId("overlap-b-2"), + itemId: asItemId("item-overlap-b"), + payload: { streamKind: "assistant_text", delta: " item" }, + }, + { + ...base, + type: "item.completed", + eventId: asEventId("overlap-b-completed"), + itemId: asItemId("item-overlap-b"), + payload: { itemType: "assistant_message", status: "completed" }, + }, + ]); + + messages = (await harness.readModel()).threads[0]?.messages ?? []; + expect(messages.find((message) => message.id === "assistant:item-overlap-b")).toMatchObject({ + text: "second item", + streaming: false, + }); + }); + it("flushes and completes buffered assistant text when an approval request opens", async () => { const harness = await createHarness(); const now = "2026-01-01T00:00:00.000Z"; @@ -2444,6 +2668,50 @@ describe("ProviderRuntimeIngestion", () => { ).toBe(false); }); + it("keeps item completion fallback text after a whitespace-only approval segment", async () => { + const harness = await createHarness(); + const base = { + provider: ProviderDriverKind.make("codex"), + createdAt: "2026-03-28T06:29:00.000Z", + threadId: asThreadId("thread-1"), + turnId: asTurnId("turn-whitespace-approval-fallback"), + }; + + await harness.emitAndDrain([ + { ...base, type: "turn.started", eventId: asEventId("whitespace-fallback-start") }, + { + ...base, + type: "content.delta", + eventId: asEventId("whitespace-fallback-delta"), + itemId: asItemId("item-whitespace-fallback"), + payload: { streamKind: "assistant_text", delta: "\n\n" }, + }, + { + ...base, + type: "request.opened", + eventId: asEventId("whitespace-fallback-request"), + requestId: ApprovalRequestId.make("req-whitespace-fallback"), + payload: { requestType: "command_execution_approval", detail: "pwd" }, + }, + { + ...base, + type: "item.completed", + eventId: asEventId("whitespace-fallback-completed"), + itemId: asItemId("item-whitespace-fallback"), + payload: { + itemType: "assistant_message", + status: "completed", + detail: "final answer", + }, + }, + ]); + + const messages = (await harness.readModel()).threads[0]?.messages ?? []; + expect(messages).toContainEqual( + expect.objectContaining({ text: "final answer", streaming: false }), + ); + }); + it("starts a new buffered assistant message segment after approval and completes without duplication", async () => { const harness = await createHarness(); const startedAt = "2026-03-28T06:07:00.000Z"; diff --git a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts index 449123d2e9da..4d577f231df0 100644 --- a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts +++ b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.ts @@ -90,7 +90,6 @@ function findTaskTitleInActivities( } interface AssistantSegmentState { - baseKey: string; nextSegmentIndex: number; activeMessageId: MessageId | null; } @@ -136,27 +135,6 @@ function sameId(left: string | null | undefined, right: string | null | undefine return left === right; } -function hasAssistantMessageForTurn( - messages: ReadonlyArray, - 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, messageId: MessageId, @@ -973,13 +951,13 @@ const make = Effect.gen(function* () { lookup: () => Effect.succeed(""), }); - const assistantSegmentStateByTurnKey = yield* Cache.make({ + const assistantSegmentStateByTurnKey = yield* Cache.make< + string, + Map + >({ capacity: TURN_MESSAGE_IDS_BY_TURN_CACHE_CAPACITY, timeToLive: TURN_MESSAGE_IDS_BY_TURN_TTL, - lookup: () => - Effect.die( - new Error("assistant segment state should be read through getOption before initialization"), - ), + lookup: () => Effect.succeed(new Map()), }); const bufferedProposedPlanById = yield* Cache.make({ @@ -1071,52 +1049,59 @@ const make = Effect.gen(function* () { 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, { - onNone: () => ({ - baseKey: input.baseKey, - 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; - }, + const segmentIndex = Option.match(existingState, { + onNone: () => 0, + onSome: (state) => state.nextSegmentIndex, }); - 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!; }), ), @@ -1132,18 +1117,20 @@ const make = Effect.gen(function* () { return assistantSegmentMessageId(assistantSegmentBaseKeyFromEvent(input.event), 0); } - const activeMessageId = yield* getActiveAssistantMessageIdForTurn( + const baseKey = assistantSegmentBaseKeyFromEvent(input.event); + const segmentState = yield* getAssistantSegmentStateForTurn( input.threadId, input.turnId, + baseKey, ); - if (Option.isSome(activeMessageId)) { - return activeMessageId.value; + if (Option.isSome(segmentState) && segmentState.value.activeMessageId !== null) { + return segmentState.value.activeMessageId; } return yield* startAssistantSegmentForTurn({ threadId: input.threadId, turnId: input.turnId, - baseKey: assistantSegmentBaseKeyFromEvent(input.event), + baseKey, }); }); @@ -1311,46 +1298,51 @@ const make = Effect.gen(function* () { 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; flushedMessageIds?: ReadonlySet; }) => Effect.gen(function* () { - const activeMessageId = yield* getActiveAssistantMessageIdForTurn( - input.threadId, - input.turnId, - ); - if (Option.isNone(activeMessageId)) { + const segmentStates = yield* getAssistantSegmentStatesForTurn(input.threadId, input.turnId); + if (Option.isNone(segmentStates)) { return; } - yield* finalizeAssistantMessage({ - event: input.event, - threadId: input.threadId, - messageId: activeMessageId.value, - turnId: input.turnId, - createdAt: input.createdAt, - commandTag: input.commandTag, - finalDeltaCommandTag: input.finalDeltaCommandTag, - hasProjectedMessage: - input.hasProjectedMessage || - (input.flushedMessageIds?.has(activeMessageId.value) ?? false), - }); - yield* forgetAssistantMessageId(input.threadId, input.turnId, activeMessageId.value); + yield* Effect.forEach( + segmentStates.value, + ([baseKey, state]) => + Effect.gen(function* () { + if (state.activeMessageId === null) { + return; + } - const state = yield* getAssistantSegmentStateForTurn(input.threadId, input.turnId); - if (Option.isSome(state)) { - yield* setAssistantSegmentStateForTurn(input.threadId, input.turnId, { - ...state.value, - activeMessageId: null, - }); - } + const messageId = state.activeMessageId; + yield* finalizeAssistantMessage({ + event: input.event, + threadId: input.threadId, + 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: { @@ -1814,7 +1806,7 @@ const make = Effect.gen(function* () { : "assistant-delta-flush-on-user-input-requested", }) : new Set(); - yield* finalizeActiveAssistantSegmentForTurn({ + yield* finalizeActiveAssistantSegmentsForTurn({ event, threadId: thread.id, turnId: pauseForUserTurnId, @@ -1827,11 +1819,13 @@ const make = Effect.gen(function* () { event.type === "request.opened" ? "assistant-delta-finalize-on-request-opened" : "assistant-delta-finalize-on-user-input-requested", - hasProjectedMessage: - detailedThread !== null && - hasAssistantMessageForTurn(detailedThread.messages, pauseForUserTurnId, { - streamingOnly: true, - }), + projectedMessageIds: new Set( + (detailedThread?.messages ?? []) + .filter( + (message) => message.role === "assistant" && message.turnId === pauseForUserTurnId, + ) + .map((message) => message.id), + ), flushedMessageIds, }); } @@ -1844,9 +1838,8 @@ const make = Effect.gen(function* () { const assistantCompletion = event.type === "item.completed" && event.payload.itemType === "assistant_message" ? { - messageId: MessageId.make( - `assistant:${event.itemId ?? event.turnId ?? event.eventId}`, - ), + baseKey: assistantSegmentBaseKeyFromEvent(event), + messageId: assistantSegmentMessageId(assistantSegmentBaseKeyFromEvent(event), 0), fallbackText: event.payload.detail, } : undefined; @@ -1863,14 +1856,25 @@ const make = Effect.gen(function* () { const detailedThread = yield* getLoadedThreadDetail(); const messages = detailedThread?.messages ?? []; const turnId = toTurnId(event.turnId); - const activeAssistantMessageId = turnId - ? yield* getActiveAssistantMessageIdForTurn(thread.id, turnId) - : Option.none(); - const hasAssistantMessagesForTurn = - turnId !== undefined ? hasAssistantMessageForTurn(messages, turnId) : false; + const segmentState = turnId + ? yield* getAssistantSegmentStateForTurn(thread.id, turnId, assistantCompletion.baseKey) + : Option.none(); + const activeAssistantMessageId = Option.flatMap(segmentState, (state) => + Option.fromNullishOr(state.activeMessageId), + ); + const assistantSegmentPrefix = `${assistantCompletion.messageId}:segment:`; + const existingAssistantMessagesForItem = messages.filter( + (message) => + message.role === "assistant" && + message.turnId === turnId && + (message.id === assistantCompletion.messageId || + message.id.startsWith(assistantSegmentPrefix)), + ); + const latestAssistantMessageForItem = + existingAssistantMessagesForItem[existingAssistantMessagesForItem.length - 1]; const assistantMessageId = Option.getOrElse( activeAssistantMessageId, - () => assistantCompletion.messageId, + () => latestAssistantMessageForItem?.id ?? assistantCompletion.messageId, ); const existingAssistantMessage = findMessageById(messages, assistantMessageId); const shouldApplyFallbackCompletionText = @@ -1878,8 +1882,7 @@ const make = Effect.gen(function* () { const shouldSkipRedundantCompletion = Option.isNone(activeAssistantMessageId) && - turnId !== undefined && - hasAssistantMessagesForTurn && + existingAssistantMessagesForItem.length > 0 && (assistantCompletion.fallbackText?.trim().length ?? 0) === 0; if (!shouldSkipRedundantCompletion) { @@ -1903,12 +1906,19 @@ const make = Effect.gen(function* () { 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) {