diff --git a/bun.lock b/bun.lock index 5435a8094..3c9e2ee93 100644 --- a/bun.lock +++ b/bun.lock @@ -52,7 +52,6 @@ "dependencies": { "@agentclientprotocol/sdk": "catalog:", "@babel/parser": "catalog:", - "@huggingface/transformers": "catalog:", "@mozilla/readability": "catalog:", "@oh-my-pi/hashline": "catalog:", "@oh-my-pi/omp-stats": "catalog:", @@ -80,6 +79,9 @@ "devDependencies": { "@types/bun": "catalog:", }, + "optionalDependencies": { + "@huggingface/transformers": "catalog:", + }, }, "packages/hashline": { "name": "@oh-my-pi/hashline", diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 3c286423b..be5a1b20c 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,12 +4,13 @@ ### Added -- Added an opt-in Providers → Tiny Model setting for session titles, with five local CPU transformers.js models racing against delayed `pi/smol` fallback while the default remains online-only. +- Added a Providers → Tiny Model setting for session titles, defaulting to local LFM2 700M with five local CPU transformers.js choices, delayed `pi/smol` fallback, in-chat download progress, and an `omp tiny-models download` prefetch command. - Added a persistent live agent roster pinned below the editor (focus it with `Ctrl+S` or `Alt+Down`), including view-as switching into delegated agent sessions with human-readable delegate names and UI pinning to suppress idle reaping while viewed. The roster stays hidden until at least one delegated agent exists and releases focus back to the editor once the last one is gone. - Recorded the originating session ID alongside each prompt in `history.db` (new `session_id` column, surfaced as `HistoryEntry.sessionId`), so recalled prompts can be traced back to the session they came from. Existing history databases gain the column automatically on next launch. - Added compact inline TUI renderers for the `retain`, `recall`, and `reflect` memory tools. `retain` now shows one themed bullet line per stored item (truncated to width) under a status header with the stored/queued count, and `recall`/`reflect` collapse to a single query header (recall reports the match count and hides recalled memories until expanded) instead of dumping the raw JSON argument tree. - Added a randomly picked tip beneath the welcome screen, sourced from an embedded `tips.txt` (one tip per line). The line is italicized with a purple `Tip:` label and a dimmed light-blue body, and the tip is chosen once per welcome instance so intro-animation and LSP re-renders don't shuffle it. - Added a Mnemosyne-only `memory_edit` agent tool for updating, forgetting, or invalidating recalled memories by id, and added `/memory stats` plus `/memory diagnose` slash commands for backend maintenance visibility. +- Added an `orchestrate` magic keyword that mirrors `ultrathink`: dropping the standalone word in a message paints it with a cool teal→violet gradient in the editor and appends a hidden system notice that switches the model into the multi-phase, parallel-subagent orchestration contract. Matching is word-bounded and case-insensitive, so `orchestrated`/`orchestrating` never trigger it. ### Changed @@ -19,9 +20,15 @@ - Changed Mnemosyne `recall` tool output to include memory ids for explicit recall results so agents can target `memory_edit`; auto-injected memory context and `reflect` remain id-free. - Changed the system prompt to advertise `memory://root` only when the local memory backend is active. +### Fixed + +- Fixed a native crash (`malloc: pointer being freed was not allocated` / `NAPI FATAL ERROR`) when quitting after the local transformers.js title model had run. The tiny-title worker no longer calls `pipeline.dispose()` on shutdown — disposing the onnxruntime session freed native memory that Bun's worker/NAPI teardown then freed again. The worker is torn down immediately after, so the OS reclaims the model memory regardless. +- Fixed the tiny-title download progress bar flashing on every first message even when the local model was already downloaded. A cached model emits the same `download`/`progress` events as a real download, so the bar is now revealed only when in-flight progress events keep arriving past a short grace window — cache hits finish (or fall silent during onnxruntime init) before then and never show the bar. + ### Removed - Removed the standalone `ask`, `task`, and `yield` tools along with their obsolete prompts, docs, and tests; delegation now routes through persistent `delegate` agents plus IRC coordination. +- Removed the `/orchestrate` slash command; orchestration is now triggered by the `orchestrate` keyword (see Added) so the contract rides alongside the user's own prompt instead of replacing it. ### Fixed diff --git a/packages/coding-agent/package.json b/packages/coding-agent/package.json index 03deb182d..930d45c3e 100644 --- a/packages/coding-agent/package.json +++ b/packages/coding-agent/package.json @@ -46,7 +46,6 @@ "dependencies": { "@agentclientprotocol/sdk": "catalog:", "@babel/parser": "catalog:", - "@huggingface/transformers": "catalog:", "@mozilla/readability": "catalog:", "@oh-my-pi/hashline": "catalog:", "@oh-my-pi/omp-stats": "catalog:", @@ -71,6 +70,9 @@ "turndown-plugin-gfm": "catalog:", "zod": "catalog:" }, + "optionalDependencies": { + "@huggingface/transformers": "catalog:" + }, "devDependencies": { "@types/bun": "catalog:" }, diff --git a/packages/coding-agent/scripts/build-binary.ts b/packages/coding-agent/scripts/build-binary.ts index b092724b0..893fa5530 100644 --- a/packages/coding-agent/scripts/build-binary.ts +++ b/packages/coding-agent/scripts/build-binary.ts @@ -56,7 +56,7 @@ async function main(): Promise { "../stats/src/sync-worker.ts", "./src/tools/browser/tab-worker-entry.ts", "./src/eval/js/worker-entry.ts", - "./src/title/tiny-title-worker.ts", + "./src/tiny/worker.ts", // Legacy pi-* extension compat entrypoints served by // `legacy-pi-compat.ts`. These are reached via computed bunfs paths // (which `--compile`'s static analyzer cannot trace), so each must be diff --git a/packages/coding-agent/src/cli-commands.ts b/packages/coding-agent/src/cli-commands.ts index a371c968e..c1d42a55b 100644 --- a/packages/coding-agent/src/cli-commands.ts +++ b/packages/coding-agent/src/cli-commands.ts @@ -28,6 +28,7 @@ export const commands: CommandEntry[] = [ { name: "ssh", load: () => import("./commands/ssh").then(m => m.default) }, { name: "stats", load: () => import("./commands/stats").then(m => m.default) }, { name: "update", load: () => import("./commands/update").then(m => m.default) }, + { name: "tiny-models", load: () => import("./commands/tiny-models").then(m => m.default) }, { name: "worktree", load: () => import("./commands/worktree").then(m => m.default), aliases: ["wt"] }, { name: "search", load: () => import("./commands/web-search").then(m => m.default), aliases: ["q"] }, ]; diff --git a/packages/coding-agent/src/cli.ts b/packages/coding-agent/src/cli.ts index 90886eba9..9ec336d90 100755 --- a/packages/coding-agent/src/cli.ts +++ b/packages/coding-agent/src/cli.ts @@ -43,7 +43,7 @@ async function showHelp(config: CliConfig): Promise { */ async function runSmokeTest(): Promise { const { smokeTestSyncWorker } = await import("@oh-my-pi/omp-stats"); - const { smokeTestTinyTitleWorker } = await import("./title/tiny-title-client"); + const { smokeTestTinyTitleWorker } = await import("./tiny/title-client"); await smokeTestSyncWorker(); await smokeTestTinyTitleWorker(); process.stdout.write("smoke-test: ok\n"); diff --git a/packages/coding-agent/src/cli/tiny-models-cli.ts b/packages/coding-agent/src/cli/tiny-models-cli.ts new file mode 100644 index 000000000..1ffd163b2 --- /dev/null +++ b/packages/coding-agent/src/cli/tiny-models-cli.ts @@ -0,0 +1,128 @@ +import { formatBytes } from "@oh-my-pi/pi-utils"; +import chalk from "chalk"; +import { + DEFAULT_TINY_TITLE_MODEL_KEY, + getTinyTitleModelSpec, + isTinyTitleLocalModelKey, + TINY_TITLE_LOCAL_MODELS, + type TinyTitleLocalModelKey, +} from "../tiny/models"; +import { shutdownTinyTitleClient, tinyTitleClient } from "../tiny/title-client"; +import type { TinyTitleProgressEvent } from "../tiny/title-protocol"; + +export type TinyModelsAction = "download" | "list"; + +export interface TinyModelsCommandArgs { + action: TinyModelsAction; + model?: string; + flags: { + json?: boolean; + }; +} + +interface ProgressReporter { + onProgress(event: TinyTitleProgressEvent): void; + finish(ok: boolean): void; +} + +interface DownloadResult { + model: TinyTitleLocalModelKey; + ok: boolean; +} + +function writeLine(text = ""): void { + process.stdout.write(`${text}\n`); +} + +function resolveModels(model: string | undefined): TinyTitleLocalModelKey[] { + if (!model) return [DEFAULT_TINY_TITLE_MODEL_KEY]; + if (model === "all") return TINY_TITLE_LOCAL_MODELS.map(spec => spec.key); + if (!isTinyTitleLocalModelKey(model)) { + const values = TINY_TITLE_LOCAL_MODELS.map(spec => spec.key).join(", "); + throw new Error(`Unknown tiny title model: ${model}. Expected one of: ${values}, all`); + } + return [model]; +} + +function listModels(json: boolean | undefined): void { + if (json) { + writeLine(JSON.stringify({ models: TINY_TITLE_LOCAL_MODELS })); + return; + } + writeLine(chalk.bold("Tiny title models")); + for (const spec of TINY_TITLE_LOCAL_MODELS) { + const defaultMark = spec.key === DEFAULT_TINY_TITLE_MODEL_KEY ? chalk.cyan(" default") : ""; + writeLine(`${chalk.cyan(spec.key)}${defaultMark}`); + writeLine(` ${spec.label} — ${spec.description}`); + } +} + +function makeProgressReporter(modelKey: TinyTitleLocalModelKey, json: boolean | undefined): ProgressReporter { + if (json || !process.stdout.isTTY) { + return { onProgress: () => undefined, finish: () => undefined }; + } + const spec = getTinyTitleModelSpec(modelKey); + let lastWidth = 0; + let lastProgress = -1; + const render = (event: TinyTitleProgressEvent): void => { + const progress = event.progress ?? lastProgress; + if (progress >= 0 && progress < lastProgress + 1 && event.status !== "ready") return; + if (progress >= 0) lastProgress = progress; + const ratio = progress >= 0 ? Math.max(0, Math.min(1, progress / 100)) : 0; + const barWidth = 30; + const filled = Math.round(ratio * barWidth); + const bar = `${"█".repeat(filled)}${"░".repeat(barWidth - filled)}`; + const pct = progress >= 0 ? `${Math.floor(progress).toString().padStart(3, " ")}%` : " --%"; + const bytes = event.loaded && event.total ? ` ${formatBytes(event.loaded)}/${formatBytes(event.total)}` : ""; + const file = event.file ? ` ${event.file.split("/").at(-1) ?? event.file}` : ""; + const label = event.status === "ready" ? "Ready" : "Downloading"; + const line = `${chalk.cyan(label)} ${spec.label} [${bar}] ${pct}${bytes}${file}`; + process.stdout.write(`\r${line.padEnd(lastWidth)}`); + lastWidth = line.length; + }; + return { + onProgress(event) { + if (event.modelKey !== modelKey) return; + render(event); + }, + finish(ok) { + const suffix = ok ? chalk.green("done") : chalk.red("failed"); + process.stdout.write(`\r${`${spec.label}: ${suffix}`.padEnd(lastWidth)}\n`); + }, + }; +} + +async function downloadOne(modelKey: TinyTitleLocalModelKey, json: boolean | undefined): Promise { + const spec = getTinyTitleModelSpec(modelKey); + if (!json && !process.stdout.isTTY) writeLine(`Downloading ${spec.label} (${modelKey})...`); + const progress = makeProgressReporter(modelKey, json); + const ok = await tinyTitleClient.downloadModel(modelKey, { onProgress: progress.onProgress }); + progress.finish(ok); + if (!json && !process.stdout.isTTY) + writeLine(ok ? `Downloaded ${spec.label}.` : `Failed to download ${spec.label}.`); + return { model: modelKey, ok }; +} + +export async function runTinyModelsCommand(command: TinyModelsCommandArgs): Promise { + if (command.action === "list") { + listModels(command.flags.json); + return; + } + + const models = resolveModels(command.model); + const results: DownloadResult[] = []; + try { + for (const model of models) { + results.push(await downloadOne(model, command.flags.json)); + } + } finally { + await shutdownTinyTitleClient(); + } + + if (command.flags.json) { + writeLine(JSON.stringify({ results })); + } + if (results.some(result => !result.ok)) { + throw new Error("One or more tiny title models failed to download"); + } +} diff --git a/packages/coding-agent/src/commands/tiny-models.ts b/packages/coding-agent/src/commands/tiny-models.ts new file mode 100644 index 000000000..b9399f568 --- /dev/null +++ b/packages/coding-agent/src/commands/tiny-models.ts @@ -0,0 +1,36 @@ +import { Args, Command, Flags } from "@oh-my-pi/pi-utils/cli"; +import { runTinyModelsCommand, type TinyModelsAction, type TinyModelsCommandArgs } from "../cli/tiny-models-cli"; + +const ACTIONS: TinyModelsAction[] = ["download", "list"]; + +export default class TinyModels extends Command { + static description = "Download tiny local title models"; + + static args = { + action: Args.string({ + description: "Action to perform", + required: false, + options: ACTIONS, + }), + model: Args.string({ + description: "Model key, or all", + required: false, + }), + }; + + static flags = { + json: Flags.boolean({ description: "Output JSON" }), + }; + + async run(): Promise { + const { args, flags } = await this.parse(TinyModels); + const command: TinyModelsCommandArgs = { + action: (args.action ?? "download") as TinyModelsAction, + model: args.model, + flags: { + json: flags.json, + }, + }; + await runTinyModelsCommand(command); + } +} diff --git a/packages/coding-agent/src/config/settings-schema.ts b/packages/coding-agent/src/config/settings-schema.ts index a5bd28bf4..58b71b78a 100644 --- a/packages/coding-agent/src/config/settings-schema.ts +++ b/packages/coding-agent/src/config/settings-schema.ts @@ -1,7 +1,7 @@ import { THINKING_EFFORTS } from "@oh-my-pi/pi-ai"; import { TASK_SIMPLE_MODES } from "../task/simple-mode"; import { getThinkingLevelMetadata } from "../thinking"; -import { TINY_TITLE_MODEL_OPTIONS, TINY_TITLE_MODEL_VALUES } from "../title/tiny-models"; +import { DEFAULT_TINY_TITLE_MODEL_KEY, TINY_TITLE_MODEL_OPTIONS, TINY_TITLE_MODEL_VALUES } from "../tiny/models"; import { EDIT_MODES } from "../utils/edit-mode"; /** Unified settings schema - single source of truth for all settings. @@ -2900,11 +2900,11 @@ export const SETTINGS_SCHEMA = { "providers.tinyModel": { type: "enum", values: TINY_TITLE_MODEL_VALUES, - default: "online", + default: DEFAULT_TINY_TITLE_MODEL_KEY, ui: { tab: "providers", label: "Tiny Model", - description: "Session-title model: online pi/smol or an opt-in local CPU model", + description: "Session-title model: local LFM2 700M by default, or online pi/smol", options: TINY_TITLE_MODEL_OPTIONS, }, }, diff --git a/packages/coding-agent/src/modes/components/index.ts b/packages/coding-agent/src/modes/components/index.ts index a67fd0576..5057566a2 100644 --- a/packages/coding-agent/src/modes/components/index.ts +++ b/packages/coding-agent/src/modes/components/index.ts @@ -26,6 +26,7 @@ export * from "./show-images-selector"; export * from "./status-line"; export * from "./theme-selector"; export * from "./thinking-selector"; +export * from "./tiny-title-download-progress"; export * from "./todo-reminder"; export * from "./tool-execution"; export * from "./tree-selector"; diff --git a/packages/coding-agent/src/modes/components/tiny-title-download-progress.ts b/packages/coding-agent/src/modes/components/tiny-title-download-progress.ts new file mode 100644 index 000000000..61e899493 --- /dev/null +++ b/packages/coding-agent/src/modes/components/tiny-title-download-progress.ts @@ -0,0 +1,90 @@ +import { type Component, truncateToWidth, visibleWidth } from "@oh-my-pi/pi-tui"; +import { formatBytes } from "@oh-my-pi/pi-utils"; +import { getTinyTitleModelSpec, type TinyTitleLocalModelKey } from "../../tiny/models"; +import type { TinyTitleProgressEvent } from "../../tiny/title-protocol"; +import { theme } from "../theme/theme"; + +const DEFAULT_BAR_WIDTH = 24; + +function padLine(line: string, width: number): string { + const visible = visibleWidth(line); + return visible >= width ? truncateToWidth(line, width) : `${line}${" ".repeat(width - visible)}`; +} + +function progressBar(progress: number | undefined, width: number): string { + const barWidth = Math.max(8, Math.min(DEFAULT_BAR_WIDTH, width)); + if (progress === undefined) return theme.fg("muted", "░".repeat(barWidth)); + const ratio = Math.max(0, Math.min(1, progress / 100)); + const filled = Math.round(ratio * barWidth); + return `${theme.fg("accent", "█".repeat(filled))}${theme.fg("muted", "░".repeat(barWidth - filled))}`; +} + +function currentFile(event: TinyTitleProgressEvent | undefined): string | undefined { + if (!event) return undefined; + if (event.file) return event.file.split("/").at(-1) ?? event.file; + if (event.files) { + let largestFile: string | undefined; + let largestLoaded = -1; + for (const file in event.files) { + const state = event.files[file]; + if (state.loaded <= largestLoaded || state.loaded >= state.total) continue; + largestFile = file; + largestLoaded = state.loaded; + } + return largestFile?.split("/").at(-1) ?? largestFile; + } + return undefined; +} + +function statusLabel(event: TinyTitleProgressEvent | undefined): string { + if (!event) return "Preparing"; + if (event.status === "error") return "Failed"; + if (event.status === "ready") return "Ready"; + if (event.status === "done") return "Downloaded"; + if (event.status === "download") return "Downloading"; + if (event.status === "progress" || event.status === "progress_total") return "Downloading"; + return "Preparing"; +} + +function byteLabel(event: TinyTitleProgressEvent | undefined): string | undefined { + if (!event?.loaded || !event.total) return undefined; + return `${formatBytes(event.loaded)} / ${formatBytes(event.total)}`; +} + +export class TinyTitleDownloadProgressComponent implements Component { + #modelKey: TinyTitleLocalModelKey; + #event: TinyTitleProgressEvent | undefined; + + constructor(modelKey: TinyTitleLocalModelKey) { + this.#modelKey = modelKey; + } + + update(event: TinyTitleProgressEvent): void { + this.#event = event; + } + + isComplete(): boolean { + return this.#event?.status === "ready" || this.#event?.status === "error"; + } + + invalidate(): void { + // No cached state. + } + + render(width: number): string[] { + width = Math.max(1, width); + const spec = getTinyTitleModelSpec(this.#modelKey); + const border = theme.fg("border", theme.boxSharp.horizontal.repeat(width)); + const status = statusLabel(this.#event); + const file = currentFile(this.#event); + const pct = + this.#event?.progress === undefined ? "" : `${Math.floor(this.#event.progress).toString().padStart(3, " ")}%`; + const bytes = byteLabel(this.#event); + const title = `${theme.fg("accent", "Tiny model")} ${theme.fg("muted", status)} ${spec.label}`; + const details = [progressBar(this.#event?.progress, Math.max(8, width - 36)), pct, bytes, file] + .filter((part): part is string => Boolean(part)) + .join(" "); + + return [border, padLine(` ${title}`, width), padLine(` ${details}`, width), border]; + } +} diff --git a/packages/coding-agent/src/modes/controllers/input-controller.ts b/packages/coding-agent/src/modes/controllers/input-controller.ts index 2aa63bd2b..d2eb186a8 100644 --- a/packages/coding-agent/src/modes/controllers/input-controller.ts +++ b/packages/coding-agent/src/modes/controllers/input-controller.ts @@ -3,6 +3,7 @@ import { type AgentMessage, ThinkingLevel } from "@oh-my-pi/pi-agent-core"; import type { AutocompleteProvider, SlashCommand } from "@oh-my-pi/pi-tui"; import { $env, sanitizeText } from "@oh-my-pi/pi-utils"; import { isSettingsInitialized, settings } from "../../config/settings"; +import { TinyTitleDownloadProgressComponent } from "../../modes/components/tiny-title-download-progress"; import { expandEmoticons } from "../../modes/emoji-autocomplete"; import { createPromptActionAutocompleteProvider } from "../../modes/prompt-action-autocomplete"; import { theme } from "../../modes/theme/theme"; @@ -10,6 +11,9 @@ import type { InteractiveModeContext } from "../../modes/types"; import type { AgentSessionEvent } from "../../session/agent-session"; import { SKILL_PROMPT_MESSAGE_TYPE, type SkillPromptDetails } from "../../session/messages"; import { executeBuiltinSlashCommand } from "../../slash-commands/builtin-registry"; +import { isTinyTitleLocalModelKey } from "../../tiny/models"; +import { tinyTitleClient } from "../../tiny/title-client"; +import type { TinyTitleProgressEvent } from "../../tiny/title-protocol"; import { copyToClipboard, readImageFromClipboard } from "../../utils/clipboard"; import { getEditorCommand, openInEditor } from "../../utils/external-editor"; import { ensureSupportedImageInput } from "../../utils/image-loading"; @@ -24,9 +28,61 @@ function isExpandable(obj: unknown): obj is Expandable { return typeof obj === "object" && obj !== null && "setExpanded" in obj && typeof obj.setExpanded === "function"; } +const TINY_TITLE_PROGRESS_DONE_TTL_MS = 3_000; +// A cached model fires its file-load events in a short burst and then goes silent +// while onnxruntime builds the session; a genuine download keeps streaming progress +// events for seconds. Only reveal the bar once a still-incomplete event arrives after +// this grace window, so an already-downloaded model never flashes the bar. +const TINY_TITLE_PROGRESS_REVEAL_DELAY_MS = 1_000; + export class InputController { constructor(private ctx: InteractiveModeContext) {} + #showTinyTitleDownloadProgress(modelKey: string): void { + if (!isTinyTitleLocalModelKey(modelKey) || this.ctx.isBackgrounded) return; + const component = new TinyTitleDownloadProgressComponent(modelKey); + let added = false; + let disposed = false; + let removeTimer: NodeJS.Timeout | undefined; + const remove = (): void => { + if (disposed) return; + disposed = true; + unsubscribe(); + if (removeTimer) { + clearTimeout(removeTimer); + removeTimer = undefined; + } + if (added) { + this.ctx.chatContainer.removeChild(component); + this.ctx.ui.requestRender(); + } + }; + const scheduleRemove = (): void => { + if (removeTimer) clearTimeout(removeTimer); + removeTimer = setTimeout(remove, TINY_TITLE_PROGRESS_DONE_TTL_MS); + removeTimer.unref?.(); + }; + let revealAt = 0; + const update = (event: TinyTitleProgressEvent): void => { + if (disposed || event.modelKey !== modelKey) return; + component.update(event); + if (revealAt === 0) revealAt = performance.now() + TINY_TITLE_PROGRESS_REVEAL_DELAY_MS; + const complete = component.isComplete(); + // Reveal only for a download still in flight past the grace window. Cache hits + // either complete or fall silent (onnx init emits no events) before this fires. + if (!added && !complete && performance.now() >= revealAt) { + this.ctx.chatContainer.addChild(component); + added = true; + } + if (added) this.ctx.ui.requestRender(); + if (complete) { + if (added) scheduleRemove(); + else remove(); + } + }; + const unsubscribe = tinyTitleClient.onProgress(update); + } + setupKeyHandlers(): void { this.ctx.editor.setActionKeys("app.interrupt", this.ctx.keybindings.getKeys("app.interrupt")); this.ctx.editor.shouldBypassAutocompleteOnEscape = () => @@ -329,6 +385,7 @@ export class InputController { // Generate session title on first message const hasUserMessages = this.ctx.session.messages.some((m: AgentMessage) => m.role === "user"); if (!hasUserMessages && !this.ctx.sessionManager.getSessionName() && !$env.PI_NO_TITLE) { + this.#showTinyTitleDownloadProgress(this.ctx.settings.get("providers.tinyModel")); const registry = this.ctx.session.modelRegistry; generateSessionTitle( text, diff --git a/packages/coding-agent/src/title/tiny-models.ts b/packages/coding-agent/src/tiny/models.ts similarity index 98% rename from packages/coding-agent/src/title/tiny-models.ts rename to packages/coding-agent/src/tiny/models.ts index e407ac642..ee679c0b8 100644 --- a/packages/coding-agent/src/title/tiny-models.ts +++ b/packages/coding-agent/src/tiny/models.ts @@ -1,4 +1,5 @@ export const ONLINE_TINY_TITLE_MODEL_KEY = "online"; +export const DEFAULT_TINY_TITLE_MODEL_KEY = "lfm2-700m"; export interface TinyTitleLocalModelSpec { key: string; diff --git a/packages/coding-agent/src/title/title-text.ts b/packages/coding-agent/src/tiny/text.ts similarity index 100% rename from packages/coding-agent/src/title/title-text.ts rename to packages/coding-agent/src/tiny/text.ts diff --git a/packages/coding-agent/src/title/tiny-title-client.ts b/packages/coding-agent/src/tiny/title-client.ts similarity index 63% rename from packages/coding-agent/src/title/tiny-title-client.ts rename to packages/coding-agent/src/tiny/title-client.ts index 93d35f925..53df0e5b3 100644 --- a/packages/coding-agent/src/title/tiny-title-client.ts +++ b/packages/coding-agent/src/tiny/title-client.ts @@ -1,30 +1,34 @@ import { isCompiledBinary, logger } from "@oh-my-pi/pi-utils"; -import { isTinyTitleLocalModelKey } from "./tiny-models"; -import type { TinyTitleWorkerInbound, TinyTitleWorkerOutbound } from "./tiny-title-protocol"; +import { isTinyTitleLocalModelKey, type TinyTitleLocalModelKey } from "./models"; +import type { TinyTitleProgressEvent, TinyTitleWorkerInbound, TinyTitleWorkerOutbound } from "./title-protocol"; interface WorkerHandle { - mode: "worker" | "inline"; send(message: TinyTitleWorkerInbound): void; onMessage(handler: (message: TinyTitleWorkerOutbound) => void): () => void; onError(handler: (error: Error) => void): () => void; terminate(): Promise; } -interface PendingRequest { - resolve(title: string | null): void; +type PendingRequest = + | { kind: "generate"; modelKey: TinyTitleLocalModelKey; resolve: (title: string | null) => void } + | { kind: "download"; modelKey: TinyTitleLocalModelKey; resolve: (ok: boolean) => void }; + +export interface TinyTitleDownloadOptions { + signal?: AbortSignal; + onProgress?: (event: TinyTitleProgressEvent) => void; } const SMOKE_TEST_TIMEOUT_MS = 5_000; export function createTinyTitleWorker(): Worker { return isCompiledBinary() - ? new Worker("./packages/coding-agent/src/title/tiny-title-worker.ts", { type: "module" }) - : new Worker(new URL("./tiny-title-worker.ts", import.meta.url).href, { type: "module" }); + ? new Worker("./packages/coding-agent/src/tiny/worker.ts", { type: "module" }) + : new Worker(new URL("./worker.ts", import.meta.url).href, { type: "module" }); } function wrapBunWorker(worker: Worker): WorkerHandle { + (worker as Worker & { unref?: () => void }).unref?.(); return { - mode: "worker", send(message) { worker.postMessage(message); }, @@ -53,7 +57,6 @@ function spawnInlineUnavailableWorker(error: unknown): WorkerHandle { for (const listener of listeners) listener(message); }; return { - mode: "inline", send(message) { queueMicrotask(() => { if (message.type === "ping") { @@ -102,8 +105,14 @@ export class TinyTitleClient { #unsubscribeMessage: (() => void) | null = null; #unsubscribeError: (() => void) | null = null; #pending = new Map(); + #progressListeners = new Set<(event: TinyTitleProgressEvent) => void>(); #nextRequestId = 0; + onProgress(listener: (event: TinyTitleProgressEvent) => void): () => void { + this.#progressListeners.add(listener); + return () => this.#progressListeners.delete(listener); + } + async generate(modelKey: string, message: string, signal?: AbortSignal): Promise { if (!isTinyTitleLocalModelKey(modelKey)) return null; if (signal?.aborted) return null; @@ -112,11 +121,12 @@ export class TinyTitleClient { const worker = this.#ensureWorker(); const id = String(++this.#nextRequestId); const { promise, resolve } = Promise.withResolvers(); - const pending: PendingRequest = { resolve }; - this.#pending.set(id, pending); + this.#pending.set(id, { kind: "generate", modelKey, resolve }); const abort = (): void => { - if (!this.#pending.delete(id)) return; - resolve(null); + const pending = this.#pending.get(id); + if (pending?.kind !== "generate") return; + this.#pending.delete(id); + pending.resolve(null); }; signal?.addEventListener("abort", abort, { once: true }); try { @@ -135,6 +145,41 @@ export class TinyTitleClient { } } + async downloadModel(modelKey: string, options: TinyTitleDownloadOptions = {}): Promise { + if (!isTinyTitleLocalModelKey(modelKey)) return false; + if (options.signal?.aborted) return false; + + const unsubscribe = options.onProgress ? this.onProgress(options.onProgress) : undefined; + try { + const worker = this.#ensureWorker(); + const id = String(++this.#nextRequestId); + const { promise, resolve } = Promise.withResolvers(); + this.#pending.set(id, { kind: "download", modelKey, resolve }); + const abort = (): void => { + const pending = this.#pending.get(id); + if (pending?.kind !== "download") return; + this.#pending.delete(id); + pending.resolve(false); + }; + options.signal?.addEventListener("abort", abort, { once: true }); + try { + worker.send({ type: "download", id, modelKey }); + return await promise; + } finally { + options.signal?.removeEventListener("abort", abort); + this.#pending.delete(id); + } + } catch (error) { + logger.debug("tiny-title: local model download failed", { + modelKey, + error: error instanceof Error ? error.message : String(error), + }); + return false; + } finally { + unsubscribe?.(); + } + } + async terminate(): Promise { const worker = this.#worker; this.#worker = null; @@ -142,15 +187,17 @@ export class TinyTitleClient { this.#unsubscribeMessage = null; this.#unsubscribeError?.(); this.#unsubscribeError = null; - for (const pending of this.#pending.values()) pending.resolve(null); + for (const pending of this.#pending.values()) { + this.#emitProgress({ modelKey: pending.modelKey, status: "error" }); + if (pending.kind === "generate") pending.resolve(null); + else pending.resolve(false); + } this.#pending.clear(); - if (!worker) return; try { - worker.send({ type: "close" }); + worker?.send({ type: "close" }); } catch { // Worker may already be gone. } - await worker.terminate().catch(() => undefined); } #ensureWorker(): WorkerHandle { @@ -167,26 +214,41 @@ export class TinyTitleClient { logWorkerMessage(message); return; } - if (message.type === "closed") { - void this.terminate(); + if (message.type === "progress") { + this.#emitProgress(message.event); return; } + if (message.type === "closed") return; if (message.type === "pong") return; const pending = this.#pending.get(message.id); if (!pending) return; this.#pending.delete(message.id); if (message.type === "title") { - pending.resolve(message.title); + if (pending.kind === "generate") pending.resolve(message.title); + return; + } + if (message.type === "downloaded") { + if (pending.kind === "download") pending.resolve(true); return; } logger.debug("tiny-title: worker returned error", { error: message.error }); - pending.resolve(null); + this.#emitProgress({ modelKey: pending.modelKey, status: "error" }); + if (pending.kind === "generate") pending.resolve(null); + else pending.resolve(false); + } + + #emitProgress(event: TinyTitleProgressEvent): void { + for (const listener of this.#progressListeners) listener(event); } #handleWorkerError(error: Error): void { logger.warn("tiny-title: worker error", { error: error.message }); - for (const pending of this.#pending.values()) pending.resolve(null); + for (const pending of this.#pending.values()) { + this.#emitProgress({ modelKey: pending.modelKey, status: "error" }); + if (pending.kind === "generate") pending.resolve(null); + else pending.resolve(false); + } this.#pending.clear(); void this.terminate(); } diff --git a/packages/coding-agent/src/tiny/title-protocol.ts b/packages/coding-agent/src/tiny/title-protocol.ts new file mode 100644 index 000000000..705e53738 --- /dev/null +++ b/packages/coding-agent/src/tiny/title-protocol.ts @@ -0,0 +1,49 @@ +import type { TinyTitleLocalModelKey } from "./models"; + +export type TinyTitleProgressStatus = + | "initiate" + | "download" + | "progress" + | "progress_total" + | "done" + | "ready" + | "error"; + +export interface TinyTitleProgressFileState { + loaded: number; + total: number; +} + +export interface TinyTitleProgressEvent { + modelKey: TinyTitleLocalModelKey; + status: TinyTitleProgressStatus; + name?: string; + file?: string; + progress?: number; + loaded?: number; + total?: number; + files?: Record; + task?: string; + model?: string; +} + +export type TinyTitleWorkerInbound = + | { type: "ping"; id: string } + | { type: "generate"; id: string; modelKey: TinyTitleLocalModelKey; message: string } + | { type: "download"; id: string; modelKey: TinyTitleLocalModelKey } + | { type: "close" }; + +export type TinyTitleWorkerOutbound = + | { type: "pong"; id: string } + | { type: "title"; id: string; title: string | null } + | { type: "downloaded"; id: string } + | { type: "error"; id: string; error: string } + | { type: "progress"; id: string; event: TinyTitleProgressEvent } + | { type: "log"; level: "debug" | "warn" | "error"; msg: string; meta?: Record } + | { type: "closed" }; + +export interface TinyTitleTransport { + send(message: TinyTitleWorkerOutbound): void; + onMessage(handler: (message: TinyTitleWorkerInbound) => void): () => void; + close(): void; +} diff --git a/packages/coding-agent/src/tiny/worker.ts b/packages/coding-agent/src/tiny/worker.ts new file mode 100644 index 000000000..01b77da4e --- /dev/null +++ b/packages/coding-agent/src/tiny/worker.ts @@ -0,0 +1,460 @@ +import * as fs from "node:fs/promises"; +import { createRequire } from "node:module"; +import * as path from "node:path"; +import { parentPort } from "node:worker_threads"; +import type { + ProgressInfo, + TextGenerationPipeline, + TextGenerationStringOutput, + StoppingCriteria as TransformersStoppingCriteria, +} from "@huggingface/transformers"; +import { getTinyModelsCacheDir, isCompiledBinary, prompt } from "@oh-my-pi/pi-utils"; +import packageJson from "../../package.json" with { type: "json" }; +import tinyTitleSystemPrompt from "../prompts/system/tiny-title-system.md" with { type: "text" }; +import { getTinyTitleModelSpec, type TinyTitleLocalModelKey } from "./models"; +import { formatTitleUserMessage, normalizeGeneratedTitle } from "./text"; +import type { + TinyTitleProgressEvent, + TinyTitleTransport, + TinyTitleWorkerInbound, + TinyTitleWorkerOutbound, +} from "./title-protocol"; + +const TITLE_PREFILL = ""; +const TITLE_CLOSE = ""; +const TITLE_MAX_NEW_TOKENS = 20; +const STOP_DECODE_WINDOW_TOKENS = 32; +const TINY_TITLE_SYSTEM_PROMPT = prompt.render(tinyTitleSystemPrompt); +const TRANSFORMERS_PACKAGE = "@huggingface/transformers"; +const sourceRequire = createRequire(import.meta.url); +const TRANSFORMERS_VERSION_SPEC = resolveTransformersVersionSpec(); +const TRANSFORMERS_RUNTIME_KEY = TRANSFORMERS_VERSION_SPEC.replace(/[^A-Za-z0-9._-]/g, "_"); +const INSTALL_LOCK_ATTEMPTS = 240; +const INSTALL_LOCK_SLEEP_MS = 250; + +interface TransformersRuntime { + env: { + cacheDir?: string; + allowLocalModels?: boolean; + logLevel?: unknown; + }; + LogLevel: { + ERROR: unknown; + }; + StoppingCriteria: new () => TransformersStoppingCriteria; + pipeline: ( + task: "text-generation", + model: string, + options: { + device: "cpu"; + dtype: "q4"; + progress_callback: (info: ProgressInfo) => void; + }, + ) => Promise; +} + +const pipelines = new Map>(); + +function resolveTransformersVersionSpec(): string { + const manifest = packageJson as { + optionalDependencies?: Record; + dependencies?: Record; + }; + const versionSpec = + manifest.optionalDependencies?.[TRANSFORMERS_PACKAGE] ?? manifest.dependencies?.[TRANSFORMERS_PACKAGE]; + if (!versionSpec) throw new Error(`${TRANSFORMERS_PACKAGE} is missing from package.json optionalDependencies`); + if (!versionSpec.startsWith("catalog:")) return versionSpec; + const installed = sourceRequire(`${TRANSFORMERS_PACKAGE}/package.json`) as { version: string }; + return installed.version; +} +let generateQueue = Promise.resolve(); +let transformersRuntime: Promise | null = null; + +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", + msg: string, + meta?: Record, +): void { + transport.send({ type: "log", level, msg, meta }); +} + +function getTinyTitleRuntimeDir(): string { + return path.join( + path.dirname(getTinyModelsCacheDir()), + "tiny-title-runtime", + `transformers-${TRANSFORMERS_RUNTIME_KEY}`, + ); +} + +async function acquireInstallLock(runtimeDir: string): Promise<() => Promise> { + const lockDir = `${runtimeDir}.lock`; + 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 { + return Bun.file(path.join(runtimeDir, "node_modules", "@huggingface", "transformers", "package.json")).exists(); +} + +async function writeRuntimeManifest(runtimeDir: string): Promise { + await fs.mkdir(runtimeDir, { recursive: true }); + await Bun.write( + path.join(runtimeDir, "package.json"), + `${JSON.stringify( + { + private: true, + type: "module", + dependencies: { + [TRANSFORMERS_PACKAGE]: TRANSFORMERS_VERSION_SPEC, + }, + trustedDependencies: ["onnxruntime-node"], + }, + null, + "\t", + )}\n`, + ); +} + +async function readPipe(stream: ReadableStream | null): Promise { + if (!stream) return ""; + return new Response(stream).text(); +} + +async function runRuntimeInstall(runtimeDir: string): Promise { + 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 | null), + readPipe(proc.stderr as ReadableStream | 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, + modelKey: TinyTitleLocalModelKey, + status: "initiate" | "download" | "done", +): void { + transport.send({ + type: "progress", + id: requestId, + event: { + modelKey, + status, + name: `${TRANSFORMERS_PACKAGE}@${TRANSFORMERS_VERSION_SPEC}`, + }, + }); +} + +async function ensureCompiledTransformersRuntime( + transport: TinyTitleTransport, + requestId: string, + modelKey: TinyTitleLocalModelKey, +): Promise { + 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(); + } +} + +function configureTransformers(transformers: TransformersRuntime): TransformersRuntime { + transformers.env.cacheDir = getTinyModelsCacheDir(); + transformers.env.allowLocalModels = false; + transformers.env.logLevel = transformers.LogLevel.ERROR; + return transformers; +} + +async function loadTransformers( + transport: TinyTitleTransport, + requestId: string, + modelKey: TinyTitleLocalModelKey, +): Promise { + if (transformersRuntime) return transformersRuntime; + transformersRuntime = (async () => { + if (!isCompiledBinary()) return configureTransformers(sourceRequire(TRANSFORMERS_PACKAGE) as TransformersRuntime); + const runtimeDir = await ensureCompiledTransformersRuntime(transport, requestId, modelKey); + const require_ = createRequire(path.join(runtimeDir, "package.json")); + return configureTransformers(require_(TRANSFORMERS_PACKAGE) as TransformersRuntime); + })().catch(error => { + transformersRuntime = null; + throw error; + }); + return transformersRuntime; +} + +function createStopOnTextCriteria( + transformers: TransformersRuntime, + tokenizer: TextGenerationPipeline["tokenizer"], + text: string, +): TransformersStoppingCriteria { + class StopOnTextCriteria extends transformers.StoppingCriteria { + #tokenizer: TextGenerationPipeline["tokenizer"]; + #text: string; + + constructor() { + super(); + this.#tokenizer = tokenizer; + this.#text = text; + } + + _call(inputIds: number[][]): boolean[] { + return inputIds.map(ids => { + const tail = ids.slice(-STOP_DECODE_WINDOW_TOKENS); + const decoded = this.#tokenizer.decode(tail, { + skip_special_tokens: false, + clean_up_tokenization_spaces: false, + }); + return decoded.includes(this.#text); + }); + } + } + return new StopOnTextCriteria(); +} + +function toProgressEvent(modelKey: TinyTitleLocalModelKey, info: ProgressInfo): TinyTitleProgressEvent { + if (info.status === "ready") { + return { modelKey, status: info.status, task: info.task, model: info.model }; + } + if (info.status === "progress_total") { + return { + modelKey, + status: info.status, + name: info.name, + progress: info.progress, + loaded: info.loaded, + total: info.total, + files: info.files, + }; + } + if (info.status === "progress") { + return { + modelKey, + status: info.status, + name: info.name, + file: info.file, + progress: info.progress, + loaded: info.loaded, + total: info.total, + }; + } + return { modelKey, status: info.status, name: info.name, file: info.file }; +} + +function sendProgress( + transport: TinyTitleTransport, + id: string, + modelKey: TinyTitleLocalModelKey, + info: ProgressInfo, +): void { + transport.send({ type: "progress", id, event: toProgressEvent(modelKey, info) }); +} + +async function loadPipeline( + modelKey: TinyTitleLocalModelKey, + transport: TinyTitleTransport, + requestId: string, +): Promise { + const cached = pipelines.get(modelKey); + if (cached) { + const spec = getTinyTitleModelSpec(modelKey); + void cached + .then(() => { + transport.send({ + type: "progress", + id: requestId, + event: { modelKey, status: "ready", task: "text-generation", model: spec.repo }, + }); + }) + .catch(() => undefined); + return cached; + } + + const spec = getTinyTitleModelSpec(modelKey); + const transformers = await loadTransformers(transport, requestId, modelKey); + const startedAt = performance.now(); + const loaded = transformers + .pipeline("text-generation", spec.repo, { + device: "cpu", + dtype: spec.dtype, + progress_callback: info => sendProgress(transport, requestId, modelKey, info), + }) + .then( + generator => { + sendLog(transport, "debug", "tiny-title: local model loaded", { + modelKey, + repo: spec.repo, + elapsedMs: Math.round(performance.now() - startedAt), + }); + transport.send({ + type: "progress", + id: requestId, + event: { modelKey, status: "ready", task: "text-generation", model: spec.repo }, + }); + return generator; + }, + error => { + pipelines.delete(modelKey); + throw error; + }, + ); + pipelines.set(modelKey, loaded); + return loaded; +} + +function buildPrompt(generator: TextGenerationPipeline, message: string): string { + const chat = [ + { role: "system", content: TINY_TITLE_SYSTEM_PROMPT }, + { role: "user", content: formatTitleUserMessage(message) }, + ]; + const chatTemplateOptions = { + add_generation_prompt: true, + tokenize: false, + enable_thinking: false, + }; + return `${generator.tokenizer.apply_chat_template(chat, chatTemplateOptions)}${TITLE_PREFILL}`; +} + +function extractTinyTitle(text: string): string | null { + const titleStart = text.lastIndexOf(TITLE_PREFILL); + const withoutPrefix = titleStart >= 0 ? text.slice(titleStart + TITLE_PREFILL.length) : text; + const closeIndex = withoutPrefix.indexOf(TITLE_CLOSE); + const withoutClose = closeIndex >= 0 ? withoutPrefix.slice(0, closeIndex) : withoutPrefix; + const tagIndex = withoutClose.indexOf("<"); + const withoutTag = tagIndex >= 0 ? withoutClose.slice(0, tagIndex) : withoutClose; + return normalizeGeneratedTitle(withoutTag); +} + +async function generateTitle( + transport: TinyTitleTransport, + requestId: string, + modelKey: TinyTitleLocalModelKey, + message: string, +): Promise { + const generator = await loadPipeline(modelKey, transport, requestId); + const promptText = buildPrompt(generator, message); + const transformers = await loadTransformers(transport, requestId, modelKey); + const output = (await generator(promptText, { + max_new_tokens: TITLE_MAX_NEW_TOKENS, + do_sample: false, + return_full_text: false, + stopping_criteria: createStopOnTextCriteria(transformers, generator.tokenizer, TITLE_CLOSE), + })) as TextGenerationStringOutput; + return extractTinyTitle(output[0]?.generated_text ?? ""); +} + +function releasePipelines(): void { + // Intentionally NOT calling `pipeline.dispose()`. transformers.js disposes the + // underlying onnxruntime InferenceSession, freeing native memory that Bun's + // worker/NAPI teardown then frees a second time — a double-free that aborts the + // process on quit ("malloc: pointer being freed was not allocated" / + // "NAPI FATAL ERROR"). The worker is torn down immediately after `close`, so the + // OS reclaims the model memory regardless; skipping dispose avoids the crash. + pipelines.clear(); +} + +function enqueueRequest( + transport: TinyTitleTransport, + request: Extract, +): void { + generateQueue = generateQueue.then( + async () => { + await handleQueuedRequest(transport, request); + }, + async () => { + await handleQueuedRequest(transport, request); + }, + ); +} + +async function handleQueuedRequest( + transport: TinyTitleTransport, + request: Extract, +): Promise { + try { + if (request.type === "download") { + await loadPipeline(request.modelKey, transport, request.id); + transport.send({ type: "downloaded", id: request.id }); + return; + } + const title = await generateTitle(transport, request.id, request.modelKey, request.message); + transport.send({ type: "title", id: request.id, title }); + } catch (error) { + transport.send({ type: "error", id: request.id, error: errorText(error) }); + } +} + +export function startTinyTitleWorker(transport: TinyTitleTransport): void { + transport.onMessage(message => { + if (message.type === "ping") { + transport.send({ type: "pong", id: message.id }); + return; + } + if (message.type === "close") { + releasePipelines(); + transport.send({ type: "closed" }); + transport.close(); + return; + } + enqueueRequest(transport, message); + }); +} + +if (!parentPort) throw new Error("tiny-title-worker: missing parentPort"); + +const port = parentPort; +const transport: TinyTitleTransport = { + send: (message: TinyTitleWorkerOutbound) => port.postMessage(message), + onMessage: handler => { + const wrap = (data: unknown): void => handler(data as TinyTitleWorkerInbound); + port.on("message", wrap); + return () => port.off("message", wrap); + }, + close: () => { + try { + port.close(); + } catch { + // Already closed. + } + }, +}; + +startTinyTitleWorker(transport); diff --git a/packages/coding-agent/src/title/tiny-title-protocol.ts b/packages/coding-agent/src/title/tiny-title-protocol.ts deleted file mode 100644 index a638ef9df..000000000 --- a/packages/coding-agent/src/title/tiny-title-protocol.ts +++ /dev/null @@ -1,19 +0,0 @@ -import type { TinyTitleLocalModelKey } from "./tiny-models"; - -export type TinyTitleWorkerInbound = - | { type: "ping"; id: string } - | { type: "generate"; id: string; modelKey: TinyTitleLocalModelKey; message: string } - | { type: "close" }; - -export type TinyTitleWorkerOutbound = - | { type: "pong"; id: string } - | { type: "title"; id: string; title: string | null } - | { type: "error"; id: string; error: string } - | { type: "log"; level: "debug" | "warn" | "error"; msg: string; meta?: Record } - | { type: "closed" }; - -export interface TinyTitleTransport { - send(message: TinyTitleWorkerOutbound): void; - onMessage(handler: (message: TinyTitleWorkerInbound) => void): () => void; - close(): void; -} diff --git a/packages/coding-agent/src/title/tiny-title-worker.ts b/packages/coding-agent/src/title/tiny-title-worker.ts deleted file mode 100644 index 613f24f87..000000000 --- a/packages/coding-agent/src/title/tiny-title-worker.ts +++ /dev/null @@ -1,198 +0,0 @@ -import { parentPort } from "node:worker_threads"; -import { - env, - LogLevel, - pipeline, - StoppingCriteria, - type TextGenerationPipeline, - type TextGenerationStringOutput, -} from "@huggingface/transformers"; -import { getTinyModelsCacheDir, prompt } from "@oh-my-pi/pi-utils"; -import tinyTitleSystemPrompt from "../prompts/system/tiny-title-system.md" with { type: "text" }; -import { getTinyTitleModelSpec, type TinyTitleLocalModelKey } from "./tiny-models"; -import type { TinyTitleTransport, TinyTitleWorkerInbound, TinyTitleWorkerOutbound } from "./tiny-title-protocol"; -import { formatTitleUserMessage, normalizeGeneratedTitle } from "./title-text"; - -const TITLE_PREFILL = ""; -const TITLE_CLOSE = ""; -const TITLE_MAX_NEW_TOKENS = 20; -const STOP_DECODE_WINDOW_TOKENS = 32; -const TINY_TITLE_SYSTEM_PROMPT = prompt.render(tinyTitleSystemPrompt); - -env.cacheDir = getTinyModelsCacheDir(); -env.allowLocalModels = false; -env.logLevel = LogLevel.ERROR; - -class StopOnTextCriteria extends StoppingCriteria { - #tokenizer: TextGenerationPipeline["tokenizer"]; - #text: string; - - constructor(tokenizer: TextGenerationPipeline["tokenizer"], text: string) { - super(); - this.#tokenizer = tokenizer; - this.#text = text; - } - - _call(inputIds: number[][]): boolean[] { - return inputIds.map(ids => { - const tail = ids.slice(-STOP_DECODE_WINDOW_TOKENS); - const text = this.#tokenizer.decode(tail, { skip_special_tokens: false, clean_up_tokenization_spaces: false }); - return text.includes(this.#text); - }); - } -} - -const pipelines = new Map>(); -let generateQueue = Promise.resolve(); - -function errorText(error: unknown): string { - return error instanceof Error ? (error.stack ?? error.message) : String(error); -} - -function sendLog( - transport: TinyTitleTransport, - level: "debug" | "warn" | "error", - msg: string, - meta?: Record, -): void { - transport.send({ type: "log", level, msg, meta }); -} - -function loadPipeline( - modelKey: TinyTitleLocalModelKey, - transport: TinyTitleTransport, -): Promise { - const cached = pipelines.get(modelKey); - if (cached) return cached; - - const spec = getTinyTitleModelSpec(modelKey); - const startedAt = performance.now(); - const loaded = pipeline("text-generation", spec.repo, { - device: "cpu", - dtype: spec.dtype, - }).then( - generator => { - sendLog(transport, "debug", "tiny-title: local model loaded", { - modelKey, - repo: spec.repo, - elapsedMs: Math.round(performance.now() - startedAt), - }); - return generator; - }, - error => { - pipelines.delete(modelKey); - throw error; - }, - ); - pipelines.set(modelKey, loaded); - return loaded; -} - -function buildPrompt(generator: TextGenerationPipeline, message: string): string { - const chat = [ - { role: "system", content: TINY_TITLE_SYSTEM_PROMPT }, - { role: "user", content: formatTitleUserMessage(message) }, - ]; - const chatTemplateOptions = { - add_generation_prompt: true, - tokenize: false, - enable_thinking: false, - }; - return `${generator.tokenizer.apply_chat_template(chat, chatTemplateOptions)}${TITLE_PREFILL}`; -} - -function extractTinyTitle(text: string): string | null { - const titleStart = text.lastIndexOf(TITLE_PREFILL); - const withoutPrefix = titleStart >= 0 ? text.slice(titleStart + TITLE_PREFILL.length) : text; - const closeIndex = withoutPrefix.indexOf(TITLE_CLOSE); - const withoutClose = closeIndex >= 0 ? withoutPrefix.slice(0, closeIndex) : withoutPrefix; - const tagIndex = withoutClose.indexOf("<"); - const withoutTag = tagIndex >= 0 ? withoutClose.slice(0, tagIndex) : withoutClose; - return normalizeGeneratedTitle(withoutTag); -} - -async function generateTitle( - transport: TinyTitleTransport, - modelKey: TinyTitleLocalModelKey, - message: string, -): Promise { - const generator = await loadPipeline(modelKey, transport); - const promptText = buildPrompt(generator, message); - const output: TextGenerationStringOutput = await generator(promptText, { - max_new_tokens: TITLE_MAX_NEW_TOKENS, - do_sample: false, - return_full_text: false, - stopping_criteria: new StopOnTextCriteria(generator.tokenizer, TITLE_CLOSE), - }); - return extractTinyTitle(output[0]?.generated_text ?? ""); -} - -async function disposePipelines(): Promise { - const settled = await Promise.allSettled([...pipelines.values()]); - pipelines.clear(); - await Promise.allSettled( - settled.map(result => (result.status === "fulfilled" ? result.value.dispose() : Promise.resolve())), - ); -} - -function handleGenerate( - transport: TinyTitleTransport, - request: Extract, -): void { - generateQueue = generateQueue.then( - async () => { - try { - const title = await generateTitle(transport, request.modelKey, request.message); - transport.send({ type: "title", id: request.id, title }); - } catch (error) { - transport.send({ type: "error", id: request.id, error: errorText(error) }); - } - }, - async () => { - try { - const title = await generateTitle(transport, request.modelKey, request.message); - transport.send({ type: "title", id: request.id, title }); - } catch (error) { - transport.send({ type: "error", id: request.id, error: errorText(error) }); - } - }, - ); -} - -export function startTinyTitleWorker(transport: TinyTitleTransport): void { - transport.onMessage(message => { - if (message.type === "ping") { - transport.send({ type: "pong", id: message.id }); - return; - } - if (message.type === "close") { - void disposePipelines().finally(() => { - transport.send({ type: "closed" }); - transport.close(); - }); - return; - } - handleGenerate(transport, message); - }); -} - -if (!parentPort) throw new Error("tiny-title-worker: missing parentPort"); - -const port = parentPort; -const transport: TinyTitleTransport = { - send: (message: TinyTitleWorkerOutbound) => port.postMessage(message), - onMessage: handler => { - const wrap = (data: unknown): void => handler(data as TinyTitleWorkerInbound); - port.on("message", wrap); - return () => port.off("message", wrap); - }, - close: () => { - try { - port.close(); - } catch { - // Already closed. - } - }, -}; - -startTinyTitleWorker(transport); diff --git a/packages/coding-agent/src/utils/title-generator.ts b/packages/coding-agent/src/utils/title-generator.ts index 1ee24d7cd..118abdb9f 100644 --- a/packages/coding-agent/src/utils/title-generator.ts +++ b/packages/coding-agent/src/utils/title-generator.ts @@ -9,9 +9,9 @@ import type { ModelRegistry } from "../config/model-registry"; import { resolveRoleSelection } from "../config/model-resolver"; import type { Settings } from "../config/settings"; import titleSystemPrompt from "../prompts/system/title-system.md" with { type: "text" }; -import { ONLINE_TINY_TITLE_MODEL_KEY } from "../title/tiny-models"; -import { tinyTitleClient } from "../title/tiny-title-client"; -import { formatTitleUserMessage, normalizeGeneratedTitle } from "../title/title-text"; +import { ONLINE_TINY_TITLE_MODEL_KEY } from "../tiny/models"; +import { formatTitleUserMessage, normalizeGeneratedTitle } from "../tiny/text"; +import { tinyTitleClient } from "../tiny/title-client"; const TITLE_SYSTEM_PROMPT = prompt.render(titleSystemPrompt); diff --git a/packages/coding-agent/test/issue-1150-repro.test.ts b/packages/coding-agent/test/issue-1150-repro.test.ts index 138d09616..ab00ecd70 100644 --- a/packages/coding-agent/test/issue-1150-repro.test.ts +++ b/packages/coding-agent/test/issue-1150-repro.test.ts @@ -35,7 +35,7 @@ describe("issue #1150 — release-build script must list all worker --compile en "./packages/stats/src/sync-worker.ts", "./packages/coding-agent/src/tools/browser/tab-worker-entry.ts", "./packages/coding-agent/src/eval/js/worker-entry.ts", - "./packages/coding-agent/src/title/tiny-title-worker.ts", + "./packages/coding-agent/src/tiny/worker.ts", ]; it("scripts/ci-release-build-binaries.ts lists every worker as an explicit --compile entrypoint", async () => { @@ -56,7 +56,7 @@ describe("issue #1150 — release-build script must list all worker --compile en "../stats/src/sync-worker.ts", "./src/tools/browser/tab-worker-entry.ts", "./src/eval/js/worker-entry.ts", - "./src/title/tiny-title-worker.ts", + "./src/tiny/worker.ts", ]; const source = await Bun.file(devScriptPath).text(); for (const entry of devEntrypoints) { diff --git a/packages/coding-agent/test/modes/orchestrate.test.ts b/packages/coding-agent/test/modes/orchestrate.test.ts new file mode 100644 index 000000000..7f6da09ff --- /dev/null +++ b/packages/coding-agent/test/modes/orchestrate.test.ts @@ -0,0 +1,70 @@ +import { beforeAll, describe, expect, it } from "bun:test"; +import { containsOrchestrate, highlightOrchestrate, ORCHESTRATE_NOTICE } from "../../src/modes/orchestrate"; +import { initTheme } from "../../src/modes/theme/theme"; +import { containsUltrathink, highlightUltrathink } from "../../src/modes/ultrathink"; +import { clearBundledCommandsCache, loadBundledCommands } from "../../src/task/commands"; + +beforeAll(() => { + // highlightOrchestrate/highlightUltrathink read the global theme's color mode. + initTheme(); +}); + +describe("orchestrate keyword detection", () => { + it("matches the standalone word in any case", () => { + expect(containsOrchestrate("orchestrate")).toBe(true); + expect(containsOrchestrate("Orchestrate")).toBe(true); + expect(containsOrchestrate("ORCHESTRATE")).toBe(true); + expect(containsOrchestrate("please orchestrate this rollout")).toBe(true); + expect(containsOrchestrate("do it. orchestrate.")).toBe(true); + }); + + it("ignores inflected forms and embedded substrings", () => { + expect(containsOrchestrate("orchestrated the build")).toBe(false); + expect(containsOrchestrate("orchestrating now")).toBe(false); + expect(containsOrchestrate("a clean orchestration")).toBe(false); + expect(containsOrchestrate("it orchestrates well")).toBe(false); + expect(containsOrchestrate("reorchestrate everything")).toBe(false); + expect(containsOrchestrate("nothing to see here")).toBe(false); + }); +}); + +describe("orchestrate keyword highlighting", () => { + it("decorates the keyword with zero-width escapes, preserving visible text", () => { + const decorated = highlightOrchestrate("please orchestrate this"); + expect(decorated).not.toBe("please orchestrate this"); + expect(decorated).toContain("\x1b"); + expect(Bun.stripANSI(decorated)).toBe("please orchestrate this"); + }); + + it("leaves text without the standalone keyword untouched", () => { + expect(highlightOrchestrate("nothing here")).toBe("nothing here"); + // Probe hits the substring but the word boundary fails — no decoration. + expect(highlightOrchestrate("orchestrated builds")).toBe("orchestrated builds"); + }); + + it("does not cross-trigger with the ultrathink highlighter", () => { + expect(highlightOrchestrate("ultrathink")).toBe("ultrathink"); + expect(highlightUltrathink("orchestrate")).toBe("orchestrate"); + expect(containsUltrathink("orchestrate")).toBe(false); + expect(containsOrchestrate("ultrathink")).toBe(false); + }); +}); + +describe("orchestrate notice", () => { + it("is a self-contained system notice carrying the orchestration contract", () => { + expect(ORCHESTRATE_NOTICE.startsWith("")).toBe(true); + expect(ORCHESTRATE_NOTICE.endsWith("")).toBe(true); + expect(ORCHESTRATE_NOTICE).toContain("orchestrator"); + // The contract must not retain the slash-command input placeholder. + expect(ORCHESTRATE_NOTICE).not.toContain("$@"); + }); +}); + +describe("orchestrate slash command removal", () => { + it("is no longer bundled as a slash command", () => { + clearBundledCommandsCache(); + const names = loadBundledCommands().map(command => command.name); + expect(names).not.toContain("orchestrate"); + expect(names).toContain("init"); + }); +}); diff --git a/packages/coding-agent/test/tiny-title-generator.test.ts b/packages/coding-agent/test/tiny-title-generator.test.ts index 0fa8c939c..c4f976f72 100644 --- a/packages/coding-agent/test/tiny-title-generator.test.ts +++ b/packages/coding-agent/test/tiny-title-generator.test.ts @@ -1,9 +1,12 @@ -import { afterEach, describe, expect, it, vi } from "bun:test"; +import { afterEach, beforeAll, describe, expect, it, vi } from "bun:test"; import * as ai from "@oh-my-pi/pi-ai"; import { type Api, type AssistantMessage, getBundledModel, type Model } from "@oh-my-pi/pi-ai"; -import { getEnumValues, getUi } from "../src/config/settings-schema"; -import { TINY_TITLE_MODEL_OPTIONS, TINY_TITLE_MODEL_VALUES } from "../src/title/tiny-models"; -import { tinyTitleClient } from "../src/title/tiny-title-client"; +import { isSubcommand } from "../src/cli-commands"; +import { getDefault, getEnumValues, getUi } from "../src/config/settings-schema"; +import { TinyTitleDownloadProgressComponent } from "../src/modes/components/tiny-title-download-progress"; +import { initTheme } from "../src/modes/theme/theme"; +import { DEFAULT_TINY_TITLE_MODEL_KEY, TINY_TITLE_MODEL_OPTIONS, TINY_TITLE_MODEL_VALUES } from "../src/tiny/models"; +import { tinyTitleClient } from "../src/tiny/title-client"; import { generateSessionTitle, raceFirstNonNull, TITLE_LOCAL_FALLBACK_DELAY_MS } from "../src/utils/title-generator"; async function flushMicrotasks(turns = 4): Promise { @@ -54,6 +57,10 @@ function mockOnlineTitle(title: string | null) { } as never); } +beforeAll(() => { + initTheme(); +}); + afterEach(() => { vi.useRealTimers(); vi.restoreAllMocks(); @@ -230,11 +237,64 @@ describe("tiny title generator routing", () => { expect(onlineSignal?.aborted).toBe(true); onlineHold.resolve({ stopReason: "abort", content: [] } as never); }); + + it("keeps local generation alive when the delayed online fallback wins", async () => { + vi.useFakeTimers(); + const model = getModelOrThrow("claude-sonnet-4-5"); + const local = Promise.withResolvers(); + let localSettled = false; + void local.promise.then(() => { + localSettled = true; + }); + vi.spyOn(tinyTitleClient, "generate").mockReturnValue(local.promise); + mockOnlineTitle("Online Title"); + + const result = generateSessionTitle( + "Investigate background download", + createRegistry(model), + createSettings(model, "lfm2-700m"), + ); + + vi.advanceTimersByTime(TITLE_LOCAL_FALLBACK_DELAY_MS); + await flushMicrotasks(); + await expect(result).resolves.toBe("Online Title"); + expect(localSettled).toBe(false); + + local.resolve("Late Local Title"); + await flushMicrotasks(); + expect(localSettled).toBe(true); + }); }); describe("providers.tinyModel schema", () => { it("keeps enum values and UI options in sync with the tiny model registry", () => { expect(getEnumValues("providers.tinyModel")).toEqual([...TINY_TITLE_MODEL_VALUES]); expect(getUi("providers.tinyModel")?.options).toEqual(TINY_TITLE_MODEL_OPTIONS); + expect(getDefault("providers.tinyModel")).toBe(DEFAULT_TINY_TITLE_MODEL_KEY); + }); +}); + +describe("tiny title download progress UI", () => { + it("renders progress updates and completion state", () => { + const component = new TinyTitleDownloadProgressComponent("lfm2-700m"); + component.update({ + modelKey: "lfm2-700m", + status: "progress_total", + name: "onnx-community/LFM2-700M-ONNX", + progress: 50, + loaded: 50, + total: 100, + files: {}, + }); + expect(component.render(80).join("\n")).toContain("LFM2 700M"); + expect(component.isComplete()).toBe(false); + component.update({ modelKey: "lfm2-700m", status: "ready", task: "text-generation", model: "repo" }); + expect(component.isComplete()).toBe(true); + }); +}); + +describe("tiny-models CLI", () => { + it("registers tiny-models as a top-level subcommand", () => { + expect(isSubcommand("tiny-models")).toBe(true); }); }); diff --git a/packages/mnemosyne/package.json b/packages/mnemosyne/package.json index 97b139dc0..c58e218bb 100644 --- a/packages/mnemosyne/package.json +++ b/packages/mnemosyne/package.json @@ -6,8 +6,7 @@ "homepage": "https://omp.sh", "author": "Can Boluk", "contributors": [ - "Abdias J", - "Mario Zechner" + "Abdias J" ], "license": "MIT", "repository": { diff --git a/scripts/ci-release-build-binaries.ts b/scripts/ci-release-build-binaries.ts index f5fe71c00..fad13b13a 100644 --- a/scripts/ci-release-build-binaries.ts +++ b/scripts/ci-release-build-binaries.ts @@ -27,7 +27,7 @@ const workerEntrypoints = [ "./packages/stats/src/sync-worker.ts", "./packages/coding-agent/src/tools/browser/tab-worker-entry.ts", "./packages/coding-agent/src/eval/js/worker-entry.ts", - "./packages/coding-agent/src/title/tiny-title-worker.ts", + "./packages/coding-agent/src/tiny/worker.ts", ]; const isDryRun = process.argv.includes("--dry-run"); const targets: BinaryTarget[] = [