remove getFinalAssistantMessage from nodejs test harness

Author: qmuntal

nodejs/test/e2e/harness/sdkTestHelper.ts

@@ -2,78 +2,7 @@
  *  Copyright (c) Microsoft Corporation. All rights reserved.
  *--------------------------------------------------------------------------------------------*/
 
-import { AssistantMessageEvent, CopilotSession, SessionEvent } from "../../../src";
-
-export async function getFinalAssistantMessage(
-    session: CopilotSession
-): Promise<AssistantMessageEvent> {
-    // We don't know whether the answer has already arrived or not, so race both possibilities
-    return new Promise<AssistantMessageEvent>(async (resolve, reject) => {
-        getFutureFinalResponse(session).then(resolve).catch(reject);
-        getExistingFinalResponse(session)
-            .then((msg) => {
-                if (msg) {
-                    resolve(msg);
-                }
-            })
-            .catch(reject);
-    });
-}
-
-function getExistingFinalResponse(
-    session: CopilotSession
-): Promise<AssistantMessageEvent | undefined> {
-    return new Promise<AssistantMessageEvent | undefined>(async (resolve, reject) => {
-        const messages = await session.getMessages();
-        const finalUserMessageIndex = messages.findLastIndex((m) => m.type === "user.message");
-        const currentTurnMessages =
-            finalUserMessageIndex < 0 ? messages : messages.slice(finalUserMessageIndex);
-
-        const currentTurnError = currentTurnMessages.find((m) => m.type === "session.error");
-        if (currentTurnError) {
-            const error = new Error(currentTurnError.data.message);
-            error.stack = currentTurnError.data.stack;
-            reject(error);
-            return;
-        }
-
-        const sessionIdleMessageIndex = currentTurnMessages.findIndex(
-            (m) => m.type === "session.idle"
-        );
-        if (sessionIdleMessageIndex !== -1) {
-            const lastAssistantMessage = currentTurnMessages
-                .slice(0, sessionIdleMessageIndex)
-                .findLast((m) => m.type === "assistant.message");
-            resolve(lastAssistantMessage as AssistantMessageEvent | undefined);
-            return;
-        }
-
-        resolve(undefined);
-    });
-}
-
-function getFutureFinalResponse(session: CopilotSession): Promise<AssistantMessageEvent> {
-    return new Promise<AssistantMessageEvent>((resolve, reject) => {
-        let finalAssistantMessage: AssistantMessageEvent | undefined;
-        session.on((event) => {
-            if (event.type === "assistant.message") {
-                finalAssistantMessage = event;
-            } else if (event.type === "session.idle") {
-                if (!finalAssistantMessage) {
-                    reject(
-                        new Error("Received session.idle without a preceding assistant.message")
-                    );
-                } else {
-                    resolve(finalAssistantMessage);
-                }
-            } else if (event.type === "session.error") {
-                const error = new Error(event.data.message);
-                error.stack = event.data.stack;
-                reject(error);
-            }
-        });
-    });
-}
+import { CopilotSession, SessionEvent } from "../../../src";
 
 export async function retry(
     message: string,

nodejs/test/e2e/session.test.ts

@@ -1,8 +1,8 @@
 import { describe, expect, it, onTestFinished } from "vitest";
 import { ParsedHttpExchange } from "../../../test/harness/replayingCapiProxy.js";
-import { CopilotClient } from "../../src/index.js";
+import { CopilotClient, SessionEvent } from "../../src/index.js";
 import { CLI_PATH, createSdkTestContext } from "./harness/sdkTestContext.js";
-import { getFinalAssistantMessage, getNextEventOfType } from "./harness/sdkTestHelper.js";
+import { getNextEventOfType } from "./harness/sdkTestHelper.js";
 
 describe("Sessions", async () => {
     const { copilotClient: client, openAiEndpoint, homeDir, env } = await createSdkTestContext();
@@ -167,7 +167,7 @@ describe("Sessions", async () => {
         expect(session2.sessionId).toBe(sessionId);
 
         // TODO: There's an inconsistency here. When resuming with a new client, we don't see
-        // the session.idle message in the history, which means we can't use getFinalAssistantMessage.
+        // the session.idle message in the history, so we can't easily identify when a turn completed.
 
         const messages = await session2.getMessages();
         expect(messages).toContainEqual(expect.objectContaining({ type: "user.message" }));
@@ -328,9 +328,8 @@ describe("Sessions", async () => {
         expect(session.sessionId).toMatch(/^[a-f0-9-]+$/);
 
         // Session should work normally with custom config dir
-        await session.send({ prompt: "What is 1+1?" });
-        const assistantMessage = await getFinalAssistantMessage(session);
-        expect(assistantMessage.data.content).toContain("2");
+        const assistantMessage = await session.sendAndWait({ prompt: "What is 1+1?" });
+        expect(assistantMessage?.data.content).toContain("2");
     });
 });
 
@@ -348,23 +347,28 @@ describe("Send Blocking Behavior", async () => {
     it("send returns immediately while events stream in background", async () => {
         const session = await client.createSession();
 
-        const events: string[] = [];
-        session.on((event) => {
-            events.push(event.type);
-        });
+        const events: Array<SessionEvent> = [];
+        session.on((event) => events.push(event));
+
+        // Set up promise to wait for idle BEFORE sending
+        const idlePromise = getNextEventOfType(session, "session.idle");
 
         // Use a slow command so we can verify send() returns before completion
         await session.send({ prompt: "Run 'sleep 2 && echo done'" });
 
         // send() should return before turn completes (no session.idle yet)
-        expect(events).not.toContain("session.idle");
+        expect(events.some((e) => e.type === "session.idle")).toBe(false);
 
         // Wait for turn to complete
-        const message = await getFinalAssistantMessage(session);
+        await idlePromise;
 
-        expect(message.data.content).toContain("done");
-        expect(events).toContain("session.idle");
-        expect(events).toContain("assistant.message");
+        // Find the last assistant message from collected events
+        const assistantMessages = events.filter((e) => e.type === "assistant.message");
+        const message = assistantMessages[assistantMessages.length - 1];
+
+        expect(message.data?.content).toContain("done");
+        expect(events.some((e) => e.type === "session.idle")).toBe(true);
+        expect(events.some((e) => e.type === "assistant.message")).toBe(true);
     });
 
     it("sendAndWait blocks until session.idle and returns final assistant message", async () => {