evOo-local/packages/adapters/src/pi-runtime-midstep-steering-stress.test.ts

import type { AgentRunRequest, AgentSteeringMessage, ConnectorTool } from "@rakazo/adapter-kit";
import { beforeEach, describe, expect, it, vi } from "vitest";

const fakeAgentState = vi.hoisted(() => ({
  mode: "normal" as
    | "normal"
    | "long_tool"
    | "rapid_steer"
    | "tool_throw_abort"
    | "tool_throw_custom"
    | "subagent"
    | "parallel_tools",
  tools: [] as Array<{
    name: string;
    execute: (toolCallId: string, params: Record<string, unknown>) => Promise<unknown>;
  }>,
  steeredMessages: [] as unknown[],
  abortCount: 0,
  toolExecutionPromise: undefined as Promise<unknown> | undefined,
  toolResolve: undefined as ((val: unknown) => void) | undefined,
  toolReject: undefined as ((err: unknown) => void) | undefined,
  agentCount: 0,
  subagentStarted: false,
  subagentResolve: undefined as (() => void) | undefined,
}));

vi.mock("@earendil-works/pi-agent-core", () => ({
  Agent: class {
    state = { errorMessage: undefined as string | undefined, messages: [] as unknown[] };
    private readonly tools: typeof fakeAgentState.tools;
    private readonly listeners: Array<(event: Record<string, unknown>) => void> = [];
    private aborted = false;
    private isSubagent = false;

    constructor(options: {
      initialState: { tools: typeof fakeAgentState.tools; messages: unknown[] };
    }) {
      this.tools = options.initialState.tools;
      fakeAgentState.tools = this.tools;
      fakeAgentState.agentCount += 1;
      if (fakeAgentState.agentCount > 1) {
        this.isSubagent = true;
      }
    }

    subscribe(listener: (event: Record<string, unknown>) => void) {
      this.listeners.push(listener);
    }

    async prompt(_prompt: string) {
      if (fakeAgentState.mode === "parallel_tools") {
        const tool1 = this.tools.find((t) => t.name === "parallel_tool_1");
        const tool2 = this.tools.find((t) => t.name === "parallel_tool_2");
        if (tool1 && tool2) {
          this.emit({ type: "tool_execution_start", toolName: tool1.name, args: { test: 1 } });
          this.emit({ type: "tool_execution_start", toolName: tool2.name, args: { test: 2 } });
          const p1 = tool1.execute("call-parallel-1", { test: 1 });
          const p2 = tool2.execute("call-parallel-2", { test: 2 });
          await Promise.allSettled([p1, p2]);
          return;
        }
      }

      if (fakeAgentState.mode === "subagent") {
        if (!this.isSubagent) {
          const target = this.tools.find((t) => t.name === "run_subagent");
          if (target) {
            this.emit({
              type: "tool_execution_start",
              toolName: target.name,
              args: { task: "subtask" },
            });
            const toolPromise = target.execute("call-subagent-1", { task: "subtask" });
            fakeAgentState.toolExecutionPromise = toolPromise;
            try {
              await toolPromise;
            } catch {
              // Handled by tool
            }
            return;
          }
        } else {
          fakeAgentState.subagentStarted = true;
          await new Promise<void>((resolve) => {
            fakeAgentState.subagentResolve = resolve;
          });
          return;
        }
      }

      const target = this.tools.find((t) => t.name === "adversarial_tool") ?? this.tools[0];
      if (!target) return;

      this.emit({ type: "tool_execution_start", toolName: target.name, args: { test: true } });
      const toolPromise = target.execute("call-adv-1", { test: true });
      fakeAgentState.toolExecutionPromise = toolPromise;
      try {
        await toolPromise;
      } catch {
        // Handled by tool
      }
    }

    async waitForIdle() {}

    steer(message: unknown) {
      fakeAgentState.steeredMessages.push(message);
    }

    abort() {
      this.aborted = true;
      fakeAgentState.abortCount += 1;
      if (this.isSubagent) {
        fakeAgentState.subagentResolve?.();
      }
    }

    private emit(event: Record<string, unknown>) {
      for (const listener of this.listeners) listener(event);
    }
  },
}));

vi.mock("@earendil-works/pi-ai/providers/all", () => ({
  builtinModels: () => ({
    getModel: (_provider: string, modelId: string) =>
      modelId === "stress-test-model" ? { provider: "test", id: modelId } : undefined,
    streamSimple: () => {
      throw new Error("fake agent must not call real provider");
    },
  }),
}));

vi.mock("./pi-local-provider.js", () => ({
  registerLocalProvider: (models: unknown) => models,
}));

vi.mock("./pi-openai-compatible-provider.js", () => ({
  OPENAI_COMPATIBLE_PROVIDER_ID: "openai-compatible",
  registerOpenAiCompatibleCatalog: (models: unknown) => models,
  registerOpenAiCompatibleRuntime: (models: unknown) => models,
}));

import { pruneCancelledToolResultsContext } from "./pi-model-config.js";
import { PiAgentRuntime, raceSignal } from "./pi-runtime.js";

const adversarialToolDef: ConnectorTool = {
  name: "adversarial_tool",
  description: "A tool used for adversarial cancellation stress testing",
  inputSchema: { type: "object", properties: { test: { type: "boolean" } } },
};

const parallelTool1Def: ConnectorTool = {
  name: "parallel_tool_1",
  description: "First parallel tool for concurrency stress testing",
  inputSchema: { type: "object", properties: { test: { type: "number" } } },
};

const parallelTool2Def: ConnectorTool = {
  name: "parallel_tool_2",
  description: "Second parallel tool for concurrency stress testing",
  inputSchema: { type: "object", properties: { test: { type: "number" } } },
};

const subagentToolDef: ConnectorTool = {
  name: "run_subagent",
  description: "Run subagent tool for mid-step cancellation testing",
  inputSchema: {
    type: "object",
    properties: {
      task: { type: "string" },
      name: { type: "string" },
    },
  },
};

describe("Adversarial Stress: Mid-Step Tool Cancellation & Conversational Steering", () => {
  let runtime: PiAgentRuntime;

  beforeEach(() => {
    runtime = new PiAgentRuntime();
    fakeAgentState.mode = "normal";
    fakeAgentState.tools = [];
    fakeAgentState.steeredMessages = [];
    fakeAgentState.abortCount = 0;
    fakeAgentState.toolExecutionPromise = undefined;
    fakeAgentState.toolResolve = undefined;
    fakeAgentState.toolReject = undefined;
    fakeAgentState.agentCount = 0;
    fakeAgentState.subagentStarted = false;
    fakeAgentState.subagentResolve = undefined;
  });

  // =========================================================================
  // Objective 1A: Mid-Step Tool Cancellation via steer() and AbortSignal
  // =========================================================================
  it("interupts long-running tool via steer(), asserting AbortSignal abortion and cancelled details", async () => {
    let capturedSignal: AbortSignal | undefined;
    let toolExecutionStarted = false;
    let toolCompletedData: unknown;

    const executeTool = vi.fn(
      async (
        _name: string,
        _args: Record<string, unknown>,
        _id: string,
        _route?: unknown,
        signal?: AbortSignal,
      ) => {
        capturedSignal = signal;
        toolExecutionStarted = true;
        return new Promise((resolve, reject) => {
          fakeAgentState.toolResolve = resolve;
          fakeAgentState.toolReject = reject;
        });
      },
    );

    const onToolCompleted = vi.fn((completion) => {
      toolCompletedData = completion;
    });

    const runRequest: AgentRunRequest = {
      botId: "bot-1",
      threadId: "thread-1",
      runId: "run-midstep-1",
      prompt: "Execute adversarial test",
      instructions: "Follow rules",
      history: [],
      tools: [adversarialToolDef],
      model: { provider: "test", id: "stress-test-model" },
      executeTool,
      onToolCompleted,
    };

    const runStream = runtime.run(runRequest);
    const iterator = runStream[Symbol.asyncIterator]();

    // Consume first event so tool execution begins
    const firstEventPromise = iterator.next();

    // Wait until tool execution is started
    while (!toolExecutionStarted) {
      await new Promise((r) => setTimeout(r, 5));
    }

    expect(capturedSignal).toBeDefined();
    expect(capturedSignal?.aborted).toBe(false);

    // Now steer mid-flight while tool is executing
    const steered = await runtime.steer("run-midstep-1", "Redirect to different task");
    expect(steered).toBe(true);

    // Verify underlying AbortSignal was immediately aborted
    expect(capturedSignal?.aborted).toBe(true);
    expect(capturedSignal?.reason).toBe("steering_redirection");

    // Let the iterator drain
    await firstEventPromise;
    await iterator.return?.();

    // Verify agent received steered message
    expect(fakeAgentState.steeredMessages).toHaveLength(1);
    expect(fakeAgentState.steeredMessages[0]).toMatchObject({
      role: "user",
      content: "Redirect to different task",
    });

    // Verify onToolCompleted received non-fatal cancelled result
    expect(toolCompletedData).toBeDefined();
    const completion = toolCompletedData as {
      result?: { content?: unknown[]; details?: Record<string, unknown> };
    };
    expect(completion.result?.details?.cancelled).toBe(true);
    expect(completion.result?.details?.reason).toBe("steering_redirection");
  });

  // =========================================================================
  // Objective 1B: Rapid Contradictory Steering Messages Under High Concurrency
  // =========================================================================
  it("handles rapid contradictory steering bursts (20 steering messages) without crash or lockup", async () => {
    let capturedSignal: AbortSignal | undefined;
    let toolExecutionStarted = false;

    const executeTool = vi.fn(
      async (
        _name: string,
        _args: Record<string, unknown>,
        _id: string,
        _route?: unknown,
        signal?: AbortSignal,
      ) => {
        capturedSignal = signal;
        toolExecutionStarted = true;
        return new Promise(() => {}); // never resolves on its own
      },
    );

    const runRequest: AgentRunRequest = {
      botId: "bot-1",
      threadId: "thread-1",
      runId: "run-rapid-burst",
      prompt: "Execute burst test",
      instructions: "Follow rules",
      history: [],
      tools: [adversarialToolDef],
      model: { provider: "test", id: "stress-test-model" },
      executeTool,
    };

    const runStream = runtime.run(runRequest);
    const iterator = runStream[Symbol.asyncIterator]();
    void iterator.next();

    while (!toolExecutionStarted) {
      await new Promise((r) => setTimeout(r, 5));
    }

    // Fire 20 rapid contradictory steering messages concurrently
    const NUM_BURSTS = 20;
    const burstPromises = Array.from({ length: NUM_BURSTS }, (_, i) =>
      runtime.steer("run-rapid-burst", `Contradictory intent #${i}`),
    );

    const steerResults = await Promise.all(burstPromises);
    expect(steerResults.every((res) => res === true)).toBe(true);

    // Verify tool was aborted
    expect(capturedSignal?.aborted).toBe(true);

    // Verify all 20 messages delivered to agent
    expect(fakeAgentState.steeredMessages).toHaveLength(NUM_BURSTS);
    for (let i = 0; i < NUM_BURSTS; i++) {
      expect(fakeAgentState.steeredMessages[i]).toMatchObject({
        role: "user",
        content: `Contradictory intent #${i}`,
      });
    }

    await iterator.return?.();
  });

  // =========================================================================
  // Objective 1C: Tool Finishes at Exact Microsecond of Abort (Microtask Races)
  // =========================================================================
  describe("Microsecond Tool Resolution vs Abort Races", () => {
    it("Microtask Race 3A: abort fires synchronously right after tool resolves", async () => {
      const controller = new AbortController();
      let resolveTool!: (val: string) => void;
      const toolPromise = new Promise<string>((res) => {
        resolveTool = res;
      });

      const raced = raceSignal(toolPromise, controller.signal);
      resolveTool("tool-output");
      controller.abort("steering_redirection");

      // Synchronous abort listener preempts the queued promise.then() microtask
      await expect(raced).rejects.toBe("steering_redirection");
    });

    it("Microtask Race 3B: tool resolves and microtask ticks before abort arrives", async () => {
      const controller = new AbortController();
      let resolveTool!: (val: string) => void;
      const toolPromise = new Promise<string>((res) => {
        resolveTool = res;
      });

      const raced = raceSignal(toolPromise, controller.signal);
      resolveTool("tool-output");
      await Promise.resolve(); // Drain microtask queue
      controller.abort("steering_redirection");

      // Promise resolution settled before abort arrived
      await expect(raced).resolves.toBe("tool-output");
    });

    it("Microtask Race 3C: fuzz test with 50 randomized simultaneous microtask races", async () => {
      const NUM_RACES = 50;
      const results = await Promise.all(
        Array.from({ length: NUM_RACES }, async (_, idx) => {
          const controller = new AbortController();
          let resolveFn!: (val: string) => void;
          const promise = new Promise<string>((resolve) => {
            resolveFn = resolve;
          });

          const raced = raceSignal(promise, controller.signal);

          // Randomize resolution and abort timing across microtasks
          const resolveMicrotasks = Math.floor(Math.random() * 3);
          const abortMicrotasks = Math.floor(Math.random() * 3);

          void (async () => {
            for (let i = 0; i < resolveMicrotasks; i++) await Promise.resolve();
            resolveFn(`result-${idx}`);
          })();

          void (async () => {
            for (let i = 0; i < abortMicrotasks; i++) await Promise.resolve();
            controller.abort(`abort-${idx}`);
          })();

          try {
            const val = await raced;
            return { status: "resolved", val };
          } catch (err) {
            return { status: "aborted", err };
          }
        }),
      );

      expect(results).toHaveLength(NUM_RACES);
      // All must have cleanly settled to either resolved or aborted without crashing
      for (const res of results) {
        if (res.status === "resolved") {
          expect(res.val).toMatch(/^result-\d+$/);
        } else {
          expect(res.err).toMatch(/^abort-\d+$/);
        }
      }
    });
  });

  // =========================================================================
  // Objective 1D: Tool Throwing Custom Exceptions or AbortError upon Cancellation
  // =========================================================================
  describe("Adversarial Exception Handling on Cancellation", () => {
    it("handles tool throwing DOMException AbortError cleanly without runtime crash", async () => {
      let toolExecuted = false;
      const executeTool = vi.fn(
        async (
          _name: string,
          _args: Record<string, unknown>,
          _id: string,
          _route?: unknown,
          signal?: AbortSignal,
        ) => {
          toolExecuted = true;
          return new Promise((_resolve, reject) => {
            signal?.addEventListener("abort", () => {
              reject(new DOMException("The operation was aborted.", "AbortError"));
            });
          });
        },
      );

      let completedResult: unknown;
      const runRequest: AgentRunRequest = {
        botId: "bot-1",
        threadId: "thread-1",
        runId: "run-abort-error",
        prompt: "Run abort error test",
        instructions: "Follow rules",
        history: [],
        tools: [adversarialToolDef],
        model: { provider: "test", id: "stress-test-model" },
        executeTool,
        onToolCompleted: (completion) => {
          completedResult = completion.result;
        },
      };

      const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
      void iterator.next();

      while (!toolExecuted) {
        await new Promise((r) => setTimeout(r, 5));
      }

      // Steer to abort the tool
      await runtime.steer("run-abort-error", "Pivot from aborted tool");
      await new Promise((r) => setTimeout(r, 50));
      await iterator.return?.();

      expect(completedResult).toBeDefined();
      const res = completedResult as { details?: Record<string, unknown> };
      expect(res.details?.cancelled).toBe(true);
    });

    it("handles tool throwing custom domain exception cleanly when aborted", async () => {
      let toolExecuted = false;
      const executeTool = vi.fn(
        async (
          _name: string,
          _args: Record<string, unknown>,
          _id: string,
          _route?: unknown,
          signal?: AbortSignal,
        ) => {
          toolExecuted = true;
          return new Promise((_resolve, reject) => {
            signal?.addEventListener("abort", () => {
              reject(new Error("FATAL_PROCESS_EXIT_SIGNAL_15"));
            });
          });
        },
      );

      let completedResult: unknown;
      const runRequest: AgentRunRequest = {
        botId: "bot-1",
        threadId: "thread-1",
        runId: "run-custom-error",
        prompt: "Run custom error test",
        instructions: "Follow rules",
        history: [],
        tools: [adversarialToolDef],
        model: { provider: "test", id: "stress-test-model" },
        executeTool,
        onToolCompleted: (completion) => {
          completedResult = completion.result;
        },
      };

      const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
      void iterator.next();

      while (!toolExecuted) {
        await new Promise((r) => setTimeout(r, 5));
      }

      await runtime.steer("run-custom-error", "Pivot from crashed tool");
      await new Promise((r) => setTimeout(r, 50));
      await iterator.return?.();

      expect(completedResult).toBeDefined();
      const res = completedResult as { details?: Record<string, unknown> };
      expect(res.details?.cancelled).toBe(true);
    });

    it("handles tool rejecting with non-Error primitive (string/object) cleanly", async () => {
      let toolExecuted = false;
      const executeTool = vi.fn(
        async (
          _name: string,
          _args: Record<string, unknown>,
          _id: string,
          _route?: unknown,
          signal?: AbortSignal,
        ) => {
          toolExecuted = true;
          return new Promise((_resolve, reject) => {
            signal?.addEventListener("abort", () => {
              reject("raw_string_rejection_504");
            });
          });
        },
      );

      let completedResult: unknown;
      const runRequest: AgentRunRequest = {
        botId: "bot-1",
        threadId: "thread-1",
        runId: "run-primitive-rejection",
        prompt: "Run primitive rejection test",
        instructions: "Follow rules",
        history: [],
        tools: [adversarialToolDef],
        model: { provider: "test", id: "stress-test-model" },
        executeTool,
        onToolCompleted: (completion) => {
          completedResult = completion.result;
        },
      };

      const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
      void iterator.next();

      while (!toolExecuted) {
        await new Promise((r) => setTimeout(r, 5));
      }

      await runtime.steer("run-primitive-rejection", "Pivot from primitive rejection");
      await new Promise((r) => setTimeout(r, 50));
      await iterator.return?.();

      expect(completedResult).toBeDefined();
      const res = completedResult as { details?: Record<string, unknown> };
      expect(res.details?.cancelled).toBe(true);
    });
  });

  // =========================================================================
  // Objective 1E: In-Flight Tool Interruption via claimSteering Polling
  // =========================================================================
  it("polls and ingests steering mid-step via claimSteering, aborting tool and deduplicating seen IDs", async () => {
    let capturedSignal: AbortSignal | undefined;
    let toolExecutionStarted = false;

    const executeTool = vi.fn(
      async (
        _name: string,
        _args: Record<string, unknown>,
        _id: string,
        _route?: unknown,
        signal?: AbortSignal,
      ) => {
        capturedSignal = signal;
        toolExecutionStarted = true;
        return new Promise(() => {}); // never finishes on its own
      },
    );

    let steeringAvailable = false;
    const steeringQueue: AgentSteeringMessage[] = [
      { id: "steer-msg-1", messageId: "steer-msg-1", text: "Pivot now via claimSteering" },
    ];

    const claimCalls: string[][] = [];
    const claimSteering = vi.fn(async (seenIds: string[]) => {
      claimCalls.push([...seenIds]);
      if (!steeringAvailable) return [];
      if (seenIds.includes("steer-msg-1")) return [];
      return steeringQueue;
    });

    const runRequest: AgentRunRequest = {
      botId: "bot-1",
      threadId: "thread-1",
      runId: "run-claim-steering",
      prompt: "Execute claim steering test",
      instructions: "Follow rules",
      history: [],
      tools: [adversarialToolDef],
      model: { provider: "test", id: "stress-test-model" },
      executeTool,
      claimSteering,
    };

    const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
    void iterator.next();

    while (!toolExecutionStarted) {
      await new Promise((r) => setTimeout(r, 5));
    }

    // Steering message arrives during active tool execution
    steeringAvailable = true;

    // Wait for the 80ms polling timer to trigger claimSteering
    await new Promise((r) => setTimeout(r, 200));

    expect(capturedSignal?.aborted).toBe(true);
    expect(capturedSignal?.reason).toBe("steering_redirection");
    expect(claimSteering).toHaveBeenCalled();

    // Verify message was steered to the agent
    expect(fakeAgentState.steeredMessages).toContainEqual(
      expect.objectContaining({
        role: "user",
        content: "Pivot now via claimSteering",
      }),
    );

    await iterator.return?.();
  });

  // =========================================================================
  // Objective 1F: Complete Context Sanitization (pruneCancelledToolResultsContext)
  // =========================================================================
  describe("Context Sanitization: pruneCancelledToolResultsContext", () => {
    it("prunes solitary cancelled tool result and turns assistant tool call into placeholder", () => {
      const messages: any[] = [
        { role: "user", content: "Search for info" },
        {
          role: "assistant",
          content: [
            { type: "toolCall", id: "call-cancelled-1", name: "adversarial_tool", arguments: {} },
          ],
        },
        {
          role: "toolResult",
          toolCallId: "call-cancelled-1",
          content: [{ type: "text", text: "[Tool cancelled]" }],
          details: { cancelled: true, reason: "steering_redirection" },
        },
        { role: "user", content: "Redirected request" },
      ];

      const pruned = pruneCancelledToolResultsContext(messages);

      expect(pruned).toHaveLength(3);
      // ToolResult must be completely stripped
      expect(pruned.find((m) => m.role === "toolResult")).toBeUndefined();

      // Assistant message must replace cancelled toolCall with [Cancelled tool step] placeholder
      const assistant = pruned.find((m) => m.role === "assistant") as any;
      expect(assistant).toBeDefined();
      expect(assistant.content).toEqual([{ type: "text", text: "[Cancelled tool step]" }]);
    });

    it("prunes only the cancelled toolCall in a parallel execution, keeping completed toolCall and result intact", () => {
      const messages: any[] = [
        { role: "user", content: "Run parallel queries" },
        {
          role: "assistant",
          content: [
            { type: "toolCall", id: "call-1", name: "tool_1", arguments: {} },
            { type: "toolCall", id: "call-2", name: "tool_2", arguments: {} },
          ],
        },
        {
          role: "toolResult",
          toolCallId: "call-1",
          content: [{ type: "text", text: "Tool 1 success" }],
          details: { cancelled: false, data: "ok" },
        },
        {
          role: "toolResult",
          toolCallId: "call-2",
          content: [{ type: "text", text: "[Cancelled]" }],
          details: { cancelled: true },
        },
      ];

      const pruned = pruneCancelledToolResultsContext(messages);

      expect(pruned).toHaveLength(3);
      // Only call-1 toolResult survives
      const results = pruned.filter((m) => m.role === "toolResult") as any[];
      expect(results).toHaveLength(1);
      expect(results[0]?.toolCallId).toBe("call-1");

      // Assistant message keeps call-1 and strips call-2
      const assistant = pruned.find((m) => m.role === "assistant") as any;
      expect(assistant).toBeDefined();
      expect(assistant.content).toEqual([
        { type: "toolCall", id: "call-1", name: "tool_1", arguments: {} },
      ]);
    });

    it("preserves assistant message commentary text when removing cancelled toolCall", () => {
      const messages: any[] = [
        {
          role: "assistant",
          content: [
            { type: "text", text: "I will now search for the latest documentation." },
            { type: "toolCall", id: "call-search", name: "search", arguments: {} },
          ],
        },
        {
          role: "toolResult",
          toolCallId: "call-search",
          content: [{ type: "text", text: "[Cancelled]" }],
          details: { cancelled: true },
        },
      ];

      const pruned = pruneCancelledToolResultsContext(messages);

      expect(pruned).toHaveLength(1);
      const assistant = pruned[0] as any;
      expect(assistant).toBeDefined();
      expect(assistant.content).toEqual([
        { type: "text", text: "I will now search for the latest documentation." },
      ]);
    });

    it("passes through messages unchanged if no tool results were cancelled", () => {
      const messages: any[] = [
        { role: "user", content: "Hello" },
        {
          role: "assistant",
          content: [{ type: "toolCall", id: "call-ok", name: "tool_ok", arguments: {} }],
        },
        {
          role: "toolResult",
          toolCallId: "call-ok",
          content: [{ type: "text", text: "Success" }],
          details: { success: true },
        },
      ];

      const pruned = pruneCancelledToolResultsContext(messages);
      expect(pruned).toBe(messages); // Strict reference equality for no-op optimization
    });

    it("prunes all tool calls when multiple tools in one assistant turn are all cancelled", () => {
      const messages: any[] = [
        {
          role: "assistant",
          content: [
            { type: "toolCall", id: "call-c1", name: "tool_c1", arguments: {} },
            { type: "toolCall", id: "call-c2", name: "tool_c2", arguments: {} },
            { type: "toolCall", id: "call-c3", name: "tool_c3", arguments: {} },
          ],
        },
        { role: "toolResult", toolCallId: "call-c1", content: [], details: { cancelled: true } },
        { role: "toolResult", toolCallId: "call-c2", content: [], details: { cancelled: true } },
        { role: "toolResult", toolCallId: "call-c3", content: [], details: { cancelled: true } },
      ];

      const pruned = pruneCancelledToolResultsContext(messages);
      expect(pruned).toHaveLength(1);
      const first = pruned[0] as any;
      expect(first).toBeDefined();
      expect(first.role).toBe("assistant");
      expect(first.content).toEqual([{ type: "text", text: "[Cancelled tool step]" }]);
    });
  });

  // =========================================================================
  // Objective 1G: Late Underlying Rejection & Legacy Tool Arity (< 5)
  // =========================================================================
  describe("Late Rejections & Legacy Tool Arity", () => {
    it("safely handles underlying tool rejecting late in the background after raceSignal cancellation", async () => {
      let lateRejectTriggered = false;
      const executeTool = vi.fn(
        async (
          _name: string,
          _args: Record<string, unknown>,
          _id: string,
          _route?: unknown,
          _signal?: AbortSignal,
        ) => {
          return new Promise((_resolve, reject) => {
            setTimeout(() => {
              lateRejectTriggered = true;
              reject(new Error("LATE_BACKGROUND_FAILURE"));
            }, 30);
          });
        },
      );

      let completedResult: unknown;
      const runRequest: AgentRunRequest = {
        botId: "bot-1",
        threadId: "thread-1",
        runId: "run-late-reject",
        prompt: "Run late reject test",
        instructions: "Follow rules",
        history: [],
        tools: [adversarialToolDef],
        model: { provider: "test", id: "stress-test-model" },
        executeTool,
        onToolCompleted: (completion) => {
          completedResult = completion.result;
        },
      };

      const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
      void iterator.next();

      await new Promise((r) => setTimeout(r, 10));
      await runtime.steer("run-late-reject", "Abort now before late reject");

      // Wait for late reject to fire
      await new Promise((r) => setTimeout(r, 60));
      expect(lateRejectTriggered).toBe(true);

      await iterator.return?.();

      expect(completedResult).toBeDefined();
      const res = completedResult as { details?: Record<string, unknown> };
      expect(res.details?.cancelled).toBe(true);
    });

    it("cancels legacy tools with executeTool arity < 5 via raceSignal without crashing", async () => {
      let legacyToolStarted = false;
      // 3-parameter signature
      const legacyExecuteTool = vi.fn(
        async (_name: string, _args: Record<string, unknown>, _id: string) => {
          legacyToolStarted = true;
          return new Promise(() => {}); // never finishes
        },
      );

      let completedResult: unknown;
      const runRequest: AgentRunRequest = {
        botId: "bot-1",
        threadId: "thread-1",
        runId: "run-legacy-tool",
        prompt: "Run legacy tool test",
        instructions: "Follow rules",
        history: [],
        tools: [adversarialToolDef],
        model: { provider: "test", id: "stress-test-model" },
        executeTool: legacyExecuteTool,
        onToolCompleted: (completion) => {
          completedResult = completion.result;
        },
      };

      const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
      void iterator.next();

      while (!legacyToolStarted) {
        await new Promise((r) => setTimeout(r, 5));
      }

      await runtime.steer("run-legacy-tool", "Abort legacy tool");
      await new Promise((r) => setTimeout(r, 30));
      await iterator.return?.();

      expect(completedResult).toBeDefined();
      const res = completedResult as { details?: Record<string, unknown> };
      expect(res.details?.cancelled).toBe(true);
    });
  });

  // =========================================================================
  // Objective 1H: High Concurrency Multi-Run Interleaved Stress Test
  // =========================================================================
  it("executes 25 parallel agent runs with randomized concurrent steering, tool execution, and cancellation", async () => {
    const NUM_PARALLEL = 25;
    const runResults = await Promise.all(
      Array.from({ length: NUM_PARALLEL }, async (_, idx) => {
        const runId = `parallel-run-${idx}`;
        let toolStarted = false;
        let toolCancelled = false;

        const executeTool = vi.fn(
          async (
            _name: string,
            _args: Record<string, unknown>,
            _id: string,
            _route?: unknown,
            signal?: AbortSignal,
          ) => {
            toolStarted = true;
            return new Promise((resolve, reject) => {
              signal?.addEventListener("abort", () => {
                toolCancelled = true;
                reject(new DOMException("Aborted", "AbortError"));
              });
              setTimeout(() => resolve({ ok: true }), 100);
            });
          },
        );

        let finalResult: unknown;
        const runRequest: AgentRunRequest = {
          botId: `bot-${idx}`,
          threadId: `thread-${idx}`,
          runId,
          prompt: `Run parallel ${idx}`,
          instructions: "Rules",
          history: [],
          tools: [adversarialToolDef],
          model: { provider: "test", id: "stress-test-model" },
          executeTool,
          onToolCompleted: (completion) => {
            finalResult = completion.result;
          },
        };

        const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
        void iterator.next();

        while (!toolStarted) {
          await new Promise((r) => setTimeout(r, 2));
        }

        // Steer on even indices, let odd indices finish
        const shouldSteer = idx % 2 === 0;
        if (shouldSteer) {
          await runtime.steer(runId, `Steer parallel ${idx}`);
        }

        await new Promise((r) => setTimeout(r, 120));
        await iterator.return?.();

        return { idx, shouldSteer, toolCancelled, finalResult };
      }),
    );

    expect(runResults).toHaveLength(NUM_PARALLEL);
    for (const res of runResults) {
      if (res.shouldSteer) {
        expect(res.toolCancelled).toBe(true);
        expect((res.finalResult as any)?.details?.cancelled).toBe(true);
      } else {
        expect(res.toolCancelled).toBe(false);
      }
    }
  });

  // =========================================================================
  // Objective 1I: Concurrent In-Flight Tool Abort via activeToolControllers
  // =========================================================================
  it("aborts all concurrent in-flight tools in a single turn when steered mid-flight", async () => {
    fakeAgentState.mode = "parallel_tools";

    let tool1Started = false;
    let tool2Started = false;
    let tool1Signal: AbortSignal | undefined;
    let tool2Signal: AbortSignal | undefined;
    const completedResults: Record<string, unknown> = {};

    const executeTool = vi.fn(
      async (
        name: string,
        _args: Record<string, unknown>,
        _id: string,
        _route?: unknown,
        signal?: AbortSignal,
      ) => {
        if (name === "parallel_tool_1") {
          tool1Started = true;
          tool1Signal = signal;
          return new Promise((resolve, reject) => {
            signal?.addEventListener("abort", () => {
              reject(new DOMException("Aborted", "AbortError"));
            });
            setTimeout(() => resolve({ name: "tool1_done" }), 1000);
          });
        }
        if (name === "parallel_tool_2") {
          tool2Started = true;
          tool2Signal = signal;
          return new Promise((resolve, reject) => {
            signal?.addEventListener("abort", () => {
              reject(new DOMException("Aborted", "AbortError"));
            });
            setTimeout(() => resolve({ name: "tool2_done" }), 1000);
          });
        }
        return { ok: true };
      },
    );

    const runId = "concurrent-steer-test";
    const runRequest: AgentRunRequest = {
      botId: "bot-concurrent",
      threadId: "thread-concurrent",
      runId,
      prompt: "Execute concurrent tools",
      instructions: "Rules",
      history: [],
      tools: [parallelTool1Def, parallelTool2Def],
      model: { provider: "test", id: "stress-test-model" },
      executeTool,
      onToolCompleted: (completion) => {
        completedResults[completion.name] = completion.result;
      },
    };

    const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
    void iterator.next();

    while (!tool1Started || !tool2Started) {
      await new Promise((r) => setTimeout(r, 5));
    }

    expect(tool1Signal?.aborted).toBe(false);
    expect(tool2Signal?.aborted).toBe(false);

    await runtime.steer(runId, "Redirect concurrent tools mid-flight");
    await new Promise((r) => setTimeout(r, 50));
    await iterator.return?.();

    expect(tool1Signal?.aborted).toBe(true);
    expect(tool1Signal?.reason).toBe("steering_redirection");
    expect(tool2Signal?.aborted).toBe(true);
    expect(tool2Signal?.reason).toBe("steering_redirection");

    const res1 = completedResults.parallel_tool_1 as
      | { details?: Record<string, unknown> }
      | undefined;
    const res2 = completedResults.parallel_tool_2 as
      | { details?: Record<string, unknown> }
      | undefined;
    expect(res1?.details?.cancelled).toBe(true);
    expect(res1?.details?.reason).toBe("steering_redirection");
    expect(res2?.details?.cancelled).toBe(true);
    expect(res2?.details?.reason).toBe("steering_redirection");

    // Verify context pruner removes both cancelled tool calls
    const messages = [
      {
        role: "assistant",
        content: [
          { type: "toolCall", id: "call-parallel-1", name: "parallel_tool_1", arguments: {} },
          { type: "toolCall", id: "call-parallel-2", name: "parallel_tool_2", arguments: {} },
        ],
      },
      {
        role: "toolResult",
        toolCallId: "call-parallel-1",
        content: [
          { type: "text", text: "[Tool parallel_tool_1 cancelled: user steering redirection]" },
        ],
        details: { cancelled: true, reason: "steering_redirection" },
      },
      {
        role: "toolResult",
        toolCallId: "call-parallel-2",
        content: [
          { type: "text", text: "[Tool parallel_tool_2 cancelled: user steering redirection]" },
        ],
        details: { cancelled: true, reason: "steering_redirection" },
      },
    ];
    const pruned = pruneCancelledToolResultsContext(messages as any);
    expect(pruned).toHaveLength(1);
    expect(pruned[0]).toEqual({
      role: "assistant",
      content: [{ type: "text", text: "[Cancelled tool step]" }],
    });
  });

  // =========================================================================
  // Objective 1J: Subagent Mid-Step Cancellation & Context Pruning
  // =========================================================================
  it("cancels subagent mid-step on steering, emits stopped event, returns cancelled details, and prunes context", async () => {
    fakeAgentState.mode = "subagent";

    const completedResults: Record<string, unknown> = {};
    const events: unknown[] = [];

    const runId = "subagent-cancellation-test";
    const runRequest: AgentRunRequest = {
      botId: "bot-subagent",
      threadId: "thread-subagent",
      runId,
      prompt: "Execute subagent delegation",
      instructions: "Rules",
      history: [],
      tools: [subagentToolDef],
      model: { provider: "test", id: "stress-test-model" },
      onToolCompleted: (completion) => {
        completedResults[completion.name] = completion.result;
      },
    };

    const iterator = runtime.run(runRequest)[Symbol.asyncIterator]();
    const eventConsumer = (async () => {
      for await (const event of { [Symbol.asyncIterator]: () => iterator }) {
        events.push(event);
      }
    })();

    while (!fakeAgentState.subagentStarted) {
      await new Promise((r) => setTimeout(r, 5));
    }

    await runtime.steer(runId, "Cancel subagent and steer");
    await new Promise((r) => setTimeout(r, 50));
    await iterator.return?.();
    await eventConsumer.catch(() => {});

    // Verify subagent event in queue was emitted with status failed / stopped
    const subagentEvents = events.filter((e: any) => e.type === "subagent");
    expect(subagentEvents.length).toBeGreaterThan(0);
    const lastSubagentEvent = subagentEvents[subagentEvents.length - 1] as any;
    expect(lastSubagentEvent.status).toBe("failed");
    expect(lastSubagentEvent.result).toBe("stopped");

    // Verify tool completion details marked cancelled
    const subResult = completedResults.run_subagent as
      | { details?: Record<string, unknown> }
      | undefined;
    expect(subResult?.details?.cancelled).toBe(true);
    expect(subResult?.details?.reason).toBe("steering_redirection");

    // Verify context pruner removes cancelled subagent
    const messages = [
      {
        role: "assistant",
        content: [
          {
            type: "toolCall",
            id: "call-subagent-1",
            name: "run_subagent",
            arguments: { task: "subtask" },
          },
        ],
      },
      {
        role: "toolResult",
        toolCallId: "call-subagent-1",
        content: [
          { type: "text", text: "[Tool run_subagent cancelled: user steering redirection]" },
        ],
        details: { cancelled: true, reason: "steering_redirection" },
      },
    ];
    const pruned = pruneCancelledToolResultsContext(messages as any);
    expect(pruned).toHaveLength(1);
    expect(pruned[0]).toEqual({
      role: "assistant",
      content: [{ type: "text", text: "[Cancelled tool step]" }],
    });
  });
});