327 lines
14 KiB
TypeScript
327 lines
14 KiB
TypeScript
import { spawn as nodeSpawn, type ChildProcessWithoutNullStreams } from "node:child_process";
|
|
import type { AppConfig } from "../config.js";
|
|
import { RpcClient } from "../rpc/rpc-client.js";
|
|
import { SessionBridge } from "../bridge/session-bridge.js";
|
|
import type { RuntimeConfigLease, ThtRunner } from "../tht/tht-runner.js";
|
|
import { buildPiChildEnv, canonicalPiProvider } from "./provider-credentials.js";
|
|
import { loadPiAuthProviders } from "./auth-providers.js";
|
|
import { secretValue } from "../config/secret-bundle.js";
|
|
import { clearPrincipalEnvironment, principalEnvironment, type PrincipalContext } from "../auth/principal.js";
|
|
import {
|
|
configuredPiProviderApiKey,
|
|
createPiRuntimeAgentSnapshot,
|
|
} from "./managed-config.js";
|
|
import {
|
|
loadRuntimeModelCatalog,
|
|
type RuntimeModelCatalog,
|
|
} from "../models/runtime-model-catalog.js";
|
|
|
|
export interface SessionRuntime {
|
|
rpc: RpcClient;
|
|
bridge: SessionBridge;
|
|
child: ChildProcessWithoutNullStreams;
|
|
ownerKey?: string;
|
|
releaseRuntimeConfig?: () => void;
|
|
}
|
|
|
|
export interface RuntimeOptions {
|
|
provider?: string;
|
|
model?: string;
|
|
thinking?: string;
|
|
author?: string;
|
|
question?: string;
|
|
mode?: "new" | "resume";
|
|
principal?: PrincipalContext;
|
|
runtimeConfig?: RuntimeConfigLease;
|
|
}
|
|
|
|
/** Injectable child-process boundary; callbacks may ignore arguments in simpler tests. */
|
|
type SpawnFn = (
|
|
command: string,
|
|
args: string[],
|
|
options: { cwd: string; env: NodeJS.ProcessEnv },
|
|
) => ChildProcessWithoutNullStreams;
|
|
|
|
export class PiProcessManager {
|
|
private runtimes = new Map<string, SessionRuntime>();
|
|
private agentSnapshotCleanups = new WeakMap<ChildProcessWithoutNullStreams, () => void>();
|
|
private spawnFn: (
|
|
sessionId: string, author: string, provider: string | undefined, model: string | undefined,
|
|
principal?: PrincipalContext, runtimeConfigPath?: string,
|
|
) => ChildProcessWithoutNullStreams;
|
|
private loadAuthProviders: (agentDir: string) => ReadonlySet<string>;
|
|
private modelCatalog: RuntimeModelCatalog;
|
|
private modelCatalogConfigured: boolean;
|
|
|
|
constructor(
|
|
private cfg: AppConfig,
|
|
opts?: {
|
|
spawnFn?: SpawnFn;
|
|
authProviders?: (agentDir: string) => ReadonlySet<string>;
|
|
modelCatalog?: RuntimeModelCatalog;
|
|
},
|
|
) {
|
|
this.modelCatalog = opts?.modelCatalog ?? loadRuntimeModelCatalog(cfg.modelCatalogFile);
|
|
this.modelCatalogConfigured = cfg.modelCatalogFile !== undefined
|
|
|| this.modelCatalog.defaultSession !== null;
|
|
this.loadAuthProviders = opts?.authProviders
|
|
?? ((agentDir) => loadPiAuthProviders({ agentDir }));
|
|
if (opts?.spawnFn) {
|
|
this.spawnFn = (sessionId, author, provider, model, principal, runtimeConfigPath) =>
|
|
this.spawnPi(opts.spawnFn!, sessionId, author, provider, model, principal, runtimeConfigPath);
|
|
} else {
|
|
this.spawnFn = (sessionId, author, provider, model, principal, runtimeConfigPath) =>
|
|
this.spawnPi(nodeSpawn, sessionId, author, provider, model, principal, runtimeConfigPath);
|
|
}
|
|
}
|
|
|
|
private cleanupAgentSnapshot(child: ChildProcessWithoutNullStreams): void {
|
|
const cleanup = this.agentSnapshotCleanups.get(child);
|
|
if (!cleanup) return;
|
|
this.agentSnapshotCleanups.delete(child);
|
|
cleanup();
|
|
}
|
|
|
|
private spawnPi(
|
|
spawnFn: SpawnFn, sessionId: string, author: string, provider: string | undefined,
|
|
model: string | undefined, principal?: PrincipalContext, runtimeConfigPath?: string,
|
|
): ChildProcessWithoutNullStreams {
|
|
// This is the final shared boundary for createFor(), spawnFor(), and resume(). Validate
|
|
// before auth-provider inspection, then make Pi consume the exact copied bytes rather than
|
|
// reopening mutable mounted auth/models files after this check.
|
|
const agent = createPiRuntimeAgentSnapshot();
|
|
let child: ChildProcessWithoutNullStreams | undefined;
|
|
try {
|
|
const catalogModel = provider && model
|
|
? this.modelCatalog.sessionModels()
|
|
.find((entry) => entry.provider === provider && entry.model === model)
|
|
: undefined;
|
|
const credentialName = catalogModel?.authentication.mode === "secret_env"
|
|
? catalogModel.authentication.apiKeyEnv
|
|
: undefined;
|
|
const projectedApiKey = configuredPiProviderApiKey(agent.models, provider)
|
|
?? (credentialName ? `$${credentialName}` : undefined);
|
|
const env = buildPiChildEnv({
|
|
provider,
|
|
authProviders: this.loadAuthProviders(agent.agentDir),
|
|
credentialValue: credentialName
|
|
? secretValue(this.cfg, credentialName)
|
|
: this.modelCatalogConfigured ? undefined : secretValue(this.cfg, "THT_MODEL_API_KEY"),
|
|
credentialFile: this.cfg.modelApiKeyFile,
|
|
configuredApiKey: projectedApiKey,
|
|
additions: { THT_SESSION: sessionId, THT_AUTHOR: author },
|
|
});
|
|
env.PI_CODING_AGENT_DIR = agent.agentDir;
|
|
env.PI_CODING_AGENT_SESSION_DIR = agent.sessionDir;
|
|
clearPrincipalEnvironment(env);
|
|
if (principal) Object.assign(env, principalEnvironment(principal));
|
|
// The Thoth gate executes the deterministic `tht` CLI as a Pi tool. Give only
|
|
// this managed session process the adapter values already loaded by the core
|
|
// entrypoint; the generic provider helper continues to scrub them by default.
|
|
for (const name of [
|
|
"THT_DWH_API_KEY", "THT_VEC_API_KEY", "THT_VEC_WRITE_API_KEY",
|
|
] as const) {
|
|
const value = secretValue(this.cfg, name) ?? process.env[name];
|
|
if (value !== undefined) env[name] = value;
|
|
}
|
|
const ca = secretValue(this.cfg, "THT_SSL_CA")
|
|
?? secretValue(this.cfg, "THT_CA")
|
|
?? process.env.THT_SSL_CA
|
|
?? process.env.THT_CA;
|
|
if (ca !== undefined) {
|
|
env.THT_CA = ca;
|
|
env.THT_SSL_CA = ca;
|
|
}
|
|
delete env.THT_DATA_ROOT;
|
|
if (this.cfg.dataRoot !== undefined) env.THT_DATA_ROOT = this.cfg.dataRoot;
|
|
if (runtimeConfigPath !== undefined) env.THT_CONFIG = runtimeConfigPath;
|
|
// pi 0.73 removed `--approve`: rpc mode is headless and its argv is intentionally minimal.
|
|
child = spawnFn(this.cfg.piBin, ["--mode", "rpc"], {
|
|
cwd: this.cfg.harnessDir,
|
|
env,
|
|
});
|
|
this.agentSnapshotCleanups.set(child, agent.cleanup);
|
|
child.once("exit", () => this.cleanupAgentSnapshot(child!));
|
|
child.once("close", () => this.cleanupAgentSnapshot(child!));
|
|
// Log stderr for debugging (was silently drained)
|
|
child.stderr.on("data", (d: Buffer) => console.error(`[pi:${sessionId}] stderr:`, d.toString().trim()));
|
|
return child;
|
|
} catch (error) {
|
|
if (child) {
|
|
try { child.kill(); } catch { /* preserve the initialization error */ }
|
|
this.cleanupAgentSnapshot(child);
|
|
} else {
|
|
agent.cleanup();
|
|
}
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
count(): number { return this.runtimes.size; }
|
|
|
|
get(id: string): SessionRuntime | undefined { return this.runtimes.get(id); }
|
|
|
|
teardownForPrincipal(principal: PrincipalContext): string[] {
|
|
const ownerKey = `${principal.issuer}\0${principal.subject}`;
|
|
const stopped: string[] = [];
|
|
for (const [id, runtime] of [...this.runtimes.entries()]) {
|
|
if (runtime.ownerKey === ownerKey && this.teardownIfCurrent(id, runtime)) stopped.push(id);
|
|
}
|
|
return stopped;
|
|
}
|
|
|
|
/** Spawn and register a runtime synchronously, without starting a model turn. */
|
|
createFor(sessionId: string, o: RuntimeOptions = {}): SessionRuntime {
|
|
// A duplicate start must never tear down a live session: that used to send
|
|
// SIGTERM to the in-flight Pi process and lose its pending gate.
|
|
const existing = this.runtimes.get(sessionId);
|
|
if (existing) {
|
|
o.runtimeConfig?.release();
|
|
throw new Error(`session runtime already active: ${sessionId}`);
|
|
}
|
|
if (o.principal) this.teardownForPrincipal(o.principal);
|
|
if (this.runtimes.size >= this.cfg.maxPiProcesses) {
|
|
o.runtimeConfig?.release();
|
|
throw new Error("max Pi processes reached");
|
|
}
|
|
const author = o.author ?? "dev@local";
|
|
const provider = canonicalPiProvider(o.provider ?? this.cfg.defaults.provider);
|
|
const model = o.model ?? this.cfg.defaults.model;
|
|
let child: ChildProcessWithoutNullStreams;
|
|
try {
|
|
child = this.spawnFn(sessionId, author, provider, model, o.principal, o.runtimeConfig?.path);
|
|
} catch (error) {
|
|
o.runtimeConfig?.release();
|
|
throw error;
|
|
}
|
|
let runtimeConfigReleased = false;
|
|
const releaseRuntimeConfig = () => {
|
|
if (runtimeConfigReleased) return;
|
|
runtimeConfigReleased = true;
|
|
o.runtimeConfig?.release();
|
|
};
|
|
child.once("exit", releaseRuntimeConfig);
|
|
child.once("close", releaseRuntimeConfig);
|
|
let rt: SessionRuntime | undefined;
|
|
try {
|
|
const rpc = new RpcClient(child);
|
|
const bridge = new SessionBridge(rpc);
|
|
const runtime: SessionRuntime = {
|
|
rpc,
|
|
bridge,
|
|
child,
|
|
ownerKey: o.principal ? `${o.principal.issuer}\0${o.principal.subject}` : undefined,
|
|
...(o.runtimeConfig ? { releaseRuntimeConfig } : {}),
|
|
};
|
|
rt = runtime;
|
|
bridge.beginTurn();
|
|
this.runtimes.set(sessionId, runtime);
|
|
// Identity-checked: a stale child's exit must not evict a newer runtime.
|
|
// Expected teardowns (teardown()/respawn) delete the runtime from the map BEFORE the
|
|
// exit event fires, so reaching this branch with `rt` still mapped means the child
|
|
// died on its own: tell the client, or the UI spins forever waiting for a turn end.
|
|
child.on("exit", (code) => {
|
|
console.error(`[pi:${sessionId}] exited code=${code ?? "?"} mapped=${this.runtimes.get(sessionId) === runtime}`);
|
|
if (this.runtimes.get(sessionId) === runtime) {
|
|
// Clean exit (code 0) while the bridge is idle means the model completed its
|
|
// turn and Pi shut down normally (e.g. after session finalization). Just clean
|
|
// up — no failure events.
|
|
if (code === 0 && runtime.bridge.turnState() === "idle") {
|
|
this.runtimes.delete(sessionId);
|
|
runtime.bridge.emitClientEvent({ type: "system_event", event: "agent_end" });
|
|
return;
|
|
}
|
|
runtime.bridge.markFailed();
|
|
this.runtimes.delete(sessionId);
|
|
runtime.bridge.emitClientEvent({
|
|
type: "info",
|
|
level: "error",
|
|
text: `Pi process exited unexpectedly (code ${code ?? "?"})`,
|
|
});
|
|
runtime.bridge.emitClientEvent({ type: "system_event", event: "session_failed" });
|
|
runtime.bridge.emitClientEvent({ type: "system_event", event: "agent_end" });
|
|
}
|
|
});
|
|
|
|
return runtime;
|
|
} catch (error) {
|
|
if (rt && this.runtimes.get(sessionId) === rt) this.runtimes.delete(sessionId);
|
|
releaseRuntimeConfig();
|
|
try { child.kill(); } catch { /* preserve the initialization error */ }
|
|
this.cleanupAgentSnapshot(child);
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
/** Configure model and thinking. Safe to run alongside deterministic retrieval. */
|
|
async configure(rt: SessionRuntime, o: RuntimeOptions = {}): Promise<void> {
|
|
const provider = canonicalPiProvider(o.provider ?? this.cfg.defaults.provider);
|
|
const model = o.model ?? this.cfg.defaults.model;
|
|
const thinking = o.thinking ?? this.cfg.defaults.thinking;
|
|
|
|
if (provider && model) {
|
|
const upstreamModel = this.modelCatalog.sessionModels()
|
|
.find((entry) => entry.provider === provider && entry.model === model)
|
|
?.upstreamModel ?? model;
|
|
const response = await rt.rpc.request(
|
|
{ type: "set_model", provider, modelId: upstreamModel } as object & { type: string },
|
|
);
|
|
rt.bridge.setContextWindow(response?.data?.contextWindow);
|
|
}
|
|
if (thinking) {
|
|
await rt.rpc.request({ type: "set_thinking_level", level: thinking } as object & { type: string });
|
|
}
|
|
}
|
|
|
|
/** Start the first turn only after callers have attached the runtime bridge. */
|
|
start(sessionId: string, rt: SessionRuntime, o: RuntimeOptions = {}): void {
|
|
if (this.runtimes.get(sessionId) !== rt) throw new Error("session runtime is no longer active");
|
|
const message = o.mode === "resume"
|
|
? `/riprendi-sessione ${sessionId}`
|
|
: `/nuova-domanda ${JSON.stringify(o.question ?? "")}`;
|
|
rt.bridge.beginTurn();
|
|
rt.rpc.send({ type: "prompt", message });
|
|
}
|
|
|
|
async spawnFor(sessionId: string, o: RuntimeOptions = {}): Promise<SessionRuntime> {
|
|
const rt = this.createFor(sessionId, o);
|
|
try {
|
|
await this.configure(rt, o);
|
|
this.start(sessionId, rt, o);
|
|
} catch (error) {
|
|
// A rejected configure/start must not leak a registered runtime with a live
|
|
// child: every later start would see "session runtime already active".
|
|
this.teardownIfCurrent(sessionId, rt);
|
|
throw error;
|
|
}
|
|
return rt;
|
|
}
|
|
|
|
async resume(sessionId: string, tht: ThtRunner): Promise<SessionRuntime> {
|
|
const manifest = await tht.sessionShow(sessionId) as { provider?: string; model?: string; thinking?: string } | null;
|
|
return this.spawnFor(sessionId, {
|
|
provider: manifest?.provider,
|
|
model: manifest?.model,
|
|
thinking: manifest?.thinking,
|
|
mode: "resume",
|
|
});
|
|
}
|
|
|
|
teardown(id: string): void {
|
|
const rt = this.runtimes.get(id);
|
|
if (rt) this.teardownIfCurrent(id, rt);
|
|
}
|
|
|
|
/** Remove only the runtime identity the caller observed. */
|
|
teardownIfCurrent(id: string, expected: SessionRuntime): boolean {
|
|
if (this.runtimes.get(id) !== expected) return false;
|
|
// Delete before signalling the child so its asynchronous exit cannot be mistaken for a
|
|
// crash, and so a replacement installed by a later lifecycle operation is never targeted.
|
|
this.runtimes.delete(id);
|
|
expected.releaseRuntimeConfig?.();
|
|
expected.child.kill();
|
|
this.cleanupAgentSnapshot(expected.child);
|
|
return true;
|
|
}
|
|
}
|