|
| 1 | +import type { ExtensionAPI } from "@earendil-works/pi-coding-agent"; |
| 2 | + |
| 3 | +import * as cli from "./cli.ts"; |
| 4 | +import { syncProjectOnce } from "./startup-sync.ts"; |
| 5 | +import { shellQuote } from "./util.ts"; |
| 6 | + |
| 7 | +const CONCURRENCY = 3; |
| 8 | +const REMOTE_DIR = "/root/workspace"; |
| 9 | + |
| 10 | +export interface FanoutScenario { |
| 11 | + name: string; |
| 12 | + command: string; |
| 13 | + environment: Array<{ name: string; value: string }>; |
| 14 | + port?: number; |
| 15 | + healthCheckPath?: string; |
| 16 | + healthCheckContains?: string; |
| 17 | +} |
| 18 | + |
| 19 | +export interface FanoutOptions { |
| 20 | + sourceDir: string; |
| 21 | + namePrefix: string; |
| 22 | + scenarios: FanoutScenario[]; |
| 23 | + shape?: string; |
| 24 | + rootfs?: string; |
| 25 | +} |
| 26 | + |
| 27 | +export interface FanoutResult { |
| 28 | + name: string; |
| 29 | + sandboxId?: string; |
| 30 | + url?: string; |
| 31 | + verified: boolean; |
| 32 | + output?: string; |
| 33 | + error?: string; |
| 34 | +} |
| 35 | + |
| 36 | +export function createScenarioCommand(scenario: FanoutScenario): string { |
| 37 | + const environment = scenario.environment.map(({ name, value }) => `${name}=${shellQuote(value)}`).join(" "); |
| 38 | + return `cd ${shellQuote(REMOTE_DIR)} && ${environment ? `${environment} ` : ""}${scenario.command}`; |
| 39 | +} |
| 40 | + |
| 41 | +export async function fanoutScenarios( |
| 42 | + pi: ExtensionAPI, |
| 43 | + options: FanoutOptions, |
| 44 | + signal?: AbortSignal, |
| 45 | +): Promise<FanoutResult[]> { |
| 46 | + const results: FanoutResult[] = []; |
| 47 | + let nextIndex = 0; |
| 48 | + |
| 49 | + async function worker(): Promise<void> { |
| 50 | + while (!signal?.aborted) { |
| 51 | + const index = nextIndex++; |
| 52 | + if (index >= options.scenarios.length) return; |
| 53 | + results[index] = await runScenario(pi, options, options.scenarios[index], signal); |
| 54 | + } |
| 55 | + } |
| 56 | + |
| 57 | + await Promise.all(Array.from({ length: Math.min(CONCURRENCY, options.scenarios.length) }, worker)); |
| 58 | + if (signal?.aborted) throw new Error("aborted"); |
| 59 | + return results; |
| 60 | +} |
| 61 | + |
| 62 | +async function runScenario( |
| 63 | + pi: ExtensionAPI, |
| 64 | + options: FanoutOptions, |
| 65 | + scenario: FanoutScenario, |
| 66 | + signal?: AbortSignal, |
| 67 | +): Promise<FanoutResult> { |
| 68 | + const name = `${options.namePrefix}-${scenario.name}`; |
| 69 | + let sandboxId: string | undefined; |
| 70 | + |
| 71 | + try { |
| 72 | + const sandbox = await cli.createSandbox(pi, { |
| 73 | + shape: options.shape, |
| 74 | + rootfs: options.rootfs, |
| 75 | + name, |
| 76 | + ingress: scenario.port !== undefined, |
| 77 | + }); |
| 78 | + sandboxId = sandbox.id; |
| 79 | + await syncProjectOnce(pi, sandbox.id, options.sourceDir, {}, signal); |
| 80 | + const command = createScenarioCommand(scenario); |
| 81 | + |
| 82 | + if (!scenario.port) return await runForeground(pi, sandbox.id, name, command, signal); |
| 83 | + |
| 84 | + const started = await cli.sandboxExec( |
| 85 | + pi, |
| 86 | + sandbox.id, |
| 87 | + `(tmux kill-session -t scenario 2>/dev/null || true) && tmux new-session -d -s scenario ${shellQuote(command)}`, |
| 88 | + signal, |
| 89 | + ); |
| 90 | + if (started.exitCode !== 0) throw new Error(started.stdout || "Scenario start command failed"); |
| 91 | + |
| 92 | + const info = await cli.getSandbox(pi, sandbox.id); |
| 93 | + const url = info.ingress_url_template?.replace("<port>", String(scenario.port)); |
| 94 | + if (!url) throw new Error("Ingress URL is unavailable"); |
| 95 | + |
| 96 | + const verified = await verifyScenario(url, scenario, signal); |
| 97 | + if (!verified) { |
| 98 | + await cli.destroySandbox(pi, sandbox.id).catch(() => undefined); |
| 99 | + return { name, sandboxId: sandbox.id, verified: false, error: "public health check failed" }; |
| 100 | + } |
| 101 | + return { name, sandboxId: sandbox.id, url, verified: true }; |
| 102 | + } catch (error) { |
| 103 | + if (sandboxId) await cli.destroySandbox(pi, sandboxId).catch(() => undefined); |
| 104 | + return { |
| 105 | + name, |
| 106 | + sandboxId, |
| 107 | + verified: false, |
| 108 | + error: error instanceof Error ? error.message : String(error), |
| 109 | + }; |
| 110 | + } |
| 111 | +} |
| 112 | + |
| 113 | +async function runForeground( |
| 114 | + pi: ExtensionAPI, |
| 115 | + sandboxId: string, |
| 116 | + name: string, |
| 117 | + command: string, |
| 118 | + signal?: AbortSignal, |
| 119 | +): Promise<FanoutResult> { |
| 120 | + const result = await cli.sandboxExec(pi, sandboxId, command, signal); |
| 121 | + try { |
| 122 | + await cli.destroySandbox(pi, sandboxId); |
| 123 | + } catch (error) { |
| 124 | + return { |
| 125 | + name, |
| 126 | + sandboxId, |
| 127 | + verified: false, |
| 128 | + output: result.stdout, |
| 129 | + error: `scenario cleanup failed: ${error instanceof Error ? error.message : String(error)}`, |
| 130 | + }; |
| 131 | + } |
| 132 | + if (result.exitCode !== 0) { |
| 133 | + return { name, sandboxId, verified: false, output: result.stdout, error: "scenario command failed" }; |
| 134 | + } |
| 135 | + return { name, sandboxId, verified: true, output: result.stdout }; |
| 136 | +} |
| 137 | + |
| 138 | +export function createHealthCheckUrl(url: string, path = "/"): URL { |
| 139 | + if (!path.startsWith("/") || path.startsWith("//")) throw new Error("health check path must be relative"); |
| 140 | + const endpoint = new URL(path, url); |
| 141 | + if (endpoint.origin !== new URL(url).origin) throw new Error("health check must use the sandbox ingress URL"); |
| 142 | + return endpoint; |
| 143 | +} |
| 144 | + |
| 145 | +async function verifyScenario(url: string, scenario: FanoutScenario, signal?: AbortSignal): Promise<boolean> { |
| 146 | + const endpoint = createHealthCheckUrl(url, scenario.healthCheckPath); |
| 147 | + for (let attempt = 0; attempt < 15; attempt += 1) { |
| 148 | + if (signal?.aborted) throw new Error("aborted"); |
| 149 | + try { |
| 150 | + const timeout = AbortSignal.timeout(5_000); |
| 151 | + const requestSignal = signal ? AbortSignal.any([signal, timeout]) : timeout; |
| 152 | + const response = await fetch(endpoint, { signal: requestSignal, redirect: "error" }); |
| 153 | + const body = await response.text(); |
| 154 | + if (response.ok && (!scenario.healthCheckContains || body.includes(scenario.healthCheckContains))) return true; |
| 155 | + } catch { |
| 156 | + // The process may still be starting. |
| 157 | + } |
| 158 | + await new Promise((resolve) => setTimeout(resolve, 1_000)); |
| 159 | + } |
| 160 | + return false; |
| 161 | +} |
0 commit comments