From f9bc96e96c1fd63c77991f3f93485032c9d9da88 Mon Sep 17 00:00:00 2001 From: Ogrodev Date: Sun, 14 Jun 2026 20:30:50 -0300 Subject: [PATCH] fix(coding-agent): harden profile auth shipping gaps --- docs/mcp-config.md | 22 +++-- packages/coding-agent/CHANGELOG.md | 2 + packages/coding-agent/src/cli.ts | 6 +- packages/coding-agent/src/cli/flag-tables.ts | 19 +++- .../coding-agent/src/cli/profile-bootstrap.ts | 35 +++++-- packages/coding-agent/src/mcp/manager.ts | 53 +++------- .../coding-agent/src/mcp/oauth-credentials.ts | 96 +++++++++++++++++++ packages/coding-agent/src/mcp/oauth-flow.ts | 21 ++-- .../controllers/mcp-command-controller.ts | 65 ++++++------- .../test/mcp-command-reauth.test.ts | 51 +++++++--- .../test/mcp-profile-auth-binding.test.ts | 78 ++++++++++++++- .../test/profile-bootstrap.test.ts | 14 +++ packages/utils/CHANGELOG.md | 2 +- packages/utils/README.md | 2 +- packages/utils/src/env.ts | 22 +---- packages/utils/src/worker-host.ts | 19 ++++ packages/utils/test/profiles.test.ts | 45 +++++++++ 17 files changed, 408 insertions(+), 144 deletions(-) create mode 100644 packages/coding-agent/src/mcp/oauth-credentials.ts create mode 100644 packages/utils/src/worker-host.ts diff --git a/docs/mcp-config.md b/docs/mcp-config.md index 28f25da84..aa2ae40e5 100644 --- a/docs/mcp-config.md +++ b/docs/mcp-config.md @@ -199,20 +199,22 @@ OMP understands two auth-related objects. Use this when OMP should remember how to rehydrate credentials for a server. You normally do not need to write this block: when OMP completes an OAuth flow -for an `http`/`sse` server it stores the credential in the active profile's -`agent.db` under a deterministic id derived from the server URL -(`mcp_oauth:`), with the refresh material embedded. Any config that points -at the same URL — including a *definition-only* entry in a shared project -`mcp.json` with no `auth` block at all — resolves the active profile's own -credential automatically. This is what makes project-scoped servers safe across +for an `http`/`sse` server it stores the credential under a deterministic id +derived from the active profile and server URL +(`mcp_oauth:profile::`), with the refresh material embedded. Any +config that points at the same URL — including a *definition-only* entry in a +shared project `mcp.json` with no `auth` block at all — resolves the active +profile's own credential automatically, including when auth storage is backed by +a shared auth broker. This is what makes project-scoped servers safe across profiles: commit the definition, and each profile authorizes (and stays signed in as) its own account via `/mcp reauth `. An explicit `credentialId` is -still honored when it resolves; if it points at another profile's row, OMP -falls back to the url-keyed binding. +still honored when it resolves; if it points at another profile's row, OMP falls +back to the profile-scoped url-keyed binding. `/mcp reauth` on a definition-only entry leaves the file untouched — the -credential (refresh material included) lives entirely in `agent.db`, so a -committed project config never picks up local auth state. An explicitly +credential (refresh material included) lives entirely in the active profile's +auth storage (local `agent.db` or broker), so a committed project config never +picks up local auth state. An explicitly configured `Authorization` header always wins over the url-keyed binding. The binding is per profile but not per project: once a profile has authorized diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index c4b3e6e25..8c23542c7 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -17,6 +17,8 @@ - Fixed selector-style UI components to honor `tui.select.up` and `tui.select.down` keybindings instead of hard-coding raw Up/Down arrow bytes ([#1535](https://github.com/can1357/oh-my-pi/issues/1535)). - Fixed a collapsed, still-streaming tool preview (an `eval`/`bash`/`ssh` box with output streaming in) reading as "weirdly truncated" — top border and head rows missing — once its box outgrew the viewport, snapping back to whole only while expanded with `ctrl+o` and breaking again when collapsed. A streaming preview was classified commit-unstable whenever collapsed, so the transcript offered none of its rows to native scrollback; once the box outgrew the window its head fell into the gap between the commit boundary and the window top, committed nowhere and repainted nowhere. The `provisionalPendingPreview` flag now applies only to the pending call preview (before any result) — once a streaming result exists the result renderer is the live, top-anchored shape and the block is commit-stable in both collapsed and expanded states, so its durable head always reaches scrollback. - Fixed a crash in subagent task execution and extensions when a string (instead of a string array) was returned or set for the system prompt. Gracefully wrap string values in arrays. +- Fixed profile bootstrap so an extension-shadowed `--plan` flag no longer swallows a following global `--profile`. +- Fixed MCP OAuth URL-keyed credentials to stay profile-scoped under shared auth-broker storage and to clear discovered definition-only server auth during `/mcp unauth`. ## [15.13.0] - 2026-06-14 diff --git a/packages/coding-agent/src/cli.ts b/packages/coding-agent/src/cli.ts index 20523022b..417f8cd66 100755 --- a/packages/coding-agent/src/cli.ts +++ b/packages/coding-agent/src/cli.ts @@ -23,6 +23,7 @@ import { setProfile, VERSION, } from "@oh-my-pi/pi-utils/dirs"; +import { declareWorkerHostEntry } from "@oh-my-pi/pi-utils/worker-host"; import { installProfileAlias, resolveProfileAliasCommandFromProcess } from "./cli/profile-alias"; import { extractProfileFlags } from "./cli/profile-bootstrap"; @@ -259,9 +260,8 @@ export async function runCli(argv: string[]): Promise { } // Declare this module as the worker-host entry now that the active profile - // is resolved — importing pi-utils/env earlier would snapshot the wrong - // agent directory's `.env`. - const { declareWorkerHostEntry } = await import("@oh-my-pi/pi-utils/env"); + // is resolved. The worker-host module is side-effect-free; importing + // `@oh-my-pi/pi-utils/env` here would snapshot the wrong agent `.env`. declareWorkerHostEntry(); if (resolvedArgv[0] === "--smoke-test") { diff --git a/packages/coding-agent/src/cli/flag-tables.ts b/packages/coding-agent/src/cli/flag-tables.ts index 9846de841..4e2354e12 100644 --- a/packages/coding-agent/src/cli/flag-tables.ts +++ b/packages/coding-agent/src/cli/flag-tables.ts @@ -85,8 +85,10 @@ const setResume: OptionalSetter = (result, value) => { }; /** - * Setters for flags that ALWAYS consume the next argv token, even when that - * token starts with `-`. + * Setters for flags with string values. Most built-ins consume the next argv + * token even when it starts with `-`; flags listed in + * {@link EXTENSION_SHADOWABLE_STRING_FLAGS} use extension-style consumption so + * a registered boolean extension can shadow them before profile bootstrap. */ export const STRING_SETTERS: Record = { "--cwd": (result, value) => { @@ -209,16 +211,25 @@ export const OPTIONAL_FLAGS: Record = { /** * Derived from {@link STRING_SETTERS}. A flag is in this set if and only if * it has a setter — by construction, drift between "the bootstrap thinks - * this flag consumes a value" and "the launch parser actually consumes one" - * is structurally impossible. + * this flag accepts a value" and "the launch parser can set one" is + * structurally impossible. */ export const STRING_VALUE_FLAGS: ReadonlySet = new Set(Object.keys(STRING_SETTERS)); +/** + * Built-in string flags known to be shadowed by bundled/common boolean + * extensions before extension metadata is available. They still accept a + * value-like successor for the built-in form (`--plan opus`), but a + * flag-looking successor remains a fresh flag (`--plan --profile work`). + */ +export const EXTENSION_SHADOWABLE_STRING_FLAGS: ReadonlySet = new Set(["--plan"]); + /** * Derived from {@link OPTIONAL_FLAGS}. Same single-source contract as * {@link STRING_VALUE_FLAGS}. */ export const OPTIONAL_VALUE_FLAGS: ReadonlySet = new Set(Object.keys(OPTIONAL_FLAGS)); + /** * Internal marker inserted by the profile bootstrap when removing `--profile` * or `--alias` would otherwise make the following value-like token become the diff --git a/packages/coding-agent/src/cli/profile-bootstrap.ts b/packages/coding-agent/src/cli/profile-bootstrap.ts index a2e2257f2..bc3b4a825 100644 --- a/packages/coding-agent/src/cli/profile-bootstrap.ts +++ b/packages/coding-agent/src/cli/profile-bootstrap.ts @@ -8,11 +8,12 @@ * crack open argv before the lazy command modules load. * * Because of that, this preparser must respect the same value-consumption - * contract as `args.ts`: known string-valued flags consume the next token - * unconditionally (so the value can legitimately start with `-`), and the - * optional-value flags (`--resume`, `--session`, `-r`, `--list-models`) - * consume the next token only when it doesn't look like another flag. Without - * this, `omp --system-prompt --profile foo` silently activates profile `foo` + * contract as `args.ts`: known string-valued flags usually consume the next + * token even when it starts with `-`, except for string flags that can be + * shadowed by preloaded boolean extensions (currently `--plan`). Optional-value + * flags (`--resume`, `--session`, `-r`) consume the next token only when it + * doesn't look like another flag. Without this, `omp --system-prompt --profile + * foo` silently activates profile `foo` * instead of passing the literal `--profile` to the system prompt and `foo` * as a positional message. * @@ -34,6 +35,7 @@ import { isSubcommand } from "../cli-commands"; import { + EXTENSION_SHADOWABLE_STRING_FLAGS, OPTIONAL_FLAGS, OPTIONAL_VALUE_FLAGS, PROFILE_BOOTSTRAP_BOUNDARY_ARG, @@ -57,7 +59,12 @@ function isUnknownLongValueCandidate(arg: string): boolean { function needsBoundaryAfterGlobalStrip(stripped: readonly string[]): boolean { const previous = stripped[stripped.length - 1]; - return previous !== undefined && (OPTIONAL_VALUE_FLAGS.has(previous) || isUnknownLongValueCandidate(previous)); + return ( + previous !== undefined && + (OPTIONAL_VALUE_FLAGS.has(previous) || + EXTENSION_SHADOWABLE_STRING_FLAGS.has(previous) || + isUnknownLongValueCandidate(previous)) + ); } export interface ProfileBootstrapResult { @@ -152,6 +159,22 @@ export function extractProfileFlags(argv: readonly string[]): ProfileBootstrapRe continue; } + // Known string flags normally consume flag-looking values (for example + // `--system-prompt --profile foo` means the system prompt is literally + // `--profile`). A small allow-list of built-ins can be shadowed by boolean + // extensions before extension metadata is loaded; those mirror extension + // consumption here so `--plan --profile work` still activates `work`. + if (EXTENSION_SHADOWABLE_STRING_FLAGS.has(arg)) { + canDispatchSubcommand = false; + stripped.push(arg); + const next = argv[index + 1]; + if (next !== undefined && !next.startsWith("-")) { + stripped.push(next); + index += 1; + } + continue; + } + // Forward both the flag and its value untouched so the downstream parser // gets exactly what the user typed. Critical for `--system-prompt // --profile foo`: the bootstrap must NOT interpret `--profile` here, it diff --git a/packages/coding-agent/src/mcp/manager.ts b/packages/coding-agent/src/mcp/manager.ts index 6a452a022..74a7c0943 100644 --- a/packages/coding-agent/src/mcp/manager.ts +++ b/packages/coding-agent/src/mcp/manager.ts @@ -27,7 +27,12 @@ import { unsubscribeFromResources, } from "./client"; import { loadAllMCPConfigs, validateServerConfig } from "./config"; -import { type MCPStoredOAuthCredential, mcpOAuthCredentialId, refreshMCPOAuthToken } from "./oauth-flow"; +import { + lookupMcpOAuthCredential, + type MCPOAuthCredentialLookup, + selectMcpOAuthRefreshMaterial, +} from "./oauth-credentials"; +import { type MCPStoredOAuthCredential, refreshMCPOAuthToken } from "./oauth-flow"; import type { MCPToolDetails } from "./tool-bridge"; import { DeferredMCPTool, MCPTool } from "./tool-bridge"; import type { MCPToolCache } from "./tool-cache"; @@ -403,7 +408,10 @@ export class MCPManager { // Gate on a resolvable managed credential, not on the auth block: // definition-only configs (url-keyed fallback) get Bearer injection // too and need the same mid-session refresh hook. - if (connection.transport instanceof HttpTransport && this.#lookupOAuthCredential(config)) { + if ( + connection.transport instanceof HttpTransport && + lookupMcpOAuthCredential(this.#authStorage, config) + ) { connection.transport.onAuthError = async () => { const refreshed = await this.#resolveAuthConfig(config, { forceRefresh: true }); if (refreshed.type === "http" || refreshed.type === "sse") { @@ -931,7 +939,7 @@ export class MCPManager { // Wire auth refresh for HTTP transports, and reconnect for any transport. // Same gate as connectServers: any resolvable managed credential. - if (connection.transport instanceof HttpTransport && this.#lookupOAuthCredential(config)) { + if (connection.transport instanceof HttpTransport && lookupMcpOAuthCredential(this.#authStorage, config)) { connection.transport.onAuthError = async () => { const refreshed = await this.#resolveAuthConfig(config, { forceRefresh: true }); if (refreshed.type === "http" || refreshed.type === "sse") { @@ -1173,40 +1181,6 @@ export class MCPManager { }; } - /** - * Look up the OAuth credential for a config: an explicit `auth.credentialId` - * wins; otherwise (or when the pointer misses this profile's storage) fall - * back to the deterministic per-URL id. The fallback is what lets a shared - * project-scope server definition resolve per-profile credentials. - */ - #lookupOAuthCredential( - config: MCPServerConfig, - ): { credentialId: string; credential: MCPStoredOAuthCredential } | undefined { - if (!this.#authStorage) return undefined; - const auth = config.auth; - if (auth && auth.type !== "oauth") return undefined; - if (auth?.credentialId) { - const credential = this.#authStorage.get(auth.credentialId); - if (credential?.type === "oauth") { - return { credentialId: auth.credentialId, credential }; - } - } - if (config.type !== "http" && config.type !== "sse") return undefined; - if (!config.url) return undefined; - // Never clobber an explicitly configured Authorization header. An auth - // block whose pointer resolved returns above (legacy semantics); the - // url-keyed fallback always yields to a pinned header. - if (Object.keys(config.headers ?? {}).some(h => h.toLowerCase() === "authorization")) { - return undefined; - } - const urlKeyId = mcpOAuthCredentialId(config.url); - const credential = this.#authStorage.get(urlKeyId); - if (credential?.type === "oauth") { - return { credentialId: urlKeyId, credential }; - } - return undefined; - } - /** * Resolve OAuth credentials and shell commands in config. * `oauth: false` skips credential injection (reauth's unauthenticated probe); @@ -1219,7 +1193,8 @@ export class MCPManager { let resolved: MCPServerConfig = { ...config }; const auth = config.auth; - const lookup = opts?.oauth !== false ? this.#lookupOAuthCredential(config) : undefined; + const lookup: MCPOAuthCredentialLookup | undefined = + opts?.oauth !== false ? lookupMcpOAuthCredential(this.#authStorage, config) : undefined; if (lookup && this.#authStorage) { const { credentialId } = lookup; try { @@ -1230,7 +1205,7 @@ export class MCPManager { // config auth block. Never mix the two: a shared file's auth block // can belong to another profile, whose client the grant is NOT // bound to. - const material = credential.tokenUrl ? credential : auth; + const material = selectMcpOAuthRefreshMaterial(credential, auth); const tokenUrl = material?.tokenUrl; const clientId = material?.clientId; const clientSecret = material?.clientSecret; diff --git a/packages/coding-agent/src/mcp/oauth-credentials.ts b/packages/coding-agent/src/mcp/oauth-credentials.ts new file mode 100644 index 000000000..f1d0a5b46 --- /dev/null +++ b/packages/coding-agent/src/mcp/oauth-credentials.ts @@ -0,0 +1,96 @@ +import { expandEnvVarsDeep } from "../discovery/helpers"; +import type { AuthStorage } from "../session/auth-storage"; +import { isManagedMCPOAuthCredentialId, type MCPStoredOAuthCredential, mcpOAuthCredentialId } from "./oauth-flow"; +import type { MCPAuthConfig, MCPServerConfig } from "./types"; + +export interface MCPOAuthCredentialLookup { + credentialId: string; + credential: MCPStoredOAuthCredential; +} + +export type MCPOAuthRefreshMaterial = MCPStoredOAuthCredential | MCPAuthConfig | undefined; + +export function mcpOAuthCredentialIdsForServerUrl(serverUrl: string | undefined): string[] { + if (!serverUrl) return []; + const ids: string[] = []; + for (const url of [expandEnvVarsDeep(serverUrl), serverUrl]) { + const id = mcpOAuthCredentialId(url); + if (!ids.includes(id)) ids.push(id); + } + return ids; +} + +export function hasMcpAuthorizationHeader(config: MCPServerConfig): boolean { + if (config.type !== "http" && config.type !== "sse") return false; + return Object.keys(config.headers ?? {}).some(header => header.toLowerCase() === "authorization"); +} + +export function lookupMcpOAuthCredentialForServer( + authStorage: AuthStorage | null | undefined, + auth: MCPAuthConfig | undefined, + serverUrl: string | undefined, + options: { allowUrlKeyedFallback?: boolean } = {}, +): MCPOAuthCredentialLookup | undefined { + if (!authStorage) return undefined; + if (auth && auth.type !== "oauth") return undefined; + const urlKeyedCredentialIds = mcpOAuthCredentialIdsForServerUrl(serverUrl); + if ( + auth?.credentialId && + (!auth.credentialId.startsWith("mcp_oauth:profile:") || urlKeyedCredentialIds.includes(auth.credentialId)) + ) { + const credential = authStorage.get(auth.credentialId); + if (credential?.type === "oauth") { + return { credentialId: auth.credentialId, credential }; + } + } + if (options.allowUrlKeyedFallback === false) return undefined; + for (const credentialId of urlKeyedCredentialIds) { + const credential = authStorage.get(credentialId); + if (credential?.type === "oauth") { + return { credentialId, credential }; + } + } + return undefined; +} + +export function lookupMcpOAuthCredential( + authStorage: AuthStorage | null | undefined, + config: MCPServerConfig, +): MCPOAuthCredentialLookup | undefined { + const auth = config.auth; + if (config.type !== "http" && config.type !== "sse") { + return lookupMcpOAuthCredentialForServer(authStorage, auth, undefined); + } + if (hasMcpAuthorizationHeader(config)) { + return lookupMcpOAuthCredentialForServer(authStorage, auth, config.url, { allowUrlKeyedFallback: false }); + } + return lookupMcpOAuthCredentialForServer(authStorage, auth, config.url); +} + +export function selectMcpOAuthRefreshMaterial( + credential: MCPStoredOAuthCredential, + auth: MCPAuthConfig | undefined, +): MCPOAuthRefreshMaterial { + return credential.tokenUrl ? credential : auth; +} + +export async function removeManagedMcpOAuthCredential( + authStorage: AuthStorage, + credentialId: string | undefined, +): Promise { + if (!isManagedMCPOAuthCredentialId(credentialId)) return false; + if (authStorage.get(credentialId)?.type !== "oauth") return false; + await authStorage.remove(credentialId); + return true; +} + +export async function removeManagedMcpOAuthCredentials( + authStorage: AuthStorage, + credentialIds: readonly (string | undefined)[], +): Promise { + let removed = false; + for (const credentialId of credentialIds) { + removed = (await removeManagedMcpOAuthCredential(authStorage, credentialId)) || removed; + } + return removed; +} diff --git a/packages/coding-agent/src/mcp/oauth-flow.ts b/packages/coding-agent/src/mcp/oauth-flow.ts index cad94cef1..7954e6f1c 100644 --- a/packages/coding-agent/src/mcp/oauth-flow.ts +++ b/packages/coding-agent/src/mcp/oauth-flow.ts @@ -9,23 +9,24 @@ import type { OAuthCallbackFlowOptions } from "@oh-my-pi/pi-ai/oauth/callback-se import { OAuthCallbackFlow } from "@oh-my-pi/pi-ai/oauth/callback-server"; import type { OAuthController, OAuthCredentials } from "@oh-my-pi/pi-ai/oauth/types"; import type { FetchImpl } from "@oh-my-pi/pi-ai/types"; +import { getActiveProfile } from "@oh-my-pi/pi-utils/dirs"; import type { OAuthCredential } from "../session/auth-storage"; -/** Credential-id prefix for OMP-managed MCP OAuth credentials keyed by server URL. */ +/** Credential-id prefix for OMP-managed MCP OAuth credentials keyed by profile and server URL. */ const MCP_OAUTH_URL_CREDENTIAL_PREFIX = "mcp_oauth:"; /** - * Deterministic credential id for an MCP server URL. + * Deterministic credential id for an MCP server URL scoped to an OMP profile. * - * The id is identical across profiles and projects while each profile's - * agent.db holds its own row under it, so a server *definition* in a shared - * project `mcp.json` resolves to per-profile credentials instead of one - * profile's random `mcp_oauth__` pointer clobbering the others. - * The URL is used verbatim (query string included) because it can carry - * tenant selectors such as `?project_ref=`. + * Local profile stores are already separate, but auth-broker storage shares one + * provider namespace across profiles. Including the profile in the provider key + * keeps a shared project `mcp.json` definition from making profile B overwrite + * or read profile A's OAuth row for the same server URL. The URL is used + * verbatim (query string included) because it can carry tenant selectors such + * as `?project_ref=`. */ -export function mcpOAuthCredentialId(serverUrl: string): string { - return `${MCP_OAUTH_URL_CREDENTIAL_PREFIX}${serverUrl}`; +export function mcpOAuthCredentialId(serverUrl: string, profile: string | undefined = getActiveProfile()): string { + return `${MCP_OAUTH_URL_CREDENTIAL_PREFIX}profile:${profile ?? "default"}:${serverUrl}`; } /** Whether a credential id was minted by OMP's MCP OAuth flows (either era). */ diff --git a/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts b/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts index e42b15ebb..9b32edae8 100644 --- a/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts +++ b/packages/coding-agent/src/modes/controllers/mcp-command-controller.ts @@ -19,11 +19,12 @@ import { updateMCPServer, } from "../../mcp/config-writer"; import { - isManagedMCPOAuthCredentialId, - MCPOAuthFlow, - type MCPStoredOAuthCredential, - mcpOAuthCredentialId, -} from "../../mcp/oauth-flow"; + lookupMcpOAuthCredentialForServer, + mcpOAuthCredentialIdsForServerUrl, + removeManagedMcpOAuthCredential, + removeManagedMcpOAuthCredentials, +} from "../../mcp/oauth-credentials"; +import { MCPOAuthFlow, type MCPStoredOAuthCredential, mcpOAuthCredentialId } from "../../mcp/oauth-flow"; import { clearSmitheryApiKey, createSmitheryCliAuthSession, @@ -871,26 +872,6 @@ export class MCPCommandController { }; } - async #removeManagedOAuthCredential(credentialId: string | undefined): Promise { - if (!isManagedMCPOAuthCredentialId(credentialId)) return; - await this.ctx.session.modelRegistry.authStorage.remove(credentialId); - } - - #lookupExistingMcpOAuthCredential( - auth: MCPAuthConfig | undefined, - serverUrl: string | undefined, - ): MCPStoredOAuthCredential | undefined { - const authStorage = this.ctx.session.modelRegistry.authStorage; - if (auth?.type === "oauth" && auth.credentialId) { - const credential = authStorage.get(auth.credentialId); - if (credential?.type === "oauth") return credential; - } - if (!serverUrl) return undefined; - const credential = authStorage.get(mcpOAuthCredentialId(serverUrl)); - if (credential?.type === "oauth") return credential; - return undefined; - } - #stripOAuthAuth(config: MCPServerConfig): MCPServerConfig { const next = { ...config } as MCPServerConfig & { auth?: MCPAuthConfig }; delete next.auth; @@ -1484,23 +1465,34 @@ export class MCPCommandController { } const currentAuth = (found.config as MCPServerConfig & { auth?: MCPAuthConfig }).auth; - if (found.discovered && currentAuth?.type !== "oauth") { - this.#showMessage(["", theme.fg("muted", `No stored OAuth auth to remove for "${name}".`), ""].join("\n")); - return; - } + const authStorage = this.ctx.session.modelRegistry.authStorage; if (currentAuth?.type === "oauth") { - await this.#removeManagedOAuthCredential(currentAuth.credentialId); + await removeManagedMcpOAuthCredential(authStorage, currentAuth.credentialId); } // Also drop this profile's url-keyed binding so the server is truly // signed out even when the config carries no auth block. Runtime // discovery expands `${...}` URL values before MCPManager looks up the // deterministic credential row, so unauth must clear that same key. + let removedUrlKeyedCredential = false; if ((found.config.type === "http" || found.config.type === "sse") && found.config.url) { - const runtimeServerUrl = expandEnvVarsDeep(found.config.url); - await this.#removeManagedOAuthCredential(mcpOAuthCredentialId(runtimeServerUrl)); - if (runtimeServerUrl !== found.config.url) { - await this.#removeManagedOAuthCredential(mcpOAuthCredentialId(found.config.url)); + removedUrlKeyedCredential = await removeManagedMcpOAuthCredentials( + authStorage, + mcpOAuthCredentialIdsForServerUrl(found.config.url), + ); + } + + if (found.discovered && currentAuth?.type !== "oauth") { + if (!removedUrlKeyedCredential) { + this.#showMessage( + ["", theme.fg("muted", `No stored OAuth auth to remove for "${name}".`), ""].join("\n"), + ); + return; } + await this.#reloadMCP(); + this.#showMessage( + ["", theme.fg("success", `- Cleared auth for "${name}" (${found.scope} config)`), ""].join("\n"), + ); + return; } const updated = this.#stripOAuthAuth(found.config); @@ -1534,6 +1526,7 @@ export class MCPCommandController { } const currentAuth = (found.config as MCPServerConfig & { auth?: MCPAuthConfig }).auth; + const authStorage = this.ctx.session.modelRegistry.authStorage; const baseConfig = this.#stripOAuthAuth(found.config); const runtimeBaseConfig = expandEnvVarsDeep(baseConfig); // Resolve endpoints first: this fails fast for stdio transports and @@ -1548,7 +1541,7 @@ export class MCPCommandController { // writes it to auth.clientSecret); DCR secrets are embedded in the // stored credential and never echoed back into config files. const configuredClientId = found.config.oauth?.clientId ?? currentAuth?.clientId; - const existingCredential = this.#lookupExistingMcpOAuthCredential(currentAuth, serverUrl); + const existingCredential = lookupMcpOAuthCredentialForServer(authStorage, currentAuth, serverUrl)?.credential; const flowClientId = oauth.clientId ?? configuredClientId ?? existingCredential?.clientId ?? ""; const storedClientSecret = existingCredential?.clientId === flowClientId ? existingCredential.clientSecret : undefined; @@ -1582,7 +1575,7 @@ export class MCPCommandController { // after success so cancelling the browser step leaves the previous // session signed in. if (currentAuth?.type === "oauth" && currentAuth.credentialId !== oauthResult.credentialId) { - await this.#removeManagedOAuthCredential(currentAuth.credentialId); + await removeManagedMcpOAuthCredential(authStorage, currentAuth.credentialId); } // Definition-only entries resolve through the url-keyed binding alone; diff --git a/packages/coding-agent/test/mcp-command-reauth.test.ts b/packages/coding-agent/test/mcp-command-reauth.test.ts index 5325a7b66..a27725a0b 100644 --- a/packages/coding-agent/test/mcp-command-reauth.test.ts +++ b/packages/coding-agent/test/mcp-command-reauth.test.ts @@ -9,7 +9,7 @@ import * as oauthFlow from "@oh-my-pi/pi-coding-agent/mcp/oauth-flow"; import type { MCPServerConfig } from "@oh-my-pi/pi-coding-agent/mcp/types"; import { MCPCommandController } from "@oh-my-pi/pi-coding-agent/modes/controllers/mcp-command-controller"; import { initTheme } from "@oh-my-pi/pi-coding-agent/modes/theme/theme"; -import { getConfigRootDir, getProjectDir, setAgentDir, setProjectDir } from "@oh-my-pi/pi-utils"; +import { getConfigRootDir, getMCPConfigPath, getProjectDir, setAgentDir, setProjectDir } from "@oh-my-pi/pi-utils"; const RAW_SERVER_URL = `https://\${MCP_HOST}/mcp`; const EXPANDED_SERVER_URL = "https://mcp.example.com/mcp"; @@ -34,9 +34,18 @@ function restoreEnvValue(name: string, value: string | undefined): void { Bun.env[name] = value; process.env[name] = value; } -function createController(authStorage: AuthStorage) { +function createController(authStorage: AuthStorage, mcpManagerOverrides: Record = {}) { const showError = vi.fn(); const prepareConfig = vi.fn(async (config: MCPServerConfig) => config); + const mcpManager = { + prepareConfig, + disconnectAll: vi.fn(async () => {}), + discoverAndConnect: vi.fn(async () => ({ errors: new Map() })), + getTools: vi.fn(() => []), + waitForConnection: vi.fn(async () => {}), + getConnectionStatus: vi.fn(() => "connected"), + ...mcpManagerOverrides, + }; const controller = new MCPCommandController({ chatContainer: { addChild: vi.fn() }, present: vi.fn(), @@ -53,17 +62,10 @@ function createController(authStorage: AuthStorage) { refreshMCPTools: vi.fn(), modelRegistry: { authStorage }, }, - mcpManager: { - prepareConfig, - disconnectAll: vi.fn(async () => {}), - discoverAndConnect: vi.fn(async () => ({ errors: new Map() })), - getTools: vi.fn(() => []), - waitForConnection: vi.fn(async () => {}), - getConnectionStatus: vi.fn(() => "connected"), - }, + mcpManager, } as never); - return { controller, showError, prepareConfig }; + return { controller, showError, prepareConfig, mcpManager }; } describe("/mcp auth commands", () => { @@ -226,4 +228,31 @@ describe("/mcp auth commands", () => { expect(savedUrl).toBe(RAW_SERVER_URL); expect(savedServer?.auth).toBeUndefined(); }); + + test("clears url-keyed auth for discovered definition-only servers", async () => { + const authStorage = new AuthStorage(new SqliteAuthCredentialStore(new Database(":memory:"))); + await authStorage.reload(); + await authStorage.set(oauthFlow.mcpOAuthCredentialId(EXPANDED_SERVER_URL), { + type: "oauth", + access: "discovered-access", + refresh: "discovered-refresh", + expires: Date.now() + 3_600_000, + }); + const { controller, showError } = createController(authStorage, { + getServerConfig: vi.fn(() => ({ type: "http", url: EXPANDED_SERVER_URL })), + getSource: vi.fn(() => ({ provider: "test", path: "/tmp/discovered.json" })), + }); + + await controller.handle("/mcp unauth discovered"); + + expect(showError).not.toHaveBeenCalled(); + expect(authStorage.get(oauthFlow.mcpOAuthCredentialId(EXPANDED_SERVER_URL))).toBeUndefined(); + const userConfigPath = getMCPConfigPath("user", projectDir); + const userConfig = JSON.parse( + await Bun.file(userConfigPath) + .text() + .catch(() => "{}"), + ) as TestConfigFile; + expect(userConfig.mcpServers?.discovered).toBeUndefined(); + }); }); diff --git a/packages/coding-agent/test/mcp-profile-auth-binding.test.ts b/packages/coding-agent/test/mcp-profile-auth-binding.test.ts index 854d1adda..fc5813a5b 100644 --- a/packages/coding-agent/test/mcp-profile-auth-binding.test.ts +++ b/packages/coding-agent/test/mcp-profile-auth-binding.test.ts @@ -3,9 +3,10 @@ * * A server *definition* may live in a shared project `mcp.json` while each * profile holds its own credential row in agent.db under the deterministic - * `mcp_oauth:` id. Before this scheme, the random `auth.credentialId` - * written into the shared file pointed at exactly one profile's row, so two - * profiles reauthorizing the same project server clobbered each other. + * `mcp_oauth:profile::` id. Before this scheme, the random + * `auth.credentialId` written into the shared file pointed at exactly one + * profile's row, so two profiles reauthorizing the same project server + * clobbered each other. */ import { Database } from "bun:sqlite"; import { afterEach, beforeEach, describe, expect, test, vi } from "bun:test"; @@ -14,6 +15,7 @@ import { MCPManager } from "@oh-my-pi/pi-coding-agent/mcp/manager"; import * as oauthFlow from "@oh-my-pi/pi-coding-agent/mcp/oauth-flow"; import { mcpOAuthCredentialId } from "@oh-my-pi/pi-coding-agent/mcp/oauth-flow"; import type { MCPServerConfig } from "@oh-my-pi/pi-coding-agent/mcp/types"; +import { getActiveProfile, setProfile } from "@oh-my-pi/pi-utils/dirs"; const SERVER_URL = "https://mcp.example.com/mcp"; const URL_KEY_ID = mcpOAuthCredentialId(SERVER_URL); @@ -26,8 +28,10 @@ function authorizationHeader(config: MCPServerConfig): string | undefined { describe("per-profile MCP OAuth binding", () => { let manager: MCPManager; let authStorage: AuthStorage; + let originalProfile: string | undefined; beforeEach(async () => { + originalProfile = getActiveProfile(); const store = new SqliteAuthCredentialStore(new Database(":memory:")); authStorage = new AuthStorage(store); await authStorage.reload(); @@ -36,9 +40,77 @@ describe("per-profile MCP OAuth binding", () => { }); afterEach(() => { + setProfile(originalProfile); vi.restoreAllMocks(); }); + test("scopes url-keyed credentials by active profile in a shared auth namespace", async () => { + const workKey = mcpOAuthCredentialId(SERVER_URL, "work"); + const personalKey = mcpOAuthCredentialId(SERVER_URL, "personal"); + expect(workKey).not.toBe(personalKey); + await authStorage.set(workKey, { + type: "oauth", + access: "work-token", + refresh: "r", + expires: Date.now() + 3_600_000, + }); + await authStorage.set(personalKey, { + type: "oauth", + access: "personal-token", + refresh: "r", + expires: Date.now() + 3_600_000, + }); + + setProfile("work"); + expect(authorizationHeader(await manager.prepareConfig({ type: "http", url: SERVER_URL }))).toBe( + "Bearer work-token", + ); + + setProfile("personal"); + expect(authorizationHeader(await manager.prepareConfig({ type: "http", url: SERVER_URL }))).toBe( + "Bearer personal-token", + ); + }); + + test("ignores another profile's explicit profile-scoped credentialId in shared storage", async () => { + const workKey = mcpOAuthCredentialId(SERVER_URL, "work"); + const personalKey = mcpOAuthCredentialId(SERVER_URL, "personal"); + await authStorage.set(workKey, { + type: "oauth", + access: "work-token", + refresh: "r", + expires: Date.now() + 3_600_000, + }); + + setProfile("personal"); + expect( + authorizationHeader( + await manager.prepareConfig({ + type: "http", + url: SERVER_URL, + auth: { type: "oauth", credentialId: workKey }, + }), + ), + ).toBeUndefined(); + + await authStorage.set(personalKey, { + type: "oauth", + access: "personal-token", + refresh: "r", + expires: Date.now() + 3_600_000, + }); + + expect( + authorizationHeader( + await manager.prepareConfig({ + type: "http", + url: SERVER_URL, + auth: { type: "oauth", credentialId: workKey }, + }), + ), + ).toBe("Bearer personal-token"); + }); + test("resolves the url-keyed credential when the file's credentialId belongs to another profile", async () => { // This profile authed the server (url-keyed row exists), but the shared // project file still carries a credentialId minted by a different profile. diff --git a/packages/coding-agent/test/profile-bootstrap.test.ts b/packages/coding-agent/test/profile-bootstrap.test.ts index a939a5bd3..d3a3255e6 100644 --- a/packages/coding-agent/test/profile-bootstrap.test.ts +++ b/packages/coding-agent/test/profile-bootstrap.test.ts @@ -36,6 +36,20 @@ describe("extractProfileFlags", () => { expect(result.argv).toEqual(["--approval-mode", "--profile", "foo", "bar"]); }); + it("honors extension-shadowed --plan before a global profile", () => { + const extracted = extractProfileFlags(["--plan", "--profile", "work", "follow up"]); + expect(extracted).toEqual({ + argv: ["--plan", PROFILE_BOOTSTRAP_BOUNDARY_ARG, "follow up"], + profile: "work", + aliasName: undefined, + }); + + const parsed = parseArgs(extracted.argv, new Map([["plan", { type: "boolean" }]])); + expect(parsed.unknownFlags.get("plan")).toBe(true); + expect(parsed.plan).toBeUndefined(); + expect(parsed.messages).toEqual(["follow up"]); + }); + it("still extracts --profile after an unrelated string-valued flag", () => { // Mirror image: when the user does mean to activate a profile *after* // a string-valued flag, we must skip past the flag's value but still diff --git a/packages/utils/CHANGELOG.md b/packages/utils/CHANGELOG.md index 153f2f689..e48d0db28 100644 --- a/packages/utils/CHANGELOG.md +++ b/packages/utils/CHANGELOG.md @@ -25,7 +25,7 @@ - Added the `path-tree` module (`buildPathTree`, `walkPathTree`, `formatGroupedPaths`, `isUrlLikePath`), moved from the coding agent's grouped file output so compaction file lists can share the same prefix-folded directory-tree rendering; `formatGroupedPaths` gains an optional `annotate` callback for per-file suffixes - Restored `PI_DEBUG_STARTUP` streaming startup markers: `logger.time` now writes a synchronous `[startup] :start` / `:done` / `:fail` stderr line per phase (independent of `PI_TIMING`), so a startup that hangs hard still names the phase it is stuck in — the `PI_TIMING` tree only prints after startup completes and is structurally unable to diagnose a hang. The CLI runner emits `cli:load:` markers around each lazily-imported command module for the same reason. - Added `logger.openSpanPath()`: ops of the currently-open timing-span chain (root → deepest), used by the coding agent's startup watchdog to name the in-flight phase of a stalled startup. -- Added `declareWorkerHostEntry()` / `workerHostEntry()` (env): self-dispatching CLI entrypoints declare `Bun.main` as the worker host so worker spawn sites can re-enter the single entry module with `WorkerOptions.argv` selectors across source, npm-bundle, and compiled distributions +- Added `declareWorkerHostEntry()` / `workerHostEntry()` in the side-effect-free `worker-host` module (also re-exported from `env`): self-dispatching CLI entrypoints declare `Bun.main` as the worker host so worker spawn sites can re-enter the single entry module with `WorkerOptions.argv` selectors across source, npm-bundle, and compiled distributions - Added `getAuthBrokerSnapshotCachePath()` with `OMP_AUTH_BROKER_SNAPSHOT_CACHE` override support for isolating the encrypted broker snapshot cache. - Added color helpers `colorLuma` (perceptual luma), `relativeLuminance` (WCAG, linearized sRGB), and `hslToHex` to the color utilities. The luminance helpers parse `#rgb`/`#rrggbb` hex and 256-color palette indices, returning `undefined` for unparseable values. - Added `peekFileEnds`, a single-open head-and-tail file peek helper that reuses the head bytes for the tail when the file fits the head window. diff --git a/packages/utils/README.md b/packages/utils/README.md index 9cdb2c4cd..fb5f308a7 100644 --- a/packages/utils/README.md +++ b/packages/utils/README.md @@ -15,7 +15,7 @@ Shared utilities for [oh-my-pi](https://github.com/can1357/oh-my-pi) packages. Z | `which` | `$which()` binary lookup with caching | | `fetch-retry` | `fetch` with retry/backoff policies | | `fs-error` | Errno guards (`isEnoent` and friends) | -| `env` | Environment plumbing, worker-host entry contract (`workerHostEntry`) | +| `env` / `worker-host` | Environment plumbing and side-effect-free worker-host entry contract (`workerHostEntry`) | | `abortable` / `async` | AbortSignal-aware stream/promise helpers | | `peek-file` | Read the first N bytes of a file with pooled buffers | | `frontmatter`, `glob`, `mime`, `temp`, `format`, `color`, `snowflake`, `tab-spacing`, `path-tree`, `sanitize-text` | Smaller single-purpose helpers | diff --git a/packages/utils/src/env.ts b/packages/utils/src/env.ts index 8ba2fc661..9e3bf1349 100644 --- a/packages/utils/src/env.ts +++ b/packages/utils/src/env.ts @@ -3,6 +3,8 @@ import * as os from "node:os"; import * as path from "node:path"; import { getAgentDir, getConfigRootDir } from "./dirs"; +export * from "./worker-host"; + const ENV_NAME_RE = /^[A-Za-z_][A-Za-z0-9_]*$/; /** @@ -172,26 +174,6 @@ export function isCompiledBinary(): boolean { return url.includes("$bunfs") || url.includes("~BUN") || url.includes("%7EBUN"); } -/** - * Main-module path declared by self-dispatching CLI entrypoints — entries - * whose top-level argv handling routes hidden `__omp_*` worker selectors. - * Worker spawn sites re-enter this module via `new Worker(entry, { argv })`, - * so every distribution (source, npm bundle, compiled binary) needs exactly - * one JavaScript entrypoint. Never set under `bun test`, SDK embedding, or - * standalone package bins — those hosts load worker modules directly. - */ -let workerHostMain: string | null = null; - -/** Called by CLI entrypoints whose main module dispatches worker argv selectors. */ -export function declareWorkerHostEntry(): void { - workerHostMain = Bun.main; -} - -/** Main-module path of the self-dispatching CLI host, or null outside it. */ -export function workerHostEntry(): string | null { - return workerHostMain; -} - const TRUTHY: Dict = { "1": true, Y: true, diff --git a/packages/utils/src/worker-host.ts b/packages/utils/src/worker-host.ts new file mode 100644 index 000000000..13eb39d1a --- /dev/null +++ b/packages/utils/src/worker-host.ts @@ -0,0 +1,19 @@ +/** + * Main-module path declared by self-dispatching CLI entrypoints — entries + * whose top-level argv handling routes hidden `__omp_*` worker selectors. + * Worker spawn sites re-enter this module via `new Worker(entry, { argv })`, + * so every distribution (source, npm bundle, compiled binary) needs exactly + * one JavaScript entrypoint. Never set under `bun test`, SDK embedding, or + * standalone package bins — those hosts load worker modules directly. + */ +let workerHostMain: string | null = null; + +/** Called by CLI entrypoints whose main module dispatches worker argv selectors. */ +export function declareWorkerHostEntry(): void { + workerHostMain = Bun.main; +} + +/** Main-module path of the self-dispatching CLI host, or null outside it. */ +export function workerHostEntry(): string | null { + return workerHostMain; +} diff --git a/packages/utils/test/profiles.test.ts b/packages/utils/test/profiles.test.ts index 223005752..7b02f0352 100644 --- a/packages/utils/test/profiles.test.ts +++ b/packages/utils/test/profiles.test.ts @@ -318,6 +318,51 @@ describe("dirs module import behavior", () => { await fs.rm(root, { recursive: true, force: true }); } }); + it("exposes worker-host without loading agent env", async () => { + const root = await fs.mkdtemp(path.join(os.tmpdir(), "pi-utils-worker-host-import-")); + try { + const workerHostUrl = import.meta.resolve("@oh-my-pi/pi-utils/worker-host"); + const agentDir = path.join(root, "agent"); + await fs.mkdir(agentDir, { recursive: true }); + await Bun.write(path.join(agentDir, ".env"), "OMP_WORKER_HOST_PROBE=from-agent-env\n"); + const probePath = path.join(root, "probe.ts"); + await Bun.write( + probePath, + [ + `import { declareWorkerHostEntry, workerHostEntry } from ${JSON.stringify(workerHostUrl)};`, + "declareWorkerHostEntry();", + "process.stdout.write(JSON.stringify({", + " envProbe: process.env.OMP_WORKER_HOST_PROBE ?? null,", + " hostDeclared: workerHostEntry() === Bun.main,", + "}));", + ].join("\n"), + ); + + const childEnv: Record = { + ...process.env, + PI_CODING_AGENT_DIR: agentDir, + }; + delete childEnv.OMP_WORKER_HOST_PROBE; + const proc = Bun.spawn([process.execPath, probePath], { + stdout: "pipe", + stderr: "pipe", + env: childEnv, + }); + const [stdout, stderr, exitCode] = await Promise.all([ + readStream(proc.stdout as ReadableStream), + readStream(proc.stderr as ReadableStream), + proc.exited, + ]); + + expect(exitCode, stderr).toBe(0); + expect(JSON.parse(stdout)).toEqual({ + envProbe: null, + hostDeclared: true, + }); + } finally { + await fs.rm(root, { recursive: true, force: true }); + } + }); it("ignores inherited profile agent dir when OMP_PROFILE explicitly selects default", async () => { const root = await fs.mkdtemp(path.join(os.tmpdir(), "pi-utils-dirs-default-profile-"));