@vellumai/assistant 0.8.6 → 0.8.7

package/src/runtime/__tests__/agent-wake.test.ts

@@ -117,7 +117,11 @@ mock.module("../../memory/llm-request-log-store.js", () => ({
   backfillMessageIdOnLogs: () => {},
 }));
 
-import type { AgentEvent } from "../../agent/loop.js";
+import type {
+  AgentEvent,
+  AgentLoopRunOptions,
+  AgentLoopRunResult,
+} from "../../agent/loop.js";
 import type { Message } from "../../providers/types.js";
 import {
   __resetWakeChainForTests,
@@ -127,6 +131,13 @@ import {
 
 // ── Test helpers ─────────────────────────────────────────────────────
 
+// Wake runs never pause at a checkpoint — their onCheckpoint always returns
+// "continue" — so the loop result always carries a null pause-reason.
+const runResult = (history: Message[]): AgentLoopRunResult => ({
+  history,
+  exitReason: null,
+});
+
 interface MockTarget extends WakeTarget {
   emittedEvents: AgentEvent[];
   pushedMessages: Message[];
@@ -134,6 +145,7 @@ interface MockTarget extends WakeTarget {
     input: Message[];
     requestId?: string;
     turnContext?: unknown;
+    allowedTools?: string[];
   }>;
   processingToggles: boolean[];
   /** Tail messages handed to `persistTailMessage`, in call order. */
@@ -152,6 +164,11 @@ interface MockTarget extends WakeTarget {
    * just inferring it from the order of recorded toggles.
    */
   processingDuringDrain: boolean[];
+  /**
+   * Tool allowlist snapshots captured whenever the wake applies/restores a
+   * scope. `undefined` means unrestricted.
+   */
+  allowedToolSnapshots: Array<string[] | undefined>;
 }
 
 function makeTarget(options: {
@@ -164,6 +181,7 @@ function makeTarget(options: {
   isProcessing?: boolean;
   /** When true, omit `drainQueue` so we can verify the wake handles its absence. */
   omitDrainQueue?: boolean;
+  initialAllowedTools?: Set<string>;
 }): MockTarget {
   const emittedEvents: AgentEvent[] = [];
   const pushedMessages: Message[] = [];
@@ -171,14 +189,19 @@ function makeTarget(options: {
     input: Message[];
     requestId?: string;
     turnContext?: unknown;
+    allowedTools?: string[];
   }> = [];
   const processingToggles: boolean[] = [];
   const persistedTailCalls: Message[] = [];
   const callSequence: string[] = [];
   const processingDuringDrain: boolean[] = [];
+  const allowedToolSnapshots: Array<string[] | undefined> = [];
   const history: Message[] = [...(options.baseline ?? [])];
   let processing = options.isProcessing ?? false;
   let drainQueueCalls = 0;
+  let activeAllowedTools = options.initialAllowedTools;
+  const snapshotAllowedTools = (): string[] | undefined =>
+    activeAllowedTools ? [...activeAllowedTools].sort() : undefined;
 
   const target: MockTarget = {
     conversationId: options.conversationId ?? "conv-test",
@@ -189,6 +212,7 @@ function makeTarget(options: {
     persistedTailCalls,
     callSequence,
     processingDuringDrain,
+    allowedToolSnapshots,
     get drainQueueCalls() {
       return drainQueueCalls;
     },
@@ -196,13 +220,14 @@ function makeTarget(options: {
       run: async (
         input: Message[],
         onEvent: (event: AgentEvent) => void | Promise<void>,
-        _signal?: AbortSignal,
-        requestId?: string,
-        _onCheckpoint?: unknown,
-        _callSite?: unknown,
-        turnContext?: unknown,
+        runOptions?: AgentLoopRunOptions,
       ) => {
-        runCalls.push({ input: [...input], requestId, turnContext });
+        runCalls.push({
+          input: [...input],
+          requestId: runOptions?.requestId,
+          turnContext: runOptions?.turnContext,
+          allowedTools: snapshotAllowedTools(),
+        });
         // Emit any scripted events the test wanted us to produce.
         for (const ev of options.scriptedEvents ?? []) {
           await onEvent(ev);
@@ -221,7 +246,7 @@ function makeTarget(options: {
             next.push(tailMsg);
           }
         }
-        return next;
+        return runResult(next);
       },
     },
     getMessages: () => history,
@@ -243,6 +268,19 @@ function makeTarget(options: {
       persistedTailCalls.push(msg);
       callSequence.push("persist");
     },
+    setWakeAllowedTools: (tools: ReadonlySet<string>) => {
+      const previous = activeAllowedTools;
+      activeAllowedTools = new Set(tools);
+      allowedToolSnapshots.push(snapshotAllowedTools());
+      callSequence.push(`tools:${snapshotAllowedTools()?.join(",") ?? "all"}`);
+      return () => {
+        activeAllowedTools = previous;
+        allowedToolSnapshots.push(snapshotAllowedTools());
+        callSequence.push(
+          `tools:${snapshotAllowedTools()?.join(",") ?? "all"}`,
+        );
+      };
+    },
     ...(options.omitDrainQueue
       ? {}
       : {
@@ -491,6 +529,69 @@ describe("wakeAgentForOpportunity", () => {
     });
   });
 
+  test("scopes allowed tools during the wake and restores before queued messages drain", async () => {
+    const target = makeTarget({
+      initialAllowedTools: new Set(["bash"]),
+      scriptedAssistant: {
+        role: "assistant",
+        content: [{ type: "text", text: "Saved." }],
+      },
+    });
+
+    const result = await wakeAgentForOpportunity(
+      {
+        conversationId: target.conversationId,
+        hint: "review for memories",
+        source: "memory-retrospective",
+        allowedTools: ["remember"],
+      },
+      { resolveTarget: async () => target },
+    );
+
+    expect(result).toEqual({ invoked: true, producedToolCalls: false });
+    expect(target.runCalls[0]!.allowedTools).toEqual(["remember"]);
+    expect(target.allowedToolSnapshots).toEqual([["remember"], ["bash"]]);
+
+    const restoreIndex = target.callSequence.indexOf("tools:bash");
+    const processingFalseIndex =
+      target.callSequence.indexOf("processing:false");
+    const drainIndex = target.callSequence.indexOf("drain");
+    expect(restoreIndex).toBeGreaterThan(-1);
+    expect(restoreIndex).toBeLessThan(processingFalseIndex);
+    expect(processingFalseIndex).toBeLessThan(drainIndex);
+  });
+
+  test("restores allowed tools before drain when the wake is a silent no-op", async () => {
+    const target = makeTarget({
+      scriptedAssistant: {
+        role: "assistant",
+        content: [{ type: "text", text: "" }],
+      },
+    });
+
+    const result = await wakeAgentForOpportunity(
+      {
+        conversationId: target.conversationId,
+        hint: "review for memories",
+        source: "memory-retrospective",
+        allowedTools: ["remember"],
+      },
+      { resolveTarget: async () => target },
+    );
+
+    expect(result).toEqual({ invoked: true, producedToolCalls: false });
+    expect(target.runCalls[0]!.allowedTools).toEqual(["remember"]);
+    expect(target.allowedToolSnapshots).toEqual([["remember"], undefined]);
+
+    const restoreIndex = target.callSequence.indexOf("tools:all");
+    const processingFalseIndex =
+      target.callSequence.indexOf("processing:false");
+    const drainIndex = target.callSequence.indexOf("drain");
+    expect(restoreIndex).toBeGreaterThan(-1);
+    expect(restoreIndex).toBeLessThan(processingFalseIndex);
+    expect(processingFalseIndex).toBeLessThan(drainIndex);
+  });
+
   test("produces tool calls when LLM emits a tool_use block", async () => {
     const assistantMessage: Message = {
       role: "assistant",
@@ -605,9 +706,13 @@ describe("wakeAgentForOpportunity", () => {
     // hold the processing flag while agentLoop.run executes.
     const observedDuringRun: boolean[] = [];
     const originalRun = target.agentLoop.run;
-    target.agentLoop.run = async (input, onEvent, signal, requestId) => {
+    target.agentLoop.run = async (
+      input: Message[],
+      onEvent: (event: AgentEvent) => void | Promise<void>,
+      runOptions?: AgentLoopRunOptions,
+    ) => {
       observedDuringRun.push(target.isProcessing());
-      return originalRun(input, onEvent, signal, requestId);
+      return originalRun(input, onEvent, runOptions);
     };
 
     await wakeAgentForOpportunity(
@@ -677,7 +782,7 @@ describe("wakeAgentForOpportunity", () => {
       agentLoop: {
         run: async (input) => {
           runCalls.push(callOrder++);
-          return [...input];
+          return runResult([...input]);
         },
       },
       getMessages: () => history,
@@ -720,7 +825,7 @@ describe("wakeAgentForOpportunity", () => {
     const target: WakeTarget = {
       conversationId: "conv-no-trust",
       agentLoop: {
-        run: async (input) => [...input],
+        run: async (input) => runResult([...input]),
       },
       getMessages: () => history,
       pushMessage: () => {},
@@ -772,7 +877,7 @@ describe("wakeAgentForOpportunity", () => {
             await gate2.promise;
           }
           runCompleteOrder.push(myIndex);
-          return input; // no assistant message → silent no-op
+          return runResult(input); // no assistant message → silent no-op
         },
       },
       getMessages: () => history,
@@ -820,7 +925,7 @@ describe("wakeAgentForOpportunity", () => {
     const target: WakeTarget & { setProcessing: (v: boolean) => void } = {
       conversationId: "conv-user-turn",
       agentLoop: {
-        run: async (input) => input,
+        run: async (input) => runResult(input),
       },
       getMessages: () => history,
       pushMessage: (msg) => {
@@ -886,7 +991,7 @@ describe("wakeAgentForOpportunity", () => {
     const history: Message[] = [];
     const target: WakeTarget = {
       conversationId: "conv-busy",
-      agentLoop: { run: async () => history },
+      agentLoop: { run: async () => runResult(history) },
       getMessages: () => history,
       pushMessage: () => {},
       emitAgentEvent: () => {},
@@ -1296,7 +1401,7 @@ describe("wakeAgentForOpportunity", () => {
       const target: WakeTarget = {
         conversationId: "conv-stream",
         agentLoop: {
-          run: async (_input, onEvent, _signal, _requestId, onCheckpoint) => {
+          run: async (_input, onEvent, runOptions) => {
             // Preamble + assistant hint + postamble (mirrors what the
             // wake injects). The agent-wake helper expects these three
             // hint messages in the input it hands to run().
@@ -1311,7 +1416,7 @@ describe("wakeAgentForOpportunity", () => {
               message: turn1Assistant,
             });
             runHistory.push(turn1ToolResult);
-            const dec1 = await onCheckpoint!({
+            const dec1 = await runOptions!.onCheckpoint!({
               turnIndex: 0,
               toolCount: 1,
               hasToolUse: true,
@@ -1328,7 +1433,7 @@ describe("wakeAgentForOpportunity", () => {
               message: turn2Assistant,
             });
             runHistory.push(turn2ToolResult);
-            const dec2 = await onCheckpoint!({
+            const dec2 = await runOptions!.onCheckpoint!({
               turnIndex: 1,
               toolCount: 1,
               hasToolUse: true,
@@ -1345,7 +1450,7 @@ describe("wakeAgentForOpportunity", () => {
               type: "message_complete",
               message: finalAssistant,
             });
-            return runHistory;
+            return runResult(runHistory);
           },
         },
         getMessages: () => history,
@@ -1432,17 +1537,17 @@ describe("wakeAgentForOpportunity", () => {
       const target: WakeTarget = {
         conversationId: "conv-card",
         agentLoop: {
-          run: async (_input, _onEvent, _signal, _requestId, onCheckpoint) => {
+          run: async (_input, _onEvent, runOptions) => {
             const runHistory: Message[] = [..._input];
             runHistory.push(firstAssistant);
             runHistory.push(toolResult);
-            await onCheckpoint!({
+            await runOptions!.onCheckpoint!({
               turnIndex: 0,
               toolCount: 1,
               hasToolUse: true,
               history: runHistory,
             });
-            return runHistory;
+            return runResult(runHistory);
           },
         },
         getMessages: () => history,
@@ -1627,7 +1732,7 @@ describe("wakeAgentForOpportunity", () => {
     let processing = false;
     return {
       conversationId,
-      agentLoop: { run: async (input) => input },
+      agentLoop: { run: async (input) => runResult(input) },
       getMessages: () => history,
       pushMessage: () => {},
       emitAgentEvent: () => {},
@@ -1704,17 +1809,17 @@ describe("wakeAgentForOpportunity", () => {
       const target: WakeTarget = {
         conversationId: "conv-suppress-surface",
         agentLoop: {
-          run: async (_input, _onEvent, _signal, _requestId, onCheckpoint) => {
+          run: async (_input, _onEvent, runOptions) => {
             const runHistory: Message[] = [..._input];
             runHistory.push(firstAssistant);
             runHistory.push(toolResult);
-            await onCheckpoint!({
+            await runOptions!.onCheckpoint!({
               turnIndex: 0,
               toolCount: 1,
               hasToolUse: true,
               history: runHistory,
             });
-            return runHistory;
+            return runResult(runHistory);
           },
         },
         getMessages: () => history,

package/src/runtime/__tests__/background-job-runner.test.ts

@@ -46,7 +46,6 @@ mock.module("../../memory/conversation-crud.js", () => ({
 let processMessageImpl: (
   conversationId: string,
   content: string,
-  attachmentIds: string[] | undefined,
   options: Record<string, unknown> | undefined,
 ) => Promise<{ messageId: string }> = async () => ({ messageId: "msg-1" });
 const processMessageCalls: Array<{
@@ -59,11 +58,10 @@ mock.module("../../daemon/process-message.js", () => ({
   processMessage: async (
     conversationId: string,
     content: string,
-    attachmentIds: string[] | undefined,
     options: Record<string, unknown> | undefined,
   ) => {
     processMessageCalls.push({ conversationId, content, options });
-    return processMessageImpl(conversationId, content, attachmentIds, options);
+    return processMessageImpl(conversationId, content, options);
   },
 }));
 

package/src/runtime/agent-wake.ts

@@ -14,7 +14,8 @@
  *     trailing user message satisfies providers that reject assistant
  *     prefill. The bookend user messages are hardcoded strings with no
  *     dynamic content, so they cannot carry injection payloads.
- *   - Invokes the agent loop with all conversation tools available.
+ *   - Invokes the agent loop with all conversation tools available unless
+ *     the caller provides an explicit `allowedTools` scope.
  *   - No tool calls AND no assistant text → silent no-op (nothing persisted,
  *     nothing emitted). Returns `{ invoked: true, producedToolCalls: false }`.
  *   - Tool calls produced → normal tool execution runs (the conversation's
@@ -164,6 +165,12 @@ export interface WakeTarget {
    * `trustContext` through the wake.
    */
   setTrustContext?(ctx: TrustContext): void;
+  /**
+   * Temporarily restrict the tools visible/executable during this wake.
+   * Returns a restore callback so the wake can release the scope before any
+   * queued user turn is drained.
+   */
+  setWakeAllowedTools?(tools: ReadonlySet<string>): () => void;
 }
 
 export interface WakeOptions {
@@ -238,6 +245,12 @@ export interface WakeOptions {
    * retrospectives whose conversation title already says "(Retrospective)").
    */
   suppressWakeSurface?: boolean;
+  /**
+   * Optional exact tool allowlist for this wake. Used by internal maintenance
+   * jobs that need the assistant's judgment but must not execute arbitrary
+   * side-effect tools.
+   */
+  allowedTools?: readonly string[];
 }
 
 /**
@@ -834,11 +847,55 @@ export async function wakeAgentForOpportunity(
       overrideProfile,
     });
 
+    let wakeToolScopeRestored = false;
+    let restoreWakeToolScope: (() => void) | null = null;
+    const restoreWakeAllowedTools = (): void => {
+      if (wakeToolScopeRestored) return;
+      wakeToolScopeRestored = true;
+      if (!restoreWakeToolScope) return;
+      try {
+        restoreWakeToolScope();
+      } catch (err) {
+        log.warn(
+          { conversationId, source, err },
+          "agent-wake: failed to restore tool allowlist; continuing",
+        );
+      }
+    };
+    const applyWakeAllowedTools = (): boolean => {
+      if (!opts.allowedTools) return true;
+      if (!target.setWakeAllowedTools) {
+        log.warn(
+          { conversationId, source, allowedTools: opts.allowedTools },
+          "agent-wake: target cannot apply requested tool allowlist; skipping",
+        );
+        return false;
+      }
+      try {
+        restoreWakeToolScope = target.setWakeAllowedTools(
+          new Set(opts.allowedTools),
+        );
+        return true;
+      } catch (err) {
+        log.warn(
+          { conversationId, source, allowedTools: opts.allowedTools, err },
+          "agent-wake: failed to apply requested tool allowlist; skipping",
+        );
+        return false;
+      }
+    };
+
     // Mark processing for the duration of the run so a concurrent user
     // send is queued by `enqueueMessage()` rather than spawning a second
     // concurrent agent loop on the same conversation (which would
-    // interleave writes to `conversation.messages`).
-    target.markProcessing(true);
+    // interleave writes to `conversation.messages`). This happens before
+    // applying a wake-scoped tool allowlist so a concurrent user turn cannot
+    // start under the wake's restricted tool set.
+    try {
+      target.markProcessing(true);
+    } catch (err) {
+      throw err;
+    }
 
     // Fires after each tool-execution turn finalizes (assistant message
     // + matching tool_result user message both in history). A single
@@ -860,25 +917,41 @@ export async function wakeAgentForOpportunity(
     let tailMessageCount = 0;
     let drainedInTry = false;
     try {
+      if (!applyWakeAllowedTools()) {
+        return {
+          invoked: false,
+          producedToolCalls: false,
+          reason: "no_resolver" as const,
+        };
+      }
+
       let updatedHistory: Message[];
       try {
-        updatedHistory = await target.agentLoop.run(
+        ({ history: updatedHistory } = await target.agentLoop.run(
           runInput,
           onEvent,
-          undefined, // no external abort signal
-          `wake:${source}`,
-          onCheckpoint,
-          // Route through the caller-supplied call site (defaults to
-          // `mainAgent` so a normal user-turn wake shares the user's chat
-          // selection). Without an explicit callSite, the resolver in
-          // `RetryProvider` and the routing in `CallSiteRoutingProvider`
-          // short-circuit and silently drop both per-callsite config and the
-          // pinned `overrideProfile` below.
-          callSite,
-          wakeTurnContext,
-          overrideProfile,
-          effectiveContextWindow.maxInputTokens,
-        );
+          {
+            requestId: `wake:${source}`,
+            onCheckpoint,
+            // Route through the caller-supplied call site (defaults to
+            // `mainAgent` so a normal user-turn wake shares the user's chat
+            // selection). Without an explicit callSite, the resolver in
+            // `RetryProvider` and the routing in `CallSiteRoutingProvider`
+            // short-circuit and silently drop both per-callsite config and the
+            // pinned `overrideProfile` below.
+            callSite,
+            turnContext: wakeTurnContext,
+            overrideProfile,
+            // Wake runs have no orchestrator-side mid-loop compaction path,
+            // so the budget gate stays disabled (`overflowRecovery.enabled =
+            // false`); `maxInputTokens` is still supplied for tool-result
+            // truncation.
+            resolveContextWindow: () => ({
+              maxInputTokens: effectiveContextWindow.maxInputTokens,
+              overflowRecovery: { enabled: false, safetyMarginRatio: 0 },
+            }),
+          },
+        ));
       } catch (err) {
         // Capture the error for post-finally logging, then short-circuit
         // the rest of the try body — no tail to push/persist when the
@@ -935,6 +1008,7 @@ export async function wakeAgentForOpportunity(
       // accepts entries while processing === true, and drain expects
       // processing to already be false). The finally block handles the
       // error/early-return paths where no tail was produced.
+      restoreWakeAllowedTools();
       try {
         target.markProcessing(false);
       } catch (err) {
@@ -963,6 +1037,7 @@ export async function wakeAgentForOpportunity(
       // exit the try body before reaching the drain block, so
       // `drainedInTry` is still false.
       if (!drainedInTry) {
+        restoreWakeAllowedTools();
         try {
           target.markProcessing(false);
         } catch (err) {

package/src/runtime/assistant-event-hub.ts

@@ -43,6 +43,7 @@ import { appendEventToStream } from "../signals/event-stream.js";
 import { getLogger } from "../util/logger.js";
 import type { AssistantEvent } from "./assistant-event.js";
 import { buildAssistantEvent } from "./assistant-event.js";
+import { stampAndBuffer } from "./conversation-stream-state.js";
 
 const log = getLogger("assistant-event-hub");
 
@@ -604,6 +605,7 @@ export function broadcastMessage(
           excludeClientId,
         }
       : undefined;
+  stampAndBuffer(event, { targeting: publishOptions });
   _hubChain = _hubChain
     .then(() => assistantEventHub.publish(event, publishOptions))
     .then(() => {
@@ -658,7 +660,6 @@ function resolveCanonicalRequestSourceType(
   return "channel";
 }
 
-
 /**
  * Lazily load heavy dependencies and create a canonical guardian request +
  * bridge for a confirmation_request message. Called fire-and-forget from
@@ -732,4 +733,3 @@ async function createCanonicalRequestForConfirmation(
     );
   }
 }
-