@vellumai/assistant 0.4.52 → 0.4.54

package/src/telemetry/usage-telemetry-reporter.test.ts

@@ -0,0 +1,366 @@
+/**
+ * Tests for UsageTelemetryReporter.
+ *
+ * Covers both auth modes (authenticated / anonymous), watermark advancement,
+ * error handling, batch recursion, installation ID persistence, and payload shape.
+ */
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+// ---------------------------------------------------------------------------
+// Module mocks (must precede production imports)
+// ---------------------------------------------------------------------------
+
+const mockGetMemoryCheckpoint = mock<(key: string) => string | null>(
+  () => null,
+);
+const mockSetMemoryCheckpoint = mock<(key: string, value: string) => void>(
+  () => {},
+);
+
+mock.module("../memory/checkpoints.js", () => ({
+  getMemoryCheckpoint: mockGetMemoryCheckpoint,
+  setMemoryCheckpoint: mockSetMemoryCheckpoint,
+}));
+
+const mockQueryUnreportedUsageEvents = mock(
+  () =>
+    [] as ReturnType<
+      typeof import("../memory/llm-usage-store.js").queryUnreportedUsageEvents
+    >,
+);
+
+mock.module("../memory/llm-usage-store.js", () => ({
+  queryUnreportedUsageEvents: mockQueryUnreportedUsageEvents,
+}));
+
+const mockResolveManagedProxyContext = mock(async () => ({
+  enabled: false,
+  platformBaseUrl: "",
+  assistantApiKey: "",
+}));
+
+mock.module("../providers/managed-proxy/context.js", () => ({
+  resolveManagedProxyContext: mockResolveManagedProxyContext,
+}));
+
+const mockGetTelemetryPlatformUrl = mock(() => "https://platform.vellum.ai");
+const mockGetTelemetryAppToken = mock(() => "");
+
+mock.module("../config/env.js", () => ({
+  getTelemetryPlatformUrl: mockGetTelemetryPlatformUrl,
+  getTelemetryAppToken: mockGetTelemetryAppToken,
+  // Re-export anything else the module might import transitively
+  str: () => undefined,
+  num: () => undefined,
+  bool: () => false,
+}));
+
+mock.module("../util/logger.js", () => ({
+  getLogger: () =>
+    new Proxy({} as Record<string, unknown>, {
+      get: () => () => {},
+    }),
+}));
+
+// ---------------------------------------------------------------------------
+// Production import (after mocks)
+// ---------------------------------------------------------------------------
+
+import type { UsageEvent } from "../usage/types.js";
+import { UsageTelemetryReporter } from "./usage-telemetry-reporter.js";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+let eventIdCounter = 0;
+
+function makeUsageEvent(overrides: Partial<UsageEvent> = {}): UsageEvent {
+  eventIdCounter += 1;
+  return {
+    id: `evt-${eventIdCounter}`,
+    createdAt: 1700000000000 + eventIdCounter * 1000,
+    provider: "anthropic",
+    model: "claude-sonnet-4-20250514",
+    inputTokens: 100,
+    outputTokens: 50,
+    cacheCreationInputTokens: 10,
+    cacheReadInputTokens: 5,
+    actor: "main_agent",
+    conversationId: "conv-1",
+    runId: null,
+    requestId: null,
+    estimatedCostUsd: 0.001,
+    pricingStatus: "priced",
+    ...overrides,
+  };
+}
+
+const originalFetch = globalThis.fetch;
+let mockFetch: ReturnType<typeof mock>;
+
+// ---------------------------------------------------------------------------
+// Setup / Teardown
+// ---------------------------------------------------------------------------
+
+beforeEach(() => {
+  eventIdCounter = 0;
+  mockGetMemoryCheckpoint.mockReset();
+  mockSetMemoryCheckpoint.mockReset();
+  mockQueryUnreportedUsageEvents.mockReset();
+  mockResolveManagedProxyContext.mockReset();
+  mockGetTelemetryPlatformUrl.mockReset();
+  mockGetTelemetryAppToken.mockReset();
+
+  // Defaults
+  mockGetMemoryCheckpoint.mockReturnValue(null);
+  mockResolveManagedProxyContext.mockResolvedValue({
+    enabled: false,
+    platformBaseUrl: "",
+    assistantApiKey: "",
+  });
+  mockGetTelemetryPlatformUrl.mockReturnValue("https://platform.vellum.ai");
+  mockGetTelemetryAppToken.mockReturnValue("default-test-token");
+
+  mockFetch = mock(() =>
+    Promise.resolve(new Response('{"accepted":0}', { status: 200 })),
+  );
+  globalThis.fetch = mockFetch as unknown as typeof fetch;
+});
+
+afterEach(() => {
+  globalThis.fetch = originalFetch;
+});
+
+// ---------------------------------------------------------------------------
+// Tests
+// ---------------------------------------------------------------------------
+
+describe("UsageTelemetryReporter", () => {
+  test("authenticated flush uses Api-Key header and proxy URL", async () => {
+    mockResolveManagedProxyContext.mockResolvedValue({
+      enabled: true,
+      platformBaseUrl: "https://test.vellum.ai",
+      assistantApiKey: "test-key",
+    });
+    const events = [makeUsageEvent(), makeUsageEvent()];
+    mockQueryUnreportedUsageEvents.mockReturnValue(events);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(
+        new Response(`{"accepted":${events.length}}`, { status: 200 }),
+      ),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+    await reporter.flush();
+
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    const [url, opts] = mockFetch.mock.calls[0] as [string, RequestInit];
+    expect(url).toBe(
+      "https://test.vellum.ai/v1/assistants/self-hosted-local/telemetry/usage/",
+    );
+    expect((opts.headers as Record<string, string>)["Authorization"]).toBe(
+      "Api-Key test-key",
+    );
+  });
+
+  test("anonymous flush uses X-Telemetry-Token and default URL", async () => {
+    mockResolveManagedProxyContext.mockResolvedValue({
+      enabled: false,
+      platformBaseUrl: "",
+      assistantApiKey: "",
+    });
+    mockGetTelemetryPlatformUrl.mockReturnValue("https://platform.test.ai");
+    mockGetTelemetryAppToken.mockReturnValue("anon-token");
+
+    const events = [makeUsageEvent()];
+    mockQueryUnreportedUsageEvents.mockReturnValue(events);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(new Response('{"accepted":1}', { status: 200 })),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+    await reporter.flush();
+
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    const [url, opts] = mockFetch.mock.calls[0] as [string, RequestInit];
+    expect(url).toStartWith("https://platform.test.ai");
+    expect((opts.headers as Record<string, string>)["X-Telemetry-Token"]).toBe(
+      "anon-token",
+    );
+  });
+
+  test("watermark advances on successful upload", async () => {
+    const events = [
+      makeUsageEvent({ id: "evt-w1", createdAt: 1700000001000 }),
+      makeUsageEvent({ id: "evt-w2", createdAt: 1700000002000 }),
+    ];
+    mockQueryUnreportedUsageEvents.mockReturnValue(events);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(new Response('{"accepted":2}', { status: 200 })),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+    await reporter.flush();
+
+    const watermarkCalls = mockSetMemoryCheckpoint.mock.calls.filter(
+      (c) => c[0] === "telemetry:usage:last_reported_at",
+    );
+    expect(watermarkCalls.length).toBeGreaterThanOrEqual(1);
+    // The watermark should be set to the createdAt of the last event
+    expect(watermarkCalls[watermarkCalls.length - 1][1]).toBe(
+      String(1700000002000),
+    );
+
+    // The compound cursor ID should also be set to the last event's id
+    const idCalls = mockSetMemoryCheckpoint.mock.calls.filter(
+      (c) => c[0] === "telemetry:usage:last_reported_id",
+    );
+    expect(idCalls.length).toBeGreaterThanOrEqual(1);
+    expect(idCalls[idCalls.length - 1][1]).toBe("evt-w2");
+  });
+
+  test("watermark stays on failed upload", async () => {
+    const events = [makeUsageEvent()];
+    mockQueryUnreportedUsageEvents.mockReturnValue(events);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(new Response("error", { status: 500 })),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+    await reporter.flush();
+
+    const watermarkCalls = mockSetMemoryCheckpoint.mock.calls.filter(
+      (c) => c[0] === "telemetry:usage:last_reported_at",
+    );
+    expect(watermarkCalls.length).toBe(0);
+  });
+
+  test("installation ID generated on first flush, reused thereafter", async () => {
+    let storedInstallId: string | null = null;
+
+    mockGetMemoryCheckpoint.mockImplementation((key: string) => {
+      if (key === "telemetry:installation_id") return storedInstallId;
+      return null;
+    });
+    mockSetMemoryCheckpoint.mockImplementation((key: string, value: string) => {
+      if (key === "telemetry:installation_id") storedInstallId = value;
+    });
+
+    const events = [makeUsageEvent()];
+    mockQueryUnreportedUsageEvents.mockReturnValue(events);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(new Response('{"accepted":1}', { status: 200 })),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+
+    // First flush — should generate and store a new installation ID
+    await reporter.flush();
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    const body1 = JSON.parse(
+      (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+    );
+    expect(body1.installation_id).toBeTruthy();
+    expect(typeof body1.installation_id).toBe("string");
+
+    // Second flush
+    mockQueryUnreportedUsageEvents.mockReturnValue([makeUsageEvent()]);
+    await reporter.flush();
+    expect(mockFetch).toHaveBeenCalledTimes(2);
+    const body2 = JSON.parse(
+      (mockFetch.mock.calls[1] as [string, RequestInit])[1].body as string,
+    );
+
+    // Both flushes should use the same installation ID
+    expect(body2.installation_id).toBe(body1.installation_id);
+  });
+
+  test("empty batch makes no HTTP call", async () => {
+    mockQueryUnreportedUsageEvents.mockReturnValue([]);
+
+    const reporter = new UsageTelemetryReporter();
+    await reporter.flush();
+
+    expect(mockFetch).not.toHaveBeenCalled();
+  });
+
+  test("batch recursion when full, capped at 10", async () => {
+    // Always return exactly 500 events (BATCH_SIZE) to trigger recursion
+    const fullBatch = Array.from({ length: 500 }, (_, i) =>
+      makeUsageEvent({ id: `evt-batch-${i}`, createdAt: 1700000000000 + i }),
+    );
+    mockQueryUnreportedUsageEvents.mockReturnValue(fullBatch);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(new Response('{"accepted":500}', { status: 200 })),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+    await reporter.flush();
+
+    // MAX_CONSECUTIVE_BATCHES = 10
+    expect(mockFetch).toHaveBeenCalledTimes(10);
+  });
+
+  test("stop() performs final flush", async () => {
+    const events = [makeUsageEvent()];
+    mockQueryUnreportedUsageEvents.mockReturnValue(events);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(new Response('{"accepted":1}', { status: 200 })),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+    reporter.start();
+
+    // Wait a tick so start()'s immediate flush settles.
+    await new Promise((resolve) => setTimeout(resolve, 50));
+    const callsBeforeStop = mockFetch.mock.calls.length;
+
+    await reporter.stop();
+
+    // stop() must trigger at least one additional flush beyond what start() did.
+    expect(mockFetch.mock.calls.length).toBeGreaterThan(callsBeforeStop);
+  });
+
+  test("payload shape is correct", async () => {
+    const event = makeUsageEvent({
+      id: "evt-shape-test",
+      provider: "anthropic",
+      model: "claude-sonnet-4-20250514",
+      inputTokens: 200,
+      outputTokens: 100,
+      cacheCreationInputTokens: 20,
+      cacheReadInputTokens: 15,
+      actor: "context_compactor",
+      createdAt: 1700000099000,
+    });
+    mockQueryUnreportedUsageEvents.mockReturnValue([event]);
+    mockFetch.mockImplementation(() =>
+      Promise.resolve(new Response('{"accepted":1}', { status: 200 })),
+    );
+
+    const reporter = new UsageTelemetryReporter();
+    await reporter.flush();
+
+    expect(mockFetch).toHaveBeenCalledTimes(1);
+    const body = JSON.parse(
+      (mockFetch.mock.calls[0] as [string, RequestInit])[1].body as string,
+    );
+
+    // Top-level: installation_id and events array
+    expect(typeof body.installation_id).toBe("string");
+    expect(Array.isArray(body.events)).toBe(true);
+    expect(body.events.length).toBe(1);
+
+    const e = body.events[0];
+    expect(e.daemon_event_id).toBe("evt-shape-test");
+    expect(e.provider).toBe("anthropic");
+    expect(e.model).toBe("claude-sonnet-4-20250514");
+    expect(e.input_tokens).toBe(200);
+    expect(e.output_tokens).toBe(100);
+    expect(e.cache_creation_input_tokens).toBe(20);
+    expect(e.cache_read_input_tokens).toBe(15);
+    expect(e.actor).toBe("context_compactor");
+    expect(e.recorded_at).toBe(1700000099000);
+  });
+});

package/src/telemetry/usage-telemetry-reporter.ts

@@ -0,0 +1,181 @@
+/**
+ * Usage telemetry reporter.
+ *
+ * Periodically flushes LLM usage events from the local SQLite
+ * `llm_usage_events` table and POSTs them to the platform telemetry endpoint.
+ *
+ * Two auth modes:
+ * - Authenticated: Api-Key header via managed proxy context
+ * - Anonymous: X-Telemetry-Token static token from env
+ */
+
+import { v4 as uuid } from "uuid";
+
+import {
+  getTelemetryAppToken,
+  getTelemetryPlatformUrl,
+} from "../config/env.js";
+import {
+  getMemoryCheckpoint,
+  setMemoryCheckpoint,
+} from "../memory/checkpoints.js";
+import { queryUnreportedUsageEvents } from "../memory/llm-usage-store.js";
+import { resolveManagedProxyContext } from "../providers/managed-proxy/context.js";
+import { getLogger } from "../util/logger.js";
+
+const log = getLogger("usage-telemetry");
+
+// ---------------------------------------------------------------------------
+// Constants
+// ---------------------------------------------------------------------------
+
+const CHECKPOINT_KEY_WATERMARK = "telemetry:usage:last_reported_at";
+const CHECKPOINT_KEY_WATERMARK_ID = "telemetry:usage:last_reported_id";
+const CHECKPOINT_KEY_INSTALL_ID = "telemetry:installation_id";
+const REPORT_INTERVAL_MS = 5 * 60 * 1000;
+const BATCH_SIZE = 500;
+const MAX_CONSECUTIVE_BATCHES = 10;
+const TELEMETRY_PATH = "/v1/assistants/self-hosted-local/telemetry/usage/";
+
+// ---------------------------------------------------------------------------
+// Installation ID
+// ---------------------------------------------------------------------------
+
+function getOrCreateInstallationId(): string {
+  const existing = getMemoryCheckpoint(CHECKPOINT_KEY_INSTALL_ID);
+  if (existing) return existing;
+
+  const id = uuid();
+  setMemoryCheckpoint(CHECKPOINT_KEY_INSTALL_ID, id);
+  return id;
+}
+
+// ---------------------------------------------------------------------------
+// Reporter
+// ---------------------------------------------------------------------------
+
+export class UsageTelemetryReporter {
+  private timer: ReturnType<typeof setInterval> | null = null;
+  private activeFlush: Promise<void> | null = null;
+
+  start(): void {
+    this.flush().catch((err) => {
+      log.warn({ err }, "Initial usage telemetry flush failed");
+    });
+    this.timer = setInterval(() => {
+      this.flush().catch((err) => {
+        log.warn({ err }, "Scheduled usage telemetry flush failed");
+      });
+    }, REPORT_INTERVAL_MS);
+  }
+
+  async stop(): Promise<void> {
+    if (this.timer) {
+      clearInterval(this.timer);
+      this.timer = null;
+    }
+    if (this.activeFlush) {
+      await this.activeFlush;
+    }
+    await this.flush();
+  }
+
+  async flush(): Promise<void> {
+    if (this.activeFlush) return; // overlap guard
+    this.activeFlush = this._doFlush();
+    try {
+      await this.activeFlush;
+    } finally {
+      this.activeFlush = null;
+    }
+  }
+
+  private async _doFlush(batchCount = 0): Promise<void> {
+    try {
+      if (batchCount >= MAX_CONSECUTIVE_BATCHES) return;
+
+      // Read watermark (compound cursor: createdAt + id)
+      const watermark = Number(
+        getMemoryCheckpoint(CHECKPOINT_KEY_WATERMARK) ?? "0",
+      );
+      const watermarkId =
+        getMemoryCheckpoint(CHECKPOINT_KEY_WATERMARK_ID) ?? undefined;
+
+      // Query unreported events
+      const events = queryUnreportedUsageEvents(
+        watermark,
+        watermarkId,
+        BATCH_SIZE,
+      );
+      if (events.length === 0) return;
+
+      // Resolve auth context — skip flush when neither auth mode is viable
+      const proxyCtx = await resolveManagedProxyContext();
+      if (!proxyCtx.enabled && !getTelemetryAppToken()) {
+        return;
+      }
+
+      let url: string;
+      let authHeaders: Record<string, string>;
+
+      if (proxyCtx.enabled) {
+        url = `${proxyCtx.platformBaseUrl}${TELEMETRY_PATH}`;
+        authHeaders = { Authorization: `Api-Key ${proxyCtx.assistantApiKey}` };
+      } else {
+        url = `${getTelemetryPlatformUrl()}${TELEMETRY_PATH}`;
+        authHeaders = { "X-Telemetry-Token": getTelemetryAppToken() };
+      }
+
+      // Build payload
+      const payload = {
+        installation_id: getOrCreateInstallationId(),
+        events: events.map((e) => ({
+          daemon_event_id: e.id,
+          provider: e.provider,
+          model: e.model,
+          input_tokens: e.inputTokens,
+          output_tokens: e.outputTokens,
+          cache_creation_input_tokens: e.cacheCreationInputTokens ?? null,
+          cache_read_input_tokens: e.cacheReadInputTokens ?? null,
+          actor: e.actor,
+          recorded_at: e.createdAt,
+        })),
+      };
+
+      // Send
+      const resp = await fetch(url, {
+        method: "POST",
+        headers: {
+          "Content-Type": "application/json",
+          ...authHeaders,
+        },
+        body: JSON.stringify(payload),
+      });
+
+      if (!resp.ok) {
+        await resp.text(); // consume body to release connection
+        log.warn(
+          { status: resp.status, url },
+          "Usage telemetry POST failed — will retry next cycle",
+        );
+        return;
+      }
+      await resp.text(); // consume body to release connection
+
+      // Advance watermark (compound cursor)
+      const lastEvent = events[events.length - 1];
+      setMemoryCheckpoint(
+        CHECKPOINT_KEY_WATERMARK,
+        String(lastEvent.createdAt),
+      );
+      setMemoryCheckpoint(CHECKPOINT_KEY_WATERMARK_ID, lastEvent.id);
+
+      // If we got a full batch, there may be more events — recurse
+      if (events.length === BATCH_SIZE) {
+        await this._doFlush(batchCount + 1);
+      }
+    } catch (err) {
+      log.warn({ err }, "Usage telemetry flush error — non-fatal, will retry");
+    }
+  }
+}

package/src/tools/claude-code/claude-code.ts

@@ -2,9 +2,9 @@ import {
   getCCCommand,
   loadCCCommandTemplate,
 } from "../../commands/cc-command-registry.js";
-import { getConfig } from "../../config/loader.js";
 import { RiskLevel } from "../../permissions/types.js";
 import type { ToolDefinition } from "../../providers/types.js";
+import { getSecureKeyAsync } from "../../security/secure-keys.js";
 import type { WorkerProfile } from "../../swarm/worker-backend.js";
 import { getProfilePolicy } from "../../swarm/worker-backend.js";
 import { getLogger } from "../../util/logger.js";
@@ -121,7 +121,7 @@ export const claudeCodeTool: Tool = {
             type: "string",
             enum: ["general", "researcher", "coder", "reviewer"],
             description:
-              "Worker profile that scopes tool access. Defaults to general (backward compatible).",
+              "Worker profile that scopes tool access. Defaults to general.",
           },
         },
       },
@@ -203,12 +203,12 @@ export const claudeCodeTool: Tool = {
     const profilePolicy = getProfilePolicy(profileName);
 
     // Validate API key
-    const config = getConfig();
-    const apiKey = config.apiKeys.anthropic ?? process.env.ANTHROPIC_API_KEY;
+    const apiKey =
+      (await getSecureKeyAsync("anthropic")) ?? process.env.ANTHROPIC_API_KEY;
     if (!apiKey) {
       return {
         content:
-          "Error: No Anthropic API key configured. Set it via config or ANTHROPIC_API_KEY environment variable.",
+          "Error: No Anthropic API key configured. Set it via `keys set anthropic <key>` or configure it from the Settings page under API Keys.",
         isError: true,
       };
     }
@@ -261,7 +261,7 @@ export const claudeCodeTool: Tool = {
           return { behavior: "allow" as const };
         }
 
-        // Auto-approve safe read-only tools (backward compat for general profile)
+        // Auto-approve safe read-only tools (general profile default)
         if (AUTO_APPROVE_TOOLS.has(toolName)) {
           return { behavior: "allow" as const };
         }

package/src/tools/credentials/broker.ts

@@ -1,7 +1,7 @@
 import { v4 as uuid } from "uuid";
 
 import { credentialKey } from "../../security/credential-key.js";
-import { getSecureKey } from "../../security/secure-keys.js";
+import { getSecureKeyAsync } from "../../security/secure-keys.js";
 import { getLogger } from "../../util/logger.js";
 import type {
   AuthorizeRequest,
@@ -217,7 +217,7 @@ export class CredentialBroker {
     // Deletion is deferred until after a successful fill so the value survives
     // transient failures (e.g. stale element, page navigation, Playwright timeout).
     const transient = this.transientValues.get(storageKey);
-    const value = transient?.value ?? getSecureKey(storageKey);
+    const value = transient?.value ?? (await getSecureKeyAsync(storageKey));
     if (!value) {
       return {
         success: false,
@@ -302,7 +302,7 @@ export class CredentialBroker {
 
     const storageKey = credentialKey(request.service, request.field);
     const transient = this.transientValues.get(storageKey);
-    const value = transient?.value ?? getSecureKey(storageKey);
+    const value = transient?.value ?? (await getSecureKeyAsync(storageKey));
     if (!value) {
       return {
         success: false,
@@ -344,7 +344,9 @@ export class CredentialBroker {
    * never included in the result — the proxy reads it separately via
    * the secure key backend at injection time.
    */
-  serverUseById(request: ServerUseByIdRequest): ServerUseByIdResult {
+  async serverUseById(
+    request: ServerUseByIdRequest,
+  ): Promise<ServerUseByIdResult> {
     const resolved = resolveById(request.credentialId);
     if (!resolved) {
       return {
@@ -383,7 +385,7 @@ export class CredentialBroker {
 
     // Fail-closed: verify the secret value actually exists in secure storage.
     // Without this, downstream proxy code would attempt unauthenticated requests.
-    const value = getSecureKey(resolved.storageKey);
+    const value = await getSecureKeyAsync(resolved.storageKey);
     if (!value) {
       return {
         success: false,

package/src/tools/credentials/vault.ts

@@ -21,7 +21,7 @@ import { credentialKey } from "../../security/credential-key.js";
 import {
   deleteSecureKeyAsync,
   getSecureKeyAsync,
-  listSecureKeys,
+  listSecureKeysAsync,
   setSecureKeyAsync,
 } from "../../security/secure-keys.js";
 import { getLogger } from "../../util/logger.js";
@@ -367,10 +367,12 @@ class CredentialStoreTool implements Tool {
 
         const allMetadata = listCredentialMetadata();
         // Verify secrets still exist by reading all key names once (instead of
-        // per-entry getSecureKey calls that each re-read/re-derive the store).
+        // per-entry getSecureKeyAsync calls that each re-read/re-derive the store).
+        // Uses the async variant to include keys from both the primary backend
+        // (e.g. keychain) and the encrypted store (legacy keys).
         let secureKeySet: Set<string> | undefined;
         try {
-          secureKeySet = new Set(listSecureKeys());
+          secureKeySet = new Set(await listSecureKeysAsync());
         } catch (err) {
           log.error(
             { err },
@@ -765,7 +767,9 @@ class CredentialStoreTool implements Tool {
           if (dbApp) {
             if (!clientId) clientId = dbApp.clientId;
             if (!clientSecret) {
-              clientSecret = await getSecureKeyAsync(dbApp.clientSecretCredentialPath);
+              clientSecret = await getSecureKeyAsync(
+                dbApp.clientSecretCredentialPath,
+              );
             }
           }
         }
@@ -900,8 +904,9 @@ class CredentialStoreTool implements Tool {
             | "loopback"
             | "gateway"
             | null) ?? "gateway";
-        if (transport === "loopback" && descProviderRow.loopbackPort) {
-          redirectUri = `http://127.0.0.1:${descProviderRow.loopbackPort}/oauth/callback`;
+        const loopbackPort = descBehavior?.loopbackPort;
+        if (transport === "loopback" && loopbackPort) {
+          redirectUri = `http://localhost:${loopbackPort}/oauth/callback`;
         } else if (transport === "loopback") {
           redirectUri =
             "(automatic — no redirect URI needed, uses random localhost port)";