@czottmann/pi-automode 1.7.0 → 1.8.1

package/extensions/auto-mode/classifier.ts

@@ -1,3 +1,4 @@
+import { createHash } from "node:crypto";
 import { complete } from "@earendil-works/pi-ai";
 import type {
   AssistantMessage,
@@ -5,9 +6,13 @@ import type {
   UserMessage,
 } from "@earendil-works/pi-ai";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
-import { CLASSIFIER_SYSTEM_PROMPT } from "./constants.ts";
+import {
+  CLASSIFIER_DETAILED_INSTRUCTION,
+  CLASSIFIER_FAST_INSTRUCTION,
+  CLASSIFIER_SYSTEM_PROMPT,
+} from "./constants.ts";
 import { formatModelSpec, parseModelSpec } from "./model.ts";
-import { buildTranscript } from "./transcript.ts";
+import { buildClassifierTranscript } from "./transcript.ts";
 import type {
   ClassificationDecision,
   ClassifyAction,
@@ -65,7 +70,9 @@ export type ClassifierCompletionFn = (
     headers?: Record<string, string>;
     signal?: AbortSignal;
     maxTokens: number;
-    temperature: number;
+    temperature?: number;
+    sessionId?: string;
+    cacheRetention?: "none" | "short" | "long";
   },
 ) => Promise<AssistantMessage>;
 
@@ -73,60 +80,144 @@ export type RetryOptions = {
   maxAttempts?: number;
   maxTokens?: number;
   temperature?: number;
+  sessionId?: string;
+  cacheRetention?: "none" | "short" | "long";
+  stage?: "fast" | "detailed";
   /** Receives each attempt's raw response (or error) and parsed decision, for observability logging. */
   onAttempt?: (attempt: ClassifierIoAttempt) => void;
 };
 
+const FAST_CLASSIFIER_MAX_TOKENS = 4;
+
+export type StagedClassifierOptions = {
+  sessionId: string;
+  onAttempt?: (attempt: ClassifierIoAttempt) => void;
+};
+
 /** Concatenate all text blocks of an assistant message into a single string. */
-function extractAssistantText(message: AssistantMessage): string {
-  return message.content
+function extractAssistantText(message: AssistantMessage, trim = true): string {
+  const text = message.content
     .filter(
       (block): block is { type: "text"; text: string } => block.type === "text",
     )
     .map((block) => block.text)
-    .join("\n")
-    .trim();
+    .join("\n");
+  return trim ? text.trim() : text;
 }
 
-/** Parse the classifier's JSON-only response. Invalid output is handled fail-closed by the caller. */
+/** Parse the exact detailed-stage JSON contract; any wrapper or shape drift fails closed. */
 export function parseClassifierDecision(
   message: AssistantMessage,
 ): ClassificationDecision | undefined {
   const text = extractAssistantText(message);
-  const fenced = text.match(/```(?:json)?\s*([\s\S]*?)```/i)?.[1];
-  const candidates = [fenced, text, text.match(/\{[\s\S]*\}/)?.[0]].filter(
-    Boolean,
-  ) as string[];
-  for (const candidate of candidates) {
-    try {
-      const parsed = JSON.parse(candidate) as Partial<ClassificationDecision>;
-      if (
-        (parsed.decision === "allow" || parsed.decision === "block") &&
-        typeof parsed.reason === "string"
-      ) {
-        return {
-          decision: parsed.decision,
-          tier: parsed.tier ?? "none",
-          reason: parsed.reason,
-        };
-      }
-    } catch {
-      // Try next candidate.
+  const validTiers = new Set<ClassificationDecision["tier"]>([
+    "hard_deny",
+    "soft_deny",
+    "allow",
+    "explicit_intent",
+    "none",
+  ]);
+  try {
+    for (const key of ["decision", "tier", "reason"]) {
+      const occurrences = text.match(new RegExp(`"${key}"\\s*:`, "g"))?.length ?? 0;
+      if (occurrences !== 1) return undefined;
+    }
+    const parsed = JSON.parse(text) as Record<string, unknown>;
+    if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
+      return undefined;
+    }
+    const keys = Object.keys(parsed).sort();
+    if (keys.join(",") !== "decision,reason,tier") return undefined;
+    if (parsed.decision !== "allow" && parsed.decision !== "block") {
+      return undefined;
+    }
+    if (!validTiers.has(parsed.tier as ClassificationDecision["tier"])) {
+      return undefined;
+    }
+    const tier = parsed.tier as ClassificationDecision["tier"];
+    if (
+      (parsed.decision === "allow" &&
+        !["allow", "explicit_intent", "none"].includes(tier)) ||
+      (parsed.decision === "block" &&
+        !["hard_deny", "soft_deny", "none"].includes(tier))
+    ) {
+      return undefined;
     }
+    if (typeof parsed.reason !== "string" || parsed.reason.trim() === "") {
+      return undefined;
+    }
+    return {
+      decision: parsed.decision,
+      tier,
+      reason: parsed.reason,
+    };
+  } catch {
+    return undefined;
+  }
+}
+
+function stageMessage(text: string): UserMessage {
+  return {
+    role: "user",
+    content: [{ type: "text", text }],
+    timestamp: Date.now(),
+  };
+}
+
+function responseAttempt(
+  stage: "fast" | "detailed",
+  attempt: number,
+  response: AssistantMessage,
+  durationMs: number,
+  parsed?: ClassificationDecision,
+  trimText = true,
+): ClassifierIoAttempt {
+  return {
+    stage,
+    attempt,
+    response: {
+      stopReason: response.stopReason,
+      text: extractAssistantText(response, trimText),
+      model: response.model,
+      timestamp: response.timestamp,
+      usage: response.usage,
+      ...(response.errorMessage === undefined
+        ? {}
+        : { errorMessage: response.errorMessage }),
+    },
+    parsed,
+    durationMs,
+  };
+}
+
+function classifierFailure(
+  response: AssistantMessage,
+  label: "Classifier" | "Fast classifier",
+  retryLength = false,
+): ClassificationDecision | undefined {
+  if (
+    response.stopReason === "stop" ||
+    (retryLength && response.stopReason === "length")
+  ) {
+    return undefined;
   }
-  return undefined;
+  const fallback = response.stopReason === "aborted"
+    ? "Classifier model request was aborted."
+    : response.stopReason === "error"
+    ? "Classifier model returned an error response."
+    : `${label} response did not stop cleanly (${response.stopReason}).`;
+  return {
+    decision: "block",
+    tier: "none",
+    reason: `${label} failed; auto mode fails closed: ${
+      response.errorMessage || fallback
+    }`,
+  };
 }
 
 /**
- * Call the classifier model and parse its decision, retrying when the model
- * returns malformed or truncated output. Small classifier models occasionally
- * ramble into the token cap (stopReason "length") or emit prose instead of
- * JSON; a single retry recovers most of these transient failures.
- *
- * Safety is preserved: an "allow" is only returned when a response actually
- * parses to a valid allow decision. If every attempt fails to parse, the
- * function fails closed with a block decision. Thrown errors (network/auth)
- * are not retried — they fail closed immediately, matching prior behavior.
+ * Call the detailed classifier and parse its decision, retrying malformed or
+ * truncated output. Provider errors and exhausted retries fail closed.
  */
 export async function classifyWithRetry(
   completeFn: ClassifierCompletionFn,
@@ -141,7 +232,8 @@ export async function classifyWithRetry(
 ): Promise<ClassificationDecision> {
   const maxAttempts = options.maxAttempts ?? 2;
   const maxTokens = options.maxTokens ?? 1200;
-  const temperature = options.temperature ?? 0;
+  const temperature = options.temperature;
+  const stage = options.stage ?? "detailed";
   const onAttempt = options.onAttempt;
   let lastReason =
     "Classifier response was not valid decision JSON; auto mode fails closed.";
@@ -157,12 +249,15 @@ export async function classifyWithRetry(
           headers: classifier.headers,
           signal,
           maxTokens,
-          temperature,
+          ...(temperature === undefined ? {} : { temperature }),
+          sessionId: options.sessionId,
+          cacheRetention: options.cacheRetention,
         },
       );
     } catch (error) {
       const message = error instanceof Error ? error.message : String(error);
       onAttempt?.({
+        stage,
         attempt: attempt + 1,
         error: message,
         durationMs: Date.now() - started,
@@ -174,16 +269,14 @@ export async function classifyWithRetry(
       };
     }
     const durationMs = Date.now() - started;
-    const decision = parseClassifierDecision(response);
-    onAttempt?.({
-      attempt: attempt + 1,
-      response: {
-        stopReason: response.stopReason,
-        text: extractAssistantText(response),
-      },
-      parsed: decision,
-      durationMs,
-    });
+    const failure = classifierFailure(response, "Classifier", true);
+    const decision = response.stopReason === "stop"
+      ? parseClassifierDecision(response)
+      : undefined;
+    onAttempt?.(
+      responseAttempt(stage, attempt + 1, response, durationMs, decision, false),
+    );
+    if (failure) return failure;
     if (decision) return decision;
     lastReason =
       response.stopReason === "length"
@@ -193,6 +286,112 @@ export async function classifyWithRetry(
   return { decision: "block", tier: "none", reason: lastReason };
 }
 
+/** Run the one-token conservative gate, then detailed review only when requested. */
+export async function classifyInStages(
+  completeFn: ClassifierCompletionFn,
+  classifier: {
+    model: Model<any>;
+    apiKey?: string;
+    headers?: Record<string, string>;
+  },
+  prompt: { systemPrompt: string; contextMessage: UserMessage },
+  signal: AbortSignal | undefined,
+  options: StagedClassifierOptions,
+): Promise<ClassificationDecision> {
+  const fastStarted = Date.now();
+  let fastResponse: AssistantMessage;
+  try {
+    fastResponse = await completeFn(
+      classifier.model,
+      {
+        systemPrompt: prompt.systemPrompt,
+        messages: [
+          prompt.contextMessage,
+          stageMessage(CLASSIFIER_FAST_INSTRUCTION),
+        ],
+      },
+      {
+        apiKey: classifier.apiKey,
+        headers: classifier.headers,
+        signal,
+        // Some OpenAI-compatible servers count an initial control token and
+        // EOS against max_tokens. Four tokens reliably permit one visible digit.
+        maxTokens: FAST_CLASSIFIER_MAX_TOKENS,
+        sessionId: options.sessionId,
+        cacheRetention: "short",
+      },
+    );
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    options.onAttempt?.({
+      stage: "fast",
+      attempt: 1,
+      error: message,
+      durationMs: Date.now() - fastStarted,
+    });
+    return {
+      decision: "block",
+      tier: "none",
+      reason: `Fast classifier failed; auto mode fails closed: ${message}`,
+    };
+  }
+
+  const fastText = extractAssistantText(fastResponse, false).trim();
+  const failure = classifierFailure(fastResponse, "Fast classifier");
+  options.onAttempt?.(
+    responseAttempt(
+      "fast",
+      1,
+      fastResponse,
+      Date.now() - fastStarted,
+      undefined,
+      false,
+    ),
+  );
+  if (failure) return failure;
+  if (fastText === "0") {
+    return {
+      decision: "allow",
+      tier: "none",
+      reason: "Fast classifier found no policy-relevant risk.",
+    };
+  }
+  if (fastText !== "1") {
+    return {
+      decision: "block",
+      tier: "none",
+      reason:
+        "Fast classifier response was not 0 or 1 after trimming whitespace; auto mode fails closed.",
+    };
+  }
+
+  return classifyWithRetry(
+    completeFn,
+    classifier,
+    {
+      systemPrompt: prompt.systemPrompt,
+      messages: [
+        prompt.contextMessage,
+        stageMessage(CLASSIFIER_DETAILED_INSTRUCTION),
+      ],
+    },
+    signal,
+    {
+      stage: "detailed",
+      sessionId: options.sessionId,
+      cacheRetention: "short",
+      onAttempt: options.onAttempt,
+    },
+  );
+}
+
+export function classifierCacheSessionId(ctx: ExtensionContext): string {
+  const source = ctx.sessionManager.getSessionId?.() ??
+    ctx.sessionManager.getSessionFile?.() ?? ctx.cwd;
+  const digest = createHash("sha256").update(source).digest("hex").slice(0, 32);
+  return `pi-automode-${digest}`;
+}
+
 export const defaultClassifyAction: ClassifyAction = async (
   ctx,
   config,
@@ -209,32 +408,44 @@ export const defaultClassifyAction: ClassifyAction = async (
   }
 
   const systemPrompt = buildClassifierPrompt(config);
-  const userText = `<loaded-project-instructions>\n${
+  const transcript = buildClassifierTranscript(ctx, {
+    maxUserTokens: config.maxUserTranscriptTokens,
+    maxToolTokens: config.maxToolTranscriptTokens,
+  });
+  const contextText = `<loaded-project-instructions>\n${
     loadedContext || "(none)"
-  }\n</loaded-project-instructions>\n\n<transcript>\n${
-    buildTranscript(ctx, config.maxTranscriptLines) || "(none)"
-  }\n</transcript>\n\nLatest action to classify:\n${action}`;
-  const userMessage: UserMessage = {
+  }\n</loaded-project-instructions>\n\n<classifier-transcript>\n${
+    transcript || "(none)"
+  }\n</classifier-transcript>\n\nLatest action to classify:\n${action}`;
+  const contextMessage: UserMessage = {
     role: "user",
-    content: [{ type: "text", text: userText }],
+    content: [{ type: "text", text: contextText }],
     timestamp: Date.now(),
   };
 
   const attempts: ClassifierIoAttempt[] = [];
   const started = Date.now();
-  const decision = await classifyWithRetry(
+  const decision = await classifyInStages(
     complete,
     classifier,
-    { systemPrompt, messages: [userMessage] },
+    { systemPrompt, contextMessage },
     ctx.signal,
-    { onAttempt: (attempt) => attempts.push(attempt) },
+    {
+      sessionId: classifierCacheSessionId(ctx),
+      onAttempt: (attempt) => attempts.push(attempt),
+    },
   );
 
   return {
     ...decision,
     io: {
       model: formatModelSpec(classifier.model),
-      prompt: { system: systemPrompt, user: userText },
+      prompt: {
+        system: systemPrompt,
+        context: contextText,
+        fastInstruction: CLASSIFIER_FAST_INSTRUCTION,
+        detailedInstruction: CLASSIFIER_DETAILED_INSTRUCTION,
+      },
       attempts,
       durationMs: Date.now() - started,
     },

package/extensions/auto-mode/config.ts

@@ -5,7 +5,8 @@ import {
   DEFAULT_ENVIRONMENT,
   DEFAULT_HARD_DENY,
   DEFAULT_LOG_CONFIG,
-  DEFAULT_MAX_TRANSCRIPT_LINES,
+  DEFAULT_MAX_TOOL_TRANSCRIPT_TOKENS,
+  DEFAULT_MAX_USER_TRANSCRIPT_TOKENS,
   DEFAULT_PROTECTED_PATHS,
   DEFAULT_SOFT_DENY,
   PI_GLOBAL_SETTINGS,
@@ -96,7 +97,8 @@ export function validateSettingsFile(
       const knownAutoMode = new Set([
         "enabled",
         "classifierModel",
-        "maxTranscriptLines",
+        "maxUserTranscriptTokens",
+        "maxToolTranscriptTokens",
         "environment",
         "allow",
         "protectedPaths",
@@ -124,14 +126,20 @@ export function validateSettingsFile(
           `${source}: autoMode.classifierModel must be a provider/model string`,
         );
       }
-      if (
-        hasOwn(autoMode, "maxTranscriptLines") &&
-        (!Number.isInteger(autoMode.maxTranscriptLines) ||
-          Number(autoMode.maxTranscriptLines) <= 0)
+      for (
+        const key of [
+          "maxUserTranscriptTokens",
+          "maxToolTranscriptTokens",
+        ] as const
       ) {
-        diagnostics.push(
-          `${source}: autoMode.maxTranscriptLines must be a positive integer`,
-        );
+        if (
+          hasOwn(autoMode, key) &&
+          (!Number.isInteger(autoMode[key]) || Number(autoMode[key]) < 32)
+        ) {
+          diagnostics.push(
+            `${source}: autoMode.${key} must be an integer of at least 32`,
+          );
+        }
       }
       validateStringArraySetting(
         autoMode.environment,
@@ -265,6 +273,10 @@ function mergeLog(
   };
 }
 
+function validTranscriptBudget(value: unknown): value is number {
+  return Number.isInteger(value) && Number(value) >= 32;
+}
+
 function applyAutoModeScalars(
   base: EffectiveConfig,
   settings: AutoModeSettings | undefined,
@@ -274,7 +286,16 @@ function applyAutoModeScalars(
     ...base,
     enabled: settings.enabled ?? base.enabled,
     classifierModel: settings.classifierModel ?? base.classifierModel,
-    maxTranscriptLines: settings.maxTranscriptLines ?? base.maxTranscriptLines,
+    maxUserTranscriptTokens: validTranscriptBudget(
+        settings.maxUserTranscriptTokens,
+      )
+      ? settings.maxUserTranscriptTokens
+      : base.maxUserTranscriptTokens,
+    maxToolTranscriptTokens: validTranscriptBudget(
+        settings.maxToolTranscriptTokens,
+      )
+      ? settings.maxToolTranscriptTokens
+      : base.maxToolTranscriptTokens,
     log: mergeLog(base.log, settings.log),
   };
 }
@@ -306,7 +327,8 @@ export function buildEffectiveConfigFromSources(
 ): EffectiveConfig {
   let config: EffectiveConfig = {
     enabled: true,
-    maxTranscriptLines: DEFAULT_MAX_TRANSCRIPT_LINES,
+    maxUserTranscriptTokens: DEFAULT_MAX_USER_TRANSCRIPT_TOKENS,
+    maxToolTranscriptTokens: DEFAULT_MAX_TOOL_TRANSCRIPT_TOKENS,
     environment: [...DEFAULT_ENVIRONMENT],
     allow: [...DEFAULT_ALLOW],
     protectedPaths: [...DEFAULT_PROTECTED_PATHS],

package/extensions/auto-mode/constants.ts

@@ -55,7 +55,8 @@ export const DEFAULT_PROTECTED_PATHS = [
   ".mcp.json",
 ];
 
-export const DEFAULT_MAX_TRANSCRIPT_LINES = 80;
+export const DEFAULT_MAX_USER_TRANSCRIPT_TOKENS = 4000;
+export const DEFAULT_MAX_TOOL_TRANSCRIPT_TOKENS = 4000;
 export const DENIAL_HISTORY_LIMIT = 12;
 
 /** Built-in trusted environment. Users extend this with `$defaults`. */
@@ -115,7 +116,7 @@ Auto mode is active. Routine permission prompts are replaced by a pre-execution
 export const CLASSIFIER_SYSTEM_PROMPT =
   `You are an auto-mode security classifier for an autonomous coding agent.
 
-You evaluate the agent's latest tool call before it runs. Return only JSON.
+You evaluate the agent's latest tool call before it runs. Follow the final output instruction exactly.
 
 Core semantics:
 1. HARD_DENY rules block unconditionally. User intent and allow exceptions do not override them.
@@ -144,9 +145,13 @@ Examples:
 - cat ~/.ssh/id_rsa or echo $API_TOKEN => block, hard_deny, credential exposure.
 - npm test, swift test, xcodebuild test, or local formatting/linting => allow unless the action includes a blocked side effect.
 - Uploading repository contents, logs with secrets, or internal data to an untrusted domain => block, hard_deny.
-- Modifying .pi/automode*, this extension, or permission rules => block, hard_deny.
+- Modifying .pi/automode*, this extension, or permission rules => block, hard_deny.`;
 
-Respond with JSON exactly matching:
+export const CLASSIFIER_FAST_INSTRUCTION =
+  `Return exactly one digit and nothing else: 0 only when the action is clearly allowed by the policy; 1 when it may need blocking or you are uncertain. Err on 1.`;
+
+export const CLASSIFIER_DETAILED_INSTRUCTION =
+  `Return only JSON exactly matching:
 {"decision":"allow"|"block","tier":"hard_deny"|"soft_deny"|"allow"|"explicit_intent"|"none","reason":"brief concrete reason"}`;
 
 export const PI_GLOBAL_SETTINGS = [resolve(HOME, ".pi/agent/automode.json")];

package/extensions/auto-mode/extension.ts

@@ -28,7 +28,6 @@ import {
 import { formatModelSpec, parseModelSpec } from "./model.ts";
 import { promptForClassifierModel } from "./model-selector.ts";
 import { matchesToolPattern } from "./permissions.ts";
-import { isProtectedPath, resolveInputPath } from "./paths.ts";
 import {
   actionSummary,
   formatDenials,
@@ -64,9 +63,25 @@ type LogCtx = {
   classifierModel?: string;
 };
 
-/** Append a classifier I/O entry when classifier logging is enabled. */
+/** Append ccusage-compatible usage and optional classifier I/O entries. */
 function logClassifierIo(decision: ClassifyResult, log: LogCtx): void {
-  if (!log.logger.enabled || !log.logger.classifierIo || !decision.io) return;
+  if (!log.logger.enabled || !decision.io) return;
+
+  for (const attempt of decision.io.attempts) {
+    const response = attempt.response;
+    if (!response) continue;
+    log.logger.append({
+      type: "message",
+      timestamp: new Date(response.timestamp).toISOString(),
+      message: {
+        role: "assistant",
+        model: response.model,
+        usage: response.usage,
+      },
+    });
+  }
+
+  if (!log.logger.classifierIo) return;
   log.logger.append({
     type: "classifier",
     ts: new Date().toISOString(),
@@ -214,10 +229,10 @@ export function createPiAutomode(options: PiAutomodeOptions = {}) {
     });
 
     pi.on("tool_call", async (event, ctx) => {
-      // Enforcement order mirrors Claude Code's documented model:
+      // Enforcement order:
       // 1. permission deny/ask rules,
       // 2. deterministic hard-deny checks that never consult the model,
-      // 3. read-only fast path,
+      // 3. read-only built-in fast path,
       // 4. classifier for every remaining action, fail-closed on setup/parse errors.
       const cfg = effectiveConfig();
       if (!cfg.enabled) return undefined;
@@ -306,34 +321,6 @@ export function createPiAutomode(options: PiAutomodeOptions = {}) {
         );
       }
 
-      // Protected paths go to the classifier regardless of allow rules.
-      if (event.toolName === "write" || event.toolName === "edit") {
-        const path = resolveInputPath(ctx.cwd, input.path);
-        if (path && isProtectedPath(path, ctx.cwd, cfg.protectedPaths)) {
-          const decision = await classify(ctx, cfg, summary, loadedContext);
-          logClassifierIo(decision, logCtx);
-          if (decision.decision === "allow") {
-            state.classifierAllowed += 1;
-            return allow(
-              ctx,
-              "classifier",
-              decision.reason,
-              event.toolName,
-              summary,
-              logCtx,
-            );
-          }
-          state.classifierDenied += 1;
-          return block(ctx, {
-            timestamp: Date.now(),
-            toolName: event.toolName,
-            reason: decision.reason,
-            action: summary,
-            kind: "classifier",
-          }, logCtx);
-        }
-      }
-
       const decision = await classify(ctx, cfg, summary, loadedContext);
       logClassifierIo(decision, logCtx);
       if (decision.decision === "allow") {

package/extensions/auto-mode/hard-deny.ts

@@ -4,6 +4,7 @@ import {
   isProfileOrAuthorizedKeysPath,
   isSafetyControlPath,
   resolveInputPath,
+  resolvePathForPolicy,
   shellPathTokenToPath,
 } from "./paths.ts";
 
@@ -331,9 +332,11 @@ export function deterministicHardDeny(
   if (toolName === "write" || toolName === "edit") {
     const path = resolveInputPath(cwd, input.path);
     if (!path) return undefined;
-    const profileReason = isProfileOrAuthorizedKeysPath(path);
+    const policyPath = resolvePathForPolicy(path) ?? path;
+    const policyCwd = resolvePathForPolicy(cwd) ?? cwd;
+    const profileReason = isProfileOrAuthorizedKeysPath(policyPath);
     if (profileReason) return profileReason;
-    if (isSafetyControlPath(path, cwd)) {
+    if (isSafetyControlPath(policyPath, policyCwd)) {
       return "auto-mode or permission safety-control modification is hard-denied";
     }
   }