From 09d38c11a68f872da10cd4e8e8e31e05e5512a85 Mon Sep 17 00:00:00 2001 From: Yiliu Date: Tue, 7 Apr 2026 05:45:01 +0800 Subject: [PATCH] feat(infra): add session account list and switch commands Why: - Users need explicit session-scoped commands to inspect configured OpenAI accounts and switch the current session to a specific account - The plugin needed session context tracking so these commands can resolve the active session and update the correct binding What: - Added codex-account-list and codex-switch-account commands and registered them as primary tools - Added session context tracking to map session IDs to prompt cache keys for session-scoped account operations - Added account lookup helpers and tests covering account commands, session context, and runtime fetch behavior Impact: - Users can list configured accounts and switch the current session to a selected OpenAI account - Session-scoped account binding now has the context needed for command-driven switching --- index.ts | 163 +++++++++++++++++++ lib/accounts/manager.ts | 13 ++ lib/session-context.ts | 16 ++ package-lock.json | 4 +- test/codex-account-commands.test.ts | 244 ++++++++++++++++++++++++++++ test/runtime-fetch-parity.test.ts | 26 ++- test/session-context.test.ts | 19 +++ 7 files changed, 481 insertions(+), 4 deletions(-) create mode 100644 lib/session-context.ts create mode 100644 test/codex-account-commands.test.ts create mode 100644 test/session-context.test.ts diff --git a/index.ts b/index.ts index abb0c12..d522df7 100644 --- a/index.ts +++ b/index.ts @@ -34,6 +34,7 @@ import { AccountManager } from "./lib/accounts/index.js"; import type { ManagedAccount } from "./lib/accounts/index.js"; import { codexStatus } from "./lib/codex-status.js"; import { prefetchModels } from "./lib/models.js"; +import { SessionContextStore } from "./lib/session-context.js"; import { SessionBindingStore } from "./lib/session-bindings.js"; function extractModelFromBody(body: string | undefined): string | undefined { @@ -60,6 +61,65 @@ function extractPromptCacheKeyFromBody(body: string | undefined): string | undef } } +function extractSessionIdFromHeaders( + headers: unknown, +): string | undefined { + if (!headers) return undefined; + + if (headers instanceof Headers) { + const sessionId = headers.get("session_id")?.trim(); + return sessionId && sessionId.length > 0 ? sessionId : undefined; + } + + if (Array.isArray(headers)) { + for (const entry of headers) { + if (!Array.isArray(entry) || entry.length !== 2) continue; + const [key, value] = entry; + if (key === "session_id" && typeof value === "string") { + const trimmed = value.trim(); + return trimmed.length > 0 ? trimmed : undefined; + } + } + return undefined; + } + + if (!isRecord(headers)) return undefined; + const value = headers.session_id; + if (typeof value !== "string") return undefined; + const trimmed = value.trim(); + return trimmed.length > 0 ? trimmed : undefined; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null; +} + +function getStringField(record: Record, key: string): string | undefined { + const value = record[key]; + if (typeof value !== "string") return undefined; + const trimmed = value.trim(); + return trimmed.length > 0 ? trimmed : undefined; +} + +function extractSessionIdFromToolContext(context: unknown): string | undefined { + if (!isRecord(context)) return undefined; + + const direct = getStringField(context, "sessionID") || getStringField(context, "sessionId"); + if (direct) return direct; + + const session = context.session; + if (!isRecord(session)) return undefined; + return ( + getStringField(session, "id") || + getStringField(session, "sessionID") || + getStringField(session, "sessionId") + ); +} + +function resolveAccountLabel(account: ManagedAccount): string { + return account.email || `Account ${account.index + 1}`; +} + let lastToastAccountIndex: number | null = null; let lastToastTime = 0; const TOAST_DEBOUNCE_MS = 5000; @@ -196,6 +256,7 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { const sessionBindingStore = new SessionBindingStore(); sessionBindingStore.loadFromDisk(); + const sessionContextStore = new SessionContextStore(); const findAccountByIndex = (index: number): ManagedAccount | null => { return accountManager.getAllAccounts().find((acc) => acc.index === index) || null; @@ -610,6 +671,10 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { typeof init?.body === "string" ? (init.body as string) : undefined; const model = extractModelFromBody(requestBody); const sessionKey = extractPromptCacheKeyFromBody(requestBody); + const sessionId = extractSessionIdFromHeaders(init?.headers); + if (sessionId && sessionKey) { + sessionContextStore.setPromptCacheKey(sessionId, sessionKey); + } const account = await getSessionBoundAccount(sessionKey, model); if (!account) { @@ -688,6 +753,17 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { }, config: async (cfg) => { cfg.command = cfg.command || {}; + cfg.command["codex-account-list"] = { + template: + "Run the codex-account-list tool and output the result EXACTLY as returned by the tool, without any additional text or commentary.", + description: + "List all configured OpenAI accounts, marking the current session account and default account.", + }; + cfg.command["codex-switch-account"] = { + template: + 'Run the codex-switch-account tool with selector "$ARGUMENTS" and output the result EXACTLY as returned by the tool, without any additional text or commentary.', + description: "Switch the current session to a configured OpenAI account.", + }; cfg.command["codex-status"] = { template: "Run the codex-status tool and output the result EXACTLY as returned by the tool, without any additional text or commentary.", @@ -696,11 +772,98 @@ export const OpenAIAuthPlugin: Plugin = async ({ client }: PluginInput) => { cfg.experimental = cfg.experimental || {}; cfg.experimental.primary_tools = cfg.experimental.primary_tools || []; + if (!cfg.experimental.primary_tools.includes("codex-account-list")) { + cfg.experimental.primary_tools.push("codex-account-list"); + } + if (!cfg.experimental.primary_tools.includes("codex-switch-account")) { + cfg.experimental.primary_tools.push("codex-switch-account"); + } if (!cfg.experimental.primary_tools.includes("codex-status")) { cfg.experimental.primary_tools.push("codex-status"); } }, tool: { + "codex-account-list": tool({ + description: + "List all configured OpenAI accounts, marking the current session account and default account.", + args: {}, + async execute(_args, context) { + const accounts = accountManager.getAllAccounts(); + if (accounts.length === 0) { + return [ + "OpenAI Accounts", + "", + " Accounts: 0", + "", + "Add accounts:", + " opencode auth login", + ].join("\n"); + } + + const activeIndex = accountManager.getActiveAccount()?.index; + const sessionId = extractSessionIdFromToolContext(context); + const currentSessionKey = + sessionId && sessionContextStore.getPromptCacheKey(sessionId); + const currentSessionIndex = + currentSessionKey !== undefined + ? sessionBindingStore.get(currentSessionKey) + : undefined; + + const lines: string[] = ["OpenAI Accounts", ""]; + for (const account of accounts) { + const markers: string[] = []; + if (account.index === currentSessionIndex) { + markers.push("CURRENT_SESSION"); + } + if (account.index === activeIndex) { + markers.push("DEFAULT"); + } + if (markers.length === 0) { + markers.push("READY"); + } + + lines.push( + `${account.index + 1}. ${markers.join(" ")} ${resolveAccountLabel(account)} [${account.planType || "Unknown"}]`, + ); + } + + return lines.join("\n"); + }, + }), + "codex-switch-account": tool({ + description: "Switch the current session to a configured OpenAI account.", + args: { + selector: tool.schema.string().describe("Account index (1-based) or email"), + }, + async execute(args, context) { + const selector = args.selector.trim(); + if (!selector) { + return "Account selector is required."; + } + + const sessionId = extractSessionIdFromToolContext(context); + if (!sessionId) { + return "Could not determine the current session. Run this command from an active OpenCode session and try again."; + } + + const sessionKey = sessionContextStore.getPromptCacheKey(sessionId); + if (!sessionKey) { + return "Current session has no known prompt_cache_key yet. Send one normal model request in this session first, then retry."; + } + + const numericSelector = Number(selector); + const targetAccount = Number.isInteger(numericSelector) + ? accountManager.getAccountByIndex(numericSelector - 1) + : accountManager.findAccountByEmail(selector); + + if (!targetAccount) { + return `No configured OpenAI account matches \"${selector}\".`; + } + + sessionBindingStore.set(sessionKey, targetAccount.index); + return `Switched current session to account ${targetAccount.index + 1} (${resolveAccountLabel(targetAccount)}). Takes effect on the next request in this session.`; + }, + }), "codex-status": tool({ description: "List all configured OpenAI accounts and their current usage status.", args: {}, diff --git a/lib/accounts/manager.ts b/lib/accounts/manager.ts index a5661e6..7e1d8ac 100644 --- a/lib/accounts/manager.ts +++ b/lib/accounts/manager.ts @@ -189,6 +189,19 @@ export class AccountManager { return this.accounts; } + getAccountByIndex(index: number): ManagedAccount | null { + return this.accounts.find((account) => account.index === index) || null; + } + + findAccountByEmail(email: string): ManagedAccount | null { + const normalizedEmail = email.trim().toLowerCase(); + if (!normalizedEmail) return null; + return ( + this.accounts.find((account) => account.email?.toLowerCase() === normalizedEmail) || + null + ); + } + getAccountCount(): number { return this.accounts.length; } diff --git a/lib/session-context.ts b/lib/session-context.ts new file mode 100644 index 0000000..320e8c8 --- /dev/null +++ b/lib/session-context.ts @@ -0,0 +1,16 @@ +export class SessionContextStore { + private readonly promptCacheKeys = new Map(); + + setPromptCacheKey(sessionId: string, promptCacheKey: string): void { + const normalizedSessionId = sessionId.trim(); + const normalizedPromptCacheKey = promptCacheKey.trim(); + if (!normalizedSessionId || !normalizedPromptCacheKey) return; + this.promptCacheKeys.set(normalizedSessionId, normalizedPromptCacheKey); + } + + getPromptCacheKey(sessionId: string): string | undefined { + const normalizedSessionId = sessionId.trim(); + if (!normalizedSessionId) return undefined; + return this.promptCacheKeys.get(normalizedSessionId); + } +} diff --git a/package-lock.json b/package-lock.json index d1f67b6..6a8ebed 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,12 +1,12 @@ { "name": "opencode-openai-multi-auth", - "version": "5.0.5", + "version": "5.0.6", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "opencode-openai-multi-auth", - "version": "5.0.5", + "version": "5.0.6", "license": "MIT", "dependencies": { "@openauthjs/openauth": "^0.4.3", diff --git a/test/codex-account-commands.test.ts b/test/codex-account-commands.test.ts new file mode 100644 index 0000000..6e3d42e --- /dev/null +++ b/test/codex-account-commands.test.ts @@ -0,0 +1,244 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +function createToolContext(sessionID: string) { + return { + sessionID, + messageID: "msg_test", + agent: "test-agent", + abort: new AbortController().signal, + }; +} + +function requireAuthLoader(plugin: Awaited>) { + if (!plugin.auth?.loader) { + throw new Error("Expected plugin auth loader to be defined"); + } + return plugin.auth.loader; +} + +function requireTool( + plugin: Awaited>, + name: "codex-account-list" | "codex-switch-account", +) { + const toolEntry = plugin.tool?.[name]; + if (!toolEntry) { + throw new Error(`Expected tool ${name} to be defined`); + } + return toolEntry; +} + +vi.mock("@opencode-ai/plugin", () => ({ + tool: Object.assign((definition: unknown) => definition, { + schema: { + string: () => ({ + describe() { + return this; + }, + }), + }, + }), +})); + +const renderStatusMock = vi.fn(async () => []); + +vi.mock("../lib/codex-status.js", () => ({ + codexStatus: { + fetchFromBackend: vi.fn(async () => {}), + renderStatus: renderStatusMock, + }, +})); + +vi.mock("../lib/models.js", () => ({ + prefetchModels: vi.fn(async () => {}), +})); + +const accounts = [ + { + index: 0, + email: "alpha@example.com", + access: "token-a", + expires: Date.now() + 60_000, + accountId: "acct_alpha", + planType: "Plus", + }, + { + index: 1, + email: "beta@example.com", + access: "token-b", + expires: Date.now() + 60_000, + accountId: "acct_beta", + planType: "Pro", + }, +]; + +const sessionBindings = new Map(); + +vi.mock("../lib/accounts/index.js", () => { + class AccountManager { + async loadFromDisk() {} + async importFromOpenCodeAuth() {} + getAllAccounts() { + return accounts; + } + getAccountCount() { + return accounts.length; + } + getActiveAccount() { + return accounts[0]; + } + getAccountByIndex(index: number) { + return accounts.find((account) => account.index === index) || null; + } + findAccountByEmail(email: string) { + return accounts.find((account) => account.email === email) || null; + } + async getNextAvailableAccount() { + return accounts[0]; + } + async getNextAvailableAccountForNewSession() { + return accounts[0]; + } + async getNextAvailableAccountExcluding() { + return accounts[0]; + } + async ensureValidToken() { + return true; + } + markRateLimited() {} + markRefreshFailed() {} + async addAccount() {} + } + + return { AccountManager }; +}); + +vi.mock("../lib/session-bindings.js", () => { + class SessionBindingStore { + loadFromDisk() {} + get(key: string) { + return sessionBindings.get(key); + } + set(key: string, value: number) { + sessionBindings.set(key, value); + } + delete(key: string) { + sessionBindings.delete(key); + } + } + + return { SessionBindingStore }; +}); + +describe("codex account commands", () => { + beforeEach(() => { + sessionBindings.clear(); + renderStatusMock.mockClear(); + globalThis.fetch = vi.fn(async () => { + return new Response('data: {"type":"response.done"}\n\n', { + status: 200, + headers: { "content-type": "text/event-stream" }, + }); + }); + }); + + it("lists accounts with current session and default markers", async () => { + const { OpenAIAuthPlugin } = await import("../index.js"); + const plugin = await OpenAIAuthPlugin({ + client: { + auth: { set: vi.fn() }, + tui: { showToast: vi.fn() }, + }, + } as never); + + const loader = await requireAuthLoader(plugin)( + async () => ({ + type: "oauth", + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }) as never, + {} as never, + ); + + await loader.fetch("https://chatgpt.com/backend-api/responses", { + method: "POST", + headers: { + "content-type": "application/json", + session_id: "opencode-session-1", + }, + body: JSON.stringify({ + model: "gpt-5.2-codex", + prompt_cache_key: "ses_prompt_1", + input: [{ type: "message", role: "user", content: "hello" }], + }), + }); + + sessionBindings.set("ses_prompt_1", 1); + + const result = await requireTool(plugin, "codex-account-list").execute( + {}, + createToolContext("opencode-session-1"), + ); + + expect(result).toContain("1. DEFAULT alpha@example.com [Plus]"); + expect(result).toContain("2. CURRENT_SESSION beta@example.com [Pro]"); + }); + + it("switches the current session by 1-based index", async () => { + const { OpenAIAuthPlugin } = await import("../index.js"); + const plugin = await OpenAIAuthPlugin({ + client: { + auth: { set: vi.fn() }, + tui: { showToast: vi.fn() }, + }, + } as never); + + const loader = await requireAuthLoader(plugin)( + async () => ({ + type: "oauth", + access: "access-token", + refresh: "refresh-token", + expires: Date.now() + 60_000, + }) as never, + {} as never, + ); + + await loader.fetch("https://chatgpt.com/backend-api/responses", { + method: "POST", + headers: { + "content-type": "application/json", + session_id: "opencode-session-2", + }, + body: JSON.stringify({ + model: "gpt-5.2-codex", + prompt_cache_key: "ses_prompt_2", + input: [{ type: "message", role: "user", content: "hello" }], + }), + }); + + const result = await requireTool(plugin, "codex-switch-account").execute( + { selector: "2" }, + createToolContext("opencode-session-2"), + ); + + expect(result).toContain("Switched current session to account 2"); + expect(sessionBindings.get("ses_prompt_2")).toBe(1); + }); + + it("fails when the current session has no known prompt cache key", async () => { + const { OpenAIAuthPlugin } = await import("../index.js"); + const plugin = await OpenAIAuthPlugin({ + client: { + auth: { set: vi.fn() }, + tui: { showToast: vi.fn() }, + }, + } as never); + + const result = await requireTool(plugin, "codex-switch-account").execute( + { selector: "beta@example.com" }, + createToolContext("unknown-session"), + ); + + expect(result).toContain("Current session has no known prompt_cache_key yet"); + }); +}); diff --git a/test/runtime-fetch-parity.test.ts b/test/runtime-fetch-parity.test.ts index 64d03e1..9e526cc 100644 --- a/test/runtime-fetch-parity.test.ts +++ b/test/runtime-fetch-parity.test.ts @@ -1,9 +1,24 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; +function requireAuthLoader(plugin: Awaited>) { + if (!plugin.auth?.loader) { + throw new Error('Expected plugin auth loader to be defined'); + } + return plugin.auth.loader; +} + const transformRequestForCodexMock = vi.fn(); vi.mock('@opencode-ai/plugin', () => ({ - tool: (definition: unknown) => definition, + tool: Object.assign((definition: unknown) => definition, { + schema: { + string: () => ({ + describe() { + return this; + }, + }), + }, + }), })); vi.mock('../lib/request/fetch-helpers.js', async () => { @@ -16,6 +31,10 @@ vi.mock('../lib/request/fetch-helpers.js', async () => { }; }); +vi.mock('../lib/models.js', () => ({ + prefetchModels: vi.fn(async () => {}), +})); + vi.mock('../lib/accounts/index.js', () => { class AccountManager { private account = { @@ -43,6 +62,9 @@ vi.mock('../lib/accounts/index.js', () => { async getNextAvailableAccountForNewSession() { return this.account; } + async getNextAvailableAccountExcluding() { + return this.account; + } async ensureValidToken() { return true; } @@ -93,7 +115,7 @@ describe('Runtime fetch parity', () => { }, } as any); - const loader = await plugin.auth.loader( + const loader = await requireAuthLoader(plugin)( async () => ({ type: 'oauth', access: 'access-token', diff --git a/test/session-context.test.ts b/test/session-context.test.ts new file mode 100644 index 0000000..e9c0a5a --- /dev/null +++ b/test/session-context.test.ts @@ -0,0 +1,19 @@ +import { describe, expect, it } from "vitest"; +import { SessionContextStore } from "../lib/session-context.js"; + +describe("SessionContextStore", () => { + it("stores and returns prompt cache keys by session id", () => { + const store = new SessionContextStore(); + store.setPromptCacheKey("session-1", "ses_prompt_1"); + + expect(store.getPromptCacheKey("session-1")).toBe("ses_prompt_1"); + }); + + it("ignores blank values", () => { + const store = new SessionContextStore(); + store.setPromptCacheKey("", "ses_prompt_1"); + store.setPromptCacheKey("session-1", ""); + + expect(store.getPromptCacheKey("session-1")).toBeUndefined(); + }); +});