Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/klaus/src/runtime/index.ts

Raw
import { mkdtemp, rm } from "node:fs/promises";
import { validateHeaderName, validateHeaderValue } from "node:http";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { dbg, diagnosticKind } from "../debug.js";
import { resolveAndroidRuntime } from "./android.js";
import { resolveStandardRuntime } from "./standard.js";

const ENV_ALLOWLIST = [
	"PATH",
	"PATHEXT",
	"SystemRoot",
	"WINDIR",
	"ComSpec",
	"TMP",
	"TEMP",
	"TMPDIR",
	"LANG",
	"LC_ALL",
	"LC_CTYPE",
	"HTTP_PROXY",
	"HTTPS_PROXY",
	"NO_PROXY",
	"ALL_PROXY",
	"SSL_CERT_FILE",
	"SSL_CERT_DIR",
	"NODE_EXTRA_CA_CERTS",
	"PREFIX",
	"LD_LIBRARY_PATH",
] as const;

const RESERVED_HEADERS = new Set([
	"authorization",
	"connection",
	"content-length",
	"content-type",
	"host",
	"proxy-authorization",
	"transfer-encoding",
	"x-api-key",
]);

export interface ChildRuntime {
	executable?: string;
	configDir: string;
	env: Record<string, string>;
	betas: Array<"context-1m-2025-08-07">;
	cleanup(): Promise<void>;
}

function customHeaders(headers: Record<string, string | null>): {
	serialized?: string;
	betas: Array<"context-1m-2025-08-07">;
} {
	dbg?.("runtime.customHeaders.start");
	const lines: string[] = [];
	const betas: Array<"context-1m-2025-08-07"> = [];
	for (const [name, value] of Object.entries(headers)) {
		if (value === null) continue;
		const lower = name.toLowerCase();
		if (RESERVED_HEADERS.has(lower) || lower.startsWith("x-klaus-")) continue;
		try {
			validateHeaderName(name);
			validateHeaderValue(name, value);
		} catch {
			continue;
		}
		if (lower === "anthropic-beta") {
			if (
				value
					.split(",")
					.map((part) => part.trim())
					.includes("context-1m-2025-08-07")
			) {
				betas.push("context-1m-2025-08-07");
			}
			continue;
		}
		if (name.includes("\n") || value.includes("\n")) continue;
		lines.push(`${name}: ${value}`);
	}
	dbg?.("runtime.customHeaders.end", { forwardedCount: lines.length });
	return { serialized: lines.length ? lines.join("\n") : undefined, betas };
}

export async function createChildRuntime(
	oauthToken: string,
	headers: Record<string, string | null>,
	overrides: Record<string, string> = {},
	maxTokens?: number,
): Promise<ChildRuntime> {
	dbg?.("runtime.create.start");
	const configDir = await mkdtemp(join(tmpdir(), "klaus-query-"));
	dbg?.("runtime.create.configDir");
	const env: Record<string, string> = {};
	for (const key of ENV_ALLOWLIST) {
		const value = overrides[key] ?? process.env[key];
		if (value !== undefined) env[key] = value;
	}
	const mapped = customHeaders(headers);
	env.HOME = configDir;
	env.CLAUDE_CONFIG_DIR = configDir;
	env.CLAUDE_CODE_OAUTH_TOKEN = oauthToken;
	env.CLAUDE_AGENT_SDK_CLIENT_APP = "klaus/1";
	/** Pi owns every tool Klaus exposes, so the child presents them under their
	 * Pi names instead of the `mcp__klaus__` MCP namespace. The bridge still
	 * strips that prefix, so resumed sessions and runtimes without this flag
	 * keep working. */
	env.CLAUDE_AGENT_SDK_MCP_NO_PREFIX = "1";
	env.CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC = "1";
	env.CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY = "0";
	/** Pi sends its whole tool list on every request, so the child must define
	 * every tool upfront instead of deferring them behind MCP tool search. */
	env.ENABLE_TOOL_SEARCH = "false";
	/** A Klaus tool call stays parked until Pi executes it, which includes human
	 * approval and has no upper bound, so neither automatic backgrounding nor the
	 * idle watchdog may abandon it. */
	env.CLAUDE_CODE_MCP_AUTO_BACKGROUND_MS = "0";
	env.CLAUDE_CODE_MCP_TOOL_IDLE_TIMEOUT = "0";
	/** Klaus only forwards `stream_event` content, so a silent non-streaming
	 * retry would hand Pi an empty but successful turn. */
	env.CLAUDE_CODE_DISABLE_NONSTREAMING_FALLBACK = "1";
	/** Text results are bounded by the per-tool `anthropic/maxResultSizeChars`
	 * the bridge declares; this raises the remaining token cap that still governs
	 * image content in Pi tool results. */
	env.MAX_MCP_OUTPUT_TOKENS = "1000000";
	/** Pi owns the system prompt and the tool list, so every ambient Claude Code
	 * context source stays off: memory files, bundled and managed skills,
	 * built-in subagents, workflows, cron, background tasks, and checkpoints. */
	env.CLAUDE_CODE_DISABLE_CLAUDE_MDS = "1";
	env.CLAUDE_CODE_DISABLE_AUTO_MEMORY = "1";
	env.CLAUDE_CODE_DISABLE_BUNDLED_SKILLS = "1";
	env.CLAUDE_CODE_DISABLE_POLICY_SKILLS = "1";
	env.CLAUDE_AGENT_SDK_DISABLE_BUILTIN_AGENTS = "1";
	env.CLAUDE_CODE_DISABLE_WORKFLOWS = "1";
	env.CLAUDE_CODE_DISABLE_CRON = "1";
	env.CLAUDE_CODE_DISABLE_BACKGROUND_TASKS = "1";
	env.CLAUDE_CODE_DISABLE_FILE_CHECKPOINTING = "1";
	/** Child temp files land inside the per-query config directory, so the
	 * existing cleanup removes them. */
	env.CLAUDE_CODE_TMPDIR = configDir;
	if (maxTokens !== undefined) {
		if (!Number.isSafeInteger(maxTokens) || maxTokens < 1) {
			throw new Error(
				`Klaus received an invalid maxTokens value ${String(maxTokens)}.`,
			);
		}
		env.CLAUDE_CODE_MAX_OUTPUT_TOKENS = String(maxTokens);
	}
	if (mapped.serialized) env.ANTHROPIC_CUSTOM_HEADERS = mapped.serialized;
	if (
		process.env.VITEST &&
		oauthToken === "fake" &&
		process.env.KLAUS_E2E_BASE_URL
	) {
		env.ANTHROPIC_BASE_URL = process.env.KLAUS_E2E_BASE_URL;
	}
	const cleanup = async () => {
		dbg?.("runtime.cleanup.start");
		await rm(configDir, {
			recursive: true,
			force: true,
			maxRetries: 10,
			retryDelay: 100,
		});
		dbg?.("runtime.cleanup.end");
	};
	try {
		const executable =
			process.platform === "android"
				? await resolveAndroidRuntime()
				: await resolveStandardRuntime();
		dbg?.("runtime.create.ready");
		return {
			executable,
			configDir,
			env,
			betas: mapped.betas,
			cleanup,
		};
	} catch (error) {
		dbg?.("runtime.create.error", { kind: diagnosticKind(error) });
		await cleanup();
		throw error;
	}
}