295 lines
11 KiB
TypeScript
295 lines
11 KiB
TypeScript
import { spawn as nodeSpawn, type ChildProcessWithoutNullStreams } from "node:child_process";
|
|
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
import type { AppConfig } from "../config.js";
|
|
import { secretValue } from "../config/secret-bundle.js";
|
|
import { clearPrincipalEnvironment } from "../auth/principal.js";
|
|
import { RpcClient } from "../rpc/rpc-client.js";
|
|
import { loadPiAuthProviders } from "./auth-providers.js";
|
|
import { buildPiChildEnv, canonicalPiProvider } from "./provider-credentials.js";
|
|
import type { PiReasoning } from "./management.js";
|
|
import {
|
|
PiManagedConfigError,
|
|
configuredPiProviderApiKey,
|
|
isPiManagedConfigError,
|
|
parsePiConfigJson,
|
|
readConfiguredPiAgentFile,
|
|
validateDeclarativePiConfig,
|
|
} from "./managed-config.js";
|
|
import {
|
|
loadRuntimeModelCatalog,
|
|
type RuntimeModelCatalog,
|
|
} from "../models/runtime-model-catalog.js";
|
|
|
|
const SMOKE_PROMPT = "Provider health check. Reply with exactly OK.";
|
|
const SMOKE_ARGS = [
|
|
"--mode", "rpc",
|
|
"--no-session",
|
|
"--no-tools",
|
|
"--no-extensions",
|
|
"--no-skills",
|
|
"--no-prompt-templates",
|
|
"--no-themes",
|
|
"--no-context-files",
|
|
"--no-approve",
|
|
] as const;
|
|
|
|
export interface PiProviderSmokeRequest {
|
|
provider: string;
|
|
model: string;
|
|
reasoning: PiReasoning;
|
|
timeoutMs: number;
|
|
}
|
|
|
|
export type PiProviderSmoke = (request: PiProviderSmokeRequest) => Promise<void>;
|
|
|
|
interface ProviderSmokeOptions {
|
|
spawnFn?: (
|
|
command: string,
|
|
args: string[],
|
|
options: { cwd: string; env: NodeJS.ProcessEnv },
|
|
) => ChildProcessWithoutNullStreams;
|
|
authProviders?: () => ReadonlySet<string>;
|
|
readAuthStore?: () => string;
|
|
readModelsStore?: () => string | undefined;
|
|
modelCatalog?: RuntimeModelCatalog;
|
|
}
|
|
|
|
export function createPiProviderSmoke(
|
|
config: AppConfig,
|
|
options: ProviderSmokeOptions = {},
|
|
): PiProviderSmoke {
|
|
const spawnFn = options.spawnFn ?? nodeSpawn;
|
|
const authProviders = options.authProviders ?? (() => loadPiAuthProviders());
|
|
|
|
return async ({ provider, model, reasoning, timeoutMs }): Promise<void> => {
|
|
let child: ChildProcessWithoutNullStreams | undefined;
|
|
let isolatedRoot: string | undefined;
|
|
let timer: NodeJS.Timeout | undefined;
|
|
try {
|
|
const canonicalProvider = canonicalPiProvider(provider);
|
|
if (!canonicalProvider || timeoutMs <= 0) throw providerFailure();
|
|
const configuredAuthProviders = authProviders();
|
|
const configuredModels = options.readModelsStore
|
|
? options.readModelsStore()
|
|
: readConfiguredPiAgentFile("models.json", true);
|
|
const catalog = options.modelCatalog ?? loadRuntimeModelCatalog(config.modelCatalogFile);
|
|
const catalogConfigured = config.modelCatalogFile !== undefined
|
|
|| catalog.defaultSession !== null;
|
|
const catalogModel = catalog.sessionModels()
|
|
.find((entry) => entry.provider === canonicalProvider && entry.model === model);
|
|
const upstreamModel = catalogModel?.upstreamModel ?? model;
|
|
const credentialName = catalogModel?.authentication.mode === "secret_env"
|
|
? catalogModel.authentication.apiKeyEnv
|
|
: undefined;
|
|
const projectedApiKey = configuredPiProviderApiKey(configuredModels, canonicalProvider)
|
|
?? (credentialName ? `$${credentialName}` : undefined);
|
|
const env = buildPiChildEnv({
|
|
provider: canonicalProvider,
|
|
authProviders: configuredAuthProviders,
|
|
credentialValue: credentialName
|
|
? secretValue(config, credentialName)
|
|
: catalogConfigured ? undefined : secretValue(config, "THT_MODEL_API_KEY"),
|
|
credentialFile: config.modelApiKeyFile,
|
|
configuredApiKey: projectedApiKey,
|
|
});
|
|
clearPrincipalEnvironment(env);
|
|
delete env.THT_DATA_ROOT;
|
|
delete env.THT_SESSION;
|
|
delete env.THT_AUTHOR;
|
|
delete env.THT_CONFIG;
|
|
delete env.PI_CODING_AGENT_SESSION_DIR;
|
|
|
|
isolatedRoot = mkdtempSync(join(tmpdir(), "thothii-pi-smoke-"));
|
|
const isolatedCwd = join(isolatedRoot, "work");
|
|
const isolatedAgentDir = join(isolatedRoot, "agent");
|
|
mkdirSync(isolatedCwd, { mode: 0o700 });
|
|
mkdirSync(isolatedAgentDir, { mode: 0o700 });
|
|
if (configuredAuthProviders.has(canonicalProvider)) {
|
|
const authStore = selectedProviderAuthStore(
|
|
options.readAuthStore?.() ?? readConfiguredPiAgentFile("auth.json"),
|
|
canonicalProvider,
|
|
);
|
|
writeDeclarativeAgentConfig(join(isolatedAgentDir, "auth.json"), authStore);
|
|
}
|
|
if (configuredModels !== undefined) {
|
|
const modelsStore = selectedProviderModelsStore(
|
|
configuredModels,
|
|
canonicalProvider,
|
|
model,
|
|
);
|
|
if (modelsStore !== undefined) {
|
|
writeDeclarativeAgentConfig(join(isolatedAgentDir, "models.json"), modelsStore);
|
|
}
|
|
}
|
|
env.PI_CODING_AGENT_DIR = isolatedAgentDir;
|
|
|
|
child = spawnFn(config.piBin, [...SMOKE_ARGS], { cwd: isolatedCwd, env });
|
|
child.stderr.resume();
|
|
const rpc = new RpcClient(child);
|
|
const capabilityGuard = failOnUnexpectedCapabilities(rpc);
|
|
const turn = async (): Promise<void> => {
|
|
requireSuccessfulResponse(await rpc.request({
|
|
type: "set_model", provider: canonicalProvider, modelId: upstreamModel,
|
|
} as object & { type: string }));
|
|
requireSuccessfulResponse(await rpc.request({
|
|
type: "set_thinking_level", level: reasoning,
|
|
} as object & { type: string }));
|
|
await waitForProviderTurn(rpc, child!);
|
|
};
|
|
await Promise.race([
|
|
turn(),
|
|
capabilityGuard,
|
|
new Promise<never>((_resolve, reject) => {
|
|
timer = setTimeout(() => reject(providerTimeout()), timeoutMs);
|
|
}),
|
|
]);
|
|
} catch (error) {
|
|
if (isProviderTimeout(error)) throw providerTimeout();
|
|
if (isPiManagedConfigError(error)) throw new PiManagedConfigError();
|
|
throw providerFailure();
|
|
} finally {
|
|
if (timer) clearTimeout(timer);
|
|
if (child) {
|
|
try { child.kill(); } catch { /* preserve the sanitized smoke outcome */ }
|
|
}
|
|
if (isolatedRoot) {
|
|
try { rmSync(isolatedRoot, { recursive: true, force: true, maxRetries: 2 }); } catch {
|
|
/* preserve the sanitized smoke outcome; the OS temp directory remains isolated */
|
|
}
|
|
}
|
|
}
|
|
};
|
|
}
|
|
|
|
function failOnUnexpectedCapabilities(rpc: RpcClient): Promise<never> {
|
|
return new Promise((_resolve, reject) => {
|
|
rpc.on("event", (event) => {
|
|
if (isUnexpectedCapabilityEvent(event)) reject(providerFailure());
|
|
});
|
|
});
|
|
}
|
|
|
|
function isUnexpectedCapabilityEvent(event: any): boolean {
|
|
const type = typeof event?.type === "string" ? event.type : "";
|
|
if (type.startsWith("tool_") || type.startsWith("toolcall_") || type.startsWith("extension_")) {
|
|
return true;
|
|
}
|
|
const updateType = event?.assistantMessageEvent?.type;
|
|
if (typeof updateType === "string" && updateType.startsWith("toolcall_")) return true;
|
|
if (Array.isArray(event?.toolResults) && event.toolResults.length > 0) return true;
|
|
if (messageUsesTool(event?.message)) return true;
|
|
return Array.isArray(event?.messages) && event.messages.some(messageUsesTool);
|
|
}
|
|
|
|
function messageUsesTool(message: any): boolean {
|
|
return message?.role === "toolResult" || message?.stopReason === "toolUse"
|
|
|| (Array.isArray(message?.content)
|
|
&& message.content.some((content: any) => content?.type === "toolCall"));
|
|
}
|
|
|
|
function writeDeclarativeAgentConfig(path: string, raw: string): void {
|
|
validateDeclarativePiConfig(raw);
|
|
writeFileSync(path, raw, { mode: 0o600, flag: "wx" });
|
|
}
|
|
|
|
function selectedProviderAuthStore(raw: string, provider: string): string {
|
|
const parsed = parsePiConfigJson(raw);
|
|
if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
|
|
throw new PiManagedConfigError();
|
|
}
|
|
const entry = Object.entries(parsed as Record<string, unknown>)
|
|
.find(([key]) => key.trim().toLowerCase() === provider);
|
|
if (!entry) throw new PiManagedConfigError();
|
|
return JSON.stringify({ [entry[0]]: entry[1] });
|
|
}
|
|
|
|
const PROVIDER_CONFIG_FIELDS = [
|
|
"name", "baseUrl", "apiKey", "api", "headers", "compat", "authHeader",
|
|
] as const;
|
|
|
|
function selectedProviderModelsStore(raw: string, provider: string, model: string): string | undefined {
|
|
const parsed = parsePiConfigJson(raw);
|
|
if (!parsed || typeof parsed !== "object" || Array.isArray(parsed)) {
|
|
throw new PiManagedConfigError();
|
|
}
|
|
const providers = (parsed as { providers?: unknown }).providers;
|
|
if (!providers || typeof providers !== "object" || Array.isArray(providers)) {
|
|
throw new PiManagedConfigError();
|
|
}
|
|
const entry = Object.entries(providers as Record<string, unknown>)
|
|
.find(([key]) => key.trim().toLowerCase() === provider);
|
|
if (!entry) return undefined;
|
|
const providerConfig = entry[1];
|
|
if (!providerConfig || typeof providerConfig !== "object" || Array.isArray(providerConfig)) {
|
|
throw new PiManagedConfigError();
|
|
}
|
|
const source = providerConfig as Record<string, unknown>;
|
|
const selected: Record<string, unknown> = {};
|
|
for (const field of PROVIDER_CONFIG_FIELDS) {
|
|
if (Object.hasOwn(source, field)) selected[field] = source[field];
|
|
}
|
|
if (Object.hasOwn(source, "models")) {
|
|
if (!Array.isArray(source.models)) throw new PiManagedConfigError();
|
|
let selectedModel: unknown;
|
|
for (const candidate of source.models) {
|
|
if (candidate && typeof candidate === "object" && !Array.isArray(candidate)
|
|
&& (candidate as { id?: unknown }).id === model) {
|
|
selectedModel = candidate;
|
|
}
|
|
}
|
|
if (selectedModel !== undefined) selected.models = [selectedModel];
|
|
}
|
|
if (Object.hasOwn(source, "modelOverrides")) {
|
|
const overrides = source.modelOverrides;
|
|
if (!overrides || typeof overrides !== "object" || Array.isArray(overrides)) {
|
|
throw new PiManagedConfigError();
|
|
}
|
|
if (Object.hasOwn(overrides, model)) {
|
|
selected.modelOverrides = { [model]: (overrides as Record<string, unknown>)[model] };
|
|
}
|
|
}
|
|
return JSON.stringify({ providers: { [entry[0]]: selected } });
|
|
}
|
|
|
|
function waitForProviderTurn(rpc: RpcClient, child: ChildProcessWithoutNullStreams): Promise<void> {
|
|
return new Promise((resolve, reject) => {
|
|
let failed = false;
|
|
rpc.on("event", (event) => {
|
|
if (event?.type === "message_end" && event.message?.role === "assistant"
|
|
&& event.message.stopReason === "error") {
|
|
failed = true;
|
|
reject(providerFailure());
|
|
return;
|
|
}
|
|
if (event?.type === "agent_end") {
|
|
const messages = Array.isArray(event.messages) ? event.messages : [];
|
|
const eventFailed = messages.some((message: any) => (
|
|
message?.role === "assistant" && message?.stopReason === "error"
|
|
));
|
|
if (failed || eventFailed) reject(providerFailure());
|
|
else resolve();
|
|
}
|
|
});
|
|
child.once("exit", () => reject(providerFailure()));
|
|
rpc.send({ type: "prompt", message: SMOKE_PROMPT });
|
|
});
|
|
}
|
|
|
|
function requireSuccessfulResponse(response: any): void {
|
|
if (!response || response.success !== true) throw providerFailure();
|
|
}
|
|
|
|
function providerFailure(): Error {
|
|
return new Error("Pi provider smoke check failed");
|
|
}
|
|
|
|
function providerTimeout(): Error {
|
|
return Object.assign(new Error("Pi smoke check timed out"), { code: "ETIMEDOUT" });
|
|
}
|
|
|
|
function isProviderTimeout(error: unknown): boolean {
|
|
return Boolean(error && typeof error === "object" && (error as { code?: unknown }).code === "ETIMEDOUT");
|
|
}
|