@vellumai/assistant 0.12.2 → 0.12.3-staging.2

package/src/__tests__/conversation-app-control-lifecycle.test.ts

@@ -13,7 +13,7 @@
  * without bringing up the full daemon stack.
  */
 
-import { describe, expect, mock, test } from "bun:test";
+import { describe, expect, mock, spyOn, test } from "bun:test";
 
 mock.module("../providers/registry.js", () => ({
   getProvider: () => ({ name: "mock-provider" }),
@@ -70,6 +70,7 @@ mock.module("../runtime/services/auto-analysis-guard.js", () => ({
 
 import { Conversation } from "../daemon/conversation.js";
 import type { HostAppControlProxy } from "../daemon/host-app-control-proxy.js";
+import { HostCuProxy } from "../daemon/host-cu-proxy.js";
 
 /**
  * Minimal stand-in for HostAppControlProxy that records dispose() calls.
@@ -111,6 +112,57 @@ function makeConversation(): Conversation {
   return conv;
 }
 
+describe("Conversation HostCuProxy lifecycle", () => {
+  test.each(["replace", "remove"] as const)(
+    "%s disposes the prior proxy even when mode-session retirement fails",
+    async (operation) => {
+      const conversation = makeConversation();
+      const previous = new HostCuProxy();
+      const replacement =
+        operation === "replace" ? new HostCuProxy() : undefined;
+      conversation.setHostCuProxy(previous);
+      const dispose = spyOn(previous, "dispose");
+      const endTask = spyOn(
+        conversation.computerUseModeSessions,
+        "endTask",
+      ).mockImplementation(() => {
+        throw new Error("session retirement unavailable");
+      });
+      const pending = previous
+        .request(
+          "computer_use_click",
+          { element_id: 1 },
+          conversation.conversationId,
+          1,
+          "Clicking the button",
+        )
+        .catch((error: unknown) => error);
+      try {
+        expect(() => conversation.setHostCuProxy(previous)).not.toThrow();
+        expect(endTask).not.toHaveBeenCalled();
+        expect(() => conversation.setHostCuProxy(replacement)).not.toThrow();
+        expect(await pending).toMatchObject({
+          message: "Host CU proxy disposed",
+        });
+        expect(endTask).toHaveBeenCalledWith({
+          turnId: conversation.currentRequestId,
+          source: {
+            sourceId: previous.sourceId,
+            generation: previous.resetGeneration,
+          },
+        });
+        expect(dispose).toHaveBeenCalledTimes(1);
+        expect(conversation.hostCuProxy).toBe(replacement);
+      } finally {
+        endTask.mockRestore();
+        dispose.mockRestore();
+        previous.dispose();
+        replacement?.dispose();
+      }
+    },
+  );
+});
+
 describe("Conversation — HostAppControlProxy lifecycle", () => {
   test("setHostAppControlProxy stores the proxy", () => {
     const conversation = makeConversation();

package/src/__tests__/conversation-attachments.test.ts

@@ -56,6 +56,38 @@ function makeBase64(bytes: number): string {
   return buf.toString("base64");
 }
 
+function estimateBase64Bytes(
+  value: string | { data?: unknown; sizeBytes?: unknown } | null | undefined,
+): number {
+  if (value == null) {
+    return 0;
+  }
+  if (typeof value !== "string") {
+    if (typeof value.sizeBytes === "number") {
+      return value.sizeBytes;
+    }
+    return typeof value.data === "string" ? estimateBase64Bytes(value.data) : 0;
+  }
+  const trimmed = value.replace(/\s/g, "");
+  const padding = trimmed.endsWith("==") ? 2 : trimmed.endsWith("=") ? 1 : 0;
+  return Math.max(0, Math.floor((trimmed.length * 3) / 4) - padding);
+}
+
+function toolImageFilename(
+  mediaType: string,
+  toolName?: string,
+  title?: string,
+): string {
+  const prefix =
+    title ||
+    toolName
+      ?.replace(/([a-z])([A-Z])/g, "$1-$2")
+      .replace(/_/g, "-")
+      .toLowerCase() ||
+    "tool-output";
+  return `${prefix}.${mediaType.split("/")[1] ?? "png"}`;
+}
+
 // ---------------------------------------------------------------------------
 // resolveAssistantAttachments — all attachments are now file-backed
 // ---------------------------------------------------------------------------
@@ -92,6 +124,8 @@ describe("resolveAssistantAttachments", () => {
         accepted: d,
         warnings: [],
       }),
+      estimateBase64Bytes,
+      toolImageFilename,
     }));
 
     // Re-import to pick up mocks
@@ -117,6 +151,7 @@ describe("resolveAssistantAttachments", () => {
     expect(result.emittedAttachments.length).toBe(1);
     const emitted = result.emittedAttachments[0];
     expect(emitted.id).toBeDefined();
+    expect(result.linkedAttachmentIds).toEqual([emitted.id!]);
 
     // All attachments are file-backed and have a file on disk
     const filePath = getFilePathForAttachment(emitted.id!);
@@ -152,6 +187,8 @@ describe("resolveAssistantAttachments", () => {
         accepted: d,
         warnings: [],
       }),
+      estimateBase64Bytes,
+      toolImageFilename,
     }));
 
     const { resolveAssistantAttachments: resolve } =
@@ -234,6 +271,8 @@ describe("resolveAssistantAttachments", () => {
         accepted: d.filter((draft) => draft.filename !== "huge.bin"),
         warnings: ['Skipped attachment "huge.bin": too large.'],
       }),
+      estimateBase64Bytes,
+      toolImageFilename,
     }));
 
     const { resolveAssistantAttachments: resolve } =
@@ -281,6 +320,10 @@ describe("resolveAssistantAttachments", () => {
         sourceType: "sandbox_file",
       },
     ]);
+    expect(result.linkedAttachmentIds).toHaveLength(2);
+    expect(result.linkedAttachmentIds).toEqual(
+      result.emittedAttachments.map((attachment) => attachment.id!),
+    );
   });
 
   test("persistedFiles is empty when there is no message to persist against", async () => {
@@ -303,6 +346,8 @@ describe("resolveAssistantAttachments", () => {
         accepted: d,
         warnings: [],
       }),
+      estimateBase64Bytes,
+      toolImageFilename,
     }));
 
     const { resolveAssistantAttachments: resolve } =
@@ -328,6 +373,67 @@ describe("resolveAssistantAttachments", () => {
 
     expect(result.emittedAttachments).toHaveLength(1);
     expect(result.persistedFiles).toEqual([]);
+    expect(result.linkedAttachmentIds).toEqual([]);
+  });
+
+  test("reports only successful links when a later upload is skipped", async () => {
+    const conv = createConversation("test-conv-partial-upload");
+    const msg = await addMessage(conv.id, "assistant", "hello");
+    const valid: AssistantAttachmentDraft = {
+      sourceType: "sandbox_file",
+      filename: "valid.txt",
+      mimeType: "text/plain",
+      dataBase64: makeBase64(16),
+      sizeBytes: 16,
+      kind: "document",
+    };
+    const invalid: AssistantAttachmentDraft = {
+      sourceType: "sandbox_file",
+      filename: "invalid.txt",
+      mimeType: "text/plain",
+      dataBase64: "%%%",
+      sizeBytes: 3,
+      kind: "document",
+    };
+
+    mock.module("../daemon/assistant-attachments.js", () => ({
+      resolveDirectives: () =>
+        Promise.resolve({ drafts: [valid, invalid], warnings: [] }),
+      contentBlocksToDrafts: () => [],
+      deduplicateDrafts: (drafts: AssistantAttachmentDraft[]) => drafts,
+      validateDrafts: (drafts: AssistantAttachmentDraft[]) => ({
+        accepted: drafts,
+        warnings: [],
+      }),
+      estimateBase64Bytes,
+      toolImageFilename,
+    }));
+
+    const { resolveAssistantAttachments: resolve } =
+      await import("../daemon/conversation-attachments.js");
+    const result = await resolve(
+      [
+        {
+          source: "sandbox" as const,
+          path: "/fake",
+          filename: "valid.txt",
+          mimeType: "text/plain",
+        },
+      ],
+      [],
+      [],
+      "/tmp",
+      async () => true,
+      msg.id,
+    );
+
+    expect(result.emittedAttachments).toHaveLength(1);
+    expect(result.linkedAttachmentIds).toEqual([
+      result.emittedAttachments[0]!.id!,
+    ]);
+    expect(result.directiveWarnings).toContain(
+      "Attachment invalid.txt skipped: Invalid base64 encoding",
+    );
   });
 });
 

package/src/__tests__/conversation-evictor.test.ts

@@ -4,18 +4,9 @@ import { beforeEach, describe, expect, mock, test } from "bun:test";
 // Stub it so eviction tests run without a real manager: no protected children
 // by default, and abortAllForParent records the conversations it was called for.
 const abortedParents: string[] = [];
-let protectedChildren: Array<{ status: string }> = [];
 mock.module("../subagent/index.js", () => ({
   getSubagentManager: () => ({
     abortAllForParent: (id: string) => abortedParents.push(id),
-    getChildrenOf: () => protectedChildren,
-    hasActiveChildren: () =>
-      protectedChildren.some(
-        (c) =>
-          c.status === "running" ||
-          c.status === "pending" ||
-          c.status === "awaiting_input",
-      ),
   }),
 }));
 
@@ -30,11 +21,8 @@ function createMockSession(
 ): EvictableConversation & { disposed: boolean } {
   return {
     disposed: false,
-    isProcessing() {
-      return processing;
-    },
-    hasQueuedMessages() {
-      return queued;
+    hasInFlightWork() {
+      return processing || queued;
     },
     dispose() {
       this.disposed = true;
@@ -48,7 +36,6 @@ describe("ConversationEvictor", () => {
 
   beforeEach(() => {
     abortedParents.length = 0;
-    protectedChildren = [];
     sessions = new Map();
     evictor = new ConversationEvictor(
       sessions as Map<string, EvictableConversation>,
@@ -124,6 +111,19 @@ describe("ConversationEvictor", () => {
       expect(sessions.has("a")).toBe(true);
       expect(s1.disposed).toBe(false);
     });
+
+    test("skips conversations with resident mode-session work", () => {
+      const session = createMockSession();
+      session.hasInFlightWork = () => true;
+      sessions.set("mode-active", session);
+
+      const result = evictor.sweep();
+
+      expect(result.ttlEvicted).toBe(0);
+      expect(result.skipped).toBe(1);
+      expect(sessions.has("mode-active")).toBe(true);
+      expect(session.disposed).toBe(false);
+    });
   });
 
   describe("LRU eviction", () => {
@@ -217,43 +217,68 @@ describe("ConversationEvictor", () => {
       expect(s0.disposed).toBe(false);
       expect(s1.disposed).toBe(true);
     });
-  });
 
-  describe("onEvict", () => {
-    test("aborts subagents for each evicted session", () => {
-      const s1 = createMockSession();
-      sessions.set("a", s1);
-      // Never touched — will be TTL evicted
+    test("skips resident mode-session work during LRU eviction", () => {
+      const protectedSession = createMockSession();
+      protectedSession.hasInFlightWork = () => true;
+      const evictable = createMockSession();
+      const fresh1 = createMockSession();
+      const fresh2 = createMockSession();
+      sessions.set("mode-active", protectedSession);
+      sessions.set("evictable", evictable);
+      sessions.set("fresh-1", fresh1);
+      sessions.set("fresh-2", fresh2);
 
-      evictor.sweep();
+      const lastAccess = (
+        evictor as unknown as { lastAccess: Map<string, number> }
+      ).lastAccess;
+      const now = Date.now();
+      lastAccess.set("mode-active", now - 50);
+      lastAccess.set("evictable", now - 40);
+      lastAccess.set("fresh-1", now);
+      lastAccess.set("fresh-2", now);
 
-      expect(abortedParents).toEqual(["a"]);
+      const result = evictor.sweep();
+
+      expect(result.lruEvicted).toBe(1);
+      expect(sessions.has("mode-active")).toBe(true);
+      expect(sessions.has("evictable")).toBe(false);
+      expect(protectedSession.disposed).toBe(false);
     });
   });
 
-  describe("shouldProtect", () => {
-    test("skips conversations with running or pending subagents", () => {
-      protectedChildren = [{ status: "running" }];
-      const s1 = createMockSession();
-      sessions.set("a", s1);
+  describe("memory-pressure eviction", () => {
+    test("keeps resident mode-session work while evicting idle conversations", () => {
+      const resident = createMockSession();
+      resident.hasInFlightWork = () => true;
+      const idle = createMockSession();
+      sessions.set("mode-active", resident);
+      sessions.set("idle", idle);
+      evictor = new ConversationEvictor(sessions, {
+        ttlMs: Number.MAX_SAFE_INTEGER,
+        maxConversations: Number.MAX_SAFE_INTEGER,
+        memoryThresholdBytes: 0,
+      });
 
       const result = evictor.sweep();
 
-      expect(result.skipped).toBe(1);
-      expect(sessions.has("a")).toBe(true);
-      expect(s1.disposed).toBe(false);
+      expect(result.memoryEvicted).toBe(1);
+      expect(sessions.has("mode-active")).toBe(true);
+      expect(sessions.has("idle")).toBe(false);
+      expect(resident.disposed).toBe(false);
+      expect(idle.disposed).toBe(true);
     });
+  });
 
-    test("skips conversations with awaiting_input subagents", () => {
-      protectedChildren = [{ status: "awaiting_input" }];
+  describe("onEvict", () => {
+    test("aborts subagents for each evicted session", () => {
       const s1 = createMockSession();
       sessions.set("a", s1);
+      // Never touched, so it will be TTL evicted.
 
-      const result = evictor.sweep();
+      evictor.sweep();
 
-      expect(result.skipped).toBe(1);
-      expect(sessions.has("a")).toBe(true);
-      expect(s1.disposed).toBe(false);
+      expect(abortedParents).toEqual(["a"]);
     });
   });
 

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", {