@librechat/agents 3.3.4 → 3.3.5

package/src/run.ts

@@ -49,8 +49,14 @@ import {
   createTitleRunnable,
 } from '@/utils/title';
 import { createTokenCounter, encodingForModel } from '@/utils/tokens';
+import { resolveMaxSeals } from '@/llm/preempt';
 import { initializeLangfuseTracing } from './instrumentation';
-import { GraphEvents, Callback, TitleMethod } from '@/common';
+import {
+  Callback,
+  GraphEvents,
+  TitleMethod,
+  DEFAULT_RECURSION_LIMIT,
+} from '@/common';
 import { MultiAgentGraph } from '@/graphs/MultiAgentGraph';
 import { getTraceIdSeed } from '@/langfuseRuntimeContext';
 import { StandardGraph } from '@/graphs/Graph';
@@ -140,6 +146,7 @@ export class Run<_T extends t.BaseGraphState> {
   private interruptingToolNames?: string[];
   private toolExecution?: t.ToolExecutionConfig;
   private subagentUsageSink?: t.SubagentUsageSink;
+  private preemption?: t.StreamPreemption;
   private indexTokenCountMap?: Record<string, number>;
   calibrationRatio: number = 1;
   graphRunnable?: t.CompiledStateWorkflow;
@@ -201,6 +208,7 @@ export class Run<_T extends t.BaseGraphState> {
     this.interruptingToolNames = config.interruptingToolNames;
     this.toolExecution = config.toolExecution;
     this.subagentUsageSink = config.subagentUsageSink;
+    this.preemption = config.preemption;
 
     if (!config.graphConfig) {
       throw new Error('Graph config not provided');
@@ -275,6 +283,7 @@ export class Run<_T extends t.BaseGraphState> {
       indexTokenCountMap: this.indexTokenCountMap,
       calibrationRatio: this.calibrationRatio,
       subagentUsageSink: this.subagentUsageSink,
+      preemption: this.preemption,
     });
     /** Propagate compile options from graph config */
     standardGraph.compileOptions = this.applyHITLCheckpointerFallback(
@@ -306,6 +315,7 @@ export class Run<_T extends t.BaseGraphState> {
       indexTokenCountMap: this.indexTokenCountMap,
       calibrationRatio: this.calibrationRatio,
       subagentUsageSink: this.subagentUsageSink,
+      preemption: this.preemption,
     });
 
     multiAgentGraph.compileOptions =
@@ -541,6 +551,15 @@ export class Run<_T extends t.BaseGraphState> {
     return this.Graph?.getResolvedInstructionOverhead();
   }
 
+  /**
+   * Cooperative-seal counters for this run. `emptyBoundaries` is the one to
+   * watch: it counts seals whose `PreemptBoundary` produced nothing to
+   * inject, which ends the turn early and leaves the answer unfinished.
+   */
+  getPreemptStats(): t.PreemptStats {
+    return this.Graph?.getPreemptStats() ?? { seals: 0, emptyBoundaries: 0 };
+  }
+
   getToolCount(): number {
     return this.Graph?.getToolCount() ?? 0;
   }
@@ -703,12 +722,38 @@ export class Run<_T extends t.BaseGraphState> {
     const isResume = inputs instanceof Command;
     const stateInputs = isResume ? undefined : (inputs as t.IState);
 
+    /**
+     * Every honored seal costs one extra superstep, so a preemption-enabled
+     * run reserves headroom for its whole seal budget. Without it, a
+     * tool-heavy agent that gets preempted could hit `GraphRecursionError` —
+     * which surfaces as a thrown stream, setting `streamThrew`, firing
+     * `StopFailure`, and wiping via `clearHeavyState()` exactly the partial
+     * content the seal existed to preserve.
+     */
+    const recursionLimit =
+      (callerConfig.recursionLimit ?? DEFAULT_RECURSION_LIMIT) +
+      (this.preemption != null ? resolveMaxSeals(this.preemption.maxSeals) : 0);
+
     const config: t.RunStreamConfig = {
-      recursionLimit: 50,
       ...callerConfig,
+      recursionLimit,
       configurable: { ...callerConfig.configurable },
     };
 
+    /**
+     * Cancellation can arrive either at graph construction or per-call through
+     * `callerConfig.signal`, and boundary hooks need to observe both — for a
+     * multi-agent run the construction signal does not exist at all, since
+     * `MultiAgentGraphConfig` exposes none. Carried on its own field, assigned
+     * unconditionally: writing into `graph.signal` would leak this call's
+     * controller into later calls (model-call config and subagent
+     * parentSignal read that field, and `clearHeavyState()` is skipped on
+     * HITL interrupts), while a conditional write would keep observing a
+     * stale controller the host has since aborted. The boundary dispatch
+     * composes both channels; see `dispatchPreemptBoundary`.
+     */
+    graph.callerSignal = callerConfig.signal;
+
     /**
      * Skip `resetValues` on resume — we're continuing an in-flight
      * run, not starting a fresh one. Resetting would wipe the
@@ -923,10 +968,17 @@ export class Run<_T extends t.BaseGraphState> {
          * graph doesn't take another model turn after the halting
          * operation completes.
          *
-         * Limitation: the current step (in-flight model call, ongoing
-         * tool batch) is not aborted — only the next step is skipped.
-         * This matches Claude Code's `continue: false` semantic where
-         * the active operation finishes before halting takes effect.
+         * This `break` is NOT graceful, despite what a `continue: false`
+         * reading suggests. Leaving the `for await` calls the iterator's
+         * `return()`, which cancels the reader
+         * (`@langchain/core/utils/stream`), and langgraph's stream wrapper
+         * turns that cancel into `_abortController.abort()`
+         * (`pregel/stream.js`). The in-flight model call or tool batch is
+         * torn down where it stands — it does not finish first.
+         *
+         * A halt is therefore the wrong tool for "stop generating but keep
+         * what you have". That is what `RunConfig.preemption` is for: it
+         * seals the stream at a provider-safe boundary and keeps the run.
          */
         const haltSignal = this.hookRegistry?.getHaltSignal(this.id);
         if (haltSignal != null) {
@@ -959,6 +1011,16 @@ export class Run<_T extends t.BaseGraphState> {
             threadId,
             agentId: graph.defaultAgentId,
             messages: graph.getRunMessages() ?? stateInputs?.messages ?? [],
+            /**
+             * A seal whose boundary ended the turn early must say so. The
+             * hook-supplied reason wins when a `PreemptBoundary` hook halted
+             * with one — a persistence/audit `Stop` hook should record the
+             * actual cause, not the generic label — and `preempt_incomplete`
+             * is reserved for the boundary that simply had nothing to inject.
+             */
+            stopReason:
+              graph.preemptHaltReason ??
+              (graph.preemptIncomplete ? 'preempt_incomplete' : undefined),
             stopHookActive: false, // will be true when stop is triggered by a hook (Phase 2)
           },
           sessionId: this.id,
@@ -966,6 +1028,29 @@ export class Run<_T extends t.BaseGraphState> {
           /* Stop hook errors must not masquerade as stream failures */
         });
       }
+
+      /**
+       * A `PreemptBoundary` hook that returned `preventContinuation` has its
+       * registry halt cleared by the graph — that is what stops the halt from
+       * cancelling the stream before the sealed turn commits — so the reason
+       * is carried across on the graph instead. Surfaced here, AFTER the
+       * `Stop` dispatch above, so the host still receives a completion signal
+       * to persist the partial answer with while `getHaltReason()` correctly
+       * reports that a hook stopped the run rather than the model finishing.
+       *
+       * An empty boundary — sealed, but nothing to inject because the host's
+       * queue was drained or cancelled in the meantime — cut the answer short
+       * just as surely, only without a hook-supplied reason. It surfaces
+       * through the same channel under the same name the `Stop` dispatch
+       * already used for its `stopReason`, so terminal consumers
+       * (`AgentSession` emits `run.halted`, not `run.completed`) cannot
+       * finalize a truncated answer as a natural finish.
+       */
+      if (this._haltedReason == null && graph.preemptHaltReason != null) {
+        this._haltedReason = graph.preemptHaltReason;
+      } else if (this._haltedReason == null && graph.preemptIncomplete) {
+        this._haltedReason = 'preempt_incomplete';
+      }
     };
 
     try {

package/src/scripts/preempt-probe.ts

@@ -0,0 +1,330 @@
+// src/scripts/preempt-probe.ts
+//
+// Live-provider probe for cooperative mid-generation preemption.
+//
+//   npx tsx --env-file=.env src/scripts/preempt-probe.ts --provider anthropic
+//
+// (tsx rather than the repo's node-loader `script` runner: the import chain
+// reaches @mistralai/mistralai, which ships ESM-only and defeats the loader.
+// Keep live credentials OUT of the worktree's .env — jest loads it — and
+// point --env-file at wherever they actually live.)
+//
+// Asks for a long, tool-free answer, arms `shouldPreempt` once text has
+// started streaming, and asserts the full seal -> PreemptBoundary -> inject ->
+// self-loop path against a real provider: the stream stops mid-answer, the
+// injected user turn lands between two assistant turns, generation resumes in
+// the same run, and CHAT_MODEL_END fires for the sealed turn so usage is still
+// recorded. Prints a JSON verdict block as its last line.
+import { config } from 'dotenv';
+config();
+import { v4 as uuidv4 } from 'uuid';
+import { HumanMessage } from '@langchain/core/messages';
+import type { BaseMessage, UsageMetadata } from '@langchain/core/messages';
+import type * as t from '@/types';
+import { GraphEvents, Providers } from '@/common';
+import { createTokenCounter } from '@/utils/tokens';
+import { getLLMConfig } from '@/utils/llmConfig';
+import { HookRegistry } from '@/hooks';
+import { Run } from '@/run';
+
+/** Deltas to let through before arming the seal. */
+const ARM_AFTER_DELTAS = 15;
+
+const PROMPT =
+  'Write a detailed history of the Byzantine Empire from Constantine to ' +
+  '1453. Cover the major emperors, the religious schisms, and the military ' +
+  'campaigns. Aim for at least 800 words of flowing prose. Do not use ' +
+  'headings or bullet points.';
+
+const STEER =
+  'Stop. Forget the essay. Answer only this, in one short sentence: what ' +
+  'year did Constantinople fall?';
+
+type Verdict = {
+  provider: string;
+  model: string;
+  ok: boolean;
+  sealed: boolean;
+  seals: number;
+  emptyBoundaries: number;
+  boundaryFired: boolean;
+  boundarySealCount: number | null;
+  modelEndEvents: number;
+  usageEventsWithTokens: number;
+  chatModelStarts: number;
+  llmEnds: number;
+  llmErrors: number;
+  runsLeftOpen: number;
+  sealedTextChars: number;
+  resumedTextChars: number;
+  sequence: string[];
+  injectedFound: boolean;
+  resumeAnswersSteer: boolean;
+  error: string | null;
+};
+
+/**
+ * Streamed assistant turns arrive as `AIMessageChunk`, which extends
+ * `BaseMessageChunk` rather than `AIMessage` — `instanceof AIMessage` is
+ * false for them. Type discrimination has to go through `getType()`.
+ */
+function isAssistant(message: BaseMessage | undefined): boolean {
+  return message?.getType() === 'ai';
+}
+
+function messageKind(message: BaseMessage): string {
+  if (message.getType() === 'human') {
+    const source = message.additional_kwargs.source;
+    return typeof source === 'string' ? `human(${source})` : 'human';
+  }
+  return message.getType();
+}
+
+function textOf(message: BaseMessage | undefined): string {
+  if (message == null) {
+    return '';
+  }
+  if (typeof message.content === 'string') {
+    return message.content;
+  }
+  let text = '';
+  for (const block of message.content) {
+    if (block.type === 'text' && typeof block.text === 'string') {
+      text += block.text;
+    }
+  }
+  return text;
+}
+
+function parseProvider(): string {
+  const index = process.argv.indexOf('--provider');
+  if (index !== -1 && process.argv[index + 1] != null) {
+    return process.argv[index + 1];
+  }
+  return Providers.ANTHROPIC;
+}
+
+/** Control mode: identical run with no `preemption` config at all. */
+function isControl(): boolean {
+  return process.argv.includes('--control');
+}
+
+async function probe(): Promise<Verdict> {
+  const providerKey = parseProvider();
+  const llmConfig = getLLMConfig(providerKey);
+  if (llmConfig == null) {
+    throw new Error(`No llmConfig entry for provider "${providerKey}"`);
+  }
+
+  let armed = false;
+  let armedOnce = false;
+  let deltaCount = 0;
+  let modelEndEvents = 0;
+  let usageEventsWithTokens = 0;
+  let chatModelStarts = 0;
+  let llmEnds = 0;
+  let llmErrors = 0;
+  let boundaryFired = false;
+  let boundarySealCount: number | null = null;
+  let injectedOnce = false;
+  const collectedUsage: UsageMetadata[] = [];
+
+  /**
+   * Deliberately NOT a `ChatModelStreamHandler` instance: that would make
+   * `getRegisteredDefaultChatStreamHandler` return a handler and route the run
+   * down the registered-handler loop, which never seals. Plain-object handlers
+   * mirror how LibreChat registers, which is the sealable path.
+   */
+  const customHandlers: Record<string, t.EventHandler> = {
+    [GraphEvents.ON_MESSAGE_DELTA]: {
+      handle: (): void => {
+        deltaCount += 1;
+        if (!armedOnce && deltaCount >= ARM_AFTER_DELTAS) {
+          armed = true;
+          armedOnce = true;
+        }
+      },
+    },
+    [GraphEvents.CHAT_MODEL_END]: {
+      handle: (_event: string, data: t.ModelEndData): void => {
+        modelEndEvents += 1;
+        const usage = data?.output?.usage_metadata;
+        if (usage != null) {
+          collectedUsage.push(usage);
+          if ((usage.output_tokens ?? 0) > 0) {
+            usageEventsWithTokens += 1;
+          }
+        }
+      },
+    },
+  };
+
+  const hooks = new HookRegistry();
+  hooks.register('PreemptBoundary', {
+    hooks: [
+      async (input) => {
+        boundaryFired = true;
+        boundarySealCount = input.sealCount;
+        /**
+         * Disarm on drain. The SDK polls `shouldPreempt` on every chunk and
+         * does not clear it for you — a host that leaves the request set
+         * armed re-seals the resumed turn immediately. LibreChat's real drain
+         * does this via `noteSteersRemoved`.
+         */
+        armed = false;
+        if (injectedOnce) {
+          return {};
+        }
+        injectedOnce = true;
+        return {
+          injectedMessages: [
+            { role: 'user', content: STEER, source: 'steer' },
+          ],
+        };
+      },
+    ],
+  });
+
+  /**
+   * A real host configures this. It is what `dispatchSealedModelEnd` falls
+   * back to when the provider never got to send its usage chunk, so running
+   * without one hides whether sealed-turn usage is recoverable. Pass
+   * `--no-token-counter` to probe the unconfigured case deliberately.
+   */
+  const tokenCounter = process.argv.includes('--no-token-counter')
+    ? undefined
+    : await createTokenCounter();
+
+  const run = await Run.create<t.IState>({
+    runId: uuidv4(),
+    graphConfig: {
+      type: 'standard',
+      llmConfig,
+      instructions: 'You are a knowledgeable history assistant.',
+    },
+    tokenCounter,
+    customHandlers,
+    hooks,
+    ...(isControl()
+      ? {}
+      : {
+        preemption: {
+          shouldPreempt: (): boolean => armed,
+          maxSeals: 2,
+        },
+      }),
+    returnContent: true,
+    skipCleanup: true,
+  });
+
+  /**
+   * Counts the RAW LangChain model-run lifecycle, which is what a tracer
+   * (LangSmith, Langfuse) sees. A seal that fails to close its run shows up
+   * here as starts > ends — an span that would hang open forever.
+   */
+  const streamConfig = {
+    runId: uuidv4(),
+    configurable: { user_id: 'probe-user', thread_id: 'preempt-probe' },
+    streamMode: 'values',
+    version: 'v2' as const,
+    callbacks: [
+      {
+        handleChatModelStart: (): void => {
+          chatModelStarts += 1;
+        },
+        handleLLMEnd: (): void => {
+          llmEnds += 1;
+        },
+        handleLLMError: (): void => {
+          llmErrors += 1;
+        },
+      },
+    ],
+  };
+
+  await run.processStream(
+    { messages: [new HumanMessage(PROMPT)] },
+    streamConfig
+  );
+
+  const messages = run.getRunMessages() ?? [];
+  const sequence = messages.map(messageKind);
+  const injectedIndex = messages.findIndex(
+    (message) =>
+      message instanceof HumanMessage &&
+      message.additional_kwargs.source === 'steer'
+  );
+  const stats = run.getPreemptStats();
+  const sealedText = textOf(messages[injectedIndex - 1]);
+  const resumedText = textOf(messages[injectedIndex + 1]);
+  const resumeAnswersSteer = /1453/.test(resumedText);
+
+  /**
+   * Control mode is the no-preemption baseline: one uninterrupted model run,
+   * nothing sealed, nothing injected. Its value is the contrast — the same
+   * lifecycle counters must balance WITHOUT the seal machinery — so it gets
+   * its own criteria instead of failing the seal-path ones by construction.
+   */
+  const ok = isControl()
+    ? stats.seals === 0 &&
+      !boundaryFired &&
+      injectedIndex === -1 &&
+      messages.length > 0 &&
+      textOf(messages[messages.length - 1]).trim().length > 0 &&
+      modelEndEvents >= 1 &&
+      usageEventsWithTokens >= 1 &&
+      llmErrors === 0 &&
+      chatModelStarts === llmEnds
+    : stats.seals === 1 &&
+      stats.emptyBoundaries === 0 &&
+      boundaryFired &&
+      injectedIndex > 0 &&
+      isAssistant(messages[injectedIndex - 1]) &&
+      isAssistant(messages[injectedIndex + 1]) &&
+      sealedText.trim().length > 0 &&
+      resumedText.trim().length > 0 &&
+      modelEndEvents >= 2 &&
+      usageEventsWithTokens >= 2 &&
+      chatModelStarts === llmEnds;
+
+  return {
+    provider: providerKey,
+    model: String((llmConfig as { model?: unknown }).model ?? 'unknown'),
+    ok,
+    sealed: stats.seals > 0,
+    seals: stats.seals,
+    emptyBoundaries: stats.emptyBoundaries,
+    boundaryFired,
+    boundarySealCount,
+    modelEndEvents,
+    usageEventsWithTokens,
+    chatModelStarts,
+    llmEnds,
+    llmErrors,
+    runsLeftOpen: chatModelStarts - llmEnds,
+    sealedTextChars: sealedText.length,
+    resumedTextChars: resumedText.length,
+    sequence,
+    injectedFound: injectedIndex >= 0,
+    resumeAnswersSteer,
+    error: null,
+  };
+}
+
+probe()
+  .then((verdict) => {
+    console.log('\n===== SEALED TAIL / RESUMED HEAD =====');
+    console.log(JSON.stringify(verdict, null, 2));
+    console.log(`\nVERDICT_JSON ${JSON.stringify(verdict)}`);
+    process.exit(verdict.ok ? 0 : 1);
+  })
+  .catch((error: unknown) => {
+    const verdict: Partial<Verdict> = {
+      provider: parseProvider(),
+      ok: false,
+      error: error instanceof Error ? error.message : String(error),
+    };
+    console.error(error);
+    console.log(`\nVERDICT_JSON ${JSON.stringify(verdict)}`);
+    process.exit(1);
+  });