feat: directive recall + pre-tool-use hook

examples/agent_loop_demo.ts

@@ -0,0 +1,118 @@
+/**
+ * Dogfood: a real agent loop with vs. without directive recall.
+ *
+ * The agent is asked to make the judge model in `scripts/memory_quality_sweep.py`
+ * configurable. A captured lesson from a real session says exactly how we like it
+ * done — "thread `judge_model` through the scoring path instead of hardcoding
+ * `JUDGE_MODEL`". With `withDirectiveRecall`, that lesson fires the moment the agent
+ * reads the file and is injected into the read result, so the agent's edit follows
+ * the team convention. Without it, the agent guesses.
+ *
+ * Run:  cd memory-sdk-ts && npx tsx examples/agent_loop_demo.ts
+ *   env: OPENAI_API_KEY (gpt-4o-mini via @ai-sdk/openai), XTRACE_API_KEY, XTRACE_ORG_ID,
+ *        optional XTRACE_USER_ID (the user whose directives to recall), XTRACE_BASE_URL.
+ */
+import { generateText, stepCountIs, tool } from "ai";
+import { openai } from "@ai-sdk/openai";
+import { z } from "zod";
+import { MemoryClient } from "../src/index.js";
+import { withDirectiveRecall } from "../src/ai-sdk/index.js";
+
+const XTRACE_API_KEY = process.env.XTRACE_API_KEY!;
+const XTRACE_ORG_ID = process.env.XTRACE_ORG_ID!;
+const USER_ID = process.env.XTRACE_USER_ID ?? "directive-real-0625";
+
+// The file the agent will work on (hardcoded judge model — the thing to refactor).
+const FILE = "scripts/memory_quality_sweep.py";
+const FILE_CONTENT = `JUDGE_MODEL = "claude-haiku-4-5"  # hardcoded
+
+def _score_one(fact, judge):
+    resp = judge.classify(fact, model=JUDGE_MODEL)
+    return resp.label
+
+def run_sweep(facts):
+    judge = AnthropicJudge()
+    return [_score_one(f, judge) for f in facts]
+`;
+
+const client = new MemoryClient({
+  apiKey: XTRACE_API_KEY,
+  orgId: XTRACE_ORG_ID,
+  baseUrl: process.env.XTRACE_BASE_URL ?? "https://api.staging.xtrace.ai",
+});
+
+function makeTools(record: { patch?: string }) {
+  return {
+    read_file: tool({
+      description: "Read a source file. Returns its full text.",
+      inputSchema: z.object({ file_path: z.string() }),
+      execute: async () => FILE_CONTENT,
+    }),
+    propose_patch: tool({
+      description: "Propose the final edited file. Call once you know the change.",
+      inputSchema: z.object({ file_path: z.string(), new_content: z.string() }),
+      execute: async ({ new_content }) => {
+        record.patch = new_content;
+        return "recorded";
+      },
+    }),
+  };
+}
+
+const TASK =
+  `Make the judge model in ${FILE} configurable so a caller can choose it per run. ` +
+  `First read the file, then call propose_patch with the full edited file.`;
+
+async function runOnce(useDirectives: boolean): Promise<string | undefined> {
+  const record: { patch?: string } = {};
+  const fired: string[] = [];
+  const base = makeTools(record);
+  const tools = useDirectives
+    ? withDirectiveRecall(base, {
+        client,
+        user_id: USER_ID,
+        mode: "retrieve",
+        onDirectives: (ds) => fired.push(...ds.map((d) => `[${d.type}] ${d.text}`)),
+      })
+    : base;
+
+  await generateText({
+    model: openai("gpt-4o-mini"),
+    tools: tools as Parameters<typeof generateText>[0]["tools"],
+    prompt: TASK,
+    stopWhen: stepCountIs(5),
+  });
+
+  if (useDirectives) {
+    console.log(`  directives that fired (${fired.length}):`);
+    fired.forEach((f) => console.log(`    ↳ ${f.slice(0, 120)}`));
+  }
+  return record.patch;
+}
+
+(async () => {
+  if (!process.env.OPENAI_API_KEY) throw new Error("OPENAI_API_KEY not found");
+  if (!XTRACE_API_KEY || !XTRACE_ORG_ID) throw new Error("set XTRACE_API_KEY + XTRACE_ORG_ID");
+
+  console.log("=".repeat(72));
+  console.log("RUN A — no directive recall (the agent guesses the convention)");
+  console.log("=".repeat(72));
+  const a = await runOnce(false);
+  console.log("\n  proposed edit:\n" + indent(a));
+
+  console.log("\n" + "=".repeat(72));
+  console.log("RUN B — withDirectiveRecall (the team's lesson fires on read)");
+  console.log("=".repeat(72));
+  const b = await runOnce(true);
+  console.log("\n  proposed edit:\n" + indent(b));
+
+  const threadsB = /def\s+\w+\([^)]*judge_model|model=judge_model|judge_model:/.test(b ?? "");
+  const threadsA = /def\s+\w+\([^)]*judge_model|model=judge_model|judge_model:/.test(a ?? "");
+  console.log("\n" + "=".repeat(72));
+  console.log(`threads judge_model through the call path?   A=${threadsA}   B=${threadsB}`);
+  console.log("(the captured lesson: thread judge_model through, don't hardcode JUDGE_MODEL)");
+})();
+
+function indent(s?: string): string {
+  return (s ?? "(no patch proposed)").split("\n").map((l) => "    " + l).join("\n");
+}

package.json

@@ -51,6 +51,7 @@
     "prepublishOnly": "npm run typecheck && npm run test && npm run build"
   },
   "devDependencies": {
+    "@ai-sdk/openai": "^2.0.109",
     "@types/node": "^20.12.0",
     "ai": "^6.0.190",
     "openapi-typescript": "^7.4.0",

src/ai-sdk/directives.ts

@@ -0,0 +1,121 @@
+/**
+ * Vercel AI SDK integration: directive recall as a pre-tool-call hook.
+ *
+ * Where {@link createXtraceMemory} recalls *semantic* memory once per generation
+ * on the user message, directives fire per *tool call* on what the agent is about
+ * to touch. `withDirectiveRecall(tools, opts)` wraps a tools object so that BEFORE
+ * each tool runs, the directives keyed to that tool's `{ name, args }` are recalled
+ * (the symbol tripwire) and prepended to the tool's result — so the model sees the
+ * relevant lesson/procedure alongside what the tool did, before its next move.
+ *
+ * Usage:
+ *
+ *   import { streamText } from 'ai';
+ *   import { openai } from '@ai-sdk/openai';
+ *   import { MemoryClient } from '@xtraceai/memory';
+ *   import { withDirectiveRecall } from '@xtraceai/memory/ai-sdk';
+ *
+ *   const client = new MemoryClient({ apiKey, orgId });
+ *   const result = streamText({
+ *     model: openai('gpt-4o-mini'),
+ *     tools: withDirectiveRecall(myTools, { client, user_id: 'alice' }),
+ *     messages,
+ *   });
+ */
+import type { MemoryClient } from "../client.js";
+import { renderDirectivesPrompt } from "../memories.js";
+import type { DirectiveMemory, SearchMode } from "../types.js";
+
+export interface DirectiveRecallScope {
+  /** Scope directive recall to this user's directives. */
+  user_id?: string;
+  /** Optional agent scope (AND-narrows). */
+  agent_id?: string;
+  /** Optional app scope (AND-narrows). */
+  app_id?: string;
+  /** Optional shared group scope (any-of). */
+  group_ids?: string[];
+}
+
+export interface DirectiveRecallOptions extends DirectiveRecallScope {
+  /** The memory client used to recall. */
+  client: MemoryClient;
+  /**
+   * `"retrieve"` (default) — the deterministic stage-1 tripwire, fast (~0.1s),
+   * right for a per-tool-call hook. `"compose"` adds the stage-2 LLM gate.
+   */
+  mode?: SearchMode;
+  /** Max directives per tool call. Default 5. */
+  limit?: number;
+  /**
+   * Prepend the recalled directives to the tool's (string) result so the model
+   * reads them. Default `true`. Set `false` to leave results untouched and rely
+   * on {@link DirectiveRecallOptions.onDirectives} for handling.
+   */
+  injectIntoResult?: boolean;
+  /** Called whenever directives fire for a tool call — for logging / telemetry. */
+  onDirectives?: (directives: DirectiveMemory[], toolName: string) => void;
+}
+
+/**
+ * Recall the directives for one in-flight tool call → rendered context (or `""`).
+ * Fails soft: a recall error returns no directives rather than throwing.
+ */
+export async function directiveContextForToolCall(
+  opts: DirectiveRecallOptions,
+  toolName: string,
+  args: Record<string, unknown>,
+): Promise<{ directives: DirectiveMemory[]; context: string }> {
+  try {
+    const { data, context } = await opts.client.memories.recallDirectives({
+      action: { tool: toolName, args },
+      user_id: opts.user_id,
+      agent_id: opts.agent_id,
+      app_id: opts.app_id,
+      group_ids: opts.group_ids,
+      mode: opts.mode ?? "retrieve",
+      limit: opts.limit ?? 5,
+    });
+    return { directives: data, context: context ?? renderDirectivesPrompt(data) };
+  } catch {
+    return { directives: [], context: "" };
+  }
+}
+
+/**
+ * Wrap a Vercel AI SDK tools object with pre-tool-call directive recall. Tools
+ * with no `execute` are passed through untouched; a recall hiccup never breaks a
+ * tool (fail-soft).
+ */
+export function withDirectiveRecall<T extends Record<string, unknown>>(
+  tools: T,
+  opts: DirectiveRecallOptions,
+): T {
+  const inject = opts.injectIntoResult ?? true;
+  const out: Record<string, unknown> = {};
+  for (const [name, t] of Object.entries(tools)) {
+    const tool = t as { execute?: (...a: unknown[]) => unknown } & Record<string, unknown>;
+    if (typeof tool?.execute !== "function") {
+      out[name] = t;
+      continue;
+    }
+    const originalExecute = tool.execute.bind(tool);
+    out[name] = {
+      ...tool,
+      execute: async (args: Record<string, unknown>, execCtx: unknown) => {
+        const { directives, context } = await directiveContextForToolCall(
+          opts,
+          name,
+          args ?? {},
+        );
+        if (directives.length && opts.onDirectives) opts.onDirectives(directives, name);
+        const result = await originalExecute(args, execCtx);
+        if (inject && context && typeof result === "string") {
+          return `${context}\n\n${result}`;
+        }
+        return result;
+      },
+    };
+  }
+  return out as T;
+}

src/ai-sdk/index.ts

@@ -20,3 +20,6 @@ export type { CreateXtraceMemoryOptions } from "./provider.js";
 
 export { memoryTools } from "./tools.js";
 export type { MemoryToolsScope, MemoryToolsOptions } from "./tools.js";
+
+export { withDirectiveRecall, directiveContextForToolCall } from "./directives.js";
+export type { DirectiveRecallScope, DirectiveRecallOptions } from "./directives.js";