diff --git a/packages/ai/.changes/prime-inference-model-catalog.md b/packages/ai/.changes/prime-inference-model-catalog.md new file mode 100644 index 0000000000..3470623079 --- /dev/null +++ b/packages/ai/.changes/prime-inference-model-catalog.md @@ -0,0 +1 @@ +- Added live Prime Inference model names, pricing, limits, modalities, and reasoning support to the bundled catalog. diff --git a/packages/ai/scripts/generate-models.ts b/packages/ai/scripts/generate-models.ts index 04c09654c7..5a942ab3a7 100644 --- a/packages/ai/scripts/generate-models.ts +++ b/packages/ai/scripts/generate-models.ts @@ -1,12 +1,16 @@ #!/usr/bin/env tsx -import { readFileSync, writeFileSync } from "fs"; -import { homedir } from "os"; +import { writeFileSync } from "fs"; import { dirname, join } from "path"; import { fileURLToPath } from "url"; import { getAnthropicCacheCosts } from "../src/cache-pricing.js"; import { COPILOT_CLIENT_HEADERS } from "../src/copilot-client-version.js"; import { getOpenRouterReasoningCapabilities } from "../src/openrouter-reasoning.js"; +import { + isPrivatePrimeInferenceModelId, + parsePrimeInferenceModelCatalog, + type PrimeInferenceCatalogEntry, +} from "../src/prime-inference-model-catalog.js"; import { CLOUDFLARE_AI_GATEWAY_ANTHROPIC_BASE_URL, CLOUDFLARE_AI_GATEWAY_COMPAT_BASE_URL, @@ -21,6 +25,7 @@ import { type OpenAICompletionsCompat, } from "../src/types.js"; import { MODELS as EXISTING_MODELS } from "../src/models.generated.js"; +import { renderModelsFile } from "./render-models.js"; const __filename = fileURLToPath(import.meta.url); const __dirname = dirname(__filename); @@ -116,15 +121,6 @@ const PRIME_INFERENCE_COMPAT: OpenAICompletionsCompat = { maxTokensField: "max_tokens", supportsStrictMode: false, }; -interface PrimeInferenceCatalogEntry { - id: string; - input: number; - output: number; - contextWindow?: number; - maxTokens?: number; - reasoning?: boolean; -} - interface PrimeInferenceModelMetadata { contextWindow?: number; maxTokens?: number; @@ -132,13 +128,9 @@ interface PrimeInferenceModelMetadata { name?: string; } -// The full Prime Inference catalog is registered (minus raw/duplicate variants). -// Prime's /models endpoint publishes pricing only, so context/output limits and -// modalities are read from OpenRouter's public catalog, used here purely as a -// published spec sheet for the same upstream models — requests always go to -// Prime's own baseUrl. Entries below override those specs where the Prime route -// enforces a different limit (verified against the live API) or fill gaps for -// models OpenRouter does not list or leaves incomplete. +// Prime's /models endpoint is authoritative for route metadata. OpenRouter and +// these overrides only fill gaps for older or incomplete endpoint entries; +// requests always go to Prime's own baseUrl. const PRIME_INFERENCE_MODEL_METADATA: Record = { // These routes accept 200k, checked against the live API 2026-07-08. The // other Claude routes take the full window their spec lists. @@ -218,22 +210,6 @@ const PRIME_INFERENCE_OPENROUTER_ALIASES: Record = { const PRIME_INFERENCE_DEFAULT_CONTEXT_WINDOW = 128000; const PRIME_INFERENCE_DEFAULT_MAX_TOKENS = 8192; -// Raw checkpoints and duplicate routes that would clutter the picker: BF16 -// exports, fine-tune outputs, zai-org/ and HF-cased twins of canonical ids. -function isPrimeInferenceRawVariant(modelId: string): boolean { - const id = modelId.toLowerCase(); - if (id.endsWith("-bf16") || id.includes(":")) { - return true; - } - const vendor = modelId.split("/")[0] ?? ""; - return vendor === "zai-org" || vendor !== vendor.toLowerCase(); -} - -function isPrimeInferencePrivateModel(modelId: string): boolean { - const id = modelId.toLowerCase(); - return id.startsWith("internal/") || id.startsWith("dev/"); -} - const OPENAI_RESPONSES_NONE_REASONING_MODELS = new Set([ "gpt-5.1", "gpt-5.2", @@ -385,51 +361,6 @@ function getOptionalNumber(value: unknown): number | undefined { return typeof value === "number" && Number.isFinite(value) ? value : undefined; } -function getOptionalBoolean(value: unknown): boolean | undefined { - return typeof value === "boolean" ? value : undefined; -} - -function readPrimeCliConfig(): Record { - try { - const parsed = JSON.parse(readFileSync(join(homedir(), ".prime", "config.json"), "utf8")); - return isRecord(parsed) ? parsed : {}; - } catch { - return {}; - } -} - -function getPrimeInferenceConfigValue( - envName: "PRIME_API_KEY" | "PRIME_TEAM_ID", - config: Record, - configKeys: readonly string[], -): string | undefined { - const fromEnv = process.env[envName]?.trim(); - if (fromEnv) { - return fromEnv; - } - - for (const key of configKeys) { - const value = config[key]; - if (typeof value === "string" && value.trim()) { - return value.trim(); - } - } - - return undefined; -} - -function getPrimeInferenceHeaders(apiKey: string | undefined, teamId: string | undefined): Record | undefined { - const headers: Record = {}; - if (apiKey) { - headers.Authorization = `Bearer ${apiKey}`; - } - if (teamId) { - headers["X-Prime-Team-ID"] = teamId; - } - - return Object.keys(headers).length > 0 ? headers : undefined; -} - function getPrimeInferenceCacheCosts(modelId: string, inputCost: number): { cacheRead: number; cacheWrite: number } { return modelId.toLowerCase().startsWith("anthropic/") ? getAnthropicCacheCosts(inputCost, "5m") @@ -439,7 +370,7 @@ function getPrimeInferenceCacheCosts(modelId: string, inputCost: number): { cach function getExistingPrimeInferenceModels(): Model<"openai-completions">[] { const models = EXISTING_MODELS["prime-inference"] as unknown as Record>; return Object.values(models) - .filter((model) => !isPrimeInferenceRawVariant(model.id) && !isPrimeInferencePrivateModel(model.id)) + .filter((model) => !isPrivatePrimeInferenceModelId(model.id)) .map((model) => ({ ...model, input: [...model.input], @@ -486,20 +417,6 @@ function refreshPrimeInferenceAliasLimits( }); } -function includesCatalogCapability(value: unknown, capabilities: readonly string[]): boolean { - if (!Array.isArray(value)) { - return false; - } - - return value.some((item) => { - if (typeof item !== "string") { - return false; - } - const normalized = item.toLowerCase(); - return capabilities.some((capability) => normalized.includes(capability)); - }); -} - function getPrimeInferenceDisplayName(modelId: string): string { const rawName = modelId.split("/").at(-1) ?? modelId; return rawName @@ -513,29 +430,6 @@ function getPrimeInferenceDisplayName(modelId: string): string { .join(" "); } -function getPrimeInferenceCatalogReasoning(item: Record): boolean | undefined { - const metadata = isRecord(item.metadata) ? item.metadata : {}; - const direct = - getOptionalBoolean(item.reasoning) ?? - getOptionalBoolean(item.supports_reasoning) ?? - getOptionalBoolean(item.supportsReasoning) ?? - getOptionalBoolean(metadata.reasoning) ?? - getOptionalBoolean(metadata.supports_reasoning) ?? - getOptionalBoolean(metadata.supportsReasoning); - if (direct !== undefined) { - return direct; - } - - return includesCatalogCapability(item.supported_parameters, ["reasoning", "thinking"]) || - includesCatalogCapability(item.capabilities, ["reasoning", "thinking"]) || - includesCatalogCapability(item.tags, ["reasoning", "thinking"]) || - includesCatalogCapability(metadata.supported_parameters, ["reasoning", "thinking"]) || - includesCatalogCapability(metadata.capabilities, ["reasoning", "thinking"]) || - includesCatalogCapability(metadata.tags, ["reasoning", "thinking"]) - ? true - : undefined; -} - function isPrimeInferenceReasoningModel(modelId: string, catalogReasoning?: boolean): boolean { if (catalogReasoning !== undefined) { return catalogReasoning; @@ -572,37 +466,6 @@ function getPrimeInferenceCompat(modelId: string): OpenAICompletionsCompat { return PRIME_INFERENCE_COMPAT; } -function parsePrimeInferenceCatalog(data: unknown): PrimeInferenceCatalogEntry[] { - if (!isRecord(data) || !Array.isArray(data.data)) { - return []; - } - - return data.data.flatMap((item): PrimeInferenceCatalogEntry[] => { - if (!isRecord(item) || typeof item.id !== "string") { - return []; - } - - const pricing = isRecord(item.pricing) ? item.pricing : {}; - const input = getOptionalNumber(pricing.input_usd_per_mtok); - const output = getOptionalNumber(pricing.output_usd_per_mtok); - if (input === undefined || output === undefined) { - return []; - } - - const limit = isRecord(item.limit) ? item.limit : {}; - return [ - { - id: item.id, - input, - output, - contextWindow: getOptionalNumber(item.context_window ?? item.contextWindow ?? limit.context), - maxTokens: getOptionalNumber(item.max_tokens ?? item.maxTokens ?? limit.output), - reasoning: getPrimeInferenceCatalogReasoning(item), - }, - ]; - }); -} - interface PrimeInferenceOpenRouterMetadata { contextWindow?: number; maxTokens?: number; @@ -651,17 +514,12 @@ function getPrimeInferenceOpenRouterMetadata( } async function fetchPrimeInferenceModels(): Promise[]> { - const primeConfig = readPrimeCliConfig(); - const apiKey = getPrimeInferenceConfigValue("PRIME_API_KEY", primeConfig, ["api_key", "apiKey"]); - const teamId = getPrimeInferenceConfigValue("PRIME_TEAM_ID", primeConfig, ["team_id", "teamId", "teamID"]); let catalog: PrimeInferenceCatalogEntry[] = []; try { - console.log("Fetching models from Prime Inference API..."); - const response = await fetch(`${PRIME_INFERENCE_BASE_URL}/models`, { - headers: getPrimeInferenceHeaders(apiKey, teamId), - }); - catalog = parsePrimeInferenceCatalog(await response.json()); + console.log("Fetching public models from Prime Inference API..."); + const response = await fetch(`${PRIME_INFERENCE_BASE_URL}/models`); + catalog = parsePrimeInferenceModelCatalog(await response.json()); } catch (error) { console.error("Failed to fetch Prime Inference models:", error); } @@ -680,7 +538,7 @@ async function fetchPrimeInferenceModels(): Promise[ } const catalogModels = catalog - .filter((entry) => !isPrimeInferenceRawVariant(entry.id) && !isPrimeInferencePrivateModel(entry.id)) + .filter((entry) => !isPrivatePrimeInferenceModelId(entry.id)) .map((entry) => createPrimeInferenceModel( entry, @@ -689,6 +547,10 @@ async function fetchPrimeInferenceModels(): Promise[ ), ); let snapshotModels = getExistingPrimeInferenceModels(); + if (catalog.length > 0 && catalogModels.length < Math.ceil(snapshotModels.length * 0.5)) { + console.error("Prime Inference catalog is severely truncated; keeping snapshot models"); + return snapshotModels; + } if (catalog.length > 0) { const liveIds = new Set(catalogModels.map((model) => model.id.toLowerCase())); snapshotModels = snapshotModels.filter((model) => liveIds.has(model.id.toLowerCase())); @@ -704,8 +566,12 @@ function createPrimeInferenceModel( override: PrimeInferenceModelMetadata | undefined, openRouter: PrimeInferenceOpenRouterMetadata | undefined, ): Model<"openai-completions"> { - const vision = override?.vision ?? openRouter?.vision ?? false; - const cacheCosts = getPrimeInferenceCacheCosts(entry.id, entry.input); + const vision = entry.vision ?? override?.vision ?? openRouter?.vision ?? false; + const fallbackCacheCosts = getPrimeInferenceCacheCosts(entry.id, entry.input); + const cacheCosts = { + cacheRead: entry.cacheRead ?? fallbackCacheCosts.cacheRead, + cacheWrite: entry.cacheWrite ?? fallbackCacheCosts.cacheWrite, + }; const contextWindow = entry.contextWindow ?? override?.contextWindow ?? @@ -721,7 +587,7 @@ function createPrimeInferenceModel( return { id: entry.id, ...(PRIME_INFERENCE_FEATURED_MODELS.has(entry.id.toLowerCase()) ? { featured: true } : {}), - name: override?.name ?? getPrimeInferenceDisplayName(entry.id), + name: entry.name ?? override?.name ?? getPrimeInferenceDisplayName(entry.id), api: "openai-completions", provider: "prime-inference", baseUrl: PRIME_INFERENCE_BASE_URL, @@ -2394,7 +2260,7 @@ async function generateModels() { } // Group by provider and deduplicate by model ID - const providers: Record>> = {}; + const providers: Record>> = {}; for (const model of allModels) { if (!providers[model.provider]) { providers[model.provider] = {}; @@ -2406,63 +2272,9 @@ async function generateModels() { } } - // Generate TypeScript file - let output = `// This file is auto-generated by scripts/generate-models.ts -// Do not edit manually - run 'npm run generate-models' to update - -import type { Model } from "./types.js"; - -export const MODELS = { -`; - - // Generate provider sections (sorted for deterministic output) - const sortedProviderIds = Object.keys(providers).sort(); - for (const providerId of sortedProviderIds) { - const models = providers[providerId]; - output += `\t${JSON.stringify(providerId)}: {\n`; - - const sortedModelIds = Object.keys(models).sort(); - for (const modelId of sortedModelIds) { - const model = models[modelId]; - output += `\t\t"${model.id}": {\n`; - output += `\t\t\tid: "${model.id}",\n`; - output += `\t\t\tname: "${model.name}",\n`; - output += `\t\t\tapi: "${model.api}",\n`; - output += `\t\t\tprovider: "${model.provider}",\n`; - if (model.baseUrl !== undefined) { - output += `\t\t\tbaseUrl: "${model.baseUrl}",\n`; - } - if (model.headers) { - output += `\t\t\theaders: ${JSON.stringify(model.headers)},\n`; - } - if (model.compat) { - output += ` compat: ${JSON.stringify(model.compat)}, -`; - } - output += `\t\t\treasoning: ${model.reasoning},\n`; - if (model.thinkingLevelMap) { - output += `\t\t\tthinkingLevelMap: ${JSON.stringify(model.thinkingLevelMap)},\n`; - } - output += `\t\t\tinput: [${model.input.map(i => `"${i}"`).join(", ")}],\n`; - output += `\t\t\tcost: {\n`; - output += `\t\t\t\tinput: ${model.cost.input},\n`; - output += `\t\t\t\toutput: ${model.cost.output},\n`; - output += `\t\t\t\tcacheRead: ${model.cost.cacheRead},\n`; - output += `\t\t\t\tcacheWrite: ${model.cost.cacheWrite},\n`; - output += `\t\t\t},\n`; - output += `\t\t\tcontextWindow: ${model.contextWindow},\n`; - output += `\t\t\tmaxTokens: ${model.maxTokens},\n`; - if (model.featured) { - output += `\t\t\tfeatured: true,\n`; - } - output += `\t\t} satisfies Model<"${model.api}">,\n`; - } - - output += `\t},\n`; - } - - output += `} as const; -`; + // Generate TypeScript file. JSON string literals prevent remote catalog + // text from becoming executable source code. + const output = renderModelsFile(providers); // Write file writeFileSync(join(packageRoot, "src/models.generated.ts"), output); diff --git a/packages/ai/scripts/render-models.ts b/packages/ai/scripts/render-models.ts new file mode 100644 index 0000000000..d906425ee6 --- /dev/null +++ b/packages/ai/scripts/render-models.ts @@ -0,0 +1,49 @@ +import type { Api, Model } from "../src/types.js"; + +export function renderModelsFile(providers: Record>>): string { + let output = `// This file is auto-generated by scripts/generate-models.ts +// Do not edit manually - run 'npm run generate-models' to update + +import type { Model } from "./types.js"; + +export const MODELS = { +`; + + for (const providerId of Object.keys(providers).sort()) { + const models = providers[providerId]; + output += `\t${JSON.stringify(providerId)}: {\n`; + + for (const modelId of Object.keys(models).sort()) { + const model = models[modelId]; + output += `\t\t${JSON.stringify(model.id)}: {\n`; + output += `\t\t\tid: ${JSON.stringify(model.id)},\n`; + output += `\t\t\tname: ${JSON.stringify(model.name)},\n`; + output += `\t\t\tapi: ${JSON.stringify(model.api)},\n`; + output += `\t\t\tprovider: ${JSON.stringify(model.provider)},\n`; + if (model.baseUrl !== undefined) { + output += `\t\t\tbaseUrl: ${JSON.stringify(model.baseUrl)},\n`; + } + if (model.headers) output += `\t\t\theaders: ${JSON.stringify(model.headers)},\n`; + if (model.compat) output += `\t\t\tcompat: ${JSON.stringify(model.compat)},\n`; + output += `\t\t\treasoning: ${model.reasoning},\n`; + if (model.thinkingLevelMap) { + output += `\t\t\tthinkingLevelMap: ${JSON.stringify(model.thinkingLevelMap)},\n`; + } + output += `\t\t\tinput: [${model.input.map((input) => JSON.stringify(input)).join(", ")}],\n`; + output += `\t\t\tcost: {\n`; + output += `\t\t\t\tinput: ${model.cost.input},\n`; + output += `\t\t\t\toutput: ${model.cost.output},\n`; + output += `\t\t\t\tcacheRead: ${model.cost.cacheRead},\n`; + output += `\t\t\t\tcacheWrite: ${model.cost.cacheWrite},\n`; + output += `\t\t\t},\n`; + output += `\t\t\tcontextWindow: ${model.contextWindow},\n`; + output += `\t\t\tmaxTokens: ${model.maxTokens},\n`; + if (model.featured) output += `\t\t\tfeatured: true,\n`; + output += `\t\t} satisfies Model<${JSON.stringify(model.api)}>,\n`; + } + + output += `\t},\n`; + } + + return `${output}} as const;\n`; +} diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 388dffd6e2..7b59b84597 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -5,6 +5,7 @@ export * from "./api-registry.js"; export * from "./env-api-keys.js"; export * from "./log.js"; export * from "./models.js"; +export * from "./prime-inference-model-catalog.js"; export type { BedrockOptions, BedrockThinkingDisplay } from "./providers/amazon-bedrock.js"; export type { AnthropicEffort, AnthropicOptions, AnthropicThinkingDisplay } from "./providers/anthropic.js"; export type { AzureOpenAIResponsesOptions } from "./providers/azure-openai-responses.js"; diff --git a/packages/ai/src/prime-inference-model-catalog.ts b/packages/ai/src/prime-inference-model-catalog.ts new file mode 100644 index 0000000000..f1dc3101c3 --- /dev/null +++ b/packages/ai/src/prime-inference-model-catalog.ts @@ -0,0 +1,93 @@ +export interface PrimeInferenceCatalogEntry { + id: string; + name?: string; + input: number; + output: number; + cacheRead?: number; + cacheWrite?: number; + contextWindow?: number; + maxTokens?: number; + vision?: boolean; + reasoning?: boolean; +} + +function isRecord(value: unknown): value is Record { + return typeof value === "object" && value !== null && !Array.isArray(value); +} + +function nonNegativeNumber(value: unknown): number | undefined { + return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined; +} + +function positiveInteger(value: unknown): number | undefined { + return typeof value === "number" && Number.isInteger(value) && value > 0 ? value : undefined; +} + +export function isPrivatePrimeInferenceModelId(modelId: string): boolean { + const normalizedId = modelId.toLowerCase(); + return normalizedId.startsWith("internal/") || normalizedId.startsWith("dev/") || normalizedId.includes(":"); +} + +export function parsePrimeInferenceModelCatalog( + value: unknown, + options: { allowEmpty?: boolean } = {}, +): PrimeInferenceCatalogEntry[] { + if (!isRecord(value) || !Array.isArray(value.data)) throw new Error("Invalid Prime Inference model catalog"); + const models: PrimeInferenceCatalogEntry[] = []; + const seen = new Set(); + for (const item of value.data) { + if (!isRecord(item) || typeof item.id !== "string" || !item.id || item.id.length > 1_024) continue; + if (/[\u0000-\u001f\u007f-\u009f]/.test(item.id)) continue; + if (seen.has(item.id)) throw new Error(`Duplicate Prime Inference model ${item.id}`); + const pricing = isRecord(item.pricing) ? item.pricing : {}; + const input = nonNegativeNumber(pricing.input_usd_per_mtok); + const output = nonNegativeNumber(pricing.output_usd_per_mtok); + if (input === undefined || output === undefined) continue; + + const name = + typeof item.display_name === "string" + ? item.display_name.replace(/[\u0000-\u001f\u007f-\u009f]/g, "").trim() + : ""; + const specs = isRecord(item.specs) ? item.specs : {}; + const modalities = isRecord(specs.modalities) ? specs.modalities : {}; + const inputModalities = + Array.isArray(modalities.input) && modalities.input.every((modality) => typeof modality === "string") + ? modalities.input + : undefined; + const outputModalities = + Array.isArray(modalities.output) && modalities.output.every((modality) => typeof modality === "string") + ? modalities.output + : undefined; + const contextWindow = positiveInteger(specs.context_window); + const maxTokens = positiveInteger(specs.max_output_tokens); + const reasoning = typeof specs.supports_reasoning === "boolean" ? specs.supports_reasoning : undefined; + const hasSpecs = + contextWindow !== undefined && + maxTokens !== undefined && + reasoning !== undefined && + inputModalities !== undefined && + outputModalities !== undefined; + const cacheRead = nonNegativeNumber(pricing.cache_read_usd_per_mtok); + const cacheWrite = nonNegativeNumber(pricing.cache_write_usd_per_mtok); + + seen.add(item.id); + models.push({ + id: item.id, + ...(name ? { name } : {}), + input, + output, + ...(cacheRead !== undefined ? { cacheRead } : {}), + ...(cacheWrite !== undefined ? { cacheWrite } : {}), + ...(hasSpecs + ? { + contextWindow, + maxTokens: Math.min(maxTokens, contextWindow), + vision: inputModalities.includes("image"), + reasoning, + } + : {}), + }); + } + if (models.length === 0 && !options.allowEmpty) throw new Error("Prime Inference model catalog is empty"); + return models; +} diff --git a/packages/ai/test/prime-inference-model-catalog.test.ts b/packages/ai/test/prime-inference-model-catalog.test.ts new file mode 100644 index 0000000000..783c7583ff --- /dev/null +++ b/packages/ai/test/prime-inference-model-catalog.test.ts @@ -0,0 +1,101 @@ +import { describe, expect, test } from "vitest"; +import { parsePrimeInferenceModelCatalog } from "../src/prime-inference-model-catalog.js"; + +function response(...data: unknown[]) { + return { object: "list", data }; +} + +describe("Prime Inference model catalog", () => { + test("parses pricing and model specs", () => { + const [model] = parsePrimeInferenceModelCatalog( + response({ + id: "vendor/model", + display_name: "Model Name", + pricing: { + input_usd_per_mtok: 1, + output_usd_per_mtok: 2, + cache_read_usd_per_mtok: 0.1, + cache_write_usd_per_mtok: 1.25, + }, + specs: { + context_window: 200_000, + max_output_tokens: 64_000, + modalities: { input: ["text", "image", "file"], output: ["text"] }, + supports_reasoning: true, + }, + }), + ); + expect(model).toEqual({ + id: "vendor/model", + name: "Model Name", + input: 1, + output: 2, + cacheRead: 0.1, + cacheWrite: 1.25, + contextWindow: 200_000, + maxTokens: 64_000, + vision: true, + reasoning: true, + }); + }); + + test("keeps priced entries without complete specs for bundled fallback", () => { + expect( + parsePrimeInferenceModelCatalog( + response( + { + id: "vendor/no-specs", + pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, + specs: null, + }, + { + id: "vendor/partial-specs", + pricing: { input_usd_per_mtok: 3, output_usd_per_mtok: 4 }, + specs: { + context_window: 100_000, + max_output_tokens: 10_000, + modalities: { output: ["text"] }, + supports_reasoning: false, + }, + }, + ), + ), + ).toEqual([ + { id: "vendor/no-specs", input: 1, output: 2 }, + { id: "vendor/partial-specs", input: 3, output: 4 }, + ]); + }); + + test("rejects control-bearing IDs without rewriting Unicode model identities", () => { + const pricing = { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }; + const id = "vendor/模型-é"; + const controls = Array.from({ length: 33 }, (_, i) => String.fromCharCode(i + 0x7f)).concat( + Array.from({ length: 32 }, (_, i) => String.fromCharCode(i)), + ); + const models = parsePrimeInferenceModelCatalog( + response({ id, pricing }, ...controls.map((control) => ({ id: `${id}${control}`, pricing }))), + ); + expect(models.map((model) => model.id)).toEqual([id]); + }); + + test("removes display-name terminal controls while preserving Unicode", () => { + const pricing = { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }; + const controls = + Array.from({ length: 32 }, (_, i) => String.fromCharCode(i)).join("") + + Array.from({ length: 33 }, (_, i) => String.fromCharCode(i + 0x7f)).join(""); + const models = parsePrimeInferenceModelCatalog( + response( + { id: "unicode", display_name: ` 模型 é 👩‍💻${controls} `, pricing }, + { id: "fallback", display_name: controls, pricing }, + ), + ); + expect(models[0].name).toBe("模型 é 👩‍💻"); + expect(models[1]).not.toHaveProperty("name"); + }); + + test("rejects empty and duplicate catalogs", () => { + expect(() => parsePrimeInferenceModelCatalog(response())).toThrow(/empty/); + const model = { id: "duplicate", pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 } }; + expect(() => parsePrimeInferenceModelCatalog(response(model, model))).toThrow(/duplicate/i); + }); +}); diff --git a/packages/ai/test/prime-inference-models.test.ts b/packages/ai/test/prime-inference-models.test.ts index cd79d4a177..58c3a03b8f 100644 --- a/packages/ai/test/prime-inference-models.test.ts +++ b/packages/ai/test/prime-inference-models.test.ts @@ -42,13 +42,10 @@ describe("Prime Inference models", () => { ); }); - it("skips private, raw, and duplicate catalog variants", () => { + it("excludes private routes from the bundled public catalog", () => { const modelIds = getModels("prime-inference").map((model) => model.id); - expect(modelIds.filter((id) => id.startsWith("internal/"))).toEqual([]); - expect(modelIds).not.toContain("zai-org/GLM-4.7"); - expect(modelIds).not.toContain("nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"); - expect(modelIds).not.toContain("Qwen/Qwen3.5-4B"); + expect(modelIds.filter((id) => id.startsWith("internal/") || id.startsWith("dev/"))).toEqual([]); expect(modelIds.filter((id) => id.includes(":"))).toEqual([]); }); @@ -94,8 +91,7 @@ describe("Prime Inference models", () => { expect(model.reasoning).toBe(true); expect(getSupportedThinkingLevels(model)).toEqual(["off", "low", "high", "max"]); expect(model.input).toEqual(["text", "image"]); - expect(model.contextWindow).toBe(1048576); - expect(model.maxTokens).toBe(1048576); + expect(model.maxTokens).toBeLessThanOrEqual(model.contextWindow); expect(model.cost.input).toBe(provider === "prime-inference" ? 3.45 : 3); expect(model.cost.output).toBe(provider === "prime-inference" ? 17.25 : 15); } @@ -111,8 +107,7 @@ describe("Prime Inference models", () => { const nemotronSuper = getModel("prime-inference", "nvidia/nemotron-3-super-120b-a12b"); expect(nemotronSuper.reasoning).toBe(true); expect(nemotronSuper.input).toEqual(["text"]); - expect(nemotronSuper.contextWindow).toBe(262144); - expect(nemotronSuper.maxTokens).toBe(4096); + expect(nemotronSuper.maxTokens).toBeLessThanOrEqual(nemotronSuper.contextWindow); const maverick = getModel("prime-inference", "meta-llama/llama-4-maverick"); expect(maverick.contextWindow).toBe(1048576); @@ -207,7 +202,6 @@ describe("Prime Inference models", () => { expect(getModel("prime-inference", "anthropic/claude-sonnet-4.6").contextWindow).toBe(1000000); expect(getModel("prime-inference", "anthropic/claude-sonnet-5").contextWindow).toBe(1000000); expect(getModel("prime-inference", "anthropic/claude-haiku-4.5").contextWindow).toBe(200000); - expect(getModel("prime-inference", "anthropic/claude-sonnet-4.5").contextWindow).toBe(200000); }); it("resolves PRIME_API_KEY from the environment", () => { diff --git a/packages/ai/test/render-models.test.ts b/packages/ai/test/render-models.test.ts new file mode 100644 index 0000000000..b0fbc5f457 --- /dev/null +++ b/packages/ai/test/render-models.test.ts @@ -0,0 +1,30 @@ +import { describe, expect, test } from "vitest"; +import { renderModelsFile } from "../scripts/render-models.js"; + +describe("generated model serialization", () => { + test("escapes remote strings before writing TypeScript source", () => { + const id = 'vendor/model";\nexport const injected = true; //'; + const name = 'Model "name"\nwith a newline'; + const output = renderModelsFile({ + "prime-inference": { + [id]: { + id, + name, + api: "openai-completions", + provider: "prime-inference", + baseUrl: "https://api.pinference.ai/api/v1", + reasoning: false, + input: ["text"], + cost: { input: 1, output: 2, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 128_000, + maxTokens: 8_192, + }, + }, + }); + + expect(output).toContain(`\t\t${JSON.stringify(id)}: {`); + expect(output).toContain(`\t\t\tid: ${JSON.stringify(id)},`); + expect(output).toContain(`\t\t\tname: ${JSON.stringify(name)},`); + expect(output).not.toContain(`name: "${name}",`); + }); +}); diff --git a/packages/coding-agent/.changes/agents-view-anchor-open-block.md b/packages/coding-agent/.changes/agents-view-anchor-open-block.md new file mode 100644 index 0000000000..d754f32828 --- /dev/null +++ b/packages/coding-agent/.changes/agents-view-anchor-open-block.md @@ -0,0 +1 @@ +- Fixed the agents view blocking Enter with "Waiting for the selected session to load" while the remembered selection was still loading; opening the visible row now always works, and entering a subagents view no longer arms that wait at all. diff --git a/packages/coding-agent/.changes/agents-view-model-everywhere.md b/packages/coding-agent/.changes/agents-view-model-everywhere.md new file mode 100644 index 0000000000..e545e9947b --- /dev/null +++ b/packages/coding-agent/.changes/agents-view-model-everywhere.md @@ -0,0 +1 @@ +- Changed the agents view to show the model label on every session row, not only on subagent rows. diff --git a/packages/coding-agent/.changes/async-bash-completion-message.md b/packages/coding-agent/.changes/async-bash-completion-message.md new file mode 100644 index 0000000000..2eff0b940b --- /dev/null +++ b/packages/coding-agent/.changes/async-bash-completion-message.md @@ -0,0 +1,2 @@ +- Added steering Shell messages when background kernel `bash()` process groups finish so agents can inspect results at the next safe turn boundary without interrupting running tools. +- Kept sessions resident while background shell process groups run and completion delivery is pending. diff --git a/packages/coding-agent/.changes/compact-agents-view.md b/packages/coding-agent/.changes/compact-agents-view.md new file mode 100644 index 0000000000..920001f097 --- /dev/null +++ b/packages/coding-agent/.changes/compact-agents-view.md @@ -0,0 +1,2 @@ +- Simplified the agents view with total cost and age, one column header, and collapsed inactive sessions while keeping the logo, startup metadata, and search. +- Kept a running-subagent count beneath collapsed agents while their subagents are working. diff --git a/packages/coding-agent/.changes/feat-queue-path-arg-highlighting.md b/packages/coding-agent/.changes/feat-queue-path-arg-highlighting.md new file mode 100644 index 0000000000..67da1a3d1c --- /dev/null +++ b/packages/coding-agent/.changes/feat-queue-path-arg-highlighting.md @@ -0,0 +1 @@ +- Highlighted `@path` file references and `--flags` in the editor, queued message previews, and sent user messages, plus the bare `--` end-of-options separator in recognized slash commands. diff --git a/packages/coding-agent/.changes/prime-inference-model-catalog.md b/packages/coding-agent/.changes/prime-inference-model-catalog.md new file mode 100644 index 0000000000..e8317479a5 --- /dev/null +++ b/packages/coding-agent/.changes/prime-inference-model-catalog.md @@ -0,0 +1 @@ +- Added live refreshes for public and authorized private Prime Inference models while retaining bundled and cached fallbacks. diff --git a/packages/coding-agent/.changes/res-1256-rlm-root-session.md b/packages/coding-agent/.changes/res-1256-rlm-root-session.md new file mode 100644 index 0000000000..793a1404e1 --- /dev/null +++ b/packages/coding-agent/.changes/res-1256-rlm-root-session.md @@ -0,0 +1,2 @@ +- Added `rlm.create_session(...)` so daemon-backed root agents can start separate top-level sessions. +- Preserved active same-provider credentials when creating a sibling session without storing them in daemon descriptors. diff --git a/packages/coding-agent/docs/providers.md b/packages/coding-agent/docs/providers.md index 5c7ff364bb..1aaa6a2957 100644 --- a/packages/coding-agent/docs/providers.md +++ b/packages/coding-agent/docs/providers.md @@ -1,6 +1,6 @@ # Providers -Prime Agent supports subscription-based providers via OAuth and API key providers via environment variables or the auth file. Its built-in model catalog is updated with each Prime Agent release. +Prime Agent supports subscription-based providers via OAuth and API key providers via environment variables or the auth file. Models for external providers are bundled with each release. Prime Inference models refresh from its `/models` endpoint, with the bundled list and a validated disk cache as fallbacks. Set `PI_OFFLINE=1` to skip network refreshes. ## Table of Contents diff --git a/packages/coding-agent/docs/rlm.md b/packages/coding-agent/docs/rlm.md index 7d558073b4..27653930fd 100644 --- a/packages/coding-agent/docs/rlm.md +++ b/packages/coding-agent/docs/rlm.md @@ -48,6 +48,17 @@ result = await bash("npm run check") print(result.output) ``` +For a long command, keep the live handle and let the turn end instead of blocking: + +```python +checks = bash("npm test") +checks.pid +``` + +When an unawaited handle's process group finishes, Prime Agent sends a Shell message with its PID and foreground exit code. A busy agent receives it as steering at the next safe turn boundary, without interrupting a running tool. An idle agent resumes to handle it. `await handle` and `handle.poll()` still return the foreground result before shell background jobs finish. The kernel stays resident until the process group is reaped, including for handles awaited in their creating cell. The message asks the agent to inspect the saved handle with `poll()`, `output()`, or `tail()` and continue the task. `await bash(...)` stays synchronous from the agent's perspective and does not send a second Shell message. + +Awaiting the original handle in its creating cell suppresses the Shell message, even after the command finishes. If an `asyncio.as_completed` wrapper task finishes before the cell starts consuming its result, that cached-result read does not mark the handle as awaited. A Shell message can still arrive. Await the original handle in the creating cell to suppress it. + Each `bash()` call is its own process, while Python state, `os.chdir(...)`, and `os.environ[...]` changes persist in the kernel and apply to later `bash()` calls. Prime Agent extensions may intentionally add custom tools, but the built-in RLM design does not require a separate model tool for every capability. ### 2. Subagents are native RLM calls diff --git a/packages/coding-agent/src/core/agent-messages.ts b/packages/coding-agent/src/core/agent-messages.ts index 0599c003c9..ba1ad59266 100644 --- a/packages/coding-agent/src/core/agent-messages.ts +++ b/packages/coding-agent/src/core/agent-messages.ts @@ -2,7 +2,7 @@ import { randomUUID } from "node:crypto"; import type { AgentMessage } from "@earendil-works/pi-agent-core"; import type { HostRequestHandler } from "./kernel/index.js"; import type { CustomMessage } from "./messages.js"; -import { HEARTBEAT_PROMPT_CUSTOM_TYPE } from "./messages.js"; +import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE, HEARTBEAT_PROMPT_CUSTOM_TYPE } from "./messages.js"; import { canonicalSessionPath } from "./session-lease.js"; export const AGENT_MESSAGE_CUSTOM_TYPE = "agent_message"; @@ -443,7 +443,9 @@ export function startsAgentRun(message: AgentMessage): boolean { return ( message.role === "user" || isAgentSessionMessage(message) || - (message.role === "custom" && message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE) + (message.role === "custom" && + (message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE || + message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE)) ); } diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index ef39ffada5..b6c93aa521 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -165,11 +165,15 @@ import { type RestoreResult, snapshotPathIn } from "./kernel/state-snapshot.js"; import type { AcpMcpServerConfig } from "./mcp/acp-mcp-types.js"; import type { McpManager } from "./mcp/mcp-manager.js"; import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ASYNC_BASH_COMPLETION_PREVIEW_LABEL, + type AsyncBashCompletionDetails, type BashExecutionMessage, type CompactionOutcome, type CompactionOutcomeReason, type CustomMessage, convertToLlm, + createAsyncBashCompletionMessage, createCompactionOutcomeMessage, createHeartbeatPromptMessage, createRefinementOutcomeMessage, @@ -223,7 +227,9 @@ import { resolveConfigValue } from "./resolve-config-value.js"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.js"; import { type CreateRlmSubagentRuntimeOptions, + createAsyncBashCompletionHostHandler, createDefaultRlmSubagentSessionName, + createRlmCreateSessionHostHandler, createRlmDeleteSubagentHostHandler, createRlmFindModelsHostHandler, createRlmListSubagentsHostHandler, @@ -232,6 +238,7 @@ import { normalizeRequestedRlmSubagentModel, normalizeRequestedRlmSubagentSessionName, normalizeRequestedRlmSubagentThinkingLevel, + type RlmCreateSessionResult, type RlmDeleteSubagentResult, type RlmFindModelsResult, type RlmListSubagentsResult, @@ -648,6 +655,8 @@ interface PreparedPromptPreparation { class DeferredSessionInputError extends Error {} +class SessionInputAdmissionPausedError extends Error {} + function oncePreflight( preflightResult: ((success: boolean, queued?: boolean) => void) | undefined, ): (success: boolean, queued?: boolean) => void { @@ -757,6 +766,12 @@ function queuedAgentMessagePreview(action: QueuedSessionAction): string { if (payload.customMessage && isAgentSessionMessage(payload.customMessage)) { return `${AGENT_MESSAGE_RECEIVED_PREVIEW_LABEL}: ${payload.customMessage.details.message}`; } + if (payload.customMessage?.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE) { + const details = payload.customMessage.details as AsyncBashCompletionDetails | undefined; + return details + ? `${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}: pid ${details.pid}, exit ${details.exitCode}` + : ASYNC_BASH_COMPLETION_PREVIEW_LABEL; + } return payload.preview ?? payload.text; } @@ -834,6 +849,8 @@ function injectedMessagePreviewLabel(message: CustomMessage): string | undefined switch (message.customType) { case HEARTBEAT_PROMPT_CUSTOM_TYPE: return HEARTBEAT_PROMPT_PREVIEW_LABEL; + case ASYNC_BASH_COMPLETION_CUSTOM_TYPE: + return ASYNC_BASH_COMPLETION_PREVIEW_LABEL; case GOAL_CONTEXT_CUSTOM_TYPE: return GOAL_CONTEXT_PREVIEW_LABEL; default: @@ -5654,7 +5671,9 @@ export class AgentSession { throw new Error("Cannot admit a session action because the session is disposing or disposed."); } if (this._sessionInputAdmissionPauses.size > 0) { - throw new Error("Cannot admit a session action while session input admission is paused."); + throw new SessionInputAdmissionPausedError( + "Cannot admit a session action while session input admission is paused.", + ); } if (this._sessionInputPumpSuspended) { throw new Error("Cannot admit a session action while queued session input is suspended."); @@ -5678,7 +5697,9 @@ export class AgentSession { throw new Error("Cannot admit a session action because the session is disposing or disposed."); } if (this._sessionInputAdmissionPauses.size > 0) { - throw new Error("Cannot admit a session action while session input admission is paused."); + throw new SessionInputAdmissionPausedError( + "Cannot admit a session action while session input admission is paused.", + ); } if ( options.restore !== true && @@ -6599,6 +6620,7 @@ export class AgentSession { get isSessionActive(): boolean { return ( + this._ipythonKernelProvisioner?.manager?.hasBackgroundWork === true || this.isStreaming || this.isCompacting || this.isRetrying || @@ -9410,6 +9432,34 @@ export class AgentSession { "rlm.run": createRlmRunHostHandler(async ({ prompt, kwargs, cellSourceCode }) => ({ ...(await this.runRlmChild(prompt, kwargs, cellSourceCode)), })), + "rlm.create_session": createRlmCreateSessionHostHandler(async ({ prompt, kwargs }) => ({ + ...(await this.createRlmSession(prompt, kwargs)), + })), + "bash.completed": createAsyncBashCompletionHostHandler(async (details) => { + const message = createAsyncBashCompletionMessage(details); + const disposeSignal = this._sessionActionCommitDisposeAbortController.signal; + while (true) { + let admissionCommitted = false; + try { + await this._promptInjectedMessage(message.content, message, { + streamingBehavior: "steer", + queueIfBusy: true, + resumeIfIdle: true, + returnAfterAccepted: true, + suppressAutonomousContinuation: true, + admissionCommitted: () => { + admissionCommitted = true; + }, + }); + return; + } catch (error) { + if (admissionCommitted || !(error instanceof SessionInputAdmissionPausedError)) throw error; + while (this._sessionInputAdmissionPauses.size > 0 && !disposeSignal.aborted) { + await this._waitForSessionActivityChange(disposeSignal); + } + } + } + }), "rlm.find_models": createRlmFindModelsHostHandler((query, limit) => this.findRlmModels(query, limit)), "rlm.list_subagents": createRlmListSubagentsHostHandler(() => this.listRlmSubagents()), "rlm.delete_subagent": createRlmDeleteSubagentHostHandler((target) => this.deleteRlmSubagent(target)), @@ -10584,7 +10634,10 @@ export class AgentSession { }; } - private async _resolveRlmSubagentModel(reference: string | undefined): Promise { + private async _resolveRlmSubagentModel( + reference: string | undefined, + target = "subagent", + ): Promise { const parentModel = this.model; if (!parentModel) { throw new Error(formatNoModelSelectedMessage()); @@ -10601,12 +10654,12 @@ export class AgentSession { (candidate) => `${candidate.provider}/${candidate.id}`.toLowerCase() === normalizedReference, ); if (!model) { - throw new Error(`Requested subagent model "${reference}" is unavailable, unauthenticated, or expired`); + throw new Error(`Requested ${target} model "${reference}" is unavailable, unauthenticated, or expired`); } const auth = await this._modelRegistry.getApiKeyAndHeaders(model); if (!auth.ok) { - throw new Error(`Requested subagent model "${reference}" failed authentication preflight`); + throw new Error(`Requested ${target} model "${reference}" failed authentication preflight`); } return { model }; } @@ -11087,6 +11140,64 @@ export class AgentSession { }; } + async createRlmSession(prompt: string, kwargs: Record = {}): Promise { + const { name: rawName, model: rawModel, thinking: rawThinking, cwd: rawCwd, ...unsupported } = kwargs; + const unsupportedKeys = Object.keys(unsupported); + if (unsupportedKeys.length > 0) { + throw new Error(`Unsupported rlm.create_session kwargs: ${unsupportedKeys.sort().join(", ")}`); + } + if (!prompt.trim()) { + throw new Error("rlm.create_session prompt must not be empty"); + } + if (this._rlmDepth !== 0) { + throw new Error("rlm.create_session is available only from a depth-0 session"); + } + if (this._disposed || this._disposing) { + throw new Error("Cannot create a top-level session after the current session was disposed"); + } + const host = this._subagentRuntimeHost; + if (!host?.createRlmRootSession) { + throw new Error("rlm.create_session requires a daemon-backed depth-0 session"); + } + + const operation = "rlm.create_session"; + const sessionName = normalizeRequestedRlmSubagentSessionName(rawName, operation); + const requestedModel = normalizeRequestedRlmSubagentModel(rawModel, operation); + const requestedThinkingLevel = normalizeRequestedRlmSubagentThinkingLevel(rawThinking, operation); + if (sessionName) { + assertDirectAgentMessageTarget(sessionName); + const controller = this._agentMessageController; + if (controller?.assertSessionNameAvailable) { + await controller.assertSessionNameAvailable({ name: sessionName, depth: 0 }); + } + } + if (rawCwd !== undefined && (typeof rawCwd !== "string" || !rawCwd.trim())) { + throw new Error("rlm.create_session cwd must be a non-empty string"); + } + const cwd = rawCwd === undefined ? this._cwd : resolve(this._cwd, rawCwd.trim()); + const modelSelection = await this._resolveRlmSubagentModel(requestedModel, "top-level session"); + if (requestedThinkingLevel !== undefined) { + const supported = getSupportedThinkingLevels(modelSelection.model) as ThinkingLevel[]; + if (!supported.includes(requestedThinkingLevel)) { + throw new Error( + `Requested thinking level "${requestedThinkingLevel}" is not supported by model "${modelSelection.model.provider}/${modelSelection.model.id}"; supported levels: ${supported.join(", ")}`, + ); + } + } + const thinkingLevel = + requestedThinkingLevel ?? (clampThinkingLevel(modelSelection.model, this.thinkingLevel) as ThinkingLevel); + if (this._disposed || this._disposing) { + throw new Error("Cannot create a top-level session after the current session was disposed"); + } + return host.createRlmRootSession({ + prompt, + sessionName, + cwd, + model: modelSelection.model, + thinkingLevel, + }); + } + async runRlmChild( prompt: string, kwargs: Record = {}, diff --git a/packages/coding-agent/src/core/kernel/bootstrap.ts b/packages/coding-agent/src/core/kernel/bootstrap.ts index d3e4ddcceb..5ab93cc64e 100644 --- a/packages/coding-agent/src/core/kernel/bootstrap.ts +++ b/packages/coding-agent/src/core/kernel/bootstrap.ts @@ -87,7 +87,7 @@ const REQUIRED_HARNESS_METHODS = [ "delete_prompt_note", "record_refinement", ]; -const RUNTIME_READY_CHECK = `import inspect; import rlm; from rlm import McpIntegration; import rlm.mcp as mcp; from rlm.harness import HarnessEntry; _harness_methods = ${JSON.stringify(REQUIRED_HARNESS_METHODS)}; assert callable(mcp.list_tools); assert callable(mcp.call_tool); assert hasattr(rlm, 'run'); assert callable(rlm); assert hasattr(rlm, 'rlm'); assert callable(rlm.rlm); assert callable(rlm.host_request); assert callable(rlm.find_models); assert callable(rlm.rlm.find_models); assert hasattr(rlm, 'harness'); assert hasattr(rlm, 'get_harness_state'); assert hasattr(rlm.rlm, 'harness'); assert hasattr(rlm.rlm, 'get_harness_state'); assert all(callable(getattr(_harness, _method, None)) for _harness in (rlm.harness, rlm.rlm.harness) for _method in _harness_methods); assert 'reference' in HarnessEntry.__dataclass_fields__; assert 'scope' in HarnessEntry.__dataclass_fields__; assert 'reference' in inspect.signature(rlm.harness.create_skill).parameters; assert 'reference' in inspect.signature(rlm.harness.update_skill).parameters; assert 'global_' in inspect.signature(rlm.harness.create_memory).parameters; assert 'global_' in inspect.signature(rlm.get_harness_state).parameters; assert not hasattr(rlm, 'background'); assert not hasattr(rlm.rlm, 'background'); from rlm.bash import BashHandle, BashResult; assert callable(rlm.bash); assert all(callable(getattr(BashHandle, _m, None)) for _m in ('tail', 'output', 'poll', 'kill')); assert {'exit_code', 'output', 'duration'} <= set(BashResult.__dataclass_fields__); import rlm.repl as _repl; assert callable(_repl.main); assert callable(_repl.emit); assert callable(_repl.host_request); assert callable(_repl.is_active); assert _repl.PROTOCOL_VERSION == 3; assert callable(rlm.emit); assert not hasattr(rlm, 'HOST_COMM_TARGET'); assert not hasattr(mcp, 'install_shutdown_hook')`; +const RUNTIME_READY_CHECK = `import inspect; import rlm; from rlm import McpIntegration; import rlm.mcp as mcp; from rlm.harness import HarnessEntry; _harness_methods = ${JSON.stringify(REQUIRED_HARNESS_METHODS)}; assert callable(mcp.list_tools); assert callable(mcp.call_tool); assert hasattr(rlm, 'run'); assert callable(rlm); assert hasattr(rlm, 'rlm'); assert callable(rlm.rlm); assert callable(rlm.host_request); assert callable(rlm.find_models); assert callable(rlm.rlm.find_models); assert callable(rlm.create_session); assert callable(rlm.rlm.create_session); assert hasattr(rlm, 'harness'); assert hasattr(rlm, 'get_harness_state'); assert hasattr(rlm.rlm, 'harness'); assert hasattr(rlm.rlm, 'get_harness_state'); assert all(callable(getattr(_harness, _method, None)) for _harness in (rlm.harness, rlm.rlm.harness) for _method in _harness_methods); assert 'reference' in HarnessEntry.__dataclass_fields__; assert 'scope' in HarnessEntry.__dataclass_fields__; assert 'reference' in inspect.signature(rlm.harness.create_skill).parameters; assert 'reference' in inspect.signature(rlm.harness.update_skill).parameters; assert 'global_' in inspect.signature(rlm.harness.create_memory).parameters; assert 'global_' in inspect.signature(rlm.get_harness_state).parameters; assert not hasattr(rlm, 'background'); assert not hasattr(rlm.rlm, 'background'); from rlm.bash import BashHandle, BashResult; assert callable(rlm.bash); assert all(callable(getattr(BashHandle, _m, None)) for _m in ('tail', 'output', 'poll', 'kill')); assert {'exit_code', 'output', 'duration'} <= set(BashResult.__dataclass_fields__); import rlm.repl as _repl; assert callable(_repl.main); assert callable(_repl.emit); assert callable(_repl.host_request); assert callable(_repl.is_active); assert _repl.PROTOCOL_VERSION == 3; assert callable(rlm.emit); assert not hasattr(rlm, 'HOST_COMM_TARGET'); assert not hasattr(mcp, 'install_shutdown_hook')`; const BOOTSTRAP_VERSION_FILE = ".bootstrap-version"; const BOOTSTRAP_LOCK_NAME = ".bootstrap.lock"; const BOOTSTRAP_LOCK_RETRY_MS = 100; @@ -896,7 +896,7 @@ async function ensureKernelPythonUncached( const missing: string[] = []; if (!(await hasPrimeAgentRuntime(python))) { missing.push( - "a current prime-agent-runtime with callable rlm.run, rlm.host_request, and explicit harness CRUD methods", + "a current prime-agent-runtime with callable rlm.run, rlm.create_session, rlm.host_request, and explicit harness CRUD methods", ); } if (missing.length === 0) { diff --git a/packages/coding-agent/src/core/kernel/repl-manager.ts b/packages/coding-agent/src/core/kernel/repl-manager.ts index 51e4bf6452..5023e48729 100644 --- a/packages/coding-agent/src/core/kernel/repl-manager.ts +++ b/packages/coding-agent/src/core/kernel/repl-manager.ts @@ -12,6 +12,7 @@ import { ensureKernelPython } from "./bootstrap.js"; import { AGENT_MESSAGE_DISPLAY_MIME, ATTACHMENT_DISPLAY_MIME, + BASH_ACTIVITY_DISPLAY_MIME, createDeferred, createKernelStartupAbortError, DEFAULT_MAX_OUTPUT_CHARS, @@ -180,6 +181,7 @@ export class ReplKernelManager { private pendingBackgroundOutput = ""; private pendingBackgroundOutputTruncated = false; private readonly inFlightHostRequests = new Set>(); + private readonly backgroundBashHandles = new Map(); private state: "idle" | "starting" | "running" | "shutdown" = "idle"; /** Bumped by every teardown so a stale in-flight doStart can never touch a newer kernel. */ private startGeneration = 0; @@ -224,6 +226,10 @@ export class ReplKernelManager { return this.options.sessionId; } + get hasBackgroundWork(): boolean { + return this.backgroundBashHandles.size > 0; + } + private appendKernelDiagnostic(message: string): void { this.appendKernelStderrText(`[kernel] ${message.endsWith("\n") ? message : `${message}\n`}`); } @@ -362,6 +368,7 @@ export class ReplKernelManager { buffered += decoder.write(buf); let newline = buffered.indexOf("\n"); while (newline !== -1) { + if (this.child !== child) return; const line = buffered.slice(0, newline); buffered = buffered.slice(newline + 1); newline = buffered.indexOf("\n"); @@ -737,6 +744,27 @@ export class ReplKernelManager { private handleEvent(event: Record): void { const type = event.event; + if (type === "display" && isRecord(event.data) && BASH_ACTIVITY_DISPLAY_MIME in event.data) { + const activity = event.data[BASH_ACTIVITY_DISPLAY_MIME]; + if ( + isRecord(activity) && + typeof activity.id === "string" && + /^[a-f0-9]{32}$/.test(activity.id) && + typeof activity.pid === "number" && + Number.isSafeInteger(activity.pid) && + activity.pid > 0 && + typeof activity.active === "boolean" + ) { + if (activity.active) { + if (!this.backgroundBashHandles.has(activity.id)) { + this.backgroundBashHandles.set(activity.id, activity.pid); + } + } else if (this.backgroundBashHandles.get(activity.id) === activity.pid) { + this.backgroundBashHandles.delete(activity.id); + } + } + return; + } if (type === "ready") { this.readyDeferred?.resolve(typeof event.protocol === "number" ? event.protocol : -1); return; @@ -1256,6 +1284,7 @@ export class ReplKernelManager { this.clearSnapshotTimer(); this.lateSentAgentMessageHandlers.clear(); this.pendingDoneWaiters.clear(); + this.backgroundBashHandles.clear(); // Stale pre-teardown background output must not surface after a restart. this.pendingBackgroundOutput = ""; this.pendingBackgroundOutputTruncated = false; diff --git a/packages/coding-agent/src/core/kernel/shared.ts b/packages/coding-agent/src/core/kernel/shared.ts index e6a6c63767..71b7ef7377 100644 --- a/packages/coding-agent/src/core/kernel/shared.ts +++ b/packages/coding-agent/src/core/kernel/shared.ts @@ -87,6 +87,9 @@ export const ATTACHMENT_DISPLAY_MIME = "application/vnd.prime-agent.attachment+j /** MIME tag the `agent-message` skill emits after sending a message. */ export const AGENT_MESSAGE_DISPLAY_MIME = "application/vnd.prime-agent.agent-message+json"; +/** Internal lifetime notices, consumed before user display rendering. */ +export const BASH_ACTIVITY_DISPLAY_MIME = "application/vnd.prime-agent.bash-activity+json"; + /** * Hard ceiling on a single attachment's base64 payload, a defensive guard * against a runaway direct display emit. The `attach-image` skill caps @@ -280,6 +283,7 @@ export interface KernelShutdownOptions { export interface KernelClient { readonly ownerSessionId: string | undefined; readonly isRunning: boolean; + readonly hasBackgroundWork: boolean; /** Terminal: the kernel died or was torn down; only a fresh manager can serve again. */ readonly isDefunct: boolean; start(options?: KernelStartOptions): Promise; diff --git a/packages/coding-agent/src/core/keybindings.ts b/packages/coding-agent/src/core/keybindings.ts index 6c451d53e6..471cb56003 100644 --- a/packages/coding-agent/src/core/keybindings.ts +++ b/packages/coding-agent/src/core/keybindings.ts @@ -48,6 +48,8 @@ export interface AppKeybindings { "app.agents.delete": true; "app.agents.program": true; "app.agents.rename": true; + "app.agents.inactiveCollapse": true; + "app.agents.expand": true; "app.tree.foldOrUp": true; "app.tree.unfoldOrDown": true; "app.tree.editLabel": true; @@ -159,6 +161,8 @@ export const KEYBINDINGS = { "app.agents.delete": { defaultKeys: "ctrl+x", description: "Stop or delete selected agent" }, "app.agents.program": { defaultKeys: "ctrl+o", description: "Show the program that spawned subagents" }, "app.agents.rename": { defaultKeys: "ctrl+r", description: "Rename selected agent session" }, + "app.agents.inactiveCollapse": { defaultKeys: "alt+i", description: "Show or hide inactive sessions" }, + "app.agents.expand": { defaultKeys: "alt+right", description: "Expand or collapse selected agent subagents" }, "app.tree.foldOrUp": { defaultKeys: ["ctrl+left", "alt+left"], description: "Fold tree branch or move up", diff --git a/packages/coding-agent/src/core/messages.ts b/packages/coding-agent/src/core/messages.ts index 81a77117f4..b1f470b56f 100644 --- a/packages/coding-agent/src/core/messages.ts +++ b/packages/coding-agent/src/core/messages.ts @@ -35,6 +35,8 @@ export const COMPACTION_OUTCOME_CUSTOM_TYPE = "compaction_outcome"; export const REFINEMENT_OUTCOME_CUSTOM_TYPE = "refinement_outcome"; export const RLM_CHILD_FAILURE_CUSTOM_TYPE = "rlm_child_failure"; export const RLM_CHILD_TERMINAL_NOTICE_CUSTOM_TYPE = "rlm_child_terminal_notice"; +export const ASYNC_BASH_COMPLETION_CUSTOM_TYPE = "async_bash_completion"; +export const ASYNC_BASH_COMPLETION_PREVIEW_LABEL = "Shell message received"; export interface SessionSlashCommandDetails { command: SessionSlashCommand; @@ -109,6 +111,36 @@ export type RlmChildTerminalNoticeDetails = lastAssistantTextPreview?: string; }; +export interface AsyncBashCompletionDetails { + pid: number; + command: string; + exitCode: number; +} + +interface AsyncBashCompletionMessage extends CustomMessage { + customType: typeof ASYNC_BASH_COMPLETION_CUSTOM_TYPE; + content: string; +} + +export function createAsyncBashCompletionMessage( + details: AsyncBashCompletionDetails, + timestamp = Date.now(), +): AsyncBashCompletionMessage { + return { + role: "custom", + customType: ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + content: `${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}. +Source: bash +Command completed (pid ${details.pid}, exit code ${details.exitCode}). +Command: ${JSON.stringify(details.command)} + +Inspect the saved BashHandle with .poll(), .output(), or .tail(), then continue the task.`, + display: true, + details, + timestamp, + }; +} + export function createRlmChildFailureMessage( details: RlmChildFailureDetails, timestamp = Date.now(), diff --git a/packages/coding-agent/src/core/model-registry.ts b/packages/coding-agent/src/core/model-registry.ts index cdca64a1de..e66aaef4ec 100644 --- a/packages/coding-agent/src/core/model-registry.ts +++ b/packages/coding-agent/src/core/model-registry.ts @@ -16,6 +16,7 @@ import { type OAuthProviderInterface, type OpenAICompletionsCompat, type OpenAIResponsesCompat, + parsePrimeInferenceModelCatalog, registerApiProvider, resetApiProviders, type SimpleStreamOptions, @@ -32,7 +33,13 @@ import { writeFileAtomicSync } from "../utils/atomic-file.js"; import type { AuthSourceToken, AuthStatus, AuthStorage } from "./auth-storage.js"; import { PRIME_INFERENCE_PROVIDER_ID } from "./prime-inference-auth.js"; import { - fetchAuthorizedPrivatePrimeInferenceModelIds, + buildPrimeInferenceModels, + mergePrimeInferenceModels, + readCachedPrimeInferenceModels, + refreshPrimeInferenceModels, +} from "./prime-inference-model-catalog.js"; +import { + fetchAuthorizedPrivatePrimeInferenceModels, getPrivatePrimeInferenceModels, isPrivatePrimeInferenceModel, } from "./prime-inference-models.js"; @@ -417,7 +424,7 @@ const PRIVATE_PRIME_BACKGROUND_REFRESH_TIMEOUT_MS = 3_000; interface PrivatePrimeAuthorizationCache { fingerprint: string; - modelIds: Set; + models: Model<"openai-completions">[]; refreshedAt: number; } @@ -442,10 +449,12 @@ export class ModelRegistry { private modelRequestHeaders: Map> = new Map(); private registeredProviders: Map = new Map(); private authorizedPrivatePrimeInferenceModelIds = new Set(); + private authorizedPrivatePrimeInferenceModels: Model<"openai-completions">[] = []; private authorizedPrivatePrimeInferenceTeamId: string | undefined; private explicitPrivatePrimeInferenceModelIds = new Set(); private openAICodexModelsCache: { authFingerprint: string; modelIds: Set; refreshedAt: number } | undefined; private backgroundPrivatePrimeAuthorization: { fingerprint: string; promise: Promise } | undefined; + private livePrimeInferenceModels: Model<"openai-completions">[] | undefined; private loadError: string | undefined = undefined; /** Re-register dynamic OAuth providers (e.g. user MCP servers) after refresh() resets the registry. */ @@ -491,6 +500,7 @@ export class ModelRegistry { teamId !== this.authorizedPrivatePrimeInferenceTeamId ) { this.authorizedPrivatePrimeInferenceModelIds.clear(); + this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; } resetApiProviders(); @@ -505,13 +515,28 @@ export class ModelRegistry { registerBuiltinMcpOAuthProviders(); this.onOAuthProvidersReset?.(); + this.reloadModelsAfterCatalogChange(); + } + + private reloadModelsAfterCatalogChange(): void { this.loadModels(); + this.reapplyRegisteredProviders(); + } + private reapplyRegisteredProviders(): void { for (const [providerName, config] of this.registeredProviders.entries()) { this.applyProviderConfig(providerName, config); } } + private primeInferenceCatalogCachePath(): string | undefined { + return this.modelsJsonPath ? join(dirname(this.modelsJsonPath), "prime-inference-models-cache.json") : undefined; + } + + private bundledPrimeInferenceModels(): Model<"openai-completions">[] { + return getModels(PRIME_INFERENCE_PROVIDER_ID) as Model<"openai-completions">[]; + } + /** * Get any error from loading models.json (undefined if no error). */ @@ -534,7 +559,20 @@ export class ModelRegistry { this.explicitPrivatePrimeInferenceModelIds = new Set( customModels.filter(isPrivatePrimeInferenceModel).map((model) => model.id), ); - const builtInModels = [...this.loadBuiltInModels(overrides, modelOverrides), ...getPrivatePrimeInferenceModels()]; + const cachePath = this.primeInferenceCatalogCachePath(); + this.livePrimeInferenceModels ??= cachePath + ? readCachedPrimeInferenceModels(cachePath, this.bundledPrimeInferenceModels()) + : undefined; + const privateModels = new Map( + [...getPrivatePrimeInferenceModels(), ...this.authorizedPrivatePrimeInferenceModels].map((model) => [ + model.id, + model, + ]), + ); + const builtInModels = [ + ...this.loadBuiltInModels(overrides, modelOverrides, this.livePrimeInferenceModels), + ...privateModels.values(), + ]; let combined = this.mergeCustomModels(builtInModels, customModels); for (const oauthProvider of this.authStorage.getOAuthProviders()) { @@ -551,30 +589,24 @@ export class ModelRegistry { private loadBuiltInModels( overrides: Map, modelOverrides: Map>, + livePrimeInferenceModels?: Model<"openai-completions">[], ): Model[] { - return getProviders().flatMap((provider) => { - const models = getModels(provider as KnownProvider) as Model[]; - const providerOverride = overrides.get(provider); - const perModelOverrides = modelOverrides.get(provider); - - return models.map((m) => { - let model = m; - - if (providerOverride) { - model = { - ...model, - baseUrl: providerOverride.baseUrl ?? model.baseUrl, - compat: mergeCompat(model.compat, providerOverride.compat), - }; - } - - const modelOverride = perModelOverrides?.get(m.id); - if (modelOverride) { - model = applyModelOverride(model, modelOverride); - } + const bundledModels = getProviders().flatMap((provider) => getModels(provider as KnownProvider) as Model[]); + return mergePrimeInferenceModels(bundledModels, livePrimeInferenceModels).map((model) => { + const providerOverride = overrides.get(model.provider); + const perModelOverrides = modelOverrides.get(model.provider); + let configuredModel = model; + + if (providerOverride) { + configuredModel = { + ...configuredModel, + baseUrl: providerOverride.baseUrl ?? configuredModel.baseUrl, + compat: mergeCompat(configuredModel.compat, providerOverride.compat), + }; + } - return model; - }); + const modelOverride = perModelOverrides?.get(model.id); + return modelOverride ? applyModelOverride(configuredModel, modelOverride) : configuredModel; }); } @@ -778,15 +810,31 @@ export class ModelRegistry { }); } + /** + * Reload local state and private authorization. Public Prime Inference models + * return from the disk/bundled fallback immediately and refresh in the background. + */ async refreshAvailableModels(): Promise[]> { - // Serialized: a concurrent call snapshotting between this call's clear - // (inside refresh()) and its restore would capture an empty entitlement - // set and lose the cached private-model authorization for good. return this.runSerializedEntitlementRefresh(async () => { const previousPrivateModelIds = new Set(this.authorizedPrivatePrimeInferenceModelIds); const previousTeamId = this.authorizedPrivatePrimeInferenceTeamId; + const previousPrivateModels = this.authorizedPrivatePrimeInferenceModels; this.refresh(); - await this.refreshPrivatePrimeInferenceAuthorization(previousPrivateModelIds, previousTeamId); + const cachePath = this.primeInferenceCatalogCachePath(); + if (cachePath) { + void refreshPrimeInferenceModels(cachePath, this.bundledPrimeInferenceModels(), { + offline: isOfflineModeEnabled(), + }).then((models) => { + if (!models) return; + this.livePrimeInferenceModels = models; + this.reloadModelsAfterCatalogChange(); + }); + } + await this.refreshPrivatePrimeInferenceAuthorization( + previousPrivateModelIds, + previousTeamId, + previousPrivateModels, + ); return this.getAvailable(); }); } @@ -805,6 +853,7 @@ export class ModelRegistry { private async refreshPrivatePrimeInferenceAuthorization( previousPrivateModelIds = new Set(this.authorizedPrivatePrimeInferenceModelIds), previousTeamId = this.authorizedPrivatePrimeInferenceTeamId, + previousPrivateModels = this.authorizedPrivatePrimeInferenceModels, ): Promise { const apiKey = await this.authStorage.getApiKey(PRIME_INFERENCE_PROVIDER_ID); const teamHeaders = this.authStorage.getProviderHeaders(PRIME_INFERENCE_PROVIDER_ID); @@ -815,41 +864,51 @@ export class ModelRegistry { // the team they were fetched for; a team switch invalidates them. if ( this.authStorage.getAuthStatus(PRIME_INFERENCE_PROVIDER_ID).source === "stale" && + teamId && teamId === previousTeamId ) { this.authorizedPrivatePrimeInferenceModelIds = previousPrivateModelIds; + this.authorizedPrivatePrimeInferenceModels = previousPrivateModels; this.authorizedPrivatePrimeInferenceTeamId = previousTeamId; + this.reloadModelsAfterCatalogChange(); return; } this.authorizedPrivatePrimeInferenceModelIds.clear(); + this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; + this.reloadModelsAfterCatalogChange(); return; } const fingerprint = privatePrimeAuthorizationFingerprint(apiKey, teamId); const cached = this.readPrivatePrimeAuthorizationCache(); if (cached?.fingerprint === fingerprint) { - // Serve the persisted authorization decision so startup and model lists - // don't block on the network. A stale cache refreshes in the background - // and the updated ids apply to subsequent lookups in this process. - this.authorizedPrivatePrimeInferenceModelIds = new Set(cached.modelIds); + // Serve the credential-scoped cache so startup and model lists don't + // block on the network. Stale entries refresh in the background. + this.authorizedPrivatePrimeInferenceModels = cached.models; + this.authorizedPrivatePrimeInferenceModelIds = new Set(cached.models.map((model) => model.id)); this.authorizedPrivatePrimeInferenceTeamId = teamId; + this.reloadModelsAfterCatalogChange(); const cacheIsFresh = Date.now() - cached.refreshedAt < PRIVATE_PRIME_AUTHORIZATION_CACHE_TTL_MS; - if (cacheIsFresh || isOfflineModeEnabled()) { - return; - } + if (isOfflineModeEnabled() || cacheIsFresh) return; this.startBackgroundPrivatePrimeAuthorizationRefresh(apiKey, teamHeaders, teamId, fingerprint); return; } if (isOfflineModeEnabled()) { this.authorizedPrivatePrimeInferenceModelIds.clear(); + this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; + this.reloadModelsAfterCatalogChange(); return; } - let authorizedIds: Set | undefined; + let authorizedModels: Model<"openai-completions">[] | undefined; try { - authorizedIds = await fetchAuthorizedPrivatePrimeInferenceModelIds(apiKey, teamHeaders); + authorizedModels = await fetchAuthorizedPrivatePrimeInferenceModels( + apiKey, + teamHeaders, + new Set((this.livePrimeInferenceModels ?? this.bundledPrimeInferenceModels()).map((model) => model.id)), + ); } catch { // Fall back to the previous authorization below. } @@ -857,16 +916,22 @@ export class ModelRegistry { if ((await this.currentPrivatePrimeAuthorizationFingerprint()) !== fingerprint) { return; } - if (authorizedIds) { - this.authorizedPrivatePrimeInferenceModelIds = authorizedIds; + if (authorizedModels) { + this.authorizedPrivatePrimeInferenceModels = authorizedModels; + this.authorizedPrivatePrimeInferenceModelIds = new Set(authorizedModels.map((model) => model.id)); this.authorizedPrivatePrimeInferenceTeamId = teamId; - this.writePrivatePrimeAuthorizationCache({ fingerprint, modelIds: authorizedIds, refreshedAt: Date.now() }); + this.reloadModelsAfterCatalogChange(); + this.writePrivatePrimeAuthorizationCache({ fingerprint, models: authorizedModels, refreshedAt: Date.now() }); } else if (teamId === previousTeamId) { this.authorizedPrivatePrimeInferenceModelIds = previousPrivateModelIds; + this.authorizedPrivatePrimeInferenceModels = previousPrivateModels; this.authorizedPrivatePrimeInferenceTeamId = teamId; + this.reloadModelsAfterCatalogChange(); } else { this.authorizedPrivatePrimeInferenceModelIds.clear(); + this.authorizedPrivatePrimeInferenceModels = []; this.authorizedPrivatePrimeInferenceTeamId = undefined; + this.reloadModelsAfterCatalogChange(); } } @@ -887,18 +952,25 @@ export class ModelRegistry { } const run = async () => { try { - const authorizedIds = await fetchAuthorizedPrivatePrimeInferenceModelIds( + const authorizedModels = await fetchAuthorizedPrivatePrimeInferenceModels( apiKey, teamHeaders, + new Set((this.livePrimeInferenceModels ?? this.bundledPrimeInferenceModels()).map((model) => model.id)), undefined, PRIVATE_PRIME_BACKGROUND_REFRESH_TIMEOUT_MS, ); if ((await this.currentPrivatePrimeAuthorizationFingerprint()) !== fingerprint) { return; } - this.authorizedPrivatePrimeInferenceModelIds = authorizedIds; + this.authorizedPrivatePrimeInferenceModels = authorizedModels; + this.authorizedPrivatePrimeInferenceModelIds = new Set(authorizedModels.map((model) => model.id)); this.authorizedPrivatePrimeInferenceTeamId = teamId; - this.writePrivatePrimeAuthorizationCache({ fingerprint, modelIds: authorizedIds, refreshedAt: Date.now() }); + this.reloadModelsAfterCatalogChange(); + this.writePrivatePrimeAuthorizationCache({ + fingerprint, + models: authorizedModels, + refreshedAt: Date.now(), + }); } catch { // Keep the cached authorization. } @@ -928,25 +1000,28 @@ export class ModelRegistry { private readPrivatePrimeAuthorizationCache(): PrivatePrimeAuthorizationCache | undefined { const cachePath = this.privatePrimeAuthorizationCachePath(); - if (!cachePath) { - return undefined; - } + if (!cachePath) return undefined; try { - const parsed = JSON.parse(readFileSync(cachePath, "utf8")) as Partial< - Omit & { modelIds: string[] } - >; + const parsed = JSON.parse(readFileSync(cachePath, "utf8")) as { + fingerprint?: unknown; + data?: unknown; + refreshedAt?: unknown; + }; if ( typeof parsed.fingerprint !== "string" || - !Array.isArray(parsed.modelIds) || + !Array.isArray(parsed.data) || typeof parsed.refreshedAt !== "number" ) { return undefined; } - return { - fingerprint: parsed.fingerprint, - modelIds: new Set(parsed.modelIds), - refreshedAt: parsed.refreshedAt, - }; + const entries = parsePrimeInferenceModelCatalog({ data: parsed.data }, { allowEmpty: true }).filter((entry) => + isPrivatePrimeInferenceModel({ provider: PRIME_INFERENCE_PROVIDER_ID, id: entry.id }), + ); + const models = buildPrimeInferenceModels(getPrivatePrimeInferenceModels(), entries, { + includePrivate: true, + minimumModels: 0, + }); + return { fingerprint: parsed.fingerprint, models: models ?? [], refreshedAt: parsed.refreshedAt }; } catch { return undefined; } @@ -954,11 +1029,29 @@ export class ModelRegistry { private writePrivatePrimeAuthorizationCache(cache: PrivatePrimeAuthorizationCache): void { const cachePath = this.privatePrimeAuthorizationCachePath(); - if (!cachePath) { - return; - } + if (!cachePath) return; + const data = cache.models.map((model) => ({ + id: model.id, + display_name: model.name, + pricing: { + input_usd_per_mtok: model.cost.input, + output_usd_per_mtok: model.cost.output, + cache_read_usd_per_mtok: model.cost.cacheRead, + cache_write_usd_per_mtok: model.cost.cacheWrite, + }, + specs: { + context_window: model.contextWindow, + max_output_tokens: model.maxTokens, + modalities: { input: model.input, output: ["text"] }, + supports_reasoning: model.reasoning, + }, + })); try { - writeFileAtomicSync(cachePath, JSON.stringify({ ...cache, modelIds: [...cache.modelIds] }), { mode: 0o600 }); + writeFileAtomicSync( + cachePath, + JSON.stringify({ fingerprint: cache.fingerprint, data, refreshedAt: cache.refreshedAt }), + { mode: 0o600 }, + ); } catch { // A failed cache write only requires a later refetch. } diff --git a/packages/coding-agent/src/core/prime-inference-model-catalog.ts b/packages/coding-agent/src/core/prime-inference-model-catalog.ts new file mode 100644 index 0000000000..8c9c2f6725 --- /dev/null +++ b/packages/coding-agent/src/core/prime-inference-model-catalog.ts @@ -0,0 +1,170 @@ +import { Buffer } from "node:buffer"; +import { existsSync, readFileSync } from "node:fs"; +import { + type Api, + isPrivatePrimeInferenceModelId, + type Model, + type OpenAICompletionsCompat, + type PrimeInferenceCatalogEntry, + parsePrimeInferenceModelCatalog, +} from "@earendil-works/pi-ai"; + +import { writeFileAtomicSync } from "../utils/atomic-file.js"; + +export const PRIME_INFERENCE_BASE_URL = "https://api.pinference.ai/api/v1"; +const FETCH_TIMEOUT_MS = 5_000; +const MAX_RESPONSE_BYTES = 2 * 1024 * 1024; +const MIN_CATALOG_COVERAGE = 0.5; +const pendingRefreshes = new Map[] | undefined>>(); + +const DEFAULT_COMPAT: OpenAICompletionsCompat = { + supportsStore: false, + supportsDeveloperRole: false, + // The endpoint does not yet describe reasoning controls. Do not send an + // unconfirmed reasoning_effort parameter for models without a bundled template. + supportsReasoningEffort: false, + maxTokensField: "max_tokens", + supportsStrictMode: false, +}; + +function cacheCosts(entry: PrimeInferenceCatalogEntry, template?: Model<"openai-completions">) { + const anthropic = entry.id.toLowerCase().startsWith("anthropic/"); + return { + cacheRead: entry.cacheRead ?? template?.cost.cacheRead ?? (anthropic ? entry.input * 0.1 : 0), + cacheWrite: entry.cacheWrite ?? template?.cost.cacheWrite ?? (anthropic ? entry.input * 1.25 : 0), + }; +} + +export function buildPrimeInferenceModels( + bundledModels: readonly Model<"openai-completions">[], + entries: readonly PrimeInferenceCatalogEntry[], + options: { includePrivate?: boolean; minimumModels?: number } = {}, +): Model<"openai-completions">[] | undefined { + const bundled = new Map(bundledModels.map((model) => [model.id.toLowerCase(), model])); + const models: Model<"openai-completions">[] = []; + for (const entry of entries) { + if (!options.includePrivate && isPrivatePrimeInferenceModelId(entry.id)) continue; + const template = bundled.get(entry.id.toLowerCase()); + if (!template && (!entry.contextWindow || !entry.maxTokens || entry.reasoning === undefined)) continue; + const contextWindow = entry.contextWindow ?? template?.contextWindow ?? 0; + const maxTokens = Math.min(entry.maxTokens ?? template?.maxTokens ?? 0, contextWindow); + models.push({ + id: entry.id, + name: entry.name ?? template?.name ?? entry.id, + api: "openai-completions", + provider: "prime-inference", + baseUrl: PRIME_INFERENCE_BASE_URL, + reasoning: entry.reasoning ?? template?.reasoning ?? false, + ...(template?.thinkingLevelMap ? { thinkingLevelMap: { ...template.thinkingLevelMap } } : {}), + input: (entry.vision ?? template?.input.includes("image")) ? ["text", "image"] : ["text"], + cost: { input: entry.input, output: entry.output, ...cacheCosts(entry, template) }, + contextWindow, + maxTokens, + ...(template?.featured ? { featured: true } : {}), + compat: structuredClone(template?.compat ?? DEFAULT_COMPAT), + }); + } + const minimumModels = options.minimumModels ?? Math.ceil(bundledModels.length * MIN_CATALOG_COVERAGE); + const coveredBundledModels = models.filter((model) => bundled.has(model.id.toLowerCase())).length; + return coveredBundledModels >= minimumModels ? models : undefined; +} + +export function mergePrimeInferenceModels( + bundledModels: readonly Model[], + livePrimeInferenceModels?: readonly Model<"openai-completions">[], +): Model[] { + if (!livePrimeInferenceModels) return [...bundledModels]; + return [...bundledModels.filter((model) => model.provider !== "prime-inference"), ...livePrimeInferenceModels]; +} + +export function readCachedPrimeInferenceModels( + cachePath: string, + bundledModels: readonly Model<"openai-completions">[], +): Model<"openai-completions">[] | undefined { + if (!existsSync(cachePath)) return undefined; + try { + return buildPrimeInferenceModels( + bundledModels, + parsePrimeInferenceModelCatalog(JSON.parse(readFileSync(cachePath, "utf8")) as unknown), + ); + } catch { + return undefined; + } +} + +function writeCache(cachePath: string, value: unknown): void { + try { + writeFileAtomicSync(cachePath, JSON.stringify(value), { mode: 0o600 }); + } catch { + // The bundled catalog remains available when the cache cannot be persisted. + } +} + +export class PrimeInferenceCatalogRequestError extends Error { + constructor(readonly status: number) { + super(`Prime Inference model catalog request failed with status ${status}`); + } +} + +async function readResponse(response: Response): Promise { + if (!response.ok) throw new PrimeInferenceCatalogRequestError(response.status); + const contentLength = Number(response.headers.get("content-length")); + if (Number.isFinite(contentLength) && contentLength > MAX_RESPONSE_BYTES) throw new Error("Response is too large"); + if (!response.body) throw new Error("Response body is empty"); + const reader = response.body.getReader(); + const chunks: Uint8Array[] = []; + let bytesRead = 0; + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + bytesRead += value.byteLength; + if (bytesRead > MAX_RESPONSE_BYTES) { + await reader.cancel().catch(() => {}); + throw new Error("Response is too large"); + } + chunks.push(value); + } + } finally { + reader.releaseLock(); + } + return JSON.parse(Buffer.concat(chunks, bytesRead).toString("utf8")) as unknown; +} + +export async function fetchPrimeInferenceModelCatalog( + options: { fetchFn?: typeof fetch; headers?: Record; timeoutMs?: number; allowEmpty?: boolean } = {}, +): Promise<{ payload: unknown; entries: PrimeInferenceCatalogEntry[] }> { + const response = await (options.fetchFn ?? fetch)(`${PRIME_INFERENCE_BASE_URL}/models`, { + headers: { accept: "application/json", ...options.headers }, + signal: AbortSignal.timeout(options.timeoutMs ?? FETCH_TIMEOUT_MS), + }); + const payload = await readResponse(response); + return { payload, entries: parsePrimeInferenceModelCatalog(payload, { allowEmpty: options.allowEmpty }) }; +} + +export async function refreshPrimeInferenceModels( + cachePath: string, + bundledModels: readonly Model<"openai-completions">[], + options: { fetchFn?: typeof fetch; offline?: boolean } = {}, +): Promise[] | undefined> { + const cached = readCachedPrimeInferenceModels(cachePath, bundledModels); + if (options.offline) return cached; + const existing = pendingRefreshes.get(cachePath); + if (existing) return existing; + const promise = (async () => { + try { + const { payload, entries } = await fetchPrimeInferenceModelCatalog({ fetchFn: options.fetchFn }); + const models = buildPrimeInferenceModels(bundledModels, entries); + if (!models) return cached; + writeCache(cachePath, payload); + return models; + } catch { + return cached; + } + })(); + pendingRefreshes.set(cachePath, promise); + void promise.finally(() => { + if (pendingRefreshes.get(cachePath) === promise) pendingRefreshes.delete(cachePath); + }); + return promise; +} diff --git a/packages/coding-agent/src/core/prime-inference-models.ts b/packages/coding-agent/src/core/prime-inference-models.ts index c9af39adfa..15c229c3fb 100644 --- a/packages/coding-agent/src/core/prime-inference-models.ts +++ b/packages/coding-agent/src/core/prime-inference-models.ts @@ -1,6 +1,13 @@ -import type { Model } from "@earendil-works/pi-ai"; +import { isPrivatePrimeInferenceModelId, type Model } from "@earendil-works/pi-ai"; +import { + buildPrimeInferenceModels, + fetchPrimeInferenceModelCatalog, + PRIME_INFERENCE_BASE_URL, + PrimeInferenceCatalogRequestError, +} from "./prime-inference-model-catalog.js"; + +export { PRIME_INFERENCE_BASE_URL }; -export const PRIME_INFERENCE_BASE_URL = "https://api.pinference.ai/api/v1"; const PRIVATE_MODEL_REFRESH_TIMEOUT_MS = 10_000; const PRIVATE_PRIME_INFERENCE_MODELS: readonly Model<"openai-completions">[] = [ @@ -24,7 +31,7 @@ const PRIVATE_PRIME_INFERENCE_MODELS: readonly Model<"openai-completions">[] = [ ]; export function isPrivatePrimeInferenceModel(model: Pick, "provider" | "id">): boolean { - return model.provider === "prime-inference" && model.id.startsWith("internal/"); + return model.provider === "prime-inference" && isPrivatePrimeInferenceModelId(model.id); } export function getPrivatePrimeInferenceModels(): Model<"openai-completions">[] { @@ -36,42 +43,46 @@ export function getPrivatePrimeInferenceModels(): Model<"openai-completions">[] })); } -export async function fetchAuthorizedPrivatePrimeInferenceModelIds( +export async function fetchAuthorizedPrivatePrimeInferenceModels( apiKey: string, teamHeaders: Record, + publicModelIds: ReadonlySet, fetchFn: typeof fetch = fetch, timeoutMs: number = PRIVATE_MODEL_REFRESH_TIMEOUT_MS, -): Promise> { - if (!teamHeaders["X-Prime-Team-ID"]) { - return new Set(); - } - - const response = await fetchFn(`${PRIME_INFERENCE_BASE_URL}/models`, { - headers: { - Authorization: `Bearer ${apiKey}`, - ...teamHeaders, - }, - signal: AbortSignal.timeout(timeoutMs), - }); - if (response.status === 401 || response.status === 403) { - return new Set(); +): Promise[]> { + if (!teamHeaders["X-Prime-Team-ID"]) return []; + try { + const { payload, entries } = await fetchPrimeInferenceModelCatalog({ + fetchFn, + timeoutMs, + allowEmpty: true, + headers: { ...teamHeaders, Authorization: `Bearer ${apiKey}` }, + }); + const publicIds = new Set([...publicModelIds].map((id) => id.toLowerCase())); + const bundledPrivateModels = getPrivatePrimeInferenceModels(); + const bundledById = new Map(bundledPrivateModels.map((model) => [model.id.toLowerCase(), model])); + const entriesById = new Map(entries.map((entry) => [entry.id.toLowerCase(), entry])); + const data = + payload && typeof payload === "object" && "data" in payload && Array.isArray(payload.data) ? payload.data : []; + const privateEntries = data.flatMap((item) => { + if (!item || typeof item !== "object" || !("id" in item) || typeof item.id !== "string") return []; + const id = item.id.toLowerCase(); + if (publicIds.has(id) || !isPrivatePrimeInferenceModelId(id)) return []; + const parsed = entriesById.get(id); + if (parsed) return [parsed]; + const template = bundledById.get(id); + return template ? [{ id: item.id, input: template.cost.input, output: template.cost.output }] : []; + }); + return ( + buildPrimeInferenceModels(bundledPrivateModels, privateEntries, { + includePrivate: true, + minimumModels: 0, + }) ?? [] + ); + } catch (error) { + if (error instanceof PrimeInferenceCatalogRequestError && (error.status === 401 || error.status === 403)) { + return []; + } + throw error; } - if (!response.ok) { - throw new Error(`Prime Inference model catalog request failed with status ${response.status}`); - } - - const payload = (await response.json()) as unknown; - if (!payload || typeof payload !== "object" || !("data" in payload) || !Array.isArray(payload.data)) { - throw new Error("Prime Inference model catalog response is invalid"); - } - - const knownPrivateIds = new Set(PRIVATE_PRIME_INFERENCE_MODELS.map((model) => model.id)); - return new Set( - payload.data.flatMap((entry) => { - if (!entry || typeof entry !== "object" || !("id" in entry) || typeof entry.id !== "string") { - return []; - } - return knownPrivateIds.has(entry.id) ? [entry.id] : []; - }), - ); } diff --git a/packages/coding-agent/src/core/prompts/rlm.ts b/packages/coding-agent/src/core/prompts/rlm.ts index ef610268e6..d7be88bd47 100644 --- a/packages/coding-agent/src/core/prompts/rlm.ts +++ b/packages/coding-agent/src/core/prompts/rlm.ts @@ -12,7 +12,7 @@ export interface RlmPromptOptions { } const LONG_RUNNING_WORK_PROMPT = [ - "For slow or independently completing work, use a nonblocking control loop: start the work, record its handle or output location, then end your turn. Read the result on a later turn or when a reply arrives.", + "For slow or independently completing work, use a nonblocking control loop: start the work, record its handle or output location, then end your turn. A `bash()` handle left running beyond its creating cell sends a completion follow-up; when it arrives, inspect the saved handle and continue.", "When delegation is available and useful, assign independent substantive tasks to separate workers. Start independent workers without waiting for each one sequentially, and let them run in parallel.", "Do not keep the turn open by polling with `time.sleep()` or shell `sleep`, and do not replace polling with a long blocking `await`. Await only the short operation needed to start work or inspect a result that is already available; otherwise end the turn.", ].join("\n"); @@ -146,6 +146,13 @@ export function buildRlmPrompt(options: RlmPromptOptions): string { ); } + if (depth === 0 && hasIpython) { + parts.push( + "", + "From a daemon-backed depth-0 session, use `await rlm.create_session('task', name='researcher')` to start a separate top-level session. The call returns after the daemon creates the session and accepts its first prompt. Inline and nested sessions cannot use it. `rlm(...)` still creates a child.", + ); + } + if (allowRecursion && hasIpython) { parts.push( "", diff --git a/packages/coding-agent/src/core/rlm-runtime.ts b/packages/coding-agent/src/core/rlm-runtime.ts index 470224cc53..837a6d77ad 100644 --- a/packages/coding-agent/src/core/rlm-runtime.ts +++ b/packages/coding-agent/src/core/rlm-runtime.ts @@ -12,6 +12,19 @@ export interface RlmRunRequest { cellSourceCode?: string; } +interface RlmCreateSessionRequest { + prompt: string; + kwargs: Record; +} + +export interface RlmCreateSessionResult { + active_session_id: string; + session_id: string; + name: string; + session_file: string; + model: string; +} + export interface RlmSpawnHandle { rlm_child_id: string; name: string; @@ -51,6 +64,15 @@ export interface RlmFindModelsResult { } export type RlmRunHandler = (request: RlmRunRequest) => Promise>; +type RlmCreateSessionHandler = (request: RlmCreateSessionRequest) => Promise; + +interface AsyncBashCompletionRequest { + pid: number; + command: string; + exitCode: number; +} + +type AsyncBashCompletionHandler = (request: AsyncBashCompletionRequest) => void | Promise; export type RlmListSubagentsHandler = () => RlmListSubagentsResult | Promise; export type RlmDeleteSubagentHandler = (target: string) => Promise; export type RlmFindModelsHandler = (query: string, limit: number) => RlmFindModelsResult | Promise; @@ -59,47 +81,50 @@ const RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH = 64; export const DEFAULT_RLM_MODEL_SEARCH_LIMIT = 8; export const MAX_RLM_MODEL_SEARCH_LIMIT = 20; -export function normalizeRequestedRlmSubagentSessionName(value: unknown): string | undefined { +export function normalizeRequestedRlmSubagentSessionName(value: unknown, operation = "rlm.run"): string | undefined { if (value === undefined) { return undefined; } if (typeof value !== "string") { - throw new Error("rlm.run name must be a string"); + throw new Error(`${operation} name must be a string`); } const name = value.trim(); if (!name) { - throw new Error("rlm.run name must not be empty"); + throw new Error(`${operation} name must not be empty`); } if (name.length > RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH) { - throw new Error(`rlm.run name must be at most ${RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH} characters`); + throw new Error(`${operation} name must be at most ${RLM_SUBAGENT_SESSION_NAME_MAX_LENGTH} characters`); } return name; } -export function normalizeRequestedRlmSubagentThinkingLevel(value: unknown): ThinkingLevel | undefined { +export function normalizeRequestedRlmSubagentThinkingLevel( + value: unknown, + operation = "rlm.run", +): ThinkingLevel | undefined { if (value === undefined) { return undefined; } if (typeof value !== "string") { - throw new Error("rlm.run thinking must be a string"); + throw new Error(`${operation} thinking must be a string`); } const level = value.trim().toLowerCase(); if (!THINKING_LEVELS.includes(level as ThinkingLevel)) { - throw new Error(`rlm.run thinking must be one of: ${THINKING_LEVELS.join(", ")}`); + throw new Error(`${operation} thinking must be one of: ${THINKING_LEVELS.join(", ")}`); } return level as ThinkingLevel; } -export function normalizeRequestedRlmSubagentModel(value: unknown): string | undefined { +export function normalizeRequestedRlmSubagentModel(value: unknown, operation = "rlm.run"): string | undefined { if (value === undefined) { return undefined; } if (typeof value !== "string") { - throw new Error("rlm.run model must be a string"); + throw new Error(`${operation} model must be a string`); } const model = value.trim(); if (!model) { - throw new Error("rlm.run model must not be empty"); + throw new Error(`${operation} model must not be empty`); } return model; } @@ -161,6 +186,17 @@ export function findRlmModelMatches(query: string, models: Model[], limit: })); } +export function createRlmCreateSessionHostHandler(handler: RlmCreateSessionHandler): HostRequestHandler { + return async (payload) => { + if (typeof payload.prompt !== "string") { + throw new Error("rlm.create_session prompt must be a string"); + } + const kwargs = isRecord(payload.kwargs) ? payload.kwargs : {}; + const result = await handler({ prompt: payload.prompt, kwargs }); + return result as unknown as Record; + }; +} + /** Adapt an RlmRunHandler into the typed `rlm.run` kernel host handler. */ export function createRlmRunHostHandler(handler: RlmRunHandler): HostRequestHandler { return async (payload) => { @@ -178,6 +214,24 @@ export function createRlmRunHostHandler(handler: RlmRunHandler): HostRequestHand }; } +/** Adapt detached kernel bash completions into a validated host notification. */ +export function createAsyncBashCompletionHostHandler(handler: AsyncBashCompletionHandler): HostRequestHandler { + return async (payload) => { + const { pid, command, exitCode } = payload; + if (typeof pid !== "number" || !Number.isInteger(pid) || pid <= 0) { + throw new Error("bash.completed pid must be a positive integer"); + } + if (typeof command !== "string" || !command) { + throw new Error("bash.completed command must be a non-empty string"); + } + if (typeof exitCode !== "number" || !Number.isInteger(exitCode)) { + throw new Error("bash.completed exitCode must be an integer"); + } + await handler({ pid, command, exitCode }); + return {}; + }; +} + /** Search a bounded authenticated model catalog without adding it to the system prompt. */ export function createRlmFindModelsHostHandler(handler: RlmFindModelsHandler): HostRequestHandler { return async (payload) => { @@ -241,8 +295,17 @@ export interface CreateRlmSubagentRuntimeOptions { onSessionPublished?: (session: AgentSession) => void; } +export interface CreateRlmRootSessionOptions { + prompt: string; + sessionName?: string; + cwd: string; + model: Model; + thinkingLevel: ThinkingLevel; +} + export interface SubagentRuntimeHost { createRlmSubagentRuntime(options: CreateRlmSubagentRuntimeOptions): Promise; + createRlmRootSession?(options: CreateRlmRootSessionOptions): Promise; /** Persist host-owned completion before the child becomes passivation-eligible. */ completeRlmSubagentRuntime?(childId: string, session: AgentSession): boolean; /** Release a host-owned child after its detached initial task settles. */ diff --git a/packages/coding-agent/src/core/tools/ipython.ts b/packages/coding-agent/src/core/tools/ipython.ts index acab77b797..bae3ba41c1 100644 --- a/packages/coding-agent/src/core/tools/ipython.ts +++ b/packages/coding-agent/src/core/tools/ipython.ts @@ -53,6 +53,9 @@ except Exception as _prime_agent_rlm_error: async def find_models(self, query="", limit=8): self._raise_missing() + async def create_session(self, prompt, **kwargs): + self._raise_missing() + async def list_subagents(self): self._raise_missing() diff --git a/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts b/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts index 450ac81ec8..b867baaf06 100644 --- a/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts +++ b/packages/coding-agent/src/modes/agents-view/agents-view-mode.ts @@ -45,7 +45,6 @@ import { listDaemonSavedSessions, renameDaemonSavedSession, } from "../daemon/saved-session-catalog.js"; -import { formatTokenCount } from "../interactive/agent-activity.js"; import { CustomEditor } from "../interactive/components/custom-editor.js"; import { keyText } from "../interactive/components/keybinding-hints.js"; import { BrandSplashHeader, InteractiveMode } from "../interactive/interactive-mode.js"; @@ -84,8 +83,6 @@ import { getAgentsViewSummaryIdentity as getSummaryIdentity, getUnifiedSessionAncestorSessionIds, hasUnifiedSessionChildren, - isEmptyAgentsViewSession, - isSubagentSummary, migrateAgentsViewIdentitySet, reconcileUnifiedSessions, resolveAgentsViewLeftResult, @@ -163,6 +160,7 @@ export type AgentsViewPersistentState = { pendingExpandedAncestorSessionIds?: string[]; expandedSubagentParents?: Set; programShownParents?: Set; + inactiveExpanded?: boolean; statusMessage?: string; // Gathered once and reused across agents-view instances so the notices survive // re-entry and render the moment they resolve, even if the first view was left early. @@ -280,12 +278,16 @@ export function createInitialAgentsViewPersistentState( options: Pick, ): AgentsViewPersistentState { const initialSession = options.initialSession; + // A scoped view excludes its root from its own rows, so anchoring the + // selection on the entered-from chat could never resolve there and would + // only arm the pending-anchor state for the whole catalog scan. + const seedSelection = initialSession && !options.initialScopeKey; return { - ...(initialSession + ...(initialSession ? { backSession: initialSession } : {}), + ...(seedSelection ? { selectedRowIdentity: getSummaryIdentity(initialSession), selectedSessionKey: getAgentsViewSelectionKey(initialSession), - backSession: initialSession, } : {}), ...(options.initialScopeKey @@ -664,6 +666,8 @@ export class AgentsViewMode implements Component, Focusable { private deleteConfirmTimer: ReturnType | undefined; private workingIconFrame = 0; private rows: AgentsViewRow[] = []; + private allRows: AgentsViewRow[] = []; + private showActions = false; private lastListedSummaries: SessionSummary[] = []; private lastVisibleSummaries: SessionSummary[] = []; private savedSessions: AgentConnectionSavedSessionInfo[] = []; @@ -919,6 +923,12 @@ export class AgentsViewMode implements Component, Focusable { handleInput(data: string): void { this.clearStickyStatusMessage(); + if (this.showActions) { + this.showActions = false; + this.ui.requestRender(); + if (this.keybindings.matches(data, "app.shortcuts") || this.keybindings.matches(data, "tui.select.cancel")) + return; + } if (this.renameTarget) { if (this.keybindings.matches(data, "tui.select.cancel")) { this.exitRenameMode(); @@ -963,6 +973,25 @@ export class AgentsViewMode implements Component, Focusable { this.cycleProgramForSelected(); return; } + if (!this.replyTarget && this.editor.getText().length === 0) { + if (this.keybindings.matches(data, "app.shortcuts")) { + this.showActions = !this.showActions; + this.ui.requestRender(); + return; + } + if (this.keybindings.matches(data, "app.agents.inactiveCollapse")) { + this.persistentState.inactiveExpanded = !this.persistentState.inactiveExpanded; + this.rebuildRows(); + this.syncSelectedRowState(); + this.ui.requestRender(); + return; + } + if (this.keybindings.matches(data, "app.agents.expand")) { + const row = this.rows[this.selectedIndex]; + if (row && row.descendantCount > 0) this.toggleSubagentList(row); + return; + } + } if (!this.replyTarget && this.keybindings.matches(data, "app.agents.open")) { if (this.editor.getText().length === 0 || this.isSearchCursorAtEnd()) { this.openSelected(); @@ -1277,9 +1306,9 @@ export class AgentsViewMode implements Component, Focusable { this.persistentState.query = this.editor.getText(); this.armSavedSearchFetch(); this.rebuildRows(); - // Typing must not claim the visible fallback row while the restored - // anchor is still waiting for its catalog row. - if (!this.selectionAnchorPending) this.syncSelectedRowState(); + // Searching is explicit user intent: claim the visible row as the new + // anchor even if a remembered one is still waiting for its catalog row. + this.syncSelectedRowState(); this.ui.requestRender(); } @@ -1291,7 +1320,7 @@ export class AgentsViewMode implements Component, Focusable { /** Rebuild rows from the last fetched summaries, keeping selection on the same row. */ private rebuildRows(): void { const selectedIdentity = this.rows[this.selectedIndex]?.identity; - this.rows = buildAgentsViewRows( + this.allRows = buildAgentsViewRows( this.getFilteredRecords(), this.expandedSubagentParents, this.programShownParents, @@ -1299,6 +1328,12 @@ export class AgentsViewMode implements Component, Focusable { computeRecursiveRollups(this.unifiedRecords, this.unifiedIndex), this.anchorSessionId, ); + this.rows = compactSessionRows( + this.allRows, + this.persistentState.inactiveExpanded === true || + ((this.replyTarget || this.renameTarget ? this.actionModeSearchQuery : this.editor.getText()) ?? "").trim() + .length > 0, + ); const index = selectedIdentity === undefined ? -1 : this.rows.findIndex((row) => row.identity === selectedIdentity); if (index >= 0) { @@ -1378,14 +1413,6 @@ export class AgentsViewMode implements Component, Focusable { if (!row?.selectable || this.isPendingDeleteRow(row)) { return; } - if (this.selectionAnchorPending) { - this.setStatusMessage("Waiting for the selected session to load"); - return; - } - if (row.kind === "subagent-summary") { - this.toggleSubagentList(row); - return; - } if (row.kind === "subagent") { this.openSelectedSubagent(row); return; @@ -1406,16 +1433,12 @@ export class AgentsViewMode implements Component, Focusable { } private toggleSubagentList(row: AgentsViewRow): void { - if (!row.parentIdentity) { - return; - } - if (this.expandedSubagentParents.has(row.parentIdentity)) { - this.expandedSubagentParents.delete(row.parentIdentity); - // A collapsed agent's revealed program collapses with it, so reopening - // starts from the hidden state rather than a stale reveal. - this.programShownParents.delete(row.parentIdentity); + const target = row.identity; + if (this.expandedSubagentParents.has(target)) { + this.expandedSubagentParents.delete(target); + this.programShownParents.delete(target); } else { - this.expandedSubagentParents.add(row.parentIdentity); + this.expandedSubagentParents.add(target); } this.rebuildRows(); this.syncSelectedRowState(); @@ -1454,7 +1477,7 @@ export class AgentsViewMode implements Component, Focusable { /** Whether any subagent under the given agent identity carries spawn code. */ private targetHasSpawnCode(target: string): boolean { - for (const row of this.rows) { + for (const row of this.allRows) { if (row.parentIdentity !== target) { continue; } @@ -1468,16 +1491,6 @@ export class AgentsViewMode implements Component, Focusable { return false; } - /** True when the selected row exposes the "show program" affordance. */ - private selectedRowCanShowProgram(): boolean { - const row = this.rows[this.selectedIndex]; - if (!row) { - return false; - } - const target = row.kind === "agent" ? row.identity : row.parentIdentity; - return target !== undefined && this.targetHasSpawnCode(target); - } - private openSelectedSubagent(row: AgentsViewRow): void { const expandedAncestorSessionIds = this.collectSubagentAncestorSessionIds(row); if (row.summary.activeSessionId || row.summary.sessionFile) { @@ -2179,7 +2192,7 @@ export class AgentsViewMode implements Component, Focusable { } } this.scopedRecords = scopeToSessionSubtree(this.unifiedRecords, this.scopeKey, this.unifiedIndex); - this.rows = buildAgentsViewRows( + this.allRows = buildAgentsViewRows( this.getFilteredRecords(), this.expandedSubagentParents, this.programShownParents, @@ -2187,6 +2200,12 @@ export class AgentsViewMode implements Component, Focusable { computeRecursiveRollups(this.unifiedRecords, this.unifiedIndex), this.anchorSessionId, ); + this.rows = compactSessionRows( + this.allRows, + this.persistentState.inactiveExpanded === true || + ((this.replyTarget || this.renameTarget ? this.actionModeSearchQuery : this.editor.getText()) ?? "").trim() + .length > 0, + ); this.applyPendingAncestorExpansion(); this.restoreSelection(); this.ui.requestRender(); @@ -2479,151 +2498,136 @@ export class AgentsViewMode implements Component, Focusable { } private getAgentCountsText(): string { - const counts = countRowsBySection(this.rows); + const counts = countRowsBySection(this.allRows); return `${counts.running} running, ${counts.idle} idle, ${counts.inactive} inactive`; } private renderSessionRows(width: number, maxRows: number): string[] { - if (maxRows <= 0) { - return []; + if (maxRows <= 0) return []; + if (this.showActions) return this.renderActions(width).slice(0, maxRows); + const layout = buildCompactAgentsViewLayout(this.rows, width); + const displayItems: DisplayItem[] = []; + const counts = countRowsBySection(this.allRows.length > 0 ? this.allRows : this.rows); + for (const section of ["running", "idle", "inactive"] as const) { + if (counts[section] === 0) continue; + if (displayItems.length > 0) displayItems.push({ type: "spacer" }); + displayItems.push({ type: "heading", section }); + for (const row of getDisplayRowsForSection(this.rows, section)) { + displayItems.push({ type: "row", row }); + if ( + (row.kind === "agent" || row.kind === "subagent") && + row.runningSubagentCount > 0 && + !this.expandedSubagentParents.has(row.identity) + ) { + displayItems.push({ type: "running-subagents", row }); + } + } } - if (this.rows.length === 0) { - const emptyLegend = buildAgentsViewUsageLayout([]).legends.get("running") ?? ""; - return [ - this.renderSectionHeading("running", width, emptyLegend), - theme.fg("dim", " No sessions match your search."), - ].slice(0, maxRows); + if (displayItems.length === 0) { + return [theme.fg("dim", "No sessions match your search.")]; } - - const displayItems = buildDisplayItems(this.rows); - const usageLayout = buildAgentsViewUsageLayout(this.rows); + // Reserve the shared column header before calculating the selection viewport. + const headerRows = maxRows > 1 ? 1 : 0; + const visibleRows = maxRows - headerRows; const selectedIdentity = this.rows[this.selectedIndex]?.identity; const selectedDisplayIndex = displayItems.findIndex( (item) => item.type === "row" && item.row.identity === selectedIdentity, ); - const visibleRows = Math.min(maxRows, this.visibleListRows()); const start = Math.max( 0, Math.min(displayItems.length - visibleRows, selectedDisplayIndex - Math.floor(visibleRows / 2)), ); - const showLeadingEllipsis = start > 0; - let showTrailingEllipsis = start + visibleRows < displayItems.length; - if ((showLeadingEllipsis ? 1 : 0) + (showTrailingEllipsis ? 1 : 0) >= visibleRows) { - showTrailingEllipsis = false; - } - const contentVisibleRows = Math.max( - 0, - visibleRows - (showLeadingEllipsis ? 1 : 0) - (showTrailingEllipsis ? 1 : 0), - ); - // The prepended ellipsis consumes a viewport line; shift the window down - // so a selection at the very end is not pushed out of the slice. - const sliceStart = - selectedDisplayIndex >= start + contentVisibleRows ? selectedDisplayIndex - contentVisibleRows + 1 : start; - const visibleItems = displayItems.slice(sliceStart, sliceStart + contentVisibleRows); - const lines = visibleItems.map((item) => { - if (item.type === "spacer") { - return ""; + const showLeadingEllipsis = start > 0 && visibleRows > 1; + const showTrailingEllipsis = start + visibleRows < displayItems.length && visibleRows > 2; + const contentRows = visibleRows - Number(showLeadingEllipsis) - Number(showTrailingEllipsis); + const sliceStart = selectedDisplayIndex >= start + contentRows ? selectedDisplayIndex - contentRows + 1 : start; + const lines = displayItems.slice(sliceStart, sliceStart + contentRows).map((item) => { + if (item.type === "spacer") return ""; + if (item.type === "running-subagents") { + const count = item.row.runningSubagentCount; + const indent = " ".repeat(item.row.depth + 1); + return theme.fg( + "success", + truncateToWidth(`${indent}${count} subagent${count === 1 ? "" : "s"} running`, width), + ); } if (item.type === "heading") { - return this.renderSectionHeading(item.section, width, usageLayout.legends.get(item.section) ?? ""); - } - if (item.type === "empty") { - return theme.fg("dim", " No agents"); + const collapsed = + item.section === "inactive" && !this.rows.some((row) => row.depth === 0 && row.section === "inactive"); + const prefix = item.section === "inactive" ? `${collapsed ? "▸" : "▾"} ` : ""; + const hint = item.section === "inactive" ? ` · ${keyText("app.agents.inactiveCollapse")}` : ""; + return theme.bold( + truncateToWidth(`${prefix}${sectionTitle(item.section)} (${counts[item.section]})${hint}`, width), + ); } - return this.renderRow(item.row, width, usageLayout.details); + return this.renderRow(item.row, width, layout); }); - if (showLeadingEllipsis) { - lines.unshift(theme.fg("dim", " ...")); - } - if (showTrailingEllipsis) { - lines.push(theme.fg("dim", " ...")); - } + if (showLeadingEllipsis) lines.unshift(theme.fg("dim", " ...")); + if (showTrailingEllipsis) lines.push(theme.fg("dim", " ...")); + if (headerRows > 0) lines.unshift(theme.fg("muted", layout.legend)); return lines; } private renderRow( row: AgentsViewRow, width: number, - rowDetails: ReadonlyMap = buildAgentsViewUsageLayout([row]).details, + layout: AgentsViewUsageLayout = buildCompactAgentsViewLayout(this.rows.length > 0 ? this.rows : [row], width), ): string { const selected = row.selectable && row.identity === this.rows[this.selectedIndex]?.identity; const markRow = (line: string): string => (selected ? `${SELECTED_ROW_MARKER}${line}` : line); - if (row.kind === "subagent-code") { - return this.renderCodeRow(row); - } - if (row.kind === "subagent-summary") { - const indent = " ".repeat(row.depth); - const hint = row.hasSpawnCode ? theme.fg("dim", ` · ${keyText("app.agents.program")} show program`) : ""; - const titleColor = row.runningSubagentCount > 0 ? ("success" as const) : ("dim" as const); - const label = `${theme.fg(titleColor, `${row.expanded ? "▾" : "▸"} ${row.title}`)}${hint}`; - const line = padLine(truncateToWidth(`${indent}${label}`, width, ""), width); - return markRow(line); - } + if (row.kind === "subagent-code") return this.renderCodeRow(row); const pendingDelete = row.kind === "agent" && this.isPendingDeleteRow(row); const pendingKill = row.kind === "subagent" && this.isPendingKillSubagentRow(row); - const rawIcon = this.getRowIcon(row.section); - const icon = this.formatRowIcon(row.section, rawIcon); - const indent = " ".repeat(row.depth); - const details = rowDetails.get(row.identity) ?? ""; - const detailsWidth = Math.max(10, visibleWidth(details)); - const heartbeatBadge = !pendingDelete && !pendingKill ? formatHeartbeatBadge(row.heartbeat) : ""; - const heartbeatPausedOnly = (row.heartbeat?.activeCount ?? 0) < 1; - const heartbeatCell = heartbeatBadge ? theme.fg(heartbeatPausedOnly ? "dim" : "error", heartbeatBadge) : ""; - const heartbeatWidth = visibleWidth(heartbeatBadge); - const titleWidth = Math.max( - 0, - width - - visibleWidth(indent) - - visibleWidth(rawIcon) - - detailsWidth - - 2 - - (heartbeatWidth > 0 ? heartbeatWidth + 1 : 0), - ); - const armedHeartbeat = row.summary.hasActiveHeartbeat === true || (row.heartbeat?.activeCount ?? 0) > 0; - const heartbeatWarning = armedHeartbeat ? "has an armed heartbeat — " : ""; - const title = pendingDelete - ? `${heartbeatWarning}${this.getPendingDeleteTitle()}` - : pendingKill - ? `${heartbeatWarning}${keyText("app.agents.delete")} again to ${hasLiveWork(row) ? "stop" : "delete"}` - : styleRowTitle(row); - // Keep stable model information ahead of the variable summary so narrow rows truncate the summary first. - const summaryText = !pendingDelete && !pendingKill ? row.summary.summary : undefined; - const modelLabel = - isSubagentSummary(row.summary) && !pendingDelete && !pendingKill && row.summary.model - ? `${row.summary.model.provider}/${row.summary.model.id}${row.summary.thinkingLevel && row.summary.thinkingLevel !== "off" ? `:${row.summary.thinkingLevel}` : ""}` - : undefined; - const statusLabel = - !pendingDelete && - !pendingKill && - (row.summary.statusLabel !== undefined || row.summary.lastHeardFromAt !== undefined) + const details = layout.details.get(row.identity) ?? ""; + if (pendingDelete || pendingKill) { + const armed = row.summary.hasActiveHeartbeat === true || (row.heartbeat?.activeCount ?? 0) > 0; + const title = + (armed ? "has an armed heartbeat — " : "") + + (pendingDelete + ? this.getPendingDeleteTitle() + : `${keyText("app.agents.delete")} again to ${hasLiveWork(row) ? "stop" : "delete"}`); + return markRow(formatTableCell(theme.fg("error", title), width)); + } + const icon = this.formatRowIcon(row.section, this.getRowIcon(row.section)); + const expand = row.descendantCount > 0 ? (this.expandedSubagentParents.has(row.identity) ? "▾" : "▸") : " "; + const badge = formatHeartbeatBadge(row.heartbeat); + const heartbeat = badge ? `${theme.fg((row.heartbeat?.activeCount ?? 0) > 0 ? "error" : "dim", badge)} ` : ""; + const title = `${" ".repeat(row.depth)}${icon}${expand} ${heartbeat}${styleRowTitle(row)}`; + const status = + row.summary.statusLabel !== undefined || row.summary.lastHeardFromAt !== undefined ? row.statusLabel : undefined; - const suffixes = [statusLabel, modelLabel, summaryText].filter( - (suffix): suffix is string => suffix !== undefined && suffix.length > 0, - ); - const titleContent = suffixes.length > 0 ? `${title} ${theme.fg("dim", `· ${suffixes.join(" · ")}`)}` : title; - const titleCell = formatTableCell(titleContent, titleWidth); + const activity = [status, row.summary.summary].filter(Boolean).join(" · "); const cells = [ - icon, - pendingDelete || pendingKill ? theme.fg("error", titleCell) : titleCell, - formatRightTableCell(details, detailsWidth), + formatTableCell(title, layout.nameWidth), + formatTableCell(theme.fg("muted", formatSessionModel(row.summary)), layout.modelWidth), ]; - const base = `${indent}${cells[0]} ${heartbeatCell ? `${heartbeatCell} ` : ""}${cells[1]} ${cells[2]}`; - const line = padLine(truncateToWidth(base, width, ""), width); - return markRow(line); + if (layout.activityWidth > 0) cells.push(formatTableCell(theme.fg("dim", activity), layout.activityWidth)); + cells.push(details); + return markRow(formatTableCell(cells.join(" "), width)); } - // Bold like the section title so the legend reads as part of the header line. - // The legend is right-aligned to the same edge as the row details cells, so - // its columns sit exactly above the row columns. - private renderSectionHeading(section: AgentsViewSection, width: number, legend: string): string { - const counts = countRowsBySection(this.rows); - const title = `${sectionTitle(section)} (${counts[section]})`; - const gap = width - visibleWidth(title) - visibleWidth(legend); - if (legend.length === 0 || gap < 2) { - return theme.bold(truncateToWidth(title, width, "")); + private renderActions(width: number): string[] { + const row = this.rows[this.selectedIndex]; + const actions = [ + `${keyText("tui.select.confirm")} open ${keyText("app.agents.open")} open ${keyText("app.agents.new")} new`, + `${keyText("app.agents.expand")} expand/collapse subagents ${keyText("app.agents.program")} program`, + `${keyText("app.agents.inactiveCollapse")} show/hide inactive ${keyText("app.shortcuts")} close actions`, + `${keyText("app.agents.reply")} reply/resume ${keyText("app.agents.rename")} rename ${keyText("app.agents.delete")} stop/delete`, + ]; + if (row) { + const model = row.summary.model; + const usage = row.summary.usage; + actions.push( + "", + row.title, + `Model: ${model ? `${model.provider}/${model.id}` : "unknown"}${row.summary.thinkingLevel ? ` · ${row.summary.thinkingLevel}` : ""}`, + `Directory: ${row.summary.cwd}`, + `Tokens: ${usage?.inputTokens ?? 0} in · ${usage?.outputTokens ?? 0} out`, + `Cost: $${(usage?.cost ?? 0).toFixed(2)} session · $${row.recursiveCost.toFixed(2)} including subagents`, + ); } - return `${theme.bold(title)}${" ".repeat(gap)}${theme.bold(legend)}`; + return actions.flatMap((line) => wrapTextWithAnsi(theme.fg("muted", line), width)); } // Spawn-code rows are read-only context. They render deemphasized — muted @@ -2705,32 +2709,7 @@ export class AgentsViewMode implements Component, Focusable { if (this.replyTarget) { return truncateToWidth(theme.fg("muted", this.renderReplyComposerHints()), width); } - // Replying is reserved for top-level agents; subagents can be stopped or deleted. - const selectedRow = this.rows[this.selectedIndex]; - const selectedAgent = selectedRow?.kind === "agent"; - const selectedSubagent = selectedRow?.kind === "subagent"; - const selectedSummary = selectedRow?.kind === "subagent-summary"; - const hints = [ - `${keyText("tui.select.up")}/${keyText("tui.select.down")} move`, - selectedSummary - ? `${keyText("tui.select.confirm")} ${selectedRow?.expanded ? "collapse" : "expand"}` - : `${keyText("tui.select.confirm")} open`, - selectedSummary ? undefined : `${keyText("app.agents.open")} open`, - selectedAgent - ? `${keyText("app.agents.reply")} ${selectedRow?.section === "inactive" ? "resume" : "reply"}` - : undefined, - `${keyText("app.agents.new")} new`, - selectedAgent ? `${keyText("app.agents.rename")} rename` : undefined, - selectedAgent - ? `${keyText("app.agents.delete")} ${selectedRow?.section === "inactive" ? "delete" : "stop/deactivate"}` - : undefined, - selectedSubagent - ? `${keyText("app.agents.delete")} ${selectedRow.section === "running" ? "stop" : "delete"}` - : undefined, - this.selectedRowCanShowProgram() ? `${keyText("app.agents.program")} program` : undefined, - ] - .filter((hint): hint is string => hint !== undefined) - .join(" "); + const hints = `${keyText("tui.select.up")}/${keyText("tui.select.down")} navigate ${keyText("tui.select.confirm")} open ${keyText("app.agents.new")} new ${keyText("app.shortcuts")} actions`; return truncateToWidth(theme.fg("muted", hints), width); } @@ -2802,27 +2781,15 @@ export class AgentsViewMode implements Component, Focusable { type DisplayItem = | { type: "spacer" } | { type: "heading"; section: AgentsViewSection } - | { type: "empty"; section: AgentsViewSection } + | { type: "running-subagents"; row: AgentsViewRow } | { type: "row"; row: AgentsViewRow }; -function buildDisplayItems(rows: readonly AgentsViewRow[]): DisplayItem[] { - const items: DisplayItem[] = []; - const sections: AgentsViewSection[] = ["running", "idle", "inactive"]; - for (const [index, section] of sections.entries()) { - if (index > 0) { - items.push({ type: "spacer" }); - } - items.push({ type: "heading", section }); - const sectionRows = getDisplayRowsForSection(rows, section); - if (sectionRows.length === 0) { - items.push({ type: "empty", section }); - continue; - } - for (const row of sectionRows) { - items.push({ type: "row", row }); - } - } - return items; +function compactSessionRows(rows: readonly AgentsViewRow[], showInactive: boolean): AgentsViewRow[] { + let visible = true; + return rows.filter((row) => { + if (row.depth === 0) visible = showInactive || row.section !== "inactive"; + return visible && row.kind !== "subagent-summary"; + }); } // Nested rows (subagent summaries and expanded subagents) always render in @@ -2864,93 +2831,43 @@ function hasLiveWork(row: AgentsViewRow): boolean { return row.section === "running" || row.runningSubagentCount > 0 || row.summary.hasRunningRlmChildren === true; } -interface AgentsViewUsageParts { - inTokens: string; - outTokens: string; - agentCost: string; - count: string; - totalCost: string; - age: string; -} - -const AGENTS_VIEW_USAGE_LABELS: AgentsViewUsageParts = { - inTokens: "↑in", - outTokens: "↓out", - agentCost: "$agent", - count: "#sub", - totalCost: "$total", - age: "age", -}; - -const AGENTS_VIEW_USAGE_COLUMNS = Object.keys(AGENTS_VIEW_USAGE_LABELS) as (keyof AgentsViewUsageParts)[]; - export interface AgentsViewUsageLayout { - /** Legend line per section, padded to that section's column widths. */ - legends: ReadonlyMap; - /** Details string per row identity, padded to its section's column widths. */ + legend: string; details: ReadonlyMap; + nameWidth: number; + modelWidth: number; + activityWidth: number; } -/** - * One shared column layout per section for the header legend and every row: - * each column is as wide as the section's widest value or its legend label, - * everything right-aligned, so the ` · ` separators land in the same terminal - * column for the legend and every row. Empty sessions render only the age, - * aligned to the age column. - */ -export function buildAgentsViewUsageLayout(rows: readonly AgentsViewRow[]): AgentsViewUsageLayout { - const rowsBySection = new Map(); - // Nested rows render inside their top-level agent's section block, so group - // by the block's section rather than each row's own. - let blockSection: AgentsViewSection = "running"; - for (const row of rows) { - if (row.depth === 0) blockSection = row.section; - if (row.kind !== "agent" && row.kind !== "subagent") continue; - const sectionRows = rowsBySection.get(blockSection) ?? []; - sectionRows.push(row); - rowsBySection.set(blockSection, sectionRows); - } - const legends = new Map(); - const details = new Map(); - for (const section of ["running", "idle", "inactive"] as const) { - const entries = (rowsBySection.get(section) ?? []).map((row) => { - const usage = row.summary.usage; - const parts: AgentsViewUsageParts = { - inTokens: `↑${formatTokenCount(usage?.inputTokens ?? 0)}`, - outTokens: `↓${formatTokenCount(usage?.outputTokens ?? 0)}`, - agentCost: `$${(usage?.cost ?? 0).toFixed(2)}`, - count: String(row.descendantCount), - totalCost: `$${row.recursiveCost.toFixed(2)}`, - age: formatSessionDuration(row.summary), - }; - return { identity: row.identity, empty: isEmptyAgentsViewSession(row.summary), parts }; - }); - const widths = {} as Record; - for (const column of AGENTS_VIEW_USAGE_COLUMNS) { - let width = visibleWidth(AGENTS_VIEW_USAGE_LABELS[column]); - for (const entry of entries) { - // Empty sessions render no usage segment; only their age takes space. - if (entry.empty && column !== "age") continue; - width = Math.max(width, visibleWidth(entry.parts[column])); - } - widths[column] = width; - } - const pad = (parts: AgentsViewUsageParts, column: keyof AgentsViewUsageParts): string => - padCellStart(parts[column], widths[column]); - const formatLine = (parts: AgentsViewUsageParts): string => - [ - `${pad(parts, "inTokens")} ${pad(parts, "outTokens")}`, - pad(parts, "agentCost"), - pad(parts, "count"), - pad(parts, "totalCost"), - pad(parts, "age"), - ].join(" · "); - legends.set(section, formatLine(AGENTS_VIEW_USAGE_LABELS)); - for (const entry of entries) { - details.set(entry.identity, entry.empty ? pad(entry.parts, "age") : formatLine(entry.parts)); - } - } - return { legends, details }; +export function buildCompactAgentsViewLayout(rows: readonly AgentsViewRow[], width = 120): AgentsViewUsageLayout { + const sessions = rows.filter((row) => row.kind === "agent" || row.kind === "subagent"); + const entries = sessions.map((row) => ({ + identity: row.identity, + cost: `$${row.recursiveCost.toFixed(2)}`, + age: formatSessionDuration(row.summary), + })); + const costWidth = entries.reduce((size, entry) => Math.max(size, visibleWidth(entry.cost)), 4); + const ageWidth = entries.reduce((size, entry) => Math.max(size, visibleWidth(entry.age)), 3); + const detailsWidth = costWidth + 2 + ageWidth; + const available = Math.max(0, width - detailsWidth - 4); + const desiredModelWidth = sessions.reduce( + (size, row) => Math.max(size, visibleWidth(formatSessionModel(row.summary))), + 12, + ); + const modelWidth = Math.min(desiredModelWidth, 32, Math.max(0, available - 12)); + const nameWidth = Math.min(28, Math.max(0, available - modelWidth)); + const activityWidth = Math.max(0, available - modelWidth - nameWidth - 2); + const detailLine = (cost: string, age: string) => `${padCellStart(cost, costWidth)} ${padCellStart(age, ageWidth)}`; + const headings = [formatTableCell("Session", nameWidth), formatTableCell("Model", modelWidth)]; + if (activityWidth > 0) headings.push(formatTableCell("Activity", activityWidth)); + headings.push(detailLine("Cost", "Age")); + return { + legend: formatTableCell(headings.join(" "), width), + details: new Map(entries.map((entry) => [entry.identity, detailLine(entry.cost, entry.age)])), + nameWidth, + modelWidth, + activityWidth, + }; } function padCellStart(value: string, width: number): string { @@ -2974,9 +2891,8 @@ function formatTableCell(value: string, width: number): string { return truncated + " ".repeat(Math.max(0, width - visibleWidth(truncated))); } -function formatRightTableCell(value: string, width: number): string { - const truncated = truncateToWidth(value, width, ""); - return " ".repeat(Math.max(0, width - visibleWidth(truncated))) + truncated; +function formatSessionModel(summary: SessionSummary): string { + return summary.model?.id ?? "-"; } function formatSessionDuration(summary: SessionSummary): string { diff --git a/packages/coding-agent/src/modes/daemon/daemon-mode.ts b/packages/coding-agent/src/modes/daemon/daemon-mode.ts index ab71ecfc60..94d456dea5 100644 --- a/packages/coding-agent/src/modes/daemon/daemon-mode.ts +++ b/packages/coding-agent/src/modes/daemon/daemon-mode.ts @@ -11,7 +11,7 @@ import { existsSync, mkdirSync, rmSync, writeFileSync } from "node:fs"; import { stat } from "node:fs/promises"; import { createConnection, createServer, type Server, type Socket } from "node:net"; import { basename, dirname, isAbsolute, join, resolve } from "node:path"; -import { type Api, getLogger, type Model } from "@earendil-works/pi-ai"; +import { type Api, findEnvKeys, getLogger, type Model } from "@earendil-works/pi-ai"; import { createCliSubprocessEnv, createCliSubprocessLaunchSpec } from "../../cli/subprocess-launch.js"; import { appendRotatingLog, @@ -89,7 +89,12 @@ import { import { ORPHAN_PROCESS_JOURNAL_ENV } from "../../core/orphan-process-journal.js"; import { PromptAdmissionCancelledError, waitForPromptAdmission } from "../../core/prompt-admission.js"; import { providerRetryPolicy } from "../../core/provider-retry.js"; -import type { CreateRlmSubagentRuntimeOptions, SubagentRuntimeHost } from "../../core/rlm-runtime.js"; +import type { + CreateRlmRootSessionOptions, + CreateRlmSubagentRuntimeOptions, + RlmCreateSessionResult, + SubagentRuntimeHost, +} from "../../core/rlm-runtime.js"; import { canPassivateSession, type IdleEvictionMinutes, @@ -107,7 +112,7 @@ import { import { resolveSessionPath } from "../../core/session-resolver.js"; import type { SessionStats } from "../../core/session-stats.js"; import { type SideQuestionRun, startSideQuestion } from "../../core/side-question.js"; -import { isProcessAlive, spawnHidden } from "../../utils/child-process.js"; +import { isProcessAlive, spawnHidden, waitForChildProcess } from "../../utils/child-process.js"; import { tryAcquireDirLock } from "../../utils/dir-lock.js"; import { killTrackedDetachedChildren } from "../../utils/shell.js"; import { @@ -140,6 +145,7 @@ import { filterClientEnv, withClientEnv } from "./daemon-client-env.js"; import { deserializeDaemonError, serializeDaemonError } from "./daemon-errors.js"; import { bindActiveSessionState } from "./daemon-extension-binding.js"; import { + collectDaemonLaunchEnv, createDaemonEventMeta, createDaemonReplayInfo, DAEMON_DEFAULT_CLIENT_CAPABILITIES, @@ -880,15 +886,48 @@ export class AgentDaemon { env: environment, stdio: "ignore", }); + const childExited = waitForChildProcess(child); + void childExited.catch(() => undefined); child.unref(); const deadline = Date.now() + 10_000; while (!this.shuttingDown && Date.now() < deadline) { - if (await this.canConnectToSupervisor(supervisorSocketPath)) { - this.log(`launched replacement supervisor on ${supervisorSocketPath}`); + if (child.pid === undefined || child.exitCode !== null || child.signalCode !== null) { + await childExited; + return; + } + for (const [client, boundClaim] of this.supervisorClaims) { + const { claim } = boundClaim; + if ( + claim.supervisorSocketPath !== supervisorSocketPath || + !client.authenticated || + client.socket.destroyed + ) + continue; + try { + await this.assertSupervisorClaimCurrent(claim); + } catch { + continue; + } + if ( + this.shuttingDown || + Date.now() >= deadline || + this.supervisorClaims.get(client) !== boundClaim || + !client.authenticated || + client.socket.destroyed + ) + continue; + if (claim.supervisorPid === child.pid) { + this.log(`launched replacement supervisor on ${supervisorSocketPath}`); + return; + } + if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL"); + await waitForPromptAdmission(childExited, AbortSignal.timeout(Math.max(1, deadline - Date.now()))); + this.log(`stopped losing replacement supervisor ${child.pid} on ${supervisorSocketPath}`); return; } await delay(50); } + this.log(`replacement supervisor ${child.pid} left running without a current authenticated supervisor`); } catch (error) { this.log(`failed to launch replacement supervisor: ${String(error)}`); } finally { @@ -2488,6 +2527,7 @@ export class AgentDaemon { private createSubagentRuntimeHost(parentState: ActiveSessionState): SubagentRuntimeHost { return { createRlmSubagentRuntime: async (options) => this.createRlmSubagentRuntime(parentState, options), + createRlmRootSession: async (options) => this.createRlmRootSession(parentState, options), completeRlmSubagentRuntime: (childId, session) => { const state = [...this.sessions.values()].find( (candidate) => @@ -2637,6 +2677,95 @@ export class AgentDaemon { }; } + private async createRlmRootSession( + parentState: ActiveSessionState, + options: CreateRlmRootSessionOptions, + ): Promise { + const supervisorSocketPath = this.supervisorSocketPathFromEnv(); + if (!this.options.worker || !supervisorSocketPath) { + throw new Error("rlm.create_session requires a daemon worker connected to its supervisor"); + } + + const client = new DaemonClient(supervisorSocketPath); + let activeSessionId: string | undefined; + try { + await client.connect(3000); + await client.waitForHello(3000); + const runtimeConfig = parentState.runtime.runtimeConfig; + const inheritsProvider = options.model.provider === parentState.runtime.session.model?.provider; + const authSource = parentState.runtime.services.authStorage.getAuthStatus(options.model.provider).source; + const apiKey = + inheritsProvider && + authSource === "runtime" && + (runtimeConfig?.provider ?? parentState.runtime.session.model?.provider) === options.model.provider + ? runtimeConfig?.apiKey + : undefined; + const launchEnv = collectDaemonLaunchEnv({ PATH: process.env.PATH }); + const envKey = authSource === "environment" ? findEnvKeys(options.model.provider)?.[0] : undefined; + if (envKey && process.env[envKey]) launchEnv[envKey] = process.env[envKey]; + if (options.model.provider === "prime-inference" && process.env.PRIME_TEAM_ID !== undefined) { + launchEnv.PRIME_TEAM_ID = process.env.PRIME_TEAM_ID; + } + const createResponse = await client.request( + { + type: "create", + lifecycle: "resident", + launchEnv, + ...(options.sessionName ? { name: options.sessionName } : {}), + config: { + cwd: options.cwd, + agentDir: parentState.runtime.services.agentDir, + ...(runtimeConfig?.sessionDir ? { sessionDir: runtimeConfig.sessionDir } : {}), + provider: options.model.provider, + model: options.model.id, + ...(apiKey ? { apiKey } : {}), + thinking: options.thinkingLevel, + ...(runtimeConfig?.telemetryDisabled ? { telemetryDisabled: true as const } : {}), + }, + }, + 120_000, + ); + if (!createResponse.success) throw deserializeDaemonError(createResponse); + const summary = createResponse.data as Partial | undefined; + activeSessionId = summary?.activeSessionId ?? summary?.id; + if ( + !activeSessionId || + typeof summary?.sessionId !== "string" || + !summary.sessionId || + typeof summary.sessionFile !== "string" || + !summary.sessionFile || + (summary.rlmDepth !== undefined && summary.rlmDepth !== 0) + ) { + throw new Error("Daemon supervisor returned an invalid depth-0 session summary"); + } + + const promptResponse = await client.request( + { + type: "prompt", + activeSessionId, + message: options.prompt, + source: "rpc", + }, + 30_000, + ); + if (!promptResponse.success) throw deserializeDaemonError(promptResponse); + return { + active_session_id: activeSessionId, + session_id: summary.sessionId, + name: summary.sessionName ?? activeSessionId, + session_file: summary.sessionFile, + model: `${options.model.provider}/${options.model.id}`, + }; + } catch (error) { + if (activeSessionId) { + await client.request({ type: "kill", activeSessionId }, 30_000).catch(() => undefined); + } + throw error; + } finally { + client.close(); + } + } + private async createRlmSubagentRuntime( parentState: ActiveSessionState, options: CreateRlmSubagentRuntimeOptions, diff --git a/packages/coding-agent/src/modes/interactive/components/custom-editor.ts b/packages/coding-agent/src/modes/interactive/components/custom-editor.ts index 058e4076be..ee505f92b0 100644 --- a/packages/coding-agent/src/modes/interactive/components/custom-editor.ts +++ b/packages/coding-agent/src/modes/interactive/components/custom-editor.ts @@ -8,6 +8,9 @@ import { visibleWidth, } from "@earendil-works/pi-tui"; import type { AppKeybinding, KeybindingsManager } from "../../../core/keybindings.js"; +import { ArgTokenHighlighter } from "./prompt-highlight.js"; + +const COMMAND_TOKEN_PATTERN = /^(\s*)\/(\S+)/; export interface CustomEditorOptions extends EditorOptions { placeholder?: string; @@ -25,6 +28,7 @@ export class CustomEditor extends Editor { private placeholder: string | undefined; private readonly placeholderColor: (text: string) => string; private readonly isArgumentCommand: (name: string) => boolean; + private readonly argTokenHighlighter = new ArgTokenHighlighter(); public actionHandlers: Map void> = new Map(); // Special handlers that can be dynamically replaced @@ -69,13 +73,29 @@ export class CustomEditor extends Editor { layoutLineIndex: number, lineText: string, cursorCol: number | undefined, + sourceLine?: number, + sourceStart?: number, + ): string { + if (sourceLine === undefined || sourceStart === undefined || this.getBashPromptInfo(this.getLines()[0] ?? "")) { + return this.styleCommandToken(displayText, layoutLineIndex, lineText, cursorCol); + } + // Arg tokens are styled first; their spans start after the command token, so the command offsets stay valid. + const highlighted = this.argTokenHighlighter.highlightLine(displayText, lineText, sourceLine, sourceStart); + return this.styleCommandToken(highlighted, layoutLineIndex, lineText, cursorCol); + } + + private styleCommandToken( + displayText: string, + layoutLineIndex: number, + lineText: string, + cursorCol: number | undefined, ): string { const commandColor = this.commandColor; if (!commandColor || layoutLineIndex !== 0) { return displayText; } - const match = /^(\s*)\/(\S+)/.exec(lineText); + const match = COMMAND_TOKEN_PATTERN.exec(lineText); if (!match) { return displayText; } @@ -123,6 +143,9 @@ export class CustomEditor extends Editor { } override render(width: number): string[] { + const commandMatch = COMMAND_TOKEN_PATTERN.exec(this.getLines()[0] ?? ""); + const isArgumentCommandLine = commandMatch !== null && this.isArgumentCommand(commandMatch[2]!); + this.argTokenHighlighter.reset(this.getLines(), isArgumentCommandLine); let lines = super.render(width); if (this.placeholder && this.getText().length === 0 && lines.length >= 2) { lines = [lines[0]!, this.renderPlaceholderLine(width), ...lines.slice(2)]; diff --git a/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts b/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts index 6893496870..4e96c827b3 100644 --- a/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts +++ b/packages/coding-agent/src/modes/interactive/components/injected-prompt-message.ts @@ -10,6 +10,9 @@ import { } from "@earendil-works/pi-tui"; import { GOAL_CONTEXT_CUSTOM_TYPE, type GoalContextDetails } from "../../../core/goals.js"; import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ASYNC_BASH_COMPLETION_PREVIEW_LABEL, + type AsyncBashCompletionDetails, type CustomMessage, HEARTBEAT_PROMPT_CUSTOM_TYPE, type HeartbeatPromptDetails, @@ -21,9 +24,11 @@ import { type RlmChildTerminalNoticeDetails, } from "../../../core/messages.js"; import { getMarkdownTheme, theme } from "../theme/theme.js"; +import { agentMessageSummaryLine } from "./agent-message.js"; import { expandCollapseHint } from "./keybinding-hints.js"; type InjectedPromptDetails = + | AsyncBashCompletionDetails | GoalContextDetails | HeartbeatPromptDetails | IpythonStateRestoredDetails @@ -34,7 +39,8 @@ type InjectedPromptMessage = CustomMessage; export function isInjectedPromptMessage(message: AgentMessage): message is InjectedPromptMessage { return ( message.role === "custom" && - (message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE || + (message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE || + message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE || message.customType === GOAL_CONTEXT_CUSTOM_TYPE || message.customType === IPYTHON_STATE_RESTORED_CUSTOM_TYPE || message.customType === RLM_CHILD_FAILURE_CUSTOM_TYPE || @@ -125,6 +131,15 @@ export class InjectedPromptMessageComponent extends Container { if (this.message.customType === HEARTBEAT_PROMPT_CUSTOM_TYPE) { return this.heartbeatHeaderText(); } + if (this.message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE) { + const details = this.message.details as AsyncBashCompletionDetails | undefined; + const participant = details ? `pid ${details.pid}` : "bash"; + const status = details ? `exit ${details.exitCode}` : undefined; + const hint = this.expanded ? "" : ` ${expandCollapseHint("app.tools.expand", false)}`; + return ( + agentMessageSummaryLine(ASYNC_BASH_COMPLETION_PREVIEW_LABEL, participant, status) + theme.fg("dim", hint) + ); + } if (this.message.customType === IPYTHON_STATE_RESTORED_CUSTOM_TYPE) { const details = this.message.details as IpythonStateRestoredDetails | undefined; const label = details?.restored === false ? "Started fresh Python kernel" : "Restored Python kernel state"; diff --git a/packages/coding-agent/src/modes/interactive/components/prompt-highlight.ts b/packages/coding-agent/src/modes/interactive/components/prompt-highlight.ts new file mode 100644 index 0000000000..d17709e6de --- /dev/null +++ b/packages/coding-agent/src/modes/interactive/components/prompt-highlight.ts @@ -0,0 +1,199 @@ +import { visibleWidth } from "@earendil-works/pi-tui"; +import { type ThemeColor, theme } from "../theme/theme.js"; + +const ARG_TOKEN_PATTERN = /@"[^"\n]*"|@(?:\\[^\s\x1b]|[^\s\x1b|])+|--[A-Za-z0-9][A-Za-z0-9-]*/g; +/** Also matches a bare `--` end-of-options separator; only used for argument-taking slash commands. */ +const ARG_TOKEN_PATTERN_WITH_SEPARATOR = + /@"[^"\n]*"|@(?:\\[^\s\x1b]|[^\s\x1b|])+|--[A-Za-z0-9][A-Za-z0-9-]*|--(?=\s|$)/g; +const FG_SGR_PATTERN = /\x1b\[(?:0|39|3[0-7]|9[0-7]|38;[0-9;]+)m/g; +/** Escape sequences the editor splices into displayed text (cursor highlight, IME marker). */ +const CURSOR_ESCAPE_PATTERN = /\x1b\[[0-9;]*m|\x1b_[^\x07]*\x07/g; + +const MASK_BASE_START = 0xe000; +/** Each masked grapheme gets its own private-use base char, so restoring is a lookup, not positional. */ +const MASK_CAPACITY = 0xf8ff - MASK_BASE_START + 1; +const MASK_EXTRA_WIDTH = "\uFF9E"; +const MASK_PATTERN = /[\uE000-\uF8FF]\uFF9E*/gu; +/** Literal mask-range characters would alias generated placeholders; messages containing them skip masking. */ +const MASK_LITERAL_PATTERN = /[\uE000-\uF8FF\uFF9E]/u; + +const graphemeSegmenter = new Intl.Segmenter(undefined, { granularity: "grapheme" }); + +interface ArgTokenSpan { + start: number; + end: number; + color: ThemeColor; +} + +function tokenColor(token: string): ThemeColor { + return token.startsWith("@") ? "success" : "mdLink"; +} + +function hasTokenBoundary(text: string, index: number): boolean { + return index === 0 || /\s/.test(text.charAt(index - 1)); +} + +function findArgTokens(text: string, fromIndex = 0, includeBareSeparator = false): ArgTokenSpan[] { + const spans: ArgTokenSpan[] = []; + const pattern = includeBareSeparator ? ARG_TOKEN_PATTERN_WITH_SEPARATOR : ARG_TOKEN_PATTERN; + for (const match of text.matchAll(pattern)) { + if (match.index < fromIndex || !hasTokenBoundary(text, match.index)) continue; + spans.push({ start: match.index, end: match.index + match[0].length, color: tokenColor(match[0]) }); + } + return spans; +} + +/** Foreground SGR active at index; theme.fg() closes with \x1b[39m, so it must be re-emitted. */ +function activeFgBefore(line: string, index: number): string { + let active = ""; + for (const sgr of line.slice(0, index).matchAll(FG_SGR_PATTERN)) { + active = sgr[0] === "\x1b[0m" ? "" : sgr[0]; + } + return active; +} + +export function styleArgumentTokens( + text: string, + styleOther: (segment: string) => string = (segment) => segment, + includeBareSeparator = false, +): string { + let result = ""; + let offset = 0; + for (const token of findArgTokens(text, 0, includeBareSeparator)) { + result += styleOther(text.slice(offset, token.start)) + theme.fg(token.color, text.slice(token.start, token.end)); + offset = token.end; + } + return result + styleOther(text.slice(offset)); +} + +/** Masks the slash command and @path/--flag tokens with same-width placeholders before markdown layout. */ +export class PromptTokenMask { + readonly text: string; + private graphemes: { segment: string; color: ThemeColor }[] = []; + + constructor(source: string, commandEnd = 0, includeBareSeparator = false) { + // Markdown turns tabs into three spaces; a masked raw tab would be restored into a three-column layout. + source = source.replace(/\t/g, " "); + if (MASK_LITERAL_PATTERN.test(source)) { + this.text = source; + return; + } + const tokens: ArgTokenSpan[] = []; + if (commandEnd > 0) { + tokens.push({ start: 0, end: commandEnd, color: "accent" }); + } + tokens.push(...findArgTokens(source, commandEnd, includeBareSeparator)); + + let text = ""; + let cursor = 0; + for (const token of tokens) { + text += source.slice(cursor, token.start); + for (const { segment } of graphemeSegmenter.segment(source.slice(token.start, token.end))) { + const width = visibleWidth(segment); + if (width === 0) { + // Zero-width graphemes stay literal: invisible either way, and extracted text stays exact. + text += segment; + continue; + } + if (this.graphemes.length === MASK_CAPACITY) { + this.text = source; + this.graphemes = []; + return; + } + text += String.fromCharCode(MASK_BASE_START + this.graphemes.length) + MASK_EXTRA_WIDTH.repeat(width - 1); + this.graphemes.push({ segment, color: token.color }); + } + cursor = token.end; + } + this.text = text + source.slice(cursor); + } + + private graphemeFor(placeholder: string): { segment: string; color: ThemeColor } | undefined { + return this.graphemes[placeholder.charCodeAt(0) - MASK_BASE_START]; + } + + /** Restores masked graphemes in text extracted from a render, e.g. selection-region cell content. */ + restoreText(text: string): string { + return text.replace(MASK_PATTERN, (placeholder) => this.graphemeFor(placeholder)?.segment ?? placeholder); + } + + restoreLine(line: string): string { + let result = ""; + let copied = 0; + let run: { start: number; end: number; color: ThemeColor; text: string } | undefined; + const flush = () => { + if (!run) return; + result += line.slice(copied, run.start) + theme.fg(run.color, run.text) + activeFgBefore(line, run.start); + copied = run.end; + run = undefined; + }; + for (const match of line.matchAll(MASK_PATTERN)) { + const grapheme = this.graphemeFor(match[0]); + if (!grapheme) continue; // literal mask-range character from an unmasked source; leave it untouched + if (run && run.color === grapheme.color && run.end === match.index) { + run.text += grapheme.segment; + run.end += match[0].length; + } else { + flush(); + run = { + start: match.index, + end: match.index + match[0].length, + color: grapheme.color, + text: grapheme.segment, + }; + } + } + flush(); + return result + line.slice(copied); + } +} + +/** Styles tokens in laid-out editor lines from spans on the logical source lines; reset() before each render pass. */ +export class ArgTokenHighlighter { + private spans: ArgTokenSpan[][] = []; + + reset(lines: readonly string[], includeBareSeparator = false): void { + this.spans = lines.map((line) => findArgTokens(line, 0, includeBareSeparator)); + } + + /** displayText is chunkText with cursor escapes spliced in; chunkText starts at sourceStart within sourceLine. */ + highlightLine(displayText: string, chunkText: string, sourceLine: number, sourceStart: number): string { + const rangeEnd = sourceStart + chunkText.length; + const spans: ArgTokenSpan[] = []; + for (const span of this.spans[sourceLine] ?? []) { + if (span.end <= sourceStart) continue; + if (span.start >= rangeEnd) break; + spans.push({ + start: Math.max(span.start, sourceStart) - sourceStart, + end: Math.min(span.end, rangeEnd) - sourceStart, + color: span.color, + }); + } + if (spans.length === 0) return displayText; + + // Maps visible code-unit offsets to displayText offsets, skipping the editor's cursor escapes. + const visibleStart: number[] = []; + let pos = 0; + for (const seq of displayText.matchAll(CURSOR_ESCAPE_PATTERN)) { + for (; pos < seq.index; pos++) visibleStart.push(pos); + pos = seq.index + seq[0].length; + } + for (; pos < displayText.length; pos++) visibleStart.push(pos); + + let result = ""; + let copied = 0; + for (const span of spans) { + const start = visibleStart[span.start] ?? displayText.length; + const end = (visibleStart[span.end - 1] ?? displayText.length - 1) + 1; + // The cursor splice may carry a full reset mid-span; wrap each segment so the token color survives it. + const styled = displayText + .slice(start, end) + .split("\x1b[0m") + .map((segment) => theme.fg(span.color, segment)) + .join("\x1b[0m"); + result += displayText.slice(copied, start) + styled; + copied = end; + } + return result + displayText.slice(copied); + } +} diff --git a/packages/coding-agent/src/modes/interactive/components/slash-command-message.ts b/packages/coding-agent/src/modes/interactive/components/slash-command-message.ts index d977028edd..1c58af8e83 100644 --- a/packages/coding-agent/src/modes/interactive/components/slash-command-message.ts +++ b/packages/coding-agent/src/modes/interactive/components/slash-command-message.ts @@ -1,6 +1,7 @@ import { Box, Container, Text } from "@earendil-works/pi-tui"; -import { parseSlashCommand } from "../../../core/slash-commands.js"; +import { builtinSlashCommandTakesArgument, parseSlashCommand } from "../../../core/slash-commands.js"; import { theme } from "../theme/theme.js"; +import { styleArgumentTokens } from "./prompt-highlight.js"; const OSC133_ZONE_START = "\x1b]133;A\x07"; const OSC133_ZONE_END = "\x1b]133;B\x07"; @@ -11,10 +12,16 @@ export function isLeadingSlashCommand(text: string, isRecognized: (name: string) return command !== undefined && isRecognized(command.name); } -export function styleSlashCommandText(text: string, styleRest: (rest: string) => string = (rest) => rest): string { +export function styleSlashCommandText( + text: string, + styleRest: (rest: string, includeBareSeparator: boolean) => string = (rest, includeBareSeparator) => + styleArgumentTokens(rest, undefined, includeBareSeparator), +): string { const parsed = parseSlashCommand(text); const commandEnd = parsed ? parsed.name.length + 1 : text.length; - return `${theme.fg("accent", text.slice(0, commandEnd))}${styleRest(text.slice(commandEnd))}`; + // Matches the editor's gate: a bare -- is only meaningful in commands that take arguments. + const includeBareSeparator = parsed !== undefined && builtinSlashCommandTakesArgument(parsed.name); + return `${theme.fg("accent", text.slice(0, commandEnd))}${styleRest(text.slice(commandEnd), includeBareSeparator)}`; } /** Renders a durable session command with the same layout as a user message. */ diff --git a/packages/coding-agent/src/modes/interactive/components/user-message.ts b/packages/coding-agent/src/modes/interactive/components/user-message.ts index 58b8766378..581de174d0 100644 --- a/packages/coding-agent/src/modes/interactive/components/user-message.ts +++ b/packages/coding-agent/src/modes/interactive/components/user-message.ts @@ -1,51 +1,39 @@ -import { Box, type Component, Container, Markdown, type MarkdownTheme, visibleWidth } from "@earendil-works/pi-tui"; -import { parseSlashCommand } from "../../../core/slash-commands.js"; +import { + Box, + type Component, + Container, + Markdown, + type MarkdownTheme, + type TableCellSelectionRegion, +} from "@earendil-works/pi-tui"; +import { builtinSlashCommandTakesArgument, parseSlashCommand } from "../../../core/slash-commands.js"; import { getMarkdownTheme, theme } from "../theme/theme.js"; -import { isLeadingSlashCommand } from "./slash-command-message.js"; +import { PromptTokenMask } from "./prompt-highlight.js"; const OSC133_ZONE_START = "\x1b]133;A\x07"; const OSC133_ZONE_END = "\x1b]133;B\x07"; const OSC133_ZONE_FINAL = "\x1b]133;C\x07"; -const COMMAND_MASK_BASE = "\uE000"; -const COMMAND_MASK_EXTRA_WIDTH = "\uFF9E"; -const COMMAND_MASK_ZERO_WIDTH = "\u2060"; -const COMMAND_MASK_PATTERN = /\u2060|\uE000\uFF9E*/gu; -const graphemeSegmenter = new Intl.Segmenter(undefined, { granularity: "grapheme" }); -class SlashCommandMarkdown implements Component { +class HighlightedMarkdown implements Component { private readonly markdown: Markdown; - private readonly commandGraphemes: string[]; + private readonly mask: PromptTokenMask; - constructor(text: string, markdownTheme: MarkdownTheme) { - const parsed = parseSlashCommand(text); - const commandEnd = parsed ? parsed.name.length + 1 : text.length; - this.commandGraphemes = [...graphemeSegmenter.segment(text.slice(0, commandEnd))].map(({ segment }) => segment); - const placeholder = this.commandGraphemes - .map((grapheme) => { - const width = visibleWidth(grapheme); - return width === 0 - ? COMMAND_MASK_ZERO_WIDTH - : COMMAND_MASK_BASE + COMMAND_MASK_EXTRA_WIDTH.repeat(width - 1); - }) - .join(""); - this.markdown = new Markdown(`${placeholder}${text.slice(commandEnd)}`, 0, 0, markdownTheme, { + constructor(text: string, markdownTheme: MarkdownTheme, commandEnd = 0, includeBareSeparator = false) { + this.mask = new PromptTokenMask(text, commandEnd, includeBareSeparator); + this.markdown = new Markdown(this.mask.text, 0, 0, markdownTheme, { color: (content: string) => theme.fg("userMessageText", content), }); } render(width: number): string[] { - let commandOffset = 0; - return this.markdown.render(width).map((line) => { - const chunks: string[] = []; - const replaced = line.replace(COMMAND_MASK_PATTERN, (placeholder) => { - const grapheme = this.commandGraphemes[commandOffset]; - if (grapheme === undefined) return placeholder; - commandOffset++; - chunks.push(grapheme); - return ""; - }); - return chunks.length === 0 ? replaced : `${theme.fg("accent", chunks.join(""))}${replaced}`; - }); + return this.markdown.render(width).map((line) => this.mask.restoreLine(line)); + } + + getSelectionRegions(): ReadonlyArray { + return this.markdown.getSelectionRegions().map((region) => ({ + ...region, + content: this.mask.restoreText(region.content), + })); } invalidate(): void { @@ -62,14 +50,12 @@ export class UserMessageComponent extends Container { isRecognizedSlashCommand: (name: string) => boolean = () => false, ) { super(); + const command = parseSlashCommand(text); + const commandEnd = command && isRecognizedSlashCommand(command.name) ? command.name.length + 1 : 0; + const includeBareSeparator = + command !== undefined && commandEnd > 0 && builtinSlashCommandTakesArgument(command.name); this.contentBox = new Box(2, 1, (content: string) => theme.getUserMessageBackgroundColor()(content)); - this.contentBox.addChild( - isLeadingSlashCommand(text, isRecognizedSlashCommand) - ? new SlashCommandMarkdown(text, markdownTheme) - : new Markdown(text, 0, 0, markdownTheme, { - color: (content: string) => theme.fg("userMessageText", content), - }), - ); + this.contentBox.addChild(new HighlightedMarkdown(text, markdownTheme, commandEnd, includeBareSeparator)); this.addChild(this.contentBox); } diff --git a/packages/coding-agent/src/modes/interactive/interactive-mode.ts b/packages/coding-agent/src/modes/interactive/interactive-mode.ts index 2ba8b292c5..d031837b30 100644 --- a/packages/coding-agent/src/modes/interactive/interactive-mode.ts +++ b/packages/coding-agent/src/modes/interactive/interactive-mode.ts @@ -99,6 +99,7 @@ import type { KernelSentAgentMessage } from "../../core/kernel/index.js"; import { type AppKeybinding, KeybindingsManager } from "../../core/keybindings.js"; import { runMcpManagementCommand } from "../../core/mcp/mcp-command.js"; import { + ASYNC_BASH_COMPLETION_PREVIEW_LABEL, bashOutputToText, COMPACTION_OUTCOME_CUSTOM_TYPE, type CustomMessage, @@ -211,6 +212,7 @@ import { formatKeyText, keyHint, keyText, rawKeyHint } from "./components/keybin import { createMermaidMarkdownTransform } from "./components/mermaid.js"; import type { AuthSelectorProvider } from "./components/oauth-selector.js"; import { PrimeOnboardingSplashComponent } from "./components/prime-onboarding-splash.js"; +import { styleArgumentTokens } from "./components/prompt-highlight.js"; import { MalformedRefinementOutcomeMessageComponent, RefinementOutcomeMessageComponent, @@ -311,7 +313,8 @@ function isLabeledQueuedPreview(message: string): boolean { return ( message.startsWith(`${HEARTBEAT_PROMPT_PREVIEW_LABEL}: `) || message.startsWith(`${GOAL_CONTEXT_PREVIEW_LABEL}: `) || - message.startsWith(`${AGENT_MESSAGE_RECEIVED_PREVIEW_LABEL}: `) + message.startsWith(`${AGENT_MESSAGE_RECEIVED_PREVIEW_LABEL}: `) || + message.startsWith(`${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}: `) ); } @@ -325,9 +328,12 @@ export function styleQueuedMessagePreview( isRecognizedSlashCommand: (name: string) => boolean, ): string { const preview = formatQueuedMessagePreview(message, label); - if (!isLeadingSlashCommand(message, isRecognizedSlashCommand)) return theme.fg("dim", preview); + const styleDim = (segment: string) => theme.fg("dim", segment); + if (!isLeadingSlashCommand(message, isRecognizedSlashCommand)) return styleArgumentTokens(preview, styleDim); const prefix = preview.slice(0, preview.length - message.length); - return `${theme.fg("dim", prefix)}${styleSlashCommandText(message, (rest) => theme.fg("dim", rest))}`; + return `${theme.fg("dim", prefix)}${styleSlashCommandText(message, (rest, includeBareSeparator) => + styleArgumentTokens(rest, styleDim, includeBareSeparator), + )}`; } function isExpandable(obj: unknown): obj is Expandable { diff --git a/packages/coding-agent/test/agent-session-recursion.test.ts b/packages/coding-agent/test/agent-session-recursion.test.ts index 598b4c59df..a946ee3bc6 100644 --- a/packages/coding-agent/test/agent-session-recursion.test.ts +++ b/packages/coding-agent/test/agent-session-recursion.test.ts @@ -23,7 +23,7 @@ import { AuthStorage } from "../src/core/auth-storage.js"; import { computeOwnAndTotalUsage } from "../src/core/context-tree.js"; import type { LoadExtensionsResult } from "../src/core/extensions/index.js"; import { type HostRequestHandlers, ReplKernelManager } from "../src/core/kernel/index.js"; -import { convertToLlm } from "../src/core/messages.js"; +import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE, convertToLlm } from "../src/core/messages.js"; import { ModelRegistry } from "../src/core/model-registry.js"; import { createDefaultRlmSubagentSessionName, @@ -866,6 +866,33 @@ describe("AgentSession rlm recursion", () => { expect(doneUpdate?.toolUseCount).toBeUndefined(); }); + it("wakes the agent with a follow-up when a detached bash handle completes", async () => { + const prompts: string[] = []; + const root = createSession({ + streamFn: (_model, context) => { + prompts.push(userText(context)); + return streamAnswer("checked shell result"); + }, + }); + const handlers = (root as unknown as InspectableRlmSession)._createKernelHostHandlers(); + const completed = handlers["bash.completed"]; + if (!completed) throw new Error("Missing bash.completed host handler"); + + await expect(completed({ pid: 42, command: "npm test", exitCode: 1 })).resolves.toEqual({}); + await root.waitForIdle(); + + expect(prompts).toEqual([ + expect.stringContaining("Inspect the saved BashHandle with .poll(), .output(), or .tail()"), + ]); + expect(root.messages).toContainEqual( + expect.objectContaining({ + role: "custom", + customType: ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + details: { pid: 42, command: "npm test", exitCode: 1 }, + }), + ); + }); + it("marks an in-cell roled send to the parent as replied", async () => { const sendAgentMessage = vi.fn(async () => ({ id: "agentmsg-reply", @@ -3113,6 +3140,117 @@ describe("AgentSession rlm recursion", () => { expect(child.rlmMaxDepth).toBe(3); }); + it("creates an independent root session through the explicit daemon host operation", async () => { + const createRlmSubagentRuntime = vi.fn(async () => { + throw new Error("unexpected child spawn"); + }); + const createRlmRootSession = vi.fn(async () => ({ + active_session_id: "root-active", + session_id: "root-session", + name: "researcher", + session_file: join(tempDir, "sessions", "root-session.jsonl"), + model: `${model.provider}/${model.id}`, + })); + const assertSessionNameAvailable = vi.fn(); + const root = createSession({ + agentMessageController: { + assertSessionNameAvailable, + listAgents: async () => ({ + current: { activeSessionId: "current-root", sessionId: "current-session" }, + agents: [], + }), + sendAgentMessage: async () => { + throw new Error("unexpected message"); + }, + }, + subagentRuntimeHost: { + createRlmSubagentRuntime, + createRlmRootSession, + deleteRlmSubagentRuntime: vi.fn(async () => {}), + }, + }); + + await expect( + root.createRlmSession("independent task", { + name: "researcher", + model: `${model.provider}/${model.id}`, + thinking: "off", + cwd: "other-project", + }), + ).resolves.toEqual({ + active_session_id: "root-active", + session_id: "root-session", + name: "researcher", + session_file: join(tempDir, "sessions", "root-session.jsonl"), + model: `${model.provider}/${model.id}`, + }); + expect(createRlmSubagentRuntime).not.toHaveBeenCalled(); + expect(assertSessionNameAvailable).toHaveBeenCalledWith({ name: "researcher", depth: 0 }); + expect(createRlmRootSession).toHaveBeenCalledWith({ + prompt: "independent task", + sessionName: "researcher", + cwd: join(tempDir, "other-project"), + model, + thinkingLevel: "off", + }); + await expect(root.createRlmSession("task", { cwd: " " })).rejects.toThrow( + "rlm.create_session cwd must be a non-empty string", + ); + await expect(root.createRlmSession(" ")).rejects.toThrow("rlm.create_session prompt must not be empty"); + expect(createRlmRootSession).toHaveBeenCalledOnce(); + }); + + it("does not create a root session after disposal during name preflight", async () => { + let releaseName: () => void = () => {}; + const nameGate = new Promise((resolve) => { + releaseName = resolve; + }); + const createRlmRootSession = vi.fn(async () => ({ + active_session_id: "new-root", + session_id: "new-session", + name: "researcher", + session_file: join(tempDir, "new-session.jsonl"), + model: `${model.provider}/${model.id}`, + })); + const root = createSession({ + agentMessageController: { + assertSessionNameAvailable: () => nameGate, + listAgents: async () => ({ current: { activeSessionId: "root", sessionId: "session" }, agents: [] }), + sendAgentMessage: vi.fn(), + }, + subagentRuntimeHost: { + createRlmSubagentRuntime: vi.fn(), + createRlmRootSession, + deleteRlmSubagentRuntime: vi.fn(), + }, + }); + const creating = root.createRlmSession("independent task", { name: "researcher" }); + root.dispose(); + releaseName(); + await expect(creating).rejects.toThrow("disposed"); + expect(createRlmRootSession).not.toHaveBeenCalled(); + }); + + it("keeps top-level session creation unavailable to nested or inline sessions", async () => { + const createRlmRootSession = vi.fn(); + const nested = createSession({ + depth: 1, + subagentRuntimeHost: { + createRlmSubagentRuntime: vi.fn(), + createRlmRootSession, + deleteRlmSubagentRuntime: vi.fn(async () => {}), + }, + }); + await expect(nested.createRlmSession("escape to root")).rejects.toThrow("available only from a depth-0 session"); + expect(createRlmRootSession).not.toHaveBeenCalled(); + nested.dispose(); + + const inline = createSession(); + await expect(inline.createRlmSession("detached root")).rejects.toThrow( + "requires a daemon-backed depth-0 session", + ); + }); + it("rejects child creation at the configured recursion depth cap", async () => { const root = createSession({ depth: 1, maxDepth: 1 }); diff --git a/packages/coding-agent/test/agents-view-mode.test.ts b/packages/coding-agent/test/agents-view-mode.test.ts index bd731ee688..0bc0bca9fc 100644 --- a/packages/coding-agent/test/agents-view-mode.test.ts +++ b/packages/coding-agent/test/agents-view-mode.test.ts @@ -1,3 +1,4 @@ +import { getModel } from "@earendil-works/pi-ai"; import { setKeybindings } from "@earendil-works/pi-tui"; import stripAnsi from "strip-ansi"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; @@ -9,7 +10,7 @@ import type { AgentConnectionSavedSessionInfo } from "../src/modes/agent-connect import { AgentsViewMode, type AgentsViewPersistentState, - buildAgentsViewUsageLayout, + buildCompactAgentsViewLayout, combineAgentsViewStartupNotices, createInitialAgentsViewPersistentState, runAgentsViewMode, @@ -105,6 +106,9 @@ describe("AgentsViewMode", () => { editor: { getText: () => "matching query" }, persistentState: { query: "" }, savedSearchFetchStarted: true, + // Searching claims the visible row even while a remembered anchor is + // still waiting for its catalog row: user intent supersedes restore. + selectionAnchorPending: true, selectedIndex: 4, rebuildRows: vi.fn(), syncSelectedRowState: vi.fn(), @@ -119,6 +123,37 @@ describe("AgentsViewMode", () => { expect(self.persistentState.query).toBe("matching query"); expect(self.rebuildRows).toHaveBeenCalledOnce(); expect(self.selectedIndex).toBe(4); + expect(self.syncSelectedRowState).toHaveBeenCalledOnce(); + }); + + it("keeps the queried selection when a remembered session arrives later", () => { + const remembered = summary({ + id: "remembered", + activeSessionId: "remembered", + sessionId: "remembered-session", + sessionFile: "/tmp/remembered.jsonl", + sessionName: "match remembered", + }); + const fallback = summary({ sessionName: "match fallback" }); + const persistentState = createInitialAgentsViewPersistentState({ initialSession: remembered }); + persistentState.savedCatalogLoaded = true; + const view = new AgentsViewMode({ config: {}, uiServices: createUiServices() }, persistentState); + try { + Reflect.set(view, "lastListedSummaries", [fallback]); + invoke("reconcileCatalogs", view); + expect(Reflect.get(view, "selectionAnchorPending")).toBe(true); + + invoke("setSearchQuery", view, "match"); + expect(Reflect.get(view, "selectionAnchorPending")).toBe(false); + expect(persistentState.selectedSessionKey?.sessionId).toBe(fallback.sessionId); + + Reflect.set(view, "lastListedSummaries", [remembered, fallback]); + invoke("reconcileCatalogs", view); + const rows = Reflect.get(view, "rows") as AgentsViewRow[]; + expect(rows[Reflect.get(view, "selectedIndex") as number]?.summary.sessionId).toBe(fallback.sessionId); + } finally { + stopThemeWatcher(); + } }); it("loads the saved catalog on view entry without a search query", () => { @@ -643,7 +678,7 @@ describe("AgentsViewMode", () => { expect(parentRow?.identity).toBe("session:root-session"); expandedSubagentParents.add(parentRow!.identity); invoke("reconcileCatalogs", self); - expect(rowsOf(self).some((row) => row.kind === "subagent-summary" && row.expanded)).toBe(true); + expect(rowsOf(self).some((row) => row.kind === "subagent")).toBe(true); } // The runtime flushes the session file; the record identity flips to file:. self.lastListedSummaries = [{ ...parent, sessionFile: "/tmp/root.jsonl" }, child]; @@ -656,7 +691,7 @@ describe("AgentsViewMode", () => { expect( expandedRows.find((row) => row.kind === "agent" && row.summary.sessionId === "root-session")?.identity, ).toBe("file:/tmp/root.jsonl"); - expect(expandedRows.some((row) => row.kind === "subagent-summary" && row.expanded)).toBe(true); + expect(expandedRows.some((row) => row.kind === "subagent-summary")).toBe(false); expect(expandedRows.some((row) => row.kind === "subagent" && row.summary.sessionId === "child-session")).toBe( true, ); @@ -669,7 +704,7 @@ describe("AgentsViewMode", () => { expect(collapsedView.expandedSubagentParents.size).toBe(0); }); - it("toggles subagent list expansion from the summary row", () => { + it("toggles subagent list expansion from the parent row", () => { const expandedSubagentParents = new Set(["root-row"]); const programShownParents = new Set(["root-row"]); const persistentState: AgentsViewPersistentState = { @@ -684,7 +719,7 @@ describe("AgentsViewMode", () => { syncSelectedRowState: vi.fn(), ui: { requestRender: vi.fn() }, }; - const summaryRow = { kind: "subagent-summary", parentIdentity: "root-row", expanded: true }; + const summaryRow = { kind: "agent", identity: "root-row", expanded: true }; invoke("toggleSubagentList", self, summaryRow); expect(expandedSubagentParents.size).toBe(0); @@ -718,144 +753,193 @@ describe("AgentsViewMode", () => { } }); - it("renders aligned usage columns with an explicit subagent count and drops the message count", () => { + it("shows each session model and aligned total cost including collapsed descendants", () => { + const created = new Date(Date.now() - 120_000).toISOString(); const parent = summary({ id: "spender", activeSessionId: "spender", sessionId: "spender-session", + sessionName: "spender", + model: { ...getModel("openai", "gpt-4o"), id: "gpt-5.6-sol" }, + created, + summary: "Analyzing runtime composition", usage: { inputTokens: 12437, outputTokens: 1234, cost: 0.42 }, }); const child = summary({ - id: "spender-child", - activeSessionId: "spender-child", - sessionId: "spender-child-session", - sessionFile: "/tmp/spender-child.jsonl", + id: "child", + activeSessionId: "child", + sessionId: "child-session", + sessionFile: "/tmp/child.jsonl", runtimeKind: "subagent", parentActiveSessionId: "spender", + model: { ...getModel("openai", "gpt-4o"), provider: "prime-inference", id: "glm-5.2-fast" }, + created, usage: { inputTokens: 500, outputTokens: 50, cost: 0.68 }, }); const inactive = summary({ - id: "saved-only", - activeSessionId: undefined, - sessionId: "saved-only-session", - sessionFile: "/tmp/saved-only.jsonl", - rosterStatus: "inactive", - messageCount: 7, - }); - const empty = summary({ - id: "empty-draft", + id: "saved", activeSessionId: undefined, - sessionId: "empty-draft-session", - sessionFile: "/tmp/empty-draft.jsonl", + sessionId: "saved-session", + sessionFile: "/tmp/saved.jsonl", rosterStatus: "inactive", - messageCount: 0, - modified: new Date(Date.now() - 120_000).toISOString(), + created, + model: { ...getModel("openai", "gpt-4o"), provider: "prime-inference", id: "glm-5.2-fast" }, + usage: { inputTokens: 900, outputTokens: 80, cost: 123.45 }, }); + const rows = buildAgentsViewRows([parent, child, inactive]); const view = new AgentsViewMode({ config: {}, uiServices: createUiServices() }, {}); - + Reflect.set(view, "rows", rows); + Reflect.set(view, "selectedIndex", -1); try { - const collapsed = buildAgentsViewRows([parent, child, inactive, empty]); - const rows = buildAgentsViewRows( - [parent, child, inactive, empty], - new Set(collapsed.map((row) => row.identity)), - ); - Reflect.set(view, "rows", rows); - const layout = buildAgentsViewUsageLayout(rows); - const line = (row: AgentsViewRow | undefined) => - stripAnsi(invoke("renderRow", view, row, 200, layout.details) as string); - const byId = (sessionId: string, kind?: string) => - rows.find((row) => row.summary.sessionId === sessionId && (!kind || row.kind === kind)); - - // Shared per-section layout: every column right-aligned to - // max(widest section value, legend label width). - expect(line(byId("spender-session"))).toContain("↑12k ↓1.2k · $0.42 · 1 · $1.10 ·"); - expect(line(byId("spender-child-session", "subagent"))).toContain("↑500 ↓50 · $0.68 · 0 · $0.68 ·"); - const inactiveLine = line(byId("saved-only-session")); - expect(inactiveLine).toContain("↑0 ↓0 · $0.00 · 0 · $0.00 ·"); - expect(inactiveLine).not.toContain("7 ·"); - // The ` · ` separators land in the same column for the legend and every - // row of its section. - const dotColumns = (text: string) => [...text].flatMap((ch, index) => (ch === "·" ? [index] : [])); - for (const [section, sessionId] of [ - ["idle", "spender-session"], - ["idle", "spender-child-session"], - ["inactive", "saved-only-session"], - ] as const) { - const detail = layout.details.get(byId(sessionId)!.identity)!; - expect(dotColumns(detail)).toEqual(dotColumns(layout.legends.get(section)!)); + const parentRow = rows.find((row) => row.summary.sessionId === parent.sessionId)!; + const savedRow = rows.find((row) => row.summary.sessionId === inactive.sessionId)!; + const render = (row: AgentsViewRow, width: number) => + stripAnsi(invoke("renderRow", view, row, width, buildCompactAgentsViewLayout(rows, width)) as string); + const parentLine = render(parentRow, 120); + const savedLine = render(savedRow, 120); + expect(parentLine).toContain("gpt-5.6-sol"); + expect(savedLine).toContain("glm-5.2-fast"); + expect(parentLine).toContain("$1.10"); + expect(parentLine).not.toContain("$0.42"); + expect(parentLine).not.toMatch(/[↑↓]/); + expect(parentLine).toMatch(/2m\s*$/); + expect(parentLine.indexOf("$1.10") + "$1.10".length).toBe(savedLine.indexOf("$123.45") + "$123.45".length); + for (const width of [60, 80]) { + const narrow = render(parentRow, width); + expect(narrow).toContain("gpt-5.6-sol"); + expect(narrow).toContain("$1.10"); + expect(narrow).toMatch(/2m\s*$/); + expect(narrow.length).toBeLessThanOrEqual(width); } - // Empty sessions keep the age but drop the whole usage segment. - const emptyLine = line(byId("empty-draft-session")); - expect(emptyLine).not.toContain("↑"); - expect(emptyLine).not.toContain("$"); - expect(emptyLine).toMatch(/\d+[smhd]\s*$/); - // Without a shared layout the row pads only against its own section of one. - const bare = { ...byId("spender-session")!, summary: { ...parent, usage: undefined } }; - expect(stripAnsi(invoke("renderRow", view, bare, 200) as string)).toContain( - " ↑0 ↓0 · $0.00 · 1 · $1.10 ·", - ); } finally { stopThemeWatcher(); } }); - it("shows the bold usage legend on every section header", () => { - const running = (id: string, created: string) => + it("renders one column header across status groups without repeating subagent hints", () => { + const summaries = [ summary({ - id, - activeSessionId: id, - sessionId: `${id}-session`, - sessionName: id, + id: "busy", + activeSessionId: "busy", + sessionId: "busy-session", + sessionName: "busy", activity: "working", isStreaming: true, - created, - }); - const parent = running("busy-parent", "2026-01-01T00:00:00Z"); - const child = summary({ - id: "busy-child", - activeSessionId: "busy-child", - sessionId: "busy-child-session", - sessionName: "busy-child", - sessionFile: "/tmp/busy-child.jsonl", - runtimeKind: "subagent", - parentActiveSessionId: "busy-parent", - }); - const summaries = [ - running("busy-solo", "2026-01-02T00:00:00Z"), - parent, - child, - summary({ id: "idle-a", activeSessionId: "idle-a", sessionId: "idle-a-session", sessionName: "idle-a" }), - summary({ id: "idle-b", activeSessionId: "idle-b", sessionId: "idle-b-session", sessionName: "idle-b" }), + }), + summary({ + id: "idle", + activeSessionId: "idle", + sessionId: "idle-session", + sessionName: "idle", + sessionFile: "/tmp/idle.jsonl", + }), + summary({ + id: "child", + sessionId: "child-session", + sessionFile: "/tmp/child.jsonl", + runtimeKind: "subagent", + parentActiveSessionId: "busy", + }), ]; - const parentIdentity = buildAgentsViewRows(summaries).find( - (row) => row.summary.sessionId === "busy-parent-session", - )!.identity; - const rows = buildAgentsViewRows(summaries, new Set([parentIdentity])); const view = new AgentsViewMode({ config: {}, uiServices: createUiServices() }, {}); - try { - Reflect.set(view, "rows", rows); + Reflect.set(view, "lastListedSummaries", summaries); + invoke("reconcileCatalogs", view); Reflect.set(view, "selectedIndex", -1); Reflect.set(view, "ui", { terminal: { rows: 60 }, requestRender: () => {} }); const rendered = invoke("renderSessionRows", view, 120, 40) as string[]; const lines = rendered.map(stripAnsi); - const headings = lines.filter((line) => /^(Running|Idle|Inactive) \(\d+\)/.test(line)); - expect(headings).toHaveLength(3); - for (const heading of headings) { - expect(heading).toMatch(/↑in\s+↓out ·\s+\$agent ·\s+#sub ·\s+\$total ·\s+age$/); + expect(lines.filter((line) => /Model/.test(line) && /Age/i.test(line))).toHaveLength(1); + expect(lines.some((line) => line.startsWith("Running"))).toBe(true); + expect(lines.some((line) => line.startsWith("Idle"))).toBe(true); + expect(lines.join("\n")).not.toMatch(/show program|#sub|\$agent|↑in|↓out/); + const rows = Reflect.get(view, "rows") as AgentsViewRow[]; + expect(rows.filter((row) => row.kind === "subagent-summary")).toHaveLength(0); + for (const line of rendered) { + expect(invoke("finalizeRenderedLine", view, line, 120)).not.toContain("\x1b[48"); } - // Same bold weight for title and legend. - const runningLegend = buildAgentsViewUsageLayout(rows).legends.get("running")!; - expect(invoke("renderSectionHeading", view, "running", 120, runningLegend)).toContain( - theme.bold(runningLegend), - ); + } finally { + stopThemeWatcher(); + } + }); - // Session rows carry no background of their own; only the selection - // highlight may paint one. - const finalized = rendered.map((line) => invoke("finalizeRenderedLine", view, line, 120) as string); - for (const line of finalized) { - expect(line).not.toContain("\x1b[48"); - } + it("keeps collapsed inactive sessions out of navigation and reveals them for search", () => { + const live = summary({ sessionName: "live" }); + const saved = summary({ + id: "saved", + activeSessionId: undefined, + sessionId: "saved-session", + sessionName: "archive-match", + sessionFile: "/tmp/saved.jsonl", + rosterStatus: "inactive", + lifecycle: "archived", + }); + const view = new AgentsViewMode({ config: {}, uiServices: createUiServices() }, { savedCatalogLoaded: true }); + const rows = () => Reflect.get(view, "rows") as AgentsViewRow[]; + try { + Reflect.set(view, "lastListedSummaries", [live]); + Reflect.set(view, "savedSessions", [ + { + path: saved.sessionFile!, + id: saved.sessionId, + cwd: saved.cwd, + name: saved.sessionName, + created: new Date(), + modified: new Date(), + messageCount: 1, + firstMessage: "archive-match", + allMessagesText: "archive-match", + }, + ]); + invoke("reconcileCatalogs", view); + expect(rows().map((row) => row.summary.sessionId)).toEqual([live.sessionId]); + invoke("moveSelection", view, 1); + expect(rows()[Reflect.get(view, "selectedIndex") as number]?.summary.sessionId).toBe(live.sessionId); + view.handleInput("\x1bi"); + expect(rows().some((row) => row.summary.sessionId === saved.sessionId)).toBe(true); + view.handleInput("\x1bi"); + expect(rows().some((row) => row.summary.sessionId === saved.sessionId)).toBe(false); + invoke("setSearchQuery", view, "archive-match"); + expect(rows().some((row) => row.summary.sessionId === saved.sessionId)).toBe(true); + invoke("setSearchQuery", view, ""); + expect(rows().some((row) => row.summary.sessionId === saved.sessionId)).toBe(false); + } finally { + stopThemeWatcher(); + } + }); + + it("opens a parent with Enter and reveals its spawn program only on request", () => { + const parent = summary({ sessionName: "parent" }); + const child = summary({ + id: "child", + activeSessionId: "child", + sessionId: "child-session", + sessionFile: "/tmp/child.jsonl", + runtimeKind: "subagent", + parentActiveSessionId: parent.activeSessionId, + spawnCode: 'await rlm("Inspect the code")', + }); + const view = new AgentsViewMode({ config: {}, uiServices: createUiServices() }, {}); + const rows = () => Reflect.get(view, "rows") as AgentsViewRow[]; + try { + Reflect.set(view, "lastListedSummaries", [parent, child]); + invoke("reconcileCatalogs", view); + expect(rows().map((row) => row.kind)).toEqual(["agent"]); + const finish = vi.fn(); + Reflect.set(view, "finish", finish); + invoke("openSelected", view); + expect(finish).toHaveBeenCalledWith( + expect.objectContaining({ + type: "open", + summary: expect.objectContaining({ sessionId: parent.sessionId }), + }), + ); + invoke("cycleProgramForSelected", view); + expect(rows().some((row) => row.kind === "subagent-code" && row.code === child.spawnCode)).toBe(true); + expect(rows().some((row) => row.kind === "subagent" && row.summary.sessionId === child.sessionId)).toBe(true); + expect(rows().some((row) => row.kind === "subagent-summary")).toBe(false); + invoke("cycleProgramForSelected", view); + expect(rows().some((row) => row.kind === "subagent-code")).toBe(false); } finally { stopThemeWatcher(); } @@ -881,7 +965,8 @@ describe("AgentsViewMode", () => { Reflect.set(view, "selectedIndex", rows.length - 1); Reflect.set(view, "ui", { terminal: { rows: 13 }, requestRender: () => {} }); const lines = (invoke("renderSessionRows", view, 120, 4) as string[]).map(stripAnsi); - expect(lines[0]).toContain("..."); + expect(lines[1]).toContain("..."); + expect(lines).toHaveLength(4); const lastTitle = rows.at(-1)!.title; expect(lines.some((line) => line.includes(lastTitle))).toBe(true); } finally { @@ -889,33 +974,77 @@ describe("AgentsViewMode", () => { } }); - it("renders a collapsed group's busy-subagent badge legibly instead of dimmed", () => { - const parent = summary({ id: "parent", activeSessionId: "parent", sessionId: "parent-session" }); - const busyChild = summary({ - id: "busy-child", - activeSessionId: "busy-child", - sessionId: "busy-child-session", - sessionFile: "/tmp/busy-child.jsonl", + it("reveals usage details through actions and closes them before searching", () => { + const view = new AgentsViewMode({ config: {}, uiServices: createUiServices() }, { savedCatalogLoaded: true }); + try { + Reflect.set(view, "lastListedSummaries", [ + summary({ sessionName: "parent", usage: { inputTokens: 1234, outputTokens: 56, cost: 1.23 } }), + ]); + invoke("reconcileCatalogs", view); + view.handleInput("?"); + const actions = (invoke("renderSessionRows", view, 120, 20) as string[]).map(stripAnsi).join("\n"); + expect(actions).toContain("1234 in"); + expect(actions).toContain("$1.23"); + view.handleInput("p"); + expect(Reflect.get(view, "showActions")).toBe(false); + const rows = (invoke("renderSessionRows", view, 120, 20) as string[]).map(stripAnsi).join("\n"); + expect(rows).toContain("parent"); + expect(rows).not.toContain("1234 in"); + } finally { + stopThemeWatcher(); + } + }); + + it("shows running-subagent counts only while collapsed and work remains", () => { + const parent = summary({ sessionName: "parent" }); + const child = summary({ + id: "child", + activeSessionId: "child", + sessionId: "child-session", + sessionFile: "/tmp/child.jsonl", runtimeKind: "subagent", - parentActiveSessionId: "parent", + parentActiveSessionId: parent.activeSessionId, activity: "working", - isSessionActive: true, isStreaming: true, }); - const idleChild = { ...busyChild, activity: "idle" as const, isSessionActive: false, isStreaming: false }; + const secondChild = { + ...child, + id: "child-2", + activeSessionId: "child-2", + sessionId: "child-session-2", + sessionFile: "/tmp/child-2.jsonl", + }; const view = new AgentsViewMode({ config: {}, uiServices: createUiServices() }, {}); - + const rows = () => Reflect.get(view, "rows") as AgentsViewRow[]; + const lines = () => (invoke("renderSessionRows", view, 120, 20) as string[]).map(stripAnsi); try { - const busyRows = buildAgentsViewRows([parent, busyChild]); - const busySummaryRow = busyRows.find((row) => row.kind === "subagent-summary"); - expect(busySummaryRow).toMatchObject({ section: "idle", title: "1 subagent running" }); - Reflect.set(view, "rows", busyRows); - expect(invoke("renderRow", view, busySummaryRow, 160)).toContain(theme.fg("success", "▸ 1 subagent running")); - - const idleRows = buildAgentsViewRows([parent, idleChild]); - const idleSummaryRow = idleRows.find((row) => row.kind === "subagent-summary"); - Reflect.set(view, "rows", idleRows); - expect(invoke("renderRow", view, idleSummaryRow, 160)).toContain(theme.fg("dim", "▸ 1 subagent")); + Reflect.set(view, "lastListedSummaries", [parent, child, secondChild]); + invoke("reconcileCatalogs", view); + expect(rows()).toHaveLength(1); + expect(invoke("renderRow", view, rows()[0], 120)).toContain("▸"); + const collapsed = lines(); + const parentIndex = collapsed.findIndex((line) => line.includes("parent")); + expect(collapsed[parentIndex + 1]).toBe(" 2 subagents running"); + invoke("moveSelection", view, 1); + expect(Reflect.get(view, "selectedIndex")).toBe(0); + view.handleInput("\x1b[1;3C"); + expect(rows().map((row) => row.kind)).toEqual(["agent", "subagent", "subagent"]); + expect(lines().join("\n")).not.toContain("subagents running"); + expect(invoke("renderRow", view, rows()[0], 120)).toContain("▾"); + view.handleInput("\x1b[1;3C"); + expect(rows()).toHaveLength(1); + expect(lines()).toContain(" 2 subagents running"); + const idleChild = { ...child, activity: "idle", isStreaming: false }; + Reflect.set(view, "lastListedSummaries", [parent, idleChild, secondChild]); + invoke("reconcileCatalogs", view); + expect(lines()).toContain(" 1 subagent running"); + Reflect.set(view, "lastListedSummaries", [ + parent, + idleChild, + { ...secondChild, activity: "idle", isStreaming: false }, + ]); + invoke("reconcileCatalogs", view); + expect(lines().join("\n")).not.toMatch(/subagents? running/); } finally { stopThemeWatcher(); } diff --git a/packages/coding-agent/test/agents-view-state.test.ts b/packages/coding-agent/test/agents-view-state.test.ts index 07b95cb956..de75aa6892 100644 --- a/packages/coding-agent/test/agents-view-state.test.ts +++ b/packages/coding-agent/test/agents-view-state.test.ts @@ -1670,6 +1670,14 @@ describe("agents view state", () => { }); const handoffFrame = persistentState.scopeFrames?.at(-1); + // The scope root is never a row of its own children view: seeding it as + // the selection anchor would only arm pending-anchor for the whole scan. + expect(persistentState.selectedRowIdentity).toBeUndefined(); + expect(persistentState.selectedSessionKey).toBeUndefined(); + expect(persistentState.backSession).toBe(chat); + const unscoped = createInitialAgentsViewPersistentState({ initialSession: chat }); + expect(unscoped.selectedRowIdentity).toBeDefined(); + expect(handoffFrame).toEqual({ scope: rootScope, returnChat: chat }); expect(createInitialAgentsViewScopeFrames(rootScope, persistentState.backSession)).toEqual([handoffFrame]); expect(createInitialAgentsViewScopeFrames(rootScope, makeSummary({ sessionId: "stale-session" }))).toEqual([ diff --git a/packages/coding-agent/test/async-bash-completion.test.ts b/packages/coding-agent/test/async-bash-completion.test.ts new file mode 100644 index 0000000000..4ce8738402 --- /dev/null +++ b/packages/coding-agent/test/async-bash-completion.test.ts @@ -0,0 +1,53 @@ +import { describe, expect, it, vi } from "vitest"; +import { startsAgentRun } from "../src/core/agent-messages.js"; +import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + convertToLlm, + createAsyncBashCompletionMessage, +} from "../src/core/messages.js"; +import { createAsyncBashCompletionHostHandler } from "../src/core/rlm-runtime.js"; + +describe("async bash completion", () => { + it("creates a model-visible instruction to inspect the saved handle", () => { + const message = createAsyncBashCompletionMessage({ + pid: 42, + command: "npm test", + exitCode: 1, + }); + + expect(message.customType).toBe(ASYNC_BASH_COMPLETION_CUSTOM_TYPE); + expect(message.content).toContain("pid 42, exit code 1"); + expect(message.content).toContain("npm test"); + expect(message.content).toContain(".poll(), .output(), or .tail()"); + expect(convertToLlm([message])).toEqual([ + { + role: "user", + content: [{ type: "text", text: message.content }], + timestamp: message.timestamp, + }, + ]); + }); + + it("starts a new agent run for a background completion follow-up", () => { + const message = createAsyncBashCompletionMessage({ pid: 42, command: "long-running-tool", exitCode: 0 }); + expect(startsAgentRun(message)).toBe(true); + }); + + it("validates and forwards kernel completion payloads", async () => { + const completion = vi.fn(); + const handler = createAsyncBashCompletionHostHandler(completion); + const payload = { pid: 42, command: "npm test", exitCode: 0 }; + + await expect(handler(payload)).resolves.toEqual({}); + expect(completion).toHaveBeenCalledWith(payload); + }); + + it.each([ + [{ pid: 0, command: "ok", exitCode: 0 }, "positive integer"], + [{ pid: 1, command: "", exitCode: 0 }, "non-empty string"], + [{ pid: 1, command: "ok", exitCode: 0.5 }, "exitCode"], + ])("rejects an invalid payload %#", async (payload, error) => { + const handler = createAsyncBashCompletionHostHandler(() => undefined); + await expect(handler(payload)).rejects.toThrow(error); + }); +}); diff --git a/packages/coding-agent/test/custom-editor.test.ts b/packages/coding-agent/test/custom-editor.test.ts index ef3336e3ce..37ee3e034e 100644 --- a/packages/coding-agent/test/custom-editor.test.ts +++ b/packages/coding-agent/test/custom-editor.test.ts @@ -3,6 +3,7 @@ import { CURSOR_MARKER, setKeybindings, visibleWidth } from "@earendil-works/pi- import { beforeEach, describe, expect, it, vi } from "vitest"; import { KeybindingsManager } from "../src/core/keybindings.js"; import { CustomEditor } from "../src/modes/interactive/components/custom-editor.js"; +import { initTheme, type ThemeColor, theme } from "../src/modes/interactive/theme/theme.js"; const passthrough = (text: string) => text; @@ -279,6 +280,64 @@ describe("CustomEditor", () => { } }); + const makeHighlightEditor = (text: string, options?: { isArgumentCommand?: (name: string) => boolean }) => { + initTheme("dark"); + const editor = new CustomEditor(fakeTui, editorTheme, new KeybindingsManager(), options); + editor.setText(text); + return editor; + }; + + it.each<{ name: string; text: string; width: number; color: ThemeColor; has?: string[]; lacks?: string[] }>([ + { name: "--flags in the input text", text: "/new --name @a.ts", width: 40, color: "mdLink", has: ["--name"] }, + { name: "@paths in the input text", text: "/new --name @a.ts", width: 40, color: "success", has: ["@a.ts"] }, + { + name: "wrapped @path fragments across editor lines", + text: "check @src/very-long-file-name.ts please", + width: 16, + color: "success", + has: ["@src/very-lon", "g-file-name.t", "s"], + }, + { + name: "quoted @paths across wrapped editor lines", + text: 'open @"docs/some very long name.txt" now', + width: 16, + color: "success", + has: ['@"docs/some ', "very long ", 'name.txt"'], + }, + { name: "no line bleed", text: "@abcde\nfoo bar", width: 11, color: "success", has: ["@abcde"], lacks: ["foo"] }, + ])("highlights $name", ({ text, width, color, has, lacks }) => { + const rendered = makeHighlightEditor(text).render(width).join("\n"); + + for (const fragment of has ?? []) expect(rendered).toContain(theme.fg(color, fragment)); + for (const fragment of lacks ?? []) expect(rendered).not.toContain(theme.fg(color, fragment)); + }); + + it("highlights a bare -- separator only for argument commands", () => { + const options = { isArgumentCommand: (name: string) => name === "new" }; + const separator = theme.fg("mdLink", "--"); + + expect(makeHighlightEditor("/new --name bla -- hello", options).render(60).join("\n")).toContain(separator); + expect(makeHighlightEditor("this -- however -- is fine", options).render(60).join("\n")).not.toContain(separator); + expect(makeHighlightEditor("/unknown -- hello", options).render(60).join("\n")).not.toContain(separator); + }); + + it("does not mis-color visible text matching a scrolled-away token", () => { + // 9 lines with the cursor at the end scroll @foo out of view; the visible plain "foo" must stay uncolored. + const rendered = makeHighlightEditor("@foo\nhidden\nfoo\nl3\nl4\nl5\nl6\nl7\nl8").render(20).join("\n"); + + expect(rendered).toContain("↑ 2 more"); + expect(rendered).not.toContain(theme.fg("success", "foo")); + }); + + it("keeps the token tail colored when the cursor sits inside the token", () => { + const editor = makeHighlightEditor("check @src/foo.ts"); + editor.handleInput("\x1b[D"); + editor.handleInput("\x1b[D"); + + // The cursor's full reset sits before the final "s"; the tail must be re-colored. + expect(editor.render(40)[1]!).toContain(`\x1b[0m${theme.fg("success", "s")}`); + }); + it("renders no header when the callback returns undefined", () => { const editor = new CustomEditor(fakeTui, editorTheme, new KeybindingsManager()); const withoutCallback = editor.render(40); diff --git a/packages/coding-agent/test/daemon-mode.test.ts b/packages/coding-agent/test/daemon-mode.test.ts index 41e53fc92e..6f7f7a6afc 100644 --- a/packages/coding-agent/test/daemon-mode.test.ts +++ b/packages/coding-agent/test/daemon-mode.test.ts @@ -49,6 +49,7 @@ import { } from "../src/core/session-manager.js"; import { SettingsManager } from "../src/core/settings-manager.js"; import type { ActiveSessionState, DaemonSocketClient } from "../src/modes/daemon/active-session-state.js"; +import { DaemonClient } from "../src/modes/daemon/daemon-client.js"; import { AgentDaemon, cancelPendingExtensionUiRequests, @@ -61,6 +62,7 @@ import { } from "../src/modes/daemon/daemon-mode.js"; import { createDaemonCommandEnvelope, + DAEMON_DEFAULT_SERVER_CAPABILITIES, DAEMON_PROTOCOL_INFO, DAEMON_SCHEMA_ID, DAEMON_SCHEMA_REVISION, @@ -1777,6 +1779,305 @@ describe("daemon mode helpers", () => { } }); + it.each([ + { + label: "parent-only runtime key", + provider: "prime-inference", + configuredProvider: "prime-inference", + runtimeKey: "parent-runtime-key", + envKey: undefined, + stale: false, + }, + { + label: "parent environment key and selected team", + provider: "prime-inference", + configuredProvider: "prime-inference", + runtimeKey: undefined, + envKey: "parent-env-key", + stale: false, + }, + { + label: "selected provider alias precedence", + provider: "anthropic", + configuredProvider: "prime-inference", + runtimeKey: "parent-runtime-key", + envKey: "parent-env-key", + stale: false, + }, + { + label: "resolved provider switch", + provider: "openai", + configuredProvider: "prime-inference", + runtimeKey: "parent-runtime-key", + envKey: "parent-env-key", + stale: false, + }, + { + label: "stale runtime key after parent provider switch", + provider: "prime-inference", + configuredProvider: "openai", + runtimeKey: "parent-runtime-key", + envKey: "parent-env-key", + stale: false, + }, + { + label: "rejected runtime credential", + provider: "prime-inference", + configuredProvider: "prime-inference", + runtimeKey: "parent-runtime-key", + envKey: undefined, + stale: true, + }, + { + label: "rejected environment credential", + provider: "prime-inference", + configuredProvider: "prime-inference", + runtimeKey: undefined, + envKey: "parent-env-key", + stale: true, + }, + ])( + "creates and prompts a resident depth-0 session with $label", + async ({ provider, configuredProvider, runtimeKey, envKey, stale }) => { + const tempDir = mkdtempSync(join(tmpdir(), "pa-root-session-")); + const socketPath = join(tempDir, "supervisor.sock"); + const commands: Array> = []; + const server: Server = createServer((socket) => { + socket.on("error", () => undefined); + socket.write( + `${JSON.stringify({ + type: "daemon_hello", + socketPath, + protocol: DAEMON_PROTOCOL_INFO, + schemaId: DAEMON_SCHEMA_ID, + schemaRevision: DAEMON_SCHEMA_REVISION, + clientId: "supervisor", + serverCapabilities: DAEMON_DEFAULT_SERVER_CAPABILITIES, + })} +`, + ); + let buffer = ""; + socket.on("data", (chunk) => { + buffer += chunk.toString(); + for (;;) { + const newline = buffer.indexOf("\n"); + if (newline === -1) return; + const wire = JSON.parse(buffer.slice(0, newline)) as { + id: string; + command?: Record; + type?: string; + }; + buffer = buffer.slice(newline + 1); + const command = wire.command ?? wire; + commands.push(command); + const type = command.type as string; + const data = + type === "create" + ? { + id: "new-root-active", + activeSessionId: "new-root-active", + sessionId: "new-root-session", + sessionFile: join(tempDir, "new-root-session.jsonl"), + sessionName: "researcher", + cwd: join(tempDir, "project"), + rlmDepth: 0, + } + : undefined; + socket.write( + `${JSON.stringify({ type: "response", id: wire.id, command: type, success: true, data })} +`, + ); + } + }); + }); + const previousSupervisorSocket = process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV]; + try { + await new Promise((resolveListen) => server.listen(socketPath, resolveListen)); + process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV] = socketPath; + vi.stubEnv("PRIME_API_KEY", envKey); + vi.stubEnv("PRIME_TEAM_ID", "parent-team"); + vi.stubEnv("OPENAI_API_KEY", "unrelated-provider-key"); + vi.stubEnv("ANTHROPIC_OAUTH_TOKEN", "selected-anthropic-token"); + vi.stubEnv("ANTHROPIC_API_KEY", "lower-priority-anthropic-key"); + vi.stubEnv("UNRELATED_SECRET", "unrelated-secret"); + vi.stubEnv("PATH", `/parent/toolchain:${process.env.PATH}`); + const daemon = new AgentDaemon("/tmp/prime-agent-worker-test.sock", { + defaultSessionConfig: { agentDir: tempDir, cwd: tempDir }, + createRuntime: vi.fn(), + worker: { authenticationToken: "worker-token" }, + }); + const authStorage = AuthStorage.inMemory( + {}, + { usePrimeCliConfig: true, primeCliConfigPath: join(tempDir, "prime-config.json") }, + ); + if (runtimeKey) authStorage.setRuntimeApiKey(configuredProvider, runtimeKey); + if (stale) expect(authStorage.markAuthStale("prime-inference")).toBe(true); + const parent = makeState("parent-root"); + parent.runtime = { + ...parent.runtime, + runtimeConfig: { + sessionDir: join(tempDir, "sessions"), + telemetryDisabled: true, + provider: configuredProvider, + apiKey: runtimeKey, + }, + session: { model: { provider: "prime-inference", id: "parent-model" } }, + services: { agentDir: join(tempDir, "agent"), authStorage }, + } as ActiveSessionState["runtime"]; + const createHost = ( + daemon as unknown as { createSubagentRuntimeHost(state: ActiveSessionState): SubagentRuntimeHost } + ).createSubagentRuntimeHost.bind(daemon); + const host = createHost(parent); + const result = await host.createRlmRootSession?.({ + prompt: "investigate independently", + sessionName: "researcher", + cwd: join(tempDir, "project"), + model: { provider, id: "model" } as Model, + thinkingLevel: "high", + }); + + expect(result).toEqual({ + active_session_id: "new-root-active", + session_id: "new-root-session", + name: "researcher", + session_file: join(tempDir, "new-root-session.jsonl"), + model: `${provider}/model`, + }); + expect(commands.filter((command) => command.type !== "ack_result")).toEqual([ + expect.objectContaining({ + type: "create", + lifecycle: "resident", + name: "researcher", + config: expect.objectContaining({ + cwd: join(tempDir, "project"), + agentDir: join(tempDir, "agent"), + sessionDir: join(tempDir, "sessions"), + provider, + model: "model", + thinking: "high", + telemetryDisabled: true, + }), + }), + expect.objectContaining({ + type: "prompt", + activeSessionId: "new-root-active", + message: "investigate independently", + source: "rpc", + }), + ]); + const createCommand = commands.find((command) => command.type === "create"); + expect(createCommand).toBeDefined(); + expect((createCommand?.launchEnv as Record | undefined)?.PATH).toBe(process.env.PATH); + const config = createCommand?.config as Record; + expect(config.apiKey).toBe( + provider === "prime-inference" && provider === configuredProvider && !stale ? runtimeKey : undefined, + ); + expect(createCommand?.launchEnv).toEqual({ + PATH: process.env.PATH, + ...(provider === "openai" ? { OPENAI_API_KEY: "unrelated-provider-key" } : {}), + ...(provider === "anthropic" ? { ANTHROPIC_OAUTH_TOKEN: "selected-anthropic-token" } : {}), + ...(provider === "prime-inference" + ? { + ...(envKey && !stale && !(runtimeKey && configuredProvider === provider) + ? { PRIME_API_KEY: envKey } + : {}), + PRIME_TEAM_ID: "parent-team", + } + : {}), + }); + if (provider === "prime-inference") { + const childAuth = AuthStorage.inMemory( + {}, + { usePrimeCliConfig: true, primeCliConfigPath: join(tempDir, "prime-config.json") }, + ); + vi.stubEnv("PRIME_TEAM_ID", (createCommand?.launchEnv as Record).PRIME_TEAM_ID); + expect(childAuth.getProviderHeaders(provider)).toEqual({ "X-Prime-Team-ID": "parent-team" }); + } + } finally { + vi.unstubAllEnvs(); + if (previousSupervisorSocket === undefined) delete process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV]; + else process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV] = previousSupervisorSocket; + await new Promise((resolveClose) => server.close(() => resolveClose())); + rmSync(tempDir, { recursive: true, force: true }); + } + }, + ); + + it("does not prompt or guess a cleanup target after a create response timeout", async () => { + const tempDir = mkdtempSync(join(tmpdir(), "pa-root-timeout-")); + const socketPath = join(tempDir, "supervisor.sock"); + const commands: Array> = []; + const server = createServer((socket) => { + socket.on("error", () => {}); + socket.write( + `${JSON.stringify({ + type: "daemon_hello", + socketPath, + protocol: DAEMON_PROTOCOL_INFO, + schemaId: DAEMON_SCHEMA_ID, + schemaRevision: DAEMON_SCHEMA_REVISION, + clientId: "supervisor", + serverCapabilities: DAEMON_DEFAULT_SERVER_CAPABILITIES, + })}\n`, + ); + let buffer = ""; + socket.on("data", (chunk) => { + buffer += chunk.toString(); + for (;;) { + const newline = buffer.indexOf("\n"); + if (newline === -1) break; + const wire = JSON.parse(buffer.slice(0, newline)); + buffer = buffer.slice(newline + 1); + commands.push(wire.command ?? wire); + } + }); + }); + const request = DaemonClient.prototype.request; + const requestSpy = vi.spyOn(DaemonClient.prototype, "request").mockImplementation(function ( + this: DaemonClient, + command, + timeout, + options, + ) { + return request.call(this, command, command.type === "create" ? 20 : timeout, options); + }); + const previousSocket = process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV]; + try { + await new Promise((resolveListen) => server.listen(socketPath, resolveListen)); + process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV] = socketPath; + const daemon = new AgentDaemon(join(tempDir, "worker.sock"), { + defaultSessionConfig: { agentDir: tempDir, cwd: tempDir }, + createRuntime: vi.fn(), + worker: { authenticationToken: "worker-token" }, + }); + const parent = makeState("parent-root"); + parent.runtime = { + ...parent.runtime, + session: { model: { provider: "prime-inference", id: "parent-model" } }, + services: { agentDir: tempDir, authStorage: AuthStorage.inMemory() }, + } as ActiveSessionState["runtime"]; + const host = ( + daemon as unknown as { createSubagentRuntimeHost(state: ActiveSessionState): SubagentRuntimeHost } + ).createSubagentRuntimeHost(parent); + await expect( + host.createRlmRootSession?.({ + prompt: "do not run without create acknowledgement", + cwd: tempDir, + model: { provider: "prime-inference", id: "model" } as Model, + thinkingLevel: "off", + }), + ).rejects.toThrow('response to "create"'); + expect(commands.map((command) => command.type)).toEqual(["create"]); + expect(commands[0]).toMatchObject({ lifecycle: "resident" }); + } finally { + requestSpy.mockRestore(); + if (previousSocket === undefined) delete process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV]; + else process.env[DAEMON_WORKER_SUPERVISOR_SOCKET_ENV] = previousSocket; + await new Promise((resolveClose) => server.close(() => resolveClose())); + rmSync(tempDir, { recursive: true, force: true }); + } + }); + it("lists and role-addresses root siblings hosted by another worker", async () => { const daemon = new AgentDaemon("/tmp/prime-agent-worker-test.sock", { defaultSessionConfig: { agentDir: "/tmp/prime-agent-test-agent", cwd: "/tmp" }, diff --git a/packages/coding-agent/test/daemon-supervisor-launch.test.ts b/packages/coding-agent/test/daemon-supervisor-launch.test.ts new file mode 100644 index 0000000000..9fb4f3d8f0 --- /dev/null +++ b/packages/coding-agent/test/daemon-supervisor-launch.test.ts @@ -0,0 +1,183 @@ +import { type ChildProcess, spawn } from "node:child_process"; +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import { getProcessStartId } from "../src/core/session-lease.js"; +import { AgentDaemon } from "../src/modes/daemon/daemon-mode.js"; +import * as childProcesses from "../src/utils/child-process.js"; + +interface Claim { + supervisorGeneration: string; + supervisorPid: number; + supervisorProcessStartId: string; + supervisorSocketPath: string; +} +interface Client { + authenticated: boolean; + socket: { destroyed: boolean }; +} +interface BoundClaim { + claim: Claim; + ownerFingerprint: string; +} +interface Launcher { + launchReplacementSupervisor(socketPath: string): Promise; + assertSupervisorClaimCurrent(claim: Claim, fingerprint?: string): Promise; + supervisorClaims: Map; + shuttingDown: boolean; + log: ReturnType; +} + +const children: ChildProcess[] = []; +let directory: string | undefined; +let launcher: Launcher | undefined; +let launched: Promise | undefined; + +function startChild(): ChildProcess { + const child = spawn(process.execPath, ["--eval", "setInterval(() => {}, 1000)"], { stdio: "ignore" }); + children.push(child); + return child; +} + +function setup() { + directory = mkdtempSync(join(tmpdir(), "prime-launch-test-")); + vi.stubEnv("PRIME_AGENT_INTERNAL_DAEMON_SUPERVISOR_REGISTRY_DIR", directory); + const socketPath = join(directory, "supervisor.sock"); + launcher = Object.assign(Object.create(AgentDaemon.prototype), { + options: { defaultSessionConfig: { cwd: directory } }, + supervisorClaims: new Map(), + shuttingDown: false, + supervisorLaunchInProgress: false, + canConnectToSupervisor: vi.fn().mockResolvedValueOnce(false).mockResolvedValue(true), + log: vi.fn(), + }) as Launcher; + vi.spyOn(childProcesses, "spawnHidden").mockImplementation(() => startChild()); + return { daemon: launcher, socketPath }; +} + +function installClaim(daemon: Launcher, socketPath: string, child: ChildProcess) { + const pid = child.pid!; + const processStartId = getProcessStartId(pid)!; + expect(processStartId).toBeDefined(); + const generation = "authenticated-winner"; + const ownerDir = join(directory!, `${generation}.owner`); + mkdirSync(ownerDir); + const owner = { + version: 1, + role: "supervisor", + token: "test-token", + generation, + pid, + processStartId, + socketPath, + descriptorDir: directory, + agentDir: directory, + appVersion: "test", + phase: "owner", + createdAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }; + writeFileSync(join(ownerDir, "owner.json"), JSON.stringify(owner)); + const claim: Claim = { + supervisorGeneration: generation, + supervisorPid: pid, + supervisorProcessStartId: processStartId, + supervisorSocketPath: socketPath, + }; + const client: Client = { authenticated: true, socket: { destroyed: false } }; + const bound = { claim, ownerFingerprint: "do-not-trust-cached-fingerprint" }; + daemon.supervisorClaims.set(client, bound); + return { client, bound, owner, ownerDir }; +} + +async function begin(daemon: Launcher, socketPath: string): Promise { + const previous = children.length; + launched = daemon.launchReplacementSupervisor(socketPath); + await vi.waitFor(() => expect(children).toHaveLength(previous + 1)); + return children[previous]!; +} + +afterEach(async () => { + if (launcher) launcher.shuttingDown = true; + await launched; + vi.restoreAllMocks(); + for (const child of children) { + const exited = childProcesses.waitForChildProcess(child); + if (child.exitCode === null && child.signalCode === null) child.kill("SIGKILL"); + await exited; + } + children.length = 0; + vi.unstubAllEnvs(); + if (directory) rmSync(directory, { recursive: true, force: true }); + directory = undefined; + launcher = undefined; + launched = undefined; +}); + +describe("supervisor replacement launch ownership", () => { + it.each([false, true])("keeps the current winner alive (own child wins: %s)", async (ownWins) => { + const { daemon, socketPath } = setup(); + const competitor = ownWins ? undefined : startChild(); + const child = await begin(daemon, socketPath); + const winner = competitor ?? child; + const kill = vi.spyOn(child, "kill"); + const winnerKill = competitor ? vi.spyOn(competitor, "kill") : undefined; + const { bound } = installClaim(daemon, socketPath, winner); + const verify = vi.spyOn(daemon, "assertSupervisorClaimCurrent"); + await launched; + expect(verify).toHaveBeenCalledWith(bound.claim); + expect(winner.exitCode).toBeNull(); + expect(winner.signalCode).toBeNull(); + if (ownWins) expect(kill).not.toHaveBeenCalled(); + else { + expect(kill).toHaveBeenCalledOnce(); + expect(child.signalCode).toBe("SIGKILL"); + expect(winnerKill).not.toHaveBeenCalled(); + } + }); + + it.each(["stale", "replaced", "disconnected"])("does not kill on a %s claim", async (condition) => { + const { daemon, socketPath } = setup(); + const competitor = startChild(); + const child = await begin(daemon, socketPath); + const { client, bound, owner, ownerDir } = installClaim(daemon, socketPath, competitor); + const verifyCurrent = daemon.assertSupervisorClaimCurrent.bind(daemon); + let finishValidation = () => {}; + const validationFinished = new Promise((resolve) => { + finishValidation = resolve; + }); + const verify = vi.spyOn(daemon, "assertSupervisorClaimCurrent").mockImplementation(async (claim, ...rest) => { + try { + const result = await verifyCurrent(claim, ...rest); + if (condition === "replaced") daemon.supervisorClaims.set(client, { ...bound }); + if (condition === "disconnected") client.socket.destroyed = true; + return result; + } finally { + finishValidation(); + } + }); + if (condition === "stale") + writeFileSync(join(ownerDir, "owner.json"), JSON.stringify({ ...owner, processStartId: "stale" })); + const kill = vi.spyOn(child, "kill"); + await validationFinished; + await new Promise((resolve) => setImmediate(resolve)); + expect(verify).toHaveBeenCalledWith(bound.claim); + expect(kill).not.toHaveBeenCalled(); + expect(child.signalCode).toBeNull(); + daemon.shuttingDown = true; + await launched; + }); + + it.each(["shutdown", "deadline"])("does not kill without an authenticated claim at %s", async (condition) => { + const { daemon, socketPath } = setup(); + const child = await begin(daemon, socketPath); + const kill = vi.spyOn(child, "kill"); + if (condition === "shutdown") daemon.shuttingDown = true; + else vi.spyOn(Date, "now").mockReturnValue(Date.now() + 20_000); + await launched; + expect(kill).not.toHaveBeenCalled(); + expect(child.signalCode).toBeNull(); + expect(daemon.log).toHaveBeenCalledWith(expect.stringContaining("without a current authenticated supervisor")); + }); +}); diff --git a/packages/coding-agent/test/model-registry.test.ts b/packages/coding-agent/test/model-registry.test.ts index 5277ba777c..87721fa085 100644 --- a/packages/coding-agent/test/model-registry.test.ts +++ b/packages/coding-agent/test/model-registry.test.ts @@ -2,7 +2,7 @@ import { existsSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node import { tmpdir } from "node:os"; import { join } from "node:path"; import type { AnthropicMessagesCompat, Api, Context, Model, OpenAICompletionsCompat } from "@earendil-works/pi-ai"; -import { getApiProvider } from "@earendil-works/pi-ai"; +import { getApiProvider, getModels } from "@earendil-works/pi-ai"; import { getOAuthProvider, registerOAuthProvider } from "@earendil-works/pi-ai/oauth"; import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; import { AuthStorage } from "../src/core/auth-storage.js"; @@ -21,6 +21,7 @@ describe("ModelRegistry", () => { }); afterEach(() => { + vi.unstubAllGlobals(); if (tempDir && existsSync(tempDir)) { rmSync(tempDir, { recursive: true }); } @@ -603,6 +604,96 @@ describe("ModelRegistry", () => { }); }); + describe("live Prime Inference models", () => { + test("loads the cache without replacing external providers and applies local overrides", () => { + const bundled = getModels("prime-inference") as Model<"openai-completions">[]; + const catalogEntries = bundled.map((model) => ({ + id: model.id, + display_name: `Live ${model.name}`, + pricing: { input_usd_per_mtok: model.cost.input, output_usd_per_mtok: model.cost.output }, + specs: { + context_window: model.contextWindow, + max_output_tokens: model.maxTokens, + modalities: { input: model.input, output: ["text"] }, + supports_reasoning: model.reasoning, + }, + })); + catalogEntries.push({ + id: "test/live-added", + display_name: "Live Added", + pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, + specs: { + context_window: 200_000, + max_output_tokens: 20_000, + modalities: { input: ["text"], output: ["text"] }, + supports_reasoning: false, + }, + }); + writeFileSync( + join(tempDir, "prime-inference-models-cache.json"), + JSON.stringify({ object: "list", data: catalogEntries }), + ); + writeRawModelsJson({ + "prime-inference": { + baseUrl: "https://local-proxy.example.com/v1", + modelOverrides: { "test/live-added": { name: "Local Added", contextWindow: 123_456 } }, + }, + }); + + const registry = ModelRegistry.create(authStorage, modelsJsonPath); + expect(registry.find("prime-inference", "test/live-added")).toMatchObject({ + name: "Local Added", + baseUrl: "https://local-proxy.example.com/v1", + contextWindow: 123_456, + cost: { input: 1, output: 2 }, + }); + expect(getModelsForProvider(registry, "openrouter")).toHaveLength(getModels("openrouter").length); + }); + + test("restores cached authorized deployment metadata without waiting for the network", async () => { + const privateRoute = { + id: "vendor/model:deployment", + display_name: "Private Deployment", + pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, + specs: { + context_window: 200_000, + max_output_tokens: 20_000, + modalities: { input: ["text"], output: ["text"] }, + supports_reasoning: false, + }, + }; + authStorage.set("prime-inference", { + type: "api_key", + key: "prime-key", + primeTeam: { teamId: "research-team", name: "Research" }, + }); + vi.stubGlobal( + "fetch", + vi.fn( + async (_url: string | URL | Request, init?: RequestInit) => + new Response( + JSON.stringify({ data: new Headers(init?.headers).has("Authorization") ? [privateRoute] : [] }), + ), + ), + ); + const firstRegistry = ModelRegistry.create(authStorage, modelsJsonPath); + expect( + (await firstRegistry.refreshAvailableModels()).find((model) => model.id === privateRoute.id), + ).toMatchObject({ name: "Private Deployment", contextWindow: 200_000 }); + + vi.stubGlobal( + "fetch", + vi.fn(async () => { + throw new Error("offline"); + }), + ); + const restoredRegistry = ModelRegistry.create(authStorage, modelsJsonPath); + expect( + (await restoredRegistry.refreshAvailableModels()).find((model) => model.id === privateRoute.id), + ).toMatchObject({ name: "Private Deployment", contextWindow: 200_000 }); + }); + }); + describe("modelOverrides (per-model customization)", () => { test("model override applies to a single built-in model", () => { writeRawModelsJson({ @@ -1517,22 +1608,39 @@ describe("ModelRegistry", () => { writeFileSync(configPath, JSON.stringify({ api_key: "prime-test-key", team_id: "team-a" })); const cliAuth = AuthStorage.inMemory({}, { primeCliConfigPath: configPath }); const registry = ModelRegistry.create(cliAuth, modelsJsonPath); - const model = registry - .getAll() - .find((candidate) => candidate.provider === "prime-inference" && candidate.id.startsWith("internal/"))!; - expect(model).toBeDefined(); - const fetchSpy = vi - .spyOn(globalThis, "fetch") - .mockImplementation(async () => new Response(JSON.stringify({ data: [{ id: model.id }] }))); + const modelId = "internal/live-private-model"; + const fetchSpy = vi.spyOn(globalThis, "fetch").mockImplementation( + async () => + new Response( + JSON.stringify({ + data: [ + { + id: modelId, + pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, + specs: { + context_window: 200_000, + max_output_tokens: 20_000, + supports_reasoning: true, + modalities: { input: ["text"], output: ["text"] }, + }, + }, + ], + }), + ), + ); try { registry.registerProvider("unrelated-extension", { baseUrl: "https://unused.invalid" }); - expect(await registry.refreshAvailableModels()).toContainEqual(model); - expect(fetchSpy).toHaveBeenCalledTimes(1); + const model = (await registry.refreshAvailableModels()).find((candidate) => candidate.id === modelId)!; + expect(model).toBeDefined(); + expect( + fetchSpy.mock.calls.filter(([, init]) => new Headers(init?.headers).has("Authorization")), + ).toHaveLength(1); expect(registry.markProviderAuthStale("prime-inference")).toBe(true); registry.unregisterProvider("unrelated-extension"); await expect(registry.canUseModel(model, { assumeAuthConfigured: true })).resolves.toBe(true); - await registry.refreshAvailableModels(); + await Promise.all([registry.refreshAvailableModels(), registry.refreshAvailableModels()]); + expect(registry.find("prime-inference", modelId)).toEqual(model); expect(cliAuth.getProviderHeaders("prime-inference")).toEqual({ "X-Prime-Team-ID": "team-a" }); expect(registry.hasConfiguredAuth(model)).toBe(false); @@ -1564,8 +1672,11 @@ describe("ModelRegistry", () => { break; } registry.refresh(); + expect(registry.find("prime-inference", modelId)).toBeUndefined(); await expect(registry.canUseModel(model, { assumeAuthConfigured: true })).resolves.toBe(false); - expect(fetchSpy).toHaveBeenCalledTimes(1); + expect( + fetchSpy.mock.calls.filter(([, init]) => new Headers(init?.headers).has("Authorization")), + ).toHaveLength(1); } finally { fetchSpy.mockRestore(); vi.unstubAllEnvs(); @@ -1598,7 +1709,12 @@ describe("ModelRegistry", () => { test("concurrent stale-auth refreshes do not drop preserved entitlements", async () => { authStorage.setRuntimeApiKey("prime-inference", "prime-key"); const registry = ModelRegistry.create(authStorage, modelsJsonPath); - const internals = registry as unknown as { authorizedPrivatePrimeInferenceModelIds: Set }; + vi.spyOn(authStorage, "getProviderHeaders").mockReturnValue({ "X-Prime-Team-ID": "team-a" }); + const internals = registry as unknown as { + authorizedPrivatePrimeInferenceModelIds: Set; + authorizedPrivatePrimeInferenceTeamId: string | undefined; + }; + internals.authorizedPrivatePrimeInferenceTeamId = "team-a"; internals.authorizedPrivatePrimeInferenceModelIds.add("internal/private-model"); expect(registry.markProviderAuthStale("prime-inference")).toBe(true); diff --git a/packages/coding-agent/test/model-selector-actions.test.ts b/packages/coding-agent/test/model-selector-actions.test.ts index a8320a4e71..88b6d0b4d5 100644 --- a/packages/coding-agent/test/model-selector-actions.test.ts +++ b/packages/coding-agent/test/model-selector-actions.test.ts @@ -1,7 +1,10 @@ -import { setKeybindings, type TUI } from "@earendil-works/pi-tui"; +import { parsePrimeInferenceModelCatalog } from "@earendil-works/pi-ai"; +import { setKeybindings, TUI } from "@earendil-works/pi-tui"; import stripAnsi from "strip-ansi"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { VirtualTerminal } from "../../tui/test/virtual-terminal.js"; import { KeybindingsManager } from "../src/core/keybindings.js"; +import { buildPrimeInferenceModels } from "../src/core/prime-inference-model-catalog.js"; import { ModelSelectorComponent } from "../src/modes/interactive/components/model-selector.js"; import { initTheme } from "../src/modes/interactive/theme/theme.js"; import { createHarness, type Harness } from "./suite/harness.js"; @@ -37,6 +40,57 @@ describe("ModelSelectorComponent", () => { } }); + it("does not write catalog OSC actions to the terminal", async () => { + const harness = await createHarness({ models: [{ id: "base", name: "Base", reasoning: true }] }); + harnesses.push(harness); + const osc = "\x1b]52;c;VFJJQUdF\x07"; + const entry = { + id: "vendor/模型", + display_name: `模型 é${osc}`, + pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, + specs: { + context_window: 1000, + max_output_tokens: 100, + supports_reasoning: false, + modalities: { input: ["text"], output: ["text"] }, + }, + }; + const models = buildPrimeInferenceModels( + [], + parsePrimeInferenceModelCatalog({ + data: [entry, { ...entry, id: `vendor/bad${osc}` }], + }), + )!; + expect(models.map((model) => model.id)).toEqual([entry.id]); + const terminal = new VirtualTerminal(120, 40); + const write = vi.spyOn(terminal, "write"); + const tui = new TUI(terminal); + const selector = new ModelSelectorComponent( + tui, + undefined, + harness.session.modelRegistry, + [], + () => {}, + () => {}, + undefined, + { + availableModels: models, + configuredProviders: new Set(["prime-inference"]), + }, + ); + tui.addChild(selector); + tui.start(); + try { + await terminal.waitForRender(); + const output = write.mock.calls.map(([data]) => data).join(""); + expect(output).not.toContain(osc); + expect(output).toContain("模型 é"); + expect(output).toContain(entry.id); + } finally { + tui.stop(); + } + }); + it("explains model authentication without a provider shortcut", async () => { const harness = await createHarness({ models: [{ id: "faux-1", name: "One", reasoning: true }], diff --git a/packages/coding-agent/test/prime-inference-model-catalog.test.ts b/packages/coding-agent/test/prime-inference-model-catalog.test.ts new file mode 100644 index 0000000000..e5a033a00e --- /dev/null +++ b/packages/coding-agent/test/prime-inference-model-catalog.test.ts @@ -0,0 +1,195 @@ +import { mkdtempSync, readFileSync, rmSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import type { Model } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, test, vi } from "vitest"; +import { + buildPrimeInferenceModels, + mergePrimeInferenceModels, + PRIME_INFERENCE_BASE_URL, + refreshPrimeInferenceModels, +} from "../src/core/prime-inference-model-catalog.js"; +import { + fetchAuthorizedPrivatePrimeInferenceModels, + isPrivatePrimeInferenceModel, +} from "../src/core/prime-inference-models.js"; + +const directories: string[] = []; +const model = (id: string, provider = "prime-inference"): Model<"openai-completions"> => ({ + id, + name: `Bundled ${id}`, + api: "openai-completions", + provider, + baseUrl: provider === "prime-inference" ? PRIME_INFERENCE_BASE_URL : "https://example.com/v1", + reasoning: true, + thinkingLevelMap: { high: "high" }, + input: ["text"], + cost: { input: 9, output: 10, cacheRead: 0.9, cacheWrite: 11.25 }, + contextWindow: 100_000, + maxTokens: 10_000, + featured: true, + compat: { supportsDeveloperRole: false, maxTokensField: "max_tokens" }, +}); + +const entry = (id: string, overrides: Record = {}) => ({ + id, + input: 1, + output: 2, + contextWindow: 200_000, + maxTokens: 20_000, + vision: true, + reasoning: false, + ...overrides, +}); + +const payloadEntry = ( + id: string, + specs: unknown = { + context_window: 200_000, + max_output_tokens: 20_000, + modalities: { input: ["text", "image"], output: ["text"] }, + supports_reasoning: false, + }, +) => ({ + id, + display_name: `Live ${id}`, + pricing: { input_usd_per_mtok: 1, output_usd_per_mtok: 2 }, + specs, +}); + +const response = (...data: unknown[]) => new Response(JSON.stringify({ object: "list", data })); + +afterEach(() => { + for (const directory of directories.splice(0)) rmSync(directory, { recursive: true, force: true }); +}); + +describe("Prime Inference model catalog", () => { + test("uses live metadata while retaining bundled client compatibility", () => { + const [live] = buildPrimeInferenceModels( + [model("vendor/model")], + [entry("vendor/model", { name: "Live Name", cacheRead: 0.1, cacheWrite: 1.25, maxTokens: 250_000 })], + ) ?? [undefined]; + expect(live).toMatchObject({ + id: "vendor/model", + name: "Live Name", + baseUrl: PRIME_INFERENCE_BASE_URL, + api: "openai-completions", + provider: "prime-inference", + reasoning: false, + input: ["text", "image"], + cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1.25 }, + contextWindow: 200_000, + maxTokens: 200_000, + thinkingLevelMap: { high: "high" }, + featured: true, + compat: { supportsDeveloperRole: false, maxTokensField: "max_tokens" }, + }); + expect(live).not.toHaveProperty("headers"); + }); + + test("adds complete new models and skips incomplete unknown models", () => { + const models = + buildPrimeInferenceModels( + [model("bundled")], + [entry("new/complete"), { id: "new/incomplete", input: 1, output: 2 }], + { minimumModels: 0 }, + ) ?? []; + expect(models.map(({ id }) => id)).toEqual(["new/complete"]); + }); + + test("retains bundled specs when an existing live entry has none", () => { + const [live] = buildPrimeInferenceModels( + [model("vendor/model")], + [{ id: "vendor/model", name: "Renamed", input: 1, output: 2 }], + ) ?? [undefined]; + expect(live).toMatchObject({ name: "Renamed", contextWindow: 100_000, maxTokens: 10_000, reasoning: true }); + }); + + test("filters private routes and measures coverage against bundled models", () => { + const bundled = [model("one"), model("two"), model("three")]; + expect( + buildPrimeInferenceModels(bundled, [ + entry("internal/private"), + entry("dev/private"), + entry("poolside/model:deployment"), + entry("one"), + ]), + ).toBeUndefined(); + expect( + buildPrimeInferenceModels(bundled, [entry("new/one"), entry("new/two"), entry("new/three")]), + ).toBeUndefined(); + }); + + test("requires authorization for private prefixes and deployment routes", () => { + for (const id of ["internal/model", "INTERNAL/model", "dev/model", "vendor/model:deployment"]) { + expect(isPrivatePrimeInferenceModel(model(id))).toBe(true); + } + expect(isPrivatePrimeInferenceModel(model("public/model"))).toBe(false); + expect(isPrivatePrimeInferenceModel(model("vendor/model:deployment", "openrouter"))).toBe(false); + }); + + test("replaces only the Prime Inference provider list", () => { + const external = model("external", "openrouter"); + const live = model("live"); + expect(mergePrimeInferenceModels([external, model("removed")], [live])).toEqual([external, live]); + }); + + test("caches valid responses and falls back to the cache", async () => { + const directory = mkdtempSync(join(tmpdir(), "prime-models-")); + directories.push(directory); + const cachePath = join(directory, "cache.json"); + const bundled = [model("vendor/model")]; + const fetched = await refreshPrimeInferenceModels(cachePath, bundled, { + fetchFn: vi.fn(async () => response(payloadEntry("vendor/model"))), + }); + expect(fetched?.[0]?.name).toBe("Live vendor/model"); + expect(JSON.parse(readFileSync(cachePath, "utf8")).data).toHaveLength(1); + const fallback = await refreshPrimeInferenceModels(cachePath, bundled, { + fetchFn: vi.fn(async () => { + throw new Error("offline"); + }), + }); + expect(fallback?.[0]?.name).toBe("Live vendor/model"); + }); + + test("uses authenticated responses only for private routes with complete metadata", async () => { + const fetchFn = vi.fn(async (_url: string | URL | Request, init?: RequestInit) => { + expect(new Headers(init?.headers).get("Authorization")).toBe("Bearer secret"); + expect(new Headers(init?.headers).get("X-Prime-Team-ID")).toBe("team"); + return response( + payloadEntry("public/model"), + payloadEntry("internal/model"), + payloadEntry("dev/model"), + payloadEntry("poolside/model:deployment"), + payloadEntry("internal/incomplete", null), + ); + }); + const models = await fetchAuthorizedPrivatePrimeInferenceModels( + "secret", + { "X-Prime-Team-ID": "team" }, + new Set(["public/model"]), + fetchFn, + ); + expect(models.map(({ id }) => id)).toEqual(["internal/model", "dev/model", "poolside/model:deployment"]); + }); + + test("uses bundled metadata to authorize an existing private route", async () => { + const models = await fetchAuthorizedPrivatePrimeInferenceModels( + "secret", + { "X-Prime-Team-ID": "team" }, + new Set(), + vi.fn(async () => response({ id: "internal/glm-5.2-fast" })), + ); + expect(models.map(({ id }) => id)).toEqual(["internal/glm-5.2-fast"]); + }); + + test("treats rejected authenticated requests as no private access", async () => { + const models = await fetchAuthorizedPrivatePrimeInferenceModels( + "bad", + { "X-Prime-Team-ID": "team" }, + new Set(), + vi.fn(async () => new Response(null, { status: 403 })), + ); + expect(models).toEqual([]); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/2053-background-bash-passivation.test.ts b/packages/coding-agent/test/suite/regressions/2053-background-bash-passivation.test.ts new file mode 100644 index 0000000000..3577a046b1 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/2053-background-bash-passivation.test.ts @@ -0,0 +1,268 @@ +import { existsSync } from "node:fs"; +import { resolve } from "node:path"; +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { AgentSession } from "../../../src/core/agent-session.js"; +import { + BASH_ACTIVITY_DISPLAY_MIME, + createDeferred, + type HostRequestHandlers, + ReplKernelManager, +} from "../../../src/core/kernel/index.js"; +import { ASYNC_BASH_COMPLETION_CUSTOM_TYPE } from "../../../src/core/messages.js"; +import { canEvictWorker, canPassivateSession } from "../../../src/core/session-action-store.js"; +import { IpythonKernelProvisioner } from "../../../src/core/tools/ipython.js"; +import { createHarness, type Harness } from "../harness.js"; + +const runtimeDir = resolve(__dirname, "../../../../../prime-agent-runtime"); +const python = resolve(runtimeDir, ".venv/bin/python"); +const describeRuntime = existsSync(python) ? describe : describe.skip; + +interface KernelSession { + _ipythonKernelProvisioner?: IpythonKernelProvisioner; + _createKernelHostHandlers(): HostRequestHandlers; +} + +function evictionSnapshot(session: AgentSession) { + return { + isSessionActive: session.isSessionActive, + attachedClients: 0, + hasRegisteredCronJob: false, + lastActivityAt: 0, + }; +} + +function passivationAllowed(session: AgentSession): boolean { + return canPassivateSession( + { + ...evictionSnapshot(session), + hasParent: true, + hasNonPassiveDescendants: false, + isHydrating: false, + }, + 1, + 120_000, + ); +} + +describeRuntime("#2053 background kernel bash residency", () => { + let harness: Harness | undefined; + let manager: ReplKernelManager | undefined; + + afterEach(async () => { + await manager?.shutdown(); + harness?.cleanup(); + vi.restoreAllMocks(); + }); + + async function start( + beforeCompletion?: () => Promise, + withConfiguredAuth = true, + ): Promise<{ session: AgentSession; kernel: ReplKernelManager }> { + harness = await createHarness({ tools: [], rlmDepth: 1, withConfiguredAuth }); + const session = harness.session; + const internals = session as unknown as KernelSession; + const hostHandlers = internals._createKernelHostHandlers(); + const completed = hostHandlers["bash.completed"]!; + hostHandlers["bash.completed"] = async (payload) => { + await beforeCompletion?.(); + return completed(payload); + }; + manager = new ReplKernelManager({ + python, + cwd: harness.tempDir, + env: { PYTHONPATH: resolve(runtimeDir, "src") }, + hostHandlers, + }); + const provisioner = new IpythonKernelProvisioner(harness.tempDir); + vi.spyOn(provisioner, "manager", "get").mockReturnValue(manager); + internals._ipythonKernelProvisioner = provisioner; + return { session, kernel: manager }; + } + + it("keeps a managed BashHandle resident after its creating cell without blocking a new turn", async () => { + const { session, kernel } = await start(); + expect(passivationAllowed(session)).toBe(true); + const started = await kernel.execute("from rlm import bash\nhandle = bash('sleep 600')\nhandle.pid"); + expect(started.status).toBe("ok"); + expect(Number(started.result)).toBeGreaterThan(0); + expect(session.isStreaming).toBe(false); + expect(session.isBashRunning).toBe(false); + expect(passivationAllowed(session)).toBe(false); + expect( + canEvictWorker( + { + lifecycle: "ready", + isConnected: true, + isStopping: false, + hasOwnerClient: false, + isPreparingUpdateRestart: false, + hasWakeBlindSchedule: false, + sessions: [evictionSnapshot(session)], + }, + 1, + 120_000, + ), + ).toBe(false); + + harness!.setResponses([fauxAssistantMessage("Other work can continue.")]); + await session.prompt("Do other work while the command runs."); + expect(session.getLastAssistantText()).toBe("Other work can continue."); + expect(passivationAllowed(session)).toBe(false); + + harness!.setResponses([fauxAssistantMessage("Inspected the completed command.")]); + await kernel.execute("handle.kill()"); + await vi.waitFor(() => expect(session.getLastAssistantText()).toBe("Inspected the completed command.")); + await session.waitForIdle(); + await vi.waitFor(() => expect(passivationAllowed(session)).toBe(true)); + expect( + session.messages.filter( + (message) => message.role === "custom" && message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ), + ).toHaveLength(1); + }); + + it("keeps concurrent handles resident until the final completion and clears on kernel teardown", async () => { + const { session, kernel } = await start(); + await kernel.execute("from rlm import bash\nfirst = bash('sleep 600')\nsecond = bash('sleep 600')"); + expect(passivationAllowed(session)).toBe(false); + harness!.setResponses([fauxAssistantMessage("First command finished.")]); + await kernel.execute("first.kill()"); + await vi.waitFor(() => expect(session.getLastAssistantText()).toBe("First command finished.")); + await session.waitForIdle(); + expect(passivationAllowed(session)).toBe(false); + + await kernel.kill(); + expect(passivationAllowed(session)).toBe(true); + expect( + session.messages.filter( + (message) => message.role === "custom" && message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ), + ).toHaveLength(1); + }); + + it("releases awaited commands without a completion follow-up", async () => { + const { session, kernel } = await start(); + const result = await kernel.execute("from rlm import bash\n(await bash('printf done')).output"); + expect(result.status).toBe("ok"); + expect(result.result).toContain("done"); + await vi.waitFor(() => expect(passivationAllowed(session)).toBe(true)); + expect(session.messages).toEqual([]); + }); + + it("defers one completion across admission pauses without issuing another host request", async () => { + const beforeCompletion = vi.fn(async () => {}); + const { session, kernel } = await start(beforeCompletion); + const firstPause = session.acquireSessionInputPause(); + const secondPause = session.acquireSessionInputPause(); + try { + await kernel.execute("from rlm import bash\nhandle = bash('printf done')"); + await vi.waitFor(() => expect(session.hasPendingAdmissionWaiters).toBe(true)); + expect(kernel.hasBackgroundWork).toBe(true); + expect(session.messages).toEqual([]); + firstPause.release(); + await kernel.execute("42"); + expect(session.hasPendingAdmissionWaiters).toBe(true); + expect(kernel.hasBackgroundWork).toBe(true); + harness!.setResponses([fauxAssistantMessage("Completion accepted after pause.")]); + secondPause.release(); + await vi.waitFor(() => expect(session.getLastAssistantText()).toBe("Completion accepted after pause.")); + await session.waitForIdle(); + await vi.waitFor(() => expect(kernel.hasBackgroundWork).toBe(false)); + expect(beforeCompletion).toHaveBeenCalledTimes(1); + expect( + session.messages.filter( + (message) => message.role === "custom" && message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ), + ).toHaveLength(1); + } finally { + firstPause.release(); + secondPause.release(); + } + }); + + it("stops waiting for admission when the session is disposed", async () => { + const beforeCompletion = vi.fn(async () => {}); + const { session, kernel } = await start(beforeCompletion); + const pause = session.acquireSessionInputPause(); + try { + await kernel.execute("from rlm import bash\nhandle = bash('printf done')"); + await vi.waitFor(() => expect(session.hasPendingAdmissionWaiters).toBe(true)); + session.dispose(); + await vi.waitFor(() => expect(session.hasPendingAdmissionWaiters).toBe(false)); + await vi.waitFor(() => expect(kernel.hasBackgroundWork).toBe(false)); + expect(beforeCompletion).toHaveBeenCalledTimes(1); + expect(session.messages).toEqual([]); + } finally { + pause.release(); + } + }); + + it("reports a terminal readiness failure without retrying or retaining completed work", async () => { + const beforeCompletion = vi.fn(async () => {}); + const { session, kernel } = await start(beforeCompletion, false); + await kernel.execute("from rlm import bash\nhandle = bash('printf done')"); + await vi.waitFor(() => expect(beforeCompletion).toHaveBeenCalledTimes(1)); + await vi.waitFor(() => expect(kernel.hasBackgroundWork).toBe(false)); + expect(session.messages).toEqual([]); + const nextCell = await kernel.execute("handle.poll().exit_code"); + expect(nextCell.result).toBe("0"); + expect(nextCell.backgroundOutput).toContain("was not accepted"); + expect(beforeCompletion).toHaveBeenCalledTimes(1); + }); + + it("holds residency until the completion follow-up is accepted", async () => { + const reached = createDeferred(); + const release = createDeferred(); + try { + const { session, kernel } = await start(async () => { + reached.resolve(); + await release.promise; + }); + await kernel.execute("from rlm import bash\nhandle = bash('sleep 600')"); + await kernel.execute("handle.kill()"); + await reached.promise; + expect(session.messages).toEqual([]); + expect(passivationAllowed(session)).toBe(false); + harness!.setResponses([fauxAssistantMessage("Completion accepted.")]); + release.resolve(); + await vi.waitFor(() => expect(session.getLastAssistantText()).toBe("Completion accepted.")); + await session.waitForIdle(); + await vi.waitFor(() => expect(passivationAllowed(session)).toBe(true)); + } finally { + release.resolve(); + } + }); +}); + +describe("kernel bash activity validation", () => { + it("ignores unrelated display data and rejects malformed or mismatched releases", async () => { + const kernel = new ReplKernelManager({}); + const deliver = (data: Record) => + (kernel as unknown as { handleEvent(event: Record): void }).handleEvent({ + event: "display", + id: "old-cell", + data, + }); + const activity = { id: "a".repeat(32), pid: 42, active: true }; + deliver({ [BASH_ACTIVITY_DISPLAY_MIME]: activity }); + expect(kernel.hasBackgroundWork).toBe(true); + for (const invalid of [ + { ...activity, pid: 0, active: false }, + { ...activity, pid: -1, active: false }, + { ...activity, pid: 42.5, active: false }, + { ...activity, pid: 43, active: false }, + { ...activity, id: "", active: false }, + { ...activity, id: "b".repeat(32), active: false }, + { ...activity, active: "false" }, + ]) { + deliver({ [BASH_ACTIVITY_DISPLAY_MIME]: invalid }); + expect(kernel.hasBackgroundWork).toBe(true); + } + deliver({ "application/json": { ...activity, active: false } }); + expect(kernel.hasBackgroundWork).toBe(true); + deliver({ [BASH_ACTIVITY_DISPLAY_MIME]: { ...activity, active: false } }); + expect(kernel.hasBackgroundWork).toBe(false); + await kernel.shutdown(); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/2068-shell-message-steering.test.ts b/packages/coding-agent/test/suite/regressions/2068-shell-message-steering.test.ts new file mode 100644 index 0000000000..013abe534a --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/2068-shell-message-steering.test.ts @@ -0,0 +1,142 @@ +import type { AgentTool } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { afterEach, beforeAll, describe, expect, it } from "vitest"; +import { isAgentSessionMessage } from "../../../src/core/agent-messages.js"; +import { createDeferred, type HostRequestHandlers } from "../../../src/core/kernel/index.js"; +import { + ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ASYNC_BASH_COMPLETION_PREVIEW_LABEL, + createAsyncBashCompletionMessage, +} from "../../../src/core/messages.js"; +import { InjectedPromptMessageComponent } from "../../../src/modes/interactive/components/injected-prompt-message.js"; +import { formatQueuedMessagePreview } from "../../../src/modes/interactive/interactive-mode.js"; +import { initTheme } from "../../../src/modes/interactive/theme/theme.js"; +import { createHarness, getMessageText, type Harness } from "../harness.js"; + +interface KernelSession { + _createKernelHostHandlers(): HostRequestHandlers; +} + +const completion = { pid: 42, command: "npm test", exitCode: 0 }; + +function completeShell(harness: Harness) { + return (harness.session as unknown as KernelSession)._createKernelHostHandlers()["bash.completed"]!(completion); +} + +function shellMessages(harness: Harness) { + return harness.session.messages.filter( + (message) => message.role === "custom" && message.customType === ASYNC_BASH_COMPLETION_CUSTOM_TYPE, + ); +} + +describe("#2068 shell message steering", () => { + const harnesses: Harness[] = []; + + beforeAll(() => initTheme("dark")); + afterEach(() => { + for (const harness of harnesses.splice(0)) harness.cleanup(); + }); + + it("consumes shell steering at the next tool boundary without waiting for the run to become idle", async () => { + const started = createDeferred(); + const release = createDeferred(); + const order: string[] = []; + let consumed = { streaming: false, completedRuns: -1, text: "" }; + const tool: AgentTool = { + name: "wait", + label: "Wait", + description: "Hold the current tool until released", + parameters: Type.Object({}), + execute: async () => { + order.push("tool-start"); + started.resolve(); + await release.promise; + order.push("tool-end"); + return { content: [{ type: "text", text: "released" }], details: {} }; + }, + }; + const harness = await createHarness({ tools: [tool] }); + harnesses.push(harness); + harness.setResponses([ + fauxAssistantMessage(fauxToolCall("wait", {}), { stopReason: "toolUse" }), + (context) => { + consumed = { + streaming: harness.session.isStreaming, + completedRuns: harness.eventsOfType("agent_end").length, + text: getMessageText(context.messages.at(-1)), + }; + order.push("shell-consumed"); + return fauxAssistantMessage("Inspected the shell result."); + }, + ]); + const original = harness.session.prompt("Continue working."); + try { + await started.promise; + await expect(completeShell(harness)).resolves.toEqual({}); + order.push("shell-queued"); + expect(order).toEqual(["tool-start", "shell-queued"]); + expect(harness.session.getFollowUpMessages()).toEqual([]); + expect(harness.session.getSteeringMessages()).toHaveLength(1); + const queued = harness.session.getSessionActionRecoverySnapshot().actions; + expect(queued).toContainEqual(expect.objectContaining({ delivery: "next_turn_boundary" })); + const preview = harness.session.getSteeringMessagePreviews()[0]!; + expect(preview).toBe("Shell message received: pid 42, exit 0"); + expect(formatQueuedMessagePreview(preview, "Steering")).toBe(preview); + } finally { + release.resolve(); + await original; + } + await harness.session.waitForIdle(); + expect(order).toEqual(["tool-start", "shell-queued", "tool-end", "shell-consumed"]); + expect(consumed).toMatchObject({ streaming: true, completedRuns: 1 }); + expect(harness.eventsOfType("agent_end")[0]!.messages.filter((message) => message.role === "assistant")).toEqual([ + expect.objectContaining({ stopReason: "toolUse" }), + ]); + expect(consumed.text).toContain("Shell message received.\nSource: bash"); + expect(consumed.text).toContain("pid 42, exit code 0"); + expect(shellMessages(harness)).toHaveLength(1); + expect(harness.eventsOfType("agent_start")).toHaveLength(2); + expect(harness.eventsOfType("agent_end")).toHaveLength(2); + }); + + it("resumes an idle session once while retaining the distinct shell message identity", async () => { + const harness = await createHarness({ tools: [] }); + harnesses.push(harness); + let requests = 0; + harness.setResponses([ + (context) => { + requests++; + expect(getMessageText(context.messages.at(-1))).toContain("Source: bash"); + return fauxAssistantMessage("Inspected the shell result."); + }, + ]); + await completeShell(harness); + await harness.session.waitForIdle(); + expect(requests).toBe(1); + expect(shellMessages(harness)).toHaveLength(1); + expect(isAgentSessionMessage(shellMessages(harness)[0]!)).toBe(false); + expect(harness.session.getSteeringMessages()).toEqual([]); + expect(harness.session.getFollowUpMessages()).toEqual([]); + }); + + it("renders shell-specific queue and transcript labels without generic delivery prefixes", () => { + const message = createAsyncBashCompletionMessage(completion); + const preview = `${ASYNC_BASH_COMPLETION_PREVIEW_LABEL}: pid 42, exit 0`; + for (const delivery of ["Steering", "Follow-up"] as const) { + expect(formatQueuedMessagePreview(preview, delivery)).toBe(preview); + } + const component = new InjectedPromptMessageComponent(message); + const render = () => + component + .render(120) + .join("\n") + .replace(/\u001b\[[0-9;]*m/g, ""); + expect(render()).toContain("◆ Shell message received · pid 42 · exit 0"); + expect(render()).not.toMatch(/Follow-up:|Steering:|Agent message received/); + component.setExpanded(true); + expect(render()).toContain("Source: bash"); + expect(render()).toContain("npm test"); + expect(render()).toContain("Inspect the saved BashHandle"); + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts b/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts index 263bdd89bb..04b65316be 100644 --- a/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts +++ b/packages/coding-agent/test/suite/regressions/4491-provider-stale-after-401.test.ts @@ -1,6 +1,6 @@ import type { AgentEvent } from "@earendil-works/pi-agent-core"; import { type AssistantMessage, fauxAssistantMessage } from "@earendil-works/pi-ai"; -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import type { AgentSessionRuntime } from "../../../src/core/agent-session-runtime.js"; import { InProcessAgentConnection } from "../../../src/modes/agent-connection/in-process-agent-connection.js"; import { createHarness, type Harness } from "../harness.js"; @@ -41,6 +41,7 @@ describe("issue #4491 provider stale after repeated 401", () => { const harnesses: Harness[] = []; afterEach(() => { + vi.restoreAllMocks(); while (harnesses.length > 0) { harnesses.pop()?.cleanup(); } @@ -295,6 +296,7 @@ describe("issue #4491 provider stale after repeated 401", () => { const harness = await createHarness(privateModelHarnessOptions); harnesses.push(harness); const registry = harness.session.modelRegistry; + vi.spyOn(harness.authStorage, "getProviderHeaders").mockReturnValue({ "X-Prime-Team-ID": "test-team" }); lockOutProvider(harness, "prime-inference"); const privateModel = harness.models.find((model) => model.id === "internal/private-model"); expect(privateModel).toBeDefined(); @@ -303,8 +305,12 @@ describe("issue #4491 provider stale after repeated 401", () => { await expect(harness.session.setModel(privateModel!)).rejects.toThrow("not available"); expect(registry.getProviderAuthStatus("prime-inference")).toMatchObject({ source: "stale" }); - const internals = registry as unknown as { authorizedPrivatePrimeInferenceModelIds: Set }; + const internals = registry as unknown as { + authorizedPrivatePrimeInferenceModelIds: Set; + authorizedPrivatePrimeInferenceTeamId: string | undefined; + }; internals.authorizedPrivatePrimeInferenceModelIds.add("internal/private-model"); + internals.authorizedPrivatePrimeInferenceTeamId = "test-team"; // Refreshes during the stale window run keyless; they must preserve the // cached entitlements the explicit re-selection validates against. await registry.refreshAvailableModels(); diff --git a/packages/coding-agent/test/suite/regressions/4645-internal-glm.test.ts b/packages/coding-agent/test/suite/regressions/4645-internal-glm.test.ts index 5a87100a7f..1e23b26c78 100644 --- a/packages/coding-agent/test/suite/regressions/4645-internal-glm.test.ts +++ b/packages/coding-agent/test/suite/regressions/4645-internal-glm.test.ts @@ -52,6 +52,7 @@ describe("ENG-4645 internal GLM configuration", () => { }); expect(fetchMock).toHaveBeenCalledWith("https://api.pinference.ai/api/v1/models", { headers: { + accept: "application/json", Authorization: "Bearer prime-key", "X-Prime-Team-ID": "engineering-team", }, diff --git a/packages/coding-agent/test/suite/regressions/502-unified-session-view.test.ts b/packages/coding-agent/test/suite/regressions/502-unified-session-view.test.ts index 9e891ebbc0..b54110daa3 100644 --- a/packages/coding-agent/test/suite/regressions/502-unified-session-view.test.ts +++ b/packages/coding-agent/test/suite/regressions/502-unified-session-view.test.ts @@ -1,6 +1,7 @@ import stripAnsi from "strip-ansi"; import { describe, expect, test, vi } from "vitest"; import { AgentsViewMode } from "../../../src/modes/agents-view/agents-view-mode.js"; +import { buildUnifiedSessionIndex } from "../../../src/modes/agents-view/agents-view-state.js"; import type { SessionSummary } from "../../../src/modes/daemon/daemon-session-list.js"; import { initTheme } from "../../../src/modes/interactive/theme/theme.js"; import { createDeferred as deferred } from "../scheduling.js"; @@ -304,7 +305,7 @@ describe("#502 unified session view regressions", () => { expect(harness.setStatusMessage).not.toHaveBeenCalled(); }); - test("a missing selection anchor blocks open only until both catalogs settle", () => { + test("a pending selection anchor never blocks opening the visible row", () => { const finish = vi.fn(); const fallback = summary("fallback"); const harness = { @@ -314,21 +315,26 @@ describe("#502 unified session view regressions", () => { selectedActiveSessionId: undefined as string | undefined, selectedRowIdentity: "identity-intended", rows: [{ selectable: true, kind: "agent", summary: fallback }], + unifiedRecords: [], + unifiedIndex: buildUnifiedSessionIndex([]), isPendingDeleteRow: () => false, setStatusMessage: vi.fn(), finish, }; + // Enter acts on the row under the cursor even while the anchor waits. privateMethod<(this: typeof harness) => void>("openSelected").call(harness); - expect(finish).not.toHaveBeenCalled(); + expect(finish).toHaveBeenCalledWith(expect.objectContaining({ type: "open", summary: fallback })); + expect(harness.setStatusMessage).not.toHaveBeenCalled(); + + // Untouched, the anchor restore still waits for the saved catalog... privateMethod<(this: typeof harness) => void>("resolveMissingSelectionAnchor").call(harness); expect(harness.selectionAnchorPending).toBe(true); harness.savedCatalogRefreshPending = false; privateMethod<(this: typeof harness) => void>("resolveMissingSelectionAnchor").call(harness); - // Open unblocks on the visible fallback row... expect(harness.selectionAnchorPending).toBe(false); expect(harness.selectedActiveSessionId).toBe(fallback.activeSessionId ?? fallback.id); - // ...but the restored anchor identity survives so a late poll can still re-anchor. + // ...and the restored anchor identity survives so a late poll can still re-anchor. expect(harness.selectedRowIdentity).toBe("identity-intended"); }); test("rename uses the captured row after refresh removes it", async () => { @@ -409,7 +415,8 @@ describe("#502 unified session view regressions", () => { expect(filtered.map((record) => record.identity)).toEqual(["match"]); }); - test("inactive rows give usage and age their full responsive cell", () => { + test("inactive rows keep total cost and age visible in a narrow row", () => { + initTheme("dark"); const inactive = { kind: "agent" as const, section: "inactive" as const, @@ -446,10 +453,10 @@ describe("#502 unified session view regressions", () => { 50, ), ); - expect(rendered).toMatch(/↑0\s+↓0 ·\s+\$0\.00 ·\s+0 ·\s+\$0\.00 ·\s+2h\s*$/); + expect(rendered).toMatch(/\$0\.00\s+2h\s*$/); }); - test("scoped subagent rows keep model and effort ahead of summaries", () => { + test("rows keep compact model IDs visible on every row kind", () => { initTheme("dark"); const subagent = { // Direct children in a scoped Agents View render as agent rows while @@ -492,26 +499,43 @@ describe("#502 unified session view regressions", () => { ); const full = render(160); - expect(full).toContain( - "Inspect agents view · prime-inference/gpt-5.6-terra:high · Investigate a variable background status", - ); + expect(full).toMatch(/Inspect agents view\s+gpt-5\.6-terra\s+Investigate a variable background status/); + for (const width of [60, 80, 120]) { + expect(render(width)).toContain("gpt-5.6-terra"); + expect(render(width)).toHaveLength(width); + } const narrow = render(100); - expect(narrow).toContain("prime-inference/gpt-5.6-terra:high"); - expect(narrow).not.toContain("Investigate a variable background status"); + expect(narrow).toContain("gpt-5.6-terra"); + expect(narrow).not.toContain("prime-inference/"); subagent.summary.summary = ""; - expect(render(100)).toContain("Inspect agents view · prime-inference/gpt-5.6-terra:high"); + expect(render(100)).toMatch(/Inspect agents view\s+gpt-5\.6-terra/); // Older daemons identify subagents through persisted linkage instead of runtimeKind. subagent.summary.runtimeKind = undefined; subagent.summary.rlmChildId = "effort-child"; - expect(render(100)).toContain("Inspect agents view · prime-inference/gpt-5.6-terra:high"); + expect(render(100)).toMatch(/Inspect agents view\s+gpt-5\.6-terra/); subagent.summary.thinkingLevel = "off"; subagent.summary.summary = "A later summary"; - expect(render(120)).toContain("Inspect agents view · prime-inference/gpt-5.6-terra · A later summary"); + expect(render(120)).toMatch(/Inspect agents view\s+gpt-5\.6-terra\s+A later summary/); expect(render(120)).not.toContain(":off"); + // Top-level sessions show the same label; the model cell is not subagent-only. + subagent.summary.runtimeKind = "top-level"; + subagent.summary.rlmChildId = undefined; + expect(render(120)).toMatch(/Inspect agents view\s+gpt-5\.6-terra\s+A later summary/); + + // Pending delete replaces the suffixes, model label included. + const pendingDelete = { ...harness, isPendingDeleteRow: () => true, getPendingDeleteTitle: () => "delete?" }; + expect( + stripAnsi( + privateMethod<(this: typeof pendingDelete, row: typeof subagent, width: number) => string>( + "renderRow", + ).call(pendingDelete, subagent, 120), + ), + ).not.toContain("gpt-5.6-terra"); + expect(render(20)).toHaveLength(20); }); }); diff --git a/packages/coding-agent/test/system-prompt.test.ts b/packages/coding-agent/test/system-prompt.test.ts index ca39619eb1..41b038c672 100644 --- a/packages/coding-agent/test/system-prompt.test.ts +++ b/packages/coding-agent/test/system-prompt.test.ts @@ -168,6 +168,7 @@ describe("buildRlmPrompt", () => { }); expect(prompt).toContain("Use `bash()` to invoke programs, not to write shell programs"); + expect(prompt).toContain("A `bash()` handle left running beyond its creating cell sends a completion follow-up"); }); test("documents preferring Python for reading and searching files when ipython is active", () => { diff --git a/packages/coding-agent/test/user-message.test.ts b/packages/coding-agent/test/user-message.test.ts index 602dcaa62e..0473bc8eab 100644 --- a/packages/coding-agent/test/user-message.test.ts +++ b/packages/coding-agent/test/user-message.test.ts @@ -1,7 +1,8 @@ import { clearDefaultTerminalColors, setDefaultTerminalColors, visibleWidth } from "@earendil-works/pi-tui"; import { afterEach, describe, expect, test } from "vitest"; +import { styleArgumentTokens } from "../src/modes/interactive/components/prompt-highlight.js"; import { UserMessageComponent } from "../src/modes/interactive/components/user-message.js"; -import { initTheme, theme } from "../src/modes/interactive/theme/theme.js"; +import { initTheme, type ThemeColor, theme } from "../src/modes/interactive/theme/theme.js"; const OSC133_ZONE_START = "\x1b]133;A\x07"; const OSC133_ZONE_END = "\x1b]133;B\x07"; @@ -83,6 +84,83 @@ describe("UserMessageComponent", () => { expect(plainLines).toEqual(expectedLines); }); + const renderMessage = (text: string, width = 60, recognized: (name: string) => boolean = () => false) => { + initTheme("dark"); + return new UserMessageComponent(text, undefined, recognized).render(width).join("\n"); + }; + + const commandText = "/new --name foo @src/foo.ts"; + + test.each<{ name: string; text: string; width?: number; color: ThemeColor; has?: string[]; lacks?: string[] }>([ + { name: "the command accent", text: commandText, color: "accent", has: ["/new"] }, + { name: "--flags in command arguments", text: commandText, color: "mdLink", has: ["--name"] }, + { name: "@paths in command arguments", text: commandText, color: "success", has: ["@src/foo.ts"] }, + { name: "@paths in plain messages", text: "check @src/foo.ts please", color: "success", has: ["@src/foo.ts"] }, + { name: "a leading @path", text: "@src/foo.ts please", color: "success", has: ["@src/foo.ts"] }, + { name: "a newline-leading @path", text: "hello\n@foo", color: "success", has: ["@foo"] }, + { + name: "every wrapped fragment of a long @path", + text: "check @src/very-long-file-name.ts please", + width: 16, + color: "success", + has: ["@src/very-lo", "ng-file-name", ".ts"], + }, + { + name: "quoted @paths across narrow wraps", + text: 'open @"docs/some very long name.txt" now', + width: 16, + color: "success", + has: ['@"docs/some ', "very long na", 'me.txt"'], + }, + { name: "no mid-word @ (emails)", text: "email me@example.com", color: "success", lacks: ["@example.com"] }, + { name: "no glued dashes", text: "a---b", color: "mdLink", lacks: ["--b"] }, + ])("styles tokens: $name", ({ text, width, color, has, lacks }) => { + const rendered = renderMessage(text, width, (name) => name === "new"); + + for (const fragment of has ?? []) expect(rendered).toContain(theme.fg(color, fragment)); + for (const fragment of lacks ?? []) expect(rendered).not.toContain(theme.fg(color, fragment)); + }); + + test("highlights a bare -- separator only in argument-taking slash commands", () => { + const recognized = (name: string) => name === "new" || name === "compact"; + const separator = theme.fg("mdLink", "--"); + + expect(renderMessage("/new --name bla -- hello", 60, recognized)).toContain(separator); + // /compact is recognized but takes no argument, so it must match the editor and show no separator. + const unhighlighted = [ + "/compact -- hello", + "this -- however -- is fine", + "/new a --- b", + "/new x-- y", + "a --- b", + ]; + for (const text of unhighlighted) expect(renderMessage(text, 60, recognized)).not.toContain(separator); + }); + + test("keeps multi-line quoted @paths on separate lines", () => { + initTheme("dark"); + const lines = new UserMessageComponent('@"a\nb"').render(60); + const plain = lines.map((line) => line.replace(/\x1b\[[0-9;]*m|\x1b\]133;[ABC]\x07/g, "")); + const content = plain.map((line) => line.trim()).filter((line) => line.length > 0); + + expect(lines.every((line) => !line.includes("\n"))).toBe(true); + expect(content).toEqual(['@"a', 'b"']); + expect(lines.join("\n")).toContain(theme.fg("success", '@"a')); + }); + + test("normalizes tabs inside quoted @paths like Markdown does", () => { + initTheme("dark"); + const lines = new UserMessageComponent('open @"a\tb.txt" now').render(30); + + expect(lines.some((line) => line.includes("\t"))).toBe(false); + expect(lines.join("\n")).toContain(theme.fg("success", '@"a b.txt"')); + }); + + test("styleArgumentTokens highlights quoted @paths", () => { + initTheme("dark"); + expect(styleArgumentTokens('open @"a b.txt" now')).toBe(`open ${theme.fg("success", '@"a b.txt"')} now`); + }); + test("preserves mask-like argument text across narrow wraps", () => { initTheme("dark"); const command = "/averyveryverylongcommand"; @@ -96,4 +174,47 @@ describe("UserMessageComponent", () => { expect(plain.replace(/\s+/g, "")).toContain(`${command}界\uE000`); expect(lines.every((line) => visibleWidth(line) === 8)).toBe(true); }); + + test("forwards table-cell selection regions with unmasked content", () => { + initTheme("dark"); + const component = new UserMessageComponent("| alpha | beta |\n| --- | --- |\n| @src/a.ts | two |"); + component.render(60); + + const contents = component.getSelectionRegions().map((region) => region.content); + + // Copying the token cell must yield its literal text, not mask placeholders. + expect(contents).toContain("@src/a.ts"); + expect(contents).toContain("two"); + }); + + test.each(["@src/foo.ts", '@"src/a|b.ts"', "@src/a\\|b.ts", "@src/a\\\\", "@src/a\\\\\\|b.ts"])( + "preserves table cells and copied paths for %s", + (path) => { + initTheme("dark"); + const component = new UserMessageComponent(`| alpha | beta |\n| --- | --- |\n| ${path}|two |`); + const rendered = component.render(80).join("\n"); + expect(component.getSelectionRegions().map((region) => region.content)).toEqual([ + "alpha", + "beta", + path, + "two", + ]); + expect(rendered).toContain(theme.fg("success", path)); + for (const includeBareSeparator of [false, true]) { + expect(styleArgumentTokens(`${path}|two`, undefined, includeBareSeparator)).toBe( + `${theme.fg("success", path)}|two`, + ); + } + }, + ); + + test("renders a literal mask-range character before an @token uncorrupted", () => { + initTheme("dark"); + const plain = new UserMessageComponent("\uE000 check @foo") + .render(60) + .map((line) => line.replace(/\x1b\[[0-9;]*m|\x1b\]133;[ABC]\x07/g, "")) + .join("\n"); + + expect(plain).toContain("\uE000 check @foo"); + }); }); diff --git a/packages/tui/.changes/feat-queue-path-arg-highlighting.md b/packages/tui/.changes/feat-queue-path-arg-highlighting.md new file mode 100644 index 0000000000..b38135b803 --- /dev/null +++ b/packages/tui/.changes/feat-queue-path-arg-highlighting.md @@ -0,0 +1 @@ +- Added source-line coordinates to editor layout lines and `styleDisplayText()` so subclasses can style wrapped or scrolled text against exact source offsets. diff --git a/packages/tui/.changes/local-file-links.md b/packages/tui/.changes/local-file-links.md new file mode 100644 index 0000000000..c6f3817d55 --- /dev/null +++ b/packages/tui/.changes/local-file-links.md @@ -0,0 +1 @@ +- Added support for opening file links from the fullscreen terminal UI. diff --git a/packages/tui/src/components/editor.ts b/packages/tui/src/components/editor.ts index 431b659186..ad9fdc8daf 100644 --- a/packages/tui/src/components/editor.ts +++ b/packages/tui/src/components/editor.ts @@ -217,6 +217,10 @@ interface LayoutLine { text: string; hasCursor: boolean; cursorPos?: number; + /** Logical source line index this layout line renders. */ + sourceLine: number; + /** Start offset of this layout line's text within the source line. */ + sourceStart: number; } export interface EditorTheme { @@ -385,6 +389,8 @@ export class Editor implements Component, Focusable { _layoutLineIndex: number, _lineText: string, _cursorCol: number | undefined, + _sourceLine?: number, + _sourceStart?: number, ): string { return displayText; } @@ -612,6 +618,8 @@ export class Editor implements Component, Focusable { absoluteLineIndex, layoutLine.text, layoutLine.hasCursor ? layoutLine.cursorPos : undefined, + layoutLine.sourceLine, + layoutLine.sourceStart, ); const padding = " ".repeat(Math.max(0, inputWidth - lineVisibleWidth)); @@ -943,6 +951,8 @@ export class Editor implements Component, Focusable { text: "", hasCursor: true, cursorPos: 0, + sourceLine: 0, + sourceStart: 0, }); return layoutLines; } @@ -959,6 +969,8 @@ export class Editor implements Component, Focusable { text: "", hasCursor: isCurrentLine, cursorPos: isCurrentLine ? 0 : undefined, + sourceLine: i, + sourceStart: hiddenPrefixLength, }); continue; } @@ -969,11 +981,15 @@ export class Editor implements Component, Focusable { text: displayLine, hasCursor: true, cursorPos: Math.max(0, this.state.cursorCol - hiddenPrefixLength), + sourceLine: i, + sourceStart: hiddenPrefixLength, }); } else { layoutLines.push({ text: displayLine, hasCursor: false, + sourceLine: i, + sourceStart: hiddenPrefixLength, }); } } else { @@ -1011,11 +1027,15 @@ export class Editor implements Component, Focusable { text: chunk.text, hasCursor: true, cursorPos: adjustedCursorPos, + sourceLine: i, + sourceStart: hiddenPrefixLength + chunk.startIndex, }); } else { layoutLines.push({ text: chunk.text, hasCursor: false, + sourceLine: i, + sourceStart: hiddenPrefixLength + chunk.startIndex, }); } } diff --git a/packages/tui/src/tui.ts b/packages/tui/src/tui.ts index ba5b42e264..a3f6e836c0 100644 --- a/packages/tui/src/tui.ts +++ b/packages/tui/src/tui.ts @@ -768,7 +768,7 @@ export class TUI extends Container { let href: string; try { const parsed = new URL(url); - if (parsed.protocol !== "http:" && parsed.protocol !== "https:") return; + if (parsed.protocol !== "http:" && parsed.protocol !== "https:" && parsed.protocol !== "file:") return; href = parsed.href; } catch { return; diff --git a/packages/tui/test/fullscreen.test.ts b/packages/tui/test/fullscreen.test.ts index 3166b47003..99b5e89440 100644 --- a/packages/tui/test/fullscreen.test.ts +++ b/packages/tui/test/fullscreen.test.ts @@ -1066,9 +1066,11 @@ describe("TUI fullscreen mode", () => { tui.stop(); }); - it("ignores clicked hyperlinks with non-http schemes", async () => { + it("opens file hyperlinks and ignores unsupported schemes", async () => { const transcript = lines(20); - transcript[12] = "\x1b]8;;file:///etc/passwd\x1b\\secrets\x1b]8;;\x1b\\"; + transcript[12] = + "\x1b]8;;file:///tmp/example%20file.txt\x1b\\file\x1b]8;;\x1b\\ " + + "\x1b]8;;ssh://example.com\x1b\\remote\x1b]8;;\x1b\\"; const { terminal, tui, chat, dock } = setup(transcript); const opened: string[] = []; tui.onOpenUrl = (url) => opened.push(url); @@ -1077,8 +1079,10 @@ describe("TUI fullscreen mode", () => { terminal.sendInput("\x1b[<0;3;1M"); terminal.sendInput("\x1b[<0;3;1m"); + terminal.sendInput("\x1b[<0;7;1M"); + terminal.sendInput("\x1b[<0;7;1m"); await terminal.waitForRender(); - assert.deepStrictEqual(opened, []); + assert.deepStrictEqual(opened, ["file:///tmp/example%20file.txt"]); tui.stop(); }); diff --git a/prime-agent-runtime/src/rlm/__init__.py b/prime-agent-runtime/src/rlm/__init__.py index e9875c4f1a..53776789ba 100644 --- a/prime-agent-runtime/src/rlm/__init__.py +++ b/prime-agent-runtime/src/rlm/__init__.py @@ -19,6 +19,15 @@ class RLMSpawnHandle: model: str +@dataclass(frozen=True) +class RLMCreateSessionHandle: + active_session_id: str + session_id: str + name: str + session_file: Path + model: str + + @dataclass(frozen=True) class RLMModel: provider: str @@ -54,6 +63,25 @@ def _spawn_handle_from_payload(payload: Any) -> RLMSpawnHandle: ) +def _create_session_handle_from_payload(payload: Any) -> RLMCreateSessionHandle: + if not isinstance(payload, dict): + raise RuntimeError("rlm.create_session returned an invalid payload") + active_session_id = payload.get("active_session_id") + session_id = payload.get("session_id") + name = payload.get("name") + session_file = payload.get("session_file") + model = payload.get("model") + if not all(isinstance(value, str) and value for value in (active_session_id, session_id, name, session_file, model)): + raise RuntimeError("rlm.create_session returned an invalid payload structure") + return RLMCreateSessionHandle( + active_session_id=active_session_id, + session_id=session_id, + name=name, + session_file=Path(session_file), + model=model, + ) + + def _parse_host_reply(request_type: str, reply: dict[str, Any]) -> dict[str, Any]: status = reply.get("status") if status == "ok": @@ -114,6 +142,33 @@ def _model_from_payload(payload: Any) -> RLMModel: return RLMModel(provider=provider, id=model_id, name=name, selector=selector) +async def create_session( + prompt: str, + name: str | None = None, + model: str | None = None, + thinking: str | None = None, + cwd: str | None = None, +) -> RLMCreateSessionHandle: + """Create and prompt a resident depth-0 daemon session. + + Only daemon-backed depth-0 sessions support this operation. The optional + arguments set the session name, model, thinking level, and working directory. + """ + if not isinstance(prompt, str): + raise TypeError(f"prompt must be str, got {type(prompt).__name__}") + kwargs: dict[str, Any] = {} + if name is not None: + kwargs["name"] = name + if model is not None: + kwargs["model"] = model + if thinking is not None: + kwargs["thinking"] = thinking + if cwd is not None: + kwargs["cwd"] = cwd + payload = await host_request("rlm.create_session", {"prompt": prompt, "kwargs": kwargs}) + return _create_session_handle_from_payload(payload) + + async def find_models(query: str = "", limit: int = 8) -> list[RLMModel]: """Search a bounded list of models backed by active user credentials.""" if not isinstance(query, str): @@ -237,6 +292,16 @@ class _RLMCallable: async def run(self, prompt: str, **kwargs: Any) -> RLMSpawnHandle: return await run(prompt, **kwargs) + async def create_session( + self, + prompt: str, + name: str | None = None, + model: str | None = None, + thinking: str | None = None, + cwd: str | None = None, + ) -> RLMCreateSessionHandle: + return await create_session(prompt, name=name, model=model, thinking=thinking, cwd=cwd) + async def find_models(self, query: str = "", limit: int = 8) -> list[RLMModel]: return await find_models(query, limit) @@ -270,9 +335,11 @@ async def __call__(self, prompt: str, **kwargs: Any) -> RLMSpawnHandle: "McpIntegration", "McpToolError", "NotEnabled", + "RLMCreateSessionHandle", "RLMModel", "RLMSpawnHandle", "RLMSubagent", + "create_session", "RefinementEvent", "bash", "delete_subagent", diff --git a/prime-agent-runtime/src/rlm/bash.py b/prime-agent-runtime/src/rlm/bash.py index 5795b6ae72..43cbcfb69d 100644 --- a/prime-agent-runtime/src/rlm/bash.py +++ b/prime-agent-runtime/src/rlm/bash.py @@ -4,6 +4,7 @@ import asyncio import atexit +import functools import json import os import secrets @@ -43,6 +44,14 @@ # wait for a confirmed group exit before CancelledError propagates. _CANCEL_TERM_GRACE = 0.5 _CANCEL_KILL_WAIT = 2.0 +_COMPLETION_NOTICE_COMMAND_CAP = 1000 +_ASYNCIO_WRAPPER_CALLBACKS = { + ("asyncio.tasks", "gather.._done_callback"), + ("asyncio.tasks", "shield.._inner_done_callback"), + ("asyncio.tasks", "_wait.._on_completion"), + ("asyncio.tasks", "as_completed.._on_completion"), + ("asyncio.tasks", "_release_waiter"), +} _live_handles: set["BashHandle"] = set() _live_lock = threading.Lock() @@ -50,6 +59,105 @@ _hook_lock = threading.Lock() +def _current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] | None] | None: + """Get the creating REPL cell's lifecycle without coupling standalone use to repl.""" + try: + from . import repl + + if repl.is_active(): + return repl.current_cell_completion_context() + except (ImportError, RuntimeError): + pass + return None + + +def _consume_notice_task(task: asyncio.Task[None]) -> None: + """Retrieve detached notifier failures so they never become loop warnings.""" + if not task.cancelled(): + task.exception() + + +def _completion_reaches( + start: asyncio.Future[Any], targets: tuple[asyncio.Future[Any], ...] +) -> bool: + """Follow asyncio's wrapper and TaskGroup ownership callbacks.""" + pending = [start] + seen_futures: set[int] = set() + seen_values: set[int] = set() + + def collect(value: Any, depth: int = 0) -> None: + if isinstance(value, asyncio.Future): + pending.append(value) + return + identity = id(value) + if depth >= 4 or identity in seen_values: + return + seen_values.add(identity) + + nested: list[Any] = [] + if isinstance(value, asyncio.Queue): + pending.extend(value._getters) + elif isinstance(value, functools.partial): + nested.extend((value.func, value.args, value.keywords)) + elif isinstance(value, dict): + nested.extend(value.keys()) + nested.extend(value.values()) + elif isinstance(value, (tuple, list, set, frozenset)): + nested.extend(value) + else: + closure = getattr(value, "__closure__", None) or () + for cell in closure: + try: + nested.append(cell.cell_contents) + except ValueError: + pass + bound_self = getattr(value, "__self__", None) + if bound_self is not None: + nested.append(bound_self) + for item in nested: + collect(item, depth + 1) + + while pending: + future = pending.pop() + if any(future is target for target in targets): + return True + if id(future) in seen_futures: + continue + seen_futures.add(id(future)) + for entry in getattr(future, "_callbacks", None) or (): + callback = entry[0] if isinstance(entry, tuple) else entry + base = callback.func if isinstance(callback, functools.partial) else callback + identity = (getattr(base, "__module__", None), getattr(base, "__qualname__", None)) + if identity in _ASYNCIO_WRAPPER_CALLBACKS: + collect(callback) + elif identity == ("asyncio.tasks", "_AsCompletedIterator._handle_completion"): + collect(base.__self__._done) + elif identity == (None, "Task.task_wakeup"): + task = getattr(callback, "__self__", None) + if isinstance(task, asyncio.Task): + pending.append(task) + elif identity == ("asyncio.taskgroups", "TaskGroup._on_task_done"): + parent = getattr(getattr(callback, "__self__", None), "_parent_task", None) + if isinstance(parent, asyncio.Future): + pending.append(parent) + return False + + +def _creating_cell_waits_for( + owner: asyncio.Task[Any] | None, awaiter: asyncio.Task[Any] | None +) -> bool: + """Return whether the cell owner directly or transitively waits for awaiter.""" + if owner is None or awaiter is None: + return False + if owner is awaiter: + return True + waiter = getattr(owner, "_fut_waiter", None) + targets: tuple[asyncio.Future[Any], ...] = (owner,) + if isinstance(waiter, asyncio.Future): + targets += (waiter,) + return _completion_reaches(awaiter, targets) + + @dataclass(frozen=True) class BashResult: exit_code: int @@ -117,6 +225,10 @@ class BashHandle: def __init__(self, command: str) -> None: self.command = command + completion_context = _current_cell_completion_context() + self._creating_cell_finished = completion_context[0] if completion_context else None + self._creating_cell_task = completion_context[1] if completion_context else None + self._awaited_by_creating_cell = False self._buffer = _BoundedBuffer() self._done = threading.Event() self._eof = threading.Event() @@ -129,6 +241,7 @@ def __init__(self, command: str) -> None: self._reaped = False self._result: BashResult | None = None self._callbacks: list[Callable[[], None]] = [] + self._reap_callback: Callable[[], None] | None = None self._callback_lock = threading.Lock() # Serializes kill/reap so a pid fallback can never outlive the process handle. self._kill_lock = threading.Lock() @@ -235,6 +348,7 @@ def __init__(self, command: str) -> None: threading.Thread(target=self._pump, daemon=True).start() threading.Thread(target=self._report, daemon=True).start() threading.Thread(target=self._watch, daemon=True).start() + self._schedule_background_completion_notice() @property def pid(self) -> int: @@ -409,6 +523,10 @@ def _watch(self) -> None: if not _IS_POSIX: # Reaped: pid fallbacks are gone, so the handle may finally close. cast("_winjob.JobProcess", self._proc).close() + with self._callback_lock: + callback, self._reap_callback = self._reap_callback, None + if callback is not None: + callback() if delivered: _record_journal(self._pid, active=False) with _live_lock: @@ -515,6 +633,87 @@ def _add_done_callback(self, callback: Callable[[], None]) -> None: return callback() + def _schedule_background_completion_notice(self) -> None: + cell_finished = self._creating_cell_finished + if cell_finished is None: + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return + from . import repl + + activity = {"id": secrets.token_hex(16), "pid": self._pid, "active": True} + # Publish synchronously before bash() returns and the creating cell can end. + repl.emit({"application/vnd.prime-agent.bash-activity+json": activity}) + notice = self._notify_background_completion(cell_finished, activity) + try: + task = loop.create_task(notice) + except BaseException: + self.kill(signal.SIGKILL if _IS_POSIX else signal.SIGTERM) + notice.close() + repl.emit({"application/vnd.prime-agent.bash-activity+json": {**activity, "active": False}}) + raise + task.add_done_callback(_consume_notice_task) + + async def _notify_background_completion( + self, cell_finished: asyncio.Event, activity: dict[str, Any] + ) -> None: + from . import repl + + try: + result = await self._wait() + await self._wait_reaped() + # The cell may do other work before awaiting this handle. Do not classify + # it as detached until that whole cell has crossed its completion barrier. + await cell_finished.wait() + if self._awaited_by_creating_cell or not repl.is_active(): + return + command = self.command + if len(command) > _COMPLETION_NOTICE_COMMAND_CAP: + command = command[:_COMPLETION_NOTICE_COMMAND_CAP] + "\n... [command truncated]" + reply = await repl.host_request( + { + "type": "bash.completed", + "pid": self._pid, + "command": command, + "exitCode": result.exit_code, + } + ) + if not isinstance(reply, dict) or reply.get("status") != "ok": + sys.stderr.write( + f"Background bash completion follow-up for pid {self._pid} was not accepted. " + "Inspect the saved handle with poll(), output(), or tail().\n" + ) + except (OSError, RuntimeError): + # Standalone runtimes have no host handler, and teardown can close + # the bridge while a process is finishing. Shell results stay usable. + return + finally: + # Reap and deliver (or report rejection) before releasing kernel residency. + repl.emit({"application/vnd.prime-agent.bash-activity+json": {**activity, "active": False}}) + + async def _wait_reaped(self) -> None: + loop = asyncio.get_running_loop() + future: asyncio.Future[None] = loop.create_future() + + def wake() -> None: + try: + loop.call_soon_threadsafe(lambda: future.done() or future.set_result(None)) + except RuntimeError: + pass + + with self._callback_lock: + if self._reaped: + return + self._reap_callback = wake + try: + await future + finally: + with self._callback_lock: + if self._reap_callback is wake: + self._reap_callback = None + async def _wait(self) -> BashResult: # Asyncio-native wakeup: no executor thread is parked for the command's # duration, so many concurrent awaits cannot exhaust the default pool. @@ -639,10 +838,25 @@ def __await__(self) -> Generator[Any, None, BashResult]: # A handle awaited before any other API use is a one-shot command tied # to the await (kill-on-cancel); touching the handle API first marks it # as a deliberate background handle whose awaits only wait. - if self._released: - return self._wait().__await__() + try: + current_task = asyncio.current_task() + except RuntimeError: + current_task = None + creating_cell_waited = _creating_cell_waits_for(self._creating_cell_task, current_task) + owned = not self._released + wait = self._wait_owned() if owned else self._wait() self._released = True - return self._wait_owned().__await__() + completed = False + try: + result = yield from wait.__await__() + completed = True + return result + finally: + if (completed or owned) and ( + creating_cell_waited + or _creating_cell_waits_for(self._creating_cell_task, current_task) + ): + self._awaited_by_creating_cell = True def __repr__(self) -> str: state = f"exit_code={self._result.exit_code}" if self._result else "running" diff --git a/prime-agent-runtime/src/rlm/repl.py b/prime-agent-runtime/src/rlm/repl.py index ea04ca1766..d3324be774 100644 --- a/prime-agent-runtime/src/rlm/repl.py +++ b/prime-agent-runtime/src/rlm/repl.py @@ -45,10 +45,20 @@ _write_lock = threading.Lock() _loop: asyncio.AbstractEventLoop | None = None _serve_task: asyncio.Task[Any] | None = None -# Attribution rides task context: asyncio tasks copy it at creation, so a -# detached task spawned by a cell keeps writing under that cell's id after -# the cell finishes. Threads start with a fresh context and emit id null. + + +class _CellExecution: + def __init__(self) -> None: + self.finished = asyncio.Event() + self.owner: asyncio.Task[Any] | None = None + + +# Asyncio tasks copy cell context at creation, so detached tasks retain their +# output attribution and completion barrier. Threads start with a fresh context. _current_cell: contextvars.ContextVar[str | None] = contextvars.ContextVar("_current_cell", default=None) +_current_cell_execution: contextvars.ContextVar[_CellExecution | None] = contextvars.ContextVar( + "_current_cell_execution", default=None +) _active: dict[str, Any] = {"task": None, "rid": None, "interrupted": False} _cell_counter = 0 _pending_host: dict[str, "asyncio.Future[dict[str, Any]]"] = {} @@ -98,6 +108,14 @@ def is_active() -> bool: return _protocol_fd >= 0 +def current_cell_completion_context() -> tuple[asyncio.Event, asyncio.Task[Any] | None] | None: + """Return the calling cell's completion barrier and owning execution task.""" + execution = _current_cell_execution.get() + if execution is None: + return None + return execution.finished, execution.owner + + async def host_request(data: dict[str, Any]) -> dict[str, Any]: """Send one typed request to the host and await its raw reply dict.""" if _loop is None: @@ -537,13 +555,14 @@ async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: cell_id = req["id"] _cell_counter += 1 filename = f"" - # The cell task (created below) copies this context, so writes made from - # the cell and from asyncio tasks it spawns carry this cell's id. - token = _current_cell.set(cell_id) + execution = _CellExecution() + cell_token = _current_cell.set(cell_id) + execution_token = _current_cell_execution.set(execution) try: codes, has_trailing = _compile_cell(req["code"], filename) assert _loop is not None task = _loop.create_task(_run_codes(codes, ns)) + execution.owner = task status, value, error = await _run_guarded(task, cell_id) result_text: str | None = None try: @@ -568,7 +587,10 @@ async def _handle_execute(req: dict[str, Any], ns: dict[str, Any]) -> None: _send(error) _send({"event": "done", "id": cell_id, "status": status}) finally: - _current_cell.reset(token) + execution.owner = None + execution.finished.set() + _current_cell_execution.reset(execution_token) + _current_cell.reset(cell_token) def _drain_output() -> None: diff --git a/prime-agent-runtime/test/test_bash.py b/prime-agent-runtime/test/test_bash.py index f7202cf68d..0edd053f6e 100644 --- a/prime-agent-runtime/test/test_bash.py +++ b/prime-agent-runtime/test/test_bash.py @@ -13,6 +13,7 @@ import time import unittest from concurrent.futures import ThreadPoolExecutor +from types import FunctionType, SimpleNamespace from unittest import mock from rlm import bash @@ -55,6 +56,38 @@ async def test_await_returns_result(self): awaited = await handle self.assertEqual(handle.poll(), awaited) + def test_construction_cleanup_uses_windows_signal_without_sigkill(self): + failure = RuntimeError("task construction failed") + loop = mock.Mock() + loop.create_task.side_effect = failure + bridge = SimpleNamespace(emit=mock.Mock()) + namespace = dict(bash_module.BashHandle._schedule_background_completion_notice.__globals__) + + def isolated_import(name, globals=None, locals=None, fromlist=(), level=0): + if level == 1 and name == "" and fromlist == ("repl",): + return SimpleNamespace(repl=bridge) + return __import__(name, globals, locals, fromlist, level) + + namespace.update( + _IS_POSIX=False, + signal=SimpleNamespace(SIGTERM=15), + asyncio=SimpleNamespace(get_running_loop=lambda: loop), + __builtins__={**vars(__import__("builtins")), "__import__": isolated_import}, + ) + schedule = FunctionType( + bash_module.BashHandle._schedule_background_completion_notice.__code__, namespace + ) + handle = mock.Mock(_pid=42) + with self.assertRaises(RuntimeError) as caught: + schedule(handle) + self.assertIs(caught.exception, failure) + handle.kill.assert_called_once_with(15) + handle._notify_background_completion.return_value.close.assert_called_once_with() + activity = bridge.emit.call_args_list[0].args[0] + mime = "application/vnd.prime-agent.bash-activity+json" + self.assertTrue(activity[mime]["active"]) + bridge.emit.assert_called_with({mime: {**activity[mime], "active": False}}) + async def test_status_pipe_survives_high_fds_and_strict_posix_shell(self): # Regression: dash rejects multi-digit fds in redirections at parse # time, so the script must never reference the raw status-pipe fd. diff --git a/prime-agent-runtime/test/test_create_session.py b/prime-agent-runtime/test/test_create_session.py new file mode 100644 index 0000000000..84b46fa819 --- /dev/null +++ b/prime-agent-runtime/test/test_create_session.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import asyncio +import importlib +import unittest +from pathlib import Path +from unittest.mock import AsyncMock, patch + + +rlm_module = importlib.import_module("rlm") + +PAYLOAD = { + "active_session_id": "active-root", + "session_id": "session-root", + "name": "researcher", + "session_file": "/tmp/sessions/researcher/session.jsonl", + "model": "prime-inference/deepseek/deepseek-v4-flash", +} + + +class RlmCreateSessionTest(unittest.TestCase): + def test_returns_typed_handle_and_forwards_options(self) -> None: + host_request = AsyncMock(return_value=PAYLOAD) + + with patch.object(rlm_module, "host_request", host_request): + handle = asyncio.run( + rlm_module.rlm.create_session( + "analyze shard 1", + name="researcher", + model="prime-inference/deepseek/deepseek-v4-flash", + thinking="off", + cwd="/tmp/project", + ) + ) + + host_request.assert_awaited_once_with( + "rlm.create_session", + { + "prompt": "analyze shard 1", + "kwargs": { + "name": "researcher", + "model": "prime-inference/deepseek/deepseek-v4-flash", + "thinking": "off", + "cwd": "/tmp/project", + }, + }, + ) + self.assertEqual( + handle, + rlm_module.RLMCreateSessionHandle( + active_session_id="active-root", + session_id="session-root", + name="researcher", + session_file=Path("/tmp/sessions/researcher/session.jsonl"), + model="prime-inference/deepseek/deepseek-v4-flash", + ), + ) + + def test_omits_unset_options(self) -> None: + host_request = AsyncMock(return_value=PAYLOAD) + + with patch.object(rlm_module, "host_request", host_request): + asyncio.run(rlm_module.create_session("check status")) + + host_request.assert_awaited_once_with( + "rlm.create_session", + {"prompt": "check status", "kwargs": {}}, + ) + + def test_rejects_non_string_prompt(self) -> None: + with self.assertRaisesRegex(TypeError, "prompt must be str"): + asyncio.run(rlm_module.create_session(42)) # type: ignore[arg-type] + + def test_rejects_invalid_host_payload(self) -> None: + host_request = AsyncMock(return_value={**PAYLOAD, "session_file": ""}) + + with patch.object(rlm_module, "host_request", host_request): + with self.assertRaisesRegex(RuntimeError, "invalid payload"): + asyncio.run(rlm_module.create_session("test")) + + +if __name__ == "__main__": + unittest.main() diff --git a/prime-agent-runtime/test/test_repl.py b/prime-agent-runtime/test/test_repl.py index 26b7aca89d..6a2b5daf5b 100644 --- a/prime-agent-runtime/test/test_repl.py +++ b/prime-agent-runtime/test/test_repl.py @@ -109,6 +109,19 @@ def one(events: list[dict], kind: str) -> dict | None: return matches[0] if matches else None +def wait_for_host_request(repl: ReplProcess, events: list[dict]) -> dict: + request = one(events, "host_request") + while request is None: + event = repl.read_event() + if event.get("event") == "host_request": + request = event + return request + + +def reply_ok(repl: ReplProcess, request: dict) -> None: + repl.send({"type": "host_reply", "id": request["id"], "data": {"status": "ok", "result": {}}}) + + class ReplTest(unittest.TestCase): def setUp(self) -> None: self.repl = ReplProcess() @@ -713,6 +726,335 @@ def test_bash_integration(self): time.sleep(0.05) self.fail(f"bash child {pid} survived runtime shutdown") + def test_async_bash_completion_notifies_only_after_an_unawaited_creating_cell(self): + direct = self.repl.execute( + "bash-direct", "from rlm import bash\n(await bash('printf direct')).exit_code" + ) + self.assertIsNone(one(direct, "host_request")) + + eventual = self.repl.execute( + "bash-eventual", + "import asyncio\nhandle = bash('printf eventual')\nawait asyncio.sleep(0.1)\n(await handle).exit_code", + ) + self.assertIsNone(one(eventual, "host_request")) + + started = self.repl.execute( + "bash-detached", "handle = bash('sleep 0.05; printf detached')\nhandle.pid" + ) + pid = int(one(started, "result")["text"]) + request = wait_for_host_request(self.repl, started) + self.assertEqual( + request["data"], + { + "type": "bash.completed", + "pid": pid, + "command": "sleep 0.05; printf detached", + "exitCode": 0, + }, + ) + reply_ok(self.repl, request) + inspected = self.repl.execute("bash-inspect", "handle.poll().output") + self.assertIn("detached", one(inspected, "result")["text"]) + + def test_wrapper_awaits_in_creating_cell_suppress_bash_completion(self): + def task_group(label: str, await_expression: str) -> str: + return "\n".join( + [ + f"handle = bash('printf {label}')", + "async def consume():", + f" return {await_expression}", + "async with asyncio.TaskGroup() as group:", + " task = group.create_task(consume())", + "task.result().output", + ] + ) + + snippets = { + "gather": "(await asyncio.gather(bash('printf gather')))[0].output", + "wait-for": "(await asyncio.wait_for(bash('printf wait-for'), 1)).output", + "shield": "(await asyncio.shield(bash('printf shield'))).output", + "nested": "(await asyncio.shield(asyncio.gather(bash('printf nested'))))[0].output", + "as-completed": ( + "[(await completed).output for completed in " + "asyncio.as_completed([bash('printf as-completed')])]" + ), + "as-completed-delayed-consumer": ( + "handle = bash('sleep 0.15; printf as-completed-delayed-consumer')\n" + "completed = next(iter(asyncio.as_completed([handle])))\n" + "await asyncio.sleep(0)\n" + "(await completed).output" + ), + "as-completed-finished-handle": ( + "handle = bash('printf as-completed-finished-handle')\n" + "task = asyncio.ensure_future(handle)\n" + "await asyncio.to_thread(handle._done.wait)\n" + "await asyncio.sleep(0.01)\n" + "assert task.done()\n" + "(await handle).output" + ), + "as-completed-nested": ( + "[(await completed)[0].output for completed in " + "asyncio.as_completed([asyncio.shield(asyncio.gather(bash('printf as-completed-nested')))])]" + ), + "wait": "\n".join( + [ + "handle = bash('printf wait')", + "task = asyncio.ensure_future(handle)", + "await asyncio.wait({task})", + "task.result().output", + ] + ), + "task-group": task_group("task-group", "await handle"), + "task-group-gather": task_group("task-group-gather", "(await asyncio.gather(handle))[0]"), + "task-group-shield": task_group("task-group-shield", "await asyncio.shield(handle)"), + } + if sys.version_info >= (3, 13): + snippets["as-completed-async"] = ( + "[completed.result().output async for completed in " + "asyncio.as_completed([bash('printf as-completed-async')])]" + ) + for label, snippet in snippets.items(): + with self.subTest(label=label): + completed = self.repl.execute( + f"bash-wrapper-{label}", + f"from rlm import bash\nimport asyncio\n{snippet}", + ) + self.assertIn(label, one(completed, "result")["text"]) + probe = self.repl.execute(f"bash-wrapper-{label}-probe", "await asyncio.sleep(0.05)") + request = one(probe, "host_request") + if request is not None: + reply_ok(self.repl, request) + self.assertIsNone(request) + + def test_as_completed_outside_creating_cell_keeps_one_bash_completion(self): + snippets = { + "no-await": ( + "iterator = asyncio.as_completed([handle])\n" + "next(iterator).close()\n" + "await asyncio.sleep(0.05)" + ), + "detached": ( + "async def consume():\n" + " for completed in asyncio.as_completed([handle]):\n" + " await completed\n" + "consumer = asyncio.create_task(consume())\n" + "await asyncio.sleep(0.05)" + ), + "later-cell": "iterator = asyncio.as_completed([handle])", + } + for label, snippet in snippets.items(): + with self.subTest(label=label): + command = f"printf {label}" + events = self.repl.execute( + f"as-completed-{label}", + f"from rlm import bash\nimport asyncio\nhandle = bash({command!r})\n{snippet}", + ) + if label == "later-cell": + events += self.repl.execute( + "as-completed-later-await", + "for completed in iterator:\n await completed", + ) + events += self.repl.execute( + f"as-completed-{label}-probe", "await asyncio.sleep(0.05)" + ) + requests = [event for event in events if event.get("event") == "host_request"] + for request in requests: + reply_ok(self.repl, request) + self.assertEqual(len(requests), 1) + self.assertEqual(requests[0]["data"]["type"], "bash.completed") + self.assertEqual(requests[0]["data"]["command"], command) + self.assertFalse(any(event.get("event") == "error" for event in events)) + + def test_background_task_await_does_not_suppress_bash_completion(self): + code = "\n".join( + [ + "from rlm import bash", + "import asyncio", + "task_handle = bash('sleep 0.05; printf background-waiter')", + "async def consume():", + " globals()['task_result'] = await task_handle", + "waiter = asyncio.create_task(consume())", + "bookkeeping = asyncio.get_running_loop().create_future()", + "def callback_for(marker):", + " return lambda _: marker.cancelled()", + "waiter.add_done_callback(callback_for(bookkeeping))", + "asyncio.get_running_loop().call_later(0.02, bookkeeping.set_result, None)", + "await bookkeeping", + "task_handle.pid", + ] + ) + started = self.repl.execute("bash-task-waiter", code) + pid = int(one(started, "result")["text"]) + request = wait_for_host_request(self.repl, started) + self.assertEqual(request["data"]["type"], "bash.completed") + self.assertEqual(request["data"]["pid"], pid) + reply_ok(self.repl, request) + + def test_bash_activity_and_notice_wait_for_group_reap(self): + mime = "application/vnd.prime-agent.bash-activity+json" + for awaited in (False, True): + with self.subTest(awaited=awaited): + events = self.repl.execute( + "group-start", + "from rlm import bash\nimport asyncio\n" + "handle = bash('sleep 30 &')\n" + + ("await handle\n" if awaited else "") + + "await asyncio.wait_for(handle._wait(), 3)\nhandle._group_alive()", + ) + self.assertEqual(one(events, "result")["text"], "True") + self.assertIsNone(one(events, "host_request")) + self.assertFalse( + any(event.get("data", {}).get(mime, {}).get("active") is False for event in events) + ) + events += self.repl.execute("group-stop", "handle.kill()") + requests = [] + cursor = 0 + while True: + for event in events[cursor:]: + if event.get("event") == "host_request": + requests.append(event) + reply_ok(self.repl, event) + cursor = len(events) + if any(event.get("data", {}).get(mime, {}).get("active") is False for event in events): + break + events.append(self.repl.read_event()) + self.assertEqual(len(requests), 0 if awaited else 1) + checked = self.repl.execute( + "group-reaped", "await asyncio.wait_for(handle._wait_reaped(), 1)\nhandle._reaped" + ) + self.assertEqual(one(checked, "result")["text"], "True") + + def test_cancelled_released_bash_await_keeps_completion_notice(self): + mime = "application/vnd.prime-agent.bash-activity+json" + for released in (False, True): + with self.subTest(released=released): + events = self.repl.execute( + "cancel-start", + "from rlm import bash\nimport asyncio\n" + "handle = bash('sleep 0.15; printf finished')\n" + + ("handle.pid\n" if released else "") + + "try:\n await asyncio.wait_for(handle, 0.01)\n" + "except TimeoutError:\n pass\nhandle._group_alive()", + ) + self.assertEqual(one(events, "result")["text"], str(released)) + requests = [] + cursor = 0 + while True: + for event in events[cursor:]: + if event.get("event") == "host_request": + requests.append(event) + reply_ok(self.repl, event) + cursor = len(events) + if any(event.get("data", {}).get(mime, {}).get("active") is False for event in events): + break + events.append(self.repl.read_event()) + self.assertEqual(len(requests), 1 if released else 0) + checked = self.repl.execute("cancel-result", "handle.poll().output") + self.assertEqual(one(checked, "result")["text"], "'finished'" if released else "''") + + def test_interrupt_during_bash_notifier_construction_rolls_back_activity(self): + code = "\n".join( + [ + "from rlm import bash, repl as bridge", + "import sys, threading, asyncio, inspect", + "def stop_at_notice_construction(frame, event, arg):", + " coro = frame.f_locals.get('coro')", + " if (event == 'call' and frame.f_code.co_name == 'create_task'", + " and getattr(getattr(coro, 'cr_code', None), 'co_name', None)", + " == '_notify_background_completion'):", + " globals()['unscheduled_notice'] = coro", + " globals()['interrupted_handle'] = coro.cr_frame.f_locals['self']", + " bridge.emit({'application/json': {'at_notice_create_task': True}})", + " threading.Event().wait(10)", + " return stop_at_notice_construction", + "sys.settrace(stop_at_notice_construction)", + "handle = bash('sleep 30')", + ] + ) + self.repl.send({"type": "execute", "id": "construct", "code": code}) + events = [] + while True: + event = self.repl.read_event() + events.append(event) + if event.get("data", {}).get("application/json", {}).get("at_notice_create_task"): + break + self.repl.send({"type": "interrupt", "id": "construct"}) + events += self.repl.until_done("construct") + self.assertEqual(one(events, "done")["status"], "error") + self.assertEqual(one(events, "error")["ename"], "KeyboardInterrupt") + mime = "application/vnd.prime-agent.bash-activity+json" + activity = [event["data"][mime] for event in events if mime in event.get("data", {})] + self.assertEqual([item["active"] for item in activity], [True, False]) + self.assertEqual(activity[0]["id"], activity[1]["id"]) + checked = self.repl.execute( + "construct-check", + "sys.settrace(None)\n" + "assert await interrupted_handle._await_group_death(3)\n" + "(inspect.getcoroutinestate(unscheduled_notice), interrupted_handle._reap_callback, " + "interrupted_handle._group_alive())", + ) + self.assertEqual(one(checked, "result")["text"], "('CORO_CLOSED', None, False)") + + def test_rejected_bash_completion_reports_safe_stderr_and_releases_activity(self): + mime = "application/vnd.prime-agent.bash-activity+json" + for reply in ({"status": "error", "error": "secret-host-error"}, {}, {"status": "ok"}): + with self.subTest(reply=reply): + events = self.repl.execute( + "ack-start", "from rlm import bash\nhandle = bash('printf secret-command')" + ) + request = wait_for_host_request(self.repl, events) + self.repl.send({"type": "host_reply", "id": request["id"], "data": reply}) + after = [] + while not any( + event.get("data", {}).get(mime, {}).get("active") is False for event in after + ): + after.append(self.repl.read_event()) + diagnostic = stream_text(after, "stderr") + if reply.get("status") == "ok": + self.assertEqual(diagnostic, "") + else: + self.assertEqual( + diagnostic, + f"Background bash completion follow-up for pid {request['data']['pid']} " + "was not accepted. Inspect the saved handle with poll(), output(), or tail().\n", + ) + self.assertIsNone(one(after, "host_request")) + checked = self.repl.execute("ack-check", "handle._reaped, handle.poll().exit_code") + self.assertEqual(one(checked, "result")["text"], "(True, 0)") + + def test_reused_request_id_does_not_capture_old_cell_bash_completion(self): + setup = "\n".join( + [ + "from rlm import bash", + "import asyncio", + "reuse_gate = asyncio.Event()", + "async def launch_after_cell():", + " await reuse_gate.wait()", + " globals()['reused_handle'] = bash('printf reused-id')", + "asyncio.create_task(launch_after_cell())", + ] + ) + self.repl.execute("reused-cell-id", setup) + + self.repl.send( + { + "type": "execute", + "id": "reused-cell-id", + "code": "reuse_gate.set()\nawait asyncio.sleep(0.2)", + } + ) + request = None + while True: + event = self.repl.read_event() + if event.get("event") == "host_request": + request = event + break + self.assertFalse(event.get("event") == "done" and event.get("id") == "reused-cell-id") + self.assertEqual(request["data"]["type"], "bash.completed") + self.assertEqual(request["data"]["command"], "printf reused-id") + reply_ok(self.repl, request) + self.assertEqual(one(self.repl.until_done("reused-cell-id"), "done")["status"], "ok") + def test_protocol_framing_under_noise(self): setup = "\n".join( [