@silmaril-security/sdk 0.7.10 → 0.8.0

package/dist/adapters/langchain.js

@@ -148,7 +148,81 @@ function extractTextFromDocuments(documents) {
   return parts.join("\n");
 }
 
+// src/adapters/thread-identity.ts
+function record(value) {
+  return value !== null && typeof value === "object" && !Array.isArray(value) ? value : void 0;
+}
+function validThreadId(value) {
+  return typeof value === "string" && value.trim().length > 0 ? value : void 0;
+}
+function nativeThreadId(runtime, metadata, inherited, currentConfig) {
+  const explicit = record(record(runtime)?.configurable);
+  const current = record(record(currentConfig)?.configurable);
+  for (const candidate of [explicit, current, record(metadata)]) {
+    if (candidate && Object.prototype.hasOwnProperty.call(candidate, "thread_id")) {
+      return validThreadId(candidate.thread_id);
+    }
+  }
+  return inherited;
+}
+function rejectConversationOptions(options) {
+  if ("conversationId" in options || "conversation_id" in options) {
+    throw new TypeError("conversationId/conversation_id was removed. Pass the native configurable.thread_id when invoking the agent or runnable instead.");
+  }
+}
+var RunThreadIdentity = class {
+  constructor(boundIdentity, currentConfig = () => void 0) {
+    this.boundIdentity = boundIdentity;
+    this.currentConfig = currentConfig;
+  }
+  boundIdentity;
+  currentConfig;
+  runs = /* @__PURE__ */ new Map();
+  children = /* @__PURE__ */ new Map();
+  start(runId, parentRunId, metadata) {
+    const discarded = this.finish(runId);
+    if (this.runs.size >= 4096) {
+      const oldestRoot = [...this.runs].find(([, run]) => run.parentRunId === void 0 || !this.runs.has(run.parentRunId))?.[0] ?? this.runs.keys().next().value;
+      if (oldestRoot !== void 0) discarded.push(...this.finish(oldestRoot));
+    }
+    const inherited = parentRunId === void 0 ? void 0 : this.runs.get(parentRunId)?.threadId;
+    const exposed = record(metadata);
+    const threadId = this.boundIdentity && (parentRunId === void 0 || !this.runs.has(parentRunId)) ? this.boundIdentity.threadId : exposed && Object.prototype.hasOwnProperty.call(exposed, "thread_id") ? typeof exposed.langgraph_node === "string" ? nativeThreadId(void 0, exposed, inherited, this.currentConfig()) : validThreadId(exposed.thread_id) : inherited;
+    this.runs.set(runId, { threadId, parentRunId });
+    if (parentRunId !== void 0 && parentRunId !== runId) {
+      const children = this.children.get(parentRunId) ?? /* @__PURE__ */ new Set();
+      children.add(runId);
+      this.children.set(parentRunId, children);
+    }
+    return discarded;
+  }
+  get(runId) {
+    const run = this.runs.get(runId);
+    return run ? run.threadId : this.boundIdentity?.threadId;
+  }
+  finish(runId) {
+    const pending = [runId];
+    const seen = /* @__PURE__ */ new Set();
+    while (pending.length) {
+      const id = pending.pop();
+      if (seen.has(id)) continue;
+      seen.add(id);
+      pending.push(...this.children.get(id) ?? []);
+      this.children.delete(id);
+      const parent = this.runs.get(id)?.parentRunId;
+      this.runs.delete(id);
+      if (parent !== void 0) {
+        const siblings = this.children.get(parent);
+        siblings?.delete(id);
+        if (siblings?.size === 0) this.children.delete(parent);
+      }
+    }
+    return [...seen];
+  }
+};
+
 // src/adapters/langchain.ts
+var handlerCleanups = /* @__PURE__ */ new WeakMap();
 var USER_ROLES = /* @__PURE__ */ new Set(["human", "user"]);
 function getMessageRole(message) {
   if (typeof message.role === "string") {
@@ -157,7 +231,8 @@ function getMessageRole(message) {
   if (typeof message.type === "string") {
     return message.type.toLowerCase();
   }
-  return "";
+  const legacyType = message.getType?.() ?? message._getType?.();
+  return typeof legacyType === "string" ? legacyType.toLowerCase() : "";
 }
 function extractMessageText(content) {
   if (content === void 0) {
@@ -186,7 +261,16 @@ function findLastUserMessage(messages) {
   return void 0;
 }
 async function createLangChainHandler(firewall, options = {}) {
-  const { BaseCallbackHandler } = await import("@langchain/core/callbacks/base");
+  return buildLangChainHandler(firewall, options);
+}
+async function buildLangChainHandler(firewall, options, boundIdentity) {
+  rejectConversationOptions(options);
+  const { ensureConfig } = await import("@langchain/core/runnables");
+  const threads = new RunThreadIdentity(boundIdentity, () => ensureConfig());
+  const { BaseTracer } = await import("@langchain/core/tracers/base");
+  let finishRun = (runId) => {
+    threads.finish(runId);
+  };
   const enabledHooks = resolveHooks(options.hooks);
   const includeTool = options.includeTool ?? true;
   const failOpen = options.failOpen ?? true;
@@ -205,14 +289,14 @@ async function createLangChainHandler(firewall, options = {}) {
       logger("onClassify callback threw", err);
     }
   };
-  const classify = async (text, hookLabel, runId, toolName) => {
+  const classifyEvent = async (text, hookLabel, runId, toolName) => {
     let result;
     try {
       result = await firewall.classify(text, {
         hook: hookLabel,
         metadata: {
           langgraph: { run_id: runId },
-          ...options.conversationId === void 0 ? {} : { conversationId: options.conversationId }
+          ...threads.get(runId) === void 0 ? {} : { conversationId: threads.get(runId) }
         },
         ...requestedMode !== void 0 ? { mode: requestedMode } : {},
         ...toolName !== void 0 ? { toolName } : {}
@@ -256,97 +340,326 @@ async function createLangChainHandler(firewall, options = {}) {
       result
     });
   };
-  class SilmarilFirewallHandler extends BaseCallbackHandler {
+  const classify = async (text, hook, runId, toolName) => {
+    try {
+      await classifyEvent(text, hook, runId, toolName);
+    } catch (error) {
+      finishRun(runId);
+      throw error;
+    }
+  };
+  function serialized(value, name) {
+    if (value && typeof value === "object" && "id" in value && Array.isArray(value.id)) return value;
+    return { lc: 1, type: "not_implemented", id: [name] };
+  }
+  class SilmarilFirewallHandler extends BaseTracer {
     name = "silmaril_firewall_handler";
     raiseError = true;
     awaitHandlers = true;
-    async handleChatModelStart(_llm, messages, runId) {
-      if (!enabledHooks.has(FirewallHook.CHAT_MODEL_START)) {
-        return;
+    // LangChain clears runnable context inside ordinary callbacks. Its synchronous
+    // tracer pre-start hooks capture native config before that boundary. Runs are
+    // memory-only and every terminal path discards both native tracer maps.
+    async persistRun() {
+    }
+    getRunById(runId) {
+      return runId === void 0 ? void 0 : this.runMap.get(runId);
+    }
+    // Only lifecycle identity is needed. Do not retain prompts, outputs, token
+    // events, or completed children when upstream omits a cancellation callback.
+    _addRunToRunMap(run) {
+      return super._addRunToRunMap({
+        ...run,
+        inputs: {},
+        extra: {},
+        serialized: {},
+        tags: [],
+        child_runs: [],
+        events: []
+      });
+    }
+    _addChildRun() {
+    }
+    async _endTrace(run) {
+      try {
+        await super._endTrace(run);
+      } finally {
+        run.inputs = {};
+        delete run.outputs;
+        delete run.error;
+        run.extra = {};
+        run.events = [];
+        run.child_runs = [];
       }
-      const batches = messages;
-      const flat = [];
-      for (const batch of batches) {
-        for (const m of batch) {
-          flat.push(m);
-        }
+    }
+    async handleText() {
+    }
+    async handleAgentAction() {
+    }
+    async handleAgentEnd() {
+    }
+    async handleLLMNewToken(_token, _idx, runId) {
+      return this.getRunById(runId) ?? this._createRunForLLMStart(serialized(void 0, "llm"), [], runId);
+    }
+    capture(runId, parentRunId, metadata) {
+      for (const id of threads.start(runId, parentRunId, metadata)) {
+        this.runMap.delete(id);
+        this.runTreeMap?.delete(id);
+      }
+    }
+    finishAll() {
+      for (const id of [...this.runMap.keys()]) this.finishRun(id);
+      this.runTreeMap?.clear();
+    }
+    finishRun(runId) {
+      for (const id of threads.finish(runId)) {
+        this.runMap.delete(id);
+        this.runTreeMap?.delete(id);
       }
-      const lastUser = findLastUserMessage(flat);
-      if (!lastUser) {
-        return;
+    }
+    _createRunForLLMStart(...args) {
+      this.capture(args[2], args[3], args[6]);
+      return super._createRunForLLMStart(...args);
+    }
+    _createRunForChatModelStart(...args) {
+      this.capture(args[2], args[3], args[6]);
+      return super._createRunForChatModelStart(...args);
+    }
+    _createRunForChainStart(...args) {
+      this.capture(args[2], args[3], args[5]);
+      return super._createRunForChainStart(...args);
+    }
+    _createRunForToolStart(...args) {
+      this.capture(args[2], args[3], args[5]);
+      return super._createRunForToolStart(...args);
+    }
+    _createRunForRetrieverStart(...args) {
+      this.capture(args[2], args[3], args[5]);
+      return super._createRunForRetrieverStart(...args);
+    }
+    async handleChatModelStart(llm, messages, runId, parentRunId, extraParams, tags, metadata) {
+      const batches = messages;
+      const run = this.getRunById(runId) ?? this._createRunForChatModelStart(
+        serialized(llm, "chat_model"),
+        messages,
+        runId,
+        parentRunId,
+        extraParams,
+        tags,
+        metadata
+      );
+      if (enabledHooks.has(FirewallHook.CHAT_MODEL_START)) {
+        const lastUser = findLastUserMessage(batches.flat());
+        const text = lastUser ? extractMessageText(lastUser.content).trim() : "";
+        if (text) await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.CHAT_MODEL_START], runId);
       }
-      const text = extractMessageText(lastUser.content).trim();
-      if (!text) {
-        return;
+      return run;
+    }
+    async handleLLMStart(llm, prompts, runId, parentRunId, extraParams, tags, metadata) {
+      const run = this.getRunById(runId) ?? this._createRunForLLMStart(
+        serialized(llm, "llm"),
+        prompts,
+        runId,
+        parentRunId,
+        extraParams,
+        tags,
+        metadata
+      );
+      if (enabledHooks.has(FirewallHook.LLM_START)) {
+        const text = extractTextFromPrompts(prompts);
+        if (text) await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.LLM_START], runId);
       }
-      await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.CHAT_MODEL_START], runId);
+      return run;
     }
-    async handleLLMStart(_llm, prompts, runId) {
-      if (!enabledHooks.has(FirewallHook.LLM_START)) {
-        return;
+    async handleToolStart(tool, inputStr, runId, parentRunId, tags, metadata) {
+      const run = this.getRunById(runId) ?? this._createRunForToolStart(
+        serialized(tool, tool?.name ?? "tool"),
+        inputStr,
+        runId,
+        parentRunId,
+        tags,
+        metadata
+      );
+      if (includeTool && enabledHooks.has(FirewallHook.TOOL_START)) {
+        const text = extractTextFromToolInput(inputStr);
+        if (text) await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.TOOL_START], runId, tool?.name);
       }
-      const text = extractTextFromPrompts(prompts);
-      if (!text) {
-        return;
+      return run;
+    }
+    async handleRetrieverStart(retriever, query, runId, parentRunId, tags, metadata) {
+      const run = this.getRunById(runId) ?? this._createRunForRetrieverStart(
+        serialized(retriever, "retriever"),
+        query,
+        runId,
+        parentRunId,
+        tags,
+        metadata
+      );
+      if (enabledHooks.has(FirewallHook.RETRIEVER_START) && query.trim()) {
+        await classify(query.trim(), FIREWALL_HOOK_TO_LABEL[FirewallHook.RETRIEVER_START], runId);
       }
-      await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.LLM_START], runId);
+      return run;
     }
-    async handleToolStart(tool, inputStr, runId) {
-      if (!includeTool || !enabledHooks.has(FirewallHook.TOOL_START)) {
-        return;
+    async handleLLMEnd(output, runId) {
+      try {
+        if (!this.getRunById(runId)) this._createRunForLLMStart(serialized(void 0, "llm"), [], runId);
+        if (enabledHooks.has(FirewallHook.LLM_END)) {
+          const text = extractTextFromLLMResult(output);
+          if (text) await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.LLM_END], runId);
+        }
+        return await super.handleLLMEnd(output, runId);
+      } finally {
+        finishRun(runId);
       }
-      const text = extractTextFromToolInput(inputStr);
-      if (!text) {
-        return;
+    }
+    async handleToolEnd(output, runId, _parentRunId, _tags, kwargs) {
+      try {
+        if (!this.getRunById(runId)) this._createRunForToolStart(serialized(void 0, "tool"), "", runId);
+        if (includeTool && enabledHooks.has(FirewallHook.TOOL_END)) {
+          const text = String(output).trim();
+          if (text) await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.TOOL_END], runId, kwargs?.name);
+        }
+        return await super.handleToolEnd(output, runId);
+      } finally {
+        finishRun(runId);
       }
-      const toolName = tool?.name;
-      await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.TOOL_START], runId, toolName);
     }
-    async handleRetrieverStart(_retriever, query, runId) {
-      if (!enabledHooks.has(FirewallHook.RETRIEVER_START)) {
-        return;
+    async handleRetrieverEnd(documents, runId) {
+      try {
+        if (!this.getRunById(runId)) this._createRunForRetrieverStart(serialized(void 0, "retriever"), "", runId);
+        if (enabledHooks.has(FirewallHook.RETRIEVER_END)) {
+          const text = extractTextFromDocuments(documents);
+          if (text) await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.RETRIEVER_END], runId);
+        }
+        return await super.handleRetrieverEnd(documents, runId);
+      } finally {
+        finishRun(runId);
       }
-      const text = query.trim();
-      if (!text) {
-        return;
+    }
+    async handleChainStart(chain, inputs, runId, parentRunId, tags, metadata, runType, name) {
+      return super.handleChainStart(serialized(chain, "chain"), inputs, runId, parentRunId, tags, metadata, runType, name);
+    }
+    async handleChainEnd(...args) {
+      if (!this.getRunById(args[1])) this._createRunForChainStart(serialized(void 0, "chain"), {}, args[1]);
+      try {
+        return await super.handleChainEnd(...args);
+      } finally {
+        finishRun(args[1]);
       }
-      await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.RETRIEVER_START], runId);
     }
-    async handleLLMEnd(output, runId) {
-      if (!enabledHooks.has(FirewallHook.LLM_END)) {
-        return;
+    async handleChainError(...args) {
+      if (!this.getRunById(args[1])) this._createRunForChainStart(serialized(void 0, "chain"), {}, args[1]);
+      try {
+        return await super.handleChainError(...args);
+      } finally {
+        finishRun(args[1]);
       }
-      const text = extractTextFromLLMResult(output);
-      if (!text) {
-        return;
+    }
+    async handleLLMError(...args) {
+      if (!this.getRunById(args[1])) this._createRunForLLMStart(serialized(void 0, "llm"), [], args[1]);
+      try {
+        return await super.handleLLMError(...args);
+      } finally {
+        finishRun(args[1]);
       }
-      await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.LLM_END], runId);
     }
-    async handleToolEnd(output, runId, _parentRunId, _tags, _kwargs) {
-      if (!includeTool || !enabledHooks.has(FirewallHook.TOOL_END)) {
-        return;
+    async handleToolError(...args) {
+      if (!this.getRunById(args[1])) this._createRunForToolStart(serialized(void 0, "tool"), "", args[1]);
+      try {
+        return await super.handleToolError(...args);
+      } finally {
+        finishRun(args[1]);
       }
-      const text = String(output).trim();
-      if (!text) {
-        return;
+    }
+    async handleRetrieverError(...args) {
+      if (!this.getRunById(args[1])) this._createRunForRetrieverStart(serialized(void 0, "retriever"), "", args[1]);
+      try {
+        return await super.handleRetrieverError(...args);
+      } finally {
+        finishRun(args[1]);
       }
-      const toolName = _kwargs?.name;
-      await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.TOOL_END], runId, toolName);
     }
-    async handleRetrieverEnd(documents, runId) {
-      if (!enabledHooks.has(FirewallHook.RETRIEVER_END)) {
-        return;
+  }
+  const handler = new SilmarilFirewallHandler();
+  finishRun = (runId) => {
+    handler.finishRun(runId);
+  };
+  handlerCleanups.set(handler, () => handler.finishAll());
+  return handler;
+}
+async function protectLangChainRunnable(firewall, runnable, options = {}) {
+  rejectConversationOptions(options);
+  const { RunnableBinding, RunnableRetry, ensureConfig, mergeConfigs } = await import("@langchain/core/runnables");
+  const { IterableReadableStream } = await import("@langchain/core/utils/stream");
+  const configFactories = [async (config) => {
+    const effective = ensureConfig(config);
+    const handler = await buildLangChainHandler(firewall, options, {
+      threadId: validThreadId(effective.configurable?.thread_id)
+    });
+    return { callbacks: [handler] };
+  }];
+  class ProtectedBinding extends RunnableBinding {
+    async scopedConfig(...configs) {
+      const created = [];
+      const scopedFactories = configFactories.map((factory) => async (config2) => {
+        const update = await factory(config2);
+        if (Array.isArray(update.callbacks)) created.push(...update.callbacks);
+        return update;
+      });
+      const scoped = new RunnableBinding({ bound: this.bound, config: this.config, configFactories: scopedFactories });
+      const config = await scoped._mergeConfig(...configs);
+      return { config, dispose: () => {
+        for (const handler of created) handlerCleanups.get(handler)?.();
+      } };
+    }
+    async invoke(input, options2) {
+      const scoped = await this.scopedConfig(options2, this.kwargs);
+      try {
+        return await this.bound.invoke(input, scoped.config);
+      } finally {
+        scoped.dispose();
       }
-      const text = extractTextFromDocuments(documents);
-      if (!text) {
-        return;
+    }
+    async stream(input, options2) {
+      const scoped = await this.scopedConfig(ensureConfig(options2), this.kwargs);
+      let upstream;
+      try {
+        upstream = await this.bound.stream(input, scoped.config);
+      } catch (error) {
+        scoped.dispose();
+        throw error;
       }
-      await classify(text, FIREWALL_HOOK_TO_LABEL[FirewallHook.RETRIEVER_END], runId);
+      const managed = async function* () {
+        try {
+          yield* upstream;
+        } finally {
+          scoped.dispose();
+        }
+      };
+      return IterableReadableStream.fromAsyncGenerator(managed());
+    }
+    async *_streamIterator(input, options2) {
+      const stream = await this.stream(input, options2);
+      yield* stream;
+    }
+    withConfig(config) {
+      return new ProtectedBinding({ bound: this.bound, config: mergeConfigs(this.config, config), configFactories });
+    }
+    withListeners(fields) {
+      return new RunnableBinding({ bound: this, config: {} }).withListeners(fields);
+    }
+    withRetry(fields) {
+      return new RunnableRetry({
+        bound: this,
+        config: {},
+        ...fields?.stopAfterAttempt === void 0 ? {} : { maxAttemptNumber: fields.stopAfterAttempt },
+        ...fields?.onFailedAttempt === void 0 ? {} : { onFailedAttempt: fields.onFailedAttempt }
+      });
     }
   }
-  return new SilmarilFirewallHandler();
+  return new ProtectedBinding({ bound: runnable, config: {}, configFactories });
 }
 export {
-  createLangChainHandler
+  createLangChainHandler,
+  protectLangChainRunnable
 };
 //# sourceMappingURL=langchain.js.map