From 461be47c9124aad8033664f5a77b3afd66dbf002 Mon Sep 17 00:00:00 2001 From: Matt Cowger Date: Sun, 12 Jul 2026 01:20:13 -0700 Subject: [PATCH] fix(server): await OpenCode event stream shutdown (#2015) * fix(server): await OpenCode event stream shutdown * Update opencode-agent.test.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- .../agent/providers/opencode-agent.test.ts | 73 ++++++++++++++++++- .../server/agent/providers/opencode-agent.ts | 26 ++++++- .../test-utils/test-opencode-harness.ts | 37 +++++++++- 3 files changed, 132 insertions(+), 4 deletions(-) diff --git a/packages/server/src/server/agent/providers/opencode-agent.test.ts b/packages/server/src/server/agent/providers/opencode-agent.test.ts index 139a63637..44902b2cc 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.test.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.test.ts @@ -1197,6 +1197,7 @@ describe("OpenCode adapter startTurn error handling", () => { test("keeps a turn active while OpenCode is retrying", async () => { vi.useFakeTimers(); const eventsGate = createTestDeferred(); + let eventStreamSignal: AbortSignal | undefined; const retryStream: AsyncIterable = { [Symbol.asyncIterator]: () => { let emitted = false; @@ -1222,14 +1223,27 @@ describe("OpenCode adapter startTurn error handling", () => { }, }; } - return new Promise(() => {}); + return new Promise>((resolve) => { + if (eventStreamSignal?.aborted) { + resolve({ done: true, value: undefined }); + return; + } + eventStreamSignal?.addEventListener( + "abort", + () => resolve({ done: true, value: undefined }), + { once: true }, + ); + }); }, }; }, }; const fakeClient = { global: { - event: vi.fn().mockResolvedValue({ stream: retryStream }), + event: vi.fn().mockImplementation(async ({ signal }: { signal: AbortSignal }) => { + eventStreamSignal = signal; + return { stream: retryStream }; + }), }, session: { abort: vi.fn().mockResolvedValue({ error: null }), @@ -1319,6 +1333,52 @@ describe("OpenCode adapter startTurn error handling", () => { expect(fakeClient.session.delete).not.toHaveBeenCalled(); }); + test("waits for the OpenCode event stream to finish after close aborts it", async () => { + const streamAborted = createTestDeferred(); + const finishStreamCleanup = createTestDeferred(); + const fakeClient = { + global: { + event: vi.fn().mockImplementation(async ({ signal }: { signal: AbortSignal }) => ({ + stream: { + [Symbol.asyncIterator]: () => ({ + next: async () => { + if (!signal.aborted) { + await waitForAbort(signal); + } + streamAborted.resolve(); + await finishStreamCleanup.promise; + return { done: true, value: undefined }; + }, + }), + }, + })), + }, + session: { + abort: vi.fn().mockResolvedValue({ error: null }), + update: vi.fn().mockResolvedValue({ error: null }), + }, + } as never; + const session = new __openCodeInternals.OpenCodeAgentSession( + { provider: "opencode", cwd: "/tmp/test" }, + fakeClient, + "ses_unit_test", + createTestLogger(), + ); + let closeSettled = false; + + const closePromise = session.close().then(() => { + closeSettled = true; + return undefined; + }); + await streamAborted.promise; + await new Promise((resolve) => setImmediate(resolve)); + + expect(closeSettled).toBe(false); + + finishStreamCleanup.resolve(); + await closePromise; + }); + test("streamHistory preserves OpenCode replay timestamps from message and part times", async () => { const fakeClient = { session: { @@ -2178,6 +2238,15 @@ function createTestDeferred(): { return { promise, resolve, reject }; } +function waitForAbort(signal: AbortSignal): Promise { + if (signal.aborted) { + return Promise.resolve(); + } + return new Promise((resolve) => { + signal.addEventListener("abort", () => resolve(), { once: true }); + }); +} + function abortableOpenCodeStream(signal: AbortSignal): AsyncIterable { return { [Symbol.asyncIterator]: () => { diff --git a/packages/server/src/server/agent/providers/opencode-agent.ts b/packages/server/src/server/agent/providers/opencode-agent.ts index 0cd170764..2d093d365 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -2789,6 +2789,7 @@ class OpenCodeAgentSession implements AgentSession { private releaseServer: (() => void) | null; private eventStreamAbortController: AbortController | null = null; private eventStreamReady: Deferred | null = null; + private eventStreamTask: Promise | null = null; private suppressTerminalUntilNextUserMessage = false; private closed = false; private readonly persistSession: boolean; @@ -3160,11 +3161,24 @@ class OpenCodeAgentSession implements AgentSession { const eventStreamReady = createDeferred(); this.eventStreamAbortController = eventStreamAbortController; this.eventStreamReady = eventStreamReady; - void this.consumeEventStream(eventStreamAbortController, eventStreamReady).finally(() => { + const eventStreamTask = this.consumeEventStream( + eventStreamAbortController, + eventStreamReady, + ).finally(() => { if (this.eventStreamAbortController === eventStreamAbortController) { this.eventStreamAbortController = null; this.eventStreamReady = null; } + if (this.eventStreamTask === eventStreamTask) { + this.eventStreamTask = null; + } + }); + this.eventStreamTask = eventStreamTask; + void eventStreamTask.catch((error) => { + this.logger.warn( + { err: error, sessionId: this.sessionId }, + "OpenCode event stream task failed", + ); }); return eventStreamReady.promise; @@ -3561,9 +3575,19 @@ class OpenCodeAgentSession implements AgentSession { // unhandled rejection in whichever test the daemon hops to next. this.closed = true; this.abortController?.abort(); + const eventStreamTask = this.eventStreamTask; this.eventStreamAbortController?.abort(); + if (eventStreamTask) { + await eventStreamTask.catch((error) => { + this.logger.debug( + { err: error, sessionId: this.sessionId }, + "OpenCode event stream failed during close", + ); + }); + } this.eventStreamAbortController = null; this.eventStreamReady = null; + this.eventStreamTask = null; this.subscribers.clear(); await reconcileOpenCodeSessionClose({ client: this.client, diff --git a/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-harness.ts b/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-harness.ts index 0410c8abb..3935b10f0 100644 --- a/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-harness.ts +++ b/packages/server/src/server/agent/providers/opencode/test-utils/test-opencode-harness.ts @@ -156,7 +156,10 @@ export class TestOpenCodeClient { global: { event: async (options: unknown) => { this.calls.globalEvent.push(options); - return { stream: this.eventStream }; + const signal = (options as { signal?: AbortSignal }).signal; + return { + stream: signal ? stopEventStreamOnAbort(this.eventStream, signal) : this.eventStream, + }; }, }, mcp: { @@ -247,6 +250,38 @@ export class TestOpenCodeClient { } } +function stopEventStreamOnAbort( + stream: AsyncIterable, + signal: AbortSignal, +): AsyncIterable { + return { + [Symbol.asyncIterator]: () => { + const iterator = stream[Symbol.asyncIterator](); + return { + next: () => { + if (signal.aborted) { + return Promise.resolve({ done: true, value: undefined }); + } + return new Promise>((resolve, reject) => { + const onAbort = () => resolve({ done: true, value: undefined }); + signal.addEventListener("abort", onAbort, { once: true }); + void iterator.next().then( + (result) => { + signal.removeEventListener("abort", onAbort); + return resolve(result); + }, + (error) => { + signal.removeEventListener("abort", onAbort); + return reject(error); + }, + ); + }); + }, + }; + }, + }; +} + export function createEventStream(events: unknown[]): AsyncGenerator { return (async function* () { for (const event of events) {