@oh-my-pi/pi-coding-agent 15.1.1 → 15.1.3

package/src/session/agent-session.ts

@@ -18,12 +18,15 @@ import * as fs from "node:fs";
 import * as path from "node:path";
 import { scheduler } from "node:timers/promises";
 import {
+	type AfterToolCallContext,
+	type AfterToolCallResult,
 	type Agent,
 	AgentBusyError,
 	type AgentEvent,
 	type AgentMessage,
 	type AgentState,
 	type AgentTool,
+	resolveTelemetry,
 	ThinkingLevel,
 } from "@oh-my-pi/pi-agent-core";
 import {
@@ -79,6 +82,7 @@ import {
 	Snowflake,
 } from "@oh-my-pi/pi-utils";
 import { type AsyncJob, AsyncJobManager } from "../async";
+import { reset as resetCapabilities } from "../capability";
 import type { Rule } from "../capability/rule";
 import { MODEL_ROLE_IDS, type ModelRegistry } from "../config/model-registry";
 import {
@@ -92,6 +96,7 @@ import {
 import { expandPromptTemplate, type PromptTemplate } from "../config/prompt-templates";
 import type { Settings, SkillsSettings } from "../config/settings";
 import { RawSseDebugBuffer } from "../debug/raw-sse-buffer";
+import { loadCapability } from "../discovery";
 import { normalizeDiff, normalizeToLF, ParseError, previewPatch, stripBom } from "../edit";
 import {
 	disposeKernelSessionsByOwner,
@@ -151,8 +156,10 @@ import planModeToolDecisionReminderPrompt from "../prompts/system/plan-mode-tool
 	type: "text",
 };
 import ttsrInterruptTemplate from "../prompts/system/ttsr-interrupt.md" with { type: "text" };
+import ttsrToolReminderTemplate from "../prompts/system/ttsr-tool-reminder.md" with { type: "text" };
 import { type AgentRegistry, MAIN_AGENT_ID } from "../registry/agent-registry";
 import { deobfuscateSessionContext, type SecretObfuscator } from "../secrets/obfuscator";
+import { invalidateHostMetadata } from "../ssh/connection-manager";
 import { resolveThinkingLevelForModel, toReasoningEffort } from "../thinking";
 import {
 	buildDiscoverableToolSearchIndex,
@@ -274,6 +281,9 @@ export interface AgentSessionConfig {
 	convertToLlm?: (messages: AgentMessage[]) => Message[] | Promise<Message[]>;
 	/** System prompt builder that can consider tool availability. Returns ordered provider-facing blocks. */
 	rebuildSystemPrompt?: (toolNames: string[], tools: Map<string, AgentTool>) => Promise<{ systemPrompt: string[] }>;
+	/** Rebuild the SSH tool from current capability discovery results. */
+	reloadSshTool?: () => Promise<AgentTool | null>;
+	requestedToolNames?: ReadonlySet<string>;
 	/**
 	 * Optional accessor for live MCP server instructions. Read by the session's
 	 * `rebuildSystemPrompt`-skip optimization to detect server-side instruction
@@ -736,6 +746,8 @@ export class AgentSession {
 		| ((toolNames: string[], tools: Map<string, AgentTool>) => Promise<{ systemPrompt: string[] }>)
 		| undefined;
 	#getMcpServerInstructions: (() => Map<string, string> | undefined) | undefined;
+	#reloadSshTool: (() => Promise<AgentTool | null>) | undefined;
+	#requestedToolNames: ReadonlySet<string> | undefined;
 	#baseSystemPrompt: string[];
 	/**
 	 * Signature of the (toolNames, tool descriptions) tuple passed to the most
@@ -759,6 +771,10 @@ export class AgentSession {
 	// TTSR manager for time-traveling stream rules
 	#ttsrManager: TtsrManager | undefined = undefined;
 	#pendingTtsrInjections: Rule[] = [];
+	/** Per-tool TTSR rules whose `interruptMode` opted out of aborting the stream.
+	 *  These are folded into the matched tool call's `toolResult` content as an
+	 *  in-band system reminder, instead of spawning a separate follow-up turn. */
+	#perToolTtsrInjections = new Map<string, Rule[]>();
 	#ttsrAbortPending = false;
 	#ttsrRetryToken = 0;
 	#ttsrResumePromise: Promise<void> | undefined = undefined;
@@ -868,24 +884,38 @@ export class AgentSession {
 		this.#modelRegistry = config.modelRegistry;
 		this.#validateRetryFallbackChains();
 		this.#toolRegistry = config.toolRegistry ?? new Map();
+		this.#requestedToolNames = config.requestedToolNames;
 		this.#transformContext = config.transformContext ?? (messages => messages);
 		this.#onPayload = config.onPayload;
 		this.rawSseDebugBuffer = config.rawSseDebugBuffer ?? new RawSseDebugBuffer();
+		// Avoid wrapping in an `async` closure when no user callback is configured: the
+		// outer await on `#onResponse` (provider-response.ts) tolerates a sync void return,
+		// and skipping the wrapper drops a per-event `newPromiseCapability` allocation that
+		// shows up as ~3.5% self time in streaming profiles.
 		const configuredOnResponse = config.onResponse;
-		this.#onResponse = async (response, model) => {
-			this.rawSseDebugBuffer.recordResponse(response, model);
-			await configuredOnResponse?.(response, model);
-		};
+		this.#onResponse = configuredOnResponse
+			? async (response, model) => {
+					this.rawSseDebugBuffer.recordResponse(response, model);
+					await configuredOnResponse(response, model);
+				}
+			: (response, model) => {
+					this.rawSseDebugBuffer.recordResponse(response, model);
+				};
 		const configuredOnSseEvent = config.onSseEvent;
-		this.#onSseEvent = (event, model) => {
-			this.rawSseDebugBuffer.recordEvent(event, model);
-			configuredOnSseEvent?.(event, model);
-		};
+		this.#onSseEvent = configuredOnSseEvent
+			? (event, model) => {
+					this.rawSseDebugBuffer.recordEvent(event, model);
+					configuredOnSseEvent(event, model);
+				}
+			: (event, model) => {
+					this.rawSseDebugBuffer.recordEvent(event, model);
+				};
 		this.agent.setProviderResponseInterceptor(this.#onResponse);
 		this.agent.setRawSseEventInterceptor(this.#onSseEvent);
 		this.#convertToLlm = config.convertToLlm ?? convertToLlm;
 		this.#rebuildSystemPrompt = config.rebuildSystemPrompt;
 		this.#getMcpServerInstructions = config.getMcpServerInstructions;
+		this.#reloadSshTool = config.reloadSshTool;
 		this.#baseSystemPrompt = this.agent.state.systemPrompt;
 		this.#mcpDiscoveryEnabled = config.mcpDiscoveryEnabled ?? false;
 		this.#setDiscoverableMCPTools(this.#collectDiscoverableMCPToolsFromRegistry());
@@ -922,6 +952,8 @@ export class AgentSession {
 			this.#preCacheStreamingEditFile(event);
 			this.#maybeAbortStreamingEdit(event);
 		});
+		// Per-tool TTSR reminders are folded into the matched tool's result via this hook.
+		this.agent.afterToolCall = ctx => this.#ttsrAfterToolCall(ctx);
 		this.agent.providerSessionState = this.#providerSessionState;
 		this.#syncAgentSessionId();
 		this.#syncTodoPhasesFromBranch();
@@ -1315,77 +1347,87 @@ export class AgentSession {
 			if (matchContext && "delta" in assistantEvent) {
 				const matches = this.#ttsrManager.checkDelta(assistantEvent.delta, matchContext);
 				if (matches.length > 0) {
-					// Queue rules for injection; mark as injected only after successful enqueue.
-
-					this.#addPendingTtsrInjections(matches);
-
-					if (this.#shouldInterruptForTtsrMatch(matches, matchContext)) {
-						// Abort the stream immediately — do not gate on extension callbacks
-						this.#ttsrAbortPending = true;
-						this.#ensureTtsrResumePromise();
-						this.agent.abort();
-						// Notify extensions (fire-and-forget, does not block abort)
+					// Decide first: a non-interrupting tool-source match attaches to the
+					// specific tool call's result instead of driving a loop-wide follow-up.
+					const shouldInterrupt = this.#shouldInterruptForTtsrMatch(matches, matchContext);
+					const perToolId = shouldInterrupt ? undefined : this.#extractTtsrToolCallId(matchContext);
+					if (perToolId) {
+						this.#addPerToolTtsrInjections(perToolId, matches);
 						this.#emitSessionEvent({ type: "ttsr_triggered", rules: matches }).catch(() => {});
-						// Schedule retry after a short delay
-						const retryToken = ++this.#ttsrRetryToken;
-						const generation = this.#promptGeneration;
-						const targetMessageTimestamp =
-							event.message.role === "assistant" ? event.message.timestamp : undefined;
-						this.#schedulePostPromptTask(
-							async () => {
-								if (this.#ttsrRetryToken !== retryToken) {
-									this.#resolveTtsrResume();
-									return;
-								}
+					} else {
+						// Queue rules for injection; mark as injected only after successful enqueue.
+						this.#addPendingTtsrInjections(matches);
+
+						if (shouldInterrupt) {
+							// Abort the stream immediately — do not gate on extension callbacks
+							this.#ttsrAbortPending = true;
+							this.#ensureTtsrResumePromise();
+							this.agent.abort();
+							// Notify extensions (fire-and-forget, does not block abort)
+							this.#emitSessionEvent({ type: "ttsr_triggered", rules: matches }).catch(() => {});
+							// Schedule retry after a short delay
+							const retryToken = ++this.#ttsrRetryToken;
+							const generation = this.#promptGeneration;
+							const targetMessageTimestamp =
+								event.message.role === "assistant" ? event.message.timestamp : undefined;
+							this.#schedulePostPromptTask(
+								async () => {
+									if (this.#ttsrRetryToken !== retryToken) {
+										this.#resolveTtsrResume();
+										return;
+									}
 
-								const targetAssistantIndex = this.#findTtsrAssistantIndex(targetMessageTimestamp);
-								if (
-									!this.#ttsrAbortPending ||
-									this.#promptGeneration !== generation ||
-									targetAssistantIndex === -1
-								) {
+									const targetAssistantIndex = this.#findTtsrAssistantIndex(targetMessageTimestamp);
+									if (
+										!this.#ttsrAbortPending ||
+										this.#promptGeneration !== generation ||
+										targetAssistantIndex === -1
+									) {
+										this.#ttsrAbortPending = false;
+										this.#pendingTtsrInjections = [];
+										this.#perToolTtsrInjections.clear();
+										this.#resolveTtsrResume();
+										return;
+									}
 									this.#ttsrAbortPending = false;
-									this.#pendingTtsrInjections = [];
-									this.#resolveTtsrResume();
-									return;
-								}
-								this.#ttsrAbortPending = false;
-								const ttsrSettings = this.#ttsrManager?.getSettings();
-								if (ttsrSettings?.contextMode === "discard") {
-									// Remove the partial/aborted assistant turn from agent state
-									this.agent.replaceMessages(this.agent.state.messages.slice(0, targetAssistantIndex));
-								}
-								// Inject TTSR rules as system reminder before retry
-								const injection = this.#getTtsrInjectionContent();
-								if (injection) {
-									const details = { rules: injection.rules.map(rule => rule.name) };
-									this.agent.appendMessage({
-										role: "custom",
-										customType: "ttsr-injection",
-										content: injection.content,
-										display: false,
-										details,
-										attribution: "agent",
-										timestamp: Date.now(),
-									});
-									this.sessionManager.appendCustomMessageEntry(
-										"ttsr-injection",
-										injection.content,
-										false,
-										details,
-										"agent",
-									);
-									this.#markTtsrInjected(details.rules);
-								}
-								try {
-									await this.agent.continue();
-								} catch {
-									this.#resolveTtsrResume();
-								}
-							},
-							{ delayMs: 50 },
-						);
-						return;
+									this.#perToolTtsrInjections.clear();
+									const ttsrSettings = this.#ttsrManager?.getSettings();
+									if (ttsrSettings?.contextMode === "discard") {
+										// Remove the partial/aborted assistant turn from agent state
+										this.agent.replaceMessages(this.agent.state.messages.slice(0, targetAssistantIndex));
+									}
+									// Inject TTSR rules as system reminder before retry
+									const injection = this.#getTtsrInjectionContent();
+									if (injection) {
+										const details = { rules: injection.rules.map(rule => rule.name) };
+										this.agent.appendMessage({
+											role: "custom",
+											customType: "ttsr-injection",
+											content: injection.content,
+											display: false,
+											details,
+											attribution: "agent",
+											timestamp: Date.now(),
+										});
+										this.sessionManager.appendCustomMessageEntry(
+											"ttsr-injection",
+											injection.content,
+											false,
+											details,
+											"agent",
+										);
+										this.#markTtsrInjected(details.rules);
+									}
+									try {
+										await this.agent.continue();
+									} catch {
+										this.#resolveTtsrResume();
+									}
+								},
+								{ delayMs: 50 },
+							);
+							return;
+						}
 					}
 				}
 			}
@@ -1794,6 +1836,61 @@ export class AgentSession {
 		}
 	}
 
+	/** Tool-call id whose argument deltas triggered a TTSR match, when known. */
+	#extractTtsrToolCallId(matchContext: TtsrMatchContext): string | undefined {
+		if (matchContext.source !== "tool") return undefined;
+		const key = matchContext.streamKey;
+		if (typeof key !== "string" || !key.startsWith("toolcall:")) return undefined;
+		const id = key.slice("toolcall:".length);
+		return id.length > 0 ? id : undefined;
+	}
+
+	#addPerToolTtsrInjections(toolCallId: string, rules: Rule[]): void {
+		const bucket = this.#perToolTtsrInjections.get(toolCallId) ?? [];
+		const seen = new Set(bucket.map(rule => rule.name));
+		// Dedupe against rules already bucketed for other tool calls in this
+		// same assistant message so one rule attaches to exactly one tool call.
+		const claimedElsewhere = new Set<string>();
+		for (const [otherId, otherBucket] of this.#perToolTtsrInjections) {
+			if (otherId === toolCallId) continue;
+			for (const rule of otherBucket) claimedElsewhere.add(rule.name);
+		}
+		const newlyAdded: string[] = [];
+		for (const rule of rules) {
+			if (seen.has(rule.name) || claimedElsewhere.has(rule.name)) continue;
+			bucket.push(rule);
+			seen.add(rule.name);
+			newlyAdded.push(rule.name);
+		}
+		if (bucket.length === 0) return;
+		this.#perToolTtsrInjections.set(toolCallId, bucket);
+		// Claim the rules in the TTSR manager so subsequent deltas in this same
+		// turn (e.g. a sibling tool call's argument stream) don't re-match them.
+		// Persistence still happens in #ttsrAfterToolCall when the tool actually
+		// produces a result we can fold the reminder into.
+		if (newlyAdded.length > 0) {
+			this.#ttsrManager?.markInjectedByNames(newlyAdded);
+		}
+	}
+
+	/** `afterToolCall` hook: fold any per-tool TTSR reminders into the result. */
+	#ttsrAfterToolCall(ctx: AfterToolCallContext): AfterToolCallResult | undefined {
+		const rules = this.#perToolTtsrInjections.get(ctx.toolCall.id);
+		if (!rules || rules.length === 0) return undefined;
+		this.#perToolTtsrInjections.delete(ctx.toolCall.id);
+		const reminder = rules
+			.map(r => prompt.render(ttsrToolReminderTemplate, { name: r.name, path: r.path, content: r.content }))
+			.join("\n\n");
+		// The TTSR manager was already claimed at bucket time; only persistence remains.
+		const ruleNames = rules.map(r => r.name.trim()).filter(n => n.length > 0);
+		if (ruleNames.length > 0) {
+			this.sessionManager.appendTtsrInjection(ruleNames);
+		}
+		return {
+			content: [{ type: "text", text: reminder }, ...ctx.result.content],
+		};
+	}
+
 	#extractTtsrRuleNames(details: unknown): string[] {
 		if (!details || typeof details !== "object" || Array.isArray(details)) {
 			return [];
@@ -1844,6 +1941,11 @@ export class AgentSession {
 	}
 
 	#queueDeferredTtsrInjectionIfNeeded(assistantMsg: AssistantMessage): void {
+		if (assistantMsg.stopReason === "aborted" || assistantMsg.stopReason === "error") {
+			// Tools that hadn't started by abort/error will never produce results to
+			// fold injections into — drop their stale per-tool entries.
+			this.#perToolTtsrInjections.clear();
+		}
 		if (this.#ttsrAbortPending || this.#pendingTtsrInjections.length === 0) {
 			return;
 		}
@@ -3000,6 +3102,45 @@ export class AgentSession {
 		}
 	}
 
+	/**
+	 * Reload the SSH tool from disk-backed capability discovery and make the
+	 * refreshed definition visible to the next model call without restarting.
+	 */
+	async refreshSshTool(options?: { activateIfAvailable?: boolean }): Promise<void> {
+		resetCapabilities();
+		if (!this.#reloadSshTool) return;
+		const previousSshTool = this.#toolRegistry.get("ssh");
+		const previousActiveToolNames = this.getActiveToolNames();
+		const hadSshTool = previousSshTool !== undefined;
+		const wasActive = previousActiveToolNames.includes("ssh");
+		const previousHostNames =
+			previousSshTool && "hostNames" in previousSshTool && Array.isArray(previousSshTool.hostNames)
+				? [...previousSshTool.hostNames]
+				: [];
+		const candidateHostNames = new Set(previousHostNames);
+		const capability = await loadCapability<{ name: string }>("ssh", { cwd: this.sessionManager.getCwd() });
+		for (const host of capability.items) {
+			if (typeof host?.name === "string") {
+				candidateHostNames.add(host.name);
+			}
+		}
+		await invalidateHostMetadata(candidateHostNames);
+		const sshAllowed = this.#requestedToolNames === undefined || this.#requestedToolNames.has("ssh");
+		const refreshedTool = await this.#reloadSshTool();
+		if (refreshedTool) {
+			this.#toolRegistry.set(refreshedTool.name, refreshedTool);
+		} else {
+			this.#toolRegistry.delete("ssh");
+			this.#selectedDiscoveredToolNames.delete("ssh");
+		}
+
+		const nextActive = previousActiveToolNames.filter(name => name !== "ssh" && this.#toolRegistry.has(name));
+		if (refreshedTool && sshAllowed && (wasActive || (options?.activateIfAvailable && !hadSshTool))) {
+			nextActive.push(refreshedTool.name);
+		}
+		await this.#applyActiveToolsByName(nextActive);
+	}
+
 	/**
 	 * Set active tools by name.
 	 * Only tools in the registry can be enabled. Unknown tool names are ignored.
@@ -5200,6 +5341,7 @@ export class AgentSession {
 					convertToLlm,
 					initiatorOverride: "agent",
 					metadata: this.agent.metadataForProvider(model.provider),
+					telemetry: resolveTelemetry(this.agent.telemetry, this.sessionId),
 				},
 				handoffSignal,
 			);
@@ -5912,6 +6054,7 @@ export class AgentSession {
 		options?: SummaryOptions,
 	): Promise<CompactionResult> {
 		const candidates = this.#getCompactionModelCandidates(this.#modelRegistry.getAvailable());
+		const telemetry = resolveTelemetry(this.agent.telemetry, this.sessionId);
 
 		for (const candidate of candidates) {
 			const apiKey = await this.#modelRegistry.getApiKey(candidate, this.sessionId);
@@ -5922,6 +6065,7 @@ export class AgentSession {
 					...options,
 					metadata: this.agent.metadataForProvider(candidate.provider),
 					convertToLlm,
+					telemetry,
 				});
 			} catch (error) {
 				if (!this.#isCompactionAuthFailure(error)) {
@@ -6158,6 +6302,7 @@ export class AgentSession {
 			} else {
 				const candidates = this.#getCompactionModelCandidates(availableModels);
 				const retrySettings = this.settings.getGroup("retry");
+				const telemetry = resolveTelemetry(this.agent.telemetry, this.sessionId);
 				let compactResult: CompactionResult | undefined;
 				let lastError: unknown;
 
@@ -6175,6 +6320,7 @@ export class AgentSession {
 								metadata: this.agent.metadataForProvider(candidate.provider),
 								initiatorOverride: "agent",
 								convertToLlm,
+								telemetry,
 							});
 							break;
 						} catch (error) {
@@ -7779,6 +7925,7 @@ export class AgentSession {
 				reserveTokens: branchSummarySettings.reserveTokens,
 				metadata: this.agent.metadataForProvider(model.provider),
 				convertToLlm,
+				telemetry: resolveTelemetry(this.agent.telemetry, this.sessionId),
 			});
 			this.#branchSummaryAbortController = undefined;
 			if (result.aborted) {
@@ -8014,11 +8161,12 @@ export class AgentSession {
 		};
 	}
 
-	async fetchUsageReports(): Promise<UsageReport[] | null> {
+	async fetchUsageReports(signal?: AbortSignal): Promise<UsageReport[] | null> {
 		const authStorage = this.#modelRegistry.authStorage;
 		if (!authStorage.fetchUsageReports) return null;
 		return authStorage.fetchUsageReports({
 			baseUrlResolver: provider => this.#modelRegistry.getProviderBaseUrl?.(provider),
+			signal,
 		});
 	}
 

package/src/session/agent-storage.ts

@@ -1,7 +1,12 @@
 import { Database, type Statement } from "bun:sqlite";
 import * as fs from "node:fs";
 import * as path from "node:path";
-import { type AuthCredential, AuthCredentialStore, type StoredAuthCredential } from "@oh-my-pi/pi-ai";
+import {
+	type AuthCredential,
+	type AuthCredentialStore,
+	SqliteAuthCredentialStore,
+	type StoredAuthCredential,
+} from "@oh-my-pi/pi-ai";
 import { getAgentDbPath, isRecord, logger } from "@oh-my-pi/pi-utils";
 import type { RawSettings as Settings } from "../config/settings";
 
@@ -57,7 +62,7 @@ export class AgentStorage {
 		this.#hardenPermissions(dbPath);
 
 		// Create AuthCredentialStore with our open database
-		this.#authStore = new AuthCredentialStore(this.#db);
+		this.#authStore = new SqliteAuthCredentialStore(this.#db);
 
 		this.#listSettingsStmt = this.#db.prepare("SELECT key, value FROM settings");
 		this.#upsertModelUsageStmt = this.#db.prepare(

package/src/session/auth-broker-config.ts

@@ -0,0 +1,102 @@
+/**
+ * Resolve auth-broker connection configuration for the local omp client.
+ *
+ * Precedence (highest first):
+ *   1. `OMP_AUTH_BROKER_URL` / `OMP_AUTH_BROKER_TOKEN` env vars.
+ *   2. `auth.broker.url` / `auth.broker.token` in `~/.omp/agent/config.yml`
+ *      (hidden from the settings UI; `!command` resolution supported).
+ *   3. Token file `~/.omp/auth-broker.token` (paired with URL from env or config).
+ *
+ * Returns null when no broker URL is configured — caller falls back to the
+ * local SQLite store.
+ *
+ * Reads config.yml directly (instead of going through `Settings.init`) because
+ * `discoverAuthStorage` runs before the settings singleton is initialized in
+ * `runRootCommand`, and we want hand-edited config entries to be honoured at
+ * boot without forcing a startup reorder.
+ */
+import * as path from "node:path";
+import { getAgentDir, getConfigRootDir, isEnoent, logger } from "@oh-my-pi/pi-utils";
+import { YAML } from "bun";
+import { resolveConfigValue } from "../config/resolve-config-value";
+
+export interface AuthBrokerClientConfig {
+	url: string;
+	token: string;
+}
+
+/** Path to the local bearer token file. Created on the broker host by `omp auth-broker token`. */
+export function getAuthBrokerTokenFilePath(): string {
+	return path.join(getConfigRootDir(), "auth-broker.token");
+}
+
+async function readTokenFile(): Promise<string | null> {
+	try {
+		const raw = await Bun.file(getAuthBrokerTokenFilePath()).text();
+		const trimmed = raw.trim();
+		return trimmed.length > 0 ? trimmed : null;
+	} catch (err) {
+		if (isEnoent(err)) return null;
+		logger.warn("auth-broker token file unreadable", { error: String(err) });
+		return null;
+	}
+}
+
+interface ConfigSnapshot {
+	url?: string;
+	token?: string;
+}
+
+async function readConfigYaml(): Promise<ConfigSnapshot> {
+	const configPath = path.join(getAgentDir(), "config.yml");
+	try {
+		const raw = await Bun.file(configPath).text();
+		const parsed = YAML.parse(raw);
+		if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) return {};
+		const record = parsed as Record<string, unknown>;
+		const url = typeof record["auth.broker.url"] === "string" ? (record["auth.broker.url"] as string) : undefined;
+		const token =
+			typeof record["auth.broker.token"] === "string" ? (record["auth.broker.token"] as string) : undefined;
+		return { url, token };
+	} catch (err) {
+		if (isEnoent(err)) return {};
+		logger.warn("auth-broker config.yml unreadable", { error: String(err) });
+		return {};
+	}
+}
+
+/**
+ * Read broker configuration. Returns null when the URL is missing
+ * (broker disabled — local store is used). Throws when URL is set but no
+ * token is available — the caller cannot fall back silently because the
+ * user explicitly asked to use the broker.
+ */
+export async function resolveAuthBrokerConfig(): Promise<AuthBrokerClientConfig | null> {
+	const envUrl = process.env.OMP_AUTH_BROKER_URL;
+	const envToken = process.env.OMP_AUTH_BROKER_TOKEN;
+
+	let url = envUrl && envUrl.length > 0 ? envUrl : undefined;
+	let configToken: string | undefined;
+	if (!url || !envToken) {
+		const fromConfig = await readConfigYaml();
+		if (!url && fromConfig.url) {
+			const resolved = await resolveConfigValue(fromConfig.url);
+			if (resolved && resolved.length > 0) url = resolved;
+		}
+		if (fromConfig.token) {
+			const resolved = await resolveConfigValue(fromConfig.token);
+			if (resolved && resolved.length > 0) configToken = resolved;
+		}
+	}
+	if (!url) return null;
+
+	const token =
+		(envToken && envToken.length > 0 ? envToken : undefined) ?? configToken ?? (await readTokenFile()) ?? undefined;
+	if (!token) {
+		throw new Error(
+			`OMP_AUTH_BROKER_URL is set (${url}) but no bearer token is available. ` +
+				`Set OMP_AUTH_BROKER_TOKEN, the \`auth.broker.token\` config entry, or place one at ${getAuthBrokerTokenFilePath()}.`,
+		);
+	}
+	return { url, token };
+}

package/src/session/auth-storage.ts

@@ -14,4 +14,10 @@ export type {
 	SerializedAuthStorage,
 	StoredAuthCredential,
 } from "@oh-my-pi/pi-ai";
-export { AuthStorage } from "@oh-my-pi/pi-ai";
+export {
+	AuthBrokerClient,
+	AuthStorage,
+	REMOTE_REFRESH_SENTINEL,
+	RemoteAuthCredentialStore,
+	SqliteAuthCredentialStore,
+} from "@oh-my-pi/pi-ai";

package/src/session/streaming-output.ts

@@ -1,5 +1,5 @@
 import type { AgentToolUpdateCallback } from "@oh-my-pi/pi-agent-core";
-import { sanitizeText } from "@oh-my-pi/pi-natives";
+import { sanitizeText } from "@oh-my-pi/pi-utils";
 import { formatBytes } from "../tools/render-utils";
 import { sanitizeWithOptionalSixelPassthrough } from "../utils/sixel";
 

package/src/slash-commands/builtin-registry.ts

@@ -1506,6 +1506,7 @@ const BUILTIN_SLASH_COMMAND_REGISTRY: ReadonlyArray<SlashCommandSpec> = [
 			const projectPath = await resolveActiveProjectRegistryPath(runtime.ctx.sessionManager.getCwd());
 			clearPluginRootsAndCaches(projectPath ? [projectPath] : undefined);
 			await runtime.ctx.refreshSlashCommandState();
+			await runtime.ctx.session.refreshSshTool({ activateIfAvailable: true });
 			runtime.ctx.showStatus("Plugins reloaded.");
 			runtime.ctx.editor.setText("");
 		},
@@ -1632,6 +1633,7 @@ export async function executeBuiltinSlashCommand(
 				const projectPath = await resolveActiveProjectRegistryPath(ctx.sessionManager.getCwd());
 				clearPluginRootsAndCaches(projectPath ? [projectPath] : undefined);
 				await ctx.refreshSlashCommandState();
+				await ctx.session.refreshSshTool({ activateIfAvailable: true });
 			},
 		};
 		const result = await command.handle(parsed, adapted);

package/src/slash-commands/helpers/ssh.ts

@@ -142,6 +142,7 @@ async function handleRemoveCommand(rest: string, runtime: SlashCommandRuntime):
 	try {
 		const filePath = getSSHConfigPath(parsed.scope, runtime.cwd);
 		await removeSSHHost(filePath, parsed.name);
+		await runtime.session.refreshSshTool();
 		await runtime.output(`Removed SSH host "${parsed.name}" from ${parsed.scope} config.`);
 		return commandConsumed();
 	} catch (err) {
@@ -162,6 +163,7 @@ async function handleAddCommand(rest: string, runtime: SlashCommandRuntime): Pro
 	try {
 		const filePath = getSSHConfigPath(parsed.scope, runtime.cwd);
 		await addSSHHost(filePath, parsed.name, hostConfig);
+		await runtime.session.refreshSshTool({ activateIfAvailable: true });
 		await runtime.output(`Added SSH host "${parsed.name}" (${parsed.scope}).`);
 		return commandConsumed();
 	} catch (err) {