smashimo/oh-my-pi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge PR #998: feat(ai,coding-agent): credential_disabled extension event via multi-subscriber AuthStorage

Browse Source 
Converts AuthStorage from a single-subscriber to multi-subscriber model
so the SDK and extensions can both observe credential changes without
clobbering each other, and emits a new credential_disabled event when
storage permanently rejects a credential. Lets extensions prompt
re-auth instead of silently failing the next call.

Reconciliation + mismatch rejection covered for both the SDK and
runSubprocess paths.

Closes #998
This commit is contained in:
can1357
2026-05-13 02:29:03 +02:00
parent b0f115a56f 6872a73977
commit 3fc8cfc444
11 changed files with 1185 additions and 250 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/coding-agent/src/extensibility/extensions/runner.ts
+49 -1
View File
@@ -2,7 +2,7 @@
* Extension runner - executes extensions and manages their lifecycle.
*/
import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
import type { ImageContent, Model, ProviderResponseMetadata } from "@oh-my-pi/pi-ai";
import type { CredentialDisabledEvent, ImageContent, Model, ProviderResponseMetadata } from "@oh-my-pi/pi-ai";
import type { KeyId } from "@oh-my-pi/pi-tui";
import { logger } from "@oh-my-pi/pi-utils";
import type { ModelRegistry } from "../../config/model-registry";
@@ -69,6 +69,8 @@ export function __test_setExtensionHandlerTimeoutMs(timeoutMs: number): void {
const EXTENSION_HANDLER_TIMEOUT = Symbol("extensionHandlerTimeout");
const MAX_PENDING_CREDENTIAL_DISABLED = 32;
/**
* Events handled by the generic emit() method.
* Events with dedicated emitXxx() methods are excluded for stronger type safety.
@@ -185,6 +187,15 @@ export class ExtensionRunner {
#reloadHandler: () => Promise<void> = async () => {};
#shutdownHandler: ShutdownHandler = () => {};
#commandDiagnostics: Array<{ type: string; message: string; path: string }> = [];
#initialized = false;
/**
* Buffer for `credential_disabled` events received via {@link emitCredentialDisabled}
* before {@link initialize} has run. Drained through {@link emit} once initialize sets
* up the runtime context, so extension handlers see a populated UI/runtime context
* rather than the constructor's no-op default. Bounded at
* {@link MAX_PENDING_CREDENTIAL_DISABLED}; oldest entries are dropped under pressure.
*/
#pendingCredentialDisabled: CredentialDisabledEvent[] = [];
constructor(
private readonly extensions: Extension[],
@@ -237,6 +248,43 @@ export class ExtensionRunner {
}
this.#uiContext = uiContext ?? noOpUIContext;
this.#initialized = true;
// Drain events buffered by emitCredentialDisabled() before initialize ran. The
// spread adds the `type` discriminator — `event` is the pi-ai shape (no `type`).
for (const event of this.#pendingCredentialDisabled.splice(0)) {
this.emit({ type: "credential_disabled", ...event }).catch((error: unknown) => {
logger.warn("credential_disabled handler threw during initialize flush", {
provider: event.provider,
error: error instanceof Error ? error.message : String(error),
});
});
}
}
/**
* Forward a `credential_disabled` event from `AuthStorage` to extension handlers.
*
* If {@link initialize} has not yet run, the event is buffered and replayed once
* initialize wires the runtime/UI context. This matters because mode controllers
* (interactive, RPC, ACP, print, subagent) call `initialize()` AFTER `createAgentSession`
* returns, but `AuthStorage` can fire `credential_disabled` during startup model probes
* inside `createAgentSession()`. Without deferral, extension handlers would observe
* `hasUI=false`, an unset model, and no-op runtime actions on exactly the headline
* "OAuth invalid_grant during startup" path the event was designed to surface.
*
* Always returns; never throws. Errors from handlers are routed through
* {@link onError} via {@link emit}'s normal isolation.
*/
async emitCredentialDisabled(event: CredentialDisabledEvent): Promise<void> {
if (!this.#initialized) {
if (this.#pendingCredentialDisabled.length >= MAX_PENDING_CREDENTIAL_DISABLED) {
this.#pendingCredentialDisabled.shift();
}
this.#pendingCredentialDisabled.push(event);
return;
}
await this.emit({ type: "credential_disabled", ...event });
}
getUIContext(): ExtensionUIContext {
packages/coding-agent/src/extensibility/extensions/types.ts
+11
View File
@@ -624,6 +624,15 @@ export interface TodoReminderEvent {
maxAttempts: number;
}
/** Fired when AuthStorage automatically soft-disables a credential (e.g. OAuth `invalid_grant`). Not fired for user-initiated `remove()` or duplicate-credential dedup. */
export interface CredentialDisabledEvent {
type: "credential_disabled";
/** Provider id whose credential was disabled (e.g. "anthropic"). */
provider: string;
/** Verbatim error captured for forensics (truncated upstream). */
disabledCause: string;
}
// ============================================================================
// User Bash Events
// ============================================================================
@@ -831,6 +840,7 @@ export type ExtensionEvent =
| AutoRetryEndEvent
| TtsrTriggeredEvent
| TodoReminderEvent
| CredentialDisabledEvent
| UserBashEvent
| UserPythonEvent
| InputEvent
@@ -1012,6 +1022,7 @@ export interface ExtensionAPI {
on(event: "auto_retry_end", handler: ExtensionHandler<AutoRetryEndEvent>): void;
on(event: "ttsr_triggered", handler: ExtensionHandler<TtsrTriggeredEvent>): void;
on(event: "todo_reminder", handler: ExtensionHandler<TodoReminderEvent>): void;
on(event: "credential_disabled", handler: ExtensionHandler<CredentialDisabledEvent>): void;
on(event: "input", handler: ExtensionHandler<InputEvent, InputEventResult>): void;
on(event: "tool_call", handler: ExtensionHandler<ToolCallEvent, ToolCallEventResult>): void;
on(event: "tool_result", handler: ExtensionHandler<ToolResultEvent, ToolResultEventResult>): void;
packages/coding-agent/src/sdk.ts
+46 -5
View File
@@ -6,7 +6,7 @@ import {
INTENT_FIELD,
type ThinkingLevel,
} from "@oh-my-pi/pi-agent-core";
import type { Message, Model, SimpleStreamOptions } from "@oh-my-pi/pi-ai";
import type { CredentialDisabledEvent, Message, Model, SimpleStreamOptions } from "@oh-my-pi/pi-ai";
import {
getOpenAICodexTransportDetails,
prewarmOpenAICodexResponses,
@@ -670,10 +670,32 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
registerSshCleanup();
registerPythonCleanup();
// Use provided or create AuthStorage and ModelRegistry
const authStorage = options.authStorage ?? (await logger.time("discoverModels", discoverAuthStorage, agentDir));
const modelRegistry = options.modelRegistry ?? new ModelRegistry(authStorage);
// Pin authStorage to modelRegistry.authStorage: ModelRegistry.getApiKey() routes refresh
// failures through that instance, so any divergent storage handed to the bridge / mcpManager
// / session would silently miss credential_disabled events.
const modelRegistry =
options.modelRegistry ??
new ModelRegistry(options.authStorage ?? (await logger.time("discoverModels", discoverAuthStorage, agentDir)));
const authStorage = modelRegistry.authStorage;
if (options.authStorage && options.authStorage !== authStorage) {
throw new Error(
"options.authStorage and options.modelRegistry.authStorage must be the same instance when both are provided",
);
}
// Subscribe before any getApiKey() call so startup model probes can't fire a
// credential_disabled event past us. An embedder's constructor handler makes the
// listener set non-empty from construction, which defeats AuthStorage's no-listener
// buffer — so we can't rely on it to catch startup events for the extension runner.
const startupCredentialDisabledEvents: CredentialDisabledEvent[] = [];
let credentialDisabledTarget: ExtensionRunner | undefined;
let unsubscribeCredentialDisabled: (() => void) | undefined = authStorage.onCredentialDisabled(event => {
if (credentialDisabledTarget) {
// Discard return: any handler error is routed through runner.onError listeners.
void credentialDisabledTarget.emitCredentialDisabled(event);
} else {
startupCredentialDisabledEvents.push(event);
}
});
const settings = options.settings ?? (await logger.time("settings", Settings.init, { cwd, agentDir }));
logger.time("initializeWithSettings", initializeWithSettings, settings);
if (!options.modelRegistry) {
@@ -1277,6 +1299,20 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
);
}
if (extensionRunner) {
credentialDisabledTarget = extensionRunner;
for (const event of startupCredentialDisabledEvents.splice(0)) {
// Discard return: any handler error is routed through runner.onError listeners.
void extensionRunner.emitCredentialDisabled(event);
}
} else {
// No runner to forward to; release our subscription. The embedder's own
// onCredentialDisabled (if any) keeps firing through its own subscription.
startupCredentialDisabledEvents.length = 0;
unsubscribeCredentialDisabled?.();
unsubscribeCredentialDisabled = undefined;
}
const getSessionContext = () => ({
sessionManager,
modelRegistry,
@@ -1766,6 +1802,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
await originalDispose();
} finally {
agentRegistry.unregister(resolvedAgentId);
unsubscribeCredentialDisabled?.();
}
};
}
@@ -1901,6 +1938,10 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
eventBus,
};
} catch (error) {
// Release the subscription if the throw happened after install but before the
// dispose-wrap took ownership. Idempotent with dispose() — Set.delete is a no-op
// for already-removed listeners.
unsubscribeCredentialDisabled?.();
try {
if (hasSession) {
await session.dispose();
packages/coding-agent/src/task/executor.ts
+10 -3
View File
@@ -947,10 +947,17 @@ export async function runSubprocess(options: ExecutorOptions): Promise<SingleRes
try {
checkAbort();
const authStorage = options.authStorage ?? (await discoverAuthStorage());
checkAbort();
// Pin authStorage to modelRegistry.authStorage — mirrors the createAgentSession invariant.
const registryFromParent = options.modelRegistry !== undefined;
const modelRegistry = options.modelRegistry ?? new ModelRegistry(authStorage);
const modelRegistry =
options.modelRegistry ?? new ModelRegistry(options.authStorage ?? (await discoverAuthStorage()));
const authStorage = modelRegistry.authStorage;
if (options.authStorage && options.authStorage !== authStorage) {
throw new Error(
"options.authStorage and options.modelRegistry.authStorage must be the same instance when both are provided",
);
}
checkAbort();
if (!registryFromParent) {
await modelRegistry.refresh();
} else {