claude-agent-acp/src/exit-plan.ts

import type { SessionNotification } from "@agentclientprotocol/sdk/unstable";
import type { SDKResultMessage } from "@anthropic-ai/claude-agent-sdk";
import {
  type ClearContextCoordinatorHost,
  type ClearContextReset,
  type ClearContextSession,
  type ClearContextTurn,
  continuePlanInFreshContext,
} from "./clear-context-coordinator.js";
import { parseToolResultMeta } from "./tool-result-meta.js";

export function acceptedPlanToolResult(
  notification: SessionNotification,
  toolUseId: string | undefined,
): SessionNotification {
  const update = notification.update;
  if (
    !toolUseId ||
    update.sessionUpdate !== "tool_call_update" ||
    update.toolCallId !== toolUseId
  ) {
    return notification;
  }
  const completed = { ...update };
  delete completed.rawOutput;
  delete completed.content;
  return { ...notification, update: { ...completed, status: "completed" } };
}

function containsToolResultFor(content: unknown, toolUseId: string): boolean {
  return (
    Array.isArray(content) &&
    content.some(
      (block) =>
        typeof block === "object" &&
        block !== null &&
        (block as { type?: unknown }).type === "tool_result" &&
        (block as { tool_use_id?: unknown }).tool_use_id === toolUseId,
    )
  );
}

type ExitPlanState = {
  toolUseCache: Record<string, { name: string } | undefined>;
  pendingExitPlanModeInterruption?: { toolUseId: string; toolResultSeen: boolean };
  pendingExitPlanContextReset?: ClearContextReset;
};

function rejectedExitPlanToolUseId(
  content: unknown,
  toolUseCache: ExitPlanState["toolUseCache"],
  rawToolResultMeta: unknown,
): string | undefined {
  if (!Array.isArray(content)) return undefined;
  const toolResultMeta = parseToolResultMeta(rawToolResultMeta);
  if (!toolResultMeta) return undefined;
  for (const block of content) {
    if (typeof block !== "object" || block === null) continue;
    const { type, tool_use_id: toolUseId, is_error: isError } = block as Record<string, unknown>;
    if (
      type === "tool_result" &&
      typeof toolUseId === "string" &&
      isError === true &&
      toolResultMeta.get(toolUseId)?.nonExecutionKind === "user-rejected" &&
      toolUseCache[toolUseId]?.name === "ExitPlanMode"
    ) {
      return toolUseId;
    }
  }
  return undefined;
}

/** Reconcile an SDK user message with the two pending ExitPlanMode lanes and
 * return the accepted plan tool id whose rendered update must be completed. */
export function observeExitPlanToolResults(
  message: { type: string; tool_result_meta?: unknown },
  content: unknown,
  state: ExitPlanState,
): string | undefined {
  if (message.type !== "user") return undefined;

  const rejectedToolUseId = rejectedExitPlanToolUseId(
    content,
    state.toolUseCache,
    message.tool_result_meta,
  );
  if (rejectedToolUseId) {
    // The stream is authoritative: resumed queries can lose the short-lived
    // marker installed by canUseTool, while metadata preserves correlation.
    state.pendingExitPlanModeInterruption = {
      toolUseId: rejectedToolUseId,
      toolResultSeen: true,
    };
  }

  const pendingInterruption = state.pendingExitPlanModeInterruption;
  if (pendingInterruption && containsToolResultFor(content, pendingInterruption.toolUseId)) {
    pendingInterruption.toolResultSeen = true;
  }

  const pendingReset = state.pendingExitPlanContextReset;
  return pendingReset && containsToolResultFor(content, pendingReset.toolUseId)
    ? pendingReset.toolUseId
    : undefined;
}

export function executionDiagnostic(message: SDKResultMessage): string | undefined {
  if (message.subtype === "success") {
    return message.result.startsWith("[ede_diagnostic]") ? message.result : undefined;
  }
  return message.errors.find((error) => error.startsWith("[ede_diagnostic]"));
}

/** Whether `message` is the error-shaped diagnostic Claude emits for the cycle
 * a recorded ExitPlanMode `interrupt: true` answer stopped. */
export function isExitPlanInterruptionResult(
  message: SDKResultMessage,
  pending: ExitPlanState["pendingExitPlanModeInterruption"],
): boolean {
  if (!pending?.toolResultSeen || !message.is_error) return false;
  const diagnostic = executionDiagnostic(message);
  return (
    !!diagnostic &&
    /(?:^|\s)result_type=user(?:\s|$)/.test(diagnostic) &&
    /(?:^|\s)stop_reason=tool_use(?:\s|$)/.test(diagnostic)
  );
}

/** The base ID of an approved plan. AIR appends the review ID to keep each
 * approval in the transcript. Other clients replace one plan per session.
 * The TodoWrite checklist is a separate plan. */
export const PLAN_MODE_PLAN_ID = "plan-mode";

/** The plan, trimmed, that a successful ExitPlanMode result approved: the plan
 * of the structured `output`, which includes the user's edits, else the plan
 * that Claude Code adds to the tool's `input`. Replay has only the input: the
 * SDK's `getSessionMessages` drops the structured output that the transcript
 * keeps. Undefined for a subagent's plan, one that waits for its team lead
 * (both only known from the output), and an empty plan. */
export function approvedPlan(input: unknown, output: unknown): string | undefined {
  const result = output as
    { plan?: unknown; isAgent?: unknown; awaitingLeaderApproval?: unknown } | undefined;
  if (result?.isAgent === true || result?.awaitingLeaderApproval === true) return undefined;
  const plan =
    typeof result?.plan === "string"
      ? result.plan
      : (input as { plan?: unknown } | undefined)?.plan;
  return typeof plan === "string" && plan.trim() !== "" ? plan.trim() : undefined;
}

/** The update that shows an approved plan to a client that supports plan
 * operations (the `plan` client capability, which every v2 client has).
 * For AIR, pass the review ID to join the plan to its approval card. */
export function planModeUpdate(
  plan: string,
  reviewToolCallId?: string,
): SessionNotification["update"] {
  return {
    sessionUpdate: "plan_update",
    plan: {
      type: "markdown",
      planId: reviewToolCallId ? `${PLAN_MODE_PLAN_ID}:${reviewToolCallId}` : PLAN_MODE_PLAN_ID,
      content: plan,
    },
  };
}

/** Claude wraps a rejected ExitPlanMode explanation in a Markdown code fence.
 * Strip exactly one complete outer fence for that tool only. */
export function exitPlanModeRawOutput(toolName: string, content: unknown): unknown {
  if (toolName !== "ExitPlanMode" || typeof content !== "string") {
    return content;
  }
  const fenced = /^\s*```[^\r\n]*\r?\n([\s\S]*?)\r?\n```\s*$/.exec(content);
  return fenced?.[1] ?? content;
}

export type ExitPlanRestartHost<
  Session extends ClearContextSession<Turn>,
  Turn extends ClearContextTurn,
> = Omit<
  ClearContextCoordinatorHost<Session, Turn>,
  "publishSessionState" | "continuationMessage"
> & {
  sessionUpdate(notification: SessionNotification): Promise<void>;
  /** Whether the client takes plan operations (see planModeUpdate). */
  supportsPlanOperations(): boolean;
  isAirClient(): boolean;
  destroyReplacement(sessionId: string, session: Session): void;
  settleCancelledTurn(oldSession: Session, turnSession: Session, turn: Turn): void;
  settleFailedTurn(turnSession: Session, turn: Turn, error: unknown): void;
};

/** Owns the lifetime of accepted-plan context replacements. In particular, a
 * session cancellation invalidates an in-progress async restart so a late
 * restartSession result cannot recreate a closed public session. */
export class ExitPlanCoordinator<
  Session extends ClearContextSession<Turn>,
  Turn extends ClearContextTurn,
> {
  private readonly restarts = new Map<string, AbortController>();

  constructor(private readonly host: ExitPlanRestartHost<Session, Turn>) {}

  cancel(sessionId: string): void {
    this.restarts.get(sessionId)?.abort();
  }

  async restart(sessionId: string, oldSession: Session, reset: ClearContextReset): Promise<void> {
    this.cancel(sessionId);
    const controller = new AbortController();
    this.restarts.set(sessionId, controller);
    try {
      const clearContextHost: ClearContextCoordinatorHost<Session, Turn> = {
        ...this.host,
        publishSessionState: async (id, mode, configOptions, contextWindowSize) => {
          await this.host.sessionUpdate({
            sessionId: id,
            update: { sessionUpdate: "current_mode_update", currentModeId: mode },
          });
          await this.host.sessionUpdate({
            sessionId: id,
            update: { sessionUpdate: "config_option_update", configOptions },
          });
          // Otherwise the client shows the old context's usage until the fresh
          // context first reports its own.
          await this.host.sessionUpdate({
            sessionId: id,
            update: { sessionUpdate: "usage_update", used: 0, size: contextWindowSize },
          });
          // The approved plan becomes the session's plan once the fresh
          // context that carries it out exists. Its ExitPlanMode ended with
          // an error result, so toAcpNotifications did not publish it.
          if (this.host.supportsPlanOperations()) {
            await this.host.sessionUpdate({
              sessionId: id,
              update: planModeUpdate(
                reset.plan,
                this.host.isAirClient() ? reset.toolUseId : undefined,
              ),
            });
          }
        },
        continuationMessage: (id, plan, promptUuid) => ({
          type: "user",
          message: {
            role: "user",
            content: [{ type: "text", text: `Implement the following plan:\n\n${plan}` }],
          },
          session_id: id,
          parent_tool_use_id: null,
          origin: { kind: "human" },
          uuid: promptUuid as `${string}-${string}-${string}-${string}-${string}`,
        }),
      };
      await continuePlanInFreshContext(
        sessionId,
        oldSession,
        reset,
        clearContextHost,
        controller.signal,
      );
    } catch (error) {
      const currentSession = this.host.currentSession(sessionId);
      const replacement = currentSession !== oldSession ? currentSession : undefined;
      if (replacement) this.host.destroyReplacement(sessionId, replacement);

      const turn = replacement?.activeTurn ?? oldSession.activeTurn;
      if (turn && !turn.settled) {
        const turnSession = replacement?.activeTurn === turn ? replacement : oldSession;
        if (controller.signal.aborted) {
          this.host.settleCancelledTurn(oldSession, turnSession, turn);
        } else {
          // A provider/session-creation failure is turn-scoped. Settling it
          // here keeps it out of the query consumer's transport-loss catch.
          this.host.settleFailedTurn(turnSession, turn, error);
        }
        oldSession.activeTurn = null;
        oldSession.turnQueue = (oldSession.turnQueue ?? []).filter((queued) => queued !== turn);
        if (replacement) {
          replacement.activeTurn = null;
          replacement.turnQueue = (replacement.turnQueue ?? []).filter((queued) => queued !== turn);
        }
      }
      oldSession.pendingExitPlanContextReset = undefined;
    } finally {
      if (this.restarts.get(sessionId) === controller) {
        this.restarts.delete(sessionId);
      }
    }
  }
}