diff --git a/src/agents/auth-profiles/oauth.fallback-to-main-agent.test.ts b/src/agents/auth-profiles/oauth.fallback-to-main-agent.test.ts new file mode 100644 index 000000000..d37d1a8c3 --- /dev/null +++ b/src/agents/auth-profiles/oauth.fallback-to-main-agent.test.ts @@ -0,0 +1,164 @@ +import fs from "node:fs/promises"; +import path from "node:path"; +import os from "node:os"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { resolveApiKeyForProfile } from "./oauth.js"; +import { ensureAuthProfileStore } from "./store.js"; +import type { AuthProfileStore } from "./types.js"; + +describe("resolveApiKeyForProfile fallback to main agent", () => { + const previousStateDir = process.env.CLAWDBOT_STATE_DIR; + const previousAgentDir = process.env.CLAWDBOT_AGENT_DIR; + const previousPiAgentDir = process.env.PI_CODING_AGENT_DIR; + let tmpDir: string; + let mainAgentDir: string; + let secondaryAgentDir: string; + + beforeEach(async () => { + tmpDir = await fs.mkdtemp(path.join(os.tmpdir(), "oauth-fallback-test-")); + mainAgentDir = path.join(tmpDir, "agents", "main", "agent"); + secondaryAgentDir = path.join(tmpDir, "agents", "kids", "agent"); + await fs.mkdir(mainAgentDir, { recursive: true }); + await fs.mkdir(secondaryAgentDir, { recursive: true }); + + // Set environment variables so resolveClawdbotAgentDir() returns mainAgentDir + process.env.CLAWDBOT_STATE_DIR = tmpDir; + process.env.CLAWDBOT_AGENT_DIR = mainAgentDir; + process.env.PI_CODING_AGENT_DIR = mainAgentDir; + }); + + afterEach(async () => { + vi.unstubAllGlobals(); + + // Restore original environment + if (previousStateDir === undefined) delete process.env.CLAWDBOT_STATE_DIR; + else process.env.CLAWDBOT_STATE_DIR = previousStateDir; + if (previousAgentDir === undefined) delete process.env.CLAWDBOT_AGENT_DIR; + else process.env.CLAWDBOT_AGENT_DIR = previousAgentDir; + if (previousPiAgentDir === undefined) delete process.env.PI_CODING_AGENT_DIR; + else process.env.PI_CODING_AGENT_DIR = previousPiAgentDir; + + await fs.rm(tmpDir, { recursive: true, force: true }); + }); + + it("falls back to main agent credentials when secondary agent token is expired and refresh fails", async () => { + const profileId = "anthropic:claude-cli"; + const now = Date.now(); + const expiredTime = now - 60 * 60 * 1000; // 1 hour ago + const freshTime = now + 60 * 60 * 1000; // 1 hour from now + + // Write expired credentials for secondary agent + const secondaryStore: AuthProfileStore = { + version: 1, + profiles: { + [profileId]: { + type: "oauth", + provider: "anthropic", + access: "expired-access-token", + refresh: "expired-refresh-token", + expires: expiredTime, + }, + }, + }; + await fs.writeFile( + path.join(secondaryAgentDir, "auth-profiles.json"), + JSON.stringify(secondaryStore), + ); + + // Write fresh credentials for main agent + const mainStore: AuthProfileStore = { + version: 1, + profiles: { + [profileId]: { + type: "oauth", + provider: "anthropic", + access: "fresh-access-token", + refresh: "fresh-refresh-token", + expires: freshTime, + }, + }, + }; + await fs.writeFile(path.join(mainAgentDir, "auth-profiles.json"), JSON.stringify(mainStore)); + + // Mock fetch to simulate OAuth refresh failure + const fetchSpy = vi.fn(async () => { + return new Response(JSON.stringify({ error: "invalid_grant" }), { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + }); + vi.stubGlobal("fetch", fetchSpy); + + // Load the secondary agent's store (will merge with main agent's store) + const loadedSecondaryStore = ensureAuthProfileStore(secondaryAgentDir); + + // Call resolveApiKeyForProfile with the secondary agent's expired credentials + // This should: + // 1. Try to refresh the expired token (fails due to mocked fetch) + // 2. Fall back to main agent's fresh credentials + // 3. Copy those credentials to the secondary agent + const result = await resolveApiKeyForProfile({ + store: loadedSecondaryStore, + profileId, + agentDir: secondaryAgentDir, + }); + + expect(result).not.toBeNull(); + expect(result?.apiKey).toBe("fresh-access-token"); + expect(result?.provider).toBe("anthropic"); + + // Verify the credentials were copied to the secondary agent + const updatedSecondaryStore = JSON.parse( + await fs.readFile(path.join(secondaryAgentDir, "auth-profiles.json"), "utf8"), + ) as AuthProfileStore; + expect(updatedSecondaryStore.profiles[profileId]).toMatchObject({ + access: "fresh-access-token", + expires: freshTime, + }); + }); + + it("throws error when both secondary and main agent credentials are expired", async () => { + const profileId = "anthropic:claude-cli"; + const now = Date.now(); + const expiredTime = now - 60 * 60 * 1000; // 1 hour ago + + // Write expired credentials for both agents + const expiredStore: AuthProfileStore = { + version: 1, + profiles: { + [profileId]: { + type: "oauth", + provider: "anthropic", + access: "expired-access-token", + refresh: "expired-refresh-token", + expires: expiredTime, + }, + }, + }; + await fs.writeFile( + path.join(secondaryAgentDir, "auth-profiles.json"), + JSON.stringify(expiredStore), + ); + await fs.writeFile(path.join(mainAgentDir, "auth-profiles.json"), JSON.stringify(expiredStore)); + + // Mock fetch to simulate OAuth refresh failure + const fetchSpy = vi.fn(async () => { + return new Response(JSON.stringify({ error: "invalid_grant" }), { + status: 400, + headers: { "Content-Type": "application/json" }, + }); + }); + vi.stubGlobal("fetch", fetchSpy); + + const loadedSecondaryStore = ensureAuthProfileStore(secondaryAgentDir); + + // Should throw because both agents have expired credentials + await expect( + resolveApiKeyForProfile({ + store: loadedSecondaryStore, + profileId, + agentDir: secondaryAgentDir, + }), + ).rejects.toThrow(/OAuth token refresh failed/); + }); +}); diff --git a/src/agents/auth-profiles/oauth.ts b/src/agents/auth-profiles/oauth.ts index 4138cda94..d7b3360de 100644 --- a/src/agents/auth-profiles/oauth.ts +++ b/src/agents/auth-profiles/oauth.ts @@ -4,7 +4,7 @@ import lockfile from "proper-lockfile"; import type { ClawdbotConfig } from "../../config/config.js"; import { refreshChutesTokens } from "../chutes-oauth.js"; import { refreshQwenPortalCredentials } from "../../providers/qwen-portal-oauth.js"; -import { AUTH_STORE_LOCK_OPTIONS } from "./constants.js"; +import { AUTH_STORE_LOCK_OPTIONS, log } from "./constants.js"; import { formatAuthDoctorHint } from "./doctor.js"; import { ensureAuthStoreFile, resolveAuthStorePath } from "./paths.js"; import { suggestOAuthProfileIdForLegacyDefault } from "./repair.js"; @@ -196,6 +196,32 @@ export async function resolveApiKeyForProfile(params: { // keep original error } } + + // Fallback: if this is a secondary agent, try using the main agent's credentials + if (params.agentDir) { + try { + const mainStore = ensureAuthProfileStore(undefined); // main agent (no agentDir) + const mainCred = mainStore.profiles[profileId]; + if (mainCred?.type === "oauth" && Date.now() < mainCred.expires) { + // Main agent has fresh credentials - copy them to this agent and use them + refreshedStore.profiles[profileId] = { ...mainCred }; + saveAuthProfileStore(refreshedStore, params.agentDir); + log.info("inherited fresh OAuth credentials from main agent", { + profileId, + agentDir: params.agentDir, + expires: new Date(mainCred.expires).toISOString(), + }); + return { + apiKey: buildOAuthApiKey(mainCred.provider, mainCred), + provider: mainCred.provider, + email: mainCred.email, + }; + } + } catch { + // keep original error if main agent fallback also fails + } + } + const message = error instanceof Error ? error.message : String(error); const hint = formatAuthDoctorHint({ cfg, diff --git a/ui/src/ui/app-gateway.ts b/ui/src/ui/app-gateway.ts index d9a267a98..0df25bbdf 100644 --- a/ui/src/ui/app-gateway.ts +++ b/ui/src/ui/app-gateway.ts @@ -127,6 +127,12 @@ export function connectGateway(host: GatewayHost) { host.lastError = null; host.hello = hello; applySnapshot(host, hello); + // Reset orphaned chat run state from before disconnect. + // Any in-flight run's final event was lost during the disconnect window. + host.chatRunId = null; + (host as unknown as { chatStream: string | null }).chatStream = null; + (host as unknown as { chatStreamStartedAt: number | null }).chatStreamStartedAt = null; + resetToolStream(host as unknown as Parameters[0]); void loadAssistantIdentity(host as unknown as ClawdbotApp); void loadAgents(host as unknown as ClawdbotApp); void loadNodes(host as unknown as ClawdbotApp, { quiet: true });