@vellumai/assistant 0.12.2 → 0.12.3-staging.1

package/src/__tests__/conversation-fork-crud.test.ts

@@ -5,6 +5,7 @@ import { beforeEach, describe, expect, test } from "bun:test";
 import { eq, like } from "drizzle-orm";
 
 import {
+  getAttachmentContent,
   getAttachmentsForMessage,
   linkAttachmentToMessage,
   uploadAttachment,
@@ -17,10 +18,17 @@ import {
 import {
   addMessage,
   createConversation,
+  deleteConversation,
   forkConversation,
   getMessages,
+  listConversationAttachments,
+  updateMessageMetadata,
 } from "../persistence/conversation-crud.js";
 import { getConversationDirPath } from "../persistence/conversation-disk-view.js";
+import {
+  COMPUTER_USE_SCREENSHOT_ATTACHMENT_IDS_KEY,
+  computerUseScreenshotAttachmentIdsFromMetadata,
+} from "../persistence/conversation-types.js";
 import {
   getDb,
   getLogsDb,
@@ -63,6 +71,7 @@ import {
   recordInjected as recordV3Injected,
 } from "../plugins/defaults/memory/v3/ever-injected-store.js";
 import { ensureMemoryV3InjectedSectionsSchema } from "../plugins/defaults/memory/v3/plugin-schema.js";
+import { handleListMessages } from "../runtime/routes/conversation-routes.js";
 
 await initializeDb();
 
@@ -748,6 +757,188 @@ describe("forkConversation", () => {
     );
   });
 
+  test("remaps automatic screenshot refs and widens fork provenance", async () => {
+    const screenshotBase64 =
+      "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8/5+hHgAHggJ/PchI7wAAAABJRU5ErkJggg==";
+    const source = createConversation("Computer use thread");
+    await addMessage(source.id, "user", "Open the example", {
+      skipIndexing: true,
+    });
+    await addMessage(
+      source.id,
+      "assistant",
+      JSON.stringify([
+        {
+          type: "tool_use",
+          id: "computer-use-call",
+          name: "computer_use_click",
+          input: { x: 10, y: 20 },
+        },
+      ]),
+      { skipIndexing: true },
+    );
+    const screenshot = await uploadAttachment(
+      "computer-use-click.png",
+      "image/png",
+      screenshotBase64,
+    );
+    const toolResult = await addMessage(
+      source.id,
+      "user",
+      JSON.stringify([
+        {
+          type: "tool_result",
+          tool_use_id: "computer-use-call",
+          content: "Clicked",
+          contentBlocks: [
+            {
+              type: "image",
+              source: {
+                type: "workspace_ref",
+                media_type: "image/png",
+                attachmentId: screenshot.id,
+                sizeBytes: Buffer.from(screenshotBase64, "base64").byteLength,
+              },
+            },
+          ],
+        },
+      ]),
+      { skipIndexing: true },
+    );
+    linkAttachmentToMessage(toolResult.id, screenshot.id, 0);
+
+    const reply = await addMessage(
+      source.id,
+      "assistant",
+      JSON.stringify([{ type: "text", text: "Done." }]),
+      { skipIndexing: true },
+    );
+    linkAttachmentToMessage(reply.id, screenshot.id, 0);
+    const explicit = await uploadAttachment(
+      "report.pdf",
+      "application/pdf",
+      "JVBERg==",
+    );
+    linkAttachmentToMessage(reply.id, explicit.id, 1);
+    updateMessageMetadata(reply.id, {
+      [COMPUTER_USE_SCREENSHOT_ATTACHMENT_IDS_KEY]: [screenshot.id],
+    });
+
+    const fork = forkConversation({ conversationId: source.id });
+    const forkRows = getMessages(fork.id);
+    const forkToolResult = forkRows.find(
+      (row) =>
+        row.role === "user" &&
+        JSON.stringify(row.content).includes("tool_result"),
+    );
+    const forkReply = forkRows.find(
+      (row) =>
+        row.role === "assistant" &&
+        JSON.stringify(row.content).includes("Done."),
+    );
+    expect(forkToolResult).toBeDefined();
+    expect(forkReply).toBeDefined();
+
+    const forkReplyAttachments = getAttachmentsForMessage(forkReply!.id);
+    const clonedScreenshot = forkReplyAttachments.find(
+      (attachment) => attachment.originalFilename === "computer-use-click.png",
+    );
+    const clonedExplicit = forkReplyAttachments.find(
+      (attachment) => attachment.originalFilename === "report.pdf",
+    );
+    expect(clonedScreenshot?.id).toBeDefined();
+    expect(clonedScreenshot?.id).not.toBe(screenshot.id);
+    expect(clonedExplicit?.id).toBeDefined();
+    expect(JSON.stringify(forkToolResult!.content)).toContain(
+      `"attachmentId":"${clonedScreenshot!.id}"`,
+    );
+    expect(JSON.stringify(forkToolResult!.content)).not.toContain(
+      `"attachmentId":"${screenshot.id}"`,
+    );
+    const persistedSourceToolResult = getMessages(source.id).find(
+      (row) => row.id === toolResult.id,
+    );
+    expect(JSON.stringify(persistedSourceToolResult?.content)).toContain(
+      `"attachmentId":"${screenshot.id}"`,
+    );
+    expect(JSON.stringify(persistedSourceToolResult?.content)).not.toContain(
+      `"attachmentId":"${clonedScreenshot!.id}"`,
+    );
+    expect(
+      getAttachmentsForMessage(forkToolResult!.id).map(
+        (attachment) => attachment.id,
+      ),
+    ).toEqual([clonedScreenshot!.id]);
+
+    deleteConversation(source.id);
+
+    const forkJsonl = readFileSync(
+      join(getConversationDirPath(fork.id, fork.createdAt), "messages.jsonl"),
+      "utf-8",
+    );
+    expect(forkJsonl).toContain(
+      '"toolResults":[{"content":"Clicked"}],"attachments":["computer-use-click.png"]',
+    );
+
+    const forkMarkerIds = computerUseScreenshotAttachmentIdsFromMetadata(
+      parseMetadata(forkReply!.metadata) as Record<string, unknown>,
+    );
+    expect(forkMarkerIds).toEqual([screenshot.id, clonedScreenshot!.id]);
+    expect(forkMarkerIds).not.toContain(clonedExplicit!.id);
+
+    const history = (await handleListMessages({
+      queryParams: { conversationId: fork.id },
+    })) as {
+      messages: Array<{
+        toolCalls?: Array<{ imageAttachmentIds?: string[] }>;
+        attachments: Array<{
+          id: string;
+          computerUseScreenshot?: boolean;
+        }>;
+      }>;
+    };
+    expect(
+      history.messages.flatMap((message) =>
+        (message.toolCalls ?? []).flatMap(
+          (toolCall) => toolCall.imageAttachmentIds ?? [],
+        ),
+      ),
+    ).toEqual([clonedScreenshot!.id]);
+    const hydratedScreenshot = history.messages
+      .flatMap((message) => message.attachments)
+      .find((attachment) => attachment.id === clonedScreenshot!.id);
+    const hydratedExplicit = history.messages
+      .flatMap((message) => message.attachments)
+      .find((attachment) => attachment.id === clonedExplicit!.id);
+    expect(hydratedScreenshot?.computerUseScreenshot).toBe(true);
+    expect(hydratedExplicit?.computerUseScreenshot).toBeUndefined();
+    expect(getAttachmentContent(screenshot.id)).toBeNull();
+    expect(getAttachmentContent(clonedScreenshot!.id)?.toString("base64")).toBe(
+      screenshotBase64,
+    );
+
+    const firstFilesPage = listConversationAttachments(fork.id, {
+      limit: 1,
+      offset: 0,
+    });
+    const secondFilesPage = listConversationAttachments(fork.id, {
+      limit: 1,
+      offset: 1,
+    });
+    expect(firstFilesPage.total).toBe(2);
+    expect(secondFilesPage.total).toBe(2);
+    expect(
+      [...firstFilesPage.attachments, ...secondFilesPage.attachments].map(
+        (attachment) => attachment.id,
+      ),
+    ).toEqual([clonedScreenshot!.id, clonedExplicit!.id]);
+    expect(
+      [...firstFilesPage.attachments, ...secondFilesPage.attachments].every(
+        (attachment) => attachment.messageId === forkReply!.id,
+      ),
+    ).toBe(true);
+  });
+
   test("inherits the source conversation's inference profile", async () => {
     const source = createConversation("Pinned profile thread");
     await addMessage(source.id, "user", "Use the balanced profile", {

package/src/__tests__/conversation-rendered-delegation-state.test.ts

@@ -0,0 +1,276 @@
+/**
+ * `Conversation.buildCurrentSystemPrompt` captures the delegation-section
+ * state the prompt it builds renders (`renderedDelegateIndependentTasks`),
+ * which the wire-surface recorder persists for a retrospective fork to
+ * replay. The loop sends the prompt built before the run and never rebuilds
+ * it mid-run, so the captured state has to be the one that build used: not a
+ * re-derivation at send time, and not the state of a build that failed while
+ * the loop kept its previous prompt.
+ */
+import { beforeEach, describe, expect, mock, test } from "bun:test";
+
+import { CompactionCircuit } from "../agent/compaction-circuit.js";
+import type { AgentEvent } from "../agent/loop.js";
+import type { AssistantEvent } from "../api/index.js";
+import type { Message, ProviderResponse } from "../providers/types.js";
+
+// ---------------------------------------------------------------------------
+// Mocks — must precede Conversation import
+// ---------------------------------------------------------------------------
+
+mock.module("../providers/registry.js", () => ({
+  getProvider: () => ({ name: "mock-provider" }),
+  initializeProviders: async () => {},
+}));
+
+mock.module("../config/assistant-feature-flags.js", () => ({
+  isAssistantFeatureFlagEnabled: () => true,
+}));
+
+// The builder under the conversation's control: records the options each
+// build receives and can be made to throw, which is what a wake's prompt
+// sync swallows before running on the previous prompt.
+let builtWith: Array<Record<string, unknown> | undefined> = [];
+let throwOnBuild = false;
+mock.module("../prompts/system-prompt.js", () => ({
+  buildSystemPrompt: (options?: Record<string, unknown>) => {
+    if (throwOnBuild) {
+      throw new Error("prompt build failed");
+    }
+    builtWith.push(options);
+    return "system prompt";
+  },
+}));
+
+mock.module("../config/skills.js", () => ({
+  loadSkillCatalog: () => [],
+  loadSkillBySelector: () => ({ skill: null }),
+  ensureSkillIcon: async () => null,
+}));
+
+mock.module("../config/skill-state.js", () => ({
+  resolveSkillStates: () => [],
+}));
+
+mock.module("../permissions/trust-store.js", () => ({
+  addRule: () => {},
+  findHighestPriorityRule: () => null,
+  clearCache: () => {},
+}));
+
+mock.module("../security/secret-allowlist.js", () => ({
+  resetAllowlist: () => {},
+}));
+
+mock.module("../persistence/conversation-crud.js", () => ({
+  setConversationProcessingStartedAt: () => {},
+  isConversationProcessing: () => false,
+  setConversationOriginChannelIfUnset: () => {},
+  updateConversationContextWindow: () => {},
+  deleteMessageById: () => {},
+  provenanceFromTrustContext: () => ({
+    source: "user",
+    trustContext: undefined,
+  }),
+  getConversationOriginInterface: () => null,
+  getConversationOriginChannel: () => null,
+  getMessages: () => [],
+  getConversation: () => ({
+    id: "conv-1",
+    contextSummary: null,
+    contextCompactedMessageCount: 0,
+    totalInputTokens: 0,
+    totalOutputTokens: 0,
+    totalEstimatedCost: 0,
+  }),
+  createConversation: () => ({ id: "conv-1" }),
+  addMessage: () => ({ id: `msg-${Date.now()}` }),
+  updateConversationUsage: () => {},
+  updateConversationTitle: () => {},
+  reserveMessage: mock(async () => ({ id: "msg-reserve" })),
+}));
+
+mock.module("../persistence/conversation-queries.js", () => ({
+  listConversations: () => [],
+}));
+
+mock.module("../persistence/attachments-store.js", () => ({
+  uploadAttachment: () => ({ id: `att-${Date.now()}` }),
+  linkAttachmentToMessage: () => {},
+}));
+
+mock.module("../plugins/defaults/compaction/window-manager.js", () => ({
+  ContextWindowManager: class {
+    estimateInputTokens() {
+      return 0;
+    }
+    get tokenCountInputs() {
+      return { systemPrompt: "", tools: undefined };
+    }
+    constructor() {}
+    updateConfig() {}
+    shouldCompact() {
+      return { needed: false, estimatedTokens: 0 };
+    }
+    async maybeCompact() {
+      return { compacted: false };
+    }
+    resetOverflowRecovery() {}
+  },
+  createContextSummaryMessage: () => ({
+    role: "user",
+    content: [{ type: "text", text: "summary" }],
+  }),
+  getSummaryFromContextMessage: () => null,
+}));
+
+mock.module("../persistence/llm-usage-store.js", () => ({
+  recordUsageEvent: () => ({ id: "mock-id", createdAt: Date.now() }),
+  listUsageEvents: () => [],
+}));
+
+mock.module("../agent/loop.js", () => ({
+  AgentLoop: class {
+    compactionCircuit = new CompactionCircuit("test-conv");
+    constructor() {}
+    getToolTokenBudget() {
+      return 0;
+    }
+    getResolvedTools() {
+      return [];
+    }
+    getActiveModel() {
+      return undefined;
+    }
+    setSystemPrompt() {}
+    async run(_options: {
+      messages: Message[];
+      onEvent: (event: AgentEvent) => void;
+    }): Promise<Message[]> {
+      return [];
+    }
+  },
+}));
+
+// ---------------------------------------------------------------------------
+// Import Conversation AFTER mocks
+// ---------------------------------------------------------------------------
+
+import { Conversation } from "../daemon/conversation.js";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function makeProvider() {
+  return {
+    name: "mock",
+    async sendMessage(): Promise<ProviderResponse> {
+      return {
+        content: [],
+        model: "mock",
+        usage: { inputTokens: 0, outputTokens: 0 },
+        stopReason: "end_turn",
+      };
+    },
+  };
+}
+
+let conversationSeq = 0;
+
+/**
+ * A conversation whose prompt is built, or, given a prompt other than the
+ * builder's, one running on a verbatim system-prompt override.
+ */
+function makeConversation(systemPrompt = "system prompt"): Conversation {
+  conversationSeq += 1;
+  return new Conversation(
+    `conv-rendered-delegation-${conversationSeq}`,
+    makeProvider(),
+    systemPrompt,
+    (_msg: AssistantEvent) => {},
+    "/tmp",
+    { maxTokens: 4096 },
+  );
+}
+
+function lastBuildDelegation(): unknown {
+  return builtWith.at(-1)?.delegateIndependentTasks;
+}
+
+beforeEach(() => {
+  builtWith = [];
+  throwOnBuild = false;
+});
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe("Conversation.buildCurrentSystemPrompt delegation-state capture", () => {
+  test("captures the state the build rendered and hands the builder the same value", () => {
+    const conversation = makeConversation();
+    expect(conversation.renderedDelegateIndependentTasks).toBeUndefined();
+
+    conversation.buildCurrentSystemPrompt();
+
+    // An unrestricted turn can spawn, so the section renders on.
+    expect(conversation.renderedDelegateIndependentTasks).toBe(true);
+    expect(lastBuildDelegation()).toBe(true);
+  });
+
+  test("a change to the gate's inputs after the build does not touch the captured state until the next build", () => {
+    const conversation = makeConversation();
+    conversation.buildCurrentSystemPrompt();
+    expect(conversation.renderedDelegateIndependentTasks).toBe(true);
+
+    // The loop is still sending the prompt built above, so the recorder must
+    // see that prompt's state, whatever the gate would derive now.
+    conversation.setSubagentAllowedTools(new Set(["remember"]));
+    expect(conversation.renderedDelegateIndependentTasks).toBe(true);
+
+    conversation.buildCurrentSystemPrompt();
+    expect(conversation.renderedDelegateIndependentTasks).toBe(false);
+    expect(lastBuildDelegation()).toBe(false);
+  });
+
+  test("a failed build leaves the previously captured state in place", () => {
+    const conversation = makeConversation();
+    conversation.buildCurrentSystemPrompt();
+    expect(conversation.renderedDelegateIndependentTasks).toBe(true);
+
+    // A wake's prompt sync swallows this and runs on the previous prompt,
+    // whose section is still on, even though a rebuild would now render it
+    // off.
+    conversation.setSubagentAllowedTools(new Set(["remember"]));
+    throwOnBuild = true;
+    expect(() => conversation.buildCurrentSystemPrompt()).toThrow(
+      "prompt build failed",
+    );
+
+    expect(conversation.renderedDelegateIndependentTasks).toBe(true);
+  });
+
+  test("a verbatim system-prompt override captures unknown", () => {
+    const conversation = makeConversation("override prompt");
+    builtWith = [];
+
+    expect(conversation.buildCurrentSystemPrompt()).toBe("override prompt");
+
+    expect(conversation.renderedDelegateIndependentTasks).toBeNull();
+    // The override is returned as is; nothing is built.
+    expect(builtWith).toEqual([]);
+  });
+
+  test("a replaying wake's recorded state wins over the derivation", () => {
+    const conversation = makeConversation();
+    // Replaying a channel-sourced conversation that recorded the section off,
+    // on a wake whose own scope could spawn.
+    conversation.delegateIndependentTasksReplay = false;
+
+    conversation.buildCurrentSystemPrompt();
+
+    expect(conversation.renderedDelegateIndependentTasks).toBe(false);
+    expect(lastBuildDelegation()).toBe(false);
+  });
+});

package/src/__tests__/credential-execution-client.test.ts

@@ -20,6 +20,7 @@ import { tmpdir } from "node:os";
 import { join, resolve } from "node:path";
 import { describe, expect, test } from "bun:test";
 
+import { resolveIpcEndpoint } from "@vellumai/ipc-server-utils";
 import { CES_PROTOCOL_VERSION } from "@vellumai/service-contracts/credential-rpc";
 
 import {
@@ -79,44 +80,83 @@ function createMockTransport(): CesTransport & {
 }
 
 // ---------------------------------------------------------------------------
-// Managed discovery — fail closed when socket is missing or handshake fails
+// CES discovery — CES_BOOTSTRAP_SOCKET_DIR, not workspace
 // ---------------------------------------------------------------------------
 
+function withBootstrapDir(dir: string): () => void {
+  const savedDir = process.env["CES_BOOTSTRAP_SOCKET_DIR"];
+  process.env["CES_BOOTSTRAP_SOCKET_DIR"] = dir;
+  return () => {
+    if (savedDir !== undefined) {
+      process.env["CES_BOOTSTRAP_SOCKET_DIR"] = savedDir;
+    } else {
+      delete process.env["CES_BOOTSTRAP_SOCKET_DIR"];
+    }
+  };
+}
+
 describe("managed CES discovery", () => {
-  test("returns unavailable when bootstrap socket does not exist", async () => {
-    // In a non-containerized test environment, the managed socket path
-    // does not exist. Verify fail-closed behavior.
-    const saved = process.env["CES_BOOTSTRAP_SOCKET"];
+  test("returns unavailable when bootstrap socket does not exist", () => {
+    const bootstrapDir = mkdtempSync(join(tmpdir(), "ces-missing-"));
+    const restore = withBootstrapDir(bootstrapDir);
     try {
-      // Point at a non-existent path to ensure predictable behavior
-      process.env["CES_BOOTSTRAP_SOCKET"] = "/tmp/ces-test-nonexistent.sock";
-      const result = await discoverManagedCes();
+      const result = discoverManagedCes();
       expect(result.mode).toBe("unavailable");
       expect((result as { reason: string }).reason).toContain(
         "CES bootstrap socket not found",
       );
     } finally {
-      if (saved !== undefined) {
-        process.env["CES_BOOTSTRAP_SOCKET"] = saved;
-      } else {
-        delete process.env["CES_BOOTSTRAP_SOCKET"];
-      }
+      restore();
+      rmSync(bootstrapDir, { recursive: true, force: true });
     }
   });
 
-  test("never returns a fallback or in-process mode", async () => {
-    const saved = process.env["CES_BOOTSTRAP_SOCKET"];
+  test("never returns a fallback or in-process mode", () => {
+    const bootstrapDir = mkdtempSync(join(tmpdir(), "ces-missing-"));
+    const restore = withBootstrapDir(bootstrapDir);
     try {
-      process.env["CES_BOOTSTRAP_SOCKET"] = "/tmp/ces-test-nonexistent.sock";
-      const result = await discoverManagedCes();
-      // Must be "managed" or "unavailable" — no fallback modes.
+      const result = discoverManagedCes();
       expect(["managed", "unavailable"]).toContain(result.mode);
     } finally {
-      if (saved !== undefined) {
-        process.env["CES_BOOTSTRAP_SOCKET"] = saved;
-      } else {
-        delete process.env["CES_BOOTSTRAP_SOCKET"];
-      }
+      restore();
+      rmSync(bootstrapDir, { recursive: true, force: true });
+    }
+  });
+});
+
+describe("CES bootstrap socket discovery", () => {
+  test("looks for the socket under CES_BOOTSTRAP_SOCKET_DIR", () => {
+    const bootstrapDir = mkdtempSync(join(tmpdir(), "ces-bootstrap-"));
+    const restore = withBootstrapDir(bootstrapDir);
+    try {
+      const socketPath = resolveIpcEndpoint("ces", {
+        workspaceDir: bootstrapDir,
+      }).path;
+      const result = discoverManagedCes();
+      expect(result.mode).toBe("unavailable");
+      expect((result as { reason: string }).reason).toContain(socketPath);
+    } finally {
+      restore();
+      rmSync(bootstrapDir, { recursive: true, force: true });
+    }
+  });
+
+  test("does not discover CES from VELLUM_WORKSPACE_DIR", () => {
+    const bootstrapDir = mkdtempSync(join(tmpdir(), "ces-bootstrap-"));
+    const restore = withBootstrapDir(bootstrapDir);
+    try {
+      const result = discoverManagedCes();
+      const workspaceDir = process.env.VELLUM_WORKSPACE_DIR;
+      expect(workspaceDir).toBeDefined();
+      const workspaceSocket = resolveIpcEndpoint("ces", {
+        workspaceDir: workspaceDir!,
+      }).path;
+      expect((result as { reason: string }).reason).not.toContain(
+        workspaceSocket,
+      );
+    } finally {
+      restore();
+      rmSync(bootstrapDir, { recursive: true, force: true });
     }
   });
 });
@@ -126,30 +166,9 @@ describe("managed CES discovery", () => {
 // ---------------------------------------------------------------------------
 
 describe("discoverCesWithRetry", () => {
-  function withManagedEnv(socketPath: string): () => void {
-    const savedSocket = process.env["CES_BOOTSTRAP_SOCKET"];
-    const savedSocketDir = process.env["CES_BOOTSTRAP_SOCKET_DIR"];
-    const savedContainerized = process.env["IS_CONTAINERIZED"];
-    delete process.env["CES_BOOTSTRAP_SOCKET_DIR"];
-    process.env["CES_BOOTSTRAP_SOCKET"] = socketPath;
-    process.env["IS_CONTAINERIZED"] = "true";
-    return () => {
-      const restore = (key: string, value: string | undefined) => {
-        if (value !== undefined) {
-          process.env[key] = value;
-        } else {
-          delete process.env[key];
-        }
-      };
-      restore("CES_BOOTSTRAP_SOCKET", savedSocket);
-      restore("CES_BOOTSTRAP_SOCKET_DIR", savedSocketDir);
-      restore("IS_CONTAINERIZED", savedContainerized);
-    };
-  }
-
   test("returns unavailable after polling when the socket never appears", async () => {
-    const socketPath = join(tmpdir(), `ces-retry-missing-${Date.now()}.sock`);
-    const restore = withManagedEnv(socketPath);
+    const bootstrapDir = mkdtempSync(join(tmpdir(), "ces-retry-missing-"));
+    const restore = withBootstrapDir(bootstrapDir);
     try {
       const start = Date.now();
       const result = await discoverCesWithRetry({
@@ -157,22 +176,19 @@ describe("discoverCesWithRetry", () => {
         intervalMs: 20,
       });
       expect(result.mode).toBe("unavailable");
-      // It must actually have waited rather than failing on the first probe.
       expect(Date.now() - start).toBeGreaterThanOrEqual(150);
     } finally {
       restore();
+      rmSync(bootstrapDir, { recursive: true, force: true });
     }
   });
 
   test("resolves to managed once the socket is re-bound mid-poll", async () => {
     const dir = mkdtempSync(join(tmpdir(), "ces-retry-"));
-    const socketPath = join(dir, "ces.sock");
-    const restore = withManagedEnv(socketPath);
+    const socketPath = resolveIpcEndpoint("ces", { workspaceDir: dir }).path;
+    const restore = withBootstrapDir(dir);
+    const timer = setTimeout(() => writeFileSync(socketPath, ""), 80);
     try {
-      // Simulate the sidecar re-binding its bootstrap socket shortly after
-      // the assistant begins probing.
-      setTimeout(() => writeFileSync(socketPath, ""), 80);
-
       const result = await discoverCesWithRetry({
         timeoutMs: 2_000,
         intervalMs: 20,
@@ -180,6 +196,7 @@ describe("discoverCesWithRetry", () => {
       expect(result.mode).toBe("managed");
       expect((result as { socketPath: string }).socketPath).toBe(socketPath);
     } finally {
+      clearTimeout(timer);
       restore();
       rmSync(dir, { recursive: true, force: true });
     }