import type { UIMessage, UIMessageChunk } from "ai";
import { resourceCatalog } from "@trigger.dev/core/v3";
import type { LocalsKey } from "@trigger.dev/core/v3";
import { runInMockTaskContext, type MockTaskContextOptions } from "@trigger.dev/core/v3/test";
import { __setSessionOpenImplForTests, __setSessionStartImplForTests } from "../sessions.js";
import {
  __setReadChatSnapshotImplForTests,
  __setReplaySessionInTailImplForTests,
  __setReplaySessionOutTailImplForTests,
  __setWriteChatSnapshotImplForTests,
  type ChatSnapshotV1,
} from "../ai.js";
import { createTestSessionHandle, type TestSessionOutState } from "./test-session-handle.js";

/** Pre-seed locals before the agent's `run()` starts. */
export type SetupLocals = (locals: {
  set<T>(key: LocalsKey<T>, value: T): void;
}) => void | Promise<void>;

// The slim wire payload shape used by chat.agent tasks. Kept loose here so we
// don't import from the backend-only ai.ts module. At most ONE message per
// record — runtime rebuilds prior history from snapshot + replay at boot.
type ChatWirePayload = {
  /** At most one message — singular under the slim wire. Set on submit-message. */
  message?: UIMessage;
  /** Bespoke escape hatch — only set on `trigger: "handover-prepare"`. */
  headStartMessages?: UIMessage[];
  chatId: string;
  trigger:
    | "submit-message"
    | "regenerate-message"
    | "preload"
    | "close"
    | "action"
    | "handover-prepare";
  messageId?: string;
  metadata?: unknown;
  action?: unknown;
  continuation?: boolean;
  previousRunId?: string;
  idleTimeoutInSeconds?: number;
  sessionId?: string;
};

/** A reference to a `chat.agent` task returned by `chat.agent({ id, ... })`. */
type ChatAgentHandle = { id: string };

/**
 * Options for `mockChatAgent`.
 */
export type MockChatAgentOptions = {
  /** The chat session id passed into every wire payload. Defaults to `"test-chat"`. */
  chatId?: string;
  /** Client-provided metadata (`clientData`) for the session. */
  clientData?: unknown;
  /** Task context overrides passed through to {@link runInMockTaskContext}. */
  taskContext?: MockTaskContextOptions;
  /**
   * Whether to start the task in preload mode. Defaults to `true` so the
   * first `sendMessage()` triggers the first turn via the preload path.
   * Set to `false` to skip preload — the first `sendMessage()` starts turn 0 directly.
   *
   * Ignored when `mode: "handover-prepare"` is set.
   */
  preload?: boolean;
  /**
   * Initial trigger the agent boots with. Defaults to `"preload"` (or
   * `"submit-message"` when `preload: false`, or `"continuation"` when
   * `continuation: true`).
   *
   * - `"preload"` — fresh chat preloaded via `transport.preload`. Fires
   *   `onPreload`, waits for the first message.
   * - `"submit-message"` — fresh chat with the first message in the boot
   *   payload (the `chat.createStartSessionAction({ basePayload: { message } })`
   *   pattern). Goes straight to turn 0.
   * - `"continuation"` — new run picking up an existing session after the
   *   prior run ended (`chat.endRun`, waitpoint timeout, `chat.requestUpgrade`).
   *   Boots with `trigger` omitted and `continuation: true` — mirrors what
   *   the server's `ensureRunForSession` / `swapSessionRun` produces in
   *   production. The SDK enters its continuation-wait branch; `onPreload`
   *   and `onChatStart` do NOT fire on this run.
   * - `"handover-prepare"` — drives the chat.handover wait branch; call
   *   `sendHandover()` / `sendHandoverSkip()` to dispatch the handover signal.
   */
  mode?: "preload" | "submit-message" | "handover-prepare" | "continuation";
  /**
   * First-turn UIMessage history shipped on the BOOT payload. Only
   * meaningful with `mode: "handover-prepare"` — mirrors the
   * `chat.headStart` route handler's `basePayload.headStartMessages`.
   */
  headStartMessages?: UIMessage[];
  /**
   * Pre-seed the snapshot the agent reads at run boot. The runtime's
   * snapshot read is replaced with one that returns this snapshot
   * (skipping the real S3 GET). Use to drive boot scenarios — fresh
   * boot with prior history, OOM-retry boot with stale snapshot, etc.
   * Pass `undefined` (the default) to start with no snapshot.
   *
   * See plan section B.3 for the boot orchestration spec.
   */
  snapshot?: ChatSnapshotV1;
  /**
   * Set `payload.continuation = true` on the initial wire payload. Used
   * to simulate a continuation-run boot (a new run picking up after a
   * prior run on the same session ended via `chat.endRun`, waitpoint
   * timeout, or `chat.requestUpgrade`).
   *
   * Setting this without specifying `mode` auto-selects `mode:
   * "continuation"` — the SDK boot path enters its continuation-wait
   * branch and waits silently on `session.in` for the first user
   * message. `onPreload` and `onChatStart` do NOT fire on this run.
   *
   * Defaults to `false` (fresh run).
   */
  continuation?: boolean;
  /**
   * Set `payload.previousRunId` on the initial wire payload. Forwarded
   * to `onChatStart` / `onTurnStart` and used by the boot gate as a
   * prior-state signal. Usually paired with `continuation: true`.
   */
  previousRunId?: string;
  /**
   * Callback that runs **before** the agent's `run()` is invoked, with a
   * `set` function for pre-seeding locals. Use this to inject server-side
   * dependencies (database clients, service stubs) that the agent reads
   * via `locals.get()` in its hooks.
   *
   * @example
   * ```ts
   * import { dbKey } from "./db";
   *
   * const harness = mockChatAgent(agent, {
   *   chatId: "test-1",
   *   setupLocals: (locals) => {
   *     locals.set(dbKey, testDb);
   *   },
   * });
   * ```
   */
  setupLocals?: SetupLocals;
};

/**
 * Result of a single turn, returned by driver methods like `sendMessage()`.
 */
export type MockChatAgentTurn = {
  /** UIMessageChunks emitted during this turn (excludes control chunks like turn-complete). */
  chunks: UIMessageChunk[];
  /** All raw chunks including control chunks (turn-complete, upgrade-required, etc.). */
  rawChunks: unknown[];
};

/**
 * Harness returned by `mockChatAgent`. Drives a `chat.agent` task end-to-end
 * without network or task runtime.
 */
export type MockChatAgentHarness = {
  /** The chat session id used by this harness. */
  readonly chatId: string;

  /**
   * Send a single user message (or tool-approval-responded assistant
   * message) and wait for the next turn-complete. Returns the chunks
   * produced during this turn.
   *
   * Slim wire: at most ONE message per send. The agent reconstructs prior
   * history from snapshot + session.out replay at run boot.
   */
  sendMessage(message: UIMessage): Promise<MockChatAgentTurn>;

  /**
   * Send a regenerate signal (no message body — slim wire). The agent
   * trims trailing assistant messages from its in-memory accumulator and
   * re-runs. Waits for turn-complete.
   */
  sendRegenerate(): Promise<MockChatAgentTurn>;

  /**
   * Drive the head-start path: sends `trigger: "handover-prepare"` with
   * `headStartMessages` carrying the first-turn UIMessage history. Used
   * only at the very first turn before any snapshot exists. The route
   * handler ships full UIMessage history through this path because the
   * customer's HTTP endpoint isn't subject to the `/in/append` cap.
   */
  sendHeadStart(args: { messages: UIMessage[] }): Promise<MockChatAgentTurn>;

  /** Send a custom action and wait for the next turn-complete. */
  sendAction(action: unknown): Promise<MockChatAgentTurn>;

  /** Fire a stop signal. Does not wait for the turn — the task keeps running. */
  sendStop(message?: string): Promise<void>;

  /**
   * Dispatch a `handover` signal — the agent picks up partial assistant
   * messages and continues the turn. Only meaningful when the harness
   * was started with `mode: "handover-prepare"`. Waits for turn-complete.
   *
   * `isFinal: false` (default) — agent runs `streamText` which executes
   * any pending tool-calls (via the approval round) and resumes from
   * step 2.
   *
   * `isFinal: true` — agent runs lifecycle hooks but skips `streamText`.
   * The partial IS the response; `onTurnComplete` fires with it.
   */
  sendHandover(args: {
    partialAssistantMessage: unknown[];
    isFinal?: boolean;
    messageId?: string;
  }): Promise<MockChatAgentTurn>;

  /**
   * Dispatch a `handover-skip` signal — the agent exits cleanly without
   * firing turn hooks. Only meaningful when the harness was started
   * with `mode: "handover-prepare"`. Awaits the run finishing.
   */
  sendHandoverSkip(): Promise<void>;

  /**
   * Pre-seed the snapshot read for the next boot. The runtime's snapshot
   * read returns this snapshot (skipping S3). Pass `undefined` to clear —
   * the boot then sees no snapshot and falls through to replay-only.
   *
   * Effective on the next run boot only. Calling mid-turn is a no-op
   * because the snapshot read happens once at run boot.
   */
  seedSnapshot(snapshot: ChatSnapshotV1 | undefined): void;

  /**
   * Pre-seed `session.out` chunks for the next boot's replay. The runtime's
   * `replaySessionOutTail` returns whatever the synthetic chunks reduce
   * to. Pass `[]` to clear (boot replay returns no messages).
   *
   * Requires `__setReplaySessionOutTailImplForTests` exported from
   * `ai.ts`. The harness throws a clear error at call time if that hook
   * isn't available.
   */
  seedSessionOutTail(chunks?: UIMessageChunk[]): void;

  /**
   * Pre-seed a trailing partial assistant message for the next boot's
   * replay. The runtime's `replaySessionOutTail` returns this as the
   * `partial` field (alongside whatever `seedSessionOutTail` reduces
   * to). Use to simulate cancel-mid-stream: an assistant message whose
   * `finish` chunk never arrived. Pass `undefined` to clear.
   *
   * Effective on the next run boot only.
   */
  seedSessionOutPartial(partial: UIMessage | undefined): void;

  /**
   * Pre-seed user messages on the `session.in` tail for the next boot's
   * replay. Each message is paired with a synthetic seq_num (`i + 1`).
   * Used to simulate in-flight users the dead predecessor was supposed
   * to process. Pass `[]` to clear.
   *
   * Effective on the next run boot only.
   */
  seedSessionInTail(messages: UIMessage[]): void;

  /**
   * The most recently written snapshot, or `undefined` if no snapshot
   * has been written yet. Updated each time `writeChatSnapshot` is
   * invoked from the run loop's snapshot-write site (plan section B.6).
   */
  getSnapshot(): ChatSnapshotV1 | undefined;

  /**
   * Close the chat session cleanly. Sends `trigger: "close"` and awaits the
   * task's `run()` function returning. Call this at the end of every test
   * (or use `await using`) so the background task isn't left dangling.
   */
  close(): Promise<void>;

  /** All UIMessageChunks emitted since the harness was created. */
  readonly allChunks: UIMessageChunk[];

  /** Every raw chunk (including control chunks) emitted since the harness was created. */
  readonly allRawChunks: unknown[];
};

const CONTROL_CHUNK_TYPES = new Set(["trigger:turn-complete", "trigger:upgrade-required"]);

function isControlChunk(chunk: unknown): boolean {
  if (typeof chunk !== "object" || chunk === null) return false;
  const type = (chunk as { type?: string }).type;
  return typeof type === "string" && CONTROL_CHUNK_TYPES.has(type);
}

/**
 * Create an offline test harness for a `chat.agent` task.
 *
 * The harness starts the agent's `run()` function in a mocked task context,
 * waits in preload for the first message, then exposes driver methods for
 * sending messages / actions / stop signals and awaiting turn completion.
 *
 * Users are responsible for mocking the language model themselves — use
 * `MockLanguageModelV3` and `simulateReadableStream` from `ai/test` inside
 * their agent's `run()` function (typically via DI through `clientData`).
 *
 * @example
 * ```ts
 * import { mockChatAgent } from "@trigger.dev/sdk/ai/test";
 * import { MockLanguageModelV3, simulateReadableStream } from "ai/test";
 * import { myAgent } from "./my-agent";
 *
 * test("says hello", async () => {
 *   const harness = mockChatAgent(myAgent, { chatId: "test-1" });
 *   try {
 *     const turn = await harness.sendMessage({
 *       id: "m1",
 *       role: "user",
 *       parts: [{ type: "text", text: "hi" }],
 *     });
 *     expect(turn.chunks).toContainEqual(
 *       expect.objectContaining({ type: "text-delta", delta: "hello" })
 *     );
 *   } finally {
 *     await harness.close();
 *   }
 * });
 * ```
 */
export function mockChatAgent(
  agent: ChatAgentHandle,
  options: MockChatAgentOptions = {}
): MockChatAgentHarness {
  const chatId = options.chatId ?? "test-chat";
  // The agent opens the session with `payload.sessionId ?? payload.chatId`.
  // We pass no sessionId, so it falls back to chatId.
  const sessionId = chatId;
  // `continuation: true` without an explicit mode auto-selects "continuation"
  // — the canonical shape for a continuation-run boot.
  const mode: "preload" | "submit-message" | "handover-prepare" | "continuation" =
    options.mode ??
    (options.continuation === true
      ? "continuation"
      : options.preload === false
        ? "submit-message"
        : "preload");
  const clientData = options.clientData;

  const taskEntry = resourceCatalog.getTask(agent.id);
  if (!taskEntry) {
    throw new Error(
      `mockChatAgent: no task registered with id "${agent.id}". ` +
        `Import "@trigger.dev/sdk/ai/test" before your agent module so tasks register correctly.`
    );
  }

  const runFn = taskEntry.fns.run;

  // Session .out state: chunks + listener registry. Shared between the
  // harness and the TestSessionOutputChannel installed via the open-override.
  const sessionOutState: TestSessionOutState = {
    chunks: [],
    listeners: new Set(),
  };

  // Buffers that survive across harness method calls
  const allRawChunks: unknown[] = [];
  const allChunks: UIMessageChunk[] = [];

  // Promise that resolves when the background task run() function returns.
  let taskFinished!: Promise<void>;
  let sendSessionInput!: (sessionId: string, data: unknown) => Promise<void>;
  let closeSessionInput: ((sessionId: string) => void) | undefined;
  let runSignal!: AbortController;

  // A latch that resolves every time `trigger:turn-complete` appears on the chat stream.
  // We use a shared pending promise and replace it after each completion.
  let turnCompleteResolvers: Array<() => void> = [];
  const waitForTurnComplete = () =>
    new Promise<void>((resolve) => {
      turnCompleteResolvers.push(resolve);
    });

  // Signal that the caller is ready to observe output
  let harnessReadyResolve!: () => void;
  const harnessReady = new Promise<void>((resolve) => {
    harnessReadyResolve = resolve;
  });

  // ── Snapshot read/write override state ───────────────────────────────
  // The runtime's snapshot read returns whatever `seededSnapshot` is at
  // boot time. The runtime's snapshot write captures into
  // `lastWrittenSnapshot` for harness consumers to assert via
  // `getSnapshot()`. Installed below alongside the session overrides;
  // cleared on close in the same finally block.
  let seededSnapshot: ChatSnapshotV1 | undefined = options.snapshot;
  let lastWrittenSnapshot: ChatSnapshotV1 | undefined;
  let seededReplayChunks: UIMessageChunk[] = [];
  let seededReplayPartial: UIMessage | undefined;
  let seededSessionInMessages: UIMessage[] = [];

  __setReadChatSnapshotImplForTests(<T extends UIMessage>(_id: string) => {
    return seededSnapshot as ChatSnapshotV1<T> | undefined;
  });
  __setWriteChatSnapshotImplForTests(
    <T extends UIMessage>(_id: string, snapshot: ChatSnapshotV1<T>) => {
      lastWrittenSnapshot = snapshot as ChatSnapshotV1;
    }
  );

  // Replay override: install a default that returns whatever
  // `seededReplayChunks` reduces to. `mockChatAgent` doesn't model the
  // settled-vs-partial split — seeded chunks always reduce to the
  // `settled` array with `partial: undefined`. Recovery-specific
  // tests can install their own override to seed a partial.
  // Cleared in the same `finally` block as the other test overrides.
  __setReplaySessionOutTailImplForTests(async () => {
    const settled =
      seededReplayChunks.length === 0
        ? []
        : ((await reduceChunksToMessages(seededReplayChunks)) as unknown[]);
    // For the mock harness, `partialRaw` is the same as `partial` — we
    // don't model cleanupAbortedParts separately. Recovery tests that
    // need a partialRaw distinct from partial install their own stub.
    return {
      settled,
      partial: seededReplayPartial,
      partialRaw: seededReplayPartial,
    } as never;
  });

  // session.in tail override: each seeded UIMessage becomes a
  // { message, metadata: undefined, seqNum: i+1 } entry. Mirrors the
  // seq-num pattern from the out-tail stub so cursor-advance logic is
  // exercised correctly. `metadata` is `undefined` for seeded users —
  // the boot path falls back to `payload.metadata` for those.
  __setReplaySessionInTailImplForTests(async () => {
    return seededSessionInMessages.map((message, i) => ({
      message,
      metadata: undefined,
      seqNum: i + 1,
    })) as never;
  });

  // Install the session open override so `sessions.open(id)` returns a
  // SessionHandle with an in-memory `.out` that captures writes. The
  // `.in` channel routes record subscriptions (`on`/`once`/`peek`)
  // through the `sessionStreams` global — the mock task context
  // installs a `TestSessionStreamManager` there — and stubs `wait()`
  // so the suspend path resolves cleanly on `runSignal.abort()` without
  // touching the api client.
  __setSessionOpenImplForTests((id) =>
    createTestSessionHandle(id, sessionOutState, () => runSignal?.signal)
  );

  // Install the session start override so any test path that invokes
  // `sessions.start()` (typically through a server action shim like
  // `chat.createStartSessionAction`) becomes a no-op fixture instead of
  // hitting a real API. Most chat.agent tests trigger the run directly
  // via `sendPayloadAndWait` and never go through this path, but the
  // stub keeps the API safe to call from inside tested code.
  __setSessionStartImplForTests((body) => {
    if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") {
      console.log("[mockChatAgent] sessions.start override:", body);
    }
    const fakeRunId = `run_test_${body.externalId ?? "anon"}`;
    return {
      id: `session_test_${body.externalId ?? "anon"}`,
      externalId: body.externalId ?? null,
      type: body.type,
      taskIdentifier: body.taskIdentifier,
      triggerConfig: body.triggerConfig,
      currentRunId: fakeRunId,
      runId: fakeRunId,
      publicAccessToken: "tr_test_session_pat",
      tags: body.tags ?? [],
      metadata: (body.metadata ?? null) as Record<string, unknown> | null,
      closedAt: null,
      closedReason: null,
      expiresAt: null,
      createdAt: new Date(0),
      updatedAt: new Date(0),
      isCached: false,
    };
  });

  taskFinished = runInMockTaskContext(async (drivers) => {
    runSignal = new AbortController();

    // For `mode: "continuation"`, omit `trigger` from the wire payload —
    // mirrors what the server's `ensureRunForSession` / `swapSessionRun`
    // produces (the continuation overrides clear `trigger` so the SDK
    // boot path falls into the continuation-wait branch instead of
    // re-firing the basePayload's stale first-run trigger). `continuation:
    // true` is set unconditionally for this mode so the boot path's
    // continuation-wait condition matches.
    const isContinuationMode = mode === "continuation";
    const initialPayload: ChatWirePayload = {
      chatId,
      ...(isContinuationMode
        ? { trigger: undefined as never, continuation: true }
        : { trigger: mode }),
      metadata: clientData,
      ...(!isContinuationMode && options.continuation ? { continuation: true } : {}),
      ...(options.previousRunId ? { previousRunId: options.previousRunId } : {}),
      ...(options.headStartMessages ? { headStartMessages: options.headStartMessages } : {}),
    };

    sendSessionInput = drivers.sessions.in.send;
    closeSessionInput = drivers.sessions.in.close;

    // Record every chunk written to session.out, detect turn-complete.
    const listener = (chunk: unknown) => {
      allRawChunks.push(chunk);
      if (!isControlChunk(chunk)) {
        allChunks.push(chunk as UIMessageChunk);
      }
      if (
        typeof chunk === "object" &&
        chunk !== null &&
        (chunk as { type?: string }).type === "trigger:turn-complete"
      ) {
        const resolvers = turnCompleteResolvers;
        turnCompleteResolvers = [];
        for (const resolve of resolvers) resolve();
      }
    };
    sessionOutState.listeners.add(listener);
    const unsubscribe = () => sessionOutState.listeners.delete(listener);

    if (options.setupLocals) {
      await options.setupLocals({ set: drivers.locals.set });
    }

    harnessReadyResolve();

    try {
      if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") {
        console.log("[mockChatAgent] Starting runFn with payload:", initialPayload);
      }
      await runFn(initialPayload, {
        ctx: drivers.ctx,
        signal: runSignal.signal,
      });
      if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") {
        console.log("[mockChatAgent] runFn returned");
      }
    } catch (err) {
      if (process.env.TRIGGER_CHAT_TEST_DEBUG === "1") {
        console.log("[mockChatAgent] runFn threw:", err);
      }
      throw err;
    } finally {
      unsubscribe();
      // Resolve any outstanding turn-complete waiters so callers don't hang
      const resolvers = turnCompleteResolvers;
      turnCompleteResolvers = [];
      for (const resolve of resolvers) resolve();
    }
  }, options.taskContext)
    .catch((err) => {
      // Propagate errors to pending turn waiters instead of dropping them
      const resolvers = turnCompleteResolvers;
      turnCompleteResolvers = [];
      for (const resolve of resolvers) resolve();
      throw err;
    })
    .finally(() => {
      // Always clear the test overrides, even if the task threw.
      __setSessionOpenImplForTests(undefined);
      __setSessionStartImplForTests(undefined);
      __setReadChatSnapshotImplForTests(undefined);
      __setWriteChatSnapshotImplForTests(undefined);
      __setReplaySessionOutTailImplForTests(undefined);
      __setReplaySessionInTailImplForTests(undefined);
    });

  const sendPayloadAndWait = async (payload: ChatWirePayload): Promise<MockChatAgentTurn> => {
    await harnessReady;
    const before = allRawChunks.length;
    const turnComplete = waitForTurnComplete();
    await sendSessionInput(sessionId, { kind: "message", payload });
    await turnComplete;
    const rawChunks = allRawChunks.slice(before);
    const chunks = rawChunks.filter((c) => !isControlChunk(c)) as UIMessageChunk[];
    return { chunks, rawChunks };
  };

  const harness: MockChatAgentHarness = {
    chatId,

    async sendMessage(message) {
      return sendPayloadAndWait({
        message,
        chatId,
        trigger: "submit-message",
        metadata: clientData,
      });
    },

    async sendRegenerate() {
      return sendPayloadAndWait({
        chatId,
        trigger: "regenerate-message",
        metadata: clientData,
      });
    },

    async sendHeadStart({ messages }) {
      return sendPayloadAndWait({
        headStartMessages: messages,
        chatId,
        trigger: "handover-prepare",
        metadata: clientData,
      });
    },

    async sendAction(action) {
      return sendPayloadAndWait({
        chatId,
        trigger: "action",
        action,
        metadata: clientData,
      });
    },

    async sendStop(message) {
      await harnessReady;
      await sendSessionInput(sessionId, { kind: "stop", message });
    },

    async sendHandover(args) {
      await harnessReady;
      const before = allRawChunks.length;
      const turnComplete = waitForTurnComplete();
      await sendSessionInput(sessionId, {
        kind: "handover",
        partialAssistantMessage: args.partialAssistantMessage,
        messageId: args.messageId,
        isFinal: args.isFinal ?? false,
      });
      await turnComplete;
      const rawChunks = allRawChunks.slice(before);
      const chunks = rawChunks.filter((c) => !isControlChunk(c)) as UIMessageChunk[];
      return { chunks, rawChunks };
    },

    async sendHandoverSkip() {
      await harnessReady;
      // No turn-complete on skip — the agent exits without firing hooks.
      // Send the chunk and wait for the run to finish.
      await sendSessionInput(sessionId, { kind: "handover-skip" });
      await Promise.race([
        taskFinished.catch(() => {}),
        new Promise<void>((resolve) => setTimeout(resolve, 1000)),
      ]);
    },

    seedSnapshot(snapshot) {
      seededSnapshot = snapshot;
    },

    seedSessionOutTail(chunks) {
      seededReplayChunks = chunks ?? [];
    },

    seedSessionOutPartial(partial) {
      seededReplayPartial = partial;
    },

    seedSessionInTail(messages) {
      seededSessionInMessages = messages;
    },

    getSnapshot() {
      return lastWrittenSnapshot;
    },

    async close() {
      await harnessReady;

      // Send a close trigger wrapped as a `kind: "message"` ChatInputChunk.
      // The turn loop checks for this after a successful turn and exits
      // cleanly. On error-recovery paths the loop just loops back with
      // the close payload, so we also close the session input below to
      // unblock any pending once() waiters.
      try {
        await sendSessionInput(sessionId, {
          kind: "message",
          payload: {
            chatId,
            trigger: "close",
          },
        });
      } catch {
        // best-effort
      }
      // Resolve any pending once() waiters on the session input with a
      // timeout error — that makes waitWithIdleTimeout return
      // `{ ok: false }` and the turn loop exits cleanly.
      closeSessionInput?.(sessionId);

      // Also abort the run signal so anything downstream (streamText,
      // deferred work) unwinds promptly.
      runSignal?.abort("close");

      // Wait for run() to return. The loop's error recovery path will
      // see !next.ok and exit. Use a bounded wait so tests never hang.
      await Promise.race([
        taskFinished.catch(() => {}),
        new Promise<void>((resolve) => setTimeout(resolve, 1000)),
      ]);
    },

    get allChunks() {
      return allChunks.slice();
    },

    get allRawChunks() {
      return allRawChunks.slice();
    },
  };

  return harness;
}

/**
 * Reduce a synthetic UIMessageChunk[] sequence into the UIMessage[] that
 * the runtime's `replaySessionOutTail` would produce. Splits chunks at
 * `start` boundaries and feeds each segment through AI SDK's
 * `readUIMessageStream`. The trailing un-finished segment goes through
 * `cleanupAbortedParts`. Mirrors the production reducer used in
 * `ai.ts:replaySessionOutTail`.
 */
async function reduceChunksToMessages(chunks: UIMessageChunk[]): Promise<UIMessage[]> {
  if (chunks.length === 0) return [];
  const aiModule = (await import("ai")) as {
    readUIMessageStream?: (args: {
      stream: ReadableStream<UIMessageChunk>;
    }) => AsyncIterable<UIMessage>;
    cleanupAbortedParts?: (msg: UIMessage) => UIMessage;
  };
  const readUIMessageStream = aiModule.readUIMessageStream;
  const cleanupAbortedParts = aiModule.cleanupAbortedParts;
  if (!readUIMessageStream) return [];

  type Segment = { chunks: UIMessageChunk[]; closed: boolean };
  const segments: Segment[] = [];
  let current: Segment | undefined;
  for (const chunk of chunks) {
    if (chunk.type === "start") {
      current = { chunks: [chunk], closed: false };
      segments.push(current);
      continue;
    }
    if (!current) {
      current = { chunks: [], closed: false };
      segments.push(current);
    }
    current.chunks.push(chunk);
    if (chunk.type === "finish") {
      current.closed = true;
      current = undefined;
    }
  }

  const out: UIMessage[] = [];
  for (let i = 0; i < segments.length; i++) {
    const seg = segments[i]!;
    const isTrailing = i === segments.length - 1 && !seg.closed;
    const segmentStream = new ReadableStream<UIMessageChunk>({
      start(controller) {
        for (const c of seg.chunks) controller.enqueue(c);
        controller.close();
      },
    });
    let last: UIMessage | undefined;
    try {
      for await (const snapshot of readUIMessageStream({ stream: segmentStream })) {
        last = snapshot;
      }
    } catch {
      // Skip malformed segment — tests can assert by inspecting what makes it through.
      continue;
    }
    if (!last) continue;
    if (isTrailing && cleanupAbortedParts) {
      const cleaned = cleanupAbortedParts(last);
      if (!cleaned.parts || cleaned.parts.length === 0) continue;
      out.push(cleaned);
    } else {
      out.push(last);
    }
  }
  return out;
}