diff --git a/apps/server/src/orchestration-v2/EventSink.ts b/apps/server/src/orchestration-v2/EventSink.ts index 81e64ddb70920..0c42ac73ff9cb 100644 --- a/apps/server/src/orchestration-v2/EventSink.ts +++ b/apps/server/src/orchestration-v2/EventSink.ts @@ -12,7 +12,6 @@ import { type ProjectId, ThreadId, } from "@t3tools/contracts"; -import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; import * as ProviderCommitBatch from "./ProviderCommitBatch.ts"; import * as Context from "effect/Context"; import * as DateTime from "effect/DateTime"; @@ -79,9 +78,10 @@ const CurrentProviderBatch = Context.Reference( - source: Stream.Stream, - process: (event: ProviderAdapter.ProviderAdapterV2Event) => Effect.Effect, + /** Items that are not provider snapshots, such as a consumer's own control items, never batch. */ + readonly consumeProviderEvents: ( + source: Stream.Stream, + process: (event: A) => Effect.Effect, after: Effect.Effect, canBatch: Effect.Effect, ) => Effect.Effect< diff --git a/apps/server/src/orchestration-v2/ProviderCommitBatch.ts b/apps/server/src/orchestration-v2/ProviderCommitBatch.ts index 60cee2b632f66..1a6ca8d28a609 100644 --- a/apps/server/src/orchestration-v2/ProviderCommitBatch.ts +++ b/apps/server/src/orchestration-v2/ProviderCommitBatch.ts @@ -1,4 +1,3 @@ -import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; import * as Cause from "effect/Cause"; import * as Clock from "effect/Clock"; import * as Config from "effect/Config"; @@ -19,8 +18,7 @@ export const enabled = Config.Boolean("T3CODE_PROVIDER_COMMIT_BATCHING").pipe( Effect.orDie, ); -type Event = ProviderAdapter.ProviderAdapterV2Event; -const isSnapshot = (event: Event) => +const isSnapshot = (event: { readonly type: string }) => event.type === "node.updated" || event.type === "message.updated" || event.type === "turn_item.updated" || @@ -232,9 +230,9 @@ export const make = Effect.fn("ProviderCommitBatch.make")(function* ( ), ); }); - const consume = Effect.fnUntraced(function* ( - source: Stream.Stream, - process: (event: Event) => Effect.Effect, + const consume = Effect.fnUntraced(function* ( + source: Stream.Stream, + process: (event: A) => Effect.Effect, after: Effect.Effect, canBatch: Effect.Effect = Effect.succeed(true), ) { diff --git a/apps/server/src/orchestration-v2/RunExecutionService.test.ts b/apps/server/src/orchestration-v2/RunExecutionService.test.ts index 5bfb4fe34dc6a..e25b4c7ab556a 100644 --- a/apps/server/src/orchestration-v2/RunExecutionService.test.ts +++ b/apps/server/src/orchestration-v2/RunExecutionService.test.ts @@ -14,6 +14,7 @@ import { type OrchestrationV2ProviderTurn, type OrchestrationV2Run, type OrchestrationV2RunAttempt, + type OrchestrationV2StoredEvent, type OrchestrationV2Subagent, type OrchestrationV2TurnItem, ProviderDriverKind, @@ -28,14 +29,19 @@ import { TurnItemId, } from "@t3tools/contracts"; import * as Cause from "effect/Cause"; +import * as Clock from "effect/Clock"; import * as Exit from "effect/Exit"; +import * as Fiber from "effect/Fiber"; import * as DateTime from "effect/DateTime"; import * as Deferred from "effect/Deferred"; +import * as Duration from "effect/Duration"; import * as Effect from "effect/Effect"; import * as Layer from "effect/Layer"; import * as Option from "effect/Option"; +import * as Queue from "effect/Queue"; import * as Ref from "effect/Ref"; import * as Stream from "effect/Stream"; +import * as TestClock from "effect/testing/TestClock"; import { CodexAppServerRequestError } from "effect-codex-app-server/errors"; import { SqlError, classifySqliteError } from "effect/sql/SqlError"; @@ -47,6 +53,7 @@ import * as EventSink from "./EventSink.ts"; import * as IdAllocator from "@t3tools/provider-core/server/IdAllocator"; import * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; import * as ProjectionStore from "./ProjectionStore.ts"; +import * as ProviderCommitBatch from "./ProviderCommitBatch.ts"; import * as ProviderEventIngestor from "./ProviderEventIngestor.ts"; import * as ProviderUsageAnalytics from "../telemetry/ProviderUsageAnalytics.ts"; import * as McpAppModelContext from "../mcpApps/McpAppModelContext.ts"; @@ -1885,6 +1892,40 @@ it.effect("keeps ingesting until queued background work clears after the probe g }), ); +it.effect("keeps ingesting while the last stored root snapshot reports pending work", () => + Effect.gen(function* () { + // The live probe is already idle, as when it reads adapter state ahead of the queued events. + const consumed = yield* Deferred.make(); + const result = yield* Effect.raceFirst( + runBackgroundItemScenario( + "root-snapshot-pending-work", + (ids) => { + const snapshot = (pendingBackgroundWork: boolean) => + ({ + type: "provider_thread.updated", + driver, + providerThread: { id: ids.providerThreadId, appThreadId: ids.threadId }, + pendingBackgroundWork, + }) as ProviderAdapter.ProviderAdapterV2Event; + return [ + snapshot(true), + rootTerminalEvent(ids, "completed"), + backgroundTurnItemEvent(ids, "dynamic_tool", "completed", 1), + snapshot(false), + ]; + }, + { + keepEventStreamOpen: true, + hasPendingBackgroundWorkForThread: () => Effect.succeed(false), + onEventsConsumed: Deferred.succeed(consumed, undefined).pipe(Effect.asVoid), + }, + ).pipe(Effect.asSome), + Deferred.await(consumed).pipe(Effect.as(Option.none())), + ); + assert.deepEqual(result, Option.some(["root-finalized", "turn_item:completed"])); + }), +); + it.effect("keeps ingesting a child thread's late background item completion", () => Effect.gen(function* () { const observed = yield* runBackgroundItemScenario("bg-child-item", (ids) => [ @@ -3692,8 +3733,18 @@ function rootTerminalEvent( function runBackgroundItemScenario( key: string, - makeEvents: (ids: BackgroundScenarioIds) => ReadonlyArray, + makeEvents: ( + ids: BackgroundScenarioIds, + ) => + | ReadonlyArray + | Stream.Stream, options?: { + readonly onIngest?: ( + event: ProviderAdapter.ProviderAdapterV2Event, + write: { readonly inSharedCommit: boolean }, + ) => Effect.Effect; + // Consumes through the real commit coordinator, as the event sink does. + readonly batched?: boolean; readonly keepEventStreamOpen?: boolean; readonly hasPendingBackgroundWorkForThread?: ( providerThread: OrchestrationV2ProviderThread, @@ -3719,6 +3770,24 @@ function runBackgroundItemScenario( const previewActivity = yield* PreviewAgentActivity.PreviewAgentActivity.pipe( Effect.provide(PreviewAgentActivity.layer), ); + let inSharedCommit = false; + const coordinator = + options?.batched === true + ? yield* ProviderCommitBatch.make( + (work) => + Effect.suspend(() => { + inSharedCommit = true; + return work; + }).pipe( + Effect.ensuring( + Effect.sync(() => { + inSharedCommit = false; + }), + ), + ), + true, + ) + : undefined; const layerTest = RunExecutionService.layer.pipe( Layer.provide( Layer.mergeAll( @@ -3727,6 +3796,7 @@ function runBackgroundItemScenario( Layer.mock(CheckpointService.CheckpointServiceV2)({ captureBaseline: () => Effect.void }), layerTurnItemPositionsTest, mockEventSink({ + ...(coordinator === undefined ? {} : { consumeProviderEvents: coordinator.consume }), write: () => Effect.succeed([]), writeWithEffects: (input) => Effect.gen(function* () { @@ -3744,9 +3814,10 @@ function runBackgroundItemScenario( }), IdAllocator.layer, Layer.mock(ProviderEventIngestor.ProviderEventIngestorV2)({ - ingestNormalized: (input) => - Effect.gen(function* () { + ingestNormalized: (input) => { + const write = Effect.gen(function* () { const event = input.event; + yield* options?.onIngest?.(event, { inSharedCommit }) ?? Effect.void; if (event.type === "turn_item.updated") { yield* Ref.update(observed, (current) => [ ...current, @@ -3759,8 +3830,14 @@ function runBackgroundItemScenario( `subagent:${event.subagent.status}`, ]); } - return []; - }), + // Only the count of stored events matters to the run. + return event.type === "provider_thread.updated" + ? [{} as OrchestrationV2StoredEvent] + : []; + }); + // A write outside a batch waits for queued snapshots, as the event sink's does. + return coordinator === undefined ? write : coordinator.barrier(write); + }, }), ServerSettings.layerTest(), ), @@ -3780,7 +3857,8 @@ function runBackgroundItemScenario( : { hasPendingBackgroundWorkForThread: options.hasPendingBackgroundWorkForThread }), subscribeEvents: Effect.gen(function* () { yield* options?.onSubscribe ?? Effect.void; - const events = Stream.fromIterable(makeEvents(ids)).pipe( + const source = makeEvents(ids); + const events = (Stream.isStream(source) ? source : Stream.fromIterable(source)).pipe( Stream.concat( Stream.fromEffect( Effect.andThen( @@ -3847,8 +3925,9 @@ function runBackgroundItemScenario( }); }).pipe(Effect.provide(layerTest)); - const closed = yield* Deferred.await(ingestionDone).pipe(Effect.timeoutOption("2 seconds")); - assert.isTrue(Option.isSome(closed), "event ingestion fiber did not finish"); + // No timeout: tests advance the TestClock, which would fire it early. + yield* Deferred.await(ingestionDone); + return yield* Ref.get(observed); }); } @@ -3900,3 +3979,395 @@ it.effect("keeps a failed root's still-running native children live while the ru ); }), ); + +it.effect.each([ + { name: "on time", wakesEarly: false }, + // Timers run on the monotonic clock, so one can fire while the wall clock still reads a + // millisecond before its deadline. + { name: "before the wall clock reaches the deadline", wakesEarly: true }, +])("delivers held text while an open-prompt provider idles, waking $name", ({ wakesEarly }) => + Effect.gen(function* () { + const key = `idle-paragraph-${wakesEarly}`; + const ids = backgroundScenarioIds(key); + const now = yield* DateTime.now; + const providerEvents = yield* Queue.unbounded(); + const delivered = yield* Ref.make>([]); + const paragraphDelivered = yield* Deferred.make(); + const nextParagraphDelivered = yield* Deferred.make(); + const clock = yield* Clock.Clock; + const assistantText = (text: string, streaming: boolean) => + [ + { + type: "message.updated", + driver, + message: { + createdBy: "agent", + creationSource: "provider", + id: MessageId.make(`message:${key}`), + threadId: ids.threadId, + runId: ids.runId, + nodeId: ids.rootNodeId, + role: "assistant", + text, + attachments: [], + streaming, + createdAt: now, + updatedAt: now, + }, + }, + { + type: "turn_item.updated", + driver, + turnItem: { + ...makeLinkedChildTurnItemFixture({ ids, driver, type: "assistant_message" }), + id: ids.itemId, + threadId: ids.threadId, + runId: ids.runId, + nodeId: ids.rootNodeId, + providerTurnId: ids.rootProviderTurnId, + text, + streaming, + }, + }, + ] as ReadonlyArray; + + const blockerStarted = yield* Deferred.make(); + const releaseBlocker = yield* Deferred.make(); + const blocker = streamedMessageEvent({ + id: `message:${key}:blocker`, + threadId: ids.threadId, + runId: ids.runId, + role: "assistant", + text: "Other", + streaming: false, + }); + const scenario = yield* runBackgroundItemScenario(key, () => Stream.fromQueue(providerEvents), { + batched: true, + onIngest: (event, write) => + event.type === "message.updated" && event.message.text === "Other" + ? Deferred.succeed(blockerStarted, undefined).pipe( + Effect.andThen(Deferred.await(releaseBlocker)), + ) + : event.type === "message.updated" || event.type === "turn_item.updated" + ? Effect.gen(function* () { + const message = event.type === "message.updated" ? event.message : event.turnItem; + const text = "text" in message ? message.text : ""; + const streaming = "streaming" in message && message.streaming; + yield* Ref.update(delivered, (current) => [ + ...current, + `${event.type}:${streaming ? "streaming" : "final"}:${JSON.stringify(text)}${ + write.inSharedCommit ? " in commit" : "" + }`, + ]); + if (event.type === "turn_item.updated") { + yield* Deferred.succeed(paragraphDelivered, undefined); + if (text === "First.\n\nWaiting on the build.\n\n") + yield* Deferred.succeed(nextParagraphDelivered, undefined); + } + }) + : Effect.void, + }).pipe( + Effect.provideService(Clock.Clock, { + ...clock, + sleep: (duration) => + clock.sleep(wakesEarly ? Duration.subtract(duration, Duration.millis(1)) : duration), + }), + Effect.forkChild, + ); + // Waking a sleep runs ingestion until it suspends, so an idle write has landed by now. + const lastShown = Ref.get(delivered).pipe(Effect.map((shown) => shown.at(-1))); + + yield* Queue.offerAll(providerEvents, assistantText("First.\n\nWaiting on the build", true)); + yield* Deferred.await(paragraphDelivered); + // The provider keeps its prompt open while a backgrounded command runs. + yield* TestClock.adjust("2 seconds"); + assert.strictEqual( + yield* lastShown, + 'turn_item.updated:streaming:"First.\\n\\nWaiting on the build"', + ); + // The next paragraph's tail still flushes on a later pause. + yield* Queue.offerAll( + providerEvents, + assistantText("First.\n\nWaiting on the build.\n\nRunning", true), + ); + yield* Deferred.await(nextParagraphDelivered); + yield* TestClock.adjust("1 second"); + assert.strictEqual( + yield* lastShown, + 'turn_item.updated:streaming:"First.\\n\\nWaiting on the build.\\n\\nRunning"', + ); + yield* Queue.offerAll(providerEvents, [ + ...assistantText("First.\n\nWaiting on the build.\n\nRunning tests", true), + blocker, + ...assistantText("First.\n\nWaiting on the build.\n\nRunning tests passed.", false), + ]); + yield* Deferred.await(blockerStarted); + // The appended text comes due while its final snapshots wait behind the blocked write. + yield* TestClock.adjust("1 second"); + yield* Deferred.succeed(releaseBlocker, undefined); + yield* Queue.offer(providerEvents, rootTerminalEvent(ids, "completed")); + yield* Fiber.join(scenario); + + // Snapshots share the batch commit; the idle write commits alone, before it counts as shown. + // The final snapshots win over the stale append that came due behind them. + assert.deepEqual(yield* Ref.get(delivered), [ + `message.updated:streaming:"First.\\n\\n" in commit`, + `turn_item.updated:streaming:"First.\\n\\n" in commit`, + 'message.updated:streaming:"First.\\n\\nWaiting on the build"', + 'turn_item.updated:streaming:"First.\\n\\nWaiting on the build"', + `message.updated:streaming:"First.\\n\\nWaiting on the build.\\n\\n" in commit`, + `turn_item.updated:streaming:"First.\\n\\nWaiting on the build.\\n\\n" in commit`, + 'message.updated:streaming:"First.\\n\\nWaiting on the build.\\n\\nRunning"', + 'turn_item.updated:streaming:"First.\\n\\nWaiting on the build.\\n\\nRunning"', + `message.updated:final:"First.\\n\\nWaiting on the build.\\n\\nRunning tests passed." in commit`, + `turn_item.updated:final:"First.\\n\\nWaiting on the build.\\n\\nRunning tests passed." in commit`, + ]); + }), +); + +const heldChildText = "partial response"; + +function childAssistantSnapshot( + ids: BackgroundScenarioIds, +): ProviderAdapter.ProviderAdapterV2Event { + return { + type: "turn_item.updated", + driver, + turnItem: { + ...makeLinkedChildTurnItemFixture({ ids, driver, type: "assistant_message" }), + text: heldChildText, + }, + } as ProviderAdapter.ProviderAdapterV2Event; +} + +function linkedSubagentEvent( + ids: BackgroundScenarioIds, + status: "running" | "completed", +): ProviderAdapter.ProviderAdapterV2Event { + return { + type: "subagent.updated", + driver, + subagent: { + ...makeRunOwnedSubagentFixture({ + ids, + providerInstanceId: ProviderInstanceId.make("codex"), + childThreadId: ids.childThreadId, + driver, + status: "running", + }), + status, + }, + }; +} + +const isHeldChildSnapshot = (event: ProviderAdapter.ProviderAdapterV2Event) => + event.type === "turn_item.updated" && + "text" in event.turnItem && + event.turnItem.text === heldChildText; + +// Feeds a held child snapshot, then waits until ingestion has read past it. +const holdChildSnapshot = Effect.fn(function* ( + ids: BackgroundScenarioIds, + events: Queue.Enqueue, + subagentIngested: Deferred.Deferred, +) { + yield* Queue.offerAll(events, [ + childThreadCreatedEvent(ids), + childAssistantSnapshot(ids), + linkedSubagentEvent(ids, "running"), + ]); + yield* Deferred.await(subagentIngested); +}); + +it.effect("settles a child message only an idle flush wrote when the root is interrupted", () => + Effect.gen(function* () { + const ids = backgroundScenarioIds("idle-child-interrupt"); + const events = yield* Queue.unbounded(); + const subagentIngested = yield* Deferred.make(); + const flushed = yield* Deferred.make(); + const finalized = yield* Ref.make>([]); + const scenario = yield* runBackgroundItemScenario( + "idle-child-interrupt", + () => Stream.fromQueue(events), + { + onIngest: (event) => + event.type === "subagent.updated" + ? Deferred.succeed(subagentIngested, undefined) + : isHeldChildSnapshot(event) + ? Deferred.succeed(flushed, undefined) + : Effect.void, + onFinalized: (written) => Ref.set(finalized, written), + }, + ).pipe(Effect.forkChild); + + yield* holdChildSnapshot(ids, events, subagentIngested); + yield* TestClock.adjust("1 second"); + yield* Deferred.await(flushed); + yield* Queue.offer(events, rootTerminalEvent(ids, "interrupted")); + yield* Fiber.join(scenario); + + const settled = (yield* Ref.get(finalized)).flatMap((event) => + event.type === "turn-item.updated" && event.payload.id === ids.childItemId + ? [event.payload] + : [], + ); + assert.equal(settled.length, 1); + assert.equal(settled[0]?.status, "interrupted"); + assert.equal("text" in settled[0]! ? settled[0].text : null, heldChildText); + assert.isFalse("streaming" in settled[0]! && settled[0].streaming); + }), +); + +it.effect.each([ + { name: "completed", terminalStatus: "completed" as const }, + { name: "failed with its children continuing", terminalStatus: "failed" as const }, +])("idle-flushes a continuing child's text after the root $name", ({ terminalStatus }) => + Effect.gen(function* () { + const key = `idle-child-after-${terminalStatus}`; + const ids = backgroundScenarioIds(key); + const events = yield* Queue.unbounded(); + const subagentIngested = yield* Deferred.make(); + const rootFinalized = yield* Deferred.make(); + const flushed = yield* Deferred.make(); + const rootWrites = yield* Ref.make>([]); + const scenario = yield* runBackgroundItemScenario(key, () => Stream.fromQueue(events), { + onIngest: (event) => + event.type === "subagent.updated" + ? Deferred.succeed(subagentIngested, undefined) + : isHeldChildSnapshot(event) + ? Deferred.succeed(flushed, undefined) + : event.type === "message.updated" && event.message.threadId === ids.threadId + ? Ref.update(rootWrites, (current) => [...current, event.message.text]) + : Effect.void, + onFinalized: () => Deferred.succeed(rootFinalized, undefined), + }).pipe(Effect.forkChild); + + yield* holdChildSnapshot(ids, events, subagentIngested); + // The root ends mid-paragraph: its tail is held, and the run settles without it. + yield* Queue.offer( + events, + streamedMessageEvent({ + id: `message:${key}:root`, + threadId: ids.threadId, + runId: ids.runId, + role: "assistant", + text: "Root.\n\nhalf", + streaming: true, + }), + ); + const terminal = rootTerminalEvent(ids, terminalStatus); + yield* Queue.offer( + events, + terminalStatus === "failed" + ? ({ ...terminal, childTurnsContinue: true } as ProviderAdapter.ProviderAdapterV2Event) + : terminal, + ); + yield* Deferred.await(rootFinalized); + yield* TestClock.adjust("1 second"); + yield* Deferred.await(flushed); + yield* Queue.offer(events, linkedSubagentEvent(ids, "completed")); + const observed = yield* Fiber.join(scenario); + + assert.deepEqual(observed.slice(-3), [ + "root-finalized", + "turn_item:running", + "subagent:completed", + ]); + // A streaming write after the run settled would reopen the root's message. + assert.deepEqual(yield* Ref.get(rootWrites), ["Root.\n\n"]); + }), +); + +function streamedMessageEvent(input: { + readonly id: string; + readonly threadId: ThreadId; + readonly runId: RunId | null; + readonly role: "assistant" | "user"; + readonly text: string; + readonly streaming: boolean; +}): ProviderAdapter.ProviderAdapterV2Event { + const now = DateTime.makeUnsafe("2026-07-21T12:00:00.000Z"); + return { + type: "message.updated", + driver, + message: { + createdBy: input.role === "user" ? "user" : "agent", + creationSource: "provider", + id: MessageId.make(input.id), + threadId: input.threadId, + runId: input.runId, + nodeId: null, + role: input.role, + text: input.text, + attachments: [], + streaming: input.streaming, + createdAt: now, + updatedAt: now, + }, + } as ProviderAdapter.ProviderAdapterV2Event; +} + +it.effect("drops a child's held text once the child is handed to another run", () => + Effect.gen(function* () { + const key = "idle-child-handoff"; + const ids = backgroundScenarioIds(key); + const events = yield* Queue.unbounded(); + const subagentIngested = yield* Deferred.make(); + const handoffRead = yield* Deferred.make(); + const rootFlushed = yield* Deferred.make(); + const staleWrites = yield* Ref.make>([]); + const scenario = yield* runBackgroundItemScenario(key, () => Stream.fromQueue(events), { + onIngest: (event) => { + if (event.type === "subagent.updated") return Deferred.succeed(subagentIngested, undefined); + if (event.type !== "message.updated" && event.type !== "turn_item.updated") + return Effect.void; + const message = event.type === "message.updated" ? event.message : event.turnItem; + if (message.threadId === ids.childThreadId) + return Ref.update(staleWrites, (current) => [...current, event.type]); + if (event.type === "message.updated" && event.message.role === "user") + return Deferred.succeed(handoffRead, undefined); + return Deferred.succeed(rootFlushed, undefined); + }, + }).pipe(Effect.forkChild); + + yield* holdChildSnapshot(ids, events, subagentIngested); + const handoff = linkedSubagentEvent(ids, "running"); + yield* Queue.offerAll(events, [ + streamedMessageEvent({ + id: `message:${key}:child`, + threadId: ids.childThreadId, + runId: null, + role: "assistant", + text: heldChildText, + streaming: true, + }), + handoff.type === "subagent.updated" + ? { ...handoff, subagent: { ...handoff.subagent, runId: RunId.make(`run:${key}:next`) } } + : handoff, + // Due in the same flush as the child's text, after it. + streamedMessageEvent({ + id: `message:${key}:root`, + threadId: ids.threadId, + runId: ids.runId, + role: "assistant", + text: "Root partial", + streaming: true, + }), + streamedMessageEvent({ + id: `message:${key}:marker`, + threadId: ids.threadId, + runId: ids.runId, + role: "user", + text: "marker", + streaming: false, + }), + ]); + yield* Deferred.await(handoffRead); + yield* TestClock.adjust("1 second"); + yield* Deferred.await(rootFlushed); + yield* Queue.end(events); + yield* Fiber.join(scenario); + + assert.deepEqual(yield* Ref.get(staleWrites), []); + }), +); diff --git a/apps/server/src/orchestration-v2/RunExecutionService.ts b/apps/server/src/orchestration-v2/RunExecutionService.ts index 66c0d2143e9aa..07811415fa312 100644 --- a/apps/server/src/orchestration-v2/RunExecutionService.ts +++ b/apps/server/src/orchestration-v2/RunExecutionService.ts @@ -32,9 +32,11 @@ import * as Effect from "effect/Effect"; import * as Exit from "effect/Exit"; import * as Fiber from "effect/Fiber"; import * as Layer from "effect/Layer"; +import * as Queue from "effect/Queue"; import * as Ref from "effect/Ref"; import * as Schedule from "effect/Schedule"; import * as Schema from "effect/Schema"; +import * as Stream from "effect/Stream"; import * as McpSessionRegistry from "../mcp/McpSessionRegistry.ts"; import * as McpAppModelContext from "../mcpApps/McpAppModelContext.ts"; @@ -174,6 +176,12 @@ function isRunOwnedSubagentTerminalStatus( return status === "interrupted" || status === "failed" || status === "cancelled"; } +function streamedThreadId(event: ProviderAdapter.ProviderAdapterV2Event): ThreadId | null { + if (event.type === "message.updated") return event.message.threadId; + if (event.type === "turn_item.updated") return event.turnItem.threadId; + return null; +} + function childTurnsContinue(terminal: ProviderTerminalEvent): boolean { return terminal.status === "failed" && terminal.childTurnsContinue === true; } @@ -1073,6 +1081,12 @@ export const layer: Layer.Layer< ReadonlySet >(new Set(inheritedBackgroundTurnItemsById.keys())); const openRunOwnedSubagents = yield* Ref.make(emptyOpenRunOwnedSubagentProjection()); + // The last stored root snapshot's pendingBackgroundWork; adapters that omit it leave it false. + let rootSnapshotPendingWork = false; + const assistantStreaming = makeAssistantStreamingFilter(responseStreamingMode); + // Set when ingestion ends or finalization settles the run's open items, which a later + // streaming snapshot must not reopen. + let idleFlushClosed = false; const finalizeRootRun = Effect.fn("RunExecutionService.finalizeRootRun")(function* ( terminal: ProviderTerminalEvent, ) { @@ -1115,6 +1129,7 @@ export const layer: Layer.Layer< ); if (isRunOwnedSubagentTerminalStatus(terminal.status) && !childTurnsContinue(terminal)) { yield* Ref.set(openRunOwnedSubagents, emptyOpenRunOwnedSubagentProjection()); + idleFlushClosed = true; } yield* Ref.set(rootRunFinalized, true); }); @@ -1301,6 +1316,10 @@ export const layer: Layer.Layer< if (yield* Ref.get(providerThreadOwnerLost)) { return true; } + // The live probe reads adapter state, which can be ahead of the events consumed here. + if (rootSnapshotPendingWork) { + return false; + } // Claude background Bash has no turn-item projection. Use only the // thread-scoped probe: session-wide pending work (siblings, wake // buffers, session subagents) must not pin this root subscription. @@ -1315,7 +1334,6 @@ export const layer: Layer.Layer< } return true; }); - const filterAssistantEvent = makeAssistantStreamingFilter(responseStreamingMode); // Held text writes nothing, so its item would take a position only // when finally written, behind every item that began after it. const reservedHeldTurnItemIds = new Set(); @@ -1398,6 +1416,121 @@ export const layer: Layer.Layer< } return { event, accepted }; }); + // Per-event ingestion avoids adding a span for every provider snapshot. + const ingestDelivered = Effect.fnUntraced(function* ( + event: ProviderAdapter.ProviderAdapterV2Event, + deliveredEvent: ProviderAdapter.ProviderAdapterV2Event, + ) { + const rootTerminalAlreadySeen = yield* Ref.get(rootTerminalSeen); + const isRootProviderThreadUpdate = + event.type === "provider_thread.updated" && + event.providerThread.id === input.providerThread.id; + const nativeChildTurn = + event.type === "provider_turn.updated" && + event.providerTurn.providerThreadId !== input.providerThread.id; + if (event.type === "provider_turn.updated" && !nativeChildTurn) { + firstRootProviderTurnId ??= event.providerTurn.id; + } + const storedEvents = yield* providerEventIngestor.ingestNormalized({ + analyticsContext: nativeChildTurn + ? { + ...analyticsContext, + origin: NATIVE_SUBAGENT_TURN_ORIGIN, + segment: "native_child", + } + : { + ...analyticsContext, + segment: + event.type === "provider_turn.updated" && + event.providerTurn.id !== firstRootProviderTurnId + ? "continuation" + : input.attempt.reason, + }, + providerSessionId: input.providerSessionId, + providerInstanceId: input.run.providerInstanceId, + threadId: input.run.threadId, + runId: input.run.id, + nodeId: input.rootNode.id, + event: deliveredEvent, + ...(isRootProviderThreadUpdate + ? rootTerminalAlreadySeen + ? { + writeIfProviderThreadOwner: { + providerThreadId: input.providerThread.id, + runId: input.run.id, + activeAttemptId: input.attempt.id, + expectedLastRunOrdinal: input.run.ordinal, + }, + } + : { + writeIfRunCurrent: { + runId: input.run.id, + activeAttemptId: input.attempt.id, + expectedStatus: "running" as const, + }, + } + : {}), + }); + if (isRootProviderThreadUpdate && rootTerminalAlreadySeen && storedEvents.length === 0) { + yield* Ref.set(providerThreadOwnerLost, true); + } + return storedEvents.length; + }); + // The timer only wakes ingestion. The flush runs as an ingestion step outside any provider + // batch, so it stays ordered with routing and snapshots, and a write has committed when + // it returns. + const idleFlushTicks = yield* Queue.dropping<{ + readonly type: "idle_flush"; + readonly at: number; + }>(1); + let idleFlushWake: { readonly at: number; readonly fiber: Fiber.Fiber } | undefined; + const scheduleIdleFlush = Effect.suspend(() => { + const at = assistantStreaming.nextIdleFlushAt(); + if ( + idleFlushClosed || + at === null || + (idleFlushWake !== undefined && idleFlushWake.at <= at) + ) + return Effect.void; + const previous = idleFlushWake?.fiber; + return Effect.gen(function* () { + if (previous !== undefined) yield* Fiber.interrupt(previous); + const now = DateTime.toEpochMillis(yield* DateTime.now); + const fiber = yield* Effect.sleep(Duration.millis(Math.max(0, at - now))).pipe( + Effect.andThen(Queue.offer(idleFlushTicks, { type: "idle_flush", at })), + Effect.asVoid, + Effect.forkDetach, + ); + idleFlushWake = { at, fiber }; + }); + }); + const closeIdleFlush = Effect.suspend(() => { + idleFlushClosed = true; + const wake = idleFlushWake; + idleFlushWake = undefined; + return wake === undefined ? Effect.void : Fiber.interrupt(wake.fiber); + }); + // Per-tick ingestion avoids adding a span for every idle wake. + const flushIdleAssistantText = Effect.fnUntraced(function* (wokeAt: number) { + // A tick means its timer slept until `wokeAt`, but the wall clock can still read up + // to a millisecond earlier; trusting it would keep the finished wake and never re-arm. + const now = Math.max(DateTime.toEpochMillis(yield* DateTime.now), wokeAt); + // A wake due by now has fired; a later deadline needs a new one. + if (idleFlushWake !== undefined && idleFlushWake.at <= now) idleFlushWake = undefined; + if (idleFlushClosed) return; + const rootSettled = yield* Ref.get(rootRunFinalized); + for (const event of assistantStreaming.takeIdleFlushes(now)) { + // Settling the root run closed its messages; only children may still stream. + if (rootSettled && streamedThreadId(event) === input.run.threadId) continue; + // The held text's child may have been handed to another run since it arrived. + if (!routeProviderEvent(event, routeIdentity, yield* Ref.get(eventRouting))[0]) + continue; + // A failed write fails ingestion like any provider write. + yield* ingestDelivered(event, event); + yield* trackChildLifecycle(event, true); + } + yield* scheduleIdleFlush; + }); const ingestEvent = Effect.fnUntraced(function* ({ event, accepted, @@ -1414,65 +1547,12 @@ export const layer: Layer.Layer< yield* Ref.set(latestRootProviderTurn, event.providerTurn); } let storedEventCount = 0; - const deliveredEvent = filterAssistantEvent( + const deliveredEvent = assistantStreaming.filter( event, DateTime.toEpochMillis(yield* DateTime.now), ); if (deliveredEvent) { - const rootTerminalAlreadySeen = yield* Ref.get(rootTerminalSeen); - const isRootProviderThreadUpdate = - event.type === "provider_thread.updated" && - event.providerThread.id === input.providerThread.id; - const nativeChildTurn = - event.type === "provider_turn.updated" && - event.providerTurn.providerThreadId !== input.providerThread.id; - if (event.type === "provider_turn.updated" && !nativeChildTurn) { - firstRootProviderTurnId ??= event.providerTurn.id; - } - const storedEvents = yield* providerEventIngestor.ingestNormalized({ - analyticsContext: nativeChildTurn - ? { - ...analyticsContext, - origin: NATIVE_SUBAGENT_TURN_ORIGIN, - segment: "native_child", - } - : { - ...analyticsContext, - segment: - event.type === "provider_turn.updated" && - event.providerTurn.id !== firstRootProviderTurnId - ? "continuation" - : input.attempt.reason, - }, - providerSessionId: input.providerSessionId, - providerInstanceId: input.run.providerInstanceId, - threadId: input.run.threadId, - runId: input.run.id, - nodeId: input.rootNode.id, - event: deliveredEvent, - ...(isRootProviderThreadUpdate - ? rootTerminalAlreadySeen - ? { - writeIfProviderThreadOwner: { - providerThreadId: input.providerThread.id, - runId: input.run.id, - activeAttemptId: input.attempt.id, - expectedLastRunOrdinal: input.run.ordinal, - }, - } - : { - writeIfRunCurrent: { - runId: input.run.id, - activeAttemptId: input.attempt.id, - expectedStatus: "running" as const, - }, - } - : {}), - }); - storedEventCount = storedEvents.length; - if (isRootProviderThreadUpdate && rootTerminalAlreadySeen && storedEventCount === 0) { - yield* Ref.set(providerThreadOwnerLost, true); - } + storedEventCount = yield* ingestDelivered(event, deliveredEvent); } else if ( event.type === "turn_item.updated" && (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning") && @@ -1489,6 +1569,7 @@ export const layer: Layer.Layer< if (event.type === "provider_thread.updated") { if (event.providerThread.id === input.providerThread.id && storedEventCount > 0) { yield* Ref.set(latestProviderThread, event.providerThread); + rootSnapshotPendingWork = event.pendingBackgroundWork === true; } } if ( @@ -1512,14 +1593,19 @@ export const layer: Layer.Layer< } } yield* trackChildLifecycle(event, deliveredEvent !== null); + yield* scheduleIdleFlush; }); const providerEventFiber = yield* eventSink .consumeProviderEvents( - eventSubscription.events, + eventSubscription.events.pipe( + Stream.merge(Stream.fromQueue(idleFlushTicks), { haltStrategy: "left" }), + ), (event) => - Effect.gen(function* () { - yield* ingestEvent(yield* routeEvent(event)); - }), + event.type === "idle_flush" + ? flushIdleAssistantText(event.at) + : Effect.gen(function* () { + yield* ingestEvent(yield* routeEvent(event)); + }), Effect.gen(function* () { if (yield* agentTurnsSettled) yield* endPreviewActivity; return yield* shouldStopProviderEventIngestion; @@ -1527,6 +1613,7 @@ export const layer: Layer.Layer< Ref.get(rootTerminalSeen).pipe(Effect.map((seen) => !seen)), ) .pipe( + Effect.ensuring(closeIdleFlush), Effect.mapError((cause) => new RunExecutionIngestError({ runId: input.run.id, cause })), Effect.flatMap(() => Effect.gen(function* () { diff --git a/apps/server/src/orchestration-v2/assistantStreaming.test.ts b/apps/server/src/orchestration-v2/assistantStreaming.test.ts index d71042312bfb1..0e8210ebd1f6a 100644 --- a/apps/server/src/orchestration-v2/assistantStreaming.test.ts +++ b/apps/server/src/orchestration-v2/assistantStreaming.test.ts @@ -1,5 +1,8 @@ -import { describe, expect, it } from "vite-plus/test"; -import { splitBufferedAssistantText } from "./assistantStreaming.ts"; +import { MessageId, ProviderDriverKind, ThreadId } from "@t3tools/contracts"; +import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; +import * as DateTime from "effect/DateTime"; +import { assert, describe, expect, it } from "vite-plus/test"; +import { makeAssistantStreamingFilter, splitBufferedAssistantText } from "./assistantStreaming.ts"; describe("V2 assistant streaming", () => { it("keeps code fences intact", () => { @@ -24,3 +27,169 @@ it("holds a streamed section heading until its content has a boundary", () => { rest: "Next", }); }); + +type Step = + // A snapshot at `at` ms, and the text it delivers at once (null when held back). + | { readonly at: number; readonly text: string; readonly shows: string | null } + | { readonly at: number; readonly final: string } + | { readonly id: string; readonly at: number; readonly text: string; readonly shows: string } + // Held snapshots due by `idle` ms, delivered whole. + | { readonly idle: number; readonly shows: ReadonlyArray } + | { readonly dueAt: number | null }; + +const createdAt = DateTime.makeUnsafe(0); +const snapshot = (id: string, text: string, streaming: boolean) => + ({ + type: "message.updated", + driver: ProviderDriverKind.make("codex"), + message: { + createdBy: "agent", + creationSource: "provider", + id: MessageId.make(id), + threadId: ThreadId.make("thread"), + runId: null, + nodeId: null, + role: "assistant", + text, + attachments: [], + streaming, + createdAt, + updatedAt: createdAt, + }, + }) satisfies ProviderAdapter.ProviderAdapterV2Event; +const shown = (event: ProviderAdapter.ProviderAdapterV2Event | null) => + event?.type === "message.updated" ? event.message.text : null; + +it.each<{ readonly name: string; readonly mode?: "turn"; readonly steps: ReadonlyArray }>([ + { + name: "an unfinished code fence", + steps: [ + { at: 0, text: "A\n\n```ts\nx()", shows: "A\n\n" }, + { dueAt: 800 }, + { idle: 799, shows: [] }, + { idle: 800, shows: ["A\n\n```ts\nx()"] }, + { dueAt: null }, + { at: 1_300, text: "A\n\n```ts\nx()\n```\n\nB", shows: "A\n\n```ts\nx()\n```\n\n" }, + ], + }, + { + name: "a list's partial last item", + steps: [ + { at: 0, text: "Steps:\n\n- one\n- tw", shows: "Steps:\n\n- one\n" }, + { idle: 800, shows: ["Steps:\n\n- one\n- tw"] }, + { at: 1_300, text: "Steps:\n\n- one\n- two\n- th", shows: "Steps:\n\n- one\n- two\n" }, + ], + }, + { + name: "an incomplete table row", + steps: [ + { at: 0, text: "Intro.\n\n| a | b |\n|---|---|\n| 1 |", shows: "Intro.\n\n" }, + { idle: 800, shows: ["Intro.\n\n| a | b |\n|---|---|\n| 1 |"] }, + { + at: 1_300, + text: "Intro.\n\n| a | b |\n|---|---|\n| 1 | 2 |\n\nAfter", + shows: "Intro.\n\n| a | b |\n|---|---|\n| 1 | 2 |\n\n", + }, + ], + }, + { + name: "more of the same paragraph after an idle delivery", + steps: [ + { at: 0, text: "A\n\nB", shows: "A\n\n" }, + { idle: 800, shows: ["A\n\nB"] }, + { at: 900, text: "A\n\nB more", shows: null }, + { idle: 1_699, shows: [] }, + { idle: 1_700, shows: ["A\n\nB more"] }, + { idle: 5_000, shows: [] }, + // Its only boundary is behind what the idle delivery showed. + { at: 5_100, text: "A\n\nB more and", shows: null }, + { at: 5_200, text: "A\n\nB more and.\n\nC", shows: "A\n\nB more and.\n\n" }, + { at: 5_300, final: "A\n\nB more and.\n\nC." }, + { idle: 10_000, shows: [] }, + ], + }, + { + name: "completion at the idle deadline", + steps: [ + { at: 0, text: "A\n\nB", shows: "A\n\n" }, + { at: 800, final: "A\n\nB." }, + { idle: 800, shows: [] }, + ], + }, + { + name: "completion right after an idle delivery at the same instant", + steps: [ + { at: 0, text: "A\n\nB", shows: "A\n\n" }, + { idle: 800, shows: ["A\n\nB"] }, + { at: 800, final: "A\n\nB." }, + { idle: 5_000, shows: [] }, + ], + }, + { + name: "several fast paragraphs inside the throttle", + steps: [ + { at: 0, text: "A\n\nB", shows: "A\n\n" }, + { at: 100, text: "A\n\nB\n\nC", shows: null }, + { at: 200, text: "A\n\nB\n\nC\n\nD", shows: null }, + { at: 300, text: "A\n\nB\n\nC\n\nD\n\nE", shows: null }, + { idle: 1_099, shows: [] }, + { idle: 1_100, shows: ["A\n\nB\n\nC\n\nD\n\nE"] }, + ], + }, + { + name: "two messages due at different times", + steps: [ + { id: "a", at: 0, text: "A\n\nx", shows: "A\n\n" }, + { id: "b", at: 300, text: "B\n\ny", shows: "B\n\n" }, + { dueAt: 800 }, + { idle: 800, shows: ["A\n\nx"] }, + { dueAt: 1_100 }, + { idle: 1_100, shows: ["B\n\ny"] }, + ], + }, + { + name: "empty text", + steps: [{ at: 0, text: "", shows: null }, { dueAt: null }], + }, + { + name: "whitespace-only text", + steps: [ + { at: 0, text: " \n", shows: null }, + { idle: 800, shows: [" \n"] }, + { idle: 5_000, shows: [] }, + ], + }, + { + // Rewrites replace delivered text: shorter at once, same length after the throttle. + name: "provider rewrites and truncation", + steps: [ + { at: 0, text: "Original long paragraph.\n\n", shows: "Original long paragraph.\n\n" }, + { at: 1_000, text: "New.\n\n", shows: "New.\n\n" }, + { at: 1_100, text: "Old.\n\n", shows: null }, + { idle: 1_900, shows: ["Old.\n\n"] }, + { at: 2_500, text: "Old", shows: null }, + { idle: 3_300, shows: ["Old"] }, + { at: 4_000, text: "", shows: null }, + { idle: 4_800, shows: [""] }, + ], + }, + { + name: "turn mode", + mode: "turn", + steps: [{ at: 0, text: "A\n\nB", shows: null }, { dueAt: null }], + }, +])("paragraph streaming holds and idle-delivers $name", ({ mode, steps }) => { + const filter = makeAssistantStreamingFilter(mode ?? "paragraph", 800); + for (const step of steps) { + if ("dueAt" in step) assert.equal(filter.nextIdleFlushAt(), step.dueAt); + else if ("idle" in step) + assert.deepEqual(filter.takeIdleFlushes(step.idle).map(shown), [...step.shows]); + else if ("final" in step) + assert.equal(shown(filter.filter(snapshot("m", step.final, false), step.at)), step.final); + else + assert.equal( + shown(filter.filter(snapshot("id" in step ? step.id : "m", step.text, true), step.at)), + step.shows, + ); + } +}); diff --git a/apps/server/src/orchestration-v2/assistantStreaming.ts b/apps/server/src/orchestration-v2/assistantStreaming.ts index 62a173d7318c7..ef6de3ab3f483 100644 --- a/apps/server/src/orchestration-v2/assistantStreaming.ts +++ b/apps/server/src/orchestration-v2/assistantStreaming.ts @@ -68,44 +68,129 @@ export function splitBufferedAssistantText(text: string): { ready: string; rest: return { ready: text.slice(0, boundary), rest: text.slice(boundary) }; } -export function makeAssistantStreamingFilter(mode: ResponseStreamingMode) { - const delivered = new Map(); - return ( - event: ProviderAdapter.ProviderAdapterV2Event, - now: number, - ): ProviderAdapter.ProviderAdapterV2Event | null => { - if ( - event.type === "node.updated" && - event.node.kind === "assistant_message" && - event.node.status === "running" - ) - return null; - const message = - event.type === "message.updated" && event.message.role === "assistant" - ? event.message - : event.type === "turn_item.updated" && - (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning") - ? event.turnItem - : null; - if (!message) return event; - const key = `${event.type}:${message.id}`; - if (!message.streaming) { - delivered.delete(key); - return event; - } - if (mode === "turn") return null; - const previous = delivered.get(key); - if (previous && now - previous.at < 400) return null; - const { ready } = splitBufferedAssistantText(message.text); - if (!ready || ready === previous?.text) return null; - delivered.set(key, { text: ready, at: now }); - if (event.type === "message.updated") - return { ...event, message: { ...event.message, text: ready } }; - if ( - event.type === "turn_item.updated" && - (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning") - ) - return { ...event, turnItem: { ...event.turnItem, text: ready } }; - return event; +// Long enough to ride out normal gaps between chunks; short enough that text written before a +// long tool call or background wait does not look lost. +const ASSISTANT_IDLE_FLUSH_MS = 800; + +type StreamedMessage = { + readonly id: string; + readonly text: string; + readonly streaming: boolean; +}; + +interface StreamTrack { + readonly text: string; + readonly at: number; + // Latest snapshot whose text has not all been delivered yet. + readonly held: ProviderAdapter.ProviderAdapterV2Event | null; + readonly dueAt: number; +} + +function streamedMessage(event: ProviderAdapter.ProviderAdapterV2Event): StreamedMessage | null { + if (event.type === "message.updated" && event.message.role === "assistant") return event.message; + if ( + event.type === "turn_item.updated" && + (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning") + ) + return event.turnItem; + return null; +} + +function withText( + event: ProviderAdapter.ProviderAdapterV2Event, + text: string, +): ProviderAdapter.ProviderAdapterV2Event { + if (event.type === "message.updated") return { ...event, message: { ...event.message, text } }; + if ( + event.type === "turn_item.updated" && + (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning") + ) + return { ...event, turnItem: { ...event.turnItem, text } }; + return event; +} + +/** + * Paragraph mode delivers streamed text at Markdown boundaries, at most every 400ms. A snapshot + * held back that way is delivered whole once its message has been idle for `idleFlushMs`, so a + * provider that pauses mid-turn still shows everything it wrote. + */ +export function makeAssistantStreamingFilter( + mode: ResponseStreamingMode, + idleFlushMs = ASSISTANT_IDLE_FLUSH_MS, +) { + const tracks = new Map(); + // Earliest deadline of any held snapshot, or earlier: a newer snapshot only pushes its own + // deadline later, so the cache stays O(1) per event and is recomputed when it comes due. + let nextDueAt: number | null = null; + const track = (key: string, value: StreamTrack) => { + tracks.set(key, value); + if (value.held !== null && (nextDueAt === null || value.dueAt < nextDueAt)) + nextDueAt = value.dueAt; + }; + return { + filter( + event: ProviderAdapter.ProviderAdapterV2Event, + now: number, + ): ProviderAdapter.ProviderAdapterV2Event | null { + if ( + event.type === "node.updated" && + event.node.kind === "assistant_message" && + event.node.status === "running" + ) + return null; + const message = streamedMessage(event); + if (!message) return event; + const key = `${event.type}:${message.id}`; + if (!message.streaming) { + tracks.delete(key); + return event; + } + if (mode === "turn") return null; + const previous = tracks.get(key); + const delivered = previous?.text ?? ""; + const hold = () => { + track(key, { + text: delivered, + at: previous?.at ?? Number.NEGATIVE_INFINITY, + held: message.text !== delivered ? event : null, + dueAt: now + idleFlushMs, + }); + return null; + }; + if (previous && now - previous.at < 400) return hold(); + const { ready } = splitBufferedAssistantText(message.text); + // An idle delivery may already show more than the next boundary; never shrink it. A + // provider rewrite or truncation is not a prefix and replaces it. + if (delivered.startsWith(ready)) return hold(); + track(key, { + text: ready, + at: now, + held: ready !== message.text ? event : null, + dueAt: now + idleFlushMs, + }); + return withText(event, ready); + }, + nextIdleFlushAt: (): number | null => nextDueAt, + /** Returns the held snapshots due by `now` and counts them as delivered. */ + takeIdleFlushes(now: number): ReadonlyArray { + const due: Array = []; + nextDueAt = null; + for (const [key, current] of tracks) { + const event = current.held; + if (event === null) continue; + if (current.dueAt > now) { + if (nextDueAt === null || current.dueAt < nextDueAt) nextDueAt = current.dueAt; + continue; + } + due.push(event); + tracks.set(key, { + text: streamedMessage(event)?.text ?? current.text, + at: now, + held: null, + dueAt: now, + }); + } + return due; + }, }; }