Refresh SSH hosts without restart

This commit is contained in:
Gerben Meijer
2026-05-15 22:35:54 +02:00
parent 0bfea39f3c
commit 694f5e9a54
13 changed files with 352 additions and 13 deletions
+4
View File
@@ -2,6 +2,10 @@
## [Unreleased]
### Fixed
- Fixed SSH host additions/removals made inside a running session not refreshing the live `ssh` tool. `/ssh add` and `/ssh remove` now update the model-visible host list immediately, while `/reload-plugins` and `/move` refresh SSH discovery for external or project-scope config changes without restart.
## [15.1.0] - 2026-05-15
### Breaking Changes
@@ -1416,6 +1416,7 @@ export class AcpAgent implements Agent {
resetCapabilities();
const fileCommands = await loadSlashCommands({ cwd });
record.session.setSlashCommands(fileCommands);
await record.session.refreshSshTool({ activateIfAvailable: true });
await this.#emitAvailableCommandsUpdate(record);
}
@@ -968,6 +968,7 @@ export class CommandController {
clearClaudePluginRootsCache(); // re-warms preloadedPluginRoots with new project dir (async)
resetCapabilities();
await this.ctx.refreshSlashCommandState(resolvedPath);
await this.ctx.session.refreshSshTool({ activateIfAvailable: true });
this.ctx.statusLine.invalidate();
this.ctx.updateEditorTopBorder();
@@ -204,6 +204,7 @@ export class SSHCommandController {
if (compat) hostConfig.compat = true;
await addSSHHost(filePath, name, hostConfig);
await this.ctx.session.refreshSshTool({ activateIfAvailable: true });
const scopeLabel = scope === "user" ? "user" : "project";
const lines = [
@@ -364,6 +365,7 @@ export class SSHCommandController {
}
await removeSSHHost(filePath, name);
await this.ctx.session.refreshSshTool();
this.#showMessage(
["", theme.fg("success", `✓ Removed SSH host "${name}" from ${scope} config`), ""].join("\n"),
+14
View File
@@ -1428,6 +1428,17 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
}
}
const reloadSshTool = async (): Promise<AgentTool | null> => {
if (!requestedToolNameSet.has("ssh")) return null;
const sshTool = (await loadSshTool({
...toolSession,
cwd: sessionManager.getCwd(),
})) as unknown as AgentTool | null;
if (!sshTool) return null;
const wrapped = wrapToolWithMetaNotice(sshTool);
return (extensionRunner ? new ExtensionToolWrapper(wrapped, extensionRunner) : wrapped) as AgentTool;
};
let cursorEventEmitter: ((event: AgentEvent) => void) | undefined;
const cursorExecHandlers = new CursorExecHandlers({
cwd,
@@ -1532,6 +1543,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
(options.toolNames ? [...new Set(options.toolNames.map(name => name.toLowerCase()))] : undefined) ??
toolNamesFromRegistry;
const normalizedRequested = requestedToolNames.filter(name => toolRegistry.has(name));
const requestedToolNameSet = new Set(normalizedRequested);
// Effective discovery mode: tools.discoveryMode takes precedence; mcp.discoveryMode is back-compat alias.
const toolsDiscoveryModeSetting = settings.get("tools.discoveryMode");
const effectiveDiscoveryMode: "off" | "mcp-only" | "all" =
@@ -1810,6 +1822,8 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
onResponse,
convertToLlm: convertToLlmFinal,
rebuildSystemPrompt,
reloadSshTool,
requestedToolNames: requestedToolNameSet,
getMcpServerInstructions: mcpManager
? () => {
const raw = mcpManager.getServerInstructions();
@@ -79,6 +79,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 +93,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,
@@ -153,6 +155,7 @@ import planModeToolDecisionReminderPrompt from "../prompts/system/plan-mode-tool
import ttsrInterruptTemplate from "../prompts/system/ttsr-interrupt.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 +277,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 +742,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
@@ -868,6 +876,7 @@ 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();
@@ -886,6 +895,7 @@ export class AgentSession {
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());
@@ -3000,6 +3010,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.
@@ -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);
@@ -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) {
@@ -55,6 +55,16 @@ function getHostInfoPath(name: string): string {
return path.join(HOST_INFO_DIR, `${sanitizeHostName(name)}.json`);
}
async function deleteHostInfoFromDisk(hostName: string): Promise<void> {
const path = getHostInfoPath(hostName);
try {
await fs.promises.unlink(path);
} catch (err) {
if (isEnoent(err)) return;
logger.warn("Failed to delete SSH host info", { host: hostName, error: String(err) });
}
}
async function validateKeyPermissions(keyPath?: string): Promise<void> {
if (!keyPath) return;
let stats: fs.Stats;
@@ -429,6 +439,23 @@ export async function ensureConnection(host: SSHConnectionTarget): Promise<void>
}
}
export async function invalidateHostMetadata(hostNames: Iterable<string>): Promise<void> {
const names = [...hostNames];
for (const hostName of names) {
hostInfoCache.delete(hostName);
await deleteHostInfoFromDisk(hostName);
}
for (const hostName of names) {
const activeHost = activeHosts.get(hostName);
if (activeHost) {
await closeConnectionInternal(activeHost);
activeHosts.delete(hostName);
continue;
}
await closeConnectionInternal({ name: hostName, host: hostName });
}
}
async function closeConnectionInternal(host: SSHConnectionTarget): Promise<void> {
if (!supportsSshControlMaster()) return;
const target = buildSshTarget(host.username, host.host);
@@ -436,13 +463,7 @@ async function closeConnectionInternal(host: SSHConnectionTarget): Promise<void>
}
export async function closeConnection(hostName: string): Promise<void> {
const host = activeHosts.get(hostName);
if (!host) {
await closeConnectionInternal({ name: hostName, host: hostName });
return;
}
await closeConnectionInternal(host);
activeHosts.delete(hostName);
await invalidateHostMetadata([hostName]);
}
export async function closeAllConnections(): Promise<void> {
@@ -153,6 +153,8 @@ class FakeAgentSession {
// no-op for tests
}
async refreshSshTool(_options?: { activateIfAvailable?: boolean }): Promise<void> {}
async setModel(model: Model): Promise<void> {
this.model = model;
}
@@ -29,6 +29,7 @@ interface FakeAcpBuiltinSession {
getTodoPhases(): Array<{ name: string; tasks: Array<{ content: string; status: string }> }>;
setTodoPhases(phases: Array<{ name: string; tasks: Array<{ content: string; status: string }> }>): void;
refreshBaseSystemPrompt(): Promise<void>;
refreshSshTool(options?: { activateIfAvailable?: boolean }): Promise<void>;
getToolByName(name: string): unknown;
compact(args?: string): Promise<void>;
getContextUsage(): { tokens?: number; contextWindow: number } | undefined;
@@ -88,6 +89,7 @@ function createRuntime() {
getContextUsage: () => undefined,
getAvailableModels: () => [] as Array<{ provider: string; id: string; contextWindow?: number }>,
async setModel(_model: unknown) {},
async refreshSshTool(_options?: { activateIfAvailable?: boolean }) {},
};
const typedSession = session as unknown as AgentSession & FakeAcpBuiltinSession;
const fakeSessionManager = {
@@ -0,0 +1,239 @@
import { afterEach, describe, expect, it, spyOn } from "bun:test";
import { Agent, type AgentTool } from "@oh-my-pi/pi-agent-core";
import type { Model } from "@oh-my-pi/pi-ai";
import { getSSHConfigPath, TempDir } from "@oh-my-pi/pi-utils";
import { reset as resetCapabilities } from "../src/capability";
import { type SSHHost, sshCapability } from "../src/capability/ssh";
import { Settings } from "../src/config/settings";
import { loadCapability } from "../src/discovery";
import { AgentSession } from "../src/session/agent-session";
import { SessionManager } from "../src/session/session-manager";
import { addSSHHost, removeSSHHost, updateSSHHost } from "../src/ssh/config-writer";
import * as connectionManager from "../src/ssh/connection-manager";
import { loadSshTool, type ToolSession } from "../src/tools";
function createModel(): Model<"openai-responses"> {
return {
id: "mock",
name: "mock",
api: "openai-responses",
provider: "openai",
baseUrl: "https://example.invalid",
reasoning: false,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 8192,
maxTokens: 2048,
};
}
describe("AgentSession SSH tool refresh", () => {
const tempDirs: TempDir[] = [];
const sessions: AgentSession[] = [];
afterEach(async () => {
for (const session of sessions.splice(0)) {
await session.dispose();
}
for (const tempDir of tempDirs.splice(0)) {
tempDir.removeSync();
}
resetCapabilities();
});
function createSession(
cwd: string,
initialTools: AgentTool[] = [],
registryTools = initialTools,
options?: { reloadSshTool?: () => Promise<AgentTool | null>; requestedToolNames?: ReadonlySet<string> },
): AgentSession {
const settings = Settings.isolated({ "compaction.enabled": false });
const sessionManager = SessionManager.inMemory(cwd);
const toolSession: ToolSession = {
cwd,
hasUI: false,
settings,
getSessionFile: () => sessionManager.getSessionFile() ?? null,
getSessionSpawns: () => "*",
};
const toolRegistry = new Map(registryTools.map(tool => [tool.name, tool]));
const agent = new Agent({
initialState: {
model: createModel(),
systemPrompt: ["initial"],
tools: initialTools,
messages: [],
},
});
const session = new AgentSession({
agent,
sessionManager,
settings,
modelRegistry: {} as never,
toolRegistry,
reloadSshTool:
options?.reloadSshTool ?? (async () => (await loadSshTool(toolSession)) as unknown as AgentTool | null),
requestedToolNames: options?.requestedToolNames,
rebuildSystemPrompt: async (toolNames, tools) => ({
systemPrompt: toolNames.map(name => `${name}:${tools.get(name)?.description ?? ""}`),
}),
});
sessions.push(session);
return session;
}
it("adds the ssh tool after a first host is written over a cached missing config", async () => {
const tempDir = TempDir.createSync("@pi-ssh-refresh-");
tempDirs.push(tempDir);
const cwd = tempDir.path();
const preWrite = await loadCapability<SSHHost>(sshCapability.id, { cwd });
expect(preWrite.items).toHaveLength(0);
const session = createSession(cwd);
await addSSHHost(getSSHConfigPath("project", cwd), "staging", { host: "192.0.2.10" });
await session.refreshSshTool({ activateIfAvailable: true });
expect(session.getAllToolNames()).toContain("ssh");
expect(session.getActiveToolNames()).toContain("ssh");
expect(session.getToolByName("ssh")?.description).toContain("staging (192.0.2.10)");
expect(session.agent.state.systemPrompt.join("\n")).toContain("staging (192.0.2.10)");
});
it("removes ssh from registry and active tools when the last host is removed", async () => {
const tempDir = TempDir.createSync("@pi-ssh-refresh-");
tempDirs.push(tempDir);
const cwd = tempDir.path();
const configPath = getSSHConfigPath("project", cwd);
await addSSHHost(configPath, "prod", { host: "203.0.113.9" });
const sshTool = await loadSshTool({
cwd,
hasUI: false,
settings: Settings.isolated({ "compaction.enabled": false }),
getSessionSpawns: () => "*",
getSessionFile: () => null,
});
expect(sshTool).not.toBeNull();
const session = createSession(cwd, [sshTool as unknown as AgentTool]);
await removeSSHHost(configPath, "prod");
await session.refreshSshTool();
expect(session.getAllToolNames()).not.toContain("ssh");
expect(session.getActiveToolNames()).not.toContain("ssh");
});
it("does not activate an existing inactive ssh tool during reload refresh", async () => {
const tempDir = TempDir.createSync("@pi-ssh-refresh-");
tempDirs.push(tempDir);
const cwd = tempDir.path();
const configPath = getSSHConfigPath("project", cwd);
await addSSHHost(configPath, "dev", { host: "192.0.2.20" });
const sshTool = await loadSshTool({
cwd,
hasUI: false,
settings: Settings.isolated({ "compaction.enabled": false }),
getSessionSpawns: () => "*",
getSessionFile: () => null,
});
expect(sshTool).not.toBeNull();
await addSSHHost(configPath, "dev2", { host: "192.0.2.21" });
const session = createSession(cwd, [], [sshTool as unknown as AgentTool]);
await session.refreshSshTool({ activateIfAvailable: true });
expect(session.getAllToolNames()).toContain("ssh");
expect(session.getActiveToolNames()).not.toContain("ssh");
expect(session.getToolByName("ssh")?.description).toContain("dev2 (192.0.2.21)");
});
it("reloads ssh from the session's current cwd after move", async () => {
const oldProject = TempDir.createSync("@pi-ssh-refresh-old-");
const newProject = TempDir.createSync("@pi-ssh-refresh-new-");
tempDirs.push(oldProject, newProject);
await SessionManager.inMemory(oldProject.path()).moveTo?.(newProject.path());
await addSSHHost(getSSHConfigPath("project", newProject.path()), "moved", { host: "198.51.100.8" });
const movedTool = await loadSshTool({
cwd: newProject.path(),
hasUI: false,
settings: Settings.isolated({ "compaction.enabled": false }),
getSessionSpawns: () => "*",
getSessionFile: () => null,
});
expect(movedTool).not.toBeNull();
const refreshedSession = createSession(oldProject.path(), [], [], {
reloadSshTool: async () => movedTool as unknown as AgentTool,
});
await refreshedSession.refreshSshTool({ activateIfAvailable: true });
expect(refreshedSession.getAllToolNames()).toContain("ssh");
expect(refreshedSession.getToolByName("ssh")?.description).toContain("moved (198.51.100.8)");
});
it("invalidates cached host metadata before rebuilding descriptions when a host config changes", async () => {
const tempDir = TempDir.createSync("@pi-ssh-refresh-");
tempDirs.push(tempDir);
const cwd = tempDir.path();
const configPath = getSSHConfigPath("project", cwd);
await addSSHHost(configPath, "prod", { host: "203.0.113.9" });
const initialTool = await loadSshTool({
cwd,
hasUI: false,
settings: Settings.isolated({ "compaction.enabled": false }),
getSessionSpawns: () => "*",
getSessionFile: () => null,
});
expect(initialTool).not.toBeNull();
const session = createSession(cwd, [initialTool as unknown as AgentTool]);
const invalidateSpy = spyOn(connectionManager, "invalidateHostMetadata").mockResolvedValue(undefined);
await updateSSHHost(configPath, "prod", { host: "203.0.113.10" });
await session.refreshSshTool({ activateIfAvailable: true });
expect(invalidateSpy).toHaveBeenNthCalledWith(1, new Set(["prod"]));
expect(session.getToolByName("ssh")?.description).toContain("prod (203.0.113.10)");
});
it("invalidates newly added host names before rebuilding the ssh tool", async () => {
const tempDir = TempDir.createSync("@pi-ssh-refresh-");
tempDirs.push(tempDir);
const cwd = tempDir.path();
const configPath = getSSHConfigPath("project", cwd);
await addSSHHost(configPath, "fresh", { host: "203.0.113.11" });
const session = createSession(cwd);
await session.refreshSshTool({ activateIfAvailable: true });
expect(session.getToolByName("ssh")?.description).toContain("fresh (203.0.113.11)");
expect(session.getToolByName("ssh")?.description).toContain("fresh (203.0.113.11)");
});
it("does not activate ssh when it was excluded from the requested tool allowlist", async () => {
const tempDir = TempDir.createSync("@pi-ssh-refresh-");
tempDirs.push(tempDir);
const cwd = tempDir.path();
const configPath = getSSHConfigPath("project", cwd);
const blockedTool: AgentTool = {
name: "ssh",
label: "SSH",
description: "blocked",
parameters: { type: "object", properties: {} },
strict: true,
execute: async () => ({ content: [{ type: "text", text: "" }] }),
};
await addSSHHost(configPath, "hidden", { host: "203.0.113.12" });
const session = createSession(cwd, [], [blockedTool], {
reloadSshTool: async () => blockedTool,
requestedToolNames: new Set(["read"]),
});
await session.refreshSshTool({ activateIfAvailable: true });
expect(session.getAllToolNames()).toContain("ssh");
expect(session.getActiveToolNames()).not.toContain("ssh");
});
});
@@ -1,9 +1,9 @@
import { describe, expect, it } from "bun:test";
import { buildRemoteCommand, supportsSshControlMaster } from "../../src/ssh/connection-manager";
import * as connectionManager from "../../src/ssh/connection-manager";
describe("buildRemoteCommand", () => {
it("includes -n and OpenSSH ControlMaster options on Unix-like platforms", async () => {
const args = await buildRemoteCommand(
const args = await connectionManager.buildRemoteCommand(
{
name: "host",
host: "192.168.3.146",
@@ -19,7 +19,7 @@ describe("buildRemoteCommand", () => {
});
it("omits OpenSSH ControlMaster options on Windows", async () => {
const args = await buildRemoteCommand(
const args = await connectionManager.buildRemoteCommand(
{
name: "host",
host: "192.168.3.146",
@@ -40,11 +40,11 @@ describe("buildRemoteCommand", () => {
describe("supportsSshControlMaster", () => {
it("disables OpenSSH connection multiplexing on native Windows", () => {
expect(supportsSshControlMaster("win32")).toBe(false);
expect(connectionManager.supportsSshControlMaster("win32")).toBe(false);
});
it("keeps OpenSSH connection multiplexing on Unix-like platforms", () => {
expect(supportsSshControlMaster("linux")).toBe(true);
expect(supportsSshControlMaster("darwin")).toBe(true);
expect(connectionManager.supportsSshControlMaster("linux")).toBe(true);
expect(connectionManager.supportsSshControlMaster("darwin")).toBe(true);
});
});