mercury-agent 0.4.5

package/src/agent/model-capabilities-core.ts

@@ -0,0 +1,23 @@
+/**
+ * Minimal model capability types + defaults (no yaml/zod).
+ * Used by container-entry; the full agent image only copies this file, not model-capabilities.ts.
+ */
+
+export type ModelCapabilities = {
+  tools: boolean;
+  vision: boolean;
+  audio_input: boolean;
+  audio_output: boolean;
+  extended_thinking: boolean;
+};
+
+export type ModelCapabilityKey = keyof ModelCapabilities;
+
+/** Fallback when no builtin / YAML / env match. */
+export const DEFAULT_CAPABILITIES: ModelCapabilities = {
+  tools: true,
+  vision: false,
+  audio_input: false,
+  audio_output: false,
+  extended_thinking: false,
+};

package/src/agent/model-capabilities.ts

@@ -0,0 +1,231 @@
+/**
+ * Model capability flags — used to adapt prompts, pi tool flags, and skill installation.
+ */
+
+import { existsSync, readFileSync } from "node:fs";
+import path from "node:path";
+import { getModels, type KnownProvider } from "@mariozechner/pi-ai";
+import { parse as parseYaml } from "yaml";
+import { z } from "zod";
+import type { ModelLeg } from "../config.js";
+import {
+  DEFAULT_CAPABILITIES,
+  type ModelCapabilities,
+  type ModelCapabilityKey,
+} from "./model-capabilities-core.js";
+
+export type {
+  ModelCapabilities,
+  ModelCapabilityKey,
+} from "./model-capabilities-core.js";
+export { DEFAULT_CAPABILITIES } from "./model-capabilities-core.js";
+
+export type CapabilityResolveSource = "env" | "yaml" | "builtin" | "default";
+
+export type ResolvedModelCapabilities = {
+  capabilities: ModelCapabilities;
+  source: CapabilityResolveSource;
+};
+
+function mergePartialCapabilities(
+  partial: Partial<Record<ModelCapabilityKey, unknown>>,
+): ModelCapabilities {
+  const out = { ...DEFAULT_CAPABILITIES };
+  for (const k of Object.keys(partial) as ModelCapabilityKey[]) {
+    const v = partial[k];
+    if (typeof v === "boolean") out[k] = v;
+  }
+  return out;
+}
+
+const partialCapsSchema = z.object({
+  tools: z.boolean().optional(),
+  vision: z.boolean().optional(),
+  audio_input: z.boolean().optional(),
+  audio_output: z.boolean().optional(),
+  extended_thinking: z.boolean().optional(),
+});
+
+const yamlFileSchema = z.object({
+  models: z.record(z.string(), partialCapsSchema),
+});
+
+export type UserModelCapabilitiesMap = Map<string, ModelCapabilities>;
+
+/** Path to optional user overrides: `<dataDir>/model-capabilities.yaml` */
+export function modelCapabilitiesYamlPath(dataDirAbsolute: string): string {
+  return path.join(dataDirAbsolute, "model-capabilities.yaml");
+}
+
+/**
+ * Load `.mercury/model-capabilities.yaml` (under resolved data dir).
+ * Returns null if missing or invalid (callers may log).
+ */
+export function loadUserModelCapabilitiesMap(
+  dataDirAbsolute: string,
+): UserModelCapabilitiesMap | null {
+  const yamlPath = modelCapabilitiesYamlPath(dataDirAbsolute);
+  if (!existsSync(yamlPath)) return null;
+  try {
+    const raw = readFileSync(yamlPath, "utf8");
+    const doc = parseYaml(raw) as unknown;
+    const parsed = yamlFileSchema.safeParse(doc);
+    if (!parsed.success) return null;
+    const map = new Map<string, ModelCapabilities>();
+    for (const [modelId, partial] of Object.entries(parsed.data.models)) {
+      map.set(modelId.trim(), mergePartialCapabilities(partial));
+    }
+    return map;
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Look up capabilities for a model using pi's MODELS registry as the source of truth.
+ * Derives: vision (image in input), audio_input (audio in input), extended_thinking (reasoning).
+ * tools defaults to true — pi has no tools field; all pi-known models support tool use.
+ * audio_output is not tracked by pi; always false.
+ */
+function matchBuiltinCapabilities(
+  provider: string,
+  modelId: string,
+): ModelCapabilities | null {
+  // Cast to KnownProvider — getModels returns [] for unrecognised providers at runtime
+  const model = getModels(provider as KnownProvider).find(
+    (m) => m.id === modelId,
+  );
+  if (!model) return null;
+  return {
+    tools: true,
+    vision: model.input.includes("image"),
+    audio_input: false, // pi types input as ("text"|"image")[]; no audio models exist yet
+    audio_output: false,
+    extended_thinking: model.reasoning,
+  };
+}
+
+/** Parse MERCURY_MODEL_CAPABILITIES JSON; returns null if unset or invalid. */
+export function parseModelCapabilitiesEnv(
+  raw: string | undefined,
+): ModelCapabilities | null {
+  const trimmed = raw?.trim();
+  if (!trimmed) return null;
+  try {
+    const json = JSON.parse(trimmed) as unknown;
+    const r = partialCapsSchema.safeParse(json);
+    if (!r.success) return null;
+    return mergePartialCapabilities(r.data);
+  } catch {
+    return null;
+  }
+}
+
+/**
+ * Resolve capabilities for one model id (single leg).
+ * Priority: env global override → YAML exact model key → pi MODELS lookup → default.
+ */
+export function resolveModelCapabilitiesWithSource(
+  modelId: string,
+  provider: string,
+  userMap: UserModelCapabilitiesMap | null,
+  envCaps: ModelCapabilities | null,
+): ResolvedModelCapabilities {
+  if (envCaps) {
+    return { capabilities: envCaps, source: "env" };
+  }
+
+  const exact = userMap?.get(modelId.trim());
+  if (exact) {
+    return { capabilities: exact, source: "yaml" };
+  }
+
+  const builtin = matchBuiltinCapabilities(provider, modelId);
+  if (builtin) {
+    return { capabilities: builtin, source: "builtin" };
+  }
+
+  return { capabilities: { ...DEFAULT_CAPABILITIES }, source: "default" };
+}
+
+export function resolveModelCapabilities(
+  modelId: string,
+  provider: string,
+  userMap: UserModelCapabilitiesMap | null,
+  envCaps: ModelCapabilities | null,
+): ModelCapabilities {
+  return resolveModelCapabilitiesWithSource(modelId, provider, userMap, envCaps)
+    .capabilities;
+}
+
+/** Capabilities for each leg in the model chain (same length as `chain`). */
+export function resolveModelChainCapabilities(
+  chain: ModelLeg[],
+  dataDirAbsolute: string,
+  envCaps: ModelCapabilities | null,
+): {
+  chainCaps: ModelCapabilities[];
+  userMap: UserModelCapabilitiesMap | null;
+} {
+  const userMap = loadUserModelCapabilitiesMap(dataDirAbsolute);
+  const chainCaps = chain.map((leg) =>
+    resolveModelCapabilities(leg.model, leg.provider, userMap, envCaps),
+  );
+  return { chainCaps, userMap };
+}
+
+export function chainSupportsRequirements(
+  requires: ModelCapabilityKey[],
+  chainCaps: ModelCapabilities[],
+): boolean {
+  if (requires.length === 0) return true;
+  return chainCaps.some((caps) => requires.every((key) => caps[key] === true));
+}
+
+/** Log warnings for models that fell back to defaults (once per distinct model id). */
+export function logUnknownModelCapabilityWarnings(
+  chain: ModelLeg[],
+  dataDirAbsolute: string,
+  envCaps: ModelCapabilities | null,
+  log: { warn: (msg: string, obj?: Record<string, unknown>) => void },
+): void {
+  if (envCaps) return;
+  const userMap = loadUserModelCapabilitiesMap(dataDirAbsolute);
+  const seen = new Set<string>();
+
+  for (const leg of chain) {
+    const id = leg.model.trim();
+    if (seen.has(id)) continue;
+    seen.add(id);
+
+    const { source } = resolveModelCapabilitiesWithSource(
+      id,
+      leg.provider,
+      userMap,
+      null,
+    );
+    if (source === "default") {
+      log.warn(
+        `Model "${id}" not in built-in capability map and not in model-capabilities.yaml; assuming default capabilities (tools=true, vision=false). Set MERCURY_MODEL_CAPABILITIES or add .mercury/model-capabilities.yaml to override.`,
+        { model: id },
+      );
+    }
+  }
+}
+
+export function logExtensionCapabilityMismatches(
+  extensions: Array<{ name: string; requires?: ModelCapabilityKey[] }>,
+  chainCaps: ModelCapabilities[],
+  log: { warn: (msg: string, obj?: Record<string, unknown>) => void },
+): void {
+  for (const ext of extensions) {
+    const req = ext.requires;
+    if (!req?.length) continue;
+    if (!chainSupportsRequirements(req, chainCaps)) {
+      log.warn(
+        `Extension "${ext.name}" requires ${req.join(", ")} but no model leg in MERCURY_MODEL_CHAIN supports it; extension skills will not be installed.`,
+        { extension: ext.name, requires: req },
+      );
+    }
+  }
+}

package/src/agent/pi-failure-class.ts

@@ -0,0 +1,83 @@
+/**
+ * Classify pi CLI failures for retry vs next-model vs fail-fast.
+ * Uses stderr/stdout text heuristics (pi does not expose structured HTTP codes here).
+ */
+export type PiFailureClass = "failFast" | "retryable" | "fallbackable";
+
+export function classifyPiFailure(text: string): PiFailureClass {
+  if (/provider ["']cursor["'] is no longer supported/i.test(text)) {
+    return "failFast";
+  }
+
+  if (
+    /\b401\b|\b403\b|invalid\s+api\s+key|incorrect\s+api\s+key|authentication\s+failed|invalid\s+authentication|unauthorized|access\s+denied/i.test(
+      text,
+    )
+  ) {
+    return "failFast";
+  }
+
+  if (
+    /context\s+length|maximum\s+context|token\s+limit|too\s+many\s+tokens|prompt\s+is\s+too\s+long|max\s+tokens|request\s+too\s+large|maximum\s+tokens/i.test(
+      text,
+    )
+  ) {
+    return "fallbackable";
+  }
+
+  if (
+    /tool\s+(use|calling)\s+not\s+supported|tools?\s+not\s+supported|function\s+calling\s+not\s+(supported|available)|does\s+not\s+support\s+tools?|model\s+does\s+not\s+support\s+(tools?|function)|no\s+tool\s+use|unsupported.*\btools?\b/i.test(
+      text,
+    )
+  ) {
+    return "fallbackable";
+  }
+
+  if (
+    /tool\s+call\s+validation\s+failed|attempted\s+to\s+call\s+tool.*which\s+was\s+not\s+in\s+request\.tools/i.test(
+      text,
+    )
+  ) {
+    return "fallbackable";
+  }
+
+  // The model/provider returned a response the runtime couldn't turn into usable
+  // output. Two forms:
+  //  1. An unmapped finish reason — pi-ai's exhaustive switch throws
+  //     "Unhandled stop reason: <X>" (e.g. Gemini's MALFORMED_RESPONSE, which is
+  //     absent from the bundled @google/genai enum).
+  //  2. A raw malformed finish-reason enum token surfaced directly.
+  // An identical retry against the same leg reproduces it and just burns the chain
+  // budget, so fall through to the next model leg instead. The enum arm is
+  // case-sensitive on the uppercase token so prose like "malformed response" in an
+  // unrelated message can't trip it.
+  if (
+    /unhandled\s+stop\s+reason/i.test(text) ||
+    /\bMALFORMED_(RESPONSE|FUNCTION_CALL|TOOL_CALL)\b/.test(text)
+  ) {
+    return "fallbackable";
+  }
+
+  if (/\b429\b|rate[_\s]+limit/i.test(text)) {
+    return "fallbackable";
+  }
+
+  if (
+    /\b402\b|insufficient\s+credits?|not\s+enough\s+credits?|purchase\s+(more\s+)?credits?|no\s+credits?/i.test(
+      text,
+    )
+  ) {
+    return "fallbackable";
+  }
+
+  if (
+    /\b502\b|\b503\b|\b504\b|timeout|timed\s+out|ETIMEDOUT|ECONNRESET|temporarily\s+unavailable|overload|try\s+again|service\s+unavailable|bad\s+gateway|gateway\s+timeout/i.test(
+      text,
+    )
+  ) {
+    return "retryable";
+  }
+
+  // Default: assume transient; bounded retries + chain limit prevent infinite spin.
+  return "retryable";
+}

package/src/agent/pi-jsonl-parser.ts

@@ -0,0 +1,306 @@
+import type { TokenUsage } from "../types.js";
+
+export interface PiJsonlParseResult {
+  reply: string;
+  usage?: TokenUsage;
+  /**
+   * When set, the Pi run failed at the model layer (e.g. HTTP 429 inside JSONL)
+   * while `pi` still exited 0. Host should treat this as a failed invocation
+   * (retry / next model leg), not as a user-visible reply string.
+   */
+  piFailureMessage?: string;
+  /**
+   * Set when the model emitted a tool call as raw text (JSON or XML format) instead
+   * of a proper tool_use block that pi can execute. The leaked call was stripped from
+   * the reply. Caller may retry with tools disabled to get a clean text-only response.
+   */
+  hadToolLeakage?: boolean;
+}
+
+const USAGE_EVENTS = new Set(["turn_end", "message_end"]);
+
+/** Pi `--print --mode json` lines — if we see these, never fall back to raw stdout. */
+const PI_STRUCTURED_EVENT_TYPES = new Set([
+  "session",
+  "agent_start",
+  "agent_end",
+  "turn_start",
+  "turn_end",
+  "message_start",
+  "message_end",
+]);
+
+function stdoutLooksLikeStructuredPiJsonl(stdout: string): boolean {
+  for (const line of stdout.trim().split("\n")) {
+    if (!line.trim()) continue;
+    let event: Record<string, unknown>;
+    try {
+      event = JSON.parse(line) as Record<string, unknown>;
+    } catch {
+      continue;
+    }
+    const t = event.type;
+    if (typeof t === "string" && PI_STRUCTURED_EVENT_TYPES.has(t)) return true;
+  }
+  return false;
+}
+
+/** Model/API failure embedded in an assistant JSONL message (pi still exits 0). */
+function extractAssistantFailure(
+  msg: Record<string, unknown>,
+): string | undefined {
+  const stop = msg.stopReason;
+  const errRaw = msg.errorMessage;
+
+  if (stop === "error" || stop === "aborted") {
+    if (typeof errRaw === "string" && errRaw.trim()) return errRaw.trim();
+    return `stopReason: ${String(stop)}`;
+  }
+
+  if (typeof errRaw === "string" && errRaw.trim()) return errRaw.trim();
+
+  return undefined;
+}
+
+function usageRecordHasSignal(u: Record<string, unknown>): boolean {
+  const input =
+    Number(u.input ?? u.input_tokens ?? u.promptTokens ?? u.prompt_tokens) || 0;
+  const output =
+    Number(
+      u.output ?? u.output_tokens ?? u.completionTokens ?? u.completion_tokens,
+    ) || 0;
+  const total = Number(u.totalTokens ?? u.total_tokens) || 0;
+  const cacheRead = Number(u.cacheRead ?? u.cache_read) || 0;
+  const cacheWrite = Number(u.cacheWrite ?? u.cache_write) || 0;
+  const cost = u.cost as Record<string, number> | undefined;
+  const costTotal = cost?.total ?? 0;
+  return (
+    input > 0 ||
+    output > 0 ||
+    total > 0 ||
+    cacheRead > 0 ||
+    cacheWrite > 0 ||
+    costTotal > 0
+  );
+}
+
+/**
+ * Some providers/models put tool-style content in the assistant *text* channel
+ * instead of structured tool calls. Two known formats:
+ *   1. JSON blob:  `bash{"command":"mrctl …"}`
+ *   2. XML format: `<function name="bash" parameters="{…}" />`
+ * Strip these so raw tool calls don't leak into Telegram/WhatsApp.
+ * Returns any meaningful text that preceded the blob, or empty string.
+ */
+export function sanitizeLeakedToolCallText(text: string): string {
+  const raw = text.trim();
+  if (!raw) return raw;
+
+  // XML format: <function name="bash" ...> or <function name="bash" .../>
+  const xmlFuncMatch = raw.match(/^([\s\S]*?)<function\s+name=["']bash["']/i);
+  if (xmlFuncMatch) {
+    return xmlFuncMatch[1].trim();
+  }
+
+  // JSON format: bash{"command":...}
+  const jsonStart = raw.indexOf('{"command"');
+  if (jsonStart < 0) return raw;
+
+  const prefix = raw.slice(0, jsonStart).trim();
+
+  if (/^(bash|sh)/i.test(prefix) || !prefix) {
+    return "";
+  }
+
+  return prefix;
+}
+
+function addUsageTotals(
+  u: Record<string, unknown>,
+  into: {
+    inputSum: number;
+    outputSum: number;
+    totalOnlySum: number;
+    cacheReadSum: number;
+    cacheWriteSum: number;
+    costSum: number;
+  },
+): void {
+  const input =
+    Number(u.input ?? u.input_tokens ?? u.promptTokens ?? u.prompt_tokens) || 0;
+  const output =
+    Number(
+      u.output ?? u.output_tokens ?? u.completionTokens ?? u.completion_tokens,
+    ) || 0;
+  const total = Number(u.totalTokens ?? u.total_tokens) || 0;
+  into.inputSum += input;
+  into.outputSum += output;
+  if (input === 0 && output === 0 && total > 0) into.totalOnlySum += total;
+  into.cacheReadSum += Number(u.cacheRead ?? u.cache_read) || 0;
+  into.cacheWriteSum += Number(u.cacheWrite ?? u.cache_write) || 0;
+  const cost = u.cost as Record<string, number> | undefined;
+  into.costSum += cost?.total ?? 0;
+}
+
+/**
+ * Parse pi `--print --mode json` stdout: JSONL events where assistant text and
+ * usage appear on `message_end` and/or `turn_end` (pi-agent-core emits both;
+ * usage is often only on `message_end`). Sums usage across matching events in
+ * one process. If both event types carry usage, only `message_end` is summed to
+ * avoid double-counting the same turn.
+ */
+export function parsePiPrintJsonlOutput(stdout: string): PiJsonlParseResult {
+  const lines = stdout.trim().split("\n").filter(Boolean);
+
+  let usageSource: "message_end" | "turn_end" | null = null;
+  for (const line of lines) {
+    let event: Record<string, unknown>;
+    try {
+      event = JSON.parse(line) as Record<string, unknown>;
+    } catch {
+      continue;
+    }
+    if (event.type !== "message_end") continue;
+    const msg = event.message as Record<string, unknown> | undefined;
+    if (msg?.role !== "assistant") continue;
+    const u = msg.usage as Record<string, unknown> | undefined;
+    if (u && usageRecordHasSignal(u)) {
+      usageSource = "message_end";
+      break;
+    }
+  }
+  if (!usageSource) {
+    for (const line of lines) {
+      let event: Record<string, unknown>;
+      try {
+        event = JSON.parse(line) as Record<string, unknown>;
+      } catch {
+        continue;
+      }
+      if (event.type !== "turn_end") continue;
+      const msg = event.message as Record<string, unknown> | undefined;
+      if (msg?.role !== "assistant") continue;
+      const u = msg.usage as Record<string, unknown> | undefined;
+      if (u && usageRecordHasSignal(u)) {
+        usageSource = "turn_end";
+        break;
+      }
+    }
+  }
+
+  let reply = "";
+  const sums = {
+    inputSum: 0,
+    outputSum: 0,
+    totalOnlySum: 0,
+    cacheReadSum: 0,
+    cacheWriteSum: 0,
+    costSum: 0,
+  };
+  let structuredTurns = 0;
+  let lastModel: string | undefined;
+  let lastProvider: string | undefined;
+  let lastAssistantFailure: string | undefined;
+
+  for (const line of lines) {
+    let event: Record<string, unknown>;
+    try {
+      event = JSON.parse(line) as Record<string, unknown>;
+    } catch {
+      continue;
+    }
+
+    const et = event.type;
+    if (!USAGE_EVENTS.has(et as string)) continue;
+
+    const msg = event.message as Record<string, unknown> | undefined;
+    if (msg?.role !== "assistant") continue;
+
+    structuredTurns += 1;
+
+    const content = msg.content as
+      | Array<{ type: string; text?: string }>
+      | undefined;
+    const textFromContent = content
+      ? content
+          .filter((c) => c.type === "text" && c.text)
+          .map((c) => c.text)
+          .join("\n")
+      : "";
+
+    if (textFromContent) {
+      reply = textFromContent;
+      lastAssistantFailure = undefined;
+    } else {
+      const fd = extractAssistantFailure(msg);
+      if (fd) lastAssistantFailure = fd;
+    }
+
+    if (usageSource && et === usageSource) {
+      const u = msg.usage as Record<string, unknown> | undefined;
+      if (u) addUsageTotals(u, sums);
+    }
+    if (typeof msg.model === "string" && msg.model) lastModel = msg.model;
+    if (typeof msg.provider === "string" && msg.provider)
+      lastProvider = msg.provider;
+  }
+
+  const structuredStream =
+    structuredTurns > 0 || stdoutLooksLikeStructuredPiJsonl(stdout);
+  let piFailureMessage: string | undefined;
+  if (lastAssistantFailure && !reply.trim()) {
+    piFailureMessage = lastAssistantFailure;
+  }
+
+  if (!reply.trim() && !piFailureMessage) {
+    if (structuredStream) {
+      reply = "Done.";
+    } else {
+      reply = stdout.trim() || "Done.";
+    }
+  }
+
+  if (piFailureMessage) {
+    return { reply: "", piFailureMessage };
+  }
+
+  const replyBeforeSanitize = reply.trim();
+  reply = sanitizeLeakedToolCallText(reply);
+  const hadToolLeakage = reply !== replyBeforeSanitize;
+  if (!reply.trim()) {
+    reply = "Done.";
+  }
+
+  const hasUsage =
+    sums.inputSum > 0 ||
+    sums.outputSum > 0 ||
+    sums.totalOnlySum > 0 ||
+    sums.cacheReadSum > 0 ||
+    sums.cacheWriteSum > 0 ||
+    sums.costSum > 0;
+
+  let usage: TokenUsage | undefined;
+  if (structuredTurns > 0 && hasUsage) {
+    const totalFromIO = sums.inputSum + sums.outputSum;
+    const totalTokens =
+      totalFromIO > 0
+        ? totalFromIO
+        : sums.totalOnlySum > 0
+          ? sums.totalOnlySum
+          : 0;
+    usage = {
+      inputTokens: sums.inputSum,
+      outputTokens: sums.outputSum,
+      ...(totalTokens > 0 ? { totalTokens } : {}),
+      ...(sums.cacheReadSum > 0 ? { cacheReadTokens: sums.cacheReadSum } : {}),
+      ...(sums.cacheWriteSum > 0
+        ? { cacheWriteTokens: sums.cacheWriteSum }
+        : {}),
+      ...(sums.costSum > 0 ? { cost: sums.costSum } : {}),
+      ...(lastModel ? { model: lastModel } : {}),
+      ...(lastProvider ? { provider: lastProvider } : {}),
+    };
+  }
+
+  return { reply, usage, ...(hadToolLeakage ? { hadToolLeakage: true } : {}) };
+}

package/src/agent/preferences-prompt.ts

@@ -0,0 +1,20 @@
+/** Escape text for XML element bodies. */
+export function escapeXmlText(s: string): string {
+  return s.replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;");
+}
+
+function escapeXmlAttr(s: string): string {
+  return escapeXmlText(s).replace(/"/g, "&quot;");
+}
+
+/** Build `<preferences>...</preferences>` block for the agent user prompt, or null if empty. */
+export function formatPreferencesXml(
+  preferences?: Array<{ key: string; value: string }>,
+): string | null {
+  if (!preferences?.length) return null;
+  const lines = preferences.map(
+    (p) =>
+      `  <pref key="${escapeXmlAttr(p.key)}">${escapeXmlText(p.value)}</pref>`,
+  );
+  return ["<preferences>", ...lines, "</preferences>"].join("\n");
+}