pi-sage 0.2.0

package/.pi/extensions/sage/index.ts

@@ -0,0 +1,659 @@
+import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "@mariozechner/pi-coding-agent";
+import { Type } from "@sinclair/typebox";
+import {
+  evaluateSoftBudget,
+  isEligibleCaller,
+  isHardCostCapExceeded,
+  makeBlockedResult
+} from "./policy.js";
+import { isSageRunnerPolicyError, runSageSingleShot } from "./runner.js";
+import {
+  getSettingsPathForScope,
+  loadSettings,
+  mergeSettings,
+  saveSettings,
+  type SageSettings,
+  type SettingsScope
+} from "./settings.js";
+import { checkVolumeCaps, getDisallowedCustomTools, resolveToolPolicy } from "./tool-policy.js";
+import type {
+  BlockCode,
+  CallerContext,
+  ReasoningLevel,
+  SageBudgetState,
+  SageConsultInput,
+  SageMode,
+  SageToolResult,
+  ToolProfile
+} from "./types.js";
+
+type InputSource = "interactive" | "rpc" | "extension";
+
+type ModelLike = {
+  provider: string;
+  id: string;
+  name?: string;
+};
+
+const ROLE_HINT_ENV_KEYS = ["TASKPLANE_ROLE", "TASKPLANE_AGENT_ROLE", "PI_AGENT_ROLE", "ORCH_AGENT_ROLE"];
+
+const REASONING_LEVELS: ReasoningLevel[] = ["minimal", "low", "medium", "high", "xhigh"];
+const TOOL_PROFILES: ToolProfile[] = ["none", "read-only-lite", "git-review-readonly", "custom-read-only"];
+
+const SAGE_GUIDANCE = [
+  "You may call `sage_consult` for high-complexity reasoning, ambiguous debugging, risky architectural decisions, or explicit second-opinion requests.",
+  "Do not overuse Sage for routine edits.",
+  "Sage is advisory-only: treat output as recommendations and implement changes yourself.",
+  "Use one focused Sage question per consultation.",
+  "If user explicitly requests Sage/second opinion, prefer at least one Sage consultation unless blocked by hard safety limits."
+].join("\n- ");
+
+export default function registerSageExtension(pi: ExtensionAPI): void {
+  if (process.env.PI_SAGE_SUBAGENT === "1") {
+    return;
+  }
+
+  const budgetState: SageBudgetState = {
+    currentTurn: 0,
+    callsThisTurn: 0,
+    sessionCalls: 0,
+    lastAutoTurn: undefined,
+    sessionCostTotal: 0
+  };
+
+  let lastInputSource: InputSource | undefined;
+
+  pi.on("input", (event) => {
+    if (event.source === "interactive" || event.source === "rpc" || event.source === "extension") {
+      lastInputSource = event.source;
+    }
+    return { action: "continue" };
+  });
+
+  pi.on("turn_start", () => {
+    budgetState.currentTurn += 1;
+    budgetState.callsThisTurn = 0;
+  });
+
+  pi.on("before_agent_start", (event, ctx) => {
+    const loaded = loadSettings(ctx.cwd);
+    if (loaded.settings.enabled === false) return undefined;
+
+    return {
+      systemPrompt: `${event.systemPrompt}\n\nSage guidance:\n- ${SAGE_GUIDANCE}`
+    };
+  });
+
+  pi.registerCommand("sage-settings", {
+    description: "Configure Sage model, limits, and policies",
+    handler: async (_args, ctx) => {
+      await runSettingsWizard(ctx);
+    }
+  });
+
+  pi.registerTool({
+    name: "sage_consult",
+    label: "Sage Consult",
+    description: "Consult an isolated high-reasoning Sage subagent for advisory analysis.",
+    promptSnippet: "Consult Sage for complex second-opinion reasoning.",
+    promptGuidelines: [
+      "Use Sage for high-impact ambiguous tasks or explicit second-opinion requests.",
+      "Set force=true when the user explicitly asks for Sage to bypass soft limits.",
+      "After Sage returns, synthesize recommendations and continue execution."
+    ],
+    parameters: Type.Object({
+      question: Type.String({ description: "Focused question to ask Sage" }),
+      context: Type.Optional(Type.String({ description: "Compressed relevant context" })),
+      files: Type.Optional(Type.Array(Type.String(), { description: "Relevant file paths" })),
+      evidence: Type.Optional(Type.Array(Type.String(), { description: "Logs, errors, or observations" })),
+      objective: Type.Optional(
+        Type.Union([
+          Type.Literal("debug"),
+          Type.Literal("design"),
+          Type.Literal("review"),
+          Type.Literal("refactor"),
+          Type.Literal("general")
+        ])
+      ),
+      urgency: Type.Optional(
+        Type.Union([Type.Literal("low"), Type.Literal("medium"), Type.Literal("high")])
+      ),
+      force: Type.Optional(
+        Type.Boolean({ description: "Bypass soft limits for explicit user-requested consultations" })
+      )
+    }),
+    async execute(_toolCallId, params, signal, _onUpdate, ctx): Promise<SageToolResult> {
+      const input = params as SageConsultInput;
+      const loaded = loadSettings(ctx.cwd);
+      const settings = loaded.settings;
+      const mode: SageMode = input.force ? "user-requested" : "autonomous";
+
+      const modelSelection = resolveModelSpec(settings.model, ctx.modelRegistry.getAvailable(), ctx.model);
+      const resolvedModel = modelSelection.modelArg;
+
+      if (!settings.enabled) {
+        return makeBlockedResult({
+          mode,
+          model: resolvedModel,
+          reasoningLevel: settings.reasoningLevel,
+          blockCode: "disabled",
+          reason: "Sage is disabled",
+          allowedByContext: true,
+          allowedByBudget: false
+        });
+      }
+
+      const callerContext = buildCallerContext(lastInputSource, ctx.hasUI);
+      const eligibility = isEligibleCaller(callerContext);
+      if (!eligibility.ok) {
+        return makeBlockedResult({
+          mode,
+          model: resolvedModel,
+          reasoningLevel: settings.reasoningLevel,
+          blockCode: eligibility.blockCode ?? "unknown-context",
+          reason: eligibility.reason,
+          allowedByContext: false,
+          allowedByBudget: false
+        });
+      }
+
+      if (!resolvedModel) {
+        return makeBlockedResult({
+          mode,
+          model: "unavailable",
+          reasoningLevel: settings.reasoningLevel,
+          blockCode: "model-unavailable",
+          reason: "No available model configured for Sage",
+          allowedByContext: true,
+          allowedByBudget: false
+        });
+      }
+
+      if (isHardCostCapExceeded(settings, budgetState)) {
+        return makeBlockedResult({
+          mode,
+          model: resolvedModel,
+          reasoningLevel: settings.reasoningLevel,
+          blockCode: "cost-cap",
+          reason: "Sage session cost cap exceeded",
+          allowedByContext: true,
+          allowedByBudget: false
+        });
+      }
+
+      const softBudget = evaluateSoftBudget({
+        settings,
+        budgetState,
+        mode,
+        force: input.force === true,
+        questionLength: input.question.length,
+        contextLength: input.context?.length ?? 0
+      });
+
+      if (!softBudget.ok) {
+        return makeBlockedResult({
+          mode,
+          model: resolvedModel,
+          reasoningLevel: settings.reasoningLevel,
+          blockCode: "soft-limit",
+          reason: softBudget.reason,
+          allowedByContext: true,
+          allowedByBudget: false
+        });
+      }
+
+      const disallowedCustomTools = getDisallowedCustomTools(settings.toolPolicy.customAllowedTools);
+      if (disallowedCustomTools.length > 0) {
+        return makeBlockedResult({
+          mode,
+          model: resolvedModel,
+          reasoningLevel: settings.reasoningLevel,
+          blockCode: "tool-disallowed",
+          reason: `Custom tool list contains disallowed tools: ${disallowedCustomTools.join(", ")}`,
+          allowedByContext: true,
+          allowedByBudget: false
+        });
+      }
+
+      const toolPolicy = resolveToolPolicy(settings.toolPolicy);
+
+      try {
+        const runnerResult = await runSageSingleShot({
+          cwd: ctx.cwd,
+          model: resolvedModel,
+          reasoningLevel: settings.reasoningLevel,
+          timeoutMs: settings.timeoutMs,
+          question: input.question,
+          context: input.context,
+          files: input.files,
+          evidence: input.evidence,
+          objective: input.objective,
+          urgency: input.urgency,
+          toolPolicy: settings.toolPolicy,
+          signal
+        });
+
+        const volume = checkVolumeCaps(runnerResult.toolUsage, toolPolicy);
+        if (!volume.ok) {
+          return makeBlockedResult({
+            mode,
+            model: resolvedModel,
+            reasoningLevel: settings.reasoningLevel,
+            blockCode: volume.blockCode ?? "volume-cap",
+            reason: volume.reason,
+            allowedByContext: true,
+            allowedByBudget: false
+          });
+        }
+
+        budgetState.callsThisTurn += 1;
+        budgetState.sessionCalls += 1;
+        if (mode === "autonomous") budgetState.lastAutoTurn = budgetState.currentTurn;
+        if (runnerResult.usage?.costTotal) budgetState.sessionCostTotal += runnerResult.usage.costTotal;
+
+        return {
+          content: [{ type: "text", text: runnerResult.text }],
+          details: {
+            model: resolvedModel,
+            reasoningLevel: settings.reasoningLevel,
+            latencyMs: runnerResult.latencyMs,
+            stopReason: runnerResult.stopReason,
+            usage: runnerResult.usage,
+            toolUsage: runnerResult.toolUsage,
+            policy: {
+              mode,
+              allowedByContext: true,
+              contextReason: "eligible",
+              allowedByBudget: true,
+              budgetReason: softBudget.reason
+            }
+          }
+        };
+      } catch (error) {
+        const message = error instanceof Error ? error.message : "Unknown Sage failure";
+
+        if (isSageRunnerPolicyError(error)) {
+          return makeBlockedResult({
+            mode,
+            model: resolvedModel,
+            reasoningLevel: settings.reasoningLevel,
+            blockCode: error.blockCode,
+            reason: error.message,
+            allowedByContext: true,
+            allowedByBudget: false
+          });
+        }
+
+        if (message.toLowerCase().includes("timed out")) {
+          return makeBlockedResult({
+            mode,
+            model: resolvedModel,
+            reasoningLevel: settings.reasoningLevel,
+            blockCode: "timeout",
+            reason: message,
+            allowedByContext: true,
+            allowedByBudget: false
+          });
+        }
+
+        const blockCode: BlockCode = "execution-error";
+        return {
+          content: [{ type: "text", text: `Sage execution failed: ${message}` }],
+          details: {
+            model: resolvedModel,
+            reasoningLevel: settings.reasoningLevel,
+            latencyMs: 0,
+            stopReason: "error",
+            policy: {
+              mode,
+              allowedByContext: true,
+              contextReason: "eligible",
+              blockCode,
+              allowedByBudget: true,
+              budgetReason: softBudget.reason
+            }
+          },
+          isError: true
+        };
+      }
+    }
+  });
+}
+
+function buildCallerContext(lastInputSource: InputSource | undefined, hasUI: boolean): CallerContext | null {
+  if (!lastInputSource) return null;
+
+  const roleHint = getRoleHint();
+  const isRpcSource = lastInputSource === "rpc";
+  const isSubagent = process.env.PI_SAGE_SUBAGENT === "1";
+  const interactive = hasUI && lastInputSource === "interactive";
+
+  return {
+    session: {
+      interactive
+    },
+    agent: {
+      role: roleHint ?? "primary",
+      isSubagent,
+      isRpcOrchestrated: isRpcSource || Boolean(roleHint && roleHint !== "primary")
+    },
+    runtime: {
+      mode: process.env.CI ? "ci" : hasUI ? (isRpcSource ? "rpc" : "interactive") : "non-interactive"
+    }
+  };
+}
+
+function getRoleHint(): string | undefined {
+  for (const key of ROLE_HINT_ENV_KEYS) {
+    const value = process.env[key];
+    if (value) return value.trim().toLowerCase();
+  }
+  return undefined;
+}
+
+async function runSettingsWizard(ctx: ExtensionCommandContext): Promise<void> {
+  if (!ctx.hasUI) {
+    ctx.ui.notify("/sage-settings requires interactive UI", "warning");
+    return;
+  }
+
+  const loaded = loadSettings(ctx.cwd);
+  let draft = mergeSettings(loaded.settings);
+  let scope: SettingsScope = loaded.source === "global" ? "global" : "project";
+
+  while (true) {
+    const action = await ctx.ui.select("Sage settings", [
+      `Enabled: ${onOff(draft.enabled)}`,
+      `Autonomous mode: ${onOff(draft.autonomousEnabled)}`,
+      `Model: ${draft.model}`,
+      `Reasoning level: ${draft.reasoningLevel}`,
+      `Timeout ms: ${draft.timeoutMs}`,
+      `Max calls/turn: ${draft.maxCallsPerTurn}`,
+      `Max calls/session: ${draft.maxCallsPerSession}`,
+      `Cooldown turns: ${draft.cooldownTurnsBetweenAutoCalls}`,
+      `Tool profile: ${draft.toolPolicy.profile}`,
+      `Custom tools: ${(draft.toolPolicy.customAllowedTools ?? []).join(",") || "(none)"}`,
+      `Max tool calls: ${draft.toolPolicy.maxToolCalls ?? 10}`,
+      `Max files read: ${draft.toolPolicy.maxFilesRead ?? 8}`,
+      `Max bytes/file: ${draft.toolPolicy.maxBytesPerFile ?? 200 * 1024}`,
+      `Max total bytes: ${draft.toolPolicy.maxTotalBytesRead ?? 1024 * 1024}`,
+      `Sensitive denylist: ${(draft.toolPolicy.sensitivePathDenylist ?? []).join(",")}`,
+      `Cost cap/session: ${draft.maxEstimatedCostPerSession ?? "(none)"}`,
+      `Save scope: ${scope}`,
+      "Test Sage call",
+      "Save and exit",
+      "Exit without saving"
+    ]);
+
+    if (!action || action === "Exit without saving") {
+      ctx.ui.notify("Sage settings unchanged", "info");
+      return;
+    }
+
+    if (action === "Save and exit") {
+      const target = getSettingsPathForScope(ctx.cwd, scope);
+      saveSettings(target, draft);
+      ctx.ui.notify(`Saved Sage settings to ${target}`, "info");
+      return;
+    }
+
+    if (action === "Test Sage call") {
+      await runSettingsTestCall(ctx, draft);
+      continue;
+    }
+
+    if (action.startsWith("Enabled:")) {
+      draft = { ...draft, enabled: !draft.enabled };
+      continue;
+    }
+    if (action.startsWith("Autonomous mode:")) {
+      draft = { ...draft, autonomousEnabled: !draft.autonomousEnabled };
+      continue;
+    }
+    if (action.startsWith("Model:")) {
+      const selected = await pickModel(ctx);
+      if (selected) draft = { ...draft, model: selected };
+      continue;
+    }
+    if (action.startsWith("Reasoning level:")) {
+      const selected = await ctx.ui.select("Reasoning level", [...REASONING_LEVELS]);
+      if (selected && REASONING_LEVELS.includes(selected as ReasoningLevel)) {
+        draft = { ...draft, reasoningLevel: selected as ReasoningLevel };
+      }
+      continue;
+    }
+    if (action.startsWith("Timeout ms:")) {
+      draft = await setNumberSetting(ctx, draft, "timeoutMs", "Timeout in milliseconds", 1000);
+      continue;
+    }
+    if (action.startsWith("Max calls/turn:")) {
+      draft = await setNumberSetting(ctx, draft, "maxCallsPerTurn", "Max Sage calls per turn", 1);
+      continue;
+    }
+    if (action.startsWith("Max calls/session:")) {
+      draft = await setNumberSetting(ctx, draft, "maxCallsPerSession", "Max Sage calls per session", 1);
+      continue;
+    }
+    if (action.startsWith("Cooldown turns:")) {
+      draft = await setNumberSetting(ctx, draft, "cooldownTurnsBetweenAutoCalls", "Cooldown turns", 0);
+      continue;
+    }
+    if (action.startsWith("Tool profile:")) {
+      const selected = await ctx.ui.select("Tool profile", [...TOOL_PROFILES]);
+      if (selected && TOOL_PROFILES.includes(selected as ToolProfile)) {
+        draft = { ...draft, toolPolicy: { ...draft.toolPolicy, profile: selected as ToolProfile } };
+      }
+      continue;
+    }
+    if (action.startsWith("Custom tools:")) {
+      const value = await ctx.ui.input("Comma-separated custom tools", "ls,glob,grep,read");
+      if (value !== undefined) {
+        const tools = value
+          .split(",")
+          .map((item) => item.trim())
+          .filter(Boolean);
+        draft = { ...draft, toolPolicy: { ...draft.toolPolicy, customAllowedTools: tools } };
+      }
+      continue;
+    }
+    if (action.startsWith("Max tool calls:")) {
+      draft = await setToolPolicyNumberSetting(ctx, draft, "maxToolCalls", "Max tool calls", 1);
+      continue;
+    }
+    if (action.startsWith("Max files read:")) {
+      draft = await setToolPolicyNumberSetting(ctx, draft, "maxFilesRead", "Max files read", 1);
+      continue;
+    }
+    if (action.startsWith("Max bytes/file:")) {
+      draft = await setToolPolicyNumberSetting(ctx, draft, "maxBytesPerFile", "Max bytes per file", 1024);
+      continue;
+    }
+    if (action.startsWith("Max total bytes:")) {
+      draft = await setToolPolicyNumberSetting(ctx, draft, "maxTotalBytesRead", "Max total bytes", 1024);
+      continue;
+    }
+    if (action.startsWith("Sensitive denylist:")) {
+      const value = await ctx.ui.input(
+        "Comma-separated sensitive path denylist",
+        (draft.toolPolicy.sensitivePathDenylist ?? []).join(",")
+      );
+      if (value !== undefined) {
+        const denylist = value
+          .split(",")
+          .map((item) => item.trim())
+          .filter(Boolean);
+        draft = { ...draft, toolPolicy: { ...draft.toolPolicy, sensitivePathDenylist: denylist } };
+      }
+      continue;
+    }
+    if (action.startsWith("Cost cap/session:")) {
+      const value = await ctx.ui.input("Cost cap per session (blank clears)", String(draft.maxEstimatedCostPerSession ?? ""));
+      if (value === undefined) continue;
+      const trimmed = value.trim();
+      if (trimmed.length === 0) {
+        const copy = { ...draft };
+        delete copy.maxEstimatedCostPerSession;
+        draft = copy;
+      } else {
+        const parsed = Number(trimmed);
+        if (Number.isFinite(parsed) && parsed >= 0) {
+          draft = { ...draft, maxEstimatedCostPerSession: parsed };
+        } else {
+          ctx.ui.notify("Invalid cost cap value", "warning");
+        }
+      }
+      continue;
+    }
+    if (action.startsWith("Save scope:")) {
+      const selected = await ctx.ui.select("Save scope", ["project", "global"]);
+      if (selected === "project" || selected === "global") scope = selected;
+      continue;
+    }
+  }
+}
+
+async function pickModel(ctx: ExtensionContext): Promise<string | undefined> {
+  const available = ctx.modelRegistry.getAvailable();
+  if (available.length === 0) {
+    ctx.ui.notify("No available models found", "warning");
+    return undefined;
+  }
+
+  const sorted = [...available].sort((a, b) => {
+    const left = `${a.provider}/${a.id}`.toLowerCase();
+    const right = `${b.provider}/${b.id}`.toLowerCase();
+    return left.localeCompare(right);
+  });
+
+  const options = sorted.slice(0, 200).map((model) => `${model.provider}/${model.id}${model.name ? ` — ${model.name}` : ""}`);
+  const selected = await ctx.ui.select("Choose Sage model", options);
+  if (!selected) return undefined;
+  return selected.split(" — ")[0];
+}
+
+function resolveModelSpec(
+  configuredModel: string,
+  availableModels: ModelLike[],
+  currentModel: ModelLike | undefined
+): { modelArg: string | undefined; reason: string } {
+  if (availableModels.length === 0) {
+    return { modelArg: undefined, reason: "no available models" };
+  }
+
+  const exactConfigured = availableModels.find((model) => `${model.provider}/${model.id}` === configuredModel);
+  if (exactConfigured) {
+    return { modelArg: `${exactConfigured.provider}/${exactConfigured.id}`, reason: "configured exact match" };
+  }
+
+  const fuzzyConfigured = configuredModel.trim().toLowerCase();
+  if (fuzzyConfigured) {
+    const fuzzyMatch = availableModels.find((model) => {
+      const composite = `${model.provider}/${model.id}`.toLowerCase();
+      const name = model.name?.toLowerCase() ?? "";
+      return composite.includes(fuzzyConfigured) || name.includes(fuzzyConfigured);
+    });
+    if (fuzzyMatch) {
+      return { modelArg: `${fuzzyMatch.provider}/${fuzzyMatch.id}`, reason: "configured fuzzy match" };
+    }
+  }
+
+  const preferredOpus = availableModels.find((model) => {
+    const id = model.id.toLowerCase();
+    const name = (model.name ?? "").toLowerCase();
+    return model.provider === "anthropic" && (id.includes("opus") || name.includes("opus"));
+  });
+
+  if (preferredOpus) {
+    return { modelArg: `${preferredOpus.provider}/${preferredOpus.id}`, reason: "anthropic opus fallback" };
+  }
+
+  if (currentModel) {
+    const present = availableModels.find(
+      (model) => model.provider === currentModel.provider && model.id === currentModel.id
+    );
+    if (present) {
+      return { modelArg: `${present.provider}/${present.id}`, reason: "current model fallback" };
+    }
+  }
+
+  const first = availableModels.at(0);
+  if (!first) return { modelArg: undefined, reason: "no available models" };
+  return { modelArg: `${first.provider}/${first.id}`, reason: "first available fallback" };
+}
+
+function onOff(value: boolean): string {
+  return value ? "on" : "off";
+}
+
+async function setNumberSetting(
+  ctx: ExtensionContext,
+  settings: SageSettings,
+  key: "timeoutMs" | "maxCallsPerTurn" | "maxCallsPerSession" | "cooldownTurnsBetweenAutoCalls",
+  title: string,
+  min: number
+): Promise<SageSettings> {
+  const current = String(settings[key]);
+  const value = await ctx.ui.input(title, current);
+  if (value === undefined) return settings;
+
+  const parsed = Number(value.trim());
+  if (Number.isFinite(parsed) === false) {
+    ctx.ui.notify("Invalid numeric value", "warning");
+    return settings;
+  }
+
+  const normalized = Math.max(min, Math.floor(parsed));
+  return { ...settings, [key]: normalized };
+}
+
+async function setToolPolicyNumberSetting(
+  ctx: ExtensionContext,
+  settings: SageSettings,
+  key: "maxToolCalls" | "maxFilesRead" | "maxBytesPerFile" | "maxTotalBytesRead",
+  title: string,
+  min: number
+): Promise<SageSettings> {
+  const current = String(settings.toolPolicy[key] ?? "");
+  const value = await ctx.ui.input(title, current);
+  if (value === undefined) return settings;
+
+  const parsed = Number(value.trim());
+  if (Number.isFinite(parsed) === false) {
+    ctx.ui.notify("Invalid numeric value", "warning");
+    return settings;
+  }
+
+  const normalized = Math.max(min, Math.floor(parsed));
+  return {
+    ...settings,
+    toolPolicy: {
+      ...settings.toolPolicy,
+      [key]: normalized
+    }
+  };
+}
+
+async function runSettingsTestCall(ctx: ExtensionContext, settings: SageSettings): Promise<void> {
+  const modelSelection = resolveModelSpec(settings.model, ctx.modelRegistry.getAvailable(), ctx.model);
+  const resolvedModel = modelSelection.modelArg;
+  if (!resolvedModel) {
+    ctx.ui.notify("Cannot test Sage: no available model", "warning");
+    return;
+  }
+
+  try {
+    const result = await runSageSingleShot({
+      cwd: ctx.cwd,
+      model: resolvedModel,
+      reasoningLevel: settings.reasoningLevel,
+      timeoutMs: Math.min(settings.timeoutMs, 45_000),
+      question: "Reply with 'Sage test OK' and one short sentence.",
+      toolPolicy: settings.toolPolicy
+    });
+
+    ctx.ui.notify(`Sage test succeeded (${result.latencyMs}ms)`, "info");
+  } catch (error) {
+    const message = error instanceof Error ? error.message : "Unknown test failure";
+    ctx.ui.notify(`Sage test failed: ${message}`, "error");
+  }
+}