@vellumai/assistant 0.8.5 → 0.8.6

package/src/__tests__/conversation-workspace-tool-tracking.test.ts

@@ -142,6 +142,8 @@ mock.module("../memory/conversation-crud.js", () => ({
   setLastNotifiedInferenceProfile: () => {},
   getConversationOverrideProfileFromRow: () => undefined,
   updateMessageMetadata: () => {},
+  reserveMessage: mock(async () => ({ id: "msg-reserve" })),
+  updateMessageContent: mock(() => {}),
 }));
 
 mock.module("../memory/conversation-queries.js", () => ({
@@ -228,6 +230,9 @@ mock.module("../agent/loop.js", () => ({
       messages: Message[],
       onEvent: (event: AgentEvent) => void,
     ): Promise<Message[]> {
+      // Prime the assistant row anchor — production code emits this from
+      // `AgentLoop.run` just before `provider.sendMessage`.
+      await onEvent({ type: "llm_call_started" });
       agentLoopScript(onEvent);
       onEvent({
         type: "usage",

package/src/__tests__/credential-broker-browser-fill.test.ts

@@ -28,8 +28,8 @@ mock.module("../util/logger.js", () => ({
 // Use encrypted backend with a temp store path
 // ---------------------------------------------------------------------------
 
-import { _setStorePath } from "../security/encrypted-store.js";
 import { _resetBackend } from "../security/secure-keys.js";
+import { setStorePathForTesting } from "./encrypted-store-test-helpers.js";
 
 const TEST_DIR = join(
   tmpdir(),
@@ -74,7 +74,7 @@ describe("CredentialBroker.browserFill", () => {
     for (const entry of readdirSync(TEST_DIR)) {
       rmSync(join(TEST_DIR, entry), { recursive: true, force: true });
     }
-    _setStorePath(STORE_PATH);
+    setStorePathForTesting(STORE_PATH);
     _resetBackend();
     _setMetadataPath(join(TEST_DIR, "metadata.json"));
     broker = new CredentialBroker();
@@ -82,7 +82,7 @@ describe("CredentialBroker.browserFill", () => {
 
   afterEach(() => {
     _setMetadataPath(null);
-    _setStorePath(null);
+    setStorePathForTesting(null);
     _resetBackend();
   });
 

package/src/__tests__/credential-broker-server-use.test.ts

@@ -27,8 +27,8 @@ mock.module("../util/logger.js", () => ({
 // Use encrypted backend with a temp store path
 // ---------------------------------------------------------------------------
 
-import { _setStorePath } from "../security/encrypted-store.js";
 import { _resetBackend } from "../security/secure-keys.js";
+import { setStorePathForTesting } from "./encrypted-store-test-helpers.js";
 
 const TEST_DIR = join(
   tmpdir(),
@@ -69,7 +69,7 @@ describe("CredentialBroker.serverUse", () => {
 
   beforeEach(() => {
     mkdirSync(TEST_DIR, { recursive: true });
-    _setStorePath(STORE_PATH);
+    setStorePathForTesting(STORE_PATH);
     _resetBackend();
     _setMetadataPath(join(TEST_DIR, "metadata.json"));
     broker = new CredentialBroker();
@@ -77,7 +77,7 @@ describe("CredentialBroker.serverUse", () => {
 
   afterEach(() => {
     _setMetadataPath(null);
-    _setStorePath(null);
+    setStorePathForTesting(null);
     _resetBackend();
     rmSync(TEST_DIR, { recursive: true, force: true });
   });
@@ -464,7 +464,7 @@ describe("CredentialBroker.serverUseById", () => {
 
   beforeEach(() => {
     mkdirSync(TEST_DIR, { recursive: true });
-    _setStorePath(STORE_PATH);
+    setStorePathForTesting(STORE_PATH);
     _resetBackend();
     _setMetadataPath(join(TEST_DIR, "metadata.json"));
     broker = new CredentialBroker();
@@ -472,7 +472,7 @@ describe("CredentialBroker.serverUseById", () => {
 
   afterEach(() => {
     _setMetadataPath(null);
-    _setStorePath(null);
+    setStorePathForTesting(null);
     _resetBackend();
     rmSync(TEST_DIR, { recursive: true, force: true });
   });

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

@@ -8,7 +8,15 @@
  * 4. The CES RPC client correctly frames requests and validates responses.
  */
 
-import { readdirSync, readFileSync, statSync } from "node:fs";
+import {
+  mkdtempSync,
+  readdirSync,
+  readFileSync,
+  rmSync,
+  statSync,
+  writeFileSync,
+} from "node:fs";
+import { tmpdir } from "node:os";
 import { join, resolve } from "node:path";
 import { describe, expect, test } from "bun:test";
 
@@ -22,6 +30,7 @@ import {
   createCesClient,
 } from "../credential-execution/client.js";
 import {
+  discoverCesWithRetry,
   discoverLocalCes,
   discoverManagedCes,
 } from "../credential-execution/executable-discovery.js";
@@ -147,6 +156,68 @@ describe("managed CES discovery", () => {
   });
 });
 
+// ---------------------------------------------------------------------------
+// Managed discovery retry — absorbs the sidecar's socket re-bind window
+// ---------------------------------------------------------------------------
+
+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);
+    try {
+      const start = Date.now();
+      const result = await discoverCesWithRetry({
+        timeoutMs: 200,
+        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();
+    }
+  });
+
+  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);
+    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,
+      });
+      expect(result.mode).toBe("managed");
+      expect((result as { socketPath: string }).socketPath).toBe(socketPath);
+    } finally {
+      restore();
+      rmSync(dir, { recursive: true, force: true });
+    }
+  });
+});
+
 // ---------------------------------------------------------------------------
 // CES client — transport and framing
 // ---------------------------------------------------------------------------

package/src/__tests__/credential-execution-feature-gates.test.ts

@@ -10,10 +10,7 @@
 
 import { afterEach, beforeEach, describe, expect, test } from "bun:test";
 
-import {
-  _setOverridesForTesting,
-  isAssistantFeatureFlagEnabled,
-} from "../config/assistant-feature-flags.js";
+import { isAssistantFeatureFlagEnabled } from "../config/assistant-feature-flags.js";
 import type { AssistantConfig } from "../config/schema.js";
 import {
   CES_GRANT_AUDIT_FLAG_KEY,
@@ -25,20 +22,21 @@ import {
   isCesShellLockdownEnabled,
   isCesToolsEnabled,
 } from "../credential-execution/feature-gates.js";
+import { setOverridesForTesting } from "./feature-flag-test-helpers.js";
 
 beforeEach(() => {
-  _setOverridesForTesting({});
+  setOverridesForTesting({});
 });
 
 afterEach(() => {
-  _setOverridesForTesting({});
+  setOverridesForTesting({});
 });
 
 // ---------------------------------------------------------------------------
 // Helpers
 // ---------------------------------------------------------------------------
 
-/** Create a minimal AssistantConfig (flag overrides are now set via _setOverridesForTesting). */
+/** Create a minimal AssistantConfig (flag overrides are now set via setOverridesForTesting). */
 function makeConfig(): AssistantConfig {
   return {} as AssistantConfig;
 }
@@ -120,13 +118,13 @@ describe("CES flags match registry defaults", () => {
 describe("CES flags can be toggled independently", () => {
   for (const { name, fn, key } of ALL_CES_PREDICATES) {
     test(`enabling ${key} makes ${name} return true`, () => {
-      _setOverridesForTesting({ [key]: true });
+      setOverridesForTesting({ [key]: true });
       const config = makeConfig();
       expect(fn(config)).toBe(true);
     });
 
     test(`enabling ${key} does not change other CES flags from their defaults`, () => {
-      _setOverridesForTesting({ [key]: true });
+      setOverridesForTesting({ [key]: true });
       const config = makeConfig();
       for (const {
         fn: otherFn,
@@ -147,7 +145,7 @@ describe("CES flags can be toggled independently", () => {
 describe("CES flags respect explicit false overrides", () => {
   for (const { name, fn, key } of ALL_CES_PREDICATES) {
     test(`${name} returns false when explicitly set to false`, () => {
-      _setOverridesForTesting({ [key]: false });
+      setOverridesForTesting({ [key]: false });
       const config = makeConfig();
       expect(fn(config)).toBe(false);
     });
@@ -164,7 +162,7 @@ describe("CES flags do not affect unrelated flags", () => {
     for (const key of ALL_CES_FLAG_KEYS) {
       overrides[key] = true;
     }
-    _setOverridesForTesting(overrides);
+    setOverridesForTesting(overrides);
     const config = makeConfig();
 
     // account-deletion defaults to true in the registry and should stay true.
@@ -178,7 +176,7 @@ describe("CES flags do not affect unrelated flags", () => {
     for (const key of ALL_CES_FLAG_KEYS) {
       overrides[key] = true;
     }
-    _setOverridesForTesting(overrides);
+    setOverridesForTesting(overrides);
     const config = makeConfig();
 
     // Unknown flags fail closed unless explicitly overridden.

package/src/__tests__/credential-health-service.test.ts

@@ -12,6 +12,7 @@ let mockProviders: Array<{
   pingMethod: string | null;
   pingHeaders: string | null;
   pingBody: string | null;
+  managedServiceConfigKey?: string;
 }> = [];
 
 let mockConnections: Map<
@@ -33,6 +34,24 @@ let mockFetchResponse: { ok: boolean; status: number } = {
 };
 let mockFetchThrows = false;
 
+// Per-test refresh outcome — drives the withValidToken mock.
+//   "ok"          → refresh succeeds (or wasn't needed); callback runs with token
+//   "refresh_failed" → refresh throws TokenExpiredError (revoked refresh token)
+type RefreshOutcome = "ok" | "refresh_failed";
+let mockRefreshOutcome: RefreshOutcome = "ok";
+
+// Managed-path mock state.
+const managedProviders = new Set<string>();
+let managedListResponse: {
+  ok: boolean;
+  status: number;
+  body?: unknown;
+} = { ok: true, status: 200, body: { results: [] } };
+let managedPingOutcome:
+  | { kind: "status"; status: number }
+  | { kind: "credential_required" }
+  | { kind: "throw"; message: string } = { kind: "status", status: 200 };
+
 // ── Module mocks ─────────────────────────────────────────────────────
 
 mock.module("../security/secure-keys.js", () => ({
@@ -81,6 +100,88 @@ mock.module("../util/logger.js", () => ({
   }),
 }));
 
+class MockTokenExpiredError extends Error {
+  constructor(service: string, message?: string) {
+    super(message ?? `Token expired for "${service}"`);
+    this.name = "TokenExpiredError";
+  }
+}
+
+mock.module("../security/token-manager.js", () => ({
+  TokenExpiredError: MockTokenExpiredError,
+  withValidToken: async (
+    service: string,
+    callback: (token: string) => Promise<unknown>,
+    opts: { connectionId: string },
+  ) => {
+    if (mockRefreshOutcome === "refresh_failed") {
+      throw new MockTokenExpiredError(service);
+    }
+    const token =
+      secureKeyValues.get(
+        `oauth_connection/${opts.connectionId}/access_token`,
+      ) ?? "refreshed-token";
+    return callback(token);
+  },
+}));
+
+// Managed-path mocks: services schema, config loader, platform client,
+// PlatformOAuthConnection. The managed code path uses dynamic imports for
+// these so they only matter when a managed test scenario is exercised.
+mock.module("../config/schemas/services.js", () => ({
+  ServicesSchema: {
+    shape: {
+      "google-oauth": {},
+      "notion-oauth": {},
+    },
+  },
+  getServiceMode: (_services: unknown, key: string) =>
+    managedProviders.has(key) ? "managed" : "your-own",
+}));
+
+mock.module("../config/loader.js", () => ({
+  getConfig: () => ({ services: {} }),
+}));
+
+mock.module("../platform/client.js", () => ({
+  VellumPlatformClient: {
+    create: async () => ({
+      platformAssistantId: "test-assistant",
+      fetch: async (_path: string) => ({
+        ok: managedListResponse.ok,
+        status: managedListResponse.status,
+        json: async () => managedListResponse.body ?? { results: [] },
+      }),
+    }),
+  },
+}));
+
+class MockCredentialRequiredError extends Error {
+  constructor() {
+    super("credential required");
+    this.name = "CredentialRequiredError";
+  }
+}
+
+mock.module("../oauth/platform-connection.js", () => ({
+  PlatformOAuthConnection: class {
+    constructor(_opts: unknown) {}
+    async request(_req: unknown) {
+      if (managedPingOutcome.kind === "credential_required") {
+        throw new MockCredentialRequiredError();
+      }
+      if (managedPingOutcome.kind === "throw") {
+        throw new Error(managedPingOutcome.message);
+      }
+      return {
+        status: managedPingOutcome.status,
+        headers: {},
+        body: {},
+      };
+    }
+  },
+}));
+
 // ── Import under test ────────────────────────────────────────────────
 
 const { checkAllCredentials, checkCredentialForProvider, _setFetchFn } =
@@ -106,6 +207,7 @@ function addProvider(
     defaultScopes?: string[];
     pingUrl?: string | null;
     pingMethod?: string | null;
+    managedServiceConfigKey?: string;
   },
 ) {
   mockProviders.push({
@@ -115,6 +217,10 @@ function addProvider(
     pingMethod: opts?.pingMethod ?? null,
     pingHeaders: null,
     pingBody: null,
+    // checkManagedProvider reads this field via OAuthProviderRow; the
+    // health-check code branches to the managed path only when set AND
+    // when managedProviders includes the corresponding config key.
+    managedServiceConfigKey: opts?.managedServiceConfigKey,
   });
 }
 
@@ -159,6 +265,10 @@ describe("credential-health-service", () => {
     mockConnections = new Map();
     mockFetchResponse = { ok: true, status: 200 };
     mockFetchThrows = false;
+    mockRefreshOutcome = "ok";
+    managedProviders.clear();
+    managedListResponse = { ok: true, status: 200, body: { results: [] } };
+    managedPingOutcome = { kind: "status", status: 200 };
   });
 
   test("returns empty report when no providers exist", async () => {
@@ -238,8 +348,12 @@ describe("credential-health-service", () => {
     expect(report.results[0]!.canAutoRecover).toBe(false);
   });
 
-  test("returns expiring when token is past expiresAt with refresh token", async () => {
-    addProvider("google");
+  test("returns healthy when expired token has a refresh token, no ping URL, and refresh succeeds via the ping path", async () => {
+    // Without a pingUrl there's no opportunity to exercise the refresh —
+    // the connection is reported as healthy on the trust that the next
+    // real API call will refresh transparently. This matches the contract
+    // for connections that have a refresh token but no liveness endpoint.
+    addProvider("google"); // no pingUrl
     addConnection("google", "conn-1", {
       expiresAt: Date.now() - 60_000, // 1 minute ago
       hasRefreshToken: true,
@@ -247,10 +361,61 @@ describe("credential-health-service", () => {
     setToken("conn-1");
 
     const report = await checkAllCredentials();
-    expect(report.results[0]!.status).toBe("expiring");
+    expect(report.results[0]!.status).toBe("healthy");
     expect(report.results[0]!.canAutoRecover).toBe(true);
   });
 
+  test("expired BYO token + successful refresh + 200 ping → healthy", async () => {
+    // The withValidToken mock returns "refreshed" by default; pingProvider
+    // sees a fresh token and the upstream returns 200.
+    mockFetchResponse = { ok: true, status: 200 };
+    addProvider("google", {
+      pingUrl: "https://www.googleapis.com/oauth2/v2/userinfo",
+    });
+    addConnection("google", "conn-1", {
+      expiresAt: Date.now() - 60_000,
+      hasRefreshToken: true,
+    });
+    setToken("conn-1");
+
+    const report = await checkAllCredentials();
+    expect(report.results[0]!.status).toBe("healthy");
+  });
+
+  test("expired BYO token + refresh fails (revoked refresh token) → revoked", async () => {
+    mockRefreshOutcome = "refresh_failed";
+    addProvider("google", {
+      pingUrl: "https://www.googleapis.com/oauth2/v2/userinfo",
+    });
+    addConnection("google", "conn-1", {
+      expiresAt: Date.now() - 60_000,
+      hasRefreshToken: true,
+    });
+    setToken("conn-1");
+
+    const report = await checkAllCredentials();
+    expect(report.results[0]!.status).toBe("revoked");
+    expect(report.results[0]!.canAutoRecover).toBe(false);
+  });
+
+  test("BYO ping that still returns 401 after a refresh attempt → revoked", async () => {
+    // Refresh succeeds (mockRefreshOutcome stays "ok") but the upstream
+    // still returns 401 — the token was genuinely revoked by the provider.
+    mockFetchResponse = { ok: false, status: 401 };
+    addProvider("google", {
+      pingUrl: "https://www.googleapis.com/oauth2/v2/userinfo",
+    });
+    addConnection("google", "conn-1", {
+      expiresAt: Date.now() + 30 * 24 * 60 * 60 * 1000,
+      hasRefreshToken: true,
+    });
+    setToken("conn-1");
+
+    const report = await checkAllCredentials();
+    expect(report.results[0]!.status).toBe("revoked");
+    expect(report.results[0]!.details).toContain("after a refresh attempt");
+  });
+
   test("returns expiring when token expires within 7 days without refresh token", async () => {
     addProvider("google");
     addConnection("google", "conn-1", {
@@ -459,6 +624,90 @@ describe("credential-health-service", () => {
     });
   });
 
+  describe("managed-provider path", () => {
+    // Distinguishing managed-path failure modes: a 424 from the platform
+    // proxy (CredentialRequiredError) means the platform tried and gave up
+    // on refresh — actionable, fire reconnect alert. A bare 401/403 from
+    // the upstream provider could be a transient platform-side miss (proxy
+    // didn't refresh-before-forward) and should NOT fire reconnect alerts
+    // without further evidence — demote to ping_failed.
+
+    function setupManagedGoogle(opts?: { hasActiveConnection?: boolean }) {
+      addProvider("google", {
+        pingUrl: "https://www.googleapis.com/oauth2/v2/userinfo",
+        managedServiceConfigKey: "google-oauth",
+      });
+      managedProviders.add("google-oauth");
+      if (opts?.hasActiveConnection !== false) {
+        managedListResponse = {
+          ok: true,
+          status: 200,
+          body: {
+            results: [
+              {
+                id: "platform-conn-1",
+                account_label: "user@example.com",
+                status: "ACTIVE",
+              },
+            ],
+          },
+        };
+      }
+    }
+
+    test("managed 2xx ping → healthy", async () => {
+      setupManagedGoogle();
+      managedPingOutcome = { kind: "status", status: 200 };
+
+      const report = await checkAllCredentials();
+      const google = report.results.find((r) => r.provider === "google");
+      expect(google!.status).toBe("healthy");
+    });
+
+    test("managed 424 (CredentialRequiredError) → revoked", async () => {
+      setupManagedGoogle();
+      managedPingOutcome = { kind: "credential_required" };
+
+      const report = await checkAllCredentials();
+      const google = report.results.find((r) => r.provider === "google");
+      expect(google!.status).toBe("revoked");
+      expect(google!.canAutoRecover).toBe(false);
+      expect(google!.details).toContain("cannot be refreshed");
+    });
+
+    test("managed 401 from upstream → ping_failed, not revoked", async () => {
+      // This is the key change: previously this would fire a user-facing
+      // reconnect alert (because "revoked" is in hardFailureStatuses). Now
+      // we treat upstream 401 as potentially transient — only a 424 from
+      // the platform proxy is trusted enough to trigger the alert.
+      setupManagedGoogle();
+      managedPingOutcome = { kind: "status", status: 401 };
+
+      const report = await checkAllCredentials();
+      const google = report.results.find((r) => r.provider === "google");
+      expect(google!.status).toBe("ping_failed");
+      expect(google!.status).not.toBe("revoked");
+    });
+
+    test("managed 403 from upstream → ping_failed, not revoked", async () => {
+      setupManagedGoogle();
+      managedPingOutcome = { kind: "status", status: 403 };
+
+      const report = await checkAllCredentials();
+      const google = report.results.find((r) => r.provider === "google");
+      expect(google!.status).toBe("ping_failed");
+    });
+
+    test("managed 500 from upstream → ping_failed", async () => {
+      setupManagedGoogle();
+      managedPingOutcome = { kind: "status", status: 500 };
+
+      const report = await checkAllCredentials();
+      const google = report.results.find((r) => r.provider === "google");
+      expect(google!.status).toBe("ping_failed");
+    });
+  });
+
   describe("checkCredentialForProvider", () => {
     test("returns null when no connections exist", async () => {
       const result = await checkCredentialForProvider("google");

package/src/__tests__/credential-security-invariants.test.ts

@@ -41,8 +41,8 @@ afterAll(() => {
   mock.restore();
 });
 
-import { _setStorePath } from "../security/encrypted-store.js";
 import { _resetBackend } from "../security/secure-keys.js";
+import { setStorePathForTesting } from "./encrypted-store-test-helpers.js";
 
 const TEST_DIR = join(
   tmpdir(),
@@ -409,7 +409,7 @@ describe("Invariant 4: credentials only used for allowed purpose", () => {
 
   beforeEach(() => {
     mkdirSync(TEST_DIR, { recursive: true });
-    _setStorePath(STORE_PATH);
+    setStorePathForTesting(STORE_PATH);
     _resetBackend();
     _setMetadataPath(join(TEST_DIR, "metadata.json"));
     broker = new CredentialBroker();
@@ -417,7 +417,7 @@ describe("Invariant 4: credentials only used for allowed purpose", () => {
 
   afterEach(() => {
     _setMetadataPath(null);
-    _setStorePath(null);
+    setStorePathForTesting(null);
     _resetBackend();
     rmSync(TEST_DIR, { recursive: true, force: true });
   });
@@ -497,14 +497,14 @@ describe("Invariant 4: credentials only used for allowed purpose", () => {
 describe("Invariant 6: oauth2ClientSecret not in metadata, only in secure store", () => {
   beforeEach(() => {
     mkdirSync(TEST_DIR, { recursive: true });
-    _setStorePath(STORE_PATH);
+    setStorePathForTesting(STORE_PATH);
     _resetBackend();
     _setMetadataPath(join(TEST_DIR, "metadata.json"));
   });
 
   afterEach(() => {
     _setMetadataPath(null);
-    _setStorePath(null);
+    setStorePathForTesting(null);
     _resetBackend();
     rmSync(TEST_DIR, { recursive: true, force: true });
   });