@czottmann/pi-automode 1.6.0 → 1.8.0

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,12 +6,18 @@ 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 { parseModelSpec } from "./model.ts";
-import { buildTranscript } from "./transcript.ts";
+import {
+  CLASSIFIER_DETAILED_INSTRUCTION,
+  CLASSIFIER_FAST_INSTRUCTION,
+  CLASSIFIER_SYSTEM_PROMPT,
+} from "./constants.ts";
+import { formatModelSpec, parseModelSpec } from "./model.ts";
+import { buildClassifierTranscript } from "./transcript.ts";
 import type {
   ClassificationDecision,
   ClassifyAction,
+  ClassifierIoAttempt,
+  ClassifyResult,
   EffectiveConfig,
 } from "./types.ts";
 
@@ -63,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>;
 
@@ -71,53 +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;
 };
 
-/** Parse the classifier's JSON-only response. Invalid output is handled fail-closed by the caller. */
-export function parseClassifierDecision(
-  message: AssistantMessage,
-): ClassificationDecision | undefined {
+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, 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();
-  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.
+    .join("\n");
+  return trim ? text.trim() : text;
+}
+
+/** 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 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;
   }
-  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;
+  }
+  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,
@@ -132,10 +232,13 @@ 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.";
   for (let attempt = 0; attempt < maxAttempts; attempt += 1) {
+    const started = Date.now();
     let response: AssistantMessage;
     try {
       response = await completeFn(
@@ -146,19 +249,34 @@ 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,
+      });
       return {
         decision: "block",
         tier: "none",
-        reason: `Classifier failed; auto mode fails closed: ${
-          error instanceof Error ? error.message : String(error)
-        }`,
+        reason: `Classifier failed; auto mode fails closed: ${message}`,
       };
     }
-    const decision = parseClassifierDecision(response);
+    const durationMs = Date.now() - started;
+    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"
@@ -168,12 +286,118 @@ 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);
+  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 exactly 0 or 1; 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,
   action,
   loadedContext,
-): Promise<ClassificationDecision> => {
+): Promise<ClassifyResult> => {
   const classifier = await resolveClassifier(ctx, config);
   if (!classifier) {
     return {
@@ -183,25 +407,47 @@ export const defaultClassifyAction: ClassifyAction = async (
     };
   }
 
-  const userMessage: UserMessage = {
+  const systemPrompt = buildClassifierPrompt(config);
+  const transcript = buildClassifierTranscript(ctx, {
+    maxUserTokens: config.maxUserTranscriptTokens,
+    maxToolTokens: config.maxToolTranscriptTokens,
+  });
+  const contextText = `<loaded-project-instructions>\n${
+    loadedContext || "(none)"
+  }\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: `<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}`,
-      },
-    ],
+    content: [{ type: "text", text: contextText }],
     timestamp: Date.now(),
   };
 
-  return classifyWithRetry(
+  const attempts: ClassifierIoAttempt[] = [];
+  const started = Date.now();
+  const decision = await classifyInStages(
     complete,
     classifier,
-    { systemPrompt: buildClassifierPrompt(config), messages: [userMessage] },
+    { systemPrompt, contextMessage },
     ctx.signal,
+    {
+      sessionId: classifierCacheSessionId(ctx),
+      onAttempt: (attempt) => attempts.push(attempt),
+    },
   );
+
+  return {
+    ...decision,
+    io: {
+      model: formatModelSpec(classifier.model),
+      prompt: {
+        system: systemPrompt,
+        context: contextText,
+        fastInstruction: CLASSIFIER_FAST_INSTRUCTION,
+        detailedInstruction: CLASSIFIER_DETAILED_INSTRUCTION,
+      },
+      attempts,
+      durationMs: Date.now() - started,
+    },
+  };
 };

package/extensions/auto-mode/config.ts

@@ -4,7 +4,9 @@ import {
   DEFAULT_ALLOW,
   DEFAULT_ENVIRONMENT,
   DEFAULT_HARD_DENY,
-  DEFAULT_MAX_TRANSCRIPT_LINES,
+  DEFAULT_LOG_CONFIG,
+  DEFAULT_MAX_TOOL_TRANSCRIPT_TOKENS,
+  DEFAULT_MAX_USER_TRANSCRIPT_TOKENS,
   DEFAULT_PROTECTED_PATHS,
   DEFAULT_SOFT_DENY,
   PI_GLOBAL_SETTINGS,
@@ -17,6 +19,7 @@ import type {
   ConfigLoadResult,
   EffectiveConfig,
   LoadedSettingsFile,
+  LogConfig,
   SettingsFile,
   SettingsSources,
   ToolPattern,
@@ -94,7 +97,8 @@ export function validateSettingsFile(
       const knownAutoMode = new Set([
         "enabled",
         "classifierModel",
-        "maxTranscriptLines",
+        "maxUserTranscriptTokens",
+        "maxToolTranscriptTokens",
         "environment",
         "allow",
         "protectedPaths",
@@ -102,6 +106,7 @@ export function validateSettingsFile(
         "softDeny",
         "hard_deny",
         "hardDeny",
+        "log",
       ]);
       for (const key of Object.keys(autoMode)) {
         if (!knownAutoMode.has(key)) {
@@ -121,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,
@@ -160,6 +171,9 @@ export function validateSettingsFile(
         "autoMode.hard_deny",
         diagnostics,
       );
+      if (hasOwn(autoMode, "log")) {
+        validateLogSetting(autoMode.log, source, diagnostics);
+      }
     }
   }
 
@@ -228,6 +242,41 @@ function finalizeRuleSetting(accumulator: RuleAccumulator): string[] {
   return [...new Set([...base, ...accumulator.entries])];
 }
 
+function validateLogSetting(
+  value: unknown,
+  source: string,
+  diagnostics: string[],
+): void {
+  if (!value || typeof value !== "object" || Array.isArray(value)) {
+    diagnostics.push(`${source}: autoMode.log must be an object`);
+    return;
+  }
+  const log = value as Record<string, unknown>;
+  if (hasOwn(log, "enabled") && typeof log.enabled !== "boolean") {
+    diagnostics.push(`${source}: autoMode.log.enabled must be a boolean`);
+  }
+  if (
+    hasOwn(log, "classifierIo") && typeof log.classifierIo !== "boolean"
+  ) {
+    diagnostics.push(`${source}: autoMode.log.classifierIo must be a boolean`);
+  }
+}
+
+function mergeLog(
+  base: LogConfig,
+  patch: Partial<LogConfig> | undefined,
+): LogConfig {
+  if (!patch) return base;
+  return {
+    enabled: patch.enabled ?? base.enabled,
+    classifierIo: patch.classifierIo ?? base.classifierIo,
+  };
+}
+
+function validTranscriptBudget(value: unknown): value is number {
+  return Number.isInteger(value) && Number(value) >= 32;
+}
+
 function applyAutoModeScalars(
   base: EffectiveConfig,
   settings: AutoModeSettings | undefined,
@@ -237,7 +286,17 @@ 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),
   };
 }
 
@@ -268,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],
@@ -276,6 +336,7 @@ export function buildEffectiveConfigFromSources(
     hardDeny: [...DEFAULT_HARD_DENY],
     permissionDeny: [],
     permissionAsk: [],
+    log: { ...DEFAULT_LOG_CONFIG },
   };
 
   const globalSettings = sources.globalSettings ?? [];

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")];
@@ -166,3 +171,9 @@ export const PROFILE_FILES = new Set([
 ]);
 
 export const READ_ONLY_TOOLS = new Set(["read", "grep", "find", "ls"]);
+
+/** Default observability log config: off, classifier I/O off. */
+export const DEFAULT_LOG_CONFIG = {
+  enabled: false,
+  classifierIo: false,
+};