diff --git a/src/backend/api/client.ts b/src/backend/api/client.ts index 2d0dd2cd5b..1076a3109b 100644 --- a/src/backend/api/client.ts +++ b/src/backend/api/client.ts @@ -2,6 +2,7 @@ import { hostname } from "node:os"; import Letta from "@letta-ai/letta-client"; import { LETTA_CLOUD_API_URL } from "@/auth/oauth"; import { refreshAccessTokenSingleFlight } from "@/auth/oauth-refresh"; +import { getApiCredential } from "@/runtime-context"; import { type Settings, settingsManager } from "@/settings-manager"; import { trackBoundaryError } from "@/telemetry/error-reporting"; import { isDebugEnabled } from "@/utils/debug"; @@ -186,16 +187,19 @@ export async function getClient() { }, refreshToken: cachedTokens.refreshToken ?? baseSettings.refreshToken, }; + const runtimeApiKey = getApiCredential(); const settings = process.env.LETTA_API_KEY || + runtimeApiKey || cachedSettings.env?.LETTA_API_KEY || cachedSettings.refreshToken ? cachedSettings : await settingsManager.getSettingsWithSecureTokens(); - let apiKey = process.env.LETTA_API_KEY || settings.env?.LETTA_API_KEY; + let apiKey = + process.env.LETTA_API_KEY || runtimeApiKey || settings.env?.LETTA_API_KEY; - if (!process.env.LETTA_API_KEY) { + if (!process.env.LETTA_API_KEY && !runtimeApiKey) { if (apiKey) { // Keep the in-process cache current on every successful keychain read. _cachedApiKey = apiKey; @@ -210,6 +214,7 @@ export async function getClient() { // Check if token is expired and refresh if needed if ( !process.env.LETTA_API_KEY && + !runtimeApiKey && settings.tokenExpiresAt && settings.refreshToken ) { diff --git a/src/backend/api/request.ts b/src/backend/api/request.ts index d35e662744..f7f2620aed 100644 --- a/src/backend/api/request.ts +++ b/src/backend/api/request.ts @@ -1,4 +1,5 @@ import { LETTA_CLOUD_API_URL } from "@/auth/oauth"; +import { getApiCredential } from "@/runtime-context"; import { settingsManager } from "@/settings-manager"; import { getLettaCodeHeaders } from "./http-headers"; @@ -37,7 +38,11 @@ export async function getApiRequestConfig(): Promise { process.env.LETTA_BASE_URL || settings.env?.LETTA_BASE_URL || LETTA_CLOUD_API_URL, - apiKey: process.env.LETTA_API_KEY || settings.env?.LETTA_API_KEY || "", + apiKey: + process.env.LETTA_API_KEY || + getApiCredential() || + settings.env?.LETTA_API_KEY || + "", }; } diff --git a/src/runtime-context.ts b/src/runtime-context.ts index e49854a8ec..2840ae5ed9 100644 --- a/src/runtime-context.ts +++ b/src/runtime-context.ts @@ -31,6 +31,7 @@ export interface RuntimeContextSnapshot { } const runtimeContextStorage = new AsyncLocalStorage(); +const apiCredentialStorage = new AsyncLocalStorage(); export function getRuntimeContext(): RuntimeContextSnapshot | undefined { return runtimeContextStorage.getStore(); @@ -57,6 +58,19 @@ export function runOutsideRuntimeContext(fn: () => T): T { return runtimeContextStorage.exit(fn); } +/** + * Scope API work to the credential that authenticated the owning runtime. + * Kept separate from RuntimeContextSnapshot so secrets are never copied into + * tool context or diagnostics that spread the public runtime snapshot. + */ +export function runWithApiCredential(apiKey: string, fn: () => T): T { + return apiCredentialStorage.run(apiKey, fn); +} + +export function getApiCredential(): string | undefined { + return apiCredentialStorage.getStore(); +} + export function updateRuntimeContext( update: Partial, ): void { diff --git a/src/websocket/listener/auth-lifecycle.test.ts b/src/websocket/listener/auth-lifecycle.test.ts index f738f58977..bc3bbc931b 100644 --- a/src/websocket/listener/auth-lifecycle.test.ts +++ b/src/websocket/listener/auth-lifecycle.test.ts @@ -4,6 +4,8 @@ import type { AddressInfo } from "node:net"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { WebSocket, WebSocketServer } from "ws"; +import { clearAvailableModelsCache } from "@/agent/available-models"; +import { __testResetRemoteModelCatalog } from "@/agent/remote-model-catalog"; import { OAuthRefreshError, type TokenResponse } from "@/auth/oauth"; import { settingsManager } from "@/settings-manager"; import type { ControlRequest } from "@/types/protocol_v2"; @@ -48,6 +50,7 @@ describe("listener auth lifecycle", () => { const originalDisableMods = process.env.LETTA_DISABLE_MODS; const originalApiKey = process.env.LETTA_API_KEY; const originalBaseUrl = process.env.LETTA_BASE_URL; + const originalFetch = globalThis.fetch; const originalGetSettingsWithSecureTokens = settingsManager.getSettingsWithSecureTokens; const originalUpdateSettings = settingsManager.updateSettings; @@ -144,6 +147,9 @@ describe("listener auth lifecycle", () => { ); server.close(); __listenerAuthTestUtils.setOAuthDepsForTests(null); + clearAvailableModelsCache(); + __testResetRemoteModelCatalog(); + globalThis.fetch = originalFetch; settingsManager.getSettingsWithSecureTokens = originalGetSettingsWithSecureTokens; settingsManager.updateSettings = originalUpdateSettings; @@ -303,6 +309,109 @@ describe("listener auth lifecycle", () => { expect(settings.refreshToken).toBe("rotated-refresh-token"); }); + test("listener API commands keep using the credential that authenticated the socket", async () => { + await startClient(); + await waitFor( + () => connections.length === 1, + "initial socket did not open", + ); + expect(authorizations[0]).toBe("Bearer initial-access-token"); + + settings = { + ...settings, + env: { + ...settings.env, + LETTA_API_KEY: "replacement-access-token", + }, + }; + process.env.LETTA_BASE_URL = "https://api.test"; + clearAvailableModelsCache(); + __testResetRemoteModelCatalog(); + + const apiRequests: Array<{ url: string; authorization: string | null }> = + []; + globalThis.fetch = mock(async (input, init) => { + const request = input instanceof Request ? input : null; + const url = request?.url ?? String(input); + const headers = new Headers(request?.headers ?? init?.headers); + apiRequests.push({ + url, + authorization: headers.get("authorization"), + }); + return new Response("[]", { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + }) as unknown as typeof fetch; + + connections[0]?.send( + JSON.stringify({ + type: "list_models", + request_id: "models-session-credential", + force: true, + }), + ); + + await waitFor( + () => + getConnectionMessages(0).some( + (message) => + typeof message === "object" && + message !== null && + (message as { type?: string; request_id?: string }).type === + "list_models_response" && + (message as { type?: string; request_id?: string }).request_id === + "models-session-credential", + ), + "listener did not answer list_models", + ); + + expect(apiRequests.length).toBeGreaterThan(0); + expect( + apiRequests.some(({ url }) => new URL(url).pathname === "/v1/providers"), + ).toBe(true); + expect( + apiRequests.some( + ({ url }) => + new URL(url).pathname.startsWith("/v1/models") && + new URL(url).pathname !== "/v1/models/catalog", + ), + ).toBe(true); + expect( + new Set(apiRequests.map(({ authorization }) => authorization)), + ).toEqual(new Set(["Bearer initial-access-token"])); + + process.env.LETTA_API_KEY = "environment-access-token"; + apiRequests.length = 0; + clearAvailableModelsCache(); + __testResetRemoteModelCatalog(); + connections[0]?.send( + JSON.stringify({ + type: "list_models", + request_id: "models-environment-credential", + force: true, + }), + ); + + await waitFor( + () => + getConnectionMessages(0).some( + (message) => + typeof message === "object" && + message !== null && + (message as { type?: string; request_id?: string }).type === + "list_models_response" && + (message as { type?: string; request_id?: string }).request_id === + "models-environment-credential", + ), + "listener did not answer list_models with an environment credential", + ); + expect(apiRequests.length).toBeGreaterThan(0); + expect( + new Set(apiRequests.map(({ authorization }) => authorization)), + ).toEqual(new Set(["Bearer environment-access-token"])); + }); + test("transient relay closes preserve turn state queues and approval resolvers", async () => { const onDisconnected = mock(() => {}); const onNeedsReregister = mock(() => {}); diff --git a/src/websocket/listener/lifecycle.ts b/src/websocket/listener/lifecycle.ts index f71a9d2c32..25f5f0f67a 100644 --- a/src/websocket/listener/lifecycle.ts +++ b/src/websocket/listener/lifecycle.ts @@ -798,7 +798,6 @@ async function connectWithRetry( return; } const apiKey = auth.apiKey; - const url = new URL(opts.wsUrl); url.searchParams.set("deviceId", opts.deviceId); url.searchParams.set("connectionName", opts.connectionName); @@ -891,6 +890,7 @@ async function connectWithRetry( runtime, socket, connectionId: opts.connectionId, + apiCredential: apiKey, opts, processQueuedTurn, fileCommandSession, diff --git a/src/websocket/listener/message-router.ts b/src/websocket/listener/message-router.ts index c82184a639..aae5289c91 100644 --- a/src/websocket/listener/message-router.ts +++ b/src/websocket/listener/message-router.ts @@ -4,6 +4,7 @@ import { estimateSystemPromptTokensFromMemoryDir, setSystemPromptDoctorState, } from "@/cli/helpers/system-prompt-warning"; +import { runWithApiCredential } from "@/runtime-context"; import { settingsManager } from "@/settings-manager"; import type { AbortMessageCommand, @@ -107,6 +108,7 @@ type MessageRouterParams = { runtime: ListenerRuntime; socket: WebSocket; connectionId?: ListenerConnectionId; + apiCredential?: string; opts: StartListenerOptions; processQueuedTurn: ProcessQueuedTurn; fileCommandSession: FileCommandSession; @@ -179,6 +181,7 @@ export function createListenerMessageHandler( runtime, socket, connectionId: explicitConnectionId, + apiCredential, opts, processQueuedTurn, fileCommandSession, @@ -196,7 +199,7 @@ export function createListenerMessageHandler( } = params; const connectionId = explicitConnectionId ?? opts.connectionId; - return async (data: WebSocket.RawData): Promise => { + const handleMessage = async (data: WebSocket.RawData): Promise => { const raw = data.toString(); let parsedScope: ParsedRuntimeScope = null; @@ -867,4 +870,10 @@ export function createListenerMessageHandler( }); } }; + + if (!apiCredential) { + return handleMessage; + } + return (data) => + runWithApiCredential(apiCredential, () => handleMessage(data)); }