abort.e2e.test.ts

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *--------------------------------------------------------------------------------------------*/

import { describe, expect, it } from "vitest";
import { z } from "zod";
import { approveAll, defineTool } from "../../src/index.js";
import { createSdkTestContext } from "./harness/sdkTestContext.js";

describe("Abort", async () => {
    const { copilotClient: client } = await createSdkTestContext();
    const TEST_TIMEOUT_MS = 120_000;

    async function withTimeout<T>(promise: Promise<T>, ms: number, label: string): Promise<T> {
        let timer: ReturnType<typeof setTimeout> | undefined;
        try {
            return await Promise.race([
                promise,
                new Promise<T>((_, reject) => {
                    timer = setTimeout(() => reject(new Error(`Timeout: ${label}`)), ms);
                }),
            ]);
        } finally {
            if (timer) clearTimeout(timer);
        }
    }

    it("should abort during active streaming", { timeout: TEST_TIMEOUT_MS }, async () => {
        const session = await client.createSession({
            onPermissionRequest: approveAll,
            streaming: true,
        });

        let firstDeltaResolve!: (value: void) => void;
        const firstDeltaReceived = new Promise<void>((resolve) => {
            firstDeltaResolve = resolve;
        });

        const events: { type: string }[] = [];
        session.on((event) => {
            events.push({ type: event.type });
            if (event.type === "assistant.message_delta") {
                firstDeltaResolve();
            }
        });

        // Fire-and-forget — we'll abort before it finishes
        void session.send({
            prompt: "Write a very long essay about the history of computing, covering every decade from the 1940s to the 2020s in great detail.",
        });

        // Wait for at least one delta to arrive (proves streaming started)
        await withTimeout(firstDeltaReceived, 60_000, "first assistant.message_delta");

        const deltaEvents = events.filter((e) => e.type === "assistant.message_delta");
        expect(deltaEvents.length).toBeGreaterThanOrEqual(1);

        // Abort mid-stream
        await session.abort();

        // Session should be usable after abort. Wait for the specific recovery
        // message rather than racing against a late idle from the aborted turn.
        let recoveryResolve!: (content: string) => void;
        const recoveryReceived = new Promise<string>((resolve) => {
            recoveryResolve = resolve;
        });
        const unsubscribeRecovery = session.on((event) => {
            if (event.type === "assistant.message") {
                const content = event.data.content ?? "";
                if (content.toLowerCase().includes("abort_recovery_ok")) {
                    recoveryResolve(content);
                }
            }
        });

        try {
            await session.send({ prompt: "Say 'abort_recovery_ok'." });
            const recoveryContent = await withTimeout(
                recoveryReceived,
                60_000,
                "assistant.message containing abort_recovery_ok"
            );
            expect(recoveryContent.toLowerCase()).toContain("abort_recovery_ok");
        } finally {
            unsubscribeRecovery();
        }

        await session.disconnect();
    });

    it("should abort during active tool execution", { timeout: TEST_TIMEOUT_MS }, async () => {
        let toolStartedResolve!: (value: string) => void;
        const toolStarted = new Promise<string>((resolve) => {
            toolStartedResolve = resolve;
        });

        let releaseToolResolve!: (value: string) => void;
        const releaseTool = new Promise<string>((resolve) => {
            releaseToolResolve = resolve;
        });

        let signalAbortedResolve!: () => void;
        const signalAborted = new Promise<void>((resolve) => {
            signalAbortedResolve = resolve;
        });

        const session = await client.createSession({
            onPermissionRequest: approveAll,
            tools: [
                defineTool("slow_analysis", {
                    description: "A slow analysis tool that blocks until released",
                    parameters: z.object({
                        value: z.string().describe("Value to analyze"),
                    }),
                    handler: async ({ value }, { signal }) => {
                        toolStartedResolve(value);
                        if (signal.aborted) {
                            signalAbortedResolve();
                        } else {
                            signal.addEventListener("abort", () => signalAbortedResolve(), {
                                once: true,
                            });
                        }
                        return await releaseTool;
                    },
                }),
            ],
        });

        // Fire-and-forget
        void session.send({
            prompt: "Use slow_analysis with value 'test_abort'. Wait for the result.",
        });

        // Wait for the tool to start executing
        const toolValue = await withTimeout(toolStarted, 60_000, "slow_analysis start");
        expect(toolValue).toBe("test_abort");

        // Abort while the tool is running
        await session.abort();

        // The handler's AbortSignal should fire as a result of session.abort()
        await withTimeout(signalAborted, 10_000, "tool handler AbortSignal");

        // Release the tool so its task doesn't leak
        releaseToolResolve("RELEASED_AFTER_ABORT");

        // Session should be usable after abort — verify with a follow-up
        let recoveryResolve!: (value: void) => void;
        const recoveryReceived = new Promise<void>((resolve) => {
            recoveryResolve = resolve;
        });

        session.on((event) => {
            if (
                event.type === "assistant.message" &&
                event.data.content?.includes("tool_abort_recovery_ok")
            ) {
                recoveryResolve();
            }
        });

        void session.send({
            prompt: "Say 'tool_abort_recovery_ok'.",
        });

        await withTimeout(recoveryReceived, 60_000, "tool abort recovery message");

        await session.disconnect();
    });

    it(
        "should cancel a single tool call via cancelToolCall",
        { timeout: TEST_TIMEOUT_MS },
        async () => {
            let toolCallIdResolve!: (value: string) => void;
            const toolCallIdReady = new Promise<string>((resolve) => {
                toolCallIdResolve = resolve;
            });

            let releaseToolResolve!: (value: string) => void;
            const releaseTool = new Promise<string>((resolve) => {
                releaseToolResolve = resolve;
            });

            let signalAbortedResolve!: () => void;
            const signalAborted = new Promise<void>((resolve) => {
                signalAbortedResolve = resolve;
            });

            const session = await client.createSession({
                onPermissionRequest: approveAll,
                tools: [
                    defineTool("slow_analysis", {
                        description: "A slow analysis tool that blocks until released",
                        parameters: z.object({
                            value: z.string().describe("Value to analyze"),
                        }),
                        handler: async ({ value: _value }, { signal, toolCallId }) => {
                            toolCallIdResolve(toolCallId);
                            if (signal.aborted) {
                                signalAbortedResolve();
                            } else {
                                signal.addEventListener("abort", () => signalAbortedResolve(), {
                                    once: true,
                                });
                            }
                            return await releaseTool;
                        },
                    }),
                ],
            });

            // Fire-and-forget
            void session.send({
                prompt: "Use slow_analysis with value 'test_cancel'. Wait for the result.",
            });

            // Wait for the tool to start executing and capture its toolCallId
            const toolCallId = await withTimeout(
                toolCallIdReady,
                60_000,
                "slow_analysis toolCallId"
            );

            // Unknown toolCallIds return false
            expect(session.cancelToolCall("nonexistent-tool-call-id")).toBe(false);

            // Cancelling the in-flight tool call returns true and fires its signal
            expect(session.cancelToolCall(toolCallId)).toBe(true);
            await withTimeout(signalAborted, 10_000, "tool handler AbortSignal via cancelToolCall");

            // A second cancel of the same (now-removed) call returns false
            expect(session.cancelToolCall(toolCallId)).toBe(false);

            // Release the tool so its task doesn't leak
            releaseToolResolve("RELEASED_AFTER_CANCEL");

            await session.disconnect();
        }
    );
});