@zhushanwen/pi-structured-output 0.2.0 → 0.2.2

package/package.json

@@ -1,12 +1,12 @@
 {
   "name": "@zhushanwen/pi-structured-output",
-  "version": "0.2.0",
+  "version": "0.2.2",
   "description": "Structured output tool for Pi — enforces JSON Schema via tool call mechanism with Ajv validation",
   "type": "module",
   "main": "index.ts",
   "pi": {
     "extensions": [
-      "./src/index.ts"
+      "./index.ts"
     ]
   },
   "keywords": [

package/src/index.ts

@@ -7,6 +7,15 @@
  *
  * Design: FR-1 to FR-5 from spec, FR-4 dual-layer enforcement.
  * Reference: Claude Code's SyntheticOutputTool (Ajv + Stop hook).
+ *
+ * Key design decisions (borrowed from Claude Code):
+ * - Schema is injected into BOTH system prompt AND tool description, so LLM
+ *   knows the exact output structure regardless of which signal it reads.
+ * - Enforcement checks "last call succeeded" (not "was called"), so validation
+ *   failures trigger retries rather than being silently skipped.
+ * - Retry cap prevents infinite enforcement loops on persistent schema mismatch.
+ * - WeakMap caches Ajv compile results for repeated calls with the same schema
+ *   object reference (mirrors Claude Code's toolCache pattern).
  */
 
 import Ajv, { type ValidateFunction } from "ajv";
@@ -18,14 +27,27 @@ type PiAPI = any;
 
 const ENV_KEY = "STRUCTURED_OUTPUT_SCHEMA";
 const TOOL_NAME = "structured-output";
-
-const SYSTEM_PROMPT =
-  "你必须在完成分析后调用 structured-output tool 来返回结构化结果。" +
-  "不要在文本回复中输出 JSON，直接调用 structured-output tool。" +
-  "这是你返回最终结果的唯一方式。";
+const MAX_RETRIES = parseInt(process.env.MAX_STRUCTURED_OUTPUT_RETRIES || "5", 10);
 
 const ENFORCEMENT_MESSAGE = "你必须调用 structured-output tool 来返回结果。";
 
+// ── Ajv WeakMap cache ─────────────────────────────────────────
+// Workflow scripts may call agent({schema}) 30-80 times per run.
+// Without caching, each call does new Ajv() + validateSchema() + compile().
+// WeakMap keyed by schema object reference brings this to near-zero overhead.
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+const ajvCache = new WeakMap<object, ValidateFunction>();
+
+function getOrCompileValidator(schema: Record<string, unknown>): ValidateFunction {
+  const cached = ajvCache.get(schema);
+  if (cached) return cached;
+
+  const ajv = new Ajv({ strict: false });
+  const validate = ajv.compile(schema);
+  ajvCache.set(schema, validate);
+  return validate;
+}
+
 export default function structuredOutputExtension(pi: PiAPI): void {
   const schemaStr = process.env[ENV_KEY];
   if (!schemaStr) return;
@@ -39,22 +61,37 @@ export default function structuredOutputExtension(pi: PiAPI): void {
     return;
   }
 
-  // Compile with Ajv
-  const ajv = new Ajv({ strict: false });
+  // Compile with Ajv (cached)
   let validate: ValidateFunction;
   try {
-    validate = ajv.compile(schema);
+    validate = getOrCompileValidator(schema);
   } catch (e) {
     console.error(`[${TOOL_NAME}] Invalid JSON Schema:`, (e as Error).message);
     return;
   }
 
-  // Register tool — passthrough parameters (schema is dynamic per session)
+  // Build prompts with schema embedded
+  const schemaJsonStr = JSON.stringify(schema, null, 2);
+  const systemPrompt =
+    "你必须在完成分析后调用 structured-output tool 来返回结构化结果。" +
+    "不要在文本回复中输出 JSON，直接调用 structured-output tool。" +
+    "这是你返回最终结果的唯一方式。\n\n" +
+    "输出必须严格符合以下 JSON Schema:\n" + schemaJsonStr;
+
+  const toolDescription =
+    "Return structured output conforming to the JSON Schema. " +
+    "You MUST call this tool exactly once to return your final result.\n\n" +
+    "The output must conform to this JSON Schema:\n" + schemaJsonStr;
+
+  // Track structured-output call state per turn
+  let turnCallCount = 0;
+  let lastCallSucceeded = false;
+
+  // Register tool
   pi.registerTool({
     name: TOOL_NAME,
     label: "Structured Output",
-    description:
-      "Return structured output conforming to the JSON Schema. You MUST call this tool to return your final result.",
+    description: toolDescription,
     promptSnippet: "Call structured-output with your final structured answer",
     promptGuidelines: [
       "You MUST call structured-output as your final action.",
@@ -69,44 +106,46 @@ export default function structuredOutputExtension(pi: PiAPI): void {
           .join("; ");
         throw new Error(`Schema validation failed: ${errors}`);
       }
+      // Mark success for enforcement check
+      lastCallSucceeded = true;
       return {
         content: [
           { type: "text" as const, text: "Structured output recorded successfully." },
         ],
         details: params,
-        terminate: true,
       };
     },
   });
 
-  // System prompt injection
+  // System prompt injection with embedded schema
   pi.on("before_agent_start", async (_event: unknown, ctx: { addSystemInstruction: (s: string) => void }) => {
-    ctx.addSystemInstruction(SYSTEM_PROMPT);
+    ctx.addSystemInstruction(systemPrompt);
   });
 
-  // Enforcement: track tool calls via tool_execution_start flag
-  let hasStructuredOutputCall = false;
-
+  // Track calls — count for retry cap, success for enforcement
   pi.on("tool_execution_start", async (event: { toolName: string }) => {
     if (event.toolName === TOOL_NAME) {
-      hasStructuredOutputCall = true;
+      turnCallCount++;
     }
   });
 
   pi.on("turn_end", async () => {
-    if (!hasStructuredOutputCall) {
-      pi.sendUserMessage(ENFORCEMENT_MESSAGE);
+    if (lastCallSucceeded) {
+      // Reset for next potential turn (shouldn't happen, but defensive)
+      lastCallSucceeded = false;
+      turnCallCount = 0;
+      return;
     }
-  });
 
-  // Block non-workflow usage
-  pi.on("tool_call", async (event: { toolName: string }) => {
-    if (event.toolName === TOOL_NAME && !process.env[ENV_KEY]) {
-      return {
-        block: true as const,
-        reason: "This tool is only available in workflow structured-output mode",
-      };
+    // Retry cap: stop enforcement after MAX_RETRIES attempts
+    if (turnCallCount >= MAX_RETRIES) {
+      console.error(
+        `[${TOOL_NAME}] Max retries (${MAX_RETRIES}) reached without valid output. Giving up enforcement.`,
+      );
+      turnCallCount = 0;
+      return;
     }
-    return undefined;
+
+    pi.sendUserMessage(ENFORCEMENT_MESSAGE);
   });
 }