@vellumai/assistant 0.12.2-staging.1 → 0.12.2-staging.3

package/src/mcp/mcp-auth-state.ts

@@ -15,14 +15,27 @@
  * connection.
  */
 type McpAuthState =
-  | { status: "pending"; authUrl: string; attemptId: string; expiresAt: number }
+  | {
+      status: "pending";
+      authUrl: string;
+      attemptId: string;
+      targetKey: string;
+      expiresAt: number;
+    }
   | {
       status: "complete";
       serverId: string;
       attemptId: string;
+      targetKey: string;
       completedAt: number;
     }
-  | { status: "error"; error: string; attemptId: string; failedAt: number };
+  | {
+      status: "error";
+      error: string;
+      attemptId: string;
+      targetKey: string;
+      failedAt: number;
+    };
 
 const activeMcpAuthFlows = new Map<string, McpAuthState>();
 
@@ -41,11 +54,13 @@ export function setMcpAuthPending(
   serverId: string,
   authUrl: string,
   attemptId: string,
+  targetKey = serverId,
 ): void {
   activeMcpAuthFlows.set(serverId, {
     status: "pending",
     authUrl,
     attemptId,
+    targetKey,
     expiresAt: Date.now() + PENDING_TTL_MS,
   });
 }
@@ -59,15 +74,20 @@ export function setMcpAuthPending(
 export function setMcpAuthComplete(
   serverId: string,
   attemptId: string,
+  targetKey = serverId,
 ): boolean {
   const current = activeMcpAuthFlows.get(serverId);
-  if (current && current.attemptId !== attemptId) {
+  if (
+    current &&
+    (current.attemptId !== attemptId || current.targetKey !== targetKey)
+  ) {
     return false; // superseded
   }
   activeMcpAuthFlows.set(serverId, {
     status: "complete",
     serverId,
     attemptId,
+    targetKey,
     completedAt: Date.now(),
   });
   return true;
@@ -81,15 +101,20 @@ export function setMcpAuthError(
   serverId: string,
   error: string,
   attemptId: string,
+  targetKey = serverId,
 ): boolean {
   const current = activeMcpAuthFlows.get(serverId);
-  if (current && current.attemptId !== attemptId) {
+  if (
+    current &&
+    (current.attemptId !== attemptId || current.targetKey !== targetKey)
+  ) {
     return false; // superseded
   }
   activeMcpAuthFlows.set(serverId, {
     status: "error",
     error,
     attemptId,
+    targetKey,
     failedAt: Date.now(),
   });
   return true;
@@ -100,7 +125,10 @@ export function setMcpAuthError(
  * expired entries on every read so a long-polling CLI can never observe
  * a stale flow past its TTL.
  */
-export function getMcpAuthState(serverId: string): McpAuthState | null {
+export function getMcpAuthState(
+  serverId: string,
+  targetKey?: string,
+): McpAuthState | null {
   const now = Date.now();
   for (const [id, state] of activeMcpAuthFlows) {
     if (state.status === "pending" && now > state.expiresAt) {
@@ -117,5 +145,9 @@ export function getMcpAuthState(serverId: string): McpAuthState | null {
       activeMcpAuthFlows.delete(id);
     }
   }
-  return activeMcpAuthFlows.get(serverId) ?? null;
+  const state = activeMcpAuthFlows.get(serverId) ?? null;
+  if (state && targetKey !== undefined && state.targetKey !== targetKey) {
+    return null;
+  }
+  return state;
 }

package/src/mcp/mcp-oauth-provider.ts

@@ -38,23 +38,16 @@ import {
 import { openInHostBrowser } from "../util/browser.js";
 import { getLogger } from "../util/logger.js";
 import { APP_VERSION } from "../version.js";
+import {
+  MCP_OAUTH_CREDENTIAL_LEAVES,
+  mcpOAuthCredentialKey,
+  type McpOAuthCredentialTarget,
+  pluginMcpOAuthCredentialPrefix,
+  workspaceMcpOAuthCredentialTarget,
+} from "./credential-target.js";
 
 const log = getLogger("mcp-oauth");
 
-// Credential store key helpers
-function tokensKey(serverId: string): string {
-  return `mcp:${serverId}:tokens`;
-}
-function clientInfoKey(serverId: string): string {
-  return `mcp:${serverId}:client_info`;
-}
-function discoveryKey(serverId: string): string {
-  return `mcp:${serverId}:discovery`;
-}
-function clientBindingKey(serverId: string): string {
-  return `mcp:${serverId}:client_binding`;
-}
-
 /**
  * Logo shown on an authorization server's consent screen.
  *
@@ -91,12 +84,15 @@ export interface McpOAuthProviderOptions {
    * URL to the IPC caller (CLI / web client).
    */
   onAuthorizationUrl?: (url: string) => void;
+  /** Credential identity. Defaults to the legacy workspace server id. */
+  credentialTarget?: McpOAuthCredentialTarget;
 }
 
 export class McpOAuthProvider implements OAuthClientProvider {
   private readonly serverId: string;
   private readonly serverUrl: string;
   private readonly interactive: boolean;
+  private readonly credentialTarget: McpOAuthCredentialTarget;
   private _codeVerifier: string | undefined;
   private _state: string | undefined;
   private _redirectUrl: string | undefined;
@@ -121,6 +117,8 @@ export class McpOAuthProvider implements OAuthClientProvider {
     this.serverUrl = serverUrl;
     this.interactive = interactive;
     this._onAuthorizationUrl = options.onAuthorizationUrl;
+    this.credentialTarget =
+      options.credentialTarget ?? workspaceMcpOAuthCredentialTarget(serverId);
   }
 
   // --- redirectUrl ---
@@ -167,7 +165,9 @@ export class McpOAuthProvider implements OAuthClientProvider {
   // --- Tokens ---
 
   async tokens(): Promise<OAuthTokens | undefined> {
-    const raw = await getSecureKeyAsync(tokensKey(this.serverId));
+    const raw = await getSecureKeyAsync(
+      mcpOAuthCredentialKey(this.credentialTarget, "tokens"),
+    );
     if (!raw) {
       return undefined;
     }
@@ -191,7 +191,9 @@ export class McpOAuthProvider implements OAuthClientProvider {
     // the previous refresh_token when the incoming response omits one.
     let toPersist: OAuthTokens = tokens;
     if (!tokens.refresh_token) {
-      const previous = await getSecureKeyAsync(tokensKey(this.serverId));
+      const previous = await getSecureKeyAsync(
+        mcpOAuthCredentialKey(this.credentialTarget, "tokens"),
+      );
       if (previous) {
         try {
           const parsed = JSON.parse(previous) as OAuthTokens;
@@ -204,7 +206,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
       }
     }
     const ok = await setSecureKeyAsync(
-      tokensKey(this.serverId),
+      mcpOAuthCredentialKey(this.credentialTarget, "tokens"),
       JSON.stringify(toPersist),
     );
     if (!ok) {
@@ -263,7 +265,9 @@ export class McpOAuthProvider implements OAuthClientProvider {
    * fronts the resource.
    */
   async clientInformation(): Promise<OAuthClientInformationMixed | undefined> {
-    const raw = await getSecureKeyAsync(clientInfoKey(this.serverId));
+    const raw = await getSecureKeyAsync(
+      mcpOAuthCredentialKey(this.credentialTarget, "client_info"),
+    );
     if (!raw) {
       return undefined;
     }
@@ -294,7 +298,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
     info: OAuthClientInformationMixed,
   ): Promise<void> {
     const ok = await setSecureKeyAsync(
-      clientInfoKey(this.serverId),
+      mcpOAuthCredentialKey(this.credentialTarget, "client_info"),
       JSON.stringify(info),
     );
     if (!ok) {
@@ -313,7 +317,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
         redirectUri: this._redirectUrl,
       };
       const boundOk = await setSecureKeyAsync(
-        clientBindingKey(this.serverId),
+        mcpOAuthCredentialKey(this.credentialTarget, "client_binding"),
         JSON.stringify(binding),
       );
       if (!boundOk) {
@@ -349,7 +353,9 @@ export class McpOAuthProvider implements OAuthClientProvider {
    * redirect check already covers the case this plugin actually changes.
    */
   private async describeStaleBinding(): Promise<string | null> {
-    const raw = await getSecureKeyAsync(clientBindingKey(this.serverId));
+    const raw = await getSecureKeyAsync(
+      mcpOAuthCredentialKey(this.credentialTarget, "client_binding"),
+    );
     if (!raw) {
       // Registered before the binding was recorded. Keep it: discarding
       // every pre-existing registration would remake all of them at once.
@@ -410,7 +416,9 @@ export class McpOAuthProvider implements OAuthClientProvider {
   // --- Discovery State ---
 
   async discoveryState(): Promise<OAuthDiscoveryState | undefined> {
-    const raw = await getSecureKeyAsync(discoveryKey(this.serverId));
+    const raw = await getSecureKeyAsync(
+      mcpOAuthCredentialKey(this.credentialTarget, "discovery"),
+    );
     if (!raw) {
       return undefined;
     }
@@ -423,7 +431,7 @@ export class McpOAuthProvider implements OAuthClientProvider {
 
   async saveDiscoveryState(state: OAuthDiscoveryState): Promise<void> {
     const ok = await setSecureKeyAsync(
-      discoveryKey(this.serverId),
+      mcpOAuthCredentialKey(this.credentialTarget, "discovery"),
       JSON.stringify(state),
     );
     if (!ok) {
@@ -498,7 +506,9 @@ export class McpOAuthProvider implements OAuthClientProvider {
     );
 
     if (scope === "all" || scope === "tokens") {
-      const result = await deleteSecureKeyAsync(tokensKey(this.serverId));
+      const result = await deleteSecureKeyAsync(
+        mcpOAuthCredentialKey(this.credentialTarget, "tokens"),
+      );
       if (result === "error") {
         log.warn(
           { serverId: this.serverId },
@@ -512,7 +522,9 @@ export class McpOAuthProvider implements OAuthClientProvider {
       }
     }
     if (scope === "all" || scope === "client") {
-      const result = await deleteSecureKeyAsync(clientInfoKey(this.serverId));
+      const result = await deleteSecureKeyAsync(
+        mcpOAuthCredentialKey(this.credentialTarget, "client_info"),
+      );
       if (result === "error") {
         log.warn(
           { serverId: this.serverId },
@@ -527,14 +539,18 @@ export class McpOAuthProvider implements OAuthClientProvider {
       // The binding describes the registration being dropped, so it goes
       // with it. Leaving it would let a later registration inherit the
       // previous one's issuer and redirect URI.
-      await deleteSecureKeyAsync(clientBindingKey(this.serverId));
+      await deleteSecureKeyAsync(
+        mcpOAuthCredentialKey(this.credentialTarget, "client_binding"),
+      );
     }
     if (scope === "all" || scope === "verifier") {
       this._codeVerifier = undefined;
       this._state = undefined;
     }
     if (scope === "all" || scope === "discovery") {
-      const result = await deleteSecureKeyAsync(discoveryKey(this.serverId));
+      const result = await deleteSecureKeyAsync(
+        mcpOAuthCredentialKey(this.credentialTarget, "discovery"),
+      );
       if (result === "error") {
         log.warn(
           { serverId: this.serverId },
@@ -617,8 +633,10 @@ export class McpOAuthProvider implements OAuthClientProvider {
 /**
  * Check whether OAuth tokens exist in the credential store for a server.
  */
-export async function hasMcpOAuthTokens(serverId: string): Promise<boolean> {
-  const raw = await getSecureKeyAsync(tokensKey(serverId));
+export async function hasMcpOAuthTokens(
+  target: McpOAuthCredentialTarget,
+): Promise<boolean> {
+  const raw = await getSecureKeyAsync(mcpOAuthCredentialKey(target, "tokens"));
   return raw != null && raw.length > 0;
 }
 
@@ -631,10 +649,14 @@ export async function deleteMcpOAuthCredentials(
 ): Promise<{ ok: boolean; failedKeys: string[] }> {
   const [tokensResult, clientResult, bindingResult, discoveryResult] =
     await Promise.all([
-      deleteSecureKeyAsync(tokensKey(serverId)),
-      deleteSecureKeyAsync(clientInfoKey(serverId)),
-      deleteSecureKeyAsync(clientBindingKey(serverId)),
-      deleteSecureKeyAsync(discoveryKey(serverId)),
+      ...MCP_OAUTH_CREDENTIAL_LEAVES.map((leaf) =>
+        deleteSecureKeyAsync(
+          mcpOAuthCredentialKey(
+            workspaceMcpOAuthCredentialTarget(serverId),
+            leaf,
+          ),
+        ),
+      ),
     ]);
   const results = [
     { key: "tokens", result: tokensResult },
@@ -660,3 +682,49 @@ export async function deleteMcpOAuthCredentials(
   );
   return { ok, failedKeys };
 }
+
+export interface PluginMcpOAuthCleanupResult {
+  readonly ok: boolean;
+  readonly unreachable: boolean;
+  readonly matchedKeys: string[];
+  readonly failedKeys: string[];
+}
+
+/** Delete every endpoint-scoped MCP OAuth record owned by one plugin. */
+export async function deletePluginMcpOAuthCredentials(
+  pluginName: string,
+): Promise<PluginMcpOAuthCleanupResult> {
+  const { listSecureKeysAsync } = await import("../security/secure-keys.js");
+  const listed = await listSecureKeysAsync();
+  if (listed.unreachable) {
+    return {
+      ok: false,
+      unreachable: true,
+      matchedKeys: [],
+      failedKeys: [],
+    };
+  }
+
+  const prefix = pluginMcpOAuthCredentialPrefix(pluginName);
+  const matchedKeys = listed.accounts.filter((key) => key.startsWith(prefix));
+  const deletionResults = await Promise.all(
+    matchedKeys.map(async (key) => {
+      try {
+        return { key, result: await deleteSecureKeyAsync(key) };
+      } catch {
+        return { key, result: "error" as const };
+      }
+    }),
+  );
+  const failedKeys = deletionResults
+    .filter(({ result }) => result === "error")
+    .map(({ key }) => key);
+  const ok = failedKeys.length === 0;
+  if (!ok) {
+    log.warn(
+      { pluginName, failedKeys },
+      "Some plugin MCP OAuth credentials could not be deleted",
+    );
+  }
+  return { ok, unreachable: false, matchedKeys, failedKeys };
+}

package/src/monitoring/__tests__/mount-watch.test.ts

@@ -0,0 +1,236 @@
+import { describe, expect, test } from "bun:test";
+
+import {
+  collectMountSnapshot,
+  diffMountSnapshots,
+  isWatchedMount,
+  mountDiffIsEmpty,
+  type MountInfoEntry,
+  type MountSnapshot,
+  parseMountinfo,
+  type PathFsSnapshot,
+  unescapeMountField,
+} from "../mount-watch.js";
+
+const SAMPLE_MOUNTINFO = `22 1 0:21 / / rw,relatime - overlay overlay rw,lowerdir=/a,upperdir=/b,workdir=/c
+36 22 0:26 / /workspace rw,relatime - virtiofs kataShared rw
+48 22 0:27 / /data rw,relatime - virtiofs kataShared rw
+25 22 0:23 / /proc rw,nosuid,nodev,noexec,relatime - proc proc rw
+`;
+
+function entry(
+  fields: Partial<MountInfoEntry> & Pick<MountInfoEntry, "mountPoint" | "fsType">,
+): MountInfoEntry {
+  return {
+    mountId: fields.mountId ?? "1",
+    parentId: fields.parentId ?? "0",
+    majorMinor: fields.majorMinor ?? "0:1",
+    root: fields.root ?? "/",
+    mountPoint: fields.mountPoint,
+    mountOptions: fields.mountOptions ?? "rw,relatime",
+    fsType: fields.fsType,
+    source: fields.source ?? "kataShared",
+    superOptions: fields.superOptions ?? "rw",
+  };
+}
+
+function pathFs(
+  fields: Partial<PathFsSnapshot> & Pick<PathFsSnapshot, "path">,
+): PathFsSnapshot {
+  return {
+    path: fields.path,
+    type: fields.type ?? 0x65735546,
+    blockSize: fields.blockSize ?? 4096,
+    blocks: fields.blocks ?? 12_884_901,
+    files: fields.files ?? 1_000_000,
+    ...(fields.error != null ? { error: fields.error } : {}),
+  };
+}
+
+function snapshot(partial: Partial<MountSnapshot>): MountSnapshot {
+  return {
+    mounts: partial.mounts ?? [],
+    paths: partial.paths ?? [],
+  };
+}
+
+describe("unescapeMountField", () => {
+  test("decodes octal-escaped spaces", () => {
+    expect(unescapeMountField("/workspace/my\\040dir")).toBe("/workspace/my dir");
+  });
+});
+
+describe("parseMountinfo", () => {
+  test("reads mount point, fstype, and source", () => {
+    const parsed = parseMountinfo(SAMPLE_MOUNTINFO);
+    expect(parsed).toEqual([
+      entry({
+        mountId: "22",
+        parentId: "1",
+        majorMinor: "0:21",
+        mountPoint: "/",
+        fsType: "overlay",
+        source: "overlay",
+        superOptions: "rw,lowerdir=/a,upperdir=/b,workdir=/c",
+      }),
+      entry({
+        mountId: "36",
+        parentId: "22",
+        majorMinor: "0:26",
+        mountPoint: "/workspace",
+        fsType: "virtiofs",
+      }),
+      entry({
+        mountId: "48",
+        parentId: "22",
+        majorMinor: "0:27",
+        mountPoint: "/data",
+        fsType: "virtiofs",
+      }),
+      entry({
+        mountId: "25",
+        parentId: "22",
+        majorMinor: "0:23",
+        mountPoint: "/proc",
+        mountOptions: "rw,nosuid,nodev,noexec,relatime",
+        fsType: "proc",
+        source: "proc",
+      }),
+    ]);
+  });
+
+  test("skips malformed lines", () => {
+    expect(parseMountinfo("not-a-mount-line\n")).toEqual([]);
+  });
+});
+
+describe("isWatchedMount", () => {
+  test("keeps virtiofs regardless of mount point", () => {
+    expect(
+      isWatchedMount(entry({ mountPoint: "/secret-bind", fsType: "virtiofs" })),
+    ).toBe(true);
+  });
+
+  test("keeps /workspace and /data even when they are overlay", () => {
+    expect(
+      isWatchedMount(entry({ mountPoint: "/workspace", fsType: "overlay" })),
+    ).toBe(true);
+    expect(isWatchedMount(entry({ mountPoint: "/data", fsType: "overlay" }))).toBe(
+      true,
+    );
+  });
+
+  test("drops proc and the root overlay", () => {
+    expect(isWatchedMount(entry({ mountPoint: "/proc", fsType: "proc" }))).toBe(
+      false,
+    );
+    expect(isWatchedMount(entry({ mountPoint: "/", fsType: "overlay" }))).toBe(
+      false,
+    );
+  });
+});
+
+describe("diffMountSnapshots", () => {
+  test("stays empty when the table is unchanged", () => {
+    const current = snapshot({
+      mounts: [entry({ mountPoint: "/workspace", fsType: "virtiofs" })],
+      paths: [pathFs({ path: "/workspace" })],
+    });
+    expect(mountDiffIsEmpty(diffMountSnapshots(current, current))).toBe(true);
+  });
+
+  test("reports a dropped virtiofs bind as removed", () => {
+    const previous = snapshot({
+      mounts: [
+        entry({ mountPoint: "/workspace", fsType: "virtiofs" }),
+        entry({ mountPoint: "/data", fsType: "virtiofs" }),
+      ],
+    });
+    const current = snapshot({
+      mounts: [entry({ mountPoint: "/data", fsType: "virtiofs" })],
+    });
+    const diff = diffMountSnapshots(previous, current);
+    expect(diff.removed.map((row) => row.mountPoint)).toEqual(["/workspace"]);
+    expect(diff.added).toEqual([]);
+    expect(diff.changed).toEqual([]);
+  });
+
+  test("reports a fstype swap at the same mount point as changed", () => {
+    const previous = snapshot({
+      mounts: [entry({ mountPoint: "/workspace", fsType: "virtiofs" })],
+    });
+    const current = snapshot({
+      mounts: [entry({ mountPoint: "/workspace", fsType: "overlay", source: "overlay" })],
+    });
+    const diff = diffMountSnapshots(previous, current);
+    expect(diff.changed).toHaveLength(1);
+    expect(diff.changed[0].previous.fsType).toBe("virtiofs");
+    expect(diff.changed[0].current.fsType).toBe("overlay");
+  });
+
+  test("reports a statfs size collapse on a watched path", () => {
+    const previous = snapshot({
+      paths: [pathFs({ path: "/workspace", blocks: 12_884_901 })],
+    });
+    const current = snapshot({
+      paths: [pathFs({ path: "/workspace", blocks: 25_165_824, type: 0x794c7630 })],
+    });
+    const diff = diffMountSnapshots(previous, current);
+    expect(diff.pathChanges).toHaveLength(1);
+    expect(diff.pathChanges[0].previous.blocks).toBe(12_884_901);
+    expect(diff.pathChanges[0].current.blocks).toBe(25_165_824);
+  });
+});
+
+describe("collectMountSnapshot", () => {
+  test("returns null on Windows and macOS", () => {
+    expect(collectMountSnapshot({ platform: "win32" })).toBeNull();
+    expect(collectMountSnapshot({ platform: "darwin" })).toBeNull();
+  });
+
+  test("filters to assistant-container volume mounts", () => {
+    const snapshot = collectMountSnapshot({
+      platform: "linux",
+      readMountinfo: () => SAMPLE_MOUNTINFO,
+      statPath: (path) => pathFs({ path }),
+      watchedMountPoints: ["/workspace", "/data"],
+    });
+    expect(snapshot).not.toBeNull();
+    expect(snapshot?.mounts.map((row) => row.mountPoint)).toEqual([
+      "/data",
+      "/workspace",
+    ]);
+    expect(snapshot?.paths.map((row) => row.path)).toEqual([
+      "/workspace",
+      "/data",
+    ]);
+  });
+
+  test("stats watched paths even when they are missing from the host", () => {
+    const snapshot = collectMountSnapshot({
+      platform: "linux",
+      readMountinfo: () => SAMPLE_MOUNTINFO,
+      statPath: (path) =>
+        pathFs({ path, type: null, blockSize: null, blocks: null, files: null, error: "statfs-failed" }),
+      watchedMountPoints: ["/workspace", "/data"],
+    });
+    expect(snapshot?.paths).toEqual([
+      pathFs({
+        path: "/workspace",
+        type: null,
+        blockSize: null,
+        blocks: null,
+        files: null,
+        error: "statfs-failed",
+      }),
+      pathFs({
+        path: "/data",
+        type: null,
+        blockSize: null,
+        blocks: null,
+        files: null,
+        error: "statfs-failed",
+      }),
+    ]);
+  });
+});