diff --git a/kun/src/adapters/tool/remote-port-forward-tool.test.ts b/kun/src/adapters/tool/remote-port-forward-tool.test.ts new file mode 100644 index 000000000..119dd4a82 --- /dev/null +++ b/kun/src/adapters/tool/remote-port-forward-tool.test.ts @@ -0,0 +1,69 @@ +import { describe, expect, it, vi } from 'vitest' +import { createRemotePortForwardRuntime } from './remote-port-forward-tool.js' +import type { ToolHostContext } from '../../ports/tool-host.js' +import type { RemoteExecutionHandle } from '../../ports/remote-execution.js' +import type { SshChildProcess } from '../../remote/ssh-executor.js' + +function fakeChild(): SshChildProcess { + return { stdout: null, stderr: null, on: () => undefined, kill: vi.fn() } as unknown as SshChildProcess +} + +function sshContext(): ToolHostContext { + const handle = { + describe: () => ({ target: { kind: 'ssh', alias: 'prod' }, status: 'connected', production: false }), + production: false + } as unknown as RemoteExecutionHandle + return { + threadId: 't', + turnId: 'tn', + workspace: '/ws', + approvalPolicy: 'auto', + abortSignal: new AbortController().signal, + awaitApproval: vi.fn(async () => 'allow' as const), + executionTarget: handle + } +} + +function localContext(): ToolHostContext { + return { + threadId: 't', + turnId: 'tn', + workspace: '/ws', + approvalPolicy: 'auto', + abortSignal: new AbortController().signal, + awaitApproval: vi.fn(async () => 'allow' as const) + } +} + +describe('remote_port_forward tool', () => { + it('opens a tunnel and returns a local preview url once it is ready', async () => { + const spawn = vi.fn(() => fakeChild()) + const runtime = createRemotePortForwardRuntime({ spawn, allocatePort: async () => 54321, probeReady: async () => true }) + const result = await runtime.provider.tools[0].execute({ action: 'open', remotePort: 3000 }, sshContext()) + const out = result.output as Record + expect(out.url).toBe('http://127.0.0.1:54321') + expect(out.status).toBe('ready') + }) + + it('refuses on a local thread', async () => { + const runtime = createRemotePortForwardRuntime({ spawn: () => fakeChild(), allocatePort: async () => 5000, probeReady: async () => true }) + const result = await runtime.provider.tools[0].execute({ action: 'open', remotePort: 3000 }, localContext()) + expect(result.isError).toBe(true) + }) + + it('validates the remote port', async () => { + const runtime = createRemotePortForwardRuntime({ spawn: () => fakeChild(), allocatePort: async () => 5000, probeReady: async () => true }) + const result = await runtime.provider.tools[0].execute({ action: 'open', remotePort: 0 }, sshContext()) + expect(result.isError).toBe(true) + }) + + it('disposes per thread', async () => { + const child = fakeChild() + const runtime = createRemotePortForwardRuntime({ spawn: () => child, allocatePort: async () => 5000, probeReady: async () => true }) + await runtime.provider.tools[0].execute({ action: 'open', remotePort: 8080 }, sshContext()) + runtime.disposeThreadResources('t') + expect(child.kill).toHaveBeenCalledWith('SIGTERM') + const listed = await runtime.provider.tools[0].execute({ action: 'list' }, sshContext()) + expect((listed.output as { forwards: unknown[] }).forwards).toHaveLength(0) + }) +}) diff --git a/kun/src/adapters/tool/remote-port-forward-tool.ts b/kun/src/adapters/tool/remote-port-forward-tool.ts new file mode 100644 index 000000000..14cc510ad --- /dev/null +++ b/kun/src/adapters/tool/remote-port-forward-tool.ts @@ -0,0 +1,173 @@ +import { spawn as nodeSpawn } from 'node:child_process' +import { connect as netConnect, createServer } from 'node:net' +import { LocalToolHost, type LocalTool } from './local-tool-host.js' +import { RemotePortForwardManager } from '../../remote/remote-port-forward-manager.js' +import type { SshChildProcess, SshSpawnFn } from '../../remote/ssh-executor.js' +import { isSshTarget } from '../../remote/remote-target.js' +import type { CapabilityToolProvider } from './capability-registry.js' +import { createApprovalRequest } from '../../domain/approval.js' + +const realSpawn: SshSpawnFn = (command, args) => + nodeSpawn(command, [...args], { shell: false }) as unknown as SshChildProcess + +function probeLocalPort(port: number, timeoutMs = 500): Promise { + return new Promise((resolve) => { + const socket = netConnect({ port, host: '127.0.0.1' }) + let settled = false + const done = (ok: boolean): void => { + if (settled) return + settled = true + socket.destroy() + resolve(ok) + } + socket.once('connect', () => done(true)) + socket.once('error', () => done(false)) + socket.setTimeout(timeoutMs, () => done(false)) + }) +} + +function allocateFreePort(): Promise { + return new Promise((resolve, reject) => { + const server = createServer() + server.once('error', reject) + server.listen(0, '127.0.0.1', () => { + const address = server.address() + const port = typeof address === 'object' && address ? address.port : 0 + server.close(() => (port ? resolve(port) : reject(new Error('could not allocate a local port')))) + }) + }) +} + +export type RemotePortForwardToolDeps = { + spawn?: SshSpawnFn + allocatePort?: () => Promise + probeReady?: (port: number) => Promise + readyTimeoutMs?: number +} + +export type RemotePortForwardRuntime = { + provider: CapabilityToolProvider + disposeThreadResources(threadId: string): void + shutdown(): void +} + +export function createRemotePortForwardRuntime(deps: RemotePortForwardToolDeps = {}): RemotePortForwardRuntime { + const readyTimeoutMs = deps.readyTimeoutMs ?? 5_000 + const probeReady = deps.probeReady ?? ((port: number) => probeLocalPort(port)) + const managers = new Map() + + const managerFor = (threadId: string): RemotePortForwardManager => { + const existing = managers.get(threadId) + if (existing) return existing + const manager = new RemotePortForwardManager({ + spawn: deps.spawn ?? realSpawn, + allocatePort: deps.allocatePort ?? allocateFreePort, + probeReady + }) + managers.set(threadId, manager) + return manager + } + + const tool = LocalToolHost.defineTool({ + name: 'remote_port_forward', + description: + 'Preview a remote service locally over an SSH tunnel (remote thread only). ' + + 'Default action is open; list and close are available for cleanup.', + inputSchema: { + type: 'object', + properties: { + action: { type: 'string', enum: ['open', 'list', 'close'] }, + remotePort: { type: 'number' }, + localPort: { type: 'number' }, + host: { type: 'string' }, + label: { type: 'string' }, + id: { type: 'string' } + }, + additionalProperties: false + }, + policy: 'on-request', + shouldAdvertise: (context) => context.executionTarget?.describe().target.kind === 'ssh', + execute: async (args, context) => { + const executionTarget = context.executionTarget + const target = executionTarget?.describe().target + if (!target || !isSshTarget(target)) { + return { output: { error: 'remote_port_forward requires a thread bound to an SSH target' }, isError: true } + } + const manager = managerFor(context.threadId) + const action = typeof args.action === 'string' ? args.action : 'open' + try { + if (action === 'list') { + return { output: { action, forwards: manager.list() } } + } + if (action === 'close') { + const id = typeof args.id === 'string' ? args.id : '' + const closed = manager.close(id) + return { output: { action, id, closed }, isError: !closed } + } + + const remotePort = typeof args.remotePort === 'number' ? args.remotePort : 0 + if (!Number.isInteger(remotePort) || remotePort < 1 || remotePort > 65535) { + return { output: { error: 'a valid remotePort (1-65535) is required' }, isError: true } + } + const rawLocalPort = typeof args.localPort === 'number' ? args.localPort : undefined + if (rawLocalPort !== undefined && rawLocalPort !== 0 && (rawLocalPort < 1024 || rawLocalPort > 65535)) { + return { output: { error: 'localPort must be omitted/0 or 1024-65535' }, isError: true } + } + if (executionTarget.production && context.approvalPolicy === 'auto') { + const approval = await context.awaitApproval(createApprovalRequest({ + id: 'appr_remote_port_forward_' + context.turnId, + threadId: context.threadId, + turnId: context.turnId, + toolName: 'remote_port_forward', + summary: 'Open SSH port forward for ' + target.alias + ':' + String(remotePort) + })) + if (approval !== 'allow') { + return { output: { error: 'remote port forward was not approved' }, isError: true } + } + } + const forward = await manager.open({ + alias: target.alias, + remotePort, + ...(typeof args.host === 'string' && args.host.trim() ? { remoteHost: args.host.trim() } : {}), + ...(typeof rawLocalPort === 'number' && rawLocalPort > 0 ? { localPort: rawLocalPort } : {}), + ...(readyTimeoutMs > 0 ? { waitForReadyMs: readyTimeoutMs } : {}) + }) + return { + output: { + target: 'ssh', + status: 'ready', + url: forward.url, + localPort: forward.localPort, + remotePort: forward.remotePort, + host: forward.remoteHost, + ...(typeof args.label === 'string' && args.label.trim() ? { label: args.label.trim() } : {}) + } + } + } catch (error) { + return { output: { error: error instanceof Error ? error.message : String(error) }, isError: true } + } + } + }) + + return { + provider: { + id: 'remote-port-forward', + kind: 'built-in', + enabled: true, + available: true, + tools: [tool] + }, + disposeThreadResources(threadId: string): void { + const manager = managers.get(threadId) + if (!manager) return + manager.closeAll() + managers.delete(threadId) + }, + shutdown(): void { + for (const manager of managers.values()) { + manager.closeAll() + } + managers.clear() + } + } +} diff --git a/kun/src/remote/remote-port-forward-manager.test.ts b/kun/src/remote/remote-port-forward-manager.test.ts new file mode 100644 index 000000000..488ea11f6 --- /dev/null +++ b/kun/src/remote/remote-port-forward-manager.test.ts @@ -0,0 +1,68 @@ +import { afterEach, describe, expect, it, vi } from 'vitest' +import { RemotePortForwardManager } from './remote-port-forward-manager.js' +import type { SshChildProcess } from './ssh-executor.js' + +function fakeChild() { + const handlers: Record void)[]> = {} + const child = { + stdout: null, + stderr: null, + on(event: string, listener: (...a: unknown[]) => void) { + ;(handlers[event] ??= []).push(listener as never) + }, + kill: vi.fn() + } as unknown as SshChildProcess + return { child, emit: (e: string, ...a: unknown[]) => handlers[e]?.forEach((h) => h(...a)) } +} + +describe('RemotePortForwardManager', () => { + afterEach(() => { + vi.useRealTimers() + }) + + it('opens a tunnel on an allocated local port and returns a preview url', async () => { + const { child } = fakeChild() + const spawn = vi.fn(() => child) + const manager = new RemotePortForwardManager({ spawn, allocatePort: async () => 54321 }) + const forward = await manager.open({ alias: 'prod', remotePort: 3000 }) + expect(forward.localPort).toBe(54321) + expect(forward.url).toBe('http://127.0.0.1:54321') + expect(spawn).toHaveBeenCalledWith('ssh', expect.arrayContaining(['-N', '-T', '-L'])) + expect(manager.list()).toHaveLength(1) + }) + + it('tears down and throws when the tunnel never becomes ready', async () => { + const { child } = fakeChild() + let now = 0 + const manager = new RemotePortForwardManager({ + spawn: () => child, + allocatePort: async () => 5000, + probeReady: async () => false, + delay: async () => { now += 100 }, + nowMs: () => now + }) + await expect(manager.open({ alias: 'prod', remotePort: 3000, waitForReadyMs: 200 })).rejects.toThrow(/did not become ready/) + expect(child.kill).toHaveBeenCalledWith('SIGTERM') + expect(manager.list()).toHaveLength(0) + }) + + it('drops a tunnel from the registry when its child exits on its own', async () => { + const { child, emit } = fakeChild() + const manager = new RemotePortForwardManager({ spawn: () => child, allocatePort: async () => 5000 }) + await manager.open({ alias: 'prod', remotePort: 8080 }) + expect(manager.list()).toHaveLength(1) + emit('close', 0, null) + expect(manager.list()).toHaveLength(0) + }) + + it('escalates from SIGTERM to SIGKILL when a tunnel does not exit', async () => { + vi.useFakeTimers() + const { child } = fakeChild() + const manager = new RemotePortForwardManager({ spawn: () => child, allocatePort: async () => 5000 }) + const forward = await manager.open({ alias: 'prod', remotePort: 8080 }) + manager.close(forward.id) + expect(child.kill).toHaveBeenCalledWith('SIGTERM') + await vi.advanceTimersByTimeAsync(2_000) + expect(child.kill).toHaveBeenCalledWith('SIGKILL') + }) +}) diff --git a/kun/src/remote/remote-port-forward-manager.ts b/kun/src/remote/remote-port-forward-manager.ts new file mode 100644 index 000000000..8fc3c9100 --- /dev/null +++ b/kun/src/remote/remote-port-forward-manager.ts @@ -0,0 +1,128 @@ +import { buildSshPortForwardArgv, type SshPortForwardOptions } from './ssh-port-forward.js' +import type { SshChildProcess, SshSpawnFn } from './ssh-executor.js' + +export type PortForward = { + id: string + localPort: number + remotePort: number + remoteHost: string + alias: string + url: string +} + +export type AllocatePortFn = () => Promise + +export type RemotePortForwardManagerDeps = { + spawn: SshSpawnFn + allocatePort: AllocatePortFn + nowMs?: () => number + probeReady?: (port: number) => Promise + delay?: (ms: number) => Promise +} + +type ActiveForward = { + forward: PortForward + child: SshChildProcess + hardKillTimer?: ReturnType +} + +const DEFAULT_READY_PROBE_INTERVAL_MS = 150 + +function defaultDelay(ms: number): Promise { + return new Promise((resolve) => { + const timer = setTimeout(resolve, ms) + if (timer && typeof timer === 'object' && 'unref' in timer) { + ;(timer as { unref: () => void }).unref() + } + }) +} + +export class RemotePortForwardManager { + private readonly forwards = new Map() + private counter = 0 + + constructor(private readonly deps: RemotePortForwardManagerDeps) {} + + async open(options: Omit & { localPort?: number; waitForReadyMs?: number }): Promise { + const localPort = options.localPort ?? (await this.deps.allocatePort()) + const argv = buildSshPortForwardArgv({ ...options, localPort }) + const child = this.deps.spawn('ssh', argv) + const id = 'pf_' + String(++this.counter) + const remoteHost = options.remoteHost ?? '127.0.0.1' + const forward: PortForward = { + id, + localPort, + remotePort: options.remotePort, + remoteHost, + alias: options.alias, + url: 'http://127.0.0.1:' + String(localPort) + } + const active: ActiveForward = { forward, child } + child.on('close', () => { + if (active.hardKillTimer) clearTimeout(active.hardKillTimer) + this.forwards.delete(id) + }) + child.on('error', () => { + if (active.hardKillTimer) clearTimeout(active.hardKillTimer) + this.forwards.delete(id) + }) + this.forwards.set(id, active) + + if (options.waitForReadyMs && options.waitForReadyMs > 0 && this.deps.probeReady) { + await this.waitForReady(id, localPort, options.waitForReadyMs) + } + return forward + } + + list(): PortForward[] { + return [...this.forwards.values()].map((entry) => entry.forward) + } + + close(id: string): boolean { + const entry = this.forwards.get(id) + if (!entry) return false + this.terminate(entry) + this.forwards.delete(id) + return true + } + + closeAll(): void { + for (const entry of this.forwards.values()) { + this.terminate(entry) + } + this.forwards.clear() + } + + private async waitForReady(id: string, port: number, timeoutMs: number): Promise { + const now = this.deps.nowMs ?? (() => Date.now()) + const delay = this.deps.delay ?? defaultDelay + const probe = this.deps.probeReady! + const deadline = now() + timeoutMs + for (;;) { + if (!this.forwards.has(id)) { + throw new Error('SSH tunnel exited before it became ready') + } + if (await probe(port)) return + if (now() >= deadline) { + this.close(id) + throw new Error('SSH tunnel did not become ready within ' + String(timeoutMs) + 'ms') + } + await delay(DEFAULT_READY_PROBE_INTERVAL_MS) + } + } + + private terminate(entry: ActiveForward): void { + entry.child.kill('SIGTERM') + if (entry.hardKillTimer) return + entry.hardKillTimer = setTimeout(() => { + try { + entry.child.kill('SIGKILL') + } catch { + // already exited + } + }, 2_000) + if (typeof entry.hardKillTimer === 'object' && 'unref' in entry.hardKillTimer) { + ;(entry.hardKillTimer as { unref: () => void }).unref() + } + } +} diff --git a/kun/src/remote/ssh-port-forward.test.ts b/kun/src/remote/ssh-port-forward.test.ts new file mode 100644 index 000000000..a6b33dd01 --- /dev/null +++ b/kun/src/remote/ssh-port-forward.test.ts @@ -0,0 +1,16 @@ +import { describe, expect, it } from 'vitest' +import { buildSshPortForwardArgv } from './ssh-port-forward.js' + +describe('ssh-port-forward', () => { + it('builds a loopback-bound -L tunnel that runs no remote command', () => { + expect(buildSshPortForwardArgv({ alias: 'prod', localPort: 5173, remotePort: 3000 })).toEqual([ + '-o', 'BatchMode=yes', '-o', 'ConnectTimeout=10', '-o', 'ExitOnForwardFailure=yes', + '-N', '-T', '-L', '127.0.0.1:5173:127.0.0.1:3000', 'prod' + ]) + }) + + it('rejects invalid ports', () => { + expect(() => buildSshPortForwardArgv({ alias: 'p', localPort: 0, remotePort: 0 })).toThrow(/invalid port/) + expect(() => buildSshPortForwardArgv({ alias: 'p', localPort: 70000, remotePort: 80 })).toThrow(/invalid port/) + }) +}) diff --git a/kun/src/remote/ssh-port-forward.ts b/kun/src/remote/ssh-port-forward.ts new file mode 100644 index 000000000..52b1235d6 --- /dev/null +++ b/kun/src/remote/ssh-port-forward.ts @@ -0,0 +1,46 @@ +import { isIP } from 'node:net' +import { assertSafeSshAlias, buildSshBaseArgs, type SshConnectionOptions } from './ssh-command.js' + +export type SshPortForwardOptions = SshConnectionOptions & { + localPort: number + remotePort: number + remoteHost?: string + localBindHost?: string +} + +export function buildSshPortForwardArgv(options: SshPortForwardOptions): string[] { + const localBindHost = options.localBindHost ?? '127.0.0.1' + const remoteHost = options.remoteHost ?? '127.0.0.1' + if (!isValidPort(options.localPort) || !isValidPort(options.remotePort)) { + throw new Error('invalid port for SSH forward (1-65535 required)') + } + const forwardSpec = formatForwardHost(localBindHost) + ':' + options.localPort + ':' + formatForwardHost(remoteHost) + ':' + options.remotePort + return [ + ...buildSshBaseArgs({ + ...options, + batchMode: options.batchMode ?? true, + connectTimeoutSec: options.connectTimeoutSec ?? 10 + }), + '-o', 'ExitOnForwardFailure=yes', + '-N', + '-T', + '-L', forwardSpec, + assertSafeSshAlias(options.alias) + ] +} + +function isValidPort(port: number): boolean { + return Number.isInteger(port) && port >= 1 && port <= 65535 +} + +function formatForwardHost(host: string): string { + const value = host.trim() + if (!value) throw new Error('SSH forward host must not be empty') + const version = isIP(value) + if (version === 4) return value + if (version === 6) return '[' + value + ']' + if (!/^(?=.{1,253}$)(?:[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?)(?:\.(?:[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?))*$/.test(value)) { + throw new Error('invalid SSH forward host: ' + host) + } + return value +} diff --git a/kun/src/server/runtime-factory.ts b/kun/src/server/runtime-factory.ts index 47bdb5b52..c13353222 100644 --- a/kun/src/server/runtime-factory.ts +++ b/kun/src/server/runtime-factory.ts @@ -26,6 +26,7 @@ import { buildDelegationToolProviders } from '../adapters/tool/delegation-tool-p import { buildWebToolProviders } from '../adapters/tool/web-tool-provider.js' import { buildImageGenToolProviders } from '../adapters/tool/image-gen-tool-provider.js' import { buildComputerUseToolProviders } from '../adapters/tool/computer-use-tool-provider.js' +import { createRemotePortForwardRuntime } from '../adapters/tool/remote-port-forward-tool.js' import { createRemoteHostsService } from '../remote/remote-hosts-service.js' import { RemoteTargetRegistry } from '../remote/remote-target-registry.js' import { @@ -274,6 +275,7 @@ export async function createKunServeRuntime( const backgroundShellTool = createBackgroundShellTool({ listBackgroundSessions: (threadId) => backgroundShellRuntime.listSessions(threadId) }) + const remotePortForwardRuntime = createRemotePortForwardRuntime() const withBackgroundShellTools = (tools: LocalTool[]): LocalTool[] => { const mapped = tools.map((tool) => tool.name === 'bash' @@ -458,6 +460,7 @@ export async function createKunServeRuntime( // Host control is available to the top-level agent only, never to // delegated subagents (which use childRegistry/baseToolProviders). ...computerUseProviders.providers, + remotePortForwardRuntime.provider, { id: 'goal', kind: 'gui' as const, @@ -605,6 +608,7 @@ export async function createKunServeRuntime( }, disposeThreadResources(threadId) { remoteTargetRegistry.evict(threadId) + remotePortForwardRuntime.disposeThreadResources(threadId) }, resumeInterruptedGoals(threadIds) { return loop.resumeInterruptedGoals(threadIds) @@ -667,6 +671,7 @@ export async function createKunServeRuntime( shutdown: async () => { try { loop.shutdownGoalResume() + remotePortForwardRuntime.shutdown() await mcpProviders.close() } finally { await stores.shutdown?.()