@prometheus-ai/agent 0.5.3 → 0.5.8

package/src/tiny/title-client.ts

@@ -1,5 +1,5 @@
 import * as path from "node:path";
-import { $env, isCompiledBinary, logger } from "@prometheus-ai/utils";
+import { $env, isBunTestRuntime, isCompiledBinary, logger, workerHostEntry } from "@prometheus-ai/utils";
 import type { Subprocess } from "bun";
 import { settings } from "../config/settings";
 import { tinyModelDeviceSettingToEnv } from "./device";
@@ -39,7 +39,31 @@ export interface TinyTitleDownloadOptions {
 	onProgress?: (event: TinyTitleProgressEvent) => void;
 }
 
-const SMOKE_TEST_TIMEOUT_MS = 5_000;
+/**
+ * Per-request controls for {@link TinyTitleClient.generate}.
+ *
+ * Carries the optional abort signal and title-system-prompt override used by
+ * callers that customize automatic session-title generation.
+ */
+export interface TinyTitleGenerateOptions {
+	signal?: AbortSignal;
+	systemPrompt?: string;
+}
+
+// Cold-starting the worker subprocess from a compiled binary (decompress + module
+// graph load) is slow on contended CI runners — the macos-15-intel release smoke
+// blew past 5s while arm64/linux/win passed. The probe only needs to prove the
+// worker spawns and ponges at all (a dead worker never ponges regardless), so a
+// generous bound removes the flake without weakening the check.
+const SMOKE_TEST_TIMEOUT_MS = 30_000;
+
+function normalizeTinyTitleGenerateOptions(
+	options: AbortSignal | TinyTitleGenerateOptions | undefined,
+): TinyTitleGenerateOptions {
+	if (!options) return {};
+	if ("aborted" in options && "addEventListener" in options) return { signal: options };
+	return options;
+}
 
 /**
  * Hidden subcommand on the main CLI that boots the tiny-model worker in the
@@ -103,21 +127,32 @@ function tinyWorkerEnv(): Record<string, string> {
 	for (const key in overlay) merged[key] = overlay[key];
 	return merged;
 }
+interface TinyWorkerSpawnCommand {
+	cmd: string[];
+	cwd?: string;
+}
 
 /**
- * Resolve the argv used to relaunch the agent CLI into tiny-worker mode. In a
- * compiled binary the entry point is the binary itself; in dev/source the
- * spawned `bun` needs the absolute path to `cli.ts` so it can resolve module
- * imports against the on-disk source tree.
+ * Resolve the command used to relaunch the agent CLI into tiny-worker mode.
+ * In a compiled binary the entry point is the binary itself (no script arg).
+ * Otherwise re-enter the declared worker-host entry (source cli.ts or
+ * npm-bundle cli.js) with a cwd-relative script path — Bun's subprocess IPC
+ * is more reliable that way than with an absolute `.ts` entry under
+ * `bun test` — and fall back to this package's own `src/cli.ts` when no host
+ * entry is declared (bun test, SDK embedding).
  */
-function tinyWorkerSpawnCmd(): string[] {
-	if (isCompiledBinary()) return [process.execPath, TINY_WORKER_ARG];
-	const cliPath = path.resolve(import.meta.dir, "..", "cli.ts");
-	return [process.execPath, cliPath, TINY_WORKER_ARG];
+function tinyWorkerSpawnCmd(): TinyWorkerSpawnCommand {
+	if (isCompiledBinary()) return { cmd: [process.execPath, TINY_WORKER_ARG] };
+	const hostEntry = workerHostEntry();
+	if (hostEntry) {
+		return { cmd: [process.execPath, path.basename(hostEntry), TINY_WORKER_ARG], cwd: path.dirname(hostEntry) };
+	}
+	const packageRoot = path.resolve(import.meta.dir, "..", "..");
+	return { cmd: [process.execPath, "src/cli.ts", TINY_WORKER_ARG], cwd: packageRoot };
 }
 
 interface SpawnedSubprocess {
-	proc: Subprocess<"ignore", "inherit", "inherit">;
+	proc: Subprocess<"ignore", "ignore", "ignore">;
 	inbound: Set<(message: TinyTitleWorkerOutbound) => void>;
 	errors: Set<(error: Error) => void>;
 	/**
@@ -138,14 +173,19 @@ export function createTinyTitleSubprocess(): SpawnedSubprocess {
 	const inbound = new Set<(message: TinyTitleWorkerOutbound) => void>();
 	const errors = new Set<(error: Error) => void>();
 	const intentionalExit = { value: false };
+	const spawnCommand = tinyWorkerSpawnCmd();
 	const proc = Bun.spawn({
-		cmd: tinyWorkerSpawnCmd(),
+		cmd: spawnCommand.cmd,
+		cwd: spawnCommand.cwd,
 		env: tinyWorkerEnv(),
 		stdin: "ignore",
-		stdout: "inherit",
-		stderr: "inherit",
+		stdout: "ignore",
+		stderr: "ignore",
 		serialization: "advanced",
 		windowsHide: true,
+		// The worker is an implementation detail of the interactive TUI. Native
+		// model runtimes may print progress or decoded text directly; never let
+		// those bytes inherit the terminal and corrupt the chat scrollback.
 		ipc(message) {
 			for (const handler of inbound) handler(message as TinyTitleWorkerOutbound);
 		},
@@ -167,7 +207,9 @@ export function createTinyTitleSubprocess(): SpawnedSubprocess {
 	});
 	// Don't keep the parent event loop alive on account of an idle worker; the
 	// agent dispose path calls `terminate()` explicitly when shutting down.
-	proc.unref();
+	// Bun's test runner can starve IPC delivery for unref'd subprocesses, so
+	// keep it referenced only under tests that assert the ping/pong contract.
+	if (!isBunTestRuntime()) proc.unref();
 	return { proc, inbound, errors, intentionalExit };
 }
 
@@ -272,9 +314,16 @@ export class TinyTitleClient {
 		return () => this.#progressListeners.delete(listener);
 	}
 
-	async generate(modelKey: string, message: string, signal?: AbortSignal): Promise<string | null> {
+	async generate(modelKey: string, message: string, signal?: AbortSignal): Promise<string | null>;
+	async generate(modelKey: string, message: string, options?: TinyTitleGenerateOptions): Promise<string | null>;
+	async generate(
+		modelKey: string,
+		message: string,
+		optionsOrSignal?: AbortSignal | TinyTitleGenerateOptions,
+	): Promise<string | null> {
+		const options = normalizeTinyTitleGenerateOptions(optionsOrSignal);
 		if (!isTinyTitleLocalModelKey(modelKey)) return null;
-		if (signal?.aborted) return null;
+		if (options.signal?.aborted) return null;
 
 		try {
 			const worker = this.#ensureWorker();
@@ -287,12 +336,15 @@ export class TinyTitleClient {
 				this.#pending.delete(id);
 				pending.resolve(null);
 			};
-			signal?.addEventListener("abort", abort, { once: true });
+			options.signal?.addEventListener("abort", abort, { once: true });
 			try {
-				worker.send({ type: "generate", id, modelKey, message });
+				const request: TinyTitleWorkerInbound = options.systemPrompt
+					? { type: "generate", id, modelKey, message, systemPrompt: options.systemPrompt }
+					: { type: "generate", id, modelKey, message };
+				worker.send(request);
 				return await promise;
 			} finally {
-				signal?.removeEventListener("abort", abort);
+				options.signal?.removeEventListener("abort", abort);
 				this.#pending.delete(id);
 			}
 		} catch (error) {

package/src/tiny/title-protocol.ts

@@ -29,7 +29,7 @@ export interface TinyTitleProgressEvent {
 
 export type TinyTitleWorkerInbound =
 	| { type: "ping"; id: string }
-	| { type: "generate"; id: string; modelKey: TinyTitleLocalModelKey; message: string }
+	| { type: "generate"; id: string; modelKey: TinyTitleLocalModelKey; message: string; systemPrompt?: string }
 	| { type: "complete"; id: string; modelKey: TinyLocalModelKey; prompt: string; maxTokens?: number }
 	| { type: "download"; id: string; modelKey: TinyLocalModelKey };
 

package/src/tiny/worker.ts

@@ -1,4 +1,3 @@
-import * as fs from "node:fs/promises";
 import { createRequire } from "node:module";
 import * as path from "node:path";
 import type {
@@ -7,10 +6,16 @@ import type {
 	TextGenerationStringOutput,
 	StoppingCriteria as TransformersStoppingCriteria,
 } from "@huggingface/transformers";
-import { getTinyModelsCacheDir, isCompiledBinary, prompt } from "@prometheus-ai/utils";
+import {
+	ensureRuntimeInstalled,
+	getTinyModelsCacheDir,
+	installRuntimeModuleResolver,
+	isCompiledBinary,
+	prompt,
+	resolveRuntimeModule,
+} from "@prometheus-ai/utils";
 import packageJson from "../../package.json" with { type: "json" };
 import tinyTitleSystemPrompt from "../prompts/system/tiny-title-system.md" with { type: "text" };
-import { installRuntimeModuleResolver, resolveRuntimeModule } from "./compiled-runtime";
 import { resolveTinyModelDevicePreference, type TinyModelDevice, tinyModelDeviceLoadOrder } from "./device";
 import { resolveTinyModelDtypeOverride, type TinyModelDtype } from "./dtype";
 import {
@@ -31,8 +36,6 @@ const TINY_TITLE_SYSTEM_PROMPT = prompt.render(tinyTitleSystemPrompt);
 const TRANSFORMERS_PACKAGE = "@huggingface/transformers";
 const COMPILED_TRANSFORMERS_VERSION = process.env.PROMETHEUS_TINY_TRANSFORMERS_VERSION;
 const sourceRequire = createRequire(import.meta.url);
-const INSTALL_LOCK_ATTEMPTS = 240;
-const INSTALL_LOCK_SLEEP_MS = 250;
 
 const tinyModelDevicePreference = resolveTinyModelDevicePreference();
 const tinyModelDtypeOverride = resolveTinyModelDtypeOverride();
@@ -97,10 +100,6 @@ function errorText(error: unknown): string {
 	return error instanceof Error ? (error.stack ?? error.message) : String(error);
 }
 
-function isErrnoCode(error: unknown, code: string): boolean {
-	return typeof error === "object" && error !== null && "code" in error && error.code === code;
-}
-
 function sendLog(
 	transport: TinyTitleTransport,
 	level: "debug" | "warn" | "error",
@@ -118,69 +117,6 @@ function getTinyTitleRuntimeDir(): string {
 	);
 }
 
-async function acquireInstallLock(runtimeDir: string): Promise<() => Promise<void>> {
-	const lockDir = `${runtimeDir}.lock`;
-	await fs.mkdir(path.dirname(lockDir), { recursive: true });
-	for (let attempt = 0; attempt < INSTALL_LOCK_ATTEMPTS; attempt++) {
-		try {
-			await fs.mkdir(lockDir);
-			return async () => {
-				await fs.rm(lockDir, { recursive: true, force: true });
-			};
-		} catch (error) {
-			if (!isErrnoCode(error, "EEXIST")) throw error;
-			await Bun.sleep(INSTALL_LOCK_SLEEP_MS);
-		}
-	}
-	throw new Error(`Timed out waiting for tiny title runtime install lock: ${lockDir}`);
-}
-
-async function isCompiledRuntimeInstalled(runtimeDir: string): Promise<boolean> {
-	return Bun.file(path.join(runtimeDir, "node_modules", "@huggingface", "transformers", "package.json")).exists();
-}
-
-async function writeRuntimeManifest(runtimeDir: string): Promise<void> {
-	await fs.mkdir(runtimeDir, { recursive: true });
-	await Bun.write(
-		path.join(runtimeDir, "package.json"),
-		`${JSON.stringify(
-			{
-				private: true,
-				type: "module",
-				dependencies: {
-					[TRANSFORMERS_PACKAGE]: getTransformersVersionSpec(),
-				},
-				trustedDependencies: ["onnxruntime-node"],
-			},
-			null,
-			"\t",
-		)}\n`,
-	);
-}
-
-async function readPipe(stream: ReadableStream<Uint8Array> | null): Promise<string> {
-	if (!stream) return "";
-	return new Response(stream).text();
-}
-
-async function runRuntimeInstall(runtimeDir: string): Promise<void> {
-	const proc = Bun.spawn([process.execPath, "install", "--cwd", runtimeDir, "--production"], {
-		env: { ...Bun.env, BUN_BE_BUN: "1" },
-		stdout: "pipe",
-		stderr: "pipe",
-	});
-	const [stdout, stderr, exitCode] = await Promise.all([
-		readPipe(proc.stdout as ReadableStream<Uint8Array> | null),
-		readPipe(proc.stderr as ReadableStream<Uint8Array> | null),
-		proc.exited,
-	]);
-	if (exitCode === 0) return;
-	const output = `${stdout}\n${stderr}`.trim();
-	throw new Error(
-		`Failed to install tiny title runtime with ${process.execPath} install (exit ${exitCode}): ${output}`,
-	);
-}
-
 function sendRuntimeInstallProgress(
 	transport: TinyTitleTransport,
 	requestId: string,
@@ -198,28 +134,6 @@ function sendRuntimeInstallProgress(
 	});
 }
 
-async function ensureCompiledTransformersRuntime(
-	transport: TinyTitleTransport,
-	requestId: string,
-	modelKey: TinyLocalModelKey,
-): Promise<string> {
-	const runtimeDir = getTinyTitleRuntimeDir();
-	if (await isCompiledRuntimeInstalled(runtimeDir)) return runtimeDir;
-
-	sendRuntimeInstallProgress(transport, requestId, modelKey, "initiate");
-	const releaseLock = await acquireInstallLock(runtimeDir);
-	try {
-		if (await isCompiledRuntimeInstalled(runtimeDir)) return runtimeDir;
-		await writeRuntimeManifest(runtimeDir);
-		sendRuntimeInstallProgress(transport, requestId, modelKey, "download");
-		await runRuntimeInstall(runtimeDir);
-		sendRuntimeInstallProgress(transport, requestId, modelKey, "done");
-		return runtimeDir;
-	} finally {
-		await releaseLock();
-	}
-}
-
 /**
  * Prepare the freshly-installed compiled runtime for loading: stub `sharp`
  * (the tiny models are text-generation only, so the native image pipeline is
@@ -252,7 +166,15 @@ async function loadTransformers(
 	if (transformersRuntime) return transformersRuntime;
 	transformersRuntime = (async () => {
 		if (!isCompiledBinary()) return configureTransformers(sourceRequire(TRANSFORMERS_PACKAGE) as TransformersRuntime);
-		const runtimeDir = await ensureCompiledTransformersRuntime(transport, requestId, modelKey);
+		const runtimeDir = await ensureRuntimeInstalled({
+			runtimeDir: getTinyTitleRuntimeDir(),
+			install: {
+				dependencies: { [TRANSFORMERS_PACKAGE]: getTransformersVersionSpec() },
+				trustedDependencies: ["onnxruntime-node"],
+			},
+			probePackage: TRANSFORMERS_PACKAGE,
+			onPhase: phase => sendRuntimeInstallProgress(transport, requestId, modelKey, phase),
+		});
 		const entry = await prepareCompiledRuntime(runtimeDir);
 		const require_ = createRequire(entry);
 		return configureTransformers(require_(entry) as TransformersRuntime);
@@ -436,9 +358,10 @@ async function loadPipeline(
 	return loaded;
 }
 
-function buildPrompt(generator: TextGenerationPipeline, message: string): string {
+function buildPrompt(generator: TextGenerationPipeline, message: string, systemPrompt?: string): string {
+	const selectedSystemPrompt = systemPrompt?.trim() || TINY_TITLE_SYSTEM_PROMPT;
 	const chat = [
-		{ role: "system", content: TINY_TITLE_SYSTEM_PROMPT },
+		{ role: "system", content: selectedSystemPrompt },
 		{ role: "user", content: formatTitleUserMessage(message) },
 	];
 	const chatTemplateOptions = {
@@ -464,9 +387,10 @@ async function generateTitle(
 	requestId: string,
 	modelKey: TinyTitleLocalModelKey,
 	message: string,
+	systemPrompt?: string,
 ): Promise<string | null> {
 	const generator = await loadPipeline(modelKey, transport, requestId);
-	const promptText = buildPrompt(generator, message);
+	const promptText = buildPrompt(generator, message, systemPrompt);
 	const transformers = await loadTransformers(transport, requestId, modelKey);
 	const output = (await generator(promptText, {
 		max_new_tokens: TITLE_MAX_NEW_TOKENS,
@@ -548,7 +472,7 @@ async function handleQueuedRequest(
 			transport.send({ type: "completion", id: request.id, text });
 			return;
 		}
-		const title = await generateTitle(transport, request.id, request.modelKey, request.message);
+		const title = await generateTitle(transport, request.id, request.modelKey, request.message, request.systemPrompt);
 		transport.send({ type: "title", id: request.id, title });
 	} catch (error) {
 		transport.send({ type: "error", id: request.id, error: errorText(error) });

package/src/tool-discovery/tool-index.ts

@@ -1,4 +1,5 @@
 import type { AgentTool } from "@prometheus-ai/agent-core";
+import { isZodSchema, zodToWireSchema } from "@prometheus-ai/ai/utils/schema";
 
 // ─── Generic Tool Discovery Types ────────────────────────────────────────────
 
@@ -65,6 +66,7 @@ export function isMCPToolName(name: string): boolean {
 }
 
 function getSchemaPropertyKeys(parameters: unknown): string[] {
+	if (isZodSchema(parameters)) parameters = zodToWireSchema(parameters);
 	if (!parameters || typeof parameters !== "object" || Array.isArray(parameters)) return [];
 	const properties = (parameters as { properties?: unknown }).properties;
 	if (!properties || typeof properties !== "object" || Array.isArray(properties)) return [];

package/src/tools/archive-reader.ts

@@ -1,7 +1,24 @@
+import * as fs from "node:fs/promises";
+import * as os from "node:os";
+import * as path from "node:path";
 import { inflateSync, strFromU8 } from "fflate";
 
+import { formatBytes } from "./render-utils";
 import { ToolError } from "./tool-errors";
 
+/**
+ * Cap on the on-disk size of tar/tar.gz archives, which are loaded fully into
+ * memory (and decompressed by `Bun.Archive`) just to index entries. ZIP is
+ * exempt: it is read via ranged central-directory access.
+ */
+const MAX_TAR_ARCHIVE_BYTES = 256 * 1024 * 1024;
+/**
+ * Cap on a single archive member's declared (uncompressed) size. The declared
+ * size is attacker-controlled metadata — a crafted ZIP entry can claim
+ * multi-GB sizes that would be allocated up front before any data inflates.
+ */
+const MAX_ARCHIVE_MEMBER_BYTES = 64 * 1024 * 1024;
+
 export type ArchiveFormat = "zip" | "tar" | "tar.gz";
 
 export interface ArchivePathCandidate {
@@ -123,11 +140,21 @@ function getArchiveFormatFromPath(filePath: string): ArchiveFormat | undefined {
 	return undefined;
 }
 
+export function formatArchiveEntryLines(entries: readonly ArchiveDirectoryEntry[]): string[] {
+	return entries.map(entry => {
+		if (entry.isDirectory) return `${entry.name}/`;
+
+		const sizeSuffix = entry.size > 0 ? ` (${formatBytes(entry.size)})` : "";
+		return `${entry.name}${sizeSuffix}`;
+	});
+}
+
 const ZIP_LOCAL_FILE_HEADER_SIGNATURE = 0x04034b50;
 const ZIP_CENTRAL_DIRECTORY_HEADER_SIGNATURE = 0x02014b50;
 const ZIP64_EOCD_SIGNATURE = 0x06064b50;
 const ZIP64_EOCD_LOCATOR_SIGNATURE = 0x07064b50;
 const ZIP_EOCD_SIGNATURE = 0x06054b50;
+const ZIP_DATA_DESCRIPTOR_SIGNATURE = 0x08074b50;
 const ZIP_EOCD_MIN_LENGTH = 22;
 const ZIP_EOCD_MAX_COMMENT_LENGTH = 0xffff;
 const ZIP64_EOCD_LOCATOR_LENGTH = 20;
@@ -167,6 +194,37 @@ function readUInt32LE(bytes: Uint8Array, offset: number): number {
 	return (bytes[offset]! | (bytes[offset + 1]! << 8) | (bytes[offset + 2]! << 16) | (bytes[offset + 3]! << 24)) >>> 0;
 }
 
+function bytesMatchAscii(bytes: Uint8Array, offset: number, value: string): boolean {
+	if (bytes.byteLength < offset + value.length) return false;
+	for (let index = 0; index < value.length; index++) {
+		if (bytes[offset + index] !== value.charCodeAt(index)) return false;
+	}
+	return true;
+}
+
+export function sniffArchiveFormat(bytes: Uint8Array): ArchiveFormat | undefined {
+	if (bytes.byteLength >= 4) {
+		const signature = readUInt32LE(bytes, 0);
+		if (
+			signature === ZIP_LOCAL_FILE_HEADER_SIGNATURE ||
+			signature === ZIP_EOCD_SIGNATURE ||
+			signature === ZIP_DATA_DESCRIPTOR_SIGNATURE
+		) {
+			return "zip";
+		}
+	}
+
+	if (bytes.byteLength >= 2 && bytes[0] === 0x1f && bytes[1] === 0x8b) {
+		return "tar.gz";
+	}
+
+	if (bytesMatchAscii(bytes, 257, "ustar")) {
+		return "tar";
+	}
+
+	return undefined;
+}
+
 function readUInt64LEAsNumber(bytes: Uint8Array, offset: number): number {
 	const value = readUInt32LE(bytes, offset) + readUInt32LE(bytes, offset + 4) * ZIP_UINT32_RANGE;
 	if (!Number.isSafeInteger(value)) {
@@ -601,6 +659,11 @@ export class ArchiveReader {
 		if (!entry.storage) {
 			throw new ToolError(`Archive file '${normalizedPath}' has no readable storage`);
 		}
+		if (entry.size > MAX_ARCHIVE_MEMBER_BYTES) {
+			throw new ToolError(
+				`Archive member '${normalizedPath}' is too large to extract in memory (${formatBytes(entry.size)} > ${formatBytes(MAX_ARCHIVE_MEMBER_BYTES)} limit)`,
+			);
+		}
 
 		const bytes =
 			entry.storage.type === "tar"
@@ -623,7 +686,36 @@ export async function openArchive(filePath: string): Promise<ArchiveReader> {
 		throw new ToolError(`Unsupported archive format: ${filePath}`);
 	}
 
-	const entries =
-		format === "zip" ? await readZipEntries(filePath) : await readTarEntries(await Bun.file(filePath).bytes());
+	if (format === "zip") {
+		return new ArchiveReader(format, await readZipEntries(filePath));
+	}
+
+	const file = Bun.file(filePath);
+	const archiveSize = file.size;
+	if (archiveSize > MAX_TAR_ARCHIVE_BYTES) {
+		throw new ToolError(
+			`Archive is too large to read in memory (${formatBytes(archiveSize)} > ${formatBytes(MAX_TAR_ARCHIVE_BYTES)} limit)`,
+		);
+	}
+	const entries = await readTarEntries(await file.bytes());
 	return new ArchiveReader(format, entries);
 }
+
+export async function listArchiveRoot(
+	bytes: Uint8Array,
+	format: ArchiveFormat,
+	opts: { limit?: number } = {},
+): Promise<string> {
+	const tempDir = await fs.mkdtemp(path.join(os.tmpdir(), "prometheus-archive-"));
+	const tempPath = path.join(tempDir, `payload.${format}`);
+	try {
+		await Bun.write(tempPath, bytes);
+		const archive = await openArchive(tempPath);
+		const entries = archive.listDirectory("");
+		const limitedEntries = opts.limit !== undefined && opts.limit > 0 ? entries.slice(0, opts.limit) : entries;
+		const lines = formatArchiveEntryLines(limitedEntries);
+		return lines.length > 0 ? lines.join("\n") : "(empty archive directory)";
+	} finally {
+		await fs.rm(tempDir, { recursive: true, force: true });
+	}
+}