diff --git a/apps/server/src/orchestration-v2/EventSink.ts b/apps/server/src/orchestration-v2/EventSink.ts index 23db6c69dbdbb..81e64ddb70920 100644 --- a/apps/server/src/orchestration-v2/EventSink.ts +++ b/apps/server/src/orchestration-v2/EventSink.ts @@ -12,6 +12,8 @@ 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"; import * as Effect from "effect/Effect"; @@ -22,6 +24,7 @@ import * as Semaphore from "effect/Semaphore"; import * as Schema from "effect/Schema"; import * as Stream from "effect/Stream"; import * as SqlClient from "effect/sql/SqlClient"; +import * as SqlError from "effect/sql/SqlError"; import { replayAndBufferProjectedLiveEvents } from "./LiveStreamBudget.ts"; import type { UnsequencedProjectEvent } from "../persistence/OrchestrationEventStore.ts"; @@ -66,7 +69,32 @@ export class EventSinkStreamError extends Schema.TaggedError>; + eventCount: number; +} +const CurrentProviderBatch = Context.Reference( + "t3/orchestration-v2/EventSink/CurrentProviderBatch", + { defaultValue: () => undefined }, +); + export interface EventSinkV2Shape { + readonly consumeProviderEvents: ( + source: Stream.Stream, + process: (event: ProviderAdapter.ProviderAdapterV2Event) => Effect.Effect, + after: Effect.Effect, + canBatch: Effect.Effect, + ) => Effect.Effect< + void, + E | E2 | ProviderCommitBatch.CommitBatchError | Schema.SchemaError, + R | R2 + >; + + /** Normalize and persist an ordered group against one transaction; publish after its commit. */ + readonly withProviderBatch: ( + effect: Effect.Effect, + ) => Effect.Effect; + readonly write: (input: { /** Return no events if another thread owns this native session at commit time. */ readonly unlessNativeThreadOwned?: Pick< @@ -256,7 +284,7 @@ const layerBase: Layer.Layer< // Serialize commit and publish; reordered publication makes clients drop earlier sequences. const publishLane = yield* Semaphore.make(1); - const commitThenPublish = ( + const commitThenPublishNow = ( transaction: Effect.Effect, publish: (committed: A) => Effect.Effect, ) => @@ -416,6 +444,8 @@ const layerBase: Layer.Layer< "orchestration_v2.thread_id": input.events[0]?.threadId ?? null, }); + const batch = yield* CurrentProviderBatch; + if (batch !== undefined) batch.eventCount += input.events.length; const { committed } = yield* commitThenPublish( Effect.gen(function* () { // Check ownership in the same transaction that claims the native session. @@ -877,7 +907,95 @@ const layerBase: Layer.Layer< ); }; + const commitThenPublish = ( + transaction: Effect.Effect, + publish: (committed: A) => Effect.Effect, + ) => + Effect.flatMap(CurrentProviderBatch, (batch) => + batch === undefined + ? Effect.flatMap(Effect.serviceOption(sql.transactionService), (existing) => + Option.isSome(existing) + ? commitThenPublishNow(transaction, publish) + : coordinator.barrier(commitThenPublishNow(transaction, publish)), + ) + : transaction.pipe( + Effect.tap((committed) => + Effect.sync(() => { + batch.publish.push(publish(committed)); + }), + ), + ), + ); + + const withProviderBatch: EventSinkV2Shape["withProviderBatch"] = Effect.fn( + "orchestrationV2.EventSink.providerBatch", + )(function* (effect: Effect.Effect) { + if ((yield* CurrentProviderBatch) !== undefined) return yield* effect; + const batch: ProviderBatchPublication = { publish: [], eventCount: 0 }; + let started = false; + const result = yield* commitThenPublishNow( + Effect.sync(() => { + started = true; + }).pipe(Effect.andThen(Effect.provideService(effect, CurrentProviderBatch, batch))), + () => Effect.forEach(batch.publish, (publish) => publish, { discard: true }), + ).pipe( + // Normalization mutates run-local state, so only retry before its transaction body starts. + Effect.retry({ + while: (cause) => + !started && SqlError.isSqlError(cause) && cause.reason._tag === "LockTimeoutError", + times: 2, + }), + Effect.mapError((cause) => + SqlError.isSqlError(cause) + ? new EventSinkWriteError({ eventCount: batch.eventCount, cause }) + : cause, + ), + ); + yield* Effect.annotateCurrentSpan({ + "orchestration_v2.event_count": batch.eventCount, + "orchestration_v2.batch_writes": batch.publish.length, + }); + return result; + }); + + const coordinator = yield* ProviderCommitBatch.make( + (work) => + withProviderBatch(work).pipe( + Effect.mapError((cause) => + ProviderCommitBatch.isCommitBatchError(cause) + ? cause + : new ProviderCommitBatch.CommitBatchError({ cause }), + ), + ), + yield* ProviderCommitBatch.enabled, + // Each snapshot owns a savepoint and its deferred publications within the shared commit. + Effect.fnUntraced(function* ( + work: Effect.Effect, + ) { + const batch = yield* CurrentProviderBatch; + const publicationCount = batch?.publish.length ?? 0; + const eventCount = batch?.eventCount ?? 0; + return yield* sql.withTransaction(work).pipe( + Effect.onError(() => + Effect.sync(() => { + if (batch !== undefined) { + batch.publish.length = publicationCount; + batch.eventCount = eventCount; + } + }), + ), + Effect.mapError((cause) => + ProviderCommitBatch.isCommitBatchError(cause) + ? cause + : new ProviderCommitBatch.CommitBatchError({ cause }), + ), + ); + }), + ); + return EventSinkV2.of({ + consumeProviderEvents: coordinator.consume, + withProviderBatch, write: (input) => writeEffect({ ...input, effects: [] }).pipe( Effect.mapError( diff --git a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts index b929fc6ece3be..e1b2e7610663c 100644 --- a/apps/server/src/orchestration-v2/FoundationPersistence.test.ts +++ b/apps/server/src/orchestration-v2/FoundationPersistence.test.ts @@ -27,9 +27,14 @@ import { TurnItemId, } from "@t3tools/contracts"; import * as NodeSqliteClient from "@t3tools/shared/nodeSqliteClient"; +import type * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; +import * as ConfigProvider from "effect/ConfigProvider"; +import * as Queue from "effect/Queue"; import * as Arr from "effect/Array"; +import * as Cause from "effect/Cause"; import * as Clock from "effect/Clock"; import * as DateTime from "effect/DateTime"; +import * as Exit from "effect/Exit"; import * as Deferred from "effect/Deferred"; import * as Duration from "effect/Duration"; import * as Effect from "effect/Effect"; @@ -325,6 +330,239 @@ it.effect("paces scheduled compaction while preserving retention state and fixed }).pipe(Effect.scoped, Effect.provide(layerTest)), ); +it.effect( + "commits an ordered provider batch atomically and publishes only after the outer commit", + () => + Effect.gen(function* () { + const sink = yield* EventSink.EventSinkV2; + const projections = yield* ProjectionStore.ProjectionStoreV2; + const store = yield* EventStore.EventStoreV2; + const now = yield* DateTime.now; + const thread = makeThread(ThreadId.make("thread:single-stream-batch"), now); + yield* sink.write({ events: [threadCreatedEvent({ id: "batch-created", thread, now })] }); + const message = (id: string): OrchestrationV2DomainEvent => ({ + id: EventId.make(id), + type: "message.updated", + threadId: thread.id, + occurredAt: now, + payload: { + id: MessageId.make("batch-message"), + threadId: thread.id, + runId: null, + nodeId: null, + createdBy: "agent", + creationSource: "provider", + role: "assistant", + text: id, + streaming: true, + attachments: [], + createdAt: now, + updatedAt: now, + }, + }); + const highWater = yield* sink.latestSequence(); + const callerError = { _tag: "CallerError" } as const; + assert.strictEqual( + yield* sink.withProviderBatch(Effect.fail(callerError)).pipe(Effect.flip), + callerError, + ); + const failed = yield* sink + .withProviderBatch( + Effect.gen(function* () { + yield* sink.write({ events: [message("batch-duplicate")] }); + yield* sink.write({ events: [message("batch-duplicate")] }); + }), + ) + .pipe(Effect.exit); + assert.isTrue(Exit.isFailure(failed)); + assert.equal(yield* sink.latestSequence(), highWater); + assert.isEmpty((yield* projections.getThreadRecords(thread.id, ["messages"])).messages); + + const firstPublished = yield* Deferred.make(); + const writesApplied = yield* Deferred.make(); + const release = yield* Deferred.make(); + const live = yield* sink.stream({ afterSequence: highWater }).pipe( + Stream.tap(() => Deferred.succeed(firstPublished, undefined)), + Stream.take(2), + Stream.runCollect, + Effect.forkScoped({ startImmediately: true }), + ); + const batch = yield* sink + .withProviderBatch( + Effect.gen(function* () { + yield* sink.write({ events: [message("batch-first")] }); + yield* sink.write({ events: [message("batch-second")] }); + assert.equal( + (yield* projections.getThreadRecords(thread.id, ["messages"])).messages[0]?.text, + "batch-second", + ); + yield* Deferred.succeed(writesApplied, undefined); + yield* Deferred.await(release); + }), + ) + .pipe(Effect.forkScoped({ startImmediately: true })); + yield* Deferred.await(writesApplied); + yield* TestClock.adjust(0); + assert.isFalse(yield* Deferred.isDone(firstPublished)); + assert.isUndefined(yield* Effect.sync(() => batch.pollUnsafe())); + yield* Deferred.succeed(release, undefined); + yield* Fiber.join(batch); + const published = yield* Fiber.join(live); + assert.deepEqual( + published.map((stored) => stored.event.id), + ["batch-first", "batch-second"], + ); + assert.deepEqual( + yield* store.read({ afterSequence: highWater }).pipe(Stream.runCollect), + published, + ); + }).pipe(Effect.provide(Layer.fresh(layerTest)), Effect.scoped), +); + +it.effect.each( + (["none", "duplicate", "after-write"] as const).flatMap((failure) => + [1, 10].map((streams) => ({ failure, streams })), + ), +)( + "shares the writer and isolates failure=$failure across $streams independent consumers", + ({ failure, streams }) => { + const rollback = failure !== "none"; + const failedIndex = streams - 1; + const jobError = { _tag: "InvalidUpdate" } as const; + const writes: number[] = []; + const tracer = Tracer.make({ + span(options) { + const span = new Tracer.NativeSpan(options); + const end = span.end.bind(span); + span.end = (time, exit) => { + end(time, exit); + if (span.name === "orchestrationV2.EventSink.providerBatch" && Exit.isSuccess(exit)) { + writes.push(Number(span.attributes.get("orchestration_v2.batch_writes"))); + } + }; + return span; + }, + }); + return Effect.gen(function* () { + const sink = yield* EventSink.EventSinkV2; + const projections = yield* ProjectionStore.ProjectionStoreV2; + const now = yield* DateTime.now; + const threads = Array.from({ length: streams }, (_, index) => + makeThread(ThreadId.make(`shared:${index}`), now), + ); + yield* sink.write({ + events: threads.map((thread, index) => + threadCreatedEvent({ id: `created:${index}`, thread, now }), + ), + }); + const highWater = yield* sink.latestSequence(); + const store = yield* EventStore.EventStoreV2; + const publications = yield* Queue.unbounded(); + yield* sink.stream({ afterSequence: highWater }).pipe( + Stream.runForEach((stored) => Queue.offer(publications, stored.event.id)), + Effect.forkScoped({ startImmediately: true }), + ); + const pulled = yield* Deferred.make(); + let arrivals = 0; + const inputs = yield* Effect.forEach(threads, () => + Queue.unbounded(), + ); + const consumers = yield* Effect.forEach(inputs, (input, index) => + sink + .consumeProviderEvents( + Stream.fromQueue(input).pipe( + Stream.tap(() => + ++arrivals === streams ? Deferred.succeed(pulled, undefined) : Effect.void, + ), + ), + (raw) => + raw.type === "message.updated" + ? sink + .write({ + events: [ + { + id: EventId.make( + failure === "duplicate" && index === failedIndex + ? "created:0" + : `shared-event:${index}`, + ), + threadId: raw.message.threadId, + occurredAt: now, + type: "message.updated", + payload: raw.message, + }, + ], + }) + .pipe( + Effect.andThen( + failure === "after-write" && index === failedIndex + ? Effect.fail(jobError) + : Effect.void, + ), + ) + : Effect.void, + Effect.succeed(false), + Effect.succeed(true), + ) + .pipe(Effect.exit, Effect.forkScoped({ startImmediately: true })), + ); + for (let index = 0; index < threads.length; index++) { + yield* Queue.offer(inputs[index]!, { + type: "message.updated", + driver: ProviderDriverKind.make("codex"), + message: { + id: MessageId.make(`message:${index}`), + threadId: threads[index]!.id, + runId: null, + nodeId: null, + createdBy: "agent", + creationSource: "provider", + role: "assistant", + text: `stream:${index}`, + streaming: true, + attachments: [], + createdAt: now, + updatedAt: now, + }, + }); + } + yield* Deferred.await(pulled); + yield* TestClock.adjust(0); + yield* sink.write({ events: [] }); + if (rollback) { + const exit = yield* Fiber.join(consumers[failedIndex]!); + assert.isTrue(Exit.isFailure(exit)); + if (failure === "after-write" && Exit.isFailure(exit)) + assert.strictEqual(Cause.squash(exit.cause), jobError); + } + yield* TestClock.adjust(0); + const replay = yield* store.read({ afterSequence: highWater }).pipe(Stream.runCollect); + assert.deepEqual( + replay.length === 0 ? [] : yield* Queue.takeN(publications, replay.length), + replay.map((stored) => stored.event.id), + ); + assert.equal(yield* Queue.size(publications), 0); + assert.deepEqual(writes, rollback && streams === 1 ? [] : [streams - (rollback ? 1 : 0)]); + assert.equal(yield* sink.latestSequence(), highWater + streams - (rollback ? 1 : 0)); + for (let index = 0; index < threads.length; index++) { + assert.lengthOf( + (yield* projections.getThreadRecords(threads[index]!.id, ["messages"])).messages, + rollback && index === failedIndex ? 0 : 1, + ); + } + }).pipe( + Effect.scoped, + Effect.provide(Layer.fresh(layerTest)), + Effect.withTracer(tracer), + Effect.withTracerEnabled(true), + Effect.provideService( + ConfigProvider.ConfigProvider, + ConfigProvider.fromUnknown({ T3CODE_PROVIDER_COMMIT_BATCHING: true }), + ), + ); + }, +); + it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { it.effect("projects oversized tool bodies once before replay and live RPC retention", () => Effect.scoped( @@ -3288,87 +3526,101 @@ it.layer(layerTest)("orchestration V2 foundation persistence", (it) => { ); }); -it.effect("publishes live events in commit order across concurrent writers", () => - Effect.gen(function* () { - const firstCommitted = yield* Deferred.make(); - const releaseFirst = yield* Deferred.make(); - const layerPausingOutbox = Layer.effect( - EffectOutbox.EffectOutboxV2, - Effect.gen(function* () { - const delegate = yield* EffectOutbox.EffectOutboxV2; - return EffectOutbox.EffectOutboxV2.of({ - ...delegate, - notifyAvailable: (count) => - Deferred.succeed(firstCommitted, undefined).pipe( - Effect.andThen(Deferred.await(releaseFirst)), - Effect.andThen(delegate.notifyAvailable(count)), - ), - }); - }), - ).pipe(Layer.provide(layerEffectOutboxProvided)); - const layerEventSink = EventSink.layerFromStores.pipe( - Layer.provide( - Layer.mergeAll( - layerStoresProvided, - layerPausingOutbox, - layerCommandReceiptStoreProvided, - ProjectStore.layer.pipe(Layer.provide(layerDatabase)), - TurnItemPositionStore.layer.pipe(Layer.provide(layerDatabase)), +it.effect.each(["direct", "provider-batch", "cancel-provider-batch"] as const)( + "publishes live events in commit order across concurrent writers with %s", + (mode) => + Effect.gen(function* () { + const firstCommitted = yield* Deferred.make(); + const releaseFirst = yield* Deferred.make(); + const layerPausingOutbox = Layer.effect( + EffectOutbox.EffectOutboxV2, + Effect.gen(function* () { + const delegate = yield* EffectOutbox.EffectOutboxV2; + return EffectOutbox.EffectOutboxV2.of({ + ...delegate, + notifyAvailable: (count) => + Deferred.succeed(firstCommitted, undefined).pipe( + Effect.andThen(Deferred.await(releaseFirst)), + Effect.andThen(delegate.notifyAvailable(count)), + ), + }); + }), + ).pipe(Layer.provide(layerEffectOutboxProvided)); + const layerEventSink = EventSink.layerFromStores.pipe( + Layer.provide( + Layer.mergeAll( + layerStoresProvided, + layerPausingOutbox, + layerCommandReceiptStoreProvided, + ProjectStore.layer.pipe(Layer.provide(layerDatabase)), + TurnItemPositionStore.layer.pipe(Layer.provide(layerDatabase)), + ), ), - ), - ); + ); - yield* Effect.gen(function* () { - const eventSink = yield* EventSink.EventSinkV2; - const now = yield* DateTime.now; - const first = makeThread(ThreadId.make("thread:foundation-publish-order:first"), now); - const second = makeThread(ThreadId.make("thread:foundation-publish-order:second"), now); - const published = yield* eventSink - .stream({ afterSequence: yield* eventSink.latestSequence() }) - .pipe(Stream.take(2), Stream.runCollect, Effect.forkScoped({ startImmediately: true })); - - const firstWrite = yield* eventSink - .writeWithEffects({ - events: [ - threadCreatedEvent({ id: "event:foundation-publish-order:first", thread: first, now }), - ], - effects: [ - { - id: "effect:foundation-publish-order:first", - commandId: CommandId.make("command:foundation-publish-order:first"), - threadId: first.id, - request: { type: "terminal.cleanup" }, - }, - ], - }) - .pipe(Effect.forkScoped); - yield* Deferred.await(firstCommitted); - const secondWrite = yield* eventSink - .write({ - events: [ - threadCreatedEvent({ - id: "event:foundation-publish-order:second", - thread: second, - now, - }), - ], - }) - .pipe( - Effect.provideService(Scheduler.MaxOpsBeforeYield, Number.POSITIVE_INFINITY), - Effect.forkScoped, - ); - yield* Effect.yieldNow; - yield* Deferred.succeed(releaseFirst, undefined); - yield* Fiber.join(firstWrite); - yield* Fiber.join(secondWrite); + yield* Effect.gen(function* () { + const eventSink = yield* EventSink.EventSinkV2; + const now = yield* DateTime.now; + const first = makeThread(ThreadId.make("thread:foundation-publish-order:first"), now); + const second = makeThread(ThreadId.make("thread:foundation-publish-order:second"), now); + const published = yield* eventSink + .stream({ afterSequence: yield* eventSink.latestSequence() }) + .pipe(Stream.take(2), Stream.runCollect, Effect.forkScoped({ startImmediately: true })); - const sequences = Array.from(yield* Fiber.join(published), (stored) => stored.sequence); - assert.deepEqual( - sequences, - [...sequences].sort((left, right) => left - right), - ); - }).pipe(Effect.provide(layerEventSink)); - }).pipe(Effect.provide(layerDatabase)), + const firstWrite = yield* eventSink + .writeWithEffects({ + events: [ + threadCreatedEvent({ + id: "event:foundation-publish-order:first", + thread: first, + now, + }), + ], + effects: [ + { + id: "effect:foundation-publish-order:first", + commandId: CommandId.make("command:foundation-publish-order:first"), + threadId: first.id, + request: { type: "terminal.cleanup" }, + }, + ], + }) + .pipe( + (write) => (mode === "direct" ? write : eventSink.withProviderBatch(write)), + Effect.forkScoped, + ); + yield* Deferred.await(firstCommitted); + const secondWrite = yield* eventSink + .write({ + events: [ + threadCreatedEvent({ + id: "event:foundation-publish-order:second", + thread: second, + now, + }), + ], + }) + .pipe( + Effect.provideService(Scheduler.MaxOpsBeforeYield, Number.POSITIVE_INFINITY), + Effect.forkScoped, + ); + yield* Effect.yieldNow; + const interrupted = + mode === "cancel-provider-batch" + ? yield* Fiber.interrupt(firstWrite).pipe(Effect.forkScoped({ startImmediately: true })) + : undefined; + yield* Deferred.succeed(releaseFirst, undefined); + if (interrupted !== undefined) yield* Fiber.join(interrupted); + else yield* Fiber.join(firstWrite); + yield* Fiber.join(secondWrite); + + const sequences = Array.from(yield* Fiber.join(published), (stored) => stored.sequence); + assert.deepEqual( + sequences, + [...sequences].sort((left, right) => left - right), + ); + }).pipe(Effect.provide(layerEventSink)); + }).pipe(Effect.provide(layerDatabase)), ); it.effect.each(["events", "receipts"] as const)( diff --git a/apps/server/src/orchestration-v2/ProviderCommitBatch.test.ts b/apps/server/src/orchestration-v2/ProviderCommitBatch.test.ts new file mode 100644 index 0000000000000..6655d1554533f --- /dev/null +++ b/apps/server/src/orchestration-v2/ProviderCommitBatch.test.ts @@ -0,0 +1,472 @@ +import { assert, it } from "@effect/vitest"; +import { MessageId, ProviderDriverKind, ProviderTurnId, ThreadId } from "@t3tools/contracts"; +import * as KeyedLock from "@t3tools/shared/KeyedLock"; +import * as ProviderAdapter from "@t3tools/provider-core/server/ProviderAdapter"; +import * as Cause from "effect/Cause"; +import * as DateTime from "effect/DateTime"; +import * as Deferred from "effect/Deferred"; +import * as Effect from "effect/Effect"; +import * as Exit from "effect/Exit"; +import * as Fiber from "effect/Fiber"; +import * as Queue from "effect/Queue"; +import * as Stream from "effect/Stream"; +import * as TestClock from "effect/testing/TestClock"; +import * as ProviderCommitBatch from "./ProviderCommitBatch.ts"; + +const driver = ProviderDriverKind.make("codex"); +function event(index: number, text = `delta:${index}`): ProviderAdapter.ProviderAdapterV2Event { + const threadId = ThreadId.make("single-thread"); + const now = DateTime.makeUnsafe("2026-11-03T00:00:00Z"); + return { + type: "message.updated", + driver, + message: { + id: MessageId.make("single-message"), + threadId, + runId: null, + nodeId: null, + createdBy: "agent", + creationSource: "provider", + role: "assistant", + text, + attachments: [], + streaming: true, + createdAt: now, + updatedAt: now, + }, + }; +} +const control: ProviderAdapter.ProviderAdapterV2Event = { + type: "background_work.updated", + driver, + providerTurnId: ProviderTurnId.make("turn"), + pending: false, +}; + +const recordBatches = Effect.fn("test.recordProviderBatches")(function* ( + beforeCommit: Effect.Effect = Effect.void, +) { + const output = yield* Queue.unbounded>(); + let current: ProviderAdapter.ProviderAdapterV2Event[] = []; + const coordinator = yield* ProviderCommitBatch.make( + (work) => + Effect.gen(function* () { + yield* beforeCommit; + current = []; + yield* work; + if (current.length > 0) yield* Queue.offer(output, current); + }), + true, + ); + const process = (raw: ProviderAdapter.ProviderAdapterV2Event) => + raw.type === "background_work.updated" + ? Queue.offer(output, [raw]).pipe(Effect.asVoid) + : Effect.sync(() => { + current.push(raw); + }); + return { ...coordinator, process, output }; +}); + +it.effect("publishes preceding snapshots before a control", () => + Effect.gen(function* () { + const recorder = yield* recordBatches(); + const input = yield* Queue.unbounded(); + const consumer = yield* recorder + .consume(Stream.fromQueue(input), recorder.process, Effect.succeed(true)) + .pipe(Effect.forkScoped); + yield* Queue.offer(input, event(1)); + assert.deepEqual(yield* Queue.take(recorder.output), [event(1)]); + yield* Queue.offer(input, event(2)); + yield* Queue.offer(input, control); + assert.deepEqual(yield* Queue.take(recorder.output), [event(2)]); + assert.deepEqual(yield* Queue.take(recorder.output), [control]); + yield* Fiber.join(consumer); + }).pipe(Effect.scoped), +); + +it.effect("publishes isolated snapshots without advancing the clock", () => + Effect.gen(function* () { + const recorder = yield* recordBatches(); + const input = yield* Queue.unbounded(); + const pulled = yield* Queue.unbounded(); + yield* recorder + .consume( + Stream.fromQueue(input).pipe(Stream.tap((raw) => Queue.offer(pulled, raw))), + recorder.process, + Effect.succeed(false), + ) + .pipe(Effect.forkScoped); + for (const raw of [event(1), event(2)]) { + yield* Queue.offer(input, raw); + yield* Queue.take(pulled); + yield* TestClock.adjust(0); + const published = yield* Queue.poll(recorder.output); + assert.equal(published._tag, "Some"); + if (published._tag === "Some") assert.deepEqual(published.value, [raw]); + } + }).pipe(Effect.scoped), +); + +it.effect("bounds transactions by count and bytes, including the final partial group", () => + Effect.gen(function* () { + const workloads = [ + { events: Array.from({ length: 129 }, (_, index) => event(index)) }, + { events: [event(0), event(1, "x".repeat(1024 * 1024)), event(2)] }, + { + events: [event(0, "x".repeat(600_000)), event(1, "y".repeat(600_000)), event(2)], + }, + ]; + for (const { events } of workloads) { + const recorder = yield* recordBatches(); + const pulled = yield* Deferred.make(); + let arrivals = 0; + const consumer = yield* recorder + .consume( + Stream.fromArray(events).pipe( + Stream.tap(() => + ++arrivals === events.length ? Deferred.succeed(pulled, undefined) : Effect.void, + ), + ), + recorder.process, + Effect.succeed(false), + ) + .pipe(Effect.forkScoped({ startImmediately: true })); + yield* Deferred.await(pulled); + yield* TestClock.adjust(0); + yield* recorder.barrier(Effect.void); + yield* Fiber.join(consumer); + const batches = yield* Queue.takeAll(recorder.output); + assert.isTrue(batches.every((batch) => batch.length <= 64)); + if (events.length === 129) assert.isTrue(batches.some((batch) => batch.length > 1)); + assert.isTrue( + batches.every( + (batch) => + batch.length === 1 || + batch.reduce((bytes, raw) => bytes + Buffer.byteLength(JSON.stringify(raw)), 0) <= + 1024 * 1024, + ), + ); + assert.deepEqual(batches.flat(), events); + } + }).pipe(Effect.scoped), +); + +it.effect("keeps updates immediate and ordered when batching is disabled", () => + Effect.gen(function* () { + const seen: ProviderAdapter.ProviderAdapterV2Event[] = []; + const coordinator = yield* ProviderCommitBatch.make(() => Effect.die("must not batch"), false); + yield* coordinator.consume( + Stream.make(event(1), control, event(2)), + (raw) => + Effect.sync(() => { + seen.push(raw); + }), + Effect.succeed(false), + ); + assert.deepEqual(seen, [event(1), control, event(2)]); + }).pipe(Effect.scoped), +); + +it.effect.each([ + { streams: 1, updates: 10 }, + { streams: 10, updates: 1 }, + { streams: 10, updates: 10 }, +])( + "shares bounded transactions across $streams streams with $updates updates each", + ({ streams, updates }) => + Effect.gen(function* () { + const coordinator = yield* recordBatches(); + const inputs = yield* Effect.forEach(Array.from({ length: streams }), () => + Queue.unbounded(), + ); + const pulled = yield* Deferred.make(); + let arrivals = 0; + const consumers = yield* Effect.forEach(inputs, (input) => + coordinator + .consume( + Stream.fromQueue(input).pipe( + Stream.tap(() => { + arrivals++; + return arrivals === streams * updates + ? Deferred.succeed(pulled, undefined) + : Effect.void; + }), + ), + coordinator.process, + Effect.succeed(true), + ) + .pipe(Effect.forkScoped({ startImmediately: true })), + ); + for (let n = 0; n < updates; n++) { + for (let stream = 0; stream < streams; stream++) + yield* Queue.offer(inputs[stream]!, event(stream * updates + n)); + } + yield* Deferred.await(pulled); + yield* TestClock.adjust(0); + + yield* coordinator.barrier(Effect.void); + const batches: Array> = []; + let count = 0; + while (count < streams * updates) { + const batch = yield* Queue.take(coordinator.output); + batches.push(batch); + count += batch.length; + } + assert.deepEqual( + batches.map((batch) => batch.length), + streams * updates > 64 ? [64, 36] : [10], + ); + assert.equal( + new Set( + batches + .flat() + .flatMap((raw) => (raw.type === "message.updated" ? [raw.message.text] : [])), + ).size, + streams * updates, + ); + for (const input of inputs) yield* Queue.offer(input, control); + yield* Effect.forEach(consumers, Fiber.join); + }).pipe(Effect.scoped), +); + +it.effect("does not hold the shared worker while one consumer checks background work", () => + Effect.gen(function* () { + const checking = yield* Deferred.make(); + const release = yield* Deferred.make(); + const coordinator = yield* ProviderCommitBatch.make((work) => work, true); + const consumer = yield* coordinator + .consume( + Stream.make(control), + () => Effect.void, + Deferred.succeed(checking, undefined).pipe( + Effect.andThen(Deferred.await(release)), + Effect.as(true), + ), + ) + .pipe(Effect.forkScoped({ startImmediately: true })); + yield* Deferred.await(checking); + assert.equal( + yield* coordinator.barrier(Effect.succeed("another run can commit")), + "another run can commit", + ); + yield* Deferred.succeed(release, undefined); + yield* Fiber.join(consumer); + }).pipe(Effect.scoped), +); + +it.effect.each(["job", "commit"] as const)( + "isolates a %s failure at its transaction boundary", + (mode) => + Effect.gen(function* () { + const coordinator = yield* ProviderCommitBatch.make( + (work) => + mode === "commit" + ? work.pipe( + Effect.andThen( + Effect.fail(new ProviderCommitBatch.CommitBatchError({ cause: "COMMIT failed" })), + ), + ) + : work, + true, + ); + const inputs = yield* Effect.forEach([0, 1], () => + Queue.unbounded(), + ); + const consumers = yield* Effect.forEach(inputs, (input, index) => + coordinator + .consume( + Stream.fromQueue(input), + () => (mode === "job" && index === 1 ? Effect.fail("invalid update") : Effect.void), + Effect.succeed(true), + ) + .pipe(Effect.exit, Effect.forkScoped({ startImmediately: true })), + ); + for (const input of inputs) yield* Queue.offer(input, event(0)); + yield* TestClock.adjust(0); + yield* coordinator.barrier(Effect.void); + assert.isTrue(Exit.isFailure(yield* Fiber.join(consumers[1]!))); + if (mode === "commit") assert.isTrue(Exit.isFailure(yield* Fiber.join(consumers[0]!))); + else { + assert.isUndefined(consumers[0]!.pollUnsafe()); + yield* Queue.offer(inputs[0]!, control); + assert.isTrue(Exit.isSuccess(yield* Fiber.join(consumers[0]!))); + } + assert.equal(yield* coordinator.barrier(Effect.succeed("recovered")), "recovered"); + assert.deepEqual( + yield* coordinator.barrier(Effect.fail("original error")).pipe(Effect.exit), + Exit.fail("original error"), + ); + }).pipe(Effect.scoped), +); + +it.effect("keeps a transaction failure when rolling back one failed snapshot", () => + Effect.gen(function* () { + const transactionError = new ProviderCommitBatch.CommitBatchError({ cause: "ROLLBACK failed" }); + const coordinator = yield* ProviderCommitBatch.make( + (work) => work.pipe(Effect.mapError(() => transactionError)), + true, + ); + const exit = yield* coordinator + .consume( + Stream.concat(Stream.make(event(1)), Stream.never), + () => Effect.fail("invalid snapshot"), + Effect.succeed(false), + ) + .pipe(Effect.exit); + assert.isTrue(Exit.isFailure(exit)); + if (Exit.isFailure(exit)) assert.strictEqual(Cause.squash(exit.cause), transactionError); + }).pipe(Effect.scoped), +); + +it.effect( + "allows a command holding the thread lock to commit while a provider control waits for it", + () => + Effect.gen(function* () { + const coordinator = yield* ProviderCommitBatch.make((work) => work, true); + const locks = yield* KeyedLock.make(); + const threadId = ThreadId.make("locked-thread"); + const held = yield* Deferred.make(); + const waiting = yield* Deferred.make(); + const committed = yield* Deferred.make(); + const command = yield* locks + .withLock( + threadId, + Deferred.succeed(held, undefined).pipe( + Effect.andThen(Deferred.await(waiting)), + Effect.andThen(coordinator.barrier(Deferred.succeed(committed, undefined))), + ), + ) + .pipe(Effect.forkScoped({ startImmediately: true })); + yield* Deferred.await(held); + const consumer = yield* coordinator + .consume( + Stream.make(control), + () => + Deferred.succeed(waiting, undefined).pipe( + Effect.andThen(locks.withLock(threadId, Effect.void)), + ), + Effect.succeed(true), + ) + .pipe(Effect.forkScoped({ startImmediately: true })); + yield* Deferred.await(waiting); + yield* TestClock.adjust(0); + assert.isTrue(yield* Deferred.isDone(committed)); + yield* Fiber.join(command); + yield* Fiber.join(consumer); + }).pipe(Effect.scoped), +); + +it.effect.each([false, true])( + "batches work accumulated behind a busy writer with cancellation=%s", + (cancel) => + Effect.gen(function* () { + const started = yield* Deferred.make(); + const release = yield* Deferred.make(); + let first = true; + const recorder = yield* recordBatches( + Effect.gen(function* () { + if (!first) return; + first = false; + yield* Deferred.succeed(started, undefined); + yield* Deferred.await(release); + }), + ); + const leading = yield* recorder + .consume(Stream.make(event(-1)), recorder.process, Effect.succeed(false)) + .pipe(Effect.forkScoped({ startImmediately: true })); + yield* Deferred.await(started); + const pulled = yield* Deferred.make(); + let arrivals = 0; + const consumers = yield* Effect.forEach(Array.from({ length: 10 }), (_, index) => + recorder + .consume( + Stream.make(event(index)).pipe( + Stream.tap(() => + ++arrivals === 10 ? Deferred.succeed(pulled, undefined) : Effect.void, + ), + ), + recorder.process, + Effect.succeed(false), + ) + .pipe(Effect.forkScoped({ startImmediately: true })), + ); + yield* Deferred.await(pulled); + yield* TestClock.adjust(0); + if (cancel) yield* Fiber.interrupt(consumers[0]!); + assert.equal(yield* Queue.size(recorder.output), 0); + yield* Deferred.succeed(release, undefined); + yield* Fiber.join(leading); + yield* Effect.forEach(consumers.slice(cancel ? 1 : 0), Fiber.join); + assert.deepEqual(yield* Queue.takeAll(recorder.output), [ + [event(-1)], + Array.from({ length: cancel ? 9 : 10 }, (_, index) => event(index + (cancel ? 1 : 0))), + ]); + }).pipe(Effect.scoped), +); + +it.effect.each(["failure", "stop"] as const)( + "flushes the prefix and respects a source %s barrier", + (mode) => + Effect.gen(function* () { + const seen: ProviderAdapter.ProviderAdapterV2Event[] = []; + const coordinator = yield* ProviderCommitBatch.make((work) => work, true); + const source = + mode === "failure" + ? Stream.concat(Stream.make(event(1), event(2)), Stream.fail("provider failed")) + : Stream.make(event(1), control, event(2)); + const exit = yield* coordinator + .consume( + source, + (raw) => + Effect.sync(() => { + seen.push(raw); + }), + Effect.succeed(mode === "stop"), + ) + .pipe(Effect.exit); + assert.deepEqual(seen, mode === "failure" ? [event(1), event(2)] : [event(1), control]); + assert.equal(Exit.isFailure(exit), mode === "failure"); + assert.equal( + yield* coordinator.barrier(Effect.succeed("still available")), + "still available", + ); + }).pipe(Effect.scoped), +); + +it.effect("does not force a separate commit for every consumer that ends", () => + Effect.gen(function* () { + const pulled = yield* Deferred.make(); + let arrivals = 0; + const coordinator = yield* recordBatches(); + for (let index = 0; index < 10; index++) { + yield* coordinator + .consume( + Stream.make(event(index)).pipe( + Stream.tap(() => + ++arrivals === 10 ? Deferred.succeed(pulled, undefined) : Effect.void, + ), + ), + coordinator.process, + Effect.succeed(false), + ) + .pipe(Effect.forkScoped({ startImmediately: true })); + } + yield* Deferred.await(pulled); + yield* TestClock.adjust(0); + yield* coordinator.barrier(Effect.void); + assert.deepEqual( + (yield* Queue.takeAll(coordinator.output)).map((batch) => batch.length), + [10], + ); + }).pipe(Effect.scoped), +); + +it.effect("settles calls made after the coordinator's scope closes", () => + Effect.gen(function* () { + const coordinator = yield* ProviderCommitBatch.make((work) => work, true).pipe(Effect.scoped); + const exit = yield* coordinator + .barrier(Effect.die("must not execute after shutdown")) + .pipe(Effect.exit); + assert.isTrue(Exit.isFailure(exit) && Cause.hasInterruptsOnly(exit.cause)); + }), +); diff --git a/apps/server/src/orchestration-v2/ProviderCommitBatch.ts b/apps/server/src/orchestration-v2/ProviderCommitBatch.ts new file mode 100644 index 0000000000000..60cee2b632f66 --- /dev/null +++ b/apps/server/src/orchestration-v2/ProviderCommitBatch.ts @@ -0,0 +1,303 @@ +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"; +import * as Effect from "effect/Effect"; +import * as Context from "effect/Context"; +import * as Deferred from "effect/Deferred"; +import * as Exit from "effect/Exit"; +import * as Queue from "effect/Queue"; +import * as Schema from "effect/Schema"; +import * as Stream from "effect/Stream"; + +const MAX_EVENTS = 64; +const MAX_BYTES = 1024 * 1024; +const encodeQueue = Schema.encodeEffect(Schema.fromJsonString(Schema.Unknown)); + +export const enabled = Config.Boolean("T3CODE_PROVIDER_COMMIT_BATCHING").pipe( + Config.withDefault(true), + Effect.orDie, +); + +type Event = ProviderAdapter.ProviderAdapterV2Event; +const isSnapshot = (event: Event) => + event.type === "node.updated" || + event.type === "message.updated" || + event.type === "turn_item.updated" || + event.type === "plan.updated"; + +export class CommitBatchError extends Schema.TaggedError()("CommitBatchError", { + cause: Schema.Defect(), +}) {} +export const isCommitBatchError = Schema.is(CommitBatchError); + +interface Job { + readonly owner: { active: boolean }; + readonly queuedAt: number; + readonly snapshot: boolean; + readonly bytes: number; + readonly execute: Effect.Effect; + readonly settle: ( + exit: Exit.Exit, + preserveFailure: boolean, + ) => Effect.Effect; + readonly cancel: Effect.Effect; +} +interface Batch { + readonly events: ReadonlyArray; + readonly bytes: number; + readonly reason: "ready" | "limit" | "control"; +} + +function collect(queue: Queue.Queue) { + return Stream.fromEffectRepeat(Queue.takeBetween(queue, 1, MAX_EVENTS)).pipe( + Stream.flatMap((ready) => { + const batches: Array = []; + let events: Array = []; + let bytes = 0; + const flush = (reason: Batch["reason"]) => { + if (events.length > 0) batches.push({ events, bytes, reason }); + events = []; + bytes = 0; + }; + for (const job of ready) { + if (!job.snapshot) { + flush("control"); + batches.push({ events: [job], bytes: 0, reason: "control" }); + } else { + if (events.length > 0 && bytes + job.bytes > MAX_BYTES) flush("limit"); + events.push(job); + bytes += job.bytes; + if (events.length >= MAX_EVENTS || bytes >= MAX_BYTES) flush("limit"); + } + } + flush("ready"); + return Stream.fromArray(batches); + }), + ); +} + +const CurrentWorker = Context.Reference( + "t3/ProviderCommitBatch/CurrentWorker", + { + defaultValue: () => undefined, + }, +); + +export const make = Effect.fn("ProviderCommitBatch.make")(function* ( + commit: (work: Effect.Effect) => Effect.Effect, + enabled: boolean, + isolate: ( + work: Effect.Effect, + ) => Effect.Effect = (work) => work, +) { + const queue = yield* Queue.bounded(MAX_EVENTS); + const pending = new Set(); + let closed = false; + const shutdown = Effect.gen(function* () { + closed = true; + yield* Queue.shutdown(queue); + for (const job of pending) yield* job.cancel; + }); + yield* Effect.addFinalizer(() => shutdown); + // These helpers run once per raw event; the shared transaction owns the tracing span. + const enqueue = Effect.fnUntraced(function* ( + effect: Effect.Effect, + token: { active: boolean }, + snapshot: boolean, + bytes: number, + ) { + const caller = yield* Effect.context(); + const reply = yield* Deferred.make(); + let ownFailure: Cause.Cause | undefined; + let result: A; + const job: Job = { + owner: token, + queuedAt: yield* Clock.currentTimeMillis, + snapshot, + bytes, + execute: effect.pipe( + Effect.tapCause((cause) => + Effect.sync(() => { + ownFailure = cause; + }), + ), + Effect.mapError((cause) => new CommitBatchError({ cause })), + Effect.provideContext(caller), + Effect.tap((value) => + Effect.sync(() => { + result = value; + }), + ), + Effect.asVoid, + ), + settle: Effect.fnUntraced(function* ( + exit: Exit.Exit, + preserveFailure: boolean, + ) { + pending.delete(job); + if (Exit.isFailure(exit)) { + token.active = false; + yield* Deferred.failCause( + reply, + preserveFailure && ownFailure !== undefined ? ownFailure : exit.cause, + ); + } else yield* Deferred.succeed(reply, result); + }), + cancel: Effect.gen(function* () { + pending.delete(job); + yield* Deferred.interrupt(reply); + }), + }; + yield* Effect.uninterruptibleMask((restore) => + Effect.gen(function* () { + pending.add(job); + const accepted = yield* restore(Queue.offer(queue, job)).pipe( + Effect.onError(() => job.cancel), + ); + if (!accepted) yield* job.cancel; + }), + ); + return Deferred.await(reply); + }); + if (enabled) { + yield* collect(queue).pipe( + Stream.runForEach((batch) => + Effect.uninterruptibleMask((restore) => + Effect.gen(function* () { + const jobs = batch.events.filter((job) => job.owner.active); + for (const job of batch.events) if (!job.owner.active) yield* job.cancel; + if (jobs.length === 0) return; + yield* Effect.annotateCurrentSpan({ + "orchestration_v2.work_item_count": jobs.length, + "orchestration_v2.consumer_count": new Set(jobs.map((job) => job.owner)).size, + "orchestration_v2.raw_json_bytes": batch.bytes, + "orchestration_v2.batch_flush_reason": batch.reason, + "orchestration_v2.backlog_batching": true, + "orchestration_v2.batch_wait_ms": + (yield* Clock.currentTimeMillis) - Math.min(...jobs.map((job) => job.queuedAt)), + }); + const outcomes = new Map>(); + const work = Effect.forEach( + jobs, + (job) => + Effect.gen(function* () { + if (!job.owner.active) return yield* job.cancel; + const outcome = yield* ( + jobs.length > 1 && job.bytes > 0 ? isolate(job.execute) : job.execute + ).pipe(Effect.exit); + outcomes.set(job, outcome); + if (Exit.isFailure(outcome)) { + job.owner.active = false; + if (jobs.length === 1) return yield* Effect.failCause(outcome.cause); + } + }), + { discard: true }, + ); + const exit = yield* restore( + (jobs.every((job) => job.bytes === 0) ? work : commit(work)).pipe( + Effect.provideService(CurrentWorker, yield* Effect.fiberId), + ), + ).pipe(Effect.exit); + for (const job of jobs) { + const outcome = outcomes.get(job) ?? Exit.void; + yield* job.settle( + Exit.isSuccess(exit) ? outcome : exit, + Exit.isSuccess(exit) || + (Exit.isFailure(exit) && + Exit.isFailure(outcome) && + Cause.squash(exit.cause) === Cause.squash(outcome.cause)), + ); + } + }), + ).pipe(Effect.withSpan("orchestrationV2.ProviderCommitBatch.flush")), + ), + Effect.ensuring(shutdown), + Effect.forkScoped, + ); + } + const barrier = Effect.fnUntraced(function* (effect: Effect.Effect) { + if (closed) return yield* Effect.interrupt; + if (!enabled || pending.size === 0 || (yield* CurrentWorker) === (yield* Effect.fiberId)) + return yield* effect; + const token = { active: true }; + return yield* enqueue(Effect.exit(effect), token, false, 0).pipe( + Effect.flatten, + Effect.orDie, + Effect.flatten, + Effect.ensuring( + Effect.sync(() => { + token.active = false; + }), + ), + ); + }); + const consume = Effect.fnUntraced(function* ( + source: Stream.Stream, + process: (event: Event) => Effect.Effect, + after: Effect.Effect, + canBatch: Effect.Effect = Effect.succeed(true), + ) { + if (!enabled) + return yield* source.pipe( + Stream.mapEffect((event) => process(event).pipe(Effect.andThen(after))), + Stream.takeUntil(Boolean), + Stream.runDrain, + ); + const token = { active: true }; + const signals = source.pipe( + Stream.map((event) => ({ _tag: "Event" as const, event })), + Stream.catchCause((cause) => Stream.succeed({ _tag: "Failure" as const, cause })), + Stream.concat(Stream.succeed({ _tag: "End" as const })), + ); + let resume: Effect.Effect = Effect.void; + yield* signals.pipe( + Stream.mapEffect((signal) => + Effect.gen(function* () { + yield* resume; + const processed = yield* Deferred.make(); + const event = signal._tag === "Event" ? signal.event : undefined; + // EOF only acknowledges prior work; it can share a commit with other consumers. + const snapshot = + signal._tag === "End" || + (event !== undefined && isSnapshot(event) && (yield* canBatch)); + const bytes = + snapshot && event !== undefined ? Buffer.byteLength(yield* encodeQueue(event)) : 0; + const work: Effect.Effect = + signal._tag === "Failure" + ? Effect.failCause(signal.cause) + : event === undefined + ? Effect.succeed(true) + : process(event).pipe(Effect.as(false)); + const reply = snapshot + ? yield* enqueue(work, token, true, bytes) + : (yield* enqueue(Effect.void, token, false, 0)).pipe(Effect.andThen(work)); + if (!snapshot) resume = Deferred.await(processed); + return { reply, processed, check: !snapshot }; + }), + ), + Stream.buffer({ capacity: MAX_EVENTS }), + Stream.mapEffect(({ reply, processed, check }) => + reply.pipe( + Effect.flatMap((end) => + end ? Effect.succeed(true) : check ? after : Effect.succeed(false), + ), + Effect.tap((stop) => + Effect.sync(() => { + if (stop) token.active = false; + }), + ), + Effect.ensuring(Deferred.succeed(processed, undefined)), + ), + ), + Stream.takeUntil(Boolean), + Stream.runDrain, + Effect.ensuring( + Effect.sync(() => { + token.active = false; + }), + ), + ); + }); + return { barrier, consume }; +}); diff --git a/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts b/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts index d35771ac3d0c7..b6adf6371fb60 100644 --- a/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts +++ b/apps/server/src/orchestration-v2/ProviderEventIngestor.test.ts @@ -546,8 +546,12 @@ layer("ProviderEventIngestorV2", (it) => { event: { type: "node.updated", driver: CODEX_DRIVER, node: sent }, }); - assert.lengthOf(yield* ingestNode(node), 1); - assert.isEmpty(yield* ingestNode({ ...node })); + yield* eventSink.withProviderBatch( + Effect.gen(function* () { + assert.lengthOf(yield* ingestNode(node), 1); + assert.isEmpty(yield* ingestNode({ ...node })); + }), + ); yield* eventSink.write({ events: [ { diff --git a/apps/server/src/orchestration-v2/RunExecutionService.test.ts b/apps/server/src/orchestration-v2/RunExecutionService.test.ts index f3f1542146c66..5bfb4fe34dc6a 100644 --- a/apps/server/src/orchestration-v2/RunExecutionService.test.ts +++ b/apps/server/src/orchestration-v2/RunExecutionService.test.ts @@ -55,6 +55,17 @@ import * as RunExecutionService from "./RunExecutionService.ts"; import * as RunFinalizationService from "./RunFinalizationService.ts"; import * as TurnItemPositionStore from "./TurnItemPositionStore.ts"; +const mockEventSink = (implementation: Partial) => + Layer.mock(EventSink.EventSinkV2)({ + consumeProviderEvents: (source, process, after) => + source.pipe( + Stream.mapEffect((event) => process(event).pipe(Effect.andThen(after))), + Stream.takeUntil(Boolean), + Stream.runDrain, + ), + ...implementation, + }); + const driver = ProviderDriverKind.make("codex"); const rootAppThread = (id: ThreadId) => ({ @@ -73,7 +84,7 @@ const layerRunExecutionTest = RunExecutionService.layer.pipe( PreviewAgentActivity.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ captureBaseline: () => Effect.void }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({}), + mockEventSink({}), IdAllocator.layer, Layer.mock(ProviderEventIngestor.ProviderEventIngestorV2)({ ingestNormalized: () => Effect.succeed([]), @@ -900,7 +911,7 @@ it.effect("passes the thread's MCP app context to the provider under a safe key" captureBaseline: () => Effect.void, }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({}), + mockEventSink({}), IdAllocator.layer, Layer.mock(ProviderEventIngestor.ProviderEventIngestorV2)({ ingestNormalized: () => Effect.succeed([]), @@ -929,7 +940,7 @@ it.effect("fails the run when its ownership check cannot be read before calling PreviewAgentActivity.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ captureBaseline: () => Effect.void }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ writeIfRunCurrent: (input) => Ref.update(writes, (current) => [...current, ...input.events]).pipe( Effect.as({ committed: true, storedEvents: [] }), @@ -1140,7 +1151,7 @@ it.effect("starts the provider when checkpoint baseline capture fails", () => ), }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ writeWithEffects: (input) => Ref.update(writes, (current) => [ ...current, @@ -1254,7 +1265,7 @@ it.effect.each(["failure", "interruption", "stale-attempt", "start-guard"] as co scenario === "start-guard" ? Effect.void : Effect.die("not reached"), }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ writeIfRunCurrent: (input) => Effect.gen(function* () { assert.equal(input.threadId, threadId); @@ -1434,7 +1445,7 @@ it.effect("ingests new member threads of an owned child after the root turn term PreviewAgentActivity.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ captureBaseline: () => Effect.void }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ write: () => Effect.succeed([]), writeWithEffects: (input) => Effect.gen(function* () { @@ -1694,7 +1705,7 @@ it.effect("ends a nested native subagent and its thread when the root run is int PreviewAgentActivity.layer, Layer.mock(CheckpointService.CheckpointServiceV2)({ captureBaseline: () => Effect.void }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ write: () => Effect.succeed([]), writeWithEffects: (input) => Effect.as( @@ -2107,7 +2118,7 @@ it.effect( captureBaseline: () => Effect.void, }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ write: () => Effect.succeed([]), writeWithEffects: (input) => Effect.gen(function* () { @@ -3193,7 +3204,7 @@ function captureRootRunTermination(input: { }), Layer.mock(CheckpointService.CheckpointServiceV2)({ captureBaseline: () => Effect.void }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ write: (payload) => Effect.gen(function* () { for (const event of payload.events) { @@ -3715,7 +3726,7 @@ function runBackgroundItemScenario( Layer.succeed(PreviewAgentActivity.PreviewAgentActivity, previewActivity), Layer.mock(CheckpointService.CheckpointServiceV2)({ captureBaseline: () => Effect.void }), layerTurnItemPositionsTest, - Layer.mock(EventSink.EventSinkV2)({ + mockEventSink({ write: () => Effect.succeed([]), writeWithEffects: (input) => Effect.gen(function* () { diff --git a/apps/server/src/orchestration-v2/RunExecutionService.ts b/apps/server/src/orchestration-v2/RunExecutionService.ts index 6a52fac447926..304eeaba43421 100644 --- a/apps/server/src/orchestration-v2/RunExecutionService.ts +++ b/apps/server/src/orchestration-v2/RunExecutionService.ts @@ -35,7 +35,6 @@ import * as Layer from "effect/Layer"; 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"; @@ -1322,350 +1321,350 @@ export const layer: Layer.Layer< // Ends once the root turn and the children it drives settle. Background commands such // as a dev server keep this consumer alive afterwards and must not pin preview focus. const endPreviewActivity = yield* previewAgentActivity.begin(input.run.threadId); - const providerEventFiber = yield* eventSubscription.events.pipe( - Stream.mapEffect((event) => - Effect.gen(function* () { - const previous = yield* Ref.get(eventRouting); - const [accepted, routing] = routeProviderEvent(event, routeIdentity, previous); - yield* Ref.set(eventRouting, routing); - if ( - event.type === "subagent.updated" && - event.subagent.childThreadId !== null && - previous.ownedThreadIds.has(event.subagent.childThreadId) && - !routing.ownedThreadIds.has(event.subagent.childThreadId) - ) { - const moved = event.subagent; - const detached = yield* Ref.modify(openRunOwnedSubagents, (current) => { - const subagentIds = new Set([moved.id]); - for (const subagent of current.subagents.values()) { - if (subagent.threadId === moved.childThreadId) subagentIds.add(subagent.id); - } - const childTurnItems = new Map( - [...current.childTurnItems].filter( - ([, item]) => item.threadId !== moved.childThreadId, - ), - ); - const turnItems = new Map(current.turnItems); - turnItems.delete(moved.id); - const detachedItems = new Set( - [...current.childTurnItems.values(), ...current.turnItems.values()] - .filter( - (item) => item.threadId === moved.childThreadId || item.nodeId === moved.id, - ) - .map((item) => item.id), - ); - return [ - { itemIds: detachedItems, subagentIds }, - { - subagents: new Map( - [...current.subagents].filter(([id]) => !subagentIds.has(id)), - ), - turnItems, - childTurnItems, - nodes: new Map( - [...current.nodes].filter( - ([id, node]) => id !== moved.id && node.threadId !== moved.childThreadId, - ), - ), - linkedChildThreadIds: new Set( - [...current.linkedChildThreadIds].filter( - (id) => id !== moved.childThreadId, - ), - ), - }, - ]; - }); - yield* Ref.update( - activeChildSubagents, - (current) => - new Map( - [...current].filter( - ([id, threadId]) => - !detached.subagentIds.has(id) && threadId !== moved.childThreadId, - ), - ), - ); - yield* Ref.update( - activeChildProviderTurns, - (current) => - new Set([...current].filter((id) => routing.ownedProviderTurns.has(id))), - ); - for (const id of pendingBackgroundProviderTurns) { - if (!routing.ownedProviderTurns.has(id)) { - pendingBackgroundProviderTurns.delete(id); - } - } - yield* Ref.update( - activeBackgroundTurnItems, - (current) => new Set([...current].filter((id) => !detached.itemIds.has(id))), - ); - } - return { event, accepted }; - }), - ), - Stream.tap(({ event, accepted }) => - Effect.gen(function* () { - if (!accepted) return; - if ( - event.type === "provider_turn.updated" && - event.providerTurn.runAttemptId === input.attempt.id && - event.providerTurn.providerThreadId === input.providerThread.id - ) { - yield* Ref.set(latestRootProviderTurn, event.providerTurn); + // Per-event routing and lifecycle bookkeeping share the run's tracing boundary. + const routeEvent = Effect.fnUntraced(function* ( + event: ProviderAdapter.ProviderAdapterV2Event, + ) { + const previous = yield* Ref.get(eventRouting); + const [accepted, routing] = routeProviderEvent(event, routeIdentity, previous); + yield* Ref.set(eventRouting, routing); + if ( + event.type === "subagent.updated" && + event.subagent.childThreadId !== null && + previous.ownedThreadIds.has(event.subagent.childThreadId) && + !routing.ownedThreadIds.has(event.subagent.childThreadId) + ) { + const moved = event.subagent; + const detached = yield* Ref.modify(openRunOwnedSubagents, (current) => { + const subagentIds = new Set([moved.id]); + for (const subagent of current.subagents.values()) { + if (subagent.threadId === moved.childThreadId) subagentIds.add(subagent.id); } - let storedEventCount = 0; - const deliveredEvent = filterAssistantEvent( - event, - DateTime.toEpochMillis(yield* DateTime.now), + const childTurnItems = new Map( + [...current.childTurnItems].filter( + ([, item]) => item.threadId !== moved.childThreadId, + ), ); - 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); - } - } else if ( - event.type === "turn_item.updated" && - (event.turnItem.type === "assistant_message" || - event.turnItem.type === "reasoning") && - event.turnItem.streaming && - !reservedHeldTurnItemIds.has(event.turnItem.id) - ) { - yield* turnItemPositions.allocate({ - threadId: event.turnItem.threadId, - turnItemId: event.turnItem.id, - runId: event.turnItem.runId, - }); - reservedHeldTurnItemIds.add(event.turnItem.id); - } - if (event.type === "provider_thread.updated") { - if (event.providerThread.id === input.providerThread.id && storedEventCount > 0) { - yield* Ref.set(latestProviderThread, event.providerThread); - } - } - if ( - event.type === "turn_item.updated" && - event.turnItem.providerTurnId === (yield* Ref.get(eventRouting)).rootProviderTurnId - ) { - yield* Ref.update(latestTurnItemOrdinal, (current) => - Math.max(current, event.turnItem.ordinal), - ); - } - if (event.type === "turn.terminal") { - yield* Ref.set(terminalEvent, event); - yield* Ref.set(rootTerminalSeen, true); - yield* finalizeRootRun(event); - } - if (event.type === "background_work.updated") { - if (event.pending) { - pendingBackgroundProviderTurns.add(event.providerTurnId); - } else { - pendingBackgroundProviderTurns.delete(event.providerTurnId); - } + const turnItems = new Map(current.turnItems); + turnItems.delete(moved.id); + const detachedItems = new Set( + [...current.childTurnItems.values(), ...current.turnItems.values()] + .filter( + (item) => item.threadId === moved.childThreadId || item.nodeId === moved.id, + ) + .map((item) => item.id), + ); + return [ + { itemIds: detachedItems, subagentIds }, + { + subagents: new Map([...current.subagents].filter(([id]) => !subagentIds.has(id))), + turnItems, + childTurnItems, + nodes: new Map( + [...current.nodes].filter( + ([id, node]) => id !== moved.id && node.threadId !== moved.childThreadId, + ), + ), + linkedChildThreadIds: new Set( + [...current.linkedChildThreadIds].filter((id) => id !== moved.childThreadId), + ), + }, + ]; + }); + yield* Ref.update( + activeChildSubagents, + (current) => + new Map( + [...current].filter( + ([id, threadId]) => + !detached.subagentIds.has(id) && threadId !== moved.childThreadId, + ), + ), + ); + yield* Ref.update( + activeChildProviderTurns, + (current) => new Set([...current].filter((id) => routing.ownedProviderTurns.has(id))), + ); + for (const id of pendingBackgroundProviderTurns) { + if (!routing.ownedProviderTurns.has(id)) { + pendingBackgroundProviderTurns.delete(id); } - yield* trackChildLifecycle(event, deliveredEvent !== null); - }), - ), - // Also after events this run does not accept, such as a child handed to another run. - Stream.tap(() => + } + yield* Ref.update( + activeBackgroundTurnItems, + (current) => new Set([...current].filter((id) => !detached.itemIds.has(id))), + ); + } + return { event, accepted }; + }); + const ingestEvent = Effect.fnUntraced(function* ({ + event, + accepted, + }: { + event: ProviderAdapter.ProviderAdapterV2Event; + accepted: boolean; + }) { + if (!accepted) return; + if ( + event.type === "provider_turn.updated" && + event.providerTurn.runAttemptId === input.attempt.id && + event.providerTurn.providerThreadId === input.providerThread.id + ) { + yield* Ref.set(latestRootProviderTurn, event.providerTurn); + } + let storedEventCount = 0; + const deliveredEvent = filterAssistantEvent( + 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); + } + } else if ( + event.type === "turn_item.updated" && + (event.turnItem.type === "assistant_message" || event.turnItem.type === "reasoning") && + event.turnItem.streaming && + !reservedHeldTurnItemIds.has(event.turnItem.id) + ) { + yield* turnItemPositions.allocate({ + threadId: event.turnItem.threadId, + turnItemId: event.turnItem.id, + runId: event.turnItem.runId, + }); + reservedHeldTurnItemIds.add(event.turnItem.id); + } + if (event.type === "provider_thread.updated") { + if (event.providerThread.id === input.providerThread.id && storedEventCount > 0) { + yield* Ref.set(latestProviderThread, event.providerThread); + } + } + if ( + event.type === "turn_item.updated" && + event.turnItem.providerTurnId === (yield* Ref.get(eventRouting)).rootProviderTurnId + ) { + yield* Ref.update(latestTurnItemOrdinal, (current) => + Math.max(current, event.turnItem.ordinal), + ); + } + if (event.type === "turn.terminal") { + yield* Ref.set(terminalEvent, event); + yield* Ref.set(rootTerminalSeen, true); + yield* finalizeRootRun(event); + } + if (event.type === "background_work.updated") { + if (event.pending) { + pendingBackgroundProviderTurns.add(event.providerTurnId); + } else { + pendingBackgroundProviderTurns.delete(event.providerTurnId); + } + } + yield* trackChildLifecycle(event, deliveredEvent !== null); + }); + const providerEventFiber = yield* eventSink + .consumeProviderEvents( + eventSubscription.events, + (event) => + Effect.gen(function* () { + yield* ingestEvent(yield* routeEvent(event)); + }), Effect.gen(function* () { if (yield* agentTurnsSettled) yield* endPreviewActivity; + return yield* shouldStopProviderEventIngestion; }), - ), - Stream.takeUntilEffect(() => shouldStopProviderEventIngestion), - Stream.runDrain, - Effect.mapError((cause) => new RunExecutionIngestError({ runId: input.run.id, cause })), - Effect.flatMap(() => - Effect.gen(function* () { - const terminal = yield* Ref.get(terminalEvent); - if (terminal === null) { - // A session release interrupts or fails the stream; a clean - // end without a terminal means the run would never settle. - return yield* new RunExecutionIngestError({ - runId: input.run.id, - cause: new ProviderAdapter.ProviderAdapterEventStreamError({ - driver: input.providerThread.driver, - providerSessionId: input.providerSessionId, - cause: "Provider event stream ended without a root terminal.", - }), - }); - } - yield* finalizeRootRun(terminal); - }), - ), - // Server shutdown interrupts the subscription and leaves the run to - // restart recovery; every other failure fails the run. - Effect.catchCauseIf( - (cause) => !Cause.hasInterruptsOnly(cause), - (cause) => + Ref.get(rootTerminalSeen).pipe(Effect.map((seen) => !seen)), + ) + .pipe( + Effect.mapError((cause) => new RunExecutionIngestError({ runId: input.run.id, cause })), + Effect.flatMap(() => Effect.gen(function* () { - const completedAt = yield* DateTime.now; - yield* Effect.logWarning("orchestration V2 provider event ingestion failed", { - runId: input.run.id, - cause, - }); - if (yield* Ref.get(rootRunFinalized)) return; - const providerThread = yield* Ref.get(latestProviderThread); - const routing = yield* Ref.get(eventRouting); - const rootProviderTurn = yield* Ref.get(latestRootProviderTurn); - // Stop background work too; the failed consumer cannot persist terminal events. - const abandonedProviderTurn = - routing.rootProviderTurnId !== null && !routing.rootTurnEnded - ? yield* input.session - .interruptTurn({ - providerThread, - providerTurnId: routing.rootProviderTurnId, - requestRuntimeRestart: true, - }) - .pipe( - Effect.as(true), - Effect.catchCauseIf( - (interruptCause) => !Cause.hasInterruptsOnly(interruptCause), - (interruptCause) => - Effect.logWarning( - "orchestration V2 failed to interrupt turn after ingestion failure", - { runId: input.run.id, cause: interruptCause }, - ).pipe(Effect.as(false)), - ), - // A failed interrupt leaves the stored turn running so the next start - // retries it. The snapshot is cached before it is published, so it may - // not be stored: replay a terminal one as-is. - Effect.map((interrupted) => - !interrupted || rootProviderTurn === null - ? null - : isTerminalProviderTurnStatus(rootProviderTurn.status) - ? rootProviderTurn - : { - ...rootProviderTurn, - status: "interrupted" as const, - completedAt, - }, - ), - ) - : null; - const latestItemOrdinal = yield* Ref.get(latestTurnItemOrdinal); - const openSubagents = yield* Ref.get(openRunOwnedSubagents); - const squashed = Cause.squash(cause); - let nextRetryWarningAt = completedAt; - yield* writeFinalRunEvents({ - run: input.run, - rootNode: input.rootNode, - checkpointScope: input.checkpointScope, - providerThread, - attempt: input.attempt, - writeIfRunCurrent: { - activeAttemptId: input.attempt.id, - expectedStatus: "running", - }, - openRunOwnedSubagents: openSubagents, - ...(abandonedProviderTurn === null ? {} : { abandonedProviderTurn }), - completedAt, - terminal: makeFailedTerminalEvent( - makeProviderFailure({ - cause: squashed, - message: - knownProviderFailureMessage(squashed) ?? - userFacingDatabaseErrorMessage(squashed), - class: "unknown", - }), - latestItemOrdinal + 1, - ), - failureItemPersisted: false, - refreshAfterTurn, - }).pipe( - // Storage failures can be defects; retry both failure channels. - Effect.catchCause((writeCause) => - Effect.gen(function* () { - if (Cause.hasInterruptsOnly(writeCause)) { - return yield* Effect.failCause(writeCause); - } - const now = yield* DateTime.now; - if (DateTime.isGreaterThanOrEqualTo(now, nextRetryWarningAt)) { - yield* Effect.logWarning( - "orchestration V2 run failure finalization retry", - { - runId: input.run.id, - attemptId: input.attempt.id, - cause: writeCause, - }, - ); - // Persistent storage failures must not flood the same disk. - nextRetryWarningAt = DateTime.add(now, { minutes: 10 }); - } - return yield* Effect.fail(writeCause); + const terminal = yield* Ref.get(terminalEvent); + if (terminal === null) { + // A session release interrupts or fails the stream; a clean + // end without a terminal means the run would never settle. + return yield* new RunExecutionIngestError({ + runId: input.run.id, + cause: new ProviderAdapter.ProviderAdapterEventStreamError({ + driver: input.providerThread.driver, + providerSessionId: input.providerSessionId, + cause: "Provider event stream ended without a root terminal.", }), - ), - Effect.retry({ - schedule: Schedule.exponential("1 second").pipe( - Schedule.modifyDelay(({ duration }) => - Effect.succeed(Duration.min(duration, Duration.seconds(30))), - ), - ), - }), - // Startup may cancel the consumer and fail its own terminal - // write. Only service shutdown or the atomic gate ends recovery. - Effect.forkIn(scope), - Effect.flatMap(Fiber.join), - ); + }); + } + yield* finalizeRootRun(terminal); }), - ), - Effect.ensuring(endPreviewActivity), - Effect.ensuring(eventSubscription.close), - Effect.forkDetach, - ); + ), + // Server shutdown interrupts the subscription and leaves the run to + // restart recovery; every other failure fails the run. + Effect.catchCauseIf( + (cause) => !Cause.hasInterruptsOnly(cause), + (cause) => + Effect.gen(function* () { + const completedAt = yield* DateTime.now; + yield* Effect.logWarning("orchestration V2 provider event ingestion failed", { + runId: input.run.id, + cause, + }); + if (yield* Ref.get(rootRunFinalized)) return; + const providerThread = yield* Ref.get(latestProviderThread); + const routing = yield* Ref.get(eventRouting); + const rootProviderTurn = yield* Ref.get(latestRootProviderTurn); + // Stop background work too; the failed consumer cannot persist terminal events. + const abandonedProviderTurn = + routing.rootProviderTurnId !== null && !routing.rootTurnEnded + ? yield* input.session + .interruptTurn({ + providerThread, + providerTurnId: routing.rootProviderTurnId, + requestRuntimeRestart: true, + }) + .pipe( + Effect.as(true), + Effect.catchCauseIf( + (interruptCause) => !Cause.hasInterruptsOnly(interruptCause), + (interruptCause) => + Effect.logWarning( + "orchestration V2 failed to interrupt turn after ingestion failure", + { runId: input.run.id, cause: interruptCause }, + ).pipe(Effect.as(false)), + ), + // A failed interrupt leaves the stored turn running so the next start + // retries it. The snapshot is cached before it is published, so it may + // not be stored: replay a terminal one as-is. + Effect.map((interrupted) => + !interrupted || rootProviderTurn === null + ? null + : isTerminalProviderTurnStatus(rootProviderTurn.status) + ? rootProviderTurn + : { + ...rootProviderTurn, + status: "interrupted" as const, + completedAt, + }, + ), + ) + : null; + const latestItemOrdinal = yield* Ref.get(latestTurnItemOrdinal); + const openSubagents = yield* Ref.get(openRunOwnedSubagents); + const squashed = Cause.squash(cause); + let nextRetryWarningAt = completedAt; + yield* writeFinalRunEvents({ + run: input.run, + rootNode: input.rootNode, + checkpointScope: input.checkpointScope, + providerThread, + attempt: input.attempt, + writeIfRunCurrent: { + activeAttemptId: input.attempt.id, + expectedStatus: "running", + }, + openRunOwnedSubagents: openSubagents, + ...(abandonedProviderTurn === null ? {} : { abandonedProviderTurn }), + completedAt, + terminal: makeFailedTerminalEvent( + makeProviderFailure({ + cause: squashed, + message: + knownProviderFailureMessage(squashed) ?? + userFacingDatabaseErrorMessage(squashed), + class: "unknown", + }), + latestItemOrdinal + 1, + ), + failureItemPersisted: false, + refreshAfterTurn, + }).pipe( + // Storage failures can be defects; retry both failure channels. + Effect.catchCause((writeCause) => + Effect.gen(function* () { + if (Cause.hasInterruptsOnly(writeCause)) { + return yield* Effect.failCause(writeCause); + } + const now = yield* DateTime.now; + if (DateTime.isGreaterThanOrEqualTo(now, nextRetryWarningAt)) { + yield* Effect.logWarning( + "orchestration V2 run failure finalization retry", + { + runId: input.run.id, + attemptId: input.attempt.id, + cause: writeCause, + }, + ); + // Persistent storage failures must not flood the same disk. + nextRetryWarningAt = DateTime.add(now, { minutes: 10 }); + } + return yield* Effect.fail(writeCause); + }), + ), + Effect.retry({ + schedule: Schedule.exponential("1 second").pipe( + Schedule.modifyDelay(({ duration }) => + Effect.succeed(Duration.min(duration, Duration.seconds(30))), + ), + ), + }), + // Startup may cancel the consumer and fail its own terminal + // write. Only service shutdown or the atomic gate ends recovery. + Effect.forkIn(scope), + Effect.flatMap(Fiber.join), + ); + }), + ), + Effect.ensuring(endPreviewActivity), + Effect.ensuring(eventSubscription.close), + Effect.forkDetach, + ); // A provider turn is a sign that the session is still alive. Keep // its already-issued MCP credential valid even when the agent goes diff --git a/apps/server/src/orchestration-v2/testkit/CodexReplayFixtures.integration.test.ts b/apps/server/src/orchestration-v2/testkit/CodexReplayFixtures.integration.test.ts index c0593351f20d7..1274c9eaee2e9 100644 --- a/apps/server/src/orchestration-v2/testkit/CodexReplayFixtures.integration.test.ts +++ b/apps/server/src/orchestration-v2/testkit/CodexReplayFixtures.integration.test.ts @@ -1,9 +1,14 @@ -import { describe } from "@effect/vitest"; +import { assert, describe, it } from "@effect/vitest"; +import * as Tracer from "effect/Tracer"; import * as Effect from "effect/Effect"; import { materializeReplayTranscriptWorkspace } from "@t3tools/provider-testing/replayTranscript"; import { CodexOrchestratorReplayHarness } from "../Adapters/CodexAdapterV2.testkit.ts"; -import { itRunsReplayFixtures, runReplayFixture } from "./ProviderReplayFixtures.ts"; +import { + findReplayFixture, + itRunsReplayFixtures, + runReplayFixture, +} from "./ProviderReplayFixtures.ts"; describe("Codex replay fixtures", () => { itRunsReplayFixtures("codex", (input) => @@ -15,3 +20,39 @@ describe("Codex replay fixtures", () => { }), ); }); + +it.effect( + "persists recorded Codex snapshots through provider transactions in the real run consumer", + () => + Effect.gen(function* () { + const fixture = findReplayFixture("multi_turn")!; + const driver = fixture.providers.find((provider) => provider.driver === "codex")!; + const writes: number[] = []; + const tracer = Tracer.make({ + span(options) { + const span = new Tracer.NativeSpan(options); + const end = span.end.bind(span); + span.end = (time, exit) => { + end(time, exit); + if (span.name === "orchestrationV2.EventSink.providerBatch") { + writes.push(Number(span.attributes.get("orchestration_v2.batch_writes"))); + } + }; + return span; + }, + }); + yield* runReplayFixture({ + fixtureName: fixture.name, + buildInput: fixture.buildInput, + driver, + harness: CodexOrchestratorReplayHarness, + materializeWorkspace: (transcript, workspace) => + Effect.succeed(materializeReplayTranscriptWorkspace(transcript, workspace)), + }).pipe(Effect.withTracer(tracer), Effect.withTracerEnabled(true)); + assert.equal( + writes.reduce((total, count) => total + count, 0), + 6, + ); + assert.isTrue(writes.every((count) => count <= 3)); + }), +); diff --git a/packages/provider-opencode/src/server/adapter.test.ts b/packages/provider-opencode/src/server/adapter.test.ts index 70462070cff51..769d3a81e63d0 100644 --- a/packages/provider-opencode/src/server/adapter.test.ts +++ b/packages/provider-opencode/src/server/adapter.test.ts @@ -712,6 +712,10 @@ describe("OpenCodeAdapterV2", () => { }), ); const events = yield* Fiber.join(received); + assert.deepEqual( + events.map((event) => event.type), + ["turn_item.updated", "runtime_request.updated"], + ); const item = events.find((event) => event.type === "turn_item.updated"); const request = events.find((event) => event.type === "runtime_request.updated"); assert.equal(item?.type, "turn_item.updated"); diff --git a/packages/provider-opencode/src/server/adapter.ts b/packages/provider-opencode/src/server/adapter.ts index 4c72c5f6ef09f..94bdb56eb29ab 100644 --- a/packages/provider-opencode/src/server/adapter.ts +++ b/packages/provider-opencode/src/server/adapter.ts @@ -1899,15 +1899,16 @@ export const makeOpenCodeAdapterV2 = Effect.fn("makeOpenCodeAdapterV2")(function }, }); yield* emitProviderEvent({ - type: "runtime_request.updated", + type: "turn_item.updated", driver: OPENCODE_PROVIDER, - threadId: turn.threadId, - runtimeRequest, + turnItem: runtimeRequestTurnItem(pending, "waiting", null, now), }); + // A response can commit as soon as the actionable request is published. yield* emitProviderEvent({ - type: "turn_item.updated", + type: "runtime_request.updated", driver: OPENCODE_PROVIDER, - turnItem: runtimeRequestTurnItem(pending, "waiting", null, now), + threadId: turn.threadId, + runtimeRequest, }); yield* updateProviderSession("waiting", null); }); diff --git a/packages/provider-opencode/src/server/v2/adapter.ts b/packages/provider-opencode/src/server/v2/adapter.ts index ddca16ed8fed3..3cbd4d893b6f1 100644 --- a/packages/provider-opencode/src/server/v2/adapter.ts +++ b/packages/provider-opencode/src/server/v2/adapter.ts @@ -1888,13 +1888,14 @@ export const make = Effect.fn("OpenCode2Adapter.make")(function* (instanceId: Pr native, }); yield* emit({ type: "node.updated", driver, node }); + yield* emit({ type: "turn_item.updated", driver, turnItem: item }); + // A response can commit as soon as the actionable request is published. yield* emit({ type: "runtime_request.updated", driver, threadId: turn.input.threadId, runtimeRequest: request, }); - yield* emit({ type: "turn_item.updated", driver, turnItem: item }); yield* setSessionStatus("waiting", null); });