diff --git a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts index dd6a0b179a8f..9b622de65f30 100644 --- a/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts +++ b/apps/server/src/orchestration/Layers/ProviderRuntimeIngestion.test.ts @@ -5,6 +5,7 @@ import * as NodePath from "node:path"; import * as NodeChildProcess from "node:child_process"; import { + CodexSettings, OrchestrationReadModel, ProviderDriverKind, ProviderRuntimeEvent, @@ -29,10 +30,12 @@ import * as Clock from "effect/Clock"; 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 Layer from "effect/Layer"; import * as ManagedRuntime from "effect/ManagedRuntime"; import * as Option from "effect/Option"; import * as PubSub from "effect/PubSub"; +import * as Schema from "effect/Schema"; import * as Scope from "effect/Scope"; import * as Stream from "effect/Stream"; import * as Tracer from "effect/Tracer"; @@ -46,6 +49,9 @@ import { ProviderService, type ProviderServiceShape, } from "../../provider/Services/ProviderService.ts"; +import type { CodexAdapterShape } from "../../provider/Services/CodexAdapter.ts"; +import { makeCodexAdapter } from "../../provider/Layers/CodexAdapter.ts"; +import codexMultiAgentWire from "../../provider/testFixtures/codexMultiAgentWire.json" with { type: "json" }; import * as RepositoryIdentityResolver from "../../project/RepositoryIdentityResolver.ts"; import * as CheckpointStore from "../../checkpointing/CheckpointStore.ts"; import * as VcsDriverRegistry from "../../vcs/VcsDriverRegistry.ts"; @@ -75,6 +81,7 @@ const asEventId = (value: string): EventId => EventId.make(value); const asMessageId = (value: string): MessageId => MessageId.make(value); const asThreadId = (value: string): ThreadId => ThreadId.make(value); const asTurnId = (value: string): TurnId => TurnId.make(value); +const decodeCodexSettings = Schema.decodeSync(CodexSettings); type LegacyProviderRuntimeEvent = { readonly type: string; @@ -203,6 +210,56 @@ function createProviderServiceHarness() { }; } +function providerServiceFromCodexAdapter( + adapter: CodexAdapterShape, + turnCompletionEnqueued: Deferred.Deferred, +): ProviderServiceShape { + const unsupported = () => + Effect.die(new Error("Unsupported provider call in Codex progress test")); + return { + startSession: (threadId, input) => adapter.startSession({ ...input, threadId }), + sendTurn: adapter.sendTurn, + compactThread: unsupported, + interruptTurn: ({ threadId, turnId }) => adapter.interruptTurn(threadId, turnId), + respondToRequest: ({ threadId, requestId, decision }) => + adapter.respondToRequest(threadId, requestId, decision), + respondToUserInput: ({ threadId, requestId, answers }) => + adapter.respondToUserInput(threadId, requestId, answers), + stopSession: ({ threadId }) => adapter.stopSession(threadId), + listSessions: adapter.listSessions, + getCapabilities: () => Effect.succeed(adapter.capabilities), + assertConversationRollbackSupported: unsupported, + getInstanceInfo: (instanceId) => + Effect.succeed({ + instanceId, + driverKind: adapter.provider, + displayName: undefined, + enabled: true, + continuationIdentity: { + driverKind: adapter.provider, + continuationKey: `${adapter.provider}:instance:${instanceId}`, + }, + }), + rollbackConversation: ({ threadId, numTurns }) => + adapter.rollbackThread(threadId, numTurns).pipe(Effect.asVoid), + uploadFeedback: adapter.uploadFeedback, + get streamEvents() { + return adapter.streamEvents.pipe( + Stream.flatMap((event) => + Stream.concat( + Stream.succeed(event), + event.type === "turn.completed" + ? Stream.fromEffect(Deferred.succeed(turnCompletionEnqueued, undefined)).pipe( + Stream.drain, + ) + : Stream.empty, + ), + ), + ); + }, + }; +} + type ProviderRuntimeTestReadModel = OrchestrationReadModel; type ProviderRuntimeTestThread = ProviderRuntimeTestReadModel["threads"][number]; type ProviderRuntimeTestMessage = ProviderRuntimeTestThread["messages"][number]; @@ -238,6 +295,7 @@ describe("ProviderRuntimeIngestion", () => { unknown > | null = null; let scope: Scope.Closeable | null = null; + const providerScopes: Scope.Closeable[] = []; const tempDirs: string[] = []; function makeTempDir(prefix: string): string { @@ -251,6 +309,9 @@ describe("ProviderRuntimeIngestion", () => { await Effect.runPromise(Scope.close(scope, Exit.void)); } scope = null; + for (const providerScope of providerScopes.splice(0)) { + await Effect.runPromise(Scope.close(providerScope, Exit.void)); + } if (runtime) { await runtime.dispose(); } @@ -262,6 +323,7 @@ describe("ProviderRuntimeIngestion", () => { async function createHarness(options?: { serverSettings?: Partial; + providerService?: ProviderServiceShape; threadTitle?: string; workspaceSubdirectory?: string; }) { @@ -294,7 +356,9 @@ describe("ProviderRuntimeIngestion", () => { Layer.provideMerge(ThreadBackgroundLiveness.layer), Layer.provideMerge(ThreadPlanProgress.layer), Layer.provideMerge(SqlitePersistenceMemory), - Layer.provideMerge(Layer.succeed(ProviderService, provider.service)), + Layer.provideMerge( + Layer.succeed(ProviderService, options?.providerService ?? provider.service), + ), Layer.provideMerge(makeTestServerSettingsLayer(options?.serverSettings)), Layer.provideMerge(CheckpointStore.layer.pipe(Layer.provide(VcsDriverRegistry.layer))), Layer.provideMerge(VcsProcess.layer), @@ -3808,6 +3872,231 @@ describe("ProviderRuntimeIngestion", () => { expect(activity?.payload).toMatchObject({ requestId: "message-compact" }); }); + effectIt.effect.each([false, true])( + "bounds a Codex child progress burst with lifecycle barrier %s", + (withLifecycleBarrier) => + Effect.gen(function* () { + const rootThreadId = codexMultiAgentWire.rootThreadId; + const childThreadId = codexMultiAgentWire.childThreadIds[0]; + const childTurnStarted = codexMultiAgentWire.notifications.find( + (entry) => + entry.method === "turn/started" && + (entry.params as { threadId?: string }).threadId === childThreadId, + ); + const childRegistration = codexMultiAgentWire.notifications.find((entry) => { + const params = entry.params as { + threadId?: string; + item?: { type?: string; agentThreadId?: string }; + }; + return ( + params.threadId === rootThreadId && + params.item?.type === "subAgentActivity" && + params.item.agentThreadId === childThreadId + ); + }); + const childTurnCompleted = codexMultiAgentWire.notifications.find( + (entry) => + entry.method === "turn/completed" && + (entry.params as { threadId?: string }).threadId === childThreadId, + ); + expect(childRegistration).toBeDefined(); + expect(childTurnStarted).toBeDefined(); + expect(childTurnCompleted).toBeDefined(); + + const childTurnId = (childTurnStarted?.params as { turn?: { id?: string } } | undefined) + ?.turn?.id; + expect(childTurnId).toBeDefined(); + const burstSize = 32; + const burst = Array.from({ length: burstSize }, (_, index) => [ + { + method: "item/completed", + params: { + threadId: childThreadId, + turnId: childTurnId, + completedAtMs: 1_785_898_350_000 + index, + item: { + type: "webSearch", + id: `child-search-${index}`, + query: `latest-query-${index}`, + results: [], + }, + }, + }, + { + method: "thread/tokenUsage/updated", + params: { + threadId: childThreadId, + turnId: childTurnId, + tokenUsage: { + total: { + totalTokens: 10_000 + index, + inputTokens: 9_000 + index, + cachedInputTokens: 8_000 + index, + cacheWriteInputTokens: 0, + outputTokens: 1_000, + reasoningOutputTokens: index, + }, + last: { + totalTokens: 100 + index, + inputTokens: 90 + index, + cachedInputTokens: 80 + index, + cacheWriteInputTokens: 0, + outputTokens: 10, + reasoningOutputTokens: index, + }, + modelContextWindow: 258_400, + }, + }, + }, + ]).flat(); + const scriptPath = NodePath.join(makeTempDir("t3-codex-progress-script-"), "script.json"); + NodeFS.writeFileSync( + scriptPath, + // @effect-diagnostics-next-line preferSchemaOverJson:off + JSON.stringify({ + rootThreadId, + notifications: [ + childRegistration, + childTurnStarted, + ...burst.slice(0, burstSize), + ...(withLifecycleBarrier + ? [ + { + method: "thread/status/changed", + params: { + threadId: childThreadId, + status: { type: "active", activeFlags: ["waitingOnApproval"] }, + }, + }, + { + method: "thread/status/changed", + params: { + threadId: childThreadId, + status: { type: "active", activeFlags: [] }, + }, + }, + ] + : []), + ...burst.slice(burstSize), + childTurnCompleted, + ], + }), + "utf8", + ); + + const providerScope = yield* Scope.make("sequential"); + providerScopes.push(providerScope); + const peerPath = NodePath.join( + import.meta.dirname, + "../../provider/testFixtures/codexCollabMockPeer.sh", + ); + const adapter = yield* makeCodexAdapter(decodeCodexSettings({ binaryPath: peerPath }), { + environment: { ...process.env, T3_CODEX_COLLAB_SCRIPT: scriptPath }, + }).pipe( + Effect.provide( + Layer.mergeAll( + Layer.succeed(Scope.Scope, providerScope), + ServerConfig.layerTest(process.cwd(), process.cwd()), + ).pipe(Layer.provideMerge(NodeServices.layer)), + ), + ); + const turnCompletionEnqueued = yield* Deferred.make(); + const harness = yield* Effect.promise(() => + createHarness({ + providerService: providerServiceFromCodexAdapter(adapter, turnCompletionEnqueued), + }), + ); + const threadId = asThreadId("thread-1"); + const stopAdapter = yield* Deferred.make(); + const adapterRun = yield* Effect.gen(function* () { + yield* adapter.startSession({ + provider: ProviderDriverKind.make("codex"), + threadId, + runtimeMode: "full-access", + }); + yield* adapter.sendTurn({ threadId, input: "fan out" }); + yield* Deferred.await(stopAdapter); + yield* adapter.stopSession(threadId); + }).pipe(Effect.forkChild); + + try { + yield* Deferred.await(turnCompletionEnqueued); + yield* Effect.promise(() => harness.drain()); + + const events = Array.from(yield* Stream.runCollect(harness.engine.readEvents(0))); + const childActivities = events.flatMap((event) => { + if (event.type !== "thread.activity-appended") return []; + const activity = event.payload.activity; + const payload = activity.payload as { taskId?: string } | undefined; + return payload?.taskId === childThreadId + ? [{ activity, durableSequence: event.sequence }] + : []; + }); + const progress = childActivities.filter( + ({ activity }) => activity.kind === "task.progress", + ); + const completion = childActivities.find(({ activity }) => { + if (activity.kind !== "task.updated") return false; + const payload = activity.payload as { status?: string } | undefined; + return payload?.status === "idle"; + }); + + expect( + progress.some(({ activity }) => { + const payload = activity.payload as { summary?: string } | undefined; + return payload?.summary === `latest-query-${burstSize - 1}`; + }), + ).toBe(true); + const latestUsage = progress.find( + ({ activity }) => + (activity.payload as { typedUsage?: { totalTokens?: number } } | undefined) + ?.typedUsage?.totalTokens === + 10_000 + burstSize - 1, + ); + expect(latestUsage?.activity.payload).toMatchObject({ + typedUsage: { + totalTokens: 10_000 + burstSize - 1, + inputTokens: 9_000 + burstSize - 1, + cachedInputTokens: 8_000 + burstSize - 1, + outputTokens: 1_000, + reasoningOutputTokens: burstSize - 1, + }, + }); + expect(completion).toBeDefined(); + expect( + progress.every( + ({ durableSequence }) => + completion !== undefined && durableSequence < completion.durableSequence, + ), + ).toBe(true); + expect(progress.length).toBe(withLifecycleBarrier ? 4 : 2); + if (withLifecycleBarrier) { + const waiting = childActivities.find( + ({ activity }) => + activity.kind === "task.updated" && + (activity.payload as { status?: string } | undefined)?.status === "waiting", + ); + expect(waiting).toBeDefined(); + expect( + progress.filter( + ({ durableSequence }) => + waiting !== undefined && durableSequence < waiting.durableSequence, + ), + ).toHaveLength(2); + expect( + progress.filter( + ({ durableSequence }) => + waiting !== undefined && durableSequence > waiting.durableSequence, + ), + ).toHaveLength(2); + } + } finally { + yield* Deferred.succeed(stopAdapter, undefined); + yield* Fiber.join(adapterRun); + } + }), + ); + it("projects Codex task lifecycle chunks into thread activities", async () => { const harness = await createHarness(); const now = "2026-01-01T00:00:00.000Z"; diff --git a/apps/server/src/provider/Layers/CodexAdapter.test.ts b/apps/server/src/provider/Layers/CodexAdapter.test.ts index 2ca44a2a5f0a..8cae076f810a 100644 --- a/apps/server/src/provider/Layers/CodexAdapter.test.ts +++ b/apps/server/src/provider/Layers/CodexAdapter.test.ts @@ -12,6 +12,7 @@ import { ProviderItemId, type ProviderApprovalDecision, type ProviderEvent, + type ProviderRuntimeEvent, type ProviderSession, type ProviderTurnStartResult, type ProviderUserInputAnswers, @@ -23,6 +24,7 @@ import * as NodeServices from "@effect/platform-node/NodeServices"; import { it, vi } from "@effect/vitest"; import * as Context from "effect/Context"; +import * as Deferred from "effect/Deferred"; import * as Effect from "effect/Effect"; import * as Exit from "effect/Exit"; import * as Fiber from "effect/Fiber"; @@ -59,6 +61,15 @@ const asTurnId = (value: string): TurnId => TurnId.make(value); const asEventId = (value: string): EventId => EventId.make(value); const asItemId = (value: string): ProviderItemId => ProviderItemId.make(value); +type CodexTaskRuntimeEvent = Extract< + ProviderRuntimeEvent, + { type: "task.progress" | "task.updated" } +>; + +function isCodexTaskRuntimeEvent(event: ProviderRuntimeEvent): event is CodexTaskRuntimeEvent { + return event.type === "task.progress" || event.type === "task.updated"; +} + class FakeCodexRuntime implements CodexSessionRuntimeShape { private readonly eventQueue = Effect.runSync(Queue.unbounded()); private readonly now = "2026-01-01T00:00:00.000Z"; @@ -2483,6 +2494,442 @@ lifecycleLayer("CodexAdapterLive lifecycle", (it) => { }), ); + it.effect("flushes both progress lanes before each child terminal lifecycle event", () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + const terminals = [ + { + childId: "child-turn-completed", + method: "collabAgent/turnCompleted", + payload: { turn: { status: "completed" } }, + status: "idle", + }, + { + childId: "child-status-idle", + method: "collabAgent/statusChanged", + payload: { status: { type: "idle" } }, + status: "idle", + }, + { + childId: "child-closed", + method: "collabAgent/closed", + payload: {}, + status: "interrupted", + }, + { + childId: "child-error", + method: "collabAgent/statusChanged", + payload: { status: { type: "systemError" } }, + status: "failed", + }, + { + childId: "child-parent-interrupted", + method: "collabAgent/activity", + payload: { activityKind: "interrupted" }, + status: "interrupted", + }, + ] as const; + const taskEventsFiber = yield* adapter.streamEvents.pipe( + Stream.filter( + (event) => + (event.type === "task.progress" || event.type === "task.updated") && + terminals.some((terminal) => terminal.childId === event.payload.taskId), + ), + Stream.take(terminals.length * 3), + Stream.runCollect, + Effect.forkChild, + ); + + for (const [index, terminal] of terminals.entries()) { + const eventBase = { + kind: "notification" as const, + provider: ProviderDriverKind.make("codex"), + threadId: asThreadId("thread-1"), + turnId: asTurnId("turn-1"), + createdAt: "2026-01-01T00:00:00.000Z", + }; + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${terminal.childId}-item-old`), + method: "collabAgent/item", + payload: { + agentThreadId: terminal.childId, + agentPath: `/root/${terminal.childId}`, + item: { type: "webSearch", query: `old-${index}` }, + }, + } satisfies ProviderEvent); + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${terminal.childId}-item-latest`), + method: "collabAgent/item", + payload: { + agentThreadId: terminal.childId, + agentPath: `/root/${terminal.childId}`, + item: { type: "webSearch", query: `latest-${index}` }, + }, + } satisfies ProviderEvent); + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${terminal.childId}-usage-old`), + method: "collabAgent/tokenUsage", + payload: { + agentThreadId: terminal.childId, + agentPath: `/root/${terminal.childId}`, + tokenUsage: { total: { totalTokens: 100 + index } }, + }, + } satisfies ProviderEvent); + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${terminal.childId}-usage-latest`), + method: "collabAgent/tokenUsage", + payload: { + agentThreadId: terminal.childId, + agentPath: `/root/${terminal.childId}`, + tokenUsage: { total: { totalTokens: 200 + index } }, + }, + } satisfies ProviderEvent); + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${terminal.childId}-terminal`), + method: terminal.method, + payload: { agentThreadId: terminal.childId, ...terminal.payload }, + } satisfies ProviderEvent); + } + + const taskEvents = Array.from(yield* Fiber.join(taskEventsFiber)).filter( + isCodexTaskRuntimeEvent, + ); + for (const [index, terminal] of terminals.entries()) { + const childEvents = taskEvents.filter((event) => event.payload.taskId === terminal.childId); + NodeAssert.deepEqual( + childEvents.map((event) => event.type), + ["task.progress", "task.progress", "task.updated"], + ); + const itemProgress = childEvents[0]; + const usageProgress = childEvents[1]; + const terminalEvent = childEvents[2]; + NodeAssert.equal(itemProgress?.type, "task.progress"); + NodeAssert.equal(usageProgress?.type, "task.progress"); + NodeAssert.equal(terminalEvent?.type, "task.updated"); + if ( + itemProgress?.type !== "task.progress" || + usageProgress?.type !== "task.progress" || + terminalEvent?.type !== "task.updated" + ) { + return; + } + NodeAssert.equal(itemProgress.payload.summary, `latest-${index}`); + NodeAssert.equal(usageProgress.payload.typedUsage?.totalTokens, 200 + index); + NodeAssert.equal(terminalEvent.payload.status, terminal.status); + } + }), + ); + + it.effect("cancels pending child progress on stop and isolates a restarted session", () => + Effect.gen(function* () { + const adapter = yield* CodexAdapter; + const firstThreadId = asThreadId("thread-progress-close-a"); + const childId = "shared-child-after-close"; + yield* adapter.startSession({ + provider: ProviderDriverKind.make("codex"), + threadId: firstThreadId, + runtimeMode: "full-access", + }); + const firstRuntime = lifecycleRuntimeFactory.lastRuntime; + NodeAssert.ok(firstRuntime); + + const firstMarkerSeen = yield* Deferred.make(); + const taskEventsFiber = yield* adapter.streamEvents.pipe( + Stream.tap((event) => + event.type === "thread.metadata.updated" && event.payload.name === "before-close" + ? Deferred.succeed(firstMarkerSeen, undefined).pipe(Effect.asVoid) + : Effect.void, + ), + Stream.takeUntil( + (event) => + event.type === "thread.metadata.updated" && event.payload.name === "after-restart", + ), + Stream.runCollect, + Effect.forkChild, + ); + const firstEventBase = { + kind: "notification" as const, + provider: ProviderDriverKind.make("codex"), + threadId: firstThreadId, + createdAt: "2026-01-01T00:00:00.000Z", + }; + yield* firstRuntime.emit({ + ...firstEventBase, + id: asEventId("evt-progress-before-close"), + method: "collabAgent/item", + payload: { + agentThreadId: childId, + item: { type: "webSearch", query: "must be cancelled" }, + }, + } satisfies ProviderEvent); + yield* firstRuntime.emit({ + ...firstEventBase, + id: asEventId("evt-marker-before-close"), + method: "thread/name/updated", + payload: { threadId: "provider-thread-1", threadName: "before-close" }, + } satisfies ProviderEvent); + yield* Deferred.await(firstMarkerSeen); + + yield* adapter.stopSession(firstThreadId); + yield* TestClock.adjust("1 second"); + + yield* adapter.startSession({ + provider: ProviderDriverKind.make("codex"), + threadId: firstThreadId, + runtimeMode: "full-access", + }); + const secondRuntime = lifecycleRuntimeFactory.lastRuntime; + NodeAssert.ok(secondRuntime); + NodeAssert.notEqual(secondRuntime, firstRuntime); + yield* firstRuntime.emit({ + ...firstEventBase, + id: asEventId("evt-old-session-after-restart"), + method: "collabAgent/item", + payload: { agentThreadId: childId, item: { type: "webSearch", query: "old runtime" } }, + }); + yield* secondRuntime.emit({ + ...firstEventBase, + id: asEventId("evt-new-session-progress"), + method: "collabAgent/item", + payload: { agentThreadId: childId, item: { type: "webSearch", query: "new runtime" } }, + }); + yield* secondRuntime.emit({ + ...firstEventBase, + id: asEventId("evt-idle-after-restart"), + method: "collabAgent/statusChanged", + payload: { + agentThreadId: childId, + status: { type: "idle" }, + }, + } satisfies ProviderEvent); + yield* TestClock.adjust("1 second"); + yield* secondRuntime.emit({ + ...firstEventBase, + id: asEventId("evt-marker-after-restart"), + method: "thread/name/updated", + payload: { threadId: "provider-thread-1", threadName: "after-restart" }, + }); + + const taskEvents = Array.from(yield* Fiber.join(taskEventsFiber)).filter( + isCodexTaskRuntimeEvent, + ); + NodeAssert.deepEqual( + taskEvents.map((event) => event.type), + ["task.progress", "task.updated"], + ); + NodeAssert.equal(taskEvents[0]?.threadId, firstThreadId); + NodeAssert.equal( + taskEvents[0]?.type === "task.progress" && taskEvents[0].payload.summary, + "new runtime", + ); + NodeAssert.equal( + taskEvents[1]?.type === "task.updated" && taskEvents[1].payload.status, + "idle", + ); + }), + ); + + it.effect.each([ + ["turn/completed", "turn.completed"], + ["turn/aborted", "turn.aborted"], + ["session/exited", "session.exited"], + ["runtime-failure", "runtime.error"], + ["collabAgent/statusChanged", "task.updated"], + ["collabAgent/metadataUpdated", "task.updated"], + ] as const)("preserves child progress ordering before %s", ([method, boundaryType]) => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + const observed = yield* adapter.streamEvents.pipe( + Stream.takeUntil((event) => event.type === boundaryType), + Stream.runCollect, + Effect.forkChild, + ); + const eventBase = { + kind: "notification" as const, + provider: ProviderDriverKind.make("codex"), + threadId: asThreadId("thread-1"), + turnId: asTurnId("turn-1"), + createdAt: "2026-01-01T00:00:00.000Z", + }; + yield* runtime.emit({ + ...eventBase, + id: asEventId("evt-item-before-root-boundary"), + method: "collabAgent/item", + payload: { + agentThreadId: "child-root-boundary", + nickname: "old-name", + item: { type: "webSearch", query: "latest child item" }, + }, + }); + yield* runtime.emit({ + ...eventBase, + id: asEventId("evt-usage-before-root-boundary"), + method: "collabAgent/tokenUsage", + payload: { + agentThreadId: "child-root-boundary", + nickname: "old-name", + tokenUsage: { total: { totalTokens: 123 } }, + }, + }); + yield* runtime.emit( + method === "turn/completed" + ? codexTurnEvent(method, "turn-1") + : { + ...eventBase, + id: asEventId("evt-root-boundary"), + kind: method === "runtime-failure" ? "error" : "notification", + method, + message: "Owned boundary fixture", + ...(method.startsWith("collabAgent/") + ? { + payload: { + agentThreadId: "child-root-boundary", + nickname: "current-name", + status: { type: "active", activeFlags: ["waitingOnApproval"] }, + }, + } + : {}), + }, + ); + const events = Array.from(yield* Fiber.join(observed)); + NodeAssert.deepEqual( + events.map((event) => event.type), + ["task.progress", "task.progress", boundaryType], + ); + const boundary = events.at(-1); + if (boundary?.type === "task.updated") { + NodeAssert.equal(boundary.payload.title, "current-name"); + if (method === "collabAgent/statusChanged") { + NodeAssert.equal(boundary.payload.status, "waiting"); + } + } + }), + ); + + it.effect( + "flushes two children before root completion and leaves the root result unchanged", + () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + const observed = yield* adapter.streamEvents.pipe( + Stream.takeUntil((event) => event.type === "turn.completed"), + Stream.runCollect, + Effect.forkChild, + ); + const eventBase = { + kind: "notification" as const, + provider: ProviderDriverKind.make("codex"), + threadId: asThreadId("thread-1"), + turnId: asTurnId("turn-1"), + createdAt: "2026-01-01T00:00:00.000Z", + }; + for (const childId of ["child-b", "child-a"]) { + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${childId}-usage`), + method: "collabAgent/tokenUsage", + payload: { agentThreadId: childId, tokenUsage: { total: { totalTokens: 321 } } }, + }); + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${childId}-item`), + method: "collabAgent/item", + payload: { agentThreadId: childId, item: { type: "webSearch", query: childId } }, + }); + } + yield* runtime.emit(codexTurnEvent("turn/completed", "turn-1")); + const events = Array.from(yield* Fiber.join(observed)); + NodeAssert.deepEqual( + events.map((event) => event.type), + ["task.progress", "task.progress", "task.progress", "task.progress", "turn.completed"], + ); + NodeAssert.deepEqual( + events.filter(isCodexTaskRuntimeEvent).map((event) => event.payload.taskId), + ["child-b", "child-b", "child-a", "child-a"], + ); + const completion = events.at(-1); + NodeAssert.equal( + completion?.type === "turn.completed" && completion.payload.state, + "completed", + ); + NodeAssert.equal(completion?.turnId, "turn-1"); + }), + ); + + it.effect("flushes unexpected session exit and ignores later child progress", () => + Effect.gen(function* () { + const { adapter, runtime } = yield* startLifecycleRuntime(); + const exitSeen = yield* Deferred.make(); + const lateInputSeen = yield* Deferred.make(); + const observed = yield* adapter.streamEvents.pipe( + Stream.tap((event) => + event.type === "session.exited" + ? Deferred.succeed(exitSeen, undefined).pipe(Effect.asVoid) + : event.type === "thread.metadata.updated" && + event.payload.name === "late-input-processed" + ? Deferred.succeed(lateInputSeen, undefined).pipe(Effect.asVoid) + : Effect.void, + ), + Stream.takeUntil( + (event) => + event.type === "thread.metadata.updated" && event.payload.name === "after-clock", + ), + Stream.runCollect, + Effect.forkChild, + ); + const eventBase = { + kind: "notification" as const, + provider: ProviderDriverKind.make("codex"), + threadId: asThreadId("thread-1"), + createdAt: "2026-01-01T00:00:00.000Z", + }; + for (const query of ["before-exit", "after-exit"]) { + yield* runtime.emit({ + ...eventBase, + id: asEventId(`evt-${query}`), + method: "collabAgent/item", + payload: { agentThreadId: "exited-child", item: { type: "webSearch", query } }, + }); + if (query === "before-exit") { + yield* runtime.emit({ + ...eventBase, + id: asEventId("evt-unexpected-exit"), + method: "session/exited", + }); + yield* Deferred.await(exitSeen); + } + } + yield* runtime.emit({ + ...eventBase, + id: asEventId("evt-late-input-processed"), + method: "thread/name/updated", + payload: { threadId: "provider-thread-1", threadName: "late-input-processed" }, + }); + yield* Deferred.await(lateInputSeen); + yield* TestClock.adjust("1 second"); + yield* runtime.emit({ + ...eventBase, + id: asEventId("evt-after-clock"), + method: "thread/name/updated", + payload: { threadId: "provider-thread-1", threadName: "after-clock" }, + }); + const events = Array.from(yield* Fiber.join(observed)); + NodeAssert.deepEqual( + events.map((event) => event.type), + ["task.progress", "session.exited", "thread.metadata.updated", "thread.metadata.updated"], + ); + NodeAssert.equal( + events[0]?.type === "task.progress" && events[0].payload.summary, + "before-exit", + ); + }), + ); + // Production calls startSession from a request fiber that finishes as soon as // the session exists. `Effect.forkChild` made the runtime event consumer a // child of that fiber, and Effect interrupts a fiber's children when it diff --git a/apps/server/src/provider/Layers/CodexAdapter.ts b/apps/server/src/provider/Layers/CodexAdapter.ts index d1981b33d47d..bbba712fe7ff 100644 --- a/apps/server/src/provider/Layers/CodexAdapter.ts +++ b/apps/server/src/provider/Layers/CodexAdapter.ts @@ -70,6 +70,10 @@ import { type CodexSessionRuntimeShape, } from "./CodexSessionRuntime.ts"; import { type EventNdjsonLogger, makeEventNdjsonLogger } from "./EventNdjsonLogger.ts"; +import { + type CodexProgressCoalescer, + makeCodexProgressCoalescer, +} from "./CodexProgressCoalescer.ts"; import { resolveCodexLaunchArgs } from "./codexLaunchArgs.ts"; import { codexRateLimitsToUpdate } from "./codexUsageLimits.ts"; const isCodexAppServerProcessExitedError = Schema.is(CodexErrors.CodexAppServerProcessExitedError); @@ -99,6 +103,7 @@ interface CodexAdapterSessionContext { readonly threadId: ThreadId; readonly scope: Scope.Closeable; readonly runtime: CodexSessionRuntimeShape; + readonly progress: CodexProgressCoalescer; readonly eventFiber: Fiber.Fiber; readonly turnTokenUsage: CodexTurnTokenUsageState; stopped: boolean; @@ -1292,6 +1297,28 @@ function mapCollabAgentEvent( } } +function collabAgentThreadId(event: ProviderEvent): string | undefined { + if (event.kind !== "notification" || !event.method.startsWith("collabAgent/")) { + return undefined; + } + const payload = + typeof event.payload === "object" && event.payload !== null + ? (event.payload as Record) + : undefined; + return typeof payload?.agentThreadId === "string" ? payload.agentThreadId : undefined; +} + +function collabProgressLane(event: ProviderEvent): "item" | "tokenUsage" | undefined { + switch (event.method) { + case "collabAgent/item": + return "item"; + case "collabAgent/tokenUsage": + return "tokenUsage"; + default: + return undefined; + } +} + function mapToRuntimeEvents( event: ProviderEvent, canonicalThreadId: ThreadId, @@ -2302,6 +2329,9 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( }), ), ); + const progress = yield* makeCodexProgressCoalescer({ + emit: (events) => Queue.offerAll(runtimeEventQueue, events).pipe(Effect.asVoid), + }).pipe(Effect.provideService(Scope.Scope, sessionScope)); // Fork into the session scope, not the calling fiber. `forkChild` makes // this a child of `startSession`, and Effect interrupts a fiber's @@ -2368,6 +2398,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( } return runtimeEvent; }); + const childThreadId = collabAgentThreadId(event); if (runtimeEvents.length === 0) { yield* Effect.logDebug("ignoring unhandled Codex provider event", { method: event.method, @@ -2377,6 +2408,34 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( }); return; } + + const progressLane = collabProgressLane(event); + const latestProgress = runtimeEvents.length === 1 ? runtimeEvents[0] : undefined; + if (childThreadId && progressLane && latestProgress?.type === "task.progress") { + if (progressLane === "item") { + yield* progress.offerItem(childThreadId, latestProgress); + } else { + yield* progress.offerTokenUsage(childThreadId, latestProgress); + } + return; + } + + if (childThreadId) { + yield* progress.flush(childThreadId); + } else if ( + runtimeEvents.some( + (runtimeEvent) => + runtimeEvent.type === "turn.completed" || + runtimeEvent.type === "turn.aborted" || + runtimeEvent.type === "runtime.error" || + runtimeEvent.type === "session.exited", + ) + ) { + yield* progress.flushAll; + } + if (runtimeEvents.some((runtimeEvent) => runtimeEvent.type === "session.exited")) { + yield* progress.close; + } yield* Queue.offerAll(runtimeEventQueue, runtimeEvents); }), ).pipe(Effect.forkIn(sessionScope)); @@ -2392,7 +2451,8 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( }), ), Effect.onError(() => - runtime.close.pipe( + progress.close.pipe( + Effect.andThen(runtime.close), Effect.andThen(Effect.ignore(Scope.close(sessionScope, Exit.void))), Effect.andThen(Fiber.interrupt(eventFiber)), Effect.ignore, @@ -2404,6 +2464,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( threadId: input.threadId, scope: sessionScope, runtime, + progress, eventFiber, turnTokenUsage, stopped: false, @@ -2612,6 +2673,7 @@ export const makeCodexAdapter = Effect.fn("makeCodexAdapter")(function* ( } session.stopped = true; sessions.delete(session.threadId); + yield* session.progress.close; yield* session.runtime.close.pipe(Effect.ignore); yield* Effect.ignore(Scope.close(session.scope, Exit.void)); yield* Fiber.interrupt(session.eventFiber).pipe(Effect.ignore); diff --git a/apps/server/src/provider/Layers/CodexProgressCoalescer.test.ts b/apps/server/src/provider/Layers/CodexProgressCoalescer.test.ts new file mode 100644 index 000000000000..3af40f5eff35 --- /dev/null +++ b/apps/server/src/provider/Layers/CodexProgressCoalescer.test.ts @@ -0,0 +1,292 @@ +import { assert, describe, it } from "@effect/vitest"; +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 Scope from "effect/Scope"; +import * as TestClock from "effect/testing/TestClock"; + +import { makeCodexProgressCoalescer } from "./CodexProgressCoalescer.ts"; + +describe("makeCodexProgressCoalescer", () => { + it.effect("uses fixed windows and emits only the newest value from each lane", () => + Effect.scoped( + Effect.gen(function* () { + const emissions: Array> = []; + const coalescer = yield* makeCodexProgressCoalescer({ + emit: (values) => + Effect.sync(() => { + emissions.push([...values]); + }), + }); + + yield* coalescer.offerItem("child", "item-1"); + yield* TestClock.adjust("100 millis"); + yield* coalescer.offerItem("child", "item-2"); + yield* coalescer.offerTokenUsage("child", "usage-1"); + yield* TestClock.adjust("149 millis"); + yield* coalescer.offerTokenUsage("child", "usage-2"); + + assert.deepStrictEqual(emissions, []); + yield* TestClock.adjust("1 milli"); + assert.deepStrictEqual(emissions, [["item-2", "usage-2"]]); + + yield* coalescer.offerItem("child", "item-3"); + yield* coalescer.offerTokenUsage("child", "usage-3"); + yield* TestClock.adjust("100 millis"); + yield* coalescer.offerItem("child", "item-4"); + yield* TestClock.adjust("150 millis"); + + assert.deepStrictEqual(emissions, [ + ["item-2", "usage-2"], + ["item-4", "usage-3"], + ]); + assert.isTrue(emissions.every((values) => values.length <= 2)); + }), + ), + ); + + it.live("does not let a blocked key prevent another key from flushing", () => + Effect.scoped( + Effect.gen(function* () { + const emissions: Array> = []; + const keyAStarted = yield* Deferred.make(); + const releaseKeyA = yield* Deferred.make(); + const keyBEmitted = yield* Deferred.make(); + const coalescer = yield* makeCodexProgressCoalescer({ + window: "1 hour", + emit: (values) => + Effect.gen(function* () { + if (values[0] === "a") { + yield* Deferred.succeed(keyAStarted, undefined); + yield* Deferred.await(releaseKeyA); + } + emissions.push([...values]); + if (values[0] === "b") { + yield* Deferred.succeed(keyBEmitted, undefined); + } + }), + }); + + yield* coalescer.offerItem("key-a", "a"); + yield* coalescer.offerItem("key-b", "b"); + const keyAFlush = yield* coalescer + .flush("key-a") + .pipe(Effect.forkChild({ startImmediately: true })); + yield* Deferred.await(keyAStarted); + + const keyBFlush = yield* coalescer + .flush("key-b") + .pipe(Effect.forkChild({ startImmediately: true })); + yield* Deferred.await(keyBEmitted); + yield* Fiber.join(keyBFlush); + assert.deepStrictEqual(emissions, [["b"]]); + + yield* Deferred.succeed(releaseKeyA, undefined); + yield* Fiber.join(keyAFlush); + assert.deepStrictEqual(emissions, [["b"], ["a"]]); + }), + ), + ); + + it.effect("closes with its owning scope and ignores all later work", () => + Effect.gen(function* () { + const emissions: Array> = []; + const ownerScope = yield* Scope.make("sequential"); + const coalescer = yield* makeCodexProgressCoalescer({ + window: "1 second", + emit: (values) => + Effect.sync(() => { + emissions.push([...values]); + }), + }).pipe(Effect.provideService(Scope.Scope, ownerScope)); + + yield* coalescer.offerItem("child", "pending-item"); + yield* coalescer.offerTokenUsage("child", "pending-usage"); + yield* Scope.close(ownerScope, Exit.void); + yield* TestClock.adjust("1 second"); + yield* coalescer.offerItem("child", "late-item"); + yield* coalescer.offerTokenUsage("child", "late-usage"); + yield* coalescer.flush("child"); + yield* coalescer.flushAll; + yield* coalescer.close; + + assert.deepStrictEqual(emissions, []); + }), + ); + + it.effect("serializes a tick with flush and invalidates the flushed timer", () => + Effect.scoped( + Effect.gen(function* () { + const emissions: Array> = []; + const tickStarted = yield* Deferred.make(); + const releaseTick = yield* Deferred.make(); + const coalescer = yield* makeCodexProgressCoalescer({ + emit: (values) => + Effect.gen(function* () { + emissions.push([...values]); + if (values[0] === "race") { + yield* Deferred.succeed(tickStarted, undefined); + yield* Deferred.await(releaseTick); + } + }), + }); + + yield* coalescer.offerItem("child", "race"); + const advanceToTick = yield* TestClock.adjust("250 millis").pipe( + Effect.forkChild({ startImmediately: true }), + ); + yield* Deferred.await(tickStarted); + const racingFlush = yield* coalescer + .flush("child") + .pipe(Effect.forkChild({ startImmediately: true })); + yield* Deferred.succeed(releaseTick, undefined); + yield* Fiber.join(advanceToTick); + yield* Fiber.join(racingFlush); + + assert.strictEqual(emissions.flat().filter((value) => value === "race").length, 1); + + yield* coalescer.offerItem("child", "terminal"); + yield* TestClock.adjust("249 millis"); + yield* coalescer.flush("child"); + yield* coalescer.offerItem("child", "post-flush"); + yield* TestClock.adjust("1 milli"); + assert.isFalse(emissions.flat().includes("post-flush")); + + yield* TestClock.adjust("249 millis"); + assert.strictEqual(emissions.flat().filter((value) => value === "post-flush").length, 1); + assert.deepStrictEqual(emissions, [["race"], ["terminal"], ["post-flush"]]); + }), + ), + ); + + it.effect("keeps pending values isolated between coalescer instances", () => + Effect.scoped( + Effect.gen(function* () { + const firstEmissions: Array> = []; + const secondEmissions: Array> = []; + const first = yield* makeCodexProgressCoalescer({ + window: "1 hour", + emit: (values) => + Effect.sync(() => { + firstEmissions.push([...values]); + }), + }); + const second = yield* makeCodexProgressCoalescer({ + window: "1 hour", + emit: (values) => + Effect.sync(() => { + secondEmissions.push([...values]); + }), + }); + + yield* first.offerItem("same-key", "first-instance"); + yield* second.offerItem("same-key", "second-instance"); + yield* Effect.all([first.flush("same-key"), second.flush("same-key")], { + concurrency: "unbounded", + }); + + assert.deepStrictEqual(firstEmissions, [["first-instance"]]); + assert.deepStrictEqual(secondEmissions, [["second-instance"]]); + }), + ), + ); + + it.effect("flushes all children in first-seen order with item before usage", () => + Effect.scoped( + Effect.gen(function* () { + const emissions: Array> = []; + const coalescer = yield* makeCodexProgressCoalescer({ + emit: (values) => + Effect.sync(() => { + emissions.push([...values]); + }), + }); + yield* coalescer.offerTokenUsage("second-name", "usage-a-old"); + yield* coalescer.offerItem("first-name", "item-b-old"); + yield* coalescer.offerItem("second-name", "item-a"); + yield* coalescer.offerTokenUsage("second-name", "usage-a"); + yield* coalescer.offerTokenUsage("first-name", "usage-b"); + yield* coalescer.offerItem("first-name", "item-b"); + + yield* coalescer.flushAll; + yield* coalescer.flushAll; + yield* TestClock.adjust("250 millis"); + assert.deepStrictEqual(emissions, [ + ["item-a", "usage-a"], + ["item-b", "usage-b"], + ]); + }), + ), + ); + + it.effect("close interrupts an in-flight timer while flushAll waits for its key", () => + Effect.scoped( + Effect.gen(function* () { + const emissions: Array> = []; + const tickStarted = yield* Deferred.make(); + const tickInterrupted = yield* Deferred.make(); + const keepTickPending = yield* Deferred.make(); + const coalescer = yield* makeCodexProgressCoalescer({ + emit: (values) => + Effect.gen(function* () { + yield* Deferred.succeed(tickStarted, undefined); + yield* Deferred.await(keepTickPending); + emissions.push([...values]); + }).pipe(Effect.ensuring(Deferred.succeed(tickInterrupted, undefined))), + }); + yield* coalescer.offerItem("child", "timer-value"); + const advanceToTick = yield* TestClock.adjust("250 millis").pipe( + Effect.forkChild({ startImmediately: true }), + ); + yield* Deferred.await(tickStarted); + const flushing = yield* coalescer.flushAll.pipe( + Effect.forkChild({ startImmediately: true }), + ); + const closing = yield* coalescer.close.pipe(Effect.forkChild({ startImmediately: true })); + const alsoClosing = yield* coalescer.close.pipe( + Effect.forkChild({ startImmediately: true }), + ); + yield* Deferred.await(tickInterrupted); + yield* Fiber.join(closing); + yield* Fiber.join(alsoClosing); + yield* Fiber.join(flushing); + yield* Fiber.join(advanceToTick); + yield* coalescer.offerItem("child", "late-value"); + yield* coalescer.flushAll; + yield* TestClock.adjust("250 millis"); + assert.deepStrictEqual(emissions, []); + }), + ), + ); + + it.effect("close waits for an explicit flush and discards the other pending child", () => + Effect.scoped( + Effect.gen(function* () { + const emissions: Array> = []; + const flushStarted = yield* Deferred.make(); + const releaseFlush = yield* Deferred.make(); + const coalescer = yield* makeCodexProgressCoalescer({ + emit: (values) => + Effect.gen(function* () { + yield* Deferred.succeed(flushStarted, undefined); + yield* Deferred.await(releaseFlush); + emissions.push([...values]); + }), + }); + yield* coalescer.offerItem("first", "flushing-value"); + yield* coalescer.offerItem("second", "discarded-value"); + const flushing = yield* coalescer.flushAll.pipe( + Effect.forkChild({ startImmediately: true }), + ); + yield* Deferred.await(flushStarted); + const closing = yield* coalescer.close.pipe(Effect.forkChild({ startImmediately: true })); + yield* Deferred.succeed(releaseFlush, undefined); + yield* Fiber.join(flushing); + yield* Fiber.join(closing); + yield* TestClock.adjust("250 millis"); + assert.deepStrictEqual(emissions, [["flushing-value"]]); + }), + ), + ); +}); diff --git a/apps/server/src/provider/Layers/CodexProgressCoalescer.ts b/apps/server/src/provider/Layers/CodexProgressCoalescer.ts new file mode 100644 index 000000000000..c0a5f9b9cf2f --- /dev/null +++ b/apps/server/src/provider/Layers/CodexProgressCoalescer.ts @@ -0,0 +1,239 @@ +import * as Deferred from "effect/Deferred"; +import * as Duration from "effect/Duration"; +import * as Effect from "effect/Effect"; +import * as Exit from "effect/Exit"; +import * as Fiber from "effect/Fiber"; +import * as Option from "effect/Option"; +import * as Scope from "effect/Scope"; +import * as Semaphore from "effect/Semaphore"; +import * as SynchronizedRef from "effect/SynchronizedRef"; + +export interface CodexProgressCoalescer { + readonly offerItem: (key: K, value: V) => Effect.Effect; + readonly offerTokenUsage: (key: K, value: V) => Effect.Effect; + readonly flush: (key: K) => Effect.Effect; + readonly flushAll: Effect.Effect; + readonly close: Effect.Effect; +} + +interface Worker { + readonly token: object; + readonly fiber: Fiber.Fiber; +} + +interface KeyState { + readonly lock: Semaphore.Semaphore; + item: Option.Option; + tokenUsage: Option.Option; + worker: Worker | null; +} + +interface RegistryState { + readonly closed: boolean; + readonly keys: Map>; +} + +export const makeCodexProgressCoalescer = Effect.fn("makeCodexProgressCoalescer")(function* < + K, + V, +>(options: { + readonly window?: Duration.Input; + readonly emit: (values: ReadonlyArray) => Effect.Effect; +}): Effect.fn.Return, never, Scope.Scope> { + const window = options.window ?? Duration.millis(250); + const workerScope = yield* Scope.make("parallel"); + const closeDone = yield* Deferred.make(); + const registryRef = yield* SynchronizedRef.make>({ + closed: false, + keys: new Map(), + }); + + const takeValues = (state: KeyState): ReadonlyArray => { + const values: Array = []; + if (Option.isSome(state.item)) { + values.push(state.item.value); + } + if (Option.isSome(state.tokenUsage)) { + values.push(state.tokenUsage.value); + } + state.item = Option.none(); + state.tokenUsage = Option.none(); + return values; + }; + + const clearWorker = (state: KeyState, token: object) => + state.lock.withPermit( + Effect.sync(() => { + if (state.worker?.token === token) { + state.worker = null; + } + }), + ); + + const tick = (state: KeyState, token: object) => + state.lock.withPermit( + Effect.gen(function* () { + if (state.worker?.token !== token) { + return; + } + if ((yield* SynchronizedRef.get(registryRef)).closed) { + takeValues(state); + return; + } + + const values = takeValues(state); + if (values.length > 0) { + yield* options.emit(values); + } + }).pipe( + Effect.ensuring( + Effect.sync(() => { + if (state.worker?.token === token) { + state.worker = null; + } + }), + ), + ), + ); + + const startWorker = Effect.fn("CodexProgressCoalescer.startWorker")(function* ( + state: KeyState, + ) { + const token = {}; + const fiber = yield* Effect.sleep(window).pipe( + Effect.andThen(tick(state, token)), + Effect.ensuring(clearWorker(state, token)), + Effect.forkIn(workerScope, { startImmediately: true }), + ); + state.worker = { token, fiber }; + }); + + const getOrCreateKeyState = (key: K): Effect.Effect | null> => + SynchronizedRef.modifyEffect(registryRef, (registry) => { + if (registry.closed) { + return Effect.succeed([null, registry] as const); + } + const existing = registry.keys.get(key); + if (existing !== undefined) { + return Effect.succeed([existing, registry] as const); + } + return Semaphore.make(1).pipe( + Effect.map((lock) => { + const state: KeyState = { + lock, + item: Option.none(), + tokenUsage: Option.none(), + worker: null, + }; + const keys = new Map(registry.keys); + keys.set(key, state); + return [state, { closed: false, keys }] as const; + }), + ); + }); + + const getExistingKeyState = (key: K): Effect.Effect | null> => + SynchronizedRef.get(registryRef).pipe( + Effect.map((registry) => (registry.closed ? null : (registry.keys.get(key) ?? null))), + ); + + const offer = + (lane: "item" | "tokenUsage") => + (key: K, value: V): Effect.Effect => + Effect.gen(function* () { + const state = yield* getOrCreateKeyState(key); + if (state === null) { + return; + } + yield* state.lock.withPermit( + Effect.gen(function* () { + if ((yield* SynchronizedRef.get(registryRef)).closed) { + return; + } + state[lane] = Option.some(value); + if (state.worker === null) { + yield* startWorker(state); + } + }), + ); + }); + + const flush = Effect.fn("CodexProgressCoalescer.flush")(function* (key: K) { + const state = yield* getExistingKeyState(key); + if (state === null) { + return; + } + + let workerToInterrupt: Fiber.Fiber | null = null; + yield* state.lock + .withPermit( + Effect.gen(function* () { + if ((yield* SynchronizedRef.get(registryRef)).closed) { + return; + } + workerToInterrupt = state.worker?.fiber ?? null; + state.worker = null; + const values = takeValues(state); + if (values.length > 0) { + yield* options.emit(values); + } + }), + ) + .pipe( + Effect.ensuring( + Effect.suspend(() => + workerToInterrupt === null ? Effect.void : Fiber.interrupt(workerToInterrupt), + ), + ), + ); + }); + + const flushAll = Effect.gen(function* () { + const registry = yield* SynchronizedRef.get(registryRef); + yield* Effect.forEach(registry.keys.keys(), flush, { concurrency: 1, discard: true }); + }); + + const close = Effect.uninterruptible( + Effect.gen(function* () { + const states = yield* SynchronizedRef.modify(registryRef, (registry) => + registry.closed + ? ([null, registry] as const) + : ([Array.from(registry.keys.values()), { ...registry, closed: true }] as const), + ); + if (states === null) { + yield* Deferred.await(closeDone); + return; + } + + yield* Scope.close(workerScope, Exit.void); + const lateWorkers = yield* Effect.forEach( + states, + (state) => + state.lock.withPermit( + Effect.sync(() => { + const worker = state.worker?.fiber ?? null; + state.worker = null; + takeValues(state); + return worker; + }), + ), + { concurrency: "unbounded" }, + ); + yield* Fiber.interruptAll(lateWorkers.filter((fiber) => fiber !== null)); + yield* SynchronizedRef.update(registryRef, (registry) => ({ + ...registry, + keys: new Map(), + })); + yield* Deferred.succeed(closeDone, undefined); + }), + ); + + yield* Effect.addFinalizer(() => close); + return { + offerItem: offer("item"), + offerTokenUsage: offer("tokenUsage"), + flush, + flushAll, + close, + }; +});