claude-agent-acp/src/elicitation.ts

import { randomUUID } from "node:crypto";
import type {
  CreateElicitationRequest,
  CreateElicitationResponse,
  ElicitationPropertySchema,
  ElicitationSchema,
  EnumOption,
} from "@agentclientprotocol/sdk/unstable";
import type {
  ElicitationRequest,
  ElicitationResult,
  PermissionResult,
  UserDialogRequest,
  UserDialogResult,
} from "@anthropic-ai/claude-agent-sdk";
import type {
  AskUserQuestionInput,
  AskUserQuestionOutput,
} from "@anthropic-ai/claude-agent-sdk/sdk-tools.js";
import {
  AIR_CUSTOM_ANSWER_KEY,
  LEGACY_AIR_CUSTOM_ANSWER_KEY,
  withAirMeta,
} from "./air-extension.js";

/**
 * Bridges between the Claude Agent SDK's elicitation/dialog callbacks and ACP's
 * (unstable) elicitation protocol.
 *
 * Two distinct SDK surfaces flow through here:
 *
 *   1. `onElicitation` — fired when an MCP server requests user input. These map
 *      directly onto ACP `session/create_elicitation` (form or url mode).
 *   2. The built-in AskUserQuestion tool — when a `canUseTool` callback is
 *      registered the SDK routes its permission check through `canUseTool`
 *      (not the interactive `permission_ask_user_question` dialog). We render
 *      its questions as an ACP form elicitation and feed the user's selections
 *      back as the tool's `updatedInput`, which the tool's own `call()` reads.
 */

/** Modes the connected client advertised support for. */
export type ElicitationSupport = {
  form: boolean;
  url: boolean;
};

/**
 * Convert an MCP elicitation request (from the SDK's `onElicitation` callback)
 * into an ACP `CreateElicitationRequest`. Returns `null` when the request can't
 * be represented (e.g. a url-mode request with no url).
 */
export function mcpElicitationToCreateRequest(
  request: ElicitationRequest,
  sessionId: string,
): CreateElicitationRequest | null {
  if (request.mode === "url") {
    if (!request.url) {
      return null;
    }
    return {
      mode: "url",
      sessionId,
      message: request.message,
      url: request.url,
      // URL elicitations need a stable id so the client can correlate the
      // later `session/complete_elicitation` notification. MCP servers usually
      // provide one; fall back to a generated id if not.
      elicitationId: request.elicitationId ?? randomUUID(),
    };
  }

  // Form mode (the default). The MCP `requestedSchema` is already a JSON Schema
  // with primitive-typed properties, which is structurally what ACP expects.
  return {
    mode: "form",
    sessionId,
    message: request.message,
    requestedSchema: normalizeElicitationSchema(request.requestedSchema),
  };
}

/**
 * Content of an accepted elicitation response, or `{}` when the response is
 * no accept or its `content` is no JSON object.
 *
 * The values are not checked. v1 allows a string, number, boolean, or string
 * array; v2 also allows any other JSON value, such as an object. An MCP
 * server checks the content against its own `requestedSchema`, so it gets the
 * content as the client sent it: before, an accept whose content held an
 * object reached the server as `{ action: "accept", content: {} }`. Callers
 * that read a field must check its type.
 */
function acceptedElicitationContent(response: CreateElicitationResponse): Record<string, unknown> {
  if (response.action !== "accept") return {};
  const content: unknown = (response as { content?: unknown }).content;
  return content !== null && typeof content === "object" && !Array.isArray(content)
    ? (content as Record<string, unknown>)
    : {};
}

/**
 * Map an ACP elicitation response back to the MCP `ElicitResult` the SDK expects
 * to hand back to the requesting server.
 */
export function createElicitationResponseToElicitResult(
  response: CreateElicitationResponse,
): ElicitationResult {
  switch (response.action) {
    case "accept":
      // `ElicitResult` types the values as v1 does. The MCP server checks them.
      return {
        action: "accept",
        content: acceptedElicitationContent(response) as ElicitationResult["content"],
      };
    case "decline":
      return { action: "decline" };
    case "cancel":
    default:
      return { action: "cancel" };
  }
}

/**
 * A single question as supplied by the AskUserQuestion tool. Derived from the
 * SDK's input type so the shape stays in sync; the SDK validates the model's
 * tool call against this schema before it reaches us.
 */
export type AskUserQuestion = AskUserQuestionInput["questions"][number];

/**
 * Pull the well-formed questions out of an AskUserQuestion tool input. Returns
 * `null` when there are no usable questions — including the case where every
 * entry is malformed and filtering leaves an empty list — so callers can treat
 * "nothing to ask" uniformly.
 */
export function extractAskUserQuestions(input: Record<string, unknown>): AskUserQuestion[] | null {
  const questions = (input as { questions?: unknown }).questions;
  if (!Array.isArray(questions)) {
    return null;
  }
  const valid = questions.filter(
    (q): q is AskUserQuestion =>
      !!q && typeof q.question === "string" && Array.isArray(q.options) && q.options.length > 0,
  );
  return valid.length > 0 ? valid : null;
}

/** Stable form-field key for the question at the given index. */
function questionFieldKey(index: number): string {
  return `question_${index}`;
}

/**
 * Form-field key for the per-question free-text "custom answer" field that sits
 * alongside `question_<n>`. Mirrors the first-party clients, where every
 * question carries its own "Other" box rather than one form-level field.
 */
function questionCustomFieldKey(index: number): string {
  return `question_${index}_custom`;
}

/**
 * `_meta` key under which a bridged enum option carries its `preview`, the one
 * option field ACP's `EnumOption` still has no slot for (descriptions are
 * first-class as of schema 1.19). Namespaced like the agent's other `_meta`
 * extensions (`_claude/...`).
 */
const OPTION_META_KEY = "_claude/askUserQuestionOption";

/**
 * Render the AskUserQuestion tool's questions as an ACP form elicitation.
 *
 * Fields are keyed by a short stable id (`question_<n>`) rather than the full
 * question text, so the question text appears in exactly one place per field.
 * Single-select questions use a titled `oneOf` enum; multi-select questions use
 * an array with a titled `anyOf` item enum. The enum `const` is always the
 * option label, since that is what the tool records as the answer; an option's
 * secondary text travels in the enum option's own `description` field.
 *
 * Each question is followed by its own optional free-text "custom answer" field
 * (`question_<n>_custom`), mirroring the CLI's per-question "Other" box: the
 * user can type their own answer instead of picking an option, add it to a
 * multi-select's picks, or attach it as a note to a single-select's pick (see
 * `applyAskElicitationResponse`), scoped to that specific question. Nothing is
 * marked required, so the user can also just skip — matching the built-in tool,
 * which always offers Skip + a free-text box.
 */
export function askUserQuestionsToCreateRequest(
  questions: AskUserQuestion[],
  sessionId: string,
  toolCallId: string | undefined,
  airClient = false,
): CreateElicitationRequest {
  const single = questions.length === 1;
  const properties: Record<string, ElicitationPropertySchema> = {};

  questions.forEach((question, index) => {
    const options: EnumOption[] = question.options.map((option) => {
      const enumOption: EnumOption = {
        const: option.label,
        title: option.label,
      };
      if (option.description) {
        enumOption.description = option.description;
      }
      // The SDK option's `preview` (mockups, code snippets, comparisons shown
      // on focus) still has no structural slot in `EnumOption`, so forward it
      // under ACP's reserved `_meta` extension point for clients that render it.
      if (option.preview) {
        enumOption._meta = { [OPTION_META_KEY]: { preview: option.preview } };
      }
      return enumOption;
    });

    // For a single question the prompt is carried by `message`, so we don't
    // repeat it in the field description. With multiple questions each field
    // needs its own question text.
    const description = single ? undefined : question.question;
    const title = question.header || undefined;

    properties[questionFieldKey(index)] = question.multiSelect
      ? { type: "array", title, description, items: { anyOf: options } }
      : { type: "string", title, description, oneOf: options };

    const customAnswer = { questionId: questionFieldKey(index), isCustomAnswer: true };
    properties[questionCustomFieldKey(index)] = {
      type: "string",
      title: "Other",
      description: question.multiSelect
        ? "Type your own answer to add to your selection above (optional)."
        : "Type your own answer, or add a note to the option you chose above (optional).",
      // Marks the field as the custom answer companion of a select question,
      // under `_meta.jetbrains.air.customAnswer` and under the root key that
      // released AIR versions read. Only AIR gets the marker.
      ...(airClient
        ? {
            _meta: withAirMeta(
              { [LEGACY_AIR_CUSTOM_ANSWER_KEY]: customAnswer },
              AIR_CUSTOM_ANSWER_KEY,
              customAnswer,
            ),
          }
        : {}),
    };
  });

  const requestedSchema: ElicitationSchema = {
    type: "object",
    properties,
  };

  const message = single ? questions[0].question : "Please answer the following questions.";

  return {
    mode: "form",
    sessionId,
    ...(toolCallId ? { toolCallId } : {}),
    message,
    requestedSchema,
  };
}

/** Outcome of an AskUserQuestion elicitation, decoupled from any transport. */
export type AskUserQuestionOutcome =
  { action: "answered"; updatedInput: Record<string, unknown> } | { action: "cancel" };

/**
 * Serialize a multi-select answer the way the CLI's own AskUserQuestion UI
 * does: comma-joined, with any item that itself contains the separator (or a
 * double quote) JSON-quoted. The tool's `call()` splits the string back on the
 * same rule, so a free-text answer like `Redis, not Memcached` stays one item
 * instead of reading as two more picks.
 */
function joinMultiSelectAnswer(items: string[]): string {
  return items
    .map((item) => (item.includes(", ") || item.includes('"') ? JSON.stringify(item) : item))
    .join(", ");
}

/**
 * Fold an ACP elicitation response into the AskUserQuestion tool's input.
 *
 * Selected labels are read back from the indexed form fields and written into
 * `answers` as a `{ [questionText]: label }` map — the key shape the tool's own
 * `call()` reads — with multi-selects comma-joined in the CLI's own quoted form
 * (see `joinMultiSelectAnswer`). A non-empty per-question custom-answer field
 * (`question_<n>_custom`) joins the selection of a multi-select question, where
 * the two fields are independent and filling both means both. For a
 * single-select question it is the answer when nothing was picked (the user
 * typed their own instead), and otherwise travels beside the pick as the
 * tool's own per-question `annotations[question].notes` — the slot the CLI
 * uses for free text attached to a selection and renders to the model as
 * `"Q"="A" notes: ...` — so a client that presents the box as a notes field
 * cannot make the selection disappear. Decline yields empty answers (the model
 * is told the user skipped rather than the turn aborting); cancel — and any
 * custom/future action we don't understand — aborts the tool call.
 *
 * With `reportUnanswered` set (the session's `system/init` listed
 * `ask_user_question_unanswered_v1`), each question left without an answer
 * also gets `annotations[question].unanswered`: `"dismissed"` for every
 * question when the user declined the form, `"skipped"` for a question left
 * blank in an accepted form. The model does not see this field; the tool
 * result carries it for clients and transcripts. A CLI without the capability
 * can refuse the whole answer over it, so it is left out there.
 */
export function applyAskElicitationResponse(
  response: CreateElicitationResponse,
  toolInput: Record<string, unknown>,
  questions: AskUserQuestion[],
  options: { reportUnanswered?: boolean } = {},
): AskUserQuestionOutcome {
  // Typed against the output schema: as of SDK 0.3.319 only it lists
  // `unanswered`, though the tool reads the same annotations from its input.
  const annotations: NonNullable<AskUserQuestionOutput["annotations"]> = {};
  const withAnnotations = (answers: AskUserQuestionOutput["answers"]) => ({
    ...toolInput,
    answers,
    ...(Object.keys(annotations).length > 0 ? { annotations } : {}),
  });

  if (response.action === "decline") {
    if (options.reportUnanswered) {
      for (const question of questions) {
        annotations[question.question] = { unanswered: "dismissed" };
      }
    }
    return { action: "answered", updatedInput: withAnnotations({}) };
  }

  if (response.action !== "accept") {
    return { action: "cancel" };
  }

  const content = acceptedElicitationContent(response);
  // Typed against the tool's own output schema so the answer/response shapes
  // stay in sync with what the built-in tool's call() expects to read back.
  const answers: AskUserQuestionOutput["answers"] = {};
  questions.forEach((question, index) => {
    const custom = content[questionCustomFieldKey(index)];
    const customText = typeof custom === "string" ? custom.trim() : "";

    // A pick is a string, number, or boolean. Other values, such as an object
    // that a v2 client may send, are dropped instead of read as "[object Object]".
    const value = content[questionFieldKey(index)];
    const picks: string[] = (Array.isArray(value) ? value : [value])
      .filter(isElicitationScalar)
      .map(String)
      .filter((item) => item !== "");

    // A multi-select is additive, and the form offers the selection and the
    // custom box as independent fields — a user who fills both means both, so
    // the typed answer joins the checked options.
    if (question.multiSelect) {
      const text = joinMultiSelectAnswer(customText === "" ? picks : [...picks, customText]);
      if (text !== "") {
        answers[question.question] = text;
      }
      return;
    }

    // A single-select question is answered by exactly one thing. With no option
    // picked, the typed text is that answer (the CLI's "Other"). With an option
    // picked as well, the pick stays the answer and the text rides along as the
    // tool's per-question `notes` annotation, so neither is lost. A single-select
    // normally holds one item; the plain join only matters if a client hands
    // back an array for it, and then mirrors the old behavior.
    const picked = picks.join(", ");
    if (picked === "") {
      if (customText !== "") {
        answers[question.question] = customText;
      }
      return;
    }
    answers[question.question] = picked;
    if (customText !== "") {
      annotations[question.question] = { notes: customText };
    }
  });

  if (options.reportUnanswered) {
    for (const question of questions) {
      if (answers[question.question] === undefined) {
        annotations[question.question] = { unanswered: "skipped" };
      }
    }
  }

  return { action: "answered", updatedInput: withAnnotations(answers) };
}

function isElicitationScalar(value: unknown): value is string | number | boolean {
  return typeof value === "string" || typeof value === "number" || typeof value === "boolean";
}

/**
 * Coerce an arbitrary MCP `requestedSchema` into an ACP `ElicitationSchema`.
 * The two are structurally compatible JSON Schemas; we just guarantee the
 * `type: "object"` discriminator and a `properties` map are present. The v2
 * schema requires both, so a client on SDK 2.2.0 or later rejects an MCP
 * schema such as `{ type: "object" }` that has no `properties`.
 */
function normalizeElicitationSchema(
  schema: Record<string, unknown> | undefined,
): ElicitationSchema {
  if (!schema || typeof schema !== "object") {
    return { type: "object", properties: {} };
  }
  const { properties } = schema as Partial<ElicitationSchema>;
  return { ...(schema as ElicitationSchema), type: "object", properties: properties ?? {} };
}

/**
 * The `request_user_dialog` kind the CLI emits when a model refusal has a
 * fallback available but needs user consent before retrying (e.g. Claude Fable
 * declining a request with Opus available as the fallback). Declaring this
 * kind in `supportedDialogKinds` is the opt-in: the CLI fails closed and never
 * emits an undeclared kind — the flow degrades to the classic refusal error
 * ending the turn.
 */
export const REFUSAL_FALLBACK_DIALOG_KIND = "refusal_fallback_prompt";

/**
 * Payload of the `refusal_fallback_prompt` dialog. The dialog protocol
 * transports payloads opaquely, so this shape is recovered from the CLI's own
 * schema (v2.1.177): `originalModel`/`fallbackModel` are required strings;
 * `apiRefusalCategory` (nullable), `guidanceText`, and
 * `retractedMessageUuids` are optional. We ignore `retractedMessageUuids` —
 * ACP has no way to retract already-streamed chunks.
 */
export type RefusalFallbackPrompt = {
  originalModel: string;
  fallbackModel: string;
  apiRefusalCategory: string | null;
  guidanceText?: string;
};

/**
 * Validate the opaque dialog payload into a {@link RefusalFallbackPrompt}.
 * Returns `null` when the required fields are missing or mistyped (a newer CLI
 * may reshape the payload), so the caller can cancel the dialog and let the
 * CLI apply its default behavior instead of rendering something misleading.
 */
export function extractRefusalFallbackPrompt(
  payload: Record<string, unknown>,
): RefusalFallbackPrompt | null {
  const { originalModel, fallbackModel, apiRefusalCategory, guidanceText } = payload;
  if (typeof originalModel !== "string" || typeof fallbackModel !== "string") {
    return null;
  }
  return {
    originalModel,
    fallbackModel,
    apiRefusalCategory: typeof apiRefusalCategory === "string" ? apiRefusalCategory : null,
    ...(typeof guidanceText === "string" && guidanceText ? { guidanceText } : {}),
  };
}

/** Form-field key carrying the user's choice in the refusal-fallback form. */
const REFUSAL_FALLBACK_CHOICE_KEY = "choice";

/** Wire values of the dialog's result enum (CLI schema). `edit_prompt` is
 *  deliberately not offered: in the CLI it prefills the composer with the
 *  refused prompt for edit-and-retry, and ACP has no composer-prefill surface
 *  — the user can simply edit and resend on their own. */
const RETRY_FALLBACK_RESULT = "retry_fallback";
const KEEP_REFUSAL_RESULT = "cancelled";

/**
 * Render the refusal-fallback consent prompt as an ACP form elicitation: a
 * single-select between retrying on the fallback model and keeping the
 * refusal. The enum `const`s are the dialog's wire result values, so the
 * response maps back without a translation table.
 */
export function refusalFallbackToCreateRequest(
  prompt: RefusalFallbackPrompt,
  sessionId: string,
): CreateElicitationRequest {
  const category = prompt.apiRefusalCategory ? ` (${prompt.apiRefusalCategory})` : "";
  const guidance = prompt.guidanceText ? `\n\n${prompt.guidanceText}` : "";
  return {
    mode: "form",
    sessionId,
    message:
      `${prompt.originalModel} declined this request${category}. ` +
      `Retry with ${prompt.fallbackModel}?` +
      guidance,
    requestedSchema: {
      type: "object",
      properties: {
        [REFUSAL_FALLBACK_CHOICE_KEY]: {
          type: "string",
          oneOf: [
            {
              const: RETRY_FALLBACK_RESULT,
              title: `Retry with ${prompt.fallbackModel}`,
              description: `The session continues on ${prompt.fallbackModel}.`,
            },
            {
              const: KEEP_REFUSAL_RESULT,
              title: "Keep the refusal",
              description: "You can send a new message.",
            },
          ],
        },
      },
    },
  };
}

/**
 * Map the elicitation response back to the dialog's result enum. Only an
 * explicit accept-with-retry resolves to `retry_fallback`; decline, cancel, a
 * skipped field, or an unrecognized value all keep the refusal — the dialog's
 * own default — so a dismissed or half-filled form can never trigger a model
 * switch the user didn't ask for.
 */
export function refusalFallbackResultFromResponse(response: CreateElicitationResponse): string {
  const choice = acceptedElicitationContent(response)[REFUSAL_FALLBACK_CHOICE_KEY];
  return choice === RETRY_FALLBACK_RESULT ? RETRY_FALLBACK_RESULT : KEEP_REFUSAL_RESULT;
}

/**
 * What became of an elicitation: the client's response, or `withdrawn` when
 * its signal aborted first (Claude Code withdrew the request, or a cancel or
 * close of the session did). A response that comes after the abort counts as
 * withdrawn too.
 */
export type ElicitationOutcome = { response: CreateElicitationResponse } | { withdrawn: true };

/**
 * Sends an elicitation with `send` and reads what became of it. A rejection
 * while `signal` is not aborted, such as the client's -32602 for a form it
 * cannot render, is rethrown.
 */
export async function elicitUnlessWithdrawn(
  signal: AbortSignal,
  send: (signal: AbortSignal) => Promise<CreateElicitationResponse>,
): Promise<ElicitationOutcome> {
  try {
    const response = await send(signal);
    return signal.aborted ? { withdrawn: true } : { response };
  } catch (error) {
    if (signal.aborted) return { withdrawn: true };
    throw error;
  }
}

/** Sends a built elicitation to the client (see {@link elicitUnlessWithdrawn}). */
export type Elicit = (request: CreateElicitationRequest) => Promise<ElicitationOutcome>;

type ElicitationLogger = { error: (...args: unknown[]) => void };

/**
 * Answers an MCP server's elicitation (the SDK's `onElicitation`) through
 * `elicit`. Only the user's own decline reaches the MCP server as `decline`.
 * A request the user never saw, because the client did not advertise its
 * mode, it cannot be represented, the client answered with an error, or it
 * was withdrawn, is a `cancel`.
 */
export async function answerMcpElicitation(
  request: ElicitationRequest,
  {
    sessionId,
    support,
    elicit,
    logger,
  }: { sessionId: string; support: ElicitationSupport; elicit: Elicit; logger: ElicitationLogger },
): Promise<ElicitationResult> {
  const isUrl = request.mode === "url";
  if ((isUrl && !support.url) || (!isUrl && !support.form)) {
    return { action: "cancel" };
  }
  const createRequest = mcpElicitationToCreateRequest(request, sessionId);
  if (!createRequest) {
    return { action: "cancel" };
  }
  try {
    const outcome = await elicit(createRequest);
    if ("withdrawn" in outcome) return { action: "cancel" };
    return createElicitationResponseToElicitResult(outcome.response);
  } catch (error) {
    // The client answered with an error, e.g. -32602 for a form it cannot
    // render, so the user never saw the request. MCP's `decline` means the
    // user refused; `cancel` means they made no explicit choice.
    logger.error(`Failed to forward MCP elicitation: ${error}`);
    return { action: "cancel" };
  }
}

/**
 * Presents the AskUserQuestion tool's questions as a form elicitation of
 * `sessionId` through `elicit`, and returns the answers as the tool's
 * `updatedInput`. A withdrawn question, or one the user cancelled, aborts the
 * tool use. `reportUnanswered` is read once the answer is in (see
 * {@link applyAskElicitationResponse}).
 */
export async function answerAskUserQuestion(
  toolInput: Record<string, unknown>,
  {
    sessionId,
    toolUseID,
    airClient = false,
    reportUnanswered,
    elicit,
    logger,
  }: {
    sessionId: string;
    toolUseID: string;
    airClient?: boolean;
    reportUnanswered: () => boolean;
    elicit: Elicit;
    logger: ElicitationLogger;
  },
): Promise<PermissionResult> {
  const questions = extractAskUserQuestions(toolInput);
  if (!questions) {
    return { behavior: "deny", message: "AskUserQuestion called with no valid questions." };
  }
  const createRequest = askUserQuestionsToCreateRequest(questions, sessionId, toolUseID, airClient);
  let outcome: ElicitationOutcome;
  try {
    outcome = await elicit(createRequest);
  } catch (error) {
    logger.error(`Failed to present AskUserQuestion elicitation: ${error}`);
    return { behavior: "deny", message: "Could not present the question to the user." };
  }
  if ("withdrawn" in outcome) {
    throw new Error("Tool use aborted");
  }
  const answer = applyAskElicitationResponse(outcome.response, toolInput, questions, {
    reportUnanswered: reportUnanswered(),
  });
  if (answer.action === "cancel") {
    throw new Error("Tool use aborted");
  }
  return { behavior: "allow", updatedInput: answer.updatedInput };
}

/**
 * Answers a `request_user_dialog` of Claude Code (the SDK's `onUserDialog`).
 * Only kinds declared in `supportedDialogKinds` are ever emitted; everything
 * unexpected is answered `cancelled` (the required answer for unrecognized
 * kinds), which applies the dialog's default behavior CLI-side. The only
 * declared kind is the refusal-fallback consent prompt, rendered as a form
 * elicitation. A failed or withdrawn elicitation is `cancelled` too: the CLI
 * then keeps the refusal rather than switching models without consent.
 */
export async function answerUserDialog(
  request: UserDialogRequest,
  { sessionId, elicit, logger }: { sessionId: string; elicit: Elicit; logger: ElicitationLogger },
): Promise<UserDialogResult> {
  if (request.dialogKind !== REFUSAL_FALLBACK_DIALOG_KIND) {
    return { behavior: "cancelled" };
  }
  const prompt = extractRefusalFallbackPrompt(request.payload);
  if (!prompt) {
    logger.error(
      `refusal_fallback_prompt payload had an unexpected shape; cancelling the dialog: ${JSON.stringify(request.payload)}`,
    );
    return { behavior: "cancelled" };
  }
  let outcome: ElicitationOutcome;
  try {
    outcome = await elicit(refusalFallbackToCreateRequest(prompt, sessionId));
  } catch (error) {
    logger.error(`Failed to present refusal fallback elicitation: ${error}`);
    return { behavior: "cancelled" };
  }
  if ("withdrawn" in outcome) {
    return { behavior: "cancelled" };
  }
  return { behavior: "completed", result: refusalFallbackResultFromResponse(outcome.response) };
}