@atbash/atbash-langgraph 0.0.5-dev.2 → 0.0.5-dev.4

package/dist/index.d.ts

@@ -1,7 +1,7 @@
 import * as _langchain_core_messages from '@langchain/core/messages';
 import * as _langchain_langgraph from '@langchain/langgraph';
 import { StateGraph } from '@langchain/langgraph';
-import { Atbash, AgentAuth, AtbashOptions } from '@atbash/sdk';
+import { Atbash, MemoryGuardManager } from '@atbash/sdk';
 import * as _langchain_core_tools from '@langchain/core/tools';
 import { z } from 'zod';
 
@@ -16,21 +16,39 @@ type AtbashState = typeof AtbashStateAnnotation.State;
 
 interface GuardNodeOptions {
     client: Atbash;
+    guardManager?: MemoryGuardManager;
 }
 declare function createGuardNode(opts: GuardNodeOptions): (state: AtbashState) => Promise<Partial<AtbashState>>;
 
 interface AuditNodeOptions {
-    agent: AgentAuth;
-    clientOpts?: AtbashOptions;
+    client: Atbash;
 }
 declare function createAuditNode(opts: AuditNodeOptions): (state: AtbashState) => Promise<Partial<AtbashState>>;
 
 interface AtbashSafetyOptions {
-    agent?: AgentAuth;
     privkey?: string;
     endpoint?: string;
     toolsNode?: string;
     agentNode?: string;
+    /** Agent workspace directory — `~` is expanded. Defaults to `process.cwd()`. */
+    workspaceDir?: string;
+    /** Explicit MEMORY.md path. Overrides `workspaceDir/MEMORY.md`. */
+    memoryFilePath?: string;
+    /** How long (ms) to trust the local pointer between chain checks. Default 30000. */
+    memorySyncTTLMs?: number;
+    /** Block memory reads when a rolled-back version scores below this (1–10). Default 1 = warn only. */
+    memoryRollbackMinScore?: number;
+    /** Custom memory path patterns. Defaults to SDK built-ins. */
+    memoryPathPatterns?: string[];
+    /** Optional judge endpoint override. */
+    judgeEndpoint?: string;
+    /** Self-hosted judge response-signing pubkey (required when judgeEndpoint uses "self-hosted" policy). */
+    judgeVerifyPubKey?: string;
+    /** Organization name for chain resolution (public vs private chain). */
+    orgName?: string;
+    /** false = monitor mode: log but never block. Default true. */
+    enforce?: boolean;
+    debug?: boolean;
 }
 declare function addAtbashSafety(builder: StateGraph<AtbashState>, opts: AtbashSafetyOptions): StateGraph<_langchain_langgraph.StateType<{
     atbashVerdict: _langchain_langgraph.BaseChannel<string | null, string | _langchain_langgraph.OverwriteValue<string | null> | null, unknown>;
@@ -52,7 +70,7 @@ declare function addAtbashSafety(builder: StateGraph<AtbashState>, opts: AtbashS
     messages: _langchain_langgraph.BaseChannel<_langchain_core_messages.BaseMessage<_langchain_core_messages.MessageStructure<_langchain_core_messages.MessageToolSet>, _langchain_core_messages.MessageType>[], _langchain_langgraph.OverwriteValue<_langchain_core_messages.BaseMessage<_langchain_core_messages.MessageStructure<_langchain_core_messages.MessageToolSet>, _langchain_core_messages.MessageType>[]> | _langchain_langgraph.Messages, unknown>;
 }>>, "__start__", _langchain_langgraph.StateDefinition, _langchain_langgraph.StateDefinition, _langchain_langgraph.StateDefinition, unknown, unknown, unknown>;
 
-declare function createJudgeTool(agent: AgentAuth, endpoint?: string): _langchain_core_tools.DynamicStructuredTool<z.ZodObject<{
+declare function createJudgeTool(client: Atbash): _langchain_core_tools.DynamicStructuredTool<z.ZodObject<{
     action: z.ZodString;
     context: z.ZodString;
 }, "strip", z.ZodTypeAny, {

package/dist/index.js

@@ -33,11 +33,63 @@ function createGuardNode(opts) {
         atbashReason: "No tool calls detected"
       };
     }
-    const actionText = toolCalls.map((toolCall) => `${toolCall.name}(${JSON.stringify(toolCall.args)})`).join("; ");
+    const memoryHandledIds = /* @__PURE__ */ new Set();
+    if (opts.guardManager) {
+      for (const tc of toolCalls) {
+        let memDecision;
+        try {
+          memDecision = await opts.guardManager.handleBeforeToolCall(
+            { toolName: tc.name, args: tc.args },
+            {}
+          );
+        } catch (err) {
+          const reason = err instanceof Error ? err.message : String(err);
+          return {
+            messages: toolCalls.map(
+              (toolCall) => new ToolMessage({
+                tool_call_id: toolCall.id,
+                content: toolCall.id === tc.id ? `Memory guard error: ${reason}` : "Blocked \u2014 memory safety check failed for another tool call"
+              })
+            ),
+            atbashVerdict: "BLOCK",
+            atbashReason: `Memory guard error: ${reason}`,
+            atbashToolCallId: null,
+            atbashConfidence: null
+          };
+        }
+        if (memDecision !== null) {
+          if (memDecision.block) {
+            return {
+              messages: toolCalls.map(
+                (toolCall) => new ToolMessage({
+                  tool_call_id: toolCall.id,
+                  content: toolCall.id === tc.id ? `Memory write blocked by Atbash: ${memDecision.blockReason ?? ""}` : "Blocked \u2014 another tool call was blocked by memory safety"
+                })
+              ),
+              atbashVerdict: "BLOCK",
+              atbashReason: memDecision.blockReason ?? "memory write blocked",
+              atbashToolCallId: null,
+              atbashConfidence: null
+            };
+          }
+          memoryHandledIds.add(tc.id);
+        }
+      }
+    }
+    const auditCalls = toolCalls.filter((tc) => !memoryHandledIds.has(tc.id));
+    if (auditCalls.length === 0) {
+      return {
+        atbashVerdict: "ALLOW",
+        atbashReason: "Memory operations validated by guard",
+        atbashToolCallId: null,
+        atbashConfidence: null
+      };
+    }
+    const actionText = auditCalls.map((toolCall) => `${toolCall.name}(${JSON.stringify(toolCall.args)})`).join("; ");
     try {
       const decision = await opts.client.auditToolCall({
-        toolName: toolCalls.map((t) => t.name).join(",") || "langgraph_batch",
-        args: toolCalls,
+        toolName: auditCalls.map((t) => t.name).join(",") || "langgraph_batch",
+        args: auditCalls,
         context: `LangGraph agent attempting: ${actionText}`
       });
       if (decision.verdict === "BLOCK" || decision.verdict === "ERROR") {
@@ -101,9 +153,6 @@ function createGuardNode(opts) {
   };
 }
 
-// src/nodes/auditNode.ts
-import { Atbash } from "@atbash/sdk";
-
 // src/utils.ts
 function truncateText(value, max = 500) {
   return value.length <= max ? value : `${value.slice(0, max - 3)}...`;
@@ -111,12 +160,12 @@ function truncateText(value, max = 500) {
 
 // src/nodes/auditNode.ts
 function createAuditNode(opts) {
+  const { client } = opts;
   return async (state) => {
     const lastMessage = state.messages[state.messages.length - 1];
     const content = typeof lastMessage?.content === "string" ? lastMessage.content : JSON.stringify(lastMessage?.content ?? "");
     try {
-      const client = new Atbash(opts.agent.privkey, opts.clientOpts);
-      await client.judgeAction(
+      await client.logToolCall(
         truncateText(content, 500),
         "LangGraph tool execution completed"
       );
@@ -128,25 +177,45 @@ function createAuditNode(opts) {
 
 // src/builder.ts
 import {
-  Atbash as Atbash2,
-  loadAgent,
+  Atbash,
+  createMemoryGuardManager,
   setupTelemetry,
   shutdownTelemetry
 } from "@atbash/sdk";
+import { homedir } from "os";
+function expandHome(p) {
+  return p.replace(/^~(?=\/|$)/, homedir());
+}
 function addAtbashSafety(builder, opts) {
   setupTelemetry({ enabled: true, source: "plugin:langgraph" });
   process.once("beforeExit", () => shutdownTelemetry());
   process.once("SIGINT", () => shutdownTelemetry().finally(() => process.exit(0)));
   process.once("SIGTERM", () => shutdownTelemetry().finally(() => process.exit(0)));
-  const privkey = opts.privkey ?? process.env.ATBASH_AGENT_PRIVKEY;
-  const agent = opts.agent ?? loadAgent(privkey ?? "");
-  const clientOpts = opts.endpoint ? { endpoint: opts.endpoint } : void 0;
-  const client = new Atbash2(agent.privkey, clientOpts);
+  const privkey = opts.privkey ?? process.env.ATBASH_AGENT_PRIVKEY ?? "";
+  const clientOpts = {};
+  if (opts.endpoint) clientOpts.endpoint = opts.endpoint;
+  const client = new Atbash(privkey, clientOpts);
+  const guard = createMemoryGuardManager({
+    auth: client.auth,
+    workspaceDir: opts.workspaceDir ? expandHome(opts.workspaceDir) : process.cwd(),
+    ...opts.memoryFilePath ? { memoryFilePath: expandHome(opts.memoryFilePath) } : {},
+    ...opts.memorySyncTTLMs !== void 0 ? { ttlMs: opts.memorySyncTTLMs } : {},
+    ...opts.memoryRollbackMinScore !== void 0 ? { rollbackMinScore: opts.memoryRollbackMinScore } : {},
+    ...opts.memoryPathPatterns ? { memoryPathPatterns: opts.memoryPathPatterns } : {},
+    ...opts.judgeEndpoint ? { judgeEndpoint: opts.judgeEndpoint } : {},
+    ...opts.judgeVerifyPubKey ? { judgeVerifyPubKey: opts.judgeVerifyPubKey } : {},
+    ...opts.orgName ? { orgName: opts.orgName } : {},
+    enforce: opts.enforce ?? true,
+    debug: opts.debug ?? false
+  });
+  guard.runBootProbe().catch((err) => {
+    console.warn("[atbash] boot probe failed:", err instanceof Error ? err.message : String(err));
+  });
   const graph = builder;
   const toolsNode = opts.toolsNode ?? "tools";
   const agentNode = opts.agentNode ?? "agent";
-  graph.addNode("atbash_guard", createGuardNode({ client }));
-  graph.addNode("atbash_audit", createAuditNode({ agent, clientOpts }));
+  graph.addNode("atbash_guard", createGuardNode({ client, guardManager: guard }));
+  graph.addNode("atbash_audit", createAuditNode({ client }));
   graph.addConditionalEdges("atbash_guard", (state) => {
     return state.atbashVerdict === "ALLOW" ? toolsNode : agentNode;
   });
@@ -156,19 +225,21 @@ function addAtbashSafety(builder, opts) {
 }
 
 // src/tools/judgeTool.ts
-import { Atbash as Atbash3 } from "@atbash/sdk";
 import { tool } from "@langchain/core/tools";
 import { z } from "zod";
-function createJudgeTool(agent, endpoint) {
+function createJudgeTool(client) {
   return tool(
     async ({ action, context }) => {
-      const client = new Atbash3(agent.privkey, endpoint ? { endpoint } : void 0);
-      const result = await client.judgeAction(action, context);
+      const decision = await client.auditToolCall({
+        toolName: "atbash_safety_check",
+        args: { action },
+        context
+      });
       return JSON.stringify({
-        verdict: result.verdict,
-        reason: result.reason,
-        confidence: result.confidence,
-        tool_call_id: result.toolCallId
+        verdict: decision.verdict,
+        allow: decision.allow,
+        reason: decision.reason,
+        tool_call_id: decision.toolCallId
       });
     },
     {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@atbash/atbash-langgraph",
-  "version": "0.0.5-dev.2",
+  "version": "0.0.5-dev.4",
   "description": "Atbash safety guard and audit nodes for LangGraph workflows",
   "type": "module",
   "main": "dist/index.js",
@@ -37,7 +37,7 @@
     "prepublishOnly": "npm run build"
   },
   "dependencies": {
-    "@atbash/sdk": "0.5.3-dev.0",
+    "@atbash/sdk": "0.10.5-dev.0",
     "zod": "^3.25.76"
   },
   "peerDependencies": {