diff --git a/nodejs/src/session.ts b/nodejs/src/session.ts index 48e41483b..d614be4a4 100644 --- a/nodejs/src/session.ts +++ b/nodejs/src/session.ts @@ -750,6 +750,21 @@ export class CopilotSession { const outcomePromise = new Promise((resolve) => { resolveOutcome = resolve; }); + let resolveShutdown: (error: Error) => void; + const shutdownPromise = new Promise((resolve) => { + resolveShutdown = resolve; + }); + let outcomeSettled = false; + const settleOutcome = (outcome: SessionOutcome, shutdown = false): void => { + if (outcomeSettled) { + return; + } + outcomeSettled = true; + resolveOutcome(outcome); + if (shutdown && outcome.kind === "error") { + resolveShutdown(outcome.error); + } + }; let lastAssistantMessage: AssistantMessageEvent | undefined; @@ -759,17 +774,29 @@ export class CopilotSession { if (event.type === "assistant.message") { lastAssistantMessage = event; } else if (event.type === "session.idle") { - resolveOutcome({ kind: "idle" }); + settleOutcome({ kind: "idle" }); } else if (event.type === "session.error") { const error = new Error(event.data.message); error.stack = event.data.stack; - resolveOutcome({ kind: "error", error }); + settleOutcome({ kind: "error", error }); + } else if (event.type === "session.shutdown") { + const reason = event.data.errorReason ? `: ${event.data.errorReason}` : ""; + const error = new Error( + `Session ${this.sessionId} shut down before becoming idle${reason}` + ); + settleOutcome({ kind: "error", error }, true); } }); let timeoutId: ReturnType | undefined; try { - await this.send(options); + const sendResult = await Promise.race([ + this.send(options).then(() => ({ kind: "sent" }) as const), + shutdownPromise.then((error) => ({ kind: "shutdown", error }) as const), + ]); + if (sendResult.kind === "shutdown") { + throw sendResult.error; + } const timeoutPromise = new Promise((_, reject) => { timeoutId = setTimeout( diff --git a/nodejs/test/session-send-and-wait.test.ts b/nodejs/test/session-send-and-wait.test.ts index 8b6e390c4..ab7c0e9f6 100644 --- a/nodejs/test/session-send-and-wait.test.ts +++ b/nodejs/test/session-send-and-wait.test.ts @@ -29,6 +29,23 @@ function errorEvent(message: string): SessionEvent { } as SessionEvent; } +function shutdownEvent(errorReason?: string): SessionEvent { + return { + type: "session.shutdown", + id: "00000000-0000-4000-8000-000000000002", + parentId: null, + timestamp: new Date().toISOString(), + data: { + codeChanges: { filesModified: [], linesAdded: 0, linesRemoved: 0 }, + errorReason, + modelMetrics: {}, + sessionStartTime: Date.now(), + shutdownType: errorReason ? "error" : "routine", + totalApiDurationMs: 0, + }, + } as SessionEvent; +} + function controlledSession(): { session: CopilotSession; sendStarted: Promise; @@ -95,6 +112,7 @@ describe("sendAndWait", () => { await sendStarted; session._dispatchEvent(sessionEvent("session.idle")); + session._dispatchEvent(shutdownEvent("later shutdown")); const stateBeforeSend = await Promise.race([ pending.then(() => "settled"), @@ -134,4 +152,16 @@ describe("sendAndWait", () => { errorFirst.resolveSend(); await expect(errorFirstPending).rejects.toThrow("first error"); }); + + it("rejects when the session shuts down before becoming idle", async () => { + const { session, sendStarted } = controlledSession(); + const pending = session.sendAndWait({ prompt: "hi" }); + await sendStarted; + + session._dispatchEvent(shutdownEvent("runtime session was lost")); + + await expect(pending).rejects.toThrow( + "Session session-1 shut down before becoming idle: runtime session was lost" + ); + }); });