Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

103 changes: 103 additions & 0 deletions agent-bridge/src/__tests__/credentials-override.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
import { afterEach, beforeEach, describe, expect, it } from 'bun:test';

import { ClaudeCredentials } from '../common/auth/credentials.js';

describe('ClaudeCredentials settings override', () => {
let savedApiKey: string | undefined;
let savedSettingsFlag: string | undefined;
let savedOAuthToken: string | undefined;
let savedAnthropicAuthToken: string | undefined;

beforeEach(() => {
savedApiKey = process.env.ANTHROPIC_API_KEY;
savedSettingsFlag = process.env.ORBIT_SETTINGS_API_KEY;
savedOAuthToken = process.env.CLAUDE_CODE_OAUTH_TOKEN;
savedAnthropicAuthToken = process.env.ANTHROPIC_AUTH_TOKEN;
ClaudeCredentials.setCredentialOverride(undefined);
ClaudeCredentials.restoreEnvApiKeyFallback();
});

afterEach(() => {
ClaudeCredentials.setCredentialOverride(undefined);
ClaudeCredentials.restoreEnvApiKeyFallback();

if (savedApiKey !== undefined) {
process.env.ANTHROPIC_API_KEY = savedApiKey;
} else {
delete process.env.ANTHROPIC_API_KEY;
}

if (savedSettingsFlag !== undefined) {
process.env.ORBIT_SETTINGS_API_KEY = savedSettingsFlag;
} else {
delete process.env.ORBIT_SETTINGS_API_KEY;
}

if (savedOAuthToken !== undefined) {
process.env.CLAUDE_CODE_OAUTH_TOKEN = savedOAuthToken;
} else {
delete process.env.CLAUDE_CODE_OAUTH_TOKEN;
}

if (savedAnthropicAuthToken !== undefined) {
process.env.ANTHROPIC_AUTH_TOKEN = savedAnthropicAuthToken;
} else {
delete process.env.ANTHROPIC_AUTH_TOKEN;
}
});

it('initializes the override from the startup env flag', async () => {
process.env.ANTHROPIC_API_KEY = 'sk-ant-settings-startup-override';
process.env.ORBIT_SETTINGS_API_KEY = '1';

ClaudeCredentials.initOverrideFromEnv();

const result = await ClaudeCredentials.getCredentials();
expect(result.type).toBe('apikey');
expect(result.hasCredentials).toBe(true);
expect(result.token).toBe('sk-ant-settings-startup-override');
});

it('replaces and clears the override without leaving stale state behind', async () => {
process.env.ANTHROPIC_API_KEY = 'sk-ant-env-fallback';
delete process.env.ORBIT_SETTINGS_API_KEY;

ClaudeCredentials.setCredentialOverride('sk-ant-manual-override');
let result = await ClaudeCredentials.getCredentials();
expect(result.type).toBe('apikey');
expect(result.token).toBe('sk-ant-manual-override');

ClaudeCredentials.setCredentialOverride(undefined);
result = await ClaudeCredentials.getCredentials();

expect(result.token).not.toBe('sk-ant-manual-override');
if (result.type === 'apikey') {
expect(result.token).toBe('sk-ant-env-fallback');
} else {
expect(result.type).toBe('oauth');
}
});

it('restores the saved shell API key fallback exactly once', () => {
const shellFallback = 'sk-ant-shell-fallback';
process.env.ANTHROPIC_API_KEY = shellFallback;
process.env.CLAUDE_CODE_OAUTH_TOKEN = 'oauth-token';
process.env.ANTHROPIC_AUTH_TOKEN = 'legacy-token';

ClaudeCredentials.saveEnvApiKeyFallback(shellFallback);

delete process.env.ANTHROPIC_API_KEY;

expect(ClaudeCredentials.restoreEnvApiKeyFallback()).toBe(true);
expect(process.env.ANTHROPIC_API_KEY ?? '').toBe(shellFallback);
expect(process.env.CLAUDE_CODE_OAUTH_TOKEN).toBeUndefined();
expect(process.env.ANTHROPIC_AUTH_TOKEN).toBeUndefined();
expect(ClaudeCredentials.restoreEnvApiKeyFallback()).toBe(false);
});

it('returns false when no saved shell API key fallback exists', () => {
delete process.env.ANTHROPIC_API_KEY;

expect(ClaudeCredentials.restoreEnvApiKeyFallback()).toBe(false);
});
});
16 changes: 15 additions & 1 deletion agent-bridge/src/agent/core/agent.ts
Original file line number Diff line number Diff line change
Expand Up @@ -372,7 +372,12 @@ export class OrbitAgent {
* Called when the consumer catch block detects an auth error and explicitly awaits refresh.
*/
markNeedsSessionRestart(): void {
this._needsSessionRestart = true;
this.setNeedsSessionRestart(true);
}

/** Force a session restart on the next send. */
setNeedsSessionRestart(value: boolean): void {
this._needsSessionRestart = value;
}

constructor(config: OrbitAgentConfig = {}) {
Expand Down Expand Up @@ -1204,6 +1209,10 @@ You may also have access to iOS simulator automation tools via MCP when the \`mc
if (credentials.token) {
process.env.CLAUDE_CODE_OAUTH_TOKEN = credentials.token;
}
const envApiKey = process.env.ANTHROPIC_API_KEY;
if (envApiKey !== undefined && envApiKey !== '') {
ClaudeCredentials.saveEnvApiKeyFallback(envApiKey);
}
delete process.env.ANTHROPIC_API_KEY;
delete process.env.ANTHROPIC_AUTH_TOKEN;

Expand Down Expand Up @@ -1282,6 +1291,11 @@ You may also have access to iOS simulator automation tools via MCP when the \`mc
// (Code review: Opus cycle 4, #6)
if (result.error !== undefined) {
logger.warn({ error: result.error }, 'Credential refresh failed in pre-send check');
if (ClaudeCredentials.restoreEnvApiKeyFallback()) {
logger.info('Restored shell API key fallback, marking session for restart');
this.setNeedsSessionRestart(true);
return true;
}
return false;
}
// No refresh was needed — credentials still valid (API key or unexpired token)
Expand Down
23 changes: 23 additions & 0 deletions agent-bridge/src/agent/session/session-manager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -965,6 +965,14 @@ export class SessionManager extends Disposable {
}, 16);
}

/** Mark all active sessions so they restart with fresh credentials on the next send. */
markAllSessionsForRestart(): void {
for (const [sessionId, agent] of this.activeSessions.entries()) {
agent.setNeedsSessionRestart(true);
logger.info({ sessionId }, 'Marked session for restart after credential change');
}
}

/**
* Re-key all internal Maps from oldId to newId.
*
Expand Down Expand Up @@ -1127,6 +1135,15 @@ export class SessionManager extends Disposable {
resumeSessionAt: config?.resumeSessionAt,
forkSession: config?.forkSession,
onAuthFailure: (message: string) => {
if (ClaudeCredentials.restoreEnvApiKeyFallback()) {
agent.setNeedsSessionRestart(true);
logger.info(
{ sessionId: agent.effectiveSessionId },
'Restored shell API key fallback after OAuth auth failure'
);
return;
}

this._onAuthError.fire({
sessionId: agent.effectiveSessionId,
category: 'REFRESH_FAILED',
Expand Down Expand Up @@ -1779,6 +1796,12 @@ export class SessionManager extends Disposable {
const refreshResult = await ClaudeCredentials.refreshIfNeeded();
if (refreshResult.refreshed) {
agent.markNeedsSessionRestart();
} else if (ClaudeCredentials.restoreEnvApiKeyFallback()) {
agent.setNeedsSessionRestart(true);
logger.info(
{ sessionId },
'Restored shell API key fallback after active-query auth failure'
);
}
}

Expand Down
86 changes: 86 additions & 0 deletions agent-bridge/src/common/auth/credentials.ts
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,11 @@ export interface TokenRefreshResult {
/** Callback invoked when a scheduled auto-refresh fails */
export type AutoRefreshFailureCallback = (error: string) => void;

// Module-level credential override set from Orbit Settings.
// When present, this takes precedence over Keychain and shell env lookup.
let credentialOverride: CredentialResult | null = null;
let savedEnvApiKeyFallback: string | null = null;

function parseExpiryMs(expiresAt: number | string): number | null {
if (typeof expiresAt === 'number') {
return Number.isFinite(expiresAt) ? expiresAt : null;
Expand Down Expand Up @@ -293,6 +298,72 @@ function getApiKeyFromEnv(): string | null {
return apiKey;
}

/**
* Save the shell-provided API key before OAuth startup clears it from process.env.
*
* [warning] TESTED: Shell env fallback save/restore behavior is covered by
* integration tests. If you modify this, run: cd agent-bridge && bun test
* Test file: src/__tests__/credentials-override.test.ts
*/
function saveEnvApiKeyFallback(apiKey: string): void {
savedEnvApiKeyFallback = apiKey;
logger.debug('Saved shell env API key as OAuth fallback');
}

/**
* Restore the saved shell API key after OAuth auth failure. Returns true when
* a fallback key was restored and clears the saved copy so restore is one-shot.
*
* [warning] TESTED: Shell env fallback save/restore behavior is covered by
* integration tests. If you modify this, run: cd agent-bridge && bun test
* Test file: src/__tests__/credentials-override.test.ts
*/
function restoreEnvApiKeyFallback(): boolean {
if (savedEnvApiKeyFallback === null) {
return false;
}

process.env.ANTHROPIC_API_KEY = savedEnvApiKeyFallback;
delete process.env.CLAUDE_CODE_OAUTH_TOKEN;
delete process.env.ANTHROPIC_AUTH_TOKEN;

savedEnvApiKeyFallback = null;
logger.info('Restored shell env API key after OAuth failure');
return true;
}

/**
* Set or clear the credential override pushed from Orbit Settings.
*
* [warning] TESTED: Settings override precedence in this module is covered by
* integration tests. If you modify this, run: cd agent-bridge && bun test
* Test file: src/__tests__/credentials-override.test.ts
*/
function setCredentialOverride(apiKey: string | undefined): void {
if (apiKey !== undefined && apiKey !== '') {
credentialOverride = { type: 'apikey', hasCredentials: true, token: apiKey };
logger.info('Credential override set from Orbit Settings');
return;
}

credentialOverride = null;
logger.info('Credential override cleared; falling back to Keychain/env');
}

/**
* Initialize the credential override from startup env injected by Rust.
*
* [warning] TESTED: Startup override initialization is covered by integration
* tests. If you modify this, run: cd agent-bridge && bun test
* Test file: src/__tests__/credentials-override.test.ts
*/
function initOverrideFromEnv(): void {
if (process.env.ORBIT_SETTINGS_API_KEY === '1' && process.env.ANTHROPIC_API_KEY !== undefined) {
setCredentialOverride(process.env.ANTHROPIC_API_KEY);
logger.info('Credential override initialized from startup env');
}
}

/**
* Read the OAuth token expiry from macOS Keychain without performing a full
* credential validation. Returns the expiry timestamp in milliseconds, or
Expand Down Expand Up @@ -471,8 +542,19 @@ function scheduleAutoRefresh(onFailure: AutoRefreshFailureCallback): () => void
* the CLI subprocess to read from Keychain.
*
* @returns Object with credential info including the token value
*
* [warning] TESTED: Credential priority, including the Settings override path,
* is covered by integration tests. If you modify this, run:
* cd agent-bridge && bun test
* Test files: src/__tests__/credentials-refresh.test.ts,
* src/__tests__/credentials-override.test.ts
*/
async function getCredentials(): Promise<CredentialResult> {
if (credentialOverride !== null) {
logger.info('Using credential override from Orbit Settings');
return credentialOverride;
}

// Try OAuth token first
const oauthToken = await getOAuthTokenFromKeychain();

Expand Down Expand Up @@ -505,6 +587,10 @@ export const ClaudeCredentials = {
getOAuthTokenFromKeychain,
getApiKeyFromEnv,
getCredentials,
saveEnvApiKeyFallback,
restoreEnvApiKeyFallback,
setCredentialOverride,
initOverrideFromEnv,
refreshOAuthToken,
getTokenExpiry,
refreshIfNeeded,
Expand Down
21 changes: 21 additions & 0 deletions agent-bridge/src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ import {
touchSession,
} from './agent/session/session-storage.js';
import { CanvasSessionManager } from './canvas/index.js';
import { ClaudeCredentials } from './common/auth/credentials.js';
import { createLogger } from './common/logging/logger.js';
import { canEnableIOS } from './ios/ios-service.js';
import { BridgeRequestSchema } from './protocol/schemas.js';
Expand Down Expand Up @@ -133,6 +134,8 @@ async function main(): Promise<void> {
// This handles cases where the bridge process was killed and restarted
invalidateSessionCache();

ClaudeCredentials.initOverrideFromEnv();

let iosService: IOSService | undefined;
if (await canEnableIOS()) {
const { IOSService } = await import('./ios/ios-service.js');
Expand Down Expand Up @@ -543,6 +546,24 @@ async function handleRequest(
break;
}

case 'update_credentials': {
ClaudeCredentials.setCredentialOverride(request.apiKey);

delete process.env.CLAUDE_CODE_OAUTH_TOKEN;
delete process.env.ANTHROPIC_AUTH_TOKEN;
delete process.env.ANTHROPIC_API_KEY;
delete process.env.ORBIT_SETTINGS_API_KEY;

if (request.apiKey !== undefined && request.apiKey !== '') {
process.env.ANTHROPIC_API_KEY = request.apiKey;
process.env.ORBIT_SETTINGS_API_KEY = '1';
}

sessionManager.markAllSessionsForRestart();
sendResponse({ type: 'success', requestType: request.type });
break;
}

// Agent Definition Operations
case 'list_agents': {
const agents = listAgents(request.workspacePath);
Expand Down
9 changes: 9 additions & 0 deletions agent-bridge/src/protocol/protocol.ts
Original file line number Diff line number Diff line change
Expand Up @@ -175,6 +175,14 @@ export interface CleanupSessionsRequest {
maxAgeDays?: number;
}

/**
* Update the credential state used by the running sidecar.
*/
export interface UpdateCredentialsRequest {
type: 'update_credentials';
apiKey?: string;
}

// ============================================================================
// Agent Definition Requests
// ============================================================================
Expand Down Expand Up @@ -451,6 +459,7 @@ export type BridgeRequest =
| GetSDKSessionIdRequest
| GetStoredSessionRequest
| CleanupSessionsRequest
| UpdateCredentialsRequest
| ListAgentsRequest
| GetAgentRequest
| CreateAgentRequest
Expand Down
Loading