feat(coding-agent): improved vibe mode tool lifecycle management

- Transitioned vibe tools to an ephemeral registration model where they are installed only when entering `/vibe` mode and removed upon exit. - Added `activateVibeTools` and `deactivateVibeTools` methods to `AgentSession` to manage these transient tool registrations. - Removed vibe tools from the default global tool registry, preventing unnecessary background exposure.
2026-07-11 06:27:32 +02:00
parent 9ebc239280
commit 46fd8c5575
9 changed files with 120 additions and 41 deletions
@@ -38,6 +38,7 @@
 - Fixed cmux `tab.screenshot({ selector })` silently returning a full-viewport capture that models consumed as an element crop. The cmux daemon has no element-clip or full-page capture; the tool still scrolls the selector into view but now labels the image as full-viewport (same for `fullPage`) instead of mislabeling it.
 - Fixed cmux `tab.evaluate()` / `elementHandle.evaluate()` errors surfacing as the daemon's opaque `js_error: A JavaScript exception occurred`. Scripts now run inside a page-side try/catch envelope that returns the real message and stack, and a Promise return (which the daemon cannot serialize) yields an actionable "evaluates synchronously" error instead of an unsupported-type failure.
 - Fixed bug where exiting `/vibe` mode failed to restore an empty toolset
+- Fixed `/vibe` tools being pre-registered as hidden session tools; they are now installed only on mode entry and removed on exit
 - Fixed `/vibe wait` reporting to properly capture settled results that have queued follow-up turns
 - Improved `/vibe` TUI responsiveness by rendering animated content (cursors, spinners) directly from mutable state at paint time

@@ -118,7 +118,6 @@ import { setAutoQaConsentHandler } from "../tools/report-tool-issue";
 import { type ResolveToolDetails, runResolveInvocation } from "../tools/resolve";
 import { formatPhaseDisplayName, todoMatchesAnyDescription } from "../tools/todo";
 import { ToolError } from "../tools/tool-errors";
-import { VIBE_TOOL_NAMES } from "../tools/vibe";
 import { vocalizer } from "../tts/vocalizer";
 import { renderTreeList } from "../tui/tree-list";
 import type { EventBus } from "../utils/event-bus";
@@ -2122,9 +2121,7 @@ export class InteractiveMode implements InteractiveModeContext {
 		}

 		if (this.vibeModeEnabled) {
-			if (this.#vibeModePreviousTools !== undefined) {
-				await this.session.setActiveToolsByName(this.#vibeModePreviousTools);
-			}
+			await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []);
 			this.session.setVibeModeState(undefined);
 			this.vibeModeEnabled = false;
 			this.#vibeModePreviousTools = undefined;
@@ -2922,10 +2919,10 @@ export class InteractiveMode implements InteractiveModeContext {
 	}

 	/**
-	 * `/vibe` toggle. Entering strips the active toolset down to `read` plus the
-	 * vibe session tools and injects the director context; exiting restores the
-	 * previous toolset and kills every worker session (boring and safe — workers
-	 * do not outlive the mode that directs them).
+	 * `/vibe` toggle. Entering installs the ephemeral vibe tools, strips the
+	 * active toolset down to `read` plus those tools, and injects the director
+	 * context. Exiting unregisters them, restores the previous toolset, and kills
+	 * every worker session so workers cannot outlive the mode that directs them.
 	 */
 	async handleVibeModeCommand(initialPrompt?: string): Promise<void> {
 		if (this.vibeModeEnabled) {
@@ -2959,12 +2956,13 @@ export class InteractiveMode implements InteractiveModeContext {
 			return;
 		}

-		this.#vibeModePreviousTools = this.session.getActiveToolNames();
+		const previousTools = this.session.getActiveToolNames();
+		await this.session.activateVibeTools(["read"]);
+		this.#vibeModePreviousTools = previousTools;
 		this.vibeModeEnabled = true;
 		// Suppress cache-miss marker on the next turn: vibe mode changes the
 		// injected context, which predictably invalidates the cache.
 		this.lastAssistantUsage = undefined;
-		await this.session.setActiveToolsByName(["read", ...VIBE_TOOL_NAMES]);
 		this.session.setVibeModeState({ enabled: true });
 		if (this.session.isStreaming) {
 			await this.session.sendVibeModeContext({ deliverAs: "steer" });
@@ -2978,10 +2976,7 @@ export class InteractiveMode implements InteractiveModeContext {
 		if (!this.vibeModeEnabled) {
 			return;
 		}
-		const previousTools = this.#vibeModePreviousTools;
-		if (previousTools !== undefined) {
-			await this.session.setActiveToolsByName(previousTools);
-		}
+		await this.session.deactivateVibeTools(this.#vibeModePreviousTools ?? []);
 		this.session.setVibeModeState(undefined);
 		this.vibeModeEnabled = false;
 		this.#vibeModePreviousTools = undefined;
@@ -161,6 +161,7 @@ import {
 	BUILTIN_TOOLS,
 	computeEssentialBuiltinNames,
 	createTools,
+	createVibeTools,
 	type DeferredDiagnosticsEntry,
 	discoverStartupLspServers,
 	EditTool,
@@ -184,7 +185,6 @@ import {
 	setPreferredSearchProvider,
 	type Tool,
 	type ToolSession,
-	VIBE_TOOL_NAMES,
 	WebSearchTool,
 	WriteTool,
 	warmupLspServers,
@@ -2240,19 +2240,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 				builtInRegistryToolNames.add(goalTool.name);
 			}
 		}
-		// Vibe tools are hidden from every default/discovery set; they exist in the
-		// registry so entering vibe mode can activate them via setActiveToolsByName.
-		// Top-level interactive sessions only — subagents never direct vibe workers.
-		if ((options.taskDepth ?? 0) === 0 && !options.parentTaskPrefix) {
-			for (const name of VIBE_TOOL_NAMES) {
-				if (toolRegistry.has(name)) continue;
-				const vibeTool = await logger.time(`createTools:${name}:session`, HIDDEN_TOOLS[name], toolSession);
-				if (vibeTool) {
-					toolRegistry.set(vibeTool.name, wrapToolWithMetaNotice(vibeTool));
-					builtInRegistryToolNames.add(vibeTool.name);
-				}
-			}
-		}
 		for (const tool of wrappedExtensionTools) {
 			toolRegistry.set(tool.name, tool);
 			builtInRegistryToolNames.delete(tool.name);
@@ -2897,6 +2884,10 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
 			skillsSettings: settings.getGroup("skills"),
 			modelRegistry,
 			toolRegistry,
+			createVibeTools:
+				(options.taskDepth ?? 0) === 0 && !options.parentTaskPrefix
+					? () => createVibeTools(toolSession)
+					: undefined,
 			builtInToolNames: builtInRegistryToolNames,
 			transformContext,
 			transformProviderContext,
@@ -714,6 +714,8 @@ export interface AgentSessionConfig {
 	modelRegistry: ModelRegistry;
 	/** Tool registry for LSP and settings */
 	toolRegistry?: Map<string, AgentTool>;
+	/** Creates the tools registered only while `/vibe` mode is active. */
+	createVibeTools?: () => AgentTool[];
 	/** Tool names whose current registry entry is still the built-in implementation. */
 	builtInToolNames?: Iterable<string>;
 	/** Update tool-session predicates that render guidance from the live active tool set. */
@@ -1747,6 +1749,8 @@ export class AgentSession {

 	// Tool registry and prompt builder for extensions
 	#toolRegistry: Map<string, AgentTool>;
+	#createVibeTools: (() => AgentTool[]) | undefined;
+	#installedVibeToolNames = new Set<string>();
 	#transformContext: (messages: AgentMessage[], signal?: AbortSignal) => AgentMessage[] | Promise<AgentMessage[]>;
 	#onPayload: SimpleStreamOptions["onPayload"] | undefined;
 	#onResponse: SimpleStreamOptions["onResponse"] | undefined;
@@ -2110,6 +2114,7 @@ export class AgentSession {
 		this.#pruneToolDescriptions = config.pruneToolDescriptions === true;
 		this.#validateRetryFallbackChains();
 		this.#toolRegistry = config.toolRegistry ?? new Map();
+		this.#createVibeTools = config.createVibeTools;
 		this.#builtInToolNames = new Set(config.builtInToolNames ?? []);
 		this.#requestedToolNames = config.requestedToolNames;
 		this.#transformContext = config.transformContext ?? (messages => messages);
@@ -6116,6 +6121,49 @@ export class AgentSession {
 		return Array.from(this.#toolRegistry.keys());
 	}

+	#wrapRuntimeTool(tool: AgentTool): AgentTool {
+		const wrapped = wrapToolWithMetaNotice(tool);
+		return this.#extensionRunner ? new ExtensionToolWrapper(wrapped, this.#extensionRunner) : wrapped;
+	}
+
+	/**
+	 * Registers the ephemeral vibe tools and activates them alongside `baseToolNames`.
+	 *
+	 * @throws When this session cannot create vibe tools or the factory returns duplicate names.
+	 */
+	async activateVibeTools(baseToolNames: string[]): Promise<void> {
+		const createVibeTools = this.#createVibeTools;
+		if (!createVibeTools) {
+			throw new Error("Vibe tools are unavailable in this session.");
+		}
+
+		const tools = createVibeTools();
+		const vibeToolNames = tools.map(tool => tool.name);
+		if (new Set(vibeToolNames).size !== vibeToolNames.length) {
+			throw new Error("Vibe tool names must be unique.");
+		}
+
+		for (const tool of tools) {
+			if (this.#toolRegistry.has(tool.name)) continue;
+			this.#toolRegistry.set(tool.name, this.#wrapRuntimeTool(tool));
+			this.#builtInToolNames.add(tool.name);
+			this.#installedVibeToolNames.add(tool.name);
+		}
+
+		await this.#applyActiveToolsByName([...new Set([...baseToolNames, ...vibeToolNames])]);
+	}
+
+	/** Removes tools installed by {@link activateVibeTools} and activates `nextToolNames`. */
+	async deactivateVibeTools(nextToolNames: string[]): Promise<void> {
+		for (const name of this.#installedVibeToolNames) {
+			this.#toolRegistry.delete(name);
+			this.#builtInToolNames.delete(name);
+			this.#selectedDiscoveredToolNames.delete(name);
+		}
+		this.#installedVibeToolNames.clear();
+		await this.#applyActiveToolsByName(nextToolNames);
+	}
+
 	#getEditModeSession() {
 		return {
 			settings: this.settings,
@@ -64,7 +64,6 @@ import { reportFindingTool } from "./review";
 import { SearchToolBm25Tool } from "./search-tool-bm25";
 import { loadSshTool } from "./ssh";
 import { type TodoPhase, TodoTool } from "./todo";
-import { VibeKillTool, VibeListTool, VibeSendTool, VibeSpawnTool, VibeWaitTool } from "./vibe";
 import { WriteTool } from "./write";
 import { YieldTool } from "./yield";

@@ -485,11 +484,6 @@ export const HIDDEN_TOOLS: Record<string, ToolFactory> = {
 	report_tool_issue: s => createReportToolIssueTool(s),
 	resolve: s => new ResolveTool(s),
 	goal: s => new GoalTool(s),
-	vibe_spawn: s => new VibeSpawnTool(s),
-	vibe_send: s => new VibeSendTool(s),
-	vibe_wait: s => new VibeWaitTool(s),
-	vibe_kill: s => new VibeKillTool(s),
-	vibe_list: s => new VibeListTool(s),
 };

 export type ToolName = BuiltinToolName;
@@ -35,7 +35,7 @@ import {
 	type VibeSessionState,
 	type VibeWaitOutcome,
 } from "../vibe/runtime";
-import type { ToolSession } from "./index";
+import type { Tool, ToolSession } from "./index";
 import {
 	Ellipsis,
 	formatBadge,
@@ -277,6 +277,17 @@ export class VibeListTool implements AgentTool<typeof vibeListSchema, VibeToolDe
 	}
 }

+/** Creates the ephemeral tools installed while `/vibe` mode is active. */
+export function createVibeTools(session: ToolSession): Tool[] {
+	return [
+		new VibeSpawnTool(session),
+		new VibeSendTool(session),
+		new VibeWaitTool(session),
+		new VibeKillTool(session),
+		new VibeListTool(session),
+	];
+}
+
 // =============================================================================
 // TUI Renderer — mini composer (spawn/send) + TV wall (wait/list)
 // =============================================================================
@@ -1,11 +1,10 @@
 /**
 * Contracts: /vibe mode toggle on InteractiveMode.
 *
- * 1. Entering vibe mode swaps the active toolset to exactly `read` + the vibe
- *    tools and marks the session's vibe-mode state.
- * 2. Exiting restores the pre-vibe toolset EXACTLY — including the legitimate
- *    empty set (regression: a truthiness/length guard used to skip restoring
- *    `[]`, leaving read/vibe tools active after exit).
+ * 1. Vibe tools do not exist in the session registry before the mode is entered.
+ * 2. Entering registers and activates exactly `read` plus the vibe tools.
+ * 3. Exiting unregisters the vibe tools and restores the pre-vibe active toolset
+ *    exactly, including the legitimate empty set.
 */
 import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "bun:test";
 import * as path from "node:path";
@@ -54,7 +53,7 @@ describe("InteractiveMode vibe mode toggle", () => {
 		const model = modelRegistry.find("anthropic", "claude-sonnet-4-5");
 		if (!model) throw new Error("Expected claude-sonnet-4-5 to exist in registry");

-		const registryTools = ["read", ...VIBE_TOOL_NAMES].map(stubTool);
+		const registryTools = [stubTool("read")];

 		session = new AgentSession({
 			agent: new Agent({
@@ -69,6 +68,7 @@ describe("InteractiveMode vibe mode toggle", () => {
 			settings: Settings.isolated({}),
 			modelRegistry,
 			toolRegistry: new Map(registryTools.map(tool => [tool.name, tool])),
+			createVibeTools: () => VIBE_TOOL_NAMES.map(stubTool),
 		});
 		mode = new InteractiveMode(session, "test", undefined, undefined, undefined, undefined, new EventBus());
 	});
@@ -83,6 +83,7 @@ describe("InteractiveMode vibe mode toggle", () => {
 	});

 	it("restores the exact pre-vibe toolset on exit, including an empty one", async () => {
+		expect(session.getAllToolNames()).toEqual(["read"]);
 		expect(session.getActiveToolNames()).toEqual([]);

 		await mode.handleVibeModeCommand();
@@ -93,11 +94,13 @@ describe("InteractiveMode vibe mode toggle", () => {
 			expect(inMode).toContain(name);
 		}
 		expect(inMode.toSorted()).toEqual(["read", ...VIBE_TOOL_NAMES].toSorted());
+		expect(session.getAllToolNames().toSorted()).toEqual(["read", ...VIBE_TOOL_NAMES].toSorted());

 		// Toggle off: the empty previous toolset must come back — vibe tools
 		// must not leak past the mode.
 		await mode.handleVibeModeCommand();
 		expect(mode.vibeModeEnabled).toBe(false);
 		expect(session.getActiveToolNames()).toEqual([]);
+		expect(session.getAllToolNames()).toEqual(["read"]);
 	});
 });
@@ -12,6 +12,7 @@ import {
 	type ExtensionFactory,
 } from "@oh-my-pi/pi-coding-agent/sdk";
 import { SessionManager } from "@oh-my-pi/pi-coding-agent/session/session-manager";
+import { VIBE_TOOL_NAMES } from "@oh-my-pi/pi-coding-agent/tools/vibe";
 import { removeSyncWithRetries, Snowflake } from "@oh-my-pi/pi-utils";
 import { type } from "arktype";

@@ -220,6 +221,32 @@ describe("createAgentSession defaultInactive tool activation", () => {
 		}
 	});

+	it("registers vibe tools only during explicit vibe activation", async () => {
+		const tempDir = makeTempDir();
+		const { session } = await createAgentSession(baseOptions(tempDir));
+		const previousActiveToolNames = session.getActiveToolNames();
+
+		try {
+			for (const name of VIBE_TOOL_NAMES) {
+				expect(session.getToolByName(name)).toBeUndefined();
+			}
+
+			await session.activateVibeTools(["read"]);
+			for (const name of VIBE_TOOL_NAMES) {
+				expect(session.getToolByName(name)).toBeDefined();
+				expect(session.getActiveToolNames()).toContain(name);
+			}
+
+			await session.deactivateVibeTools(previousActiveToolNames);
+			for (const name of VIBE_TOOL_NAMES) {
+				expect(session.getToolByName(name)).toBeUndefined();
+			}
+			expect(session.getActiveToolNames()).toEqual(previousActiveToolNames);
+		} finally {
+			await session.dispose();
+		}
+	});
+
 	it("does not register the xAI TTS tool unless enabled", async () => {
 		const tempDir = makeTempDir();

@@ -11,6 +11,7 @@ import {
 } from "@oh-my-pi/pi-ai/utils/schema";
 import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings";
 import { createTools, HIDDEN_TOOLS, type ToolSession } from "@oh-my-pi/pi-coding-agent/tools";
+import { createVibeTools } from "@oh-my-pi/pi-coding-agent/tools/vibe";

 interface ToolSchemaEntry {
 	name: string;
@@ -58,6 +59,14 @@ async function collectToolSchemas(): Promise<ToolSchemaEntry[]> {
 		byToolName.set(name, schema);
 	}

+	for (const tool of createVibeTools(session)) {
+		const schema = toolWireSchema(tool);
+		if (!asSchemaObject(schema)) {
+			continue;
+		}
+		byToolName.set(tool.name, schema);
+	}
+
 	return [...byToolName.entries()]
 		.sort(([left], [right]) => left.localeCompare(right))
 		.map(([name, schema]) => ({ name, schema }));