From 361da19bf26b65e992c2d81a5ad378bfc696d1c5 Mon Sep 17 00:00:00 2001 From: David Date: Sun, 26 Jul 2026 17:47:11 -0700 Subject: [PATCH] =?UTF-8?q?feat:=20team=20roles=20=E2=80=94=20cloud=20driv?= =?UTF-8?q?er=20+=20verified=20local=20executors,=20with=20a=20coding=20be?= =?UTF-8?q?nch=20as=20the=20promotion=20gate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /team lets the cloud driver (e.g. pipenetwork glm-5.2) delegate execution to local models with zero endpoint typing: pick a name (or just `local`) and hi does backend detection, quiet download, hi-local serve, health wait, and role wiring, narrating phases with an in-place progress line. Bare `/team delegate|explore` opens a picker over the catalog, annotated per machine. Catalog entries now carry pipenetwork's whole MLX quant ladders (Laguna S 2.1 at 2/3/4/6/8bit, Nemotron Nano 4B at 4/8bit, GLM-5.2 REAP50, Nemotron Ultra 550B): auto picks the family by preference order and the highest quality quant that fits; `name@quant` forces a rung; explicit oversized picks fall to the smallest published quant. CUDA machines skip MLX-only families instead of failing at setup. `hi team-bench` benchmarks local models on verifiable coding tasks (codegen/bugfix/edit compiled + asserted by rustc, strict-JSON tool-call fidelity) and is the catalog's promotion gate: auto only serves models that passed it on real hardware. Live verification on a 64GB M-series Mini: coder-32b 4/4 (~7 tok/s), nemotron-4b@8bit 2/4 (~36 tok/s, over-reasons on codegen), nemotron-30b 1/4 (~17 tok/s) — and it caught laguna-s generating ~0 tok/s through the MoE expert-streaming path and qwen3.6-35b unable to load (nvfp4 unsupported by hi-mlx). Unverified/broken entries are explicit-pick only, honestly labeled. Hardening that fell out of live testing: the model cache is shared at ~/.hi/models (cwd-relative downloads duplicated 50GB checkpoints per project; legacy dirs keep working), model_present rejects partial downloads (aria2 control files, missing index shards), hi-local's default output ceiling rises 2048 → 8192 (it truncated real work mid-function), and delegate/explore routes ride SubagentRoute with cloud keys never leaked to local endpoints. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 1 + crates/hi-agent/src/agent/background_task.rs | 24 +- crates/hi-agent/src/agent/delegate_turn.rs | 17 +- crates/hi-agent/src/agent/explore_turn.rs | 104 ++- crates/hi-agent/src/agent/lifecycle.rs | 131 ++- crates/hi-agent/src/command.rs | 15 + crates/hi-agent/src/config.rs | 23 + crates/hi-agent/src/lib.rs | 26 +- crates/hi-agent/src/local_skeptic.rs | 920 ++++++++++++++++++- crates/hi-agent/src/subagent.rs | 43 + crates/hi-agent/src/tests/explore.rs | 75 ++ crates/hi-agent/src/tests/local_skeptic.rs | 7 +- crates/hi-cli/Cargo.toml | 1 + crates/hi-cli/src/agent_build.rs | 13 + crates/hi-cli/src/commands.rs | 196 ++++ crates/hi-cli/src/delegate.rs | 55 +- crates/hi-cli/src/delegate_tests.rs | 63 ++ crates/hi-cli/src/main.rs | 4 + crates/hi-cli/src/sync.rs | 88 +- crates/hi-cli/src/sync_store.rs | 79 +- crates/hi-cli/src/team_bench.rs | 681 ++++++++++++++ crates/hi-local-core/src/model.rs | 6 +- crates/hi-sqlite-journal/src/lib.rs | 62 +- crates/hi-tools/src/background.rs | 8 + crates/hi-tools/src/hf.rs | 41 + crates/hi-tools/src/lib.rs | 8 +- crates/hi-tools/src/local_server.rs | 157 +++- crates/hi-tui/src/app/commands.rs | 3 + crates/hi-tui/src/app/lifecycle.rs | 2 + crates/hi-tui/src/app/models.rs | 21 +- crates/hi-tui/src/app/run/mod.rs | 3 + crates/hi-tui/src/app/sync_commands.rs | 665 ++++++++++++++ crates/hi-tui/src/app/transcript.rs | 27 + crates/hi-tui/src/lib.rs | 32 + crates/hi-tui/src/tests.rs | 247 +++++ 35 files changed, 3764 insertions(+), 84 deletions(-) create mode 100644 crates/hi-cli/src/team_bench.rs diff --git a/Cargo.lock b/Cargo.lock index 8c17301a..07eef1eb 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2609,6 +2609,7 @@ dependencies = [ "hi-ai", "hi-crash-handler", "hi-rsi-runtime", + "hi-sqlite-journal", "hi-tools", "hi-trace", "hi-tui", diff --git a/crates/hi-agent/src/agent/background_task.rs b/crates/hi-agent/src/agent/background_task.rs index 18c61dc5..7777b665 100644 --- a/crates/hi-agent/src/agent/background_task.rs +++ b/crates/hi-agent/src/agent/background_task.rs @@ -162,8 +162,14 @@ impl crate::Agent { let summary: String = description.chars().take(72).collect(); ui.subagent_note(&format!("↳ background {subagent_type} task {n}: {summary}")); - // Build the future factory and spawn the task. - let provider = self.provider.clone(); + // Build the future factory and spawn the task. Each role runs on its + // configured route (team roles): explore/delegate children may use a + // different model or endpoint than the driver. + let provider = if is_explore { + self.explore_child_provider() + } else { + self.delegate_child_provider() + }; let child_config = if is_explore { self.build_bg_explore_config(n) } else { @@ -355,10 +361,7 @@ impl crate::Agent { /// Build a child config for a background explore subagent. fn build_bg_explore_config(&self, n: u32) -> AgentConfig { - let explore_model = std::env::var("HI_EXPLORE_MODEL") - .ok() - .filter(|model| !model.trim().is_empty()) - .unwrap_or_else(|| self.config.routing.model.clone()); + let explore_model = crate::agent::explore_turn::explore_child_model(&self.config); AgentConfig { paths: crate::AgentPaths { workspace_root: self.runtime.root().to_path_buf(), @@ -405,6 +408,13 @@ impl crate::Agent { /// Build a child config for a background delegate subagent. fn build_bg_delegate_config(&self, n: u32) -> AgentConfig { + let delegate_model = self + .config + .subagents + .delegate_model + .clone() + .filter(|model| !model.trim().is_empty()) + .unwrap_or_else(|| self.config.routing.model.clone()); AgentConfig { paths: crate::AgentPaths { workspace_root: self.runtime.root().to_path_buf(), @@ -415,7 +425,7 @@ impl crate::Agent { .join(format!("bg-delegate-{n}")), }, routing: crate::AgentRouting { - model: self.config.routing.model.clone(), + model: delegate_model, requested_max_tokens: self.config.routing.requested_max_tokens, max_tokens: self.config.routing.max_tokens, max_tokens_explicit: self.config.routing.max_tokens_explicit, diff --git a/crates/hi-agent/src/agent/delegate_turn.rs b/crates/hi-agent/src/agent/delegate_turn.rs index 4e70a591..f32bdd49 100644 --- a/crates/hi-agent/src/agent/delegate_turn.rs +++ b/crates/hi-agent/src/agent/delegate_turn.rs @@ -73,6 +73,8 @@ pub(crate) struct DelegateJob { pub(crate) task: String, pub(crate) verify: Option, pub(crate) runner: std::sync::Arc, + /// Team-role route override for this executor (all-`None` = driver route). + pub(crate) route: crate::SubagentRoute, pub(crate) cancellation: crate::TurnCancellation, /// File paths extracted from the task description (best-effort). Used to /// detect overlap between parallel delegates — only disjoint file sets @@ -242,6 +244,7 @@ impl crate::Agent { task, verify, runner, + route: self.delegate_route(), cancellation: crate::TurnCancellation::new(), file_set, }, @@ -249,6 +252,16 @@ impl crate::Agent { )) } + /// The configured executor route for `delegate` children (team roles). + /// All-`None` inherits the driver's provider/model. + pub(crate) fn delegate_route(&self) -> crate::SubagentRoute { + crate::SubagentRoute { + model: self.config.subagents.delegate_model.clone(), + base_url: self.config.subagents.delegate_endpoint.clone(), + api_key: self.config.subagents.delegate_endpoint_key.clone(), + } + } + /// Run one write-capable `delegate` subagent and return a summary. The runner /// isolates it in a worktree and applies its changes back only if verification /// passes; on failure nothing touches the real tree (spatial isolation). @@ -487,11 +500,12 @@ pub(crate) async fn run_delegate_job(job: DelegateJob) -> DelegateJobResult { task, verify, runner, + route, cancellation, file_set: _, } = job; let outcome = runner - .run_cancellable(&task, verify.as_deref(), cancellation) + .run_routed(&task, verify.as_deref(), &route, cancellation) .await; DelegateJobResult { slot, outcome } } @@ -613,6 +627,7 @@ mod tests { task: format!("update src/module-{slot}.rs"), verify: None, runner: runner_for_job, + route: crate::SubagentRoute::default(), cancellation: crate::TurnCancellation::new(), file_set: std::collections::BTreeSet::from([format!("src/module-{slot}.rs")]), }))); diff --git a/crates/hi-agent/src/agent/explore_turn.rs b/crates/hi-agent/src/agent/explore_turn.rs index 427b7cfb..f558ad65 100644 --- a/crates/hi-agent/src/agent/explore_turn.rs +++ b/crates/hi-agent/src/agent/explore_turn.rs @@ -79,10 +79,7 @@ impl crate::Agent { let n = self .subagents .try_begin_explore(MAX_EXPLORE_SUBAGENTS_PER_TURN)?; - let child_model = std::env::var("HI_EXPLORE_MODEL") - .ok() - .filter(|model| !model.trim().is_empty()) - .unwrap_or_else(|| self.config.routing.model.clone()); + let child_model = explore_child_model(&self.config); let child_project_context = self .config .memory @@ -150,11 +147,35 @@ impl crate::Agent { Some(ExploreJob { slot: n, task, - provider: self.provider.clone(), + provider: self.explore_child_provider(), child_config, }) } + /// The provider explore children run on. Shares the driver's connection + /// unless an `explore_endpoint` is configured (team roles), in which case + /// recon runs on its own OpenAI-compatible route — typically a local + /// model, so read-heavy fan-out costs nothing. + pub(crate) fn explore_child_provider(&self) -> std::sync::Arc { + routed_provider( + self.config.subagents.explore_endpoint.as_deref(), + self.config.subagents.explore_endpoint_key.as_deref(), + &self.provider, + ) + } + + /// The provider in-process background `delegate` tasks run on — the + /// delegate route when configured, else the driver's provider. (The + /// synchronous delegate path applies the same route in its child-process + /// runner instead.) + pub(crate) fn delegate_child_provider(&self) -> std::sync::Arc { + routed_provider( + self.config.subagents.delegate_endpoint.as_deref(), + self.config.subagents.delegate_endpoint_key.as_deref(), + &self.provider, + ) + } + /// Run one read-only `explore` subagent for the `{task}` argument and return /// its answer as the tool result. Best-effort: a provider/parse error becomes /// an error string fed back to the model, never fatal to the parent turn. @@ -360,10 +381,83 @@ impl Ui for BufferingUi { fn turn_end(&mut self, _summary: &str) {} } +/// The model explore children run: `HI_EXPLORE_MODEL` env (highest, a live +/// escape hatch) → `subagents.explore_model` (team roles) → the driver model. +pub(crate) fn explore_child_model(config: &crate::AgentConfig) -> String { + std::env::var("HI_EXPLORE_MODEL") + .ok() + .filter(|model| !model.trim().is_empty()) + .or_else(|| { + config + .subagents + .explore_model + .clone() + .filter(|model| !model.trim().is_empty()) + }) + .unwrap_or_else(|| config.routing.model.clone()) +} + +/// Build the provider for a routed subagent role: a dedicated +/// OpenAI-compatible client when an endpoint override is set, else the +/// driver's shared provider. Construction is cheap (one HTTP client), so +/// routed children build per spawn rather than caching. +pub(crate) fn routed_provider( + endpoint: Option<&str>, + api_key: Option<&str>, + parent: &std::sync::Arc, +) -> std::sync::Arc { + match endpoint.map(str::trim).filter(|url| !url.is_empty()) { + Some(url) => std::sync::Arc::new(hi_ai::OpenAiProvider::new( + url.to_string(), + api_key.unwrap_or_default().to_string(), + )), + None => parent.clone(), + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn explore_child_model_prefers_config_route_over_driver() { + // Env is deliberately not exercised here (global, races other tests); + // it keeps the highest precedence as a live escape hatch. + let mut config = crate::AgentConfig::default(); + config.routing.model = "pipe/glm-5.2".into(); + assert_eq!(explore_child_model(&config), "pipe/glm-5.2", "inherits"); + config.subagents.explore_model = Some("qwen3-4b".into()); + assert_eq!(explore_child_model(&config), "qwen3-4b", "team route wins"); + config.subagents.explore_model = Some(" ".into()); + assert_eq!( + explore_child_model(&config), + "pipe/glm-5.2", + "blank override is ignored" + ); + } + + #[test] + fn routed_provider_shares_the_driver_unless_an_endpoint_is_set() { + let parent: std::sync::Arc = std::sync::Arc::new( + hi_ai::OpenAiProvider::new("http://127.0.0.1:1/v1".into(), "k".into()), + ); + let inherited = routed_provider(None, None, &parent); + assert!( + std::sync::Arc::ptr_eq(&parent, &inherited), + "no endpoint → the driver's shared connection" + ); + let routed = routed_provider(Some("http://127.0.0.1:18080/v1"), None, &parent); + assert!( + !std::sync::Arc::ptr_eq(&parent, &routed), + "an endpoint override gets its own provider" + ); + let blank = routed_provider(Some(" "), None, &parent); + assert!( + std::sync::Arc::ptr_eq(&parent, &blank), + "blank endpoint is ignored" + ); + } + #[test] fn child_prompt_stays_plain_but_has_a_read_only_task_contract() { let prompt = explore_child_prompt("count the Rust source lines"); diff --git a/crates/hi-agent/src/agent/lifecycle.rs b/crates/hi-agent/src/agent/lifecycle.rs index 54035422..dc3c0847 100644 --- a/crates/hi-agent/src/agent/lifecycle.rs +++ b/crates/hi-agent/src/agent/lifecycle.rs @@ -133,6 +133,7 @@ impl crate::Agent { provider, skeptic_provider, local_skeptic: None, + team_local_servers: Vec::new(), config, runtime, task: crate::domain::TaskContextState::default(), @@ -827,10 +828,14 @@ impl crate::Agent { } /// Stop every background process owned by this agent runtime, plus any - /// auto-managed local skeptic server, on session shutdown. + /// auto-managed local skeptic server and team-role model servers, on + /// session shutdown. pub fn kill_background_processes(&self) { self.runtime.background().kill_all(); self.stop_local_skeptic_server(); + for server in &self.team_local_servers { + hi_tools::stop_local_server(&server.process_id); + } // Background subagent tasks are cleaned up via BackgroundTaskRegistry's // Drop impl when the agent is dropped. The async `kill_all` method can // be called from async cleanup paths if needed. @@ -1750,6 +1755,123 @@ impl crate::Agent { self.config.loop_limits.max_steps = u32::MAX; } + /// The team-role table for `/team`: each role with the model and route it + /// currently runs on. Roles that inherit the driver say so explicitly. + pub fn team_roles(&self) -> Vec { + let driver_model = self.config.routing.model.clone(); + let driver_route = self + .config + .routing + .provider_route + .clone() + .unwrap_or_else(|| "driver provider".to_string()); + let sub = &self.config.subagents; + let role = |role: &'static str, + model: &Option, + endpoint: &Option| + -> crate::TeamRole { + let inherited = model.is_none() && endpoint.is_none(); + crate::TeamRole { + role, + model: model.clone().unwrap_or_else(|| driver_model.clone()), + route: endpoint + .clone() + .unwrap_or_else(|| driver_route.clone()), + inherited, + } + }; + vec![ + crate::TeamRole { + role: "driver", + model: driver_model.clone(), + route: driver_route.clone(), + inherited: false, + }, + role("explore", &sub.explore_model, &sub.explore_endpoint), + role("delegate", &sub.delegate_model, &sub.delegate_endpoint), + role("skeptic", &sub.skeptic_model, &sub.skeptic_endpoint), + role("planner", &sub.planner_model, &None), + ] + } + + /// Point the write-capable `delegate` executors at a different model + /// and/or OpenAI-compatible endpoint (`None`s inherit the driver). + /// Applies to delegates started after the call. + pub fn set_delegate_route( + &mut self, + model: Option, + endpoint: Option, + api_key: Option, + ) { + self.config.subagents.delegate_model = normalized(model); + self.config.subagents.delegate_endpoint = normalized(endpoint); + self.config.subagents.delegate_endpoint_key = normalized(api_key); + } + + /// Point read-only `explore` recon children at a different model and/or + /// endpoint (`None`s inherit the driver). Applies to explores started + /// after the call. + pub fn set_explore_route( + &mut self, + model: Option, + endpoint: Option, + api_key: Option, + ) { + self.config.subagents.explore_model = normalized(model); + self.config.subagents.explore_endpoint = normalized(endpoint); + self.config.subagents.explore_endpoint_key = normalized(api_key); + } + + /// Set or clear the goal-decomposition planner model (`/team planner`). + pub fn set_planner_model(&mut self, model: Option) { + self.config.subagents.planner_model = normalized(model); + } + + /// The auto-managed local model server, when one is running (started by + /// `/config skeptic-local on`): `(base_url, model_id)`. `/team + /// local` reuses it so a role can move on-device with one command. + pub fn managed_local_route(&self) -> Option<(String, String)> { + self.local_skeptic + .as_ref() + .map(|state| (state.endpoint.clone(), state.model_id.clone())) + } + + /// A running managed server (skeptic or team) already serving `model_id`, + /// if any — `/team` reuses it instead of spawning a duplicate. + pub fn running_local_model_server(&self, model_id: &str) -> Option<(String, String)> { + if let Some(state) = &self.local_skeptic + && state.model_id == model_id + { + return Some((state.endpoint.clone(), state.model_id.clone())); + } + self.team_local_servers + .iter() + .find(|server| server.model_id == model_id) + .map(|server| (server.endpoint.clone(), server.model_id.clone())) + } + + /// Record a provisioned team-role server so later `/team` picks of the + /// same model reuse it and session teardown can stop it. + pub fn register_team_local_server( + &mut self, + endpoint: String, + model_id: String, + process_id: String, + ) { + if self + .team_local_servers + .iter() + .any(|server| server.process_id == process_id) + { + return; + } + self.team_local_servers.push(crate::TeamLocalServer { + process_id, + endpoint, + model_id, + }); + } + pub fn rsi_status(&self) -> (&'static str, &'static str, Option) { let requested = if self.config.rsi.enabled { "on" } else { "off" }; let mode = if self.config.rsi.managed { @@ -1906,3 +2028,10 @@ impl crate::Agent { self.messages.mutate_slice() } } + +/// Trim a role-route input; empty strings mean "inherit" (`None`). +fn normalized(value: Option) -> Option { + value + .map(|v| v.trim().to_string()) + .filter(|v| !v.is_empty()) +} diff --git a/crates/hi-agent/src/command.rs b/crates/hi-agent/src/command.rs index 7a6a1f8b..4a46477a 100644 --- a/crates/hi-agent/src/command.rs +++ b/crates/hi-agent/src/command.rs @@ -17,6 +17,10 @@ pub enum Command { Rsi(String), /// Run exactly one turn through the conservative MoA virtual route. Moa(String), + /// `/team` — show or assign per-role model routes (driver plans on the + /// big model; explore/delegate executors can run locally). Empty arg + /// shows the role table. + Team(String), /// Use a provider/profile for subsequent turns (empty = report current). /// Named profiles are resolved from the config; the model can be set later /// with `/model` when the profile does not configure one. @@ -232,6 +236,7 @@ pub fn parse(line: &str) -> Option { "config" | "cfg" | "set" => Command::Config(arg), "rsi" => Command::Rsi(arg), "moa" => Command::Moa(arg), + "team" | "roles" => Command::Team(arg), "provider" | "prov" => Command::Provider(arg), "login" | "signin" => Command::Login(arg), "logout" | "signout" => Command::Logout(arg), @@ -1454,6 +1459,16 @@ pub const COMMANDS: &[CommandSpec] = &[ help: "run one prompt through moa/conservative, then restore the current model", arg_values: &[], }, + CommandSpec { + name: "team", + args: "[role] [model|local|off] [base-url] [api-key]", + help: "role routes: driver plans, explore/delegate can run on other models (e.g. local)", + arg_values: &[ + ("explore", "route read-only recon children"), + ("delegate", "route write-capable executors"), + ("planner", "route goal decomposition"), + ], + }, CommandSpec { name: "provider", args: "[name|add|edit|remove]", diff --git a/crates/hi-agent/src/config.rs b/crates/hi-agent/src/config.rs index 72134b58..3c81ebd2 100644 --- a/crates/hi-agent/src/config.rs +++ b/crates/hi-agent/src/config.rs @@ -587,6 +587,23 @@ pub struct AgentSubagents { pub skeptic_endpoint: Option, /// API key sent to `skeptic_endpoint`. pub skeptic_endpoint_key: Option, + /// Model id for write-capable `delegate` executors (`None` = the driver's + /// model). With `delegate_endpoint`, lets a big cloud driver dispatch + /// execution to a local model (team roles: big brain plans, local hands + /// type). + pub delegate_model: Option, + /// Optional OpenAI-compatible base URL for delegate executors only. + pub delegate_endpoint: Option, + /// API key sent to `delegate_endpoint`. + pub delegate_endpoint_key: Option, + /// Model id for read-only `explore` recon children (`None` = the driver's + /// model). The `HI_EXPLORE_MODEL` env var still wins when set. + pub explore_model: Option, + /// Optional OpenAI-compatible base URL for explore children only. When + /// unset, explore children share the driver's provider connection. + pub explore_endpoint: Option, + /// API key sent to `explore_endpoint`. + pub explore_endpoint_key: Option, } impl Default for AgentSubagents { @@ -600,6 +617,12 @@ impl Default for AgentSubagents { skeptic_model: None, skeptic_endpoint: None, skeptic_endpoint_key: None, + delegate_model: None, + delegate_endpoint: None, + delegate_endpoint_key: None, + explore_model: None, + explore_endpoint: None, + explore_endpoint_key: None, } } } diff --git a/crates/hi-agent/src/lib.rs b/crates/hi-agent/src/lib.rs index dcfc004a..9b358b6b 100644 --- a/crates/hi-agent/src/lib.rs +++ b/crates/hi-agent/src/lib.rs @@ -109,7 +109,7 @@ pub use skills::{ build_learn_prompt, build_skill_use_prompt, learned_skills_context, list_skills, read_skill, skill_roots, }; -pub use subagent::{DelegateOutcome, DelegateRunner}; +pub use subagent::{DelegateOutcome, DelegateRunner, SubagentRoute}; pub use task_contract::{RiskLevel, TaskContract, TaskIntent}; pub use ui::{ ConfirmationFuture, ConfirmationRequest, ConfirmationResult, Ui, classify_error, tool_label, @@ -245,6 +245,25 @@ pub struct ConfigSnapshot { pub moe_streaming: String, } +/// A managed local model server provisioned for a team role. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TeamLocalServer { + pub process_id: String, + pub endpoint: String, + pub model_id: String, +} + +/// One row of the `/team` role table: which model and route a role runs on. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TeamRole { + /// `driver`, `explore`, `delegate`, `skeptic`, or `planner`. + pub role: &'static str, + pub model: String, + pub route: String, + /// True when the role has no override and follows the driver. + pub inherited: bool, +} + /// Provider-neutral wall-clock latency buckets for one turn. #[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct TurnPhaseLatencies { @@ -708,6 +727,11 @@ pub struct Agent { /// stopped, the prior skeptic settings restored, and the process killed on /// session shutdown. pub(crate) local_skeptic: Option, + /// Managed local model servers provisioned for team roles (`/team + /// delegate coder-14b` etc.). Reused across roles that pick the same + /// model; torn down with the session (the frontend's blanket + /// `stop_all_local_servers` guard also covers them). + pub(crate) team_local_servers: Vec, pub(crate) config: AgentConfig, pub(crate) runtime: WorkspaceRuntime, /// Per-turn ranked task/memory prompt assembly. diff --git a/crates/hi-agent/src/local_skeptic.rs b/crates/hi-agent/src/local_skeptic.rs index d4901fd3..d1f081ee 100644 --- a/crates/hi-agent/src/local_skeptic.rs +++ b/crates/hi-agent/src/local_skeptic.rs @@ -35,6 +35,17 @@ impl LocalBackend { LocalBackend::Cuda => "cuda", } } + + /// The `hi-local` cargo feature that compiles this backend in. The + /// backends are opt-in features (`default = []`), so a plain + /// `cargo build -p hi-local` produces a binary that instantly rejects + /// `--backend mlx` — the exact failure users saw as a silent dead server. + pub fn cargo_feature(self) -> &'static str { + match self { + LocalBackend::Mlx => "mlx", + LocalBackend::Cuda => "native-cuda", + } + } } /// Choose a backend from probed hardware facts. Pure so it can be unit tested; @@ -50,8 +61,15 @@ pub fn pick_backend(is_apple_silicon: bool, has_nvidia: bool) -> Option Option { + match std::env::var("HI_LOCAL_BACKEND").ok().as_deref() { + Some("mlx") => return Some(LocalBackend::Mlx), + Some("cuda") => return Some(LocalBackend::Cuda), + Some("none") => return None, + _ => {} + } let is_apple_silicon = cfg!(all(target_os = "macos", target_arch = "aarch64")); let has_nvidia = !is_apple_silicon && nvidia_present(); pick_backend(is_apple_silicon, has_nvidia) @@ -70,7 +88,7 @@ fn nvidia_present() -> bool { /// `detect_backend` runs a blocking `nvidia-smi` subprocess; offload it so it /// doesn't stall the async executor when called from an async context. -async fn detect_backend_offload() -> Option { +pub async fn detect_backend_offload() -> Option { tokio::task::spawn_blocking(detect_backend) .await .unwrap_or(None) @@ -140,14 +158,63 @@ pub fn resolve_model(backend: LocalBackend) -> LocalModelSpec { /// Whether the model's weights are already cached in `dir`. pub fn model_present(dir: &Path, spec: &LocalModelSpec) -> bool { + // An in-flight or interrupted fetch leaves `.aria2` control files + // behind; their presence means the weights can't be trusted yet, no + // matter what else already landed. + if aria2_remnants(dir) { + return false; + } match &spec.gguf_file { - // MLX: a loadable model directory carries a config.json (matches `/hf`). - None => dir.join("config.json").exists(), + // MLX: a loadable model directory carries a config.json (matches + // `/hf`) — and when the repo ships a safetensors index, every weight + // shard it names. config.json downloads first, so checking it alone + // would bless a directory whose later shards never started. + None => { + if !dir.join("config.json").exists() { + return false; + } + match indexed_weight_files(dir) { + Some(files) => files.iter().all(|file| dir.join(file).exists()), + None => any_safetensors(dir), + } + } // CUDA: the specific GGUF file must be on disk. Some(file) => dir.join(file).exists(), } } +/// Whether `dir` holds any aria2c control files (partial downloads). +fn aria2_remnants(dir: &Path) -> bool { + std::fs::read_dir(dir).is_ok_and(|entries| { + entries + .flatten() + .any(|entry| entry.path().extension().is_some_and(|ext| ext == "aria2")) + }) +} + +/// The unique weight files named by `model.safetensors.index.json`, when the +/// repo ships one (multi-shard models do). +fn indexed_weight_files(dir: &Path) -> Option> { + let raw = std::fs::read_to_string(dir.join("model.safetensors.index.json")).ok()?; + let index: serde_json::Value = serde_json::from_str(&raw).ok()?; + let map = index.get("weight_map")?.as_object()?; + let mut files: Vec = map + .values() + .filter_map(|value| value.as_str().map(str::to_string)) + .collect(); + files.sort(); + files.dedup(); + Some(files) +} + +fn any_safetensors(dir: &Path) -> bool { + std::fs::read_dir(dir).is_ok_and(|entries| { + entries + .flatten() + .any(|entry| entry.path().extension().is_some_and(|ext| ext == "safetensors")) + }) +} + /// The path passed to `hi-local serve`: the model *directory* for MLX, the GGUF /// *file* for CUDA. pub fn serve_model_path(dir: &Path, spec: &LocalModelSpec) -> PathBuf { @@ -304,19 +371,22 @@ impl Agent { } let abs_dir = std::fs::canonicalize(&dir).unwrap_or(dir); let model_path = serve_model_path(&abs_dir, &spec); - let bin = find_hi_local(); + let bin = ensure_hi_local_binary(backend).await?; let host = "127.0.0.1"; let port = pick_free_port(); let args = serve_args(&model_path, &spec, host, port); - let handle = hi_tools::start_local_server(&bin, &args, host, port) - .await - .with_context(|| { - format!( - "starting hi-local ({}) — is the `hi-local` binary built with the {} backend?", - bin.display(), - spec.backend.serve_flag() - ) - })?; + let deadline = health_deadline_for_model(model_dir_bytes(&abs_dir)); + let handle = + hi_tools::start_local_server_with_deadline(&bin, &args, host, port, deadline) + .await + .with_context(|| { + format!( + "hi-local ({}) did not become ready within {}s — is it built with the {} backend?", + bin.display(), + deadline.as_secs(), + spec.backend.serve_flag() + ) + })?; let prev_skeptic_model = self.config.subagents.skeptic_model.clone(); let prev_endpoint = self.config.subagents.skeptic_endpoint.clone(); @@ -361,3 +431,825 @@ impl Agent { } } } + +// ---- Supported team-role local models ------------------------------------ +// +// `/team ` provisioning: the user picks a short supported name +// (or just `local`), and hi does everything — hardware-sized selection, +// download, server spawn, health wait, role wiring. Nobody types endpoints. + +/// One quantization of a model's MLX form. pipenetwork publishes whole quant +/// ladders (2bit→8bit) per model, so the same model fits very different +/// machines — a 3bit Laguna runs where the 4bit can't. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct MlxQuant { + /// Short quant tag users can force with `name@quant` (e.g. `4bit`). + pub quant: &'static str, + /// Comfortable-fit floor for this quant (weights + KV/OS headroom). + pub min_ram_gb: u64, + pub repo: &'static str, + pub model_id: &'static str, +} + +/// A verified GGUF form for the CUDA backend. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct CudaGguf { + pub min_ram_gb: u64, + pub repo: &'static str, + pub gguf_file: &'static str, + pub model_id: &'static str, +} + +/// One curated local model users can select by short name. MLX is the +/// primary form (verified pipenetwork/mlx-community conversions), carried as +/// a quant ladder ordered best-quality-first; the CUDA GGUF form is optional +/// — entries without a verified GGUF are honestly MLX-only rather than +/// pointing at guessed repos. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct SupportedLocalModel { + /// What the user types: `/team delegate ` (or `name@quant`). + pub name: &'static str, + /// One-line human description for the picker. + pub label: &'static str, + /// MLX quant ladder, **best quality first** — selection takes the first + /// quant whose floor fits, so bigger machines get better weights. + pub mlx: &'static [MlxQuant], + pub cuda: Option, +} + +impl SupportedLocalModel { + /// Highest-quality MLX quant that fits `ram_gb`. + pub fn pick_mlx(&self, ram_gb: u64) -> Option<&'static MlxQuant> { + self.mlx.iter().find(|quant| ram_gb >= quant.min_ram_gb) + } + + /// The smallest MLX quant — what an explicit oversized pick falls to. + pub fn smallest_mlx(&self) -> Option<&'static MlxQuant> { + self.mlx.iter().min_by_key(|quant| quant.min_ram_gb) + } + + /// Smallest workable floor on `backend` (`None` = backend unknown: + /// smallest floor across all published forms). + pub fn min_ram_gb(&self, backend: Option) -> u64 { + let mlx_floor = self.smallest_mlx().map(|quant| quant.min_ram_gb); + let cuda_floor = self.cuda.map(|cuda| cuda.min_ram_gb); + match backend { + Some(LocalBackend::Mlx) => mlx_floor.unwrap_or(u64::MAX), + Some(LocalBackend::Cuda) => cuda_floor.unwrap_or(u64::MAX), + None => mlx_floor.min(cuda_floor).unwrap_or(u64::MAX), + } + } + + /// Whether any published form of this model fits `ram_gb` on `backend`. + pub fn fits(&self, ram_gb: u64, backend: Option) -> bool { + ram_gb >= self.min_ram_gb(backend) + } + + /// Compact ladder summary for the picker, e.g. `8/6/4/3/2bit`. + pub fn quant_summary(&self) -> String { + let tags: Vec<&str> = self + .mlx + .iter() + .map(|quant| quant.quant.strip_suffix("bit").unwrap_or(quant.quant)) + .collect(); + let all_bits = self.mlx.iter().all(|quant| quant.quant.ends_with("bit")); + if all_bits { format!("{}bit", tags.join("/")) } else { tags.join("/") } + } +} + +/// A catalog entry resolved for a specific machine: the family plus the MLX +/// quant chosen for it (highest quality that fits, an explicit `@quant`, or +/// the smallest quant when nothing fits an explicit pick). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ResolvedLocalModel { + pub entry: &'static SupportedLocalModel, + pub mlx: Option<&'static MlxQuant>, +} + +impl ResolvedLocalModel { + /// User-facing name: `laguna-s@3bit` when the family has a real ladder, + /// the bare family name when there is only one form. + pub fn display(&self) -> String { + match self.mlx { + Some(quant) if self.entry.mlx.len() > 1 => { + format!("{}@{}", self.entry.name, quant.quant) + } + _ => self.entry.name.to_string(), + } + } +} + +/// The supported catalog. **Order = auto-selection preference**: `local` +/// takes the first entry with any quant that fits the machine, so +/// newer/faster models come before older peers. Within an entry the quant +/// ladder is best-quality-first. Entries after `mini` (which fits almost +/// everywhere) are explicit-pick / picker-only. Every repo below was +/// verified against the pipenetwork / mlx-community HF listings — no +/// guessed names. +pub const SUPPORTED_LOCAL_MODELS: &[SupportedLocalModel] = &[ + SupportedLocalModel { + name: "coder-32b", + label: "Qwen2.5-Coder 32B — dense coder (~19GB)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 40, + repo: "mlx-community/Qwen2.5-Coder-32B-Instruct-4bit", + model_id: "Qwen2.5-Coder-32B-Instruct-4bit", + }], + cuda: Some(CudaGguf { + min_ram_gb: 40, + repo: "Qwen/Qwen2.5-Coder-32B-Instruct-GGUF", + gguf_file: "qwen2.5-coder-32b-instruct-q4_k_m.gguf", + model_id: "qwen2.5-coder-32b-instruct", + }), + }, + SupportedLocalModel { + name: "coder-14b", + label: "Qwen2.5-Coder 14B — dense coder (~9GB)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 24, + repo: "mlx-community/Qwen2.5-Coder-14B-Instruct-4bit", + model_id: "Qwen2.5-Coder-14B-Instruct-4bit", + }], + cuda: Some(CudaGguf { + min_ram_gb: 24, + repo: "Qwen/Qwen2.5-Coder-14B-Instruct-GGUF", + gguf_file: "qwen2.5-coder-14b-instruct-q4_k_m.gguf", + model_id: "qwen2.5-coder-14b-instruct", + }), + }, + SupportedLocalModel { + name: "coder-7b", + label: "Qwen2.5-Coder 7B — dense coder (~5GB)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 12, + repo: "mlx-community/Qwen2.5-Coder-7B-Instruct-4bit", + model_id: "Qwen2.5-Coder-7B-Instruct-4bit", + }], + cuda: Some(CudaGguf { + min_ram_gb: 12, + repo: "Qwen/Qwen2.5-Coder-7B-Instruct-GGUF", + gguf_file: "qwen2.5-coder-7b-instruct-q4_k_m.gguf", + model_id: "qwen2.5-coder-7b-instruct", + }), + }, + SupportedLocalModel { + name: "nemotron-4b", + label: "NVIDIA Nemotron 3 Nano 4B — tiny modern executor (~2-4GB)", + mlx: &[ + MlxQuant { + quant: "8bit", + min_ram_gb: 8, + repo: "pipenetwork/NVIDIA-Nemotron-3-Nano-4B-MLX-8bit", + model_id: "NVIDIA-Nemotron-3-Nano-4B-MLX-8bit", + }, + MlxQuant { + quant: "4bit", + min_ram_gb: 6, + repo: "pipenetwork/NVIDIA-Nemotron-3-Nano-4B-MLX-4bit", + model_id: "NVIDIA-Nemotron-3-Nano-4B-MLX-4bit", + }, + ], + cuda: None, + }, + SupportedLocalModel { + name: "mini", + label: "Qwen2.5 3B — tiny generalist for recon/review (~2GB)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 6, + repo: "mlx-community/Qwen2.5-3B-Instruct-4bit", + model_id: "Qwen2.5-3B-Instruct-4bit", + }], + cuda: Some(CudaGguf { + min_ram_gb: 6, + repo: "Qwen/Qwen2.5-3B-Instruct-GGUF", + gguf_file: "qwen2.5-3b-instruct-q4_k_m.gguf", + model_id: "qwen2.5-3b-instruct", + }), + }, + // Explicit-pick / picker-only from here down (mini already fits + // almost every machine, so auto never reaches these). A model enters the + // auto section above only after `hi team-bench` has served it and passed + // tasks on real hardware — live verification found laguna-s generating + // ~0 tok/s through the MoE path and qwen3.6's nvfp4 quant unsupported, + // so they wait here, picker-visible and honestly labeled, until the + // serving gaps close. + SupportedLocalModel { + name: "nemotron-30b", + label: "NVIDIA Nemotron 3 Nano 30B-A3B — fast but scored 1/4 on team-bench (testing only)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 24, + repo: "pipenetwork/Nemotron-3-Nano-30B-A3B-context-mlx-4bit", + model_id: "Nemotron-3-Nano-30B-A3B-context-mlx-4bit", + }], + cuda: None, + }, + SupportedLocalModel { + name: "laguna-s", + label: "Poolside Laguna S 2.1 — 118B MoE; hi-local serves it too slowly yet (testing only)", + mlx: &[ + MlxQuant { + quant: "8bit", + min_ram_gb: 192, + repo: "pipenetwork/Laguna-S-2.1-MLX-8bit", + model_id: "Laguna-S-2.1-MLX-8bit", + }, + MlxQuant { + quant: "6bit", + min_ram_gb: 128, + repo: "pipenetwork/Laguna-S-2.1-MLX-6bit", + model_id: "Laguna-S-2.1-MLX-6bit", + }, + MlxQuant { + quant: "4bit", + min_ram_gb: 72, + repo: "pipenetwork/Laguna-S-2.1-MLX-4bit", + model_id: "Laguna-S-2.1-MLX-4bit", + }, + MlxQuant { + quant: "3bit", + min_ram_gb: 64, + repo: "pipenetwork/Laguna-S-2.1-MLX-3bit", + model_id: "Laguna-S-2.1-MLX-3bit", + }, + MlxQuant { + quant: "2bit", + min_ram_gb: 48, + repo: "pipenetwork/Laguna-S-2.1-MLX-2bit", + model_id: "Laguna-S-2.1-MLX-2bit", + }, + ], + cuda: None, + }, + SupportedLocalModel { + name: "qwen3.6-35b", + label: "Qwen3.6 35B-A3B — needs nvfp4 support hi-local doesn't have yet", + mlx: &[MlxQuant { + quant: "nvfp4", + min_ram_gb: 32, + repo: "pipenetwork/Qwen3.6-35B-A3B-mlx-nvfp4", + model_id: "Qwen3.6-35B-A3B-mlx-nvfp4", + }], + cuda: None, + }, + SupportedLocalModel { + name: "glm-5.2-reap50", + label: "GLM-5.2 with half the experts REAP-pruned (~195GB; 256GB Macs)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 256, + repo: "pipenetwork/GLM-5.2-REAP50-MLX-4bit", + model_id: "GLM-5.2-REAP50-MLX-4bit", + }], + cuda: None, + }, + SupportedLocalModel { + name: "glm-5.2", + label: "GLM-5.2 — the cloud driver's own family, local (~390GB; 512GB Macs)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 448, + repo: "pipenetwork/GLM-5.2-MLX-4bit", + model_id: "GLM-5.2-MLX-4bit", + }], + cuda: None, + }, + SupportedLocalModel { + name: "nemotron-550b", + label: "NVIDIA Nemotron 3 Ultra 550B-A55B — frontier MoE (~275GB; 512GB Macs)", + mlx: &[MlxQuant { + quant: "4bit", + min_ram_gb: 512, + repo: "pipenetwork/NVIDIA-Nemotron-3-Ultra-550B-A55B-MLX-4bit", + model_id: "NVIDIA-Nemotron-3-Ultra-550B-A55B-MLX-4bit", + }], + cuda: None, + }, +]; + +/// Physical memory in GiB, best-effort (0 when unknown). +pub fn system_ram_gb() -> u64 { + #[cfg(target_os = "macos")] + { + let out = std::process::Command::new("sysctl") + .args(["-n", "hw.memsize"]) + .output(); + if let Ok(out) = out + && let Ok(text) = String::from_utf8(out.stdout) + && let Ok(bytes) = text.trim().parse::() + { + return bytes / (1024 * 1024 * 1024); + } + } + #[cfg(target_os = "linux")] + { + if let Ok(text) = std::fs::read_to_string("/proc/meminfo") + && let Some(kb) = parse_meminfo_total_kb(&text) + { + return kb / (1024 * 1024); + } + } + 0 +} + +/// Parse `MemTotal: 16384 kB` from /proc/meminfo. Pure for testing. +pub fn parse_meminfo_total_kb(meminfo: &str) -> Option { + meminfo + .lines() + .find(|line| line.starts_with("MemTotal:"))? + .split_whitespace() + .nth(1)? + .parse() + .ok() +} + +/// Resolve a `/team` model selection to a catalog entry plus the MLX quant +/// chosen for this machine. `local`/`coder`/`auto` take the first catalog +/// entry with any form that fits `ram_gb` on `backend` (quant = highest +/// quality that fits); exact names are honored even oversized (quant = the +/// smallest published, the closest to feasible); `name@quant` forces an +/// exact quant. `None` = not a supported local model (the caller treats the +/// input as a cloud model id instead). +pub fn resolve_team_local_model( + name: &str, + ram_gb: u64, + backend: Option, +) -> Option { + let name = name.trim().to_ascii_lowercase(); + if matches!(name.as_str(), "local" | "coder" | "auto") { + let entry = SUPPORTED_LOCAL_MODELS + .iter() + .find(|entry| entry.fits(ram_gb, backend)) + // Nothing fits → the smallest supported model on this backend, + // so tiny machines still get a working executor rather than an + // arbitrary entry. + .or_else(|| { + SUPPORTED_LOCAL_MODELS + .iter() + .min_by_key(|entry| entry.min_ram_gb(backend)) + })?; + return Some(ResolvedLocalModel { + entry, + mlx: entry.pick_mlx(ram_gb).or_else(|| entry.smallest_mlx()), + }); + } + // `name@quant` forces an exact rung of the ladder. + if let Some((family, quant)) = name.split_once('@') { + let entry = SUPPORTED_LOCAL_MODELS.iter().find(|entry| entry.name == family)?; + let quant = entry.mlx.iter().find(|candidate| candidate.quant == quant)?; + return Some(ResolvedLocalModel { entry, mlx: Some(quant) }); + } + let entry = SUPPORTED_LOCAL_MODELS.iter().find(|entry| entry.name == name)?; + Some(ResolvedLocalModel { + entry, + mlx: entry.pick_mlx(ram_gb).or_else(|| entry.smallest_mlx()), + }) +} + +/// The backend-specific serve spec for a resolved catalog selection. Errors +/// on entries that have no verified GGUF for the CUDA backend. +pub fn team_model_spec( + resolved: ResolvedLocalModel, + backend: LocalBackend, +) -> Result { + match backend { + LocalBackend::Mlx => { + let Some(quant) = resolved.mlx else { + bail!("{} isn't packaged for MLX", resolved.entry.name); + }; + Ok(LocalModelSpec { + repo: quant.repo.to_string(), + model_id: quant.model_id.to_string(), + gguf_file: None, + backend, + }) + } + LocalBackend::Cuda => { + let Some(cuda) = resolved.entry.cuda else { + bail!( + "{} isn't packaged for CUDA yet — pick coder-14b, coder-7b, coder-32b, or mini", + resolved.entry.name + ); + }; + Ok(LocalModelSpec { + repo: cuda.repo.to_string(), + model_id: cuda.model_id.to_string(), + gguf_file: Some(cuda.gguf_file.to_string()), + backend, + }) + } + } +} + +/// Where an in-flight `/team` local-model setup currently is. Published on a +/// watch channel so the UI can narrate honestly: minutes of silence during a +/// 19 GB download or a multi-minute model load are indistinguishable from a +/// hang without this. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum ProvisionPhase { + /// Backend detection / cache checks. + Resolving, + /// Fetching weights from the hub (quiet; the UI reports bytes on disk). + Downloading, + /// Compiling `hi-local` from a development checkout. + BuildingServer, + /// Server spawned; weights are loading into memory. `deadline_secs` is + /// the readiness window; `server_handle` lets the UI sample the server's + /// memory growth against `expected_bytes` for a real progress bar. + LoadingModel { + deadline_secs: u64, + server_handle: String, + expected_bytes: u64, + }, +} + +/// Fully provision a supported local model: detect the backend, download the +/// weights when absent, spawn a managed `hi-local` server, and wait for +/// health. Returns `(endpoint, model_id, process_id)`. Standalone (no Agent +/// borrow) so frontends can run it on a background task and wire the role +/// when it completes — a 15 GB download must never block the UI. Progress +/// phases are published on `progress` as they begin. +pub async fn provision_team_local_model( + resolved: ResolvedLocalModel, + progress: tokio::sync::watch::Sender, +) -> Result<(String, String, String)> { + let _ = progress.send(ProvisionPhase::Resolving); + let Some(backend) = detect_backend_offload().await else { + bail!( + "no local-inference backend detected (needs Apple Silicon MLX or an NVIDIA CUDA runtime)" + ); + }; + let spec = team_model_spec(resolved, backend)?; + let dir = hi_tools::skeptic_model_dir(&spec.repo); + if !model_present(&dir, &spec) { + let _ = progress.send(ProvisionPhase::Downloading); + // Quiet by contract: this runs behind a live TUI — raw downloader + // output painted over the alternate screen once; never again. + hi_tools::download_repo_keep_quiet(&spec.repo, &dir) + .await + .with_context(|| format!("downloading {}", spec.repo))?; + if !model_present(&dir, &spec) { + bail!( + "downloaded {} but its weights are still missing under {}", + spec.repo, + dir.display() + ); + } + } + let abs_dir = std::fs::canonicalize(&dir).unwrap_or(dir); + let model_path = serve_model_path(&abs_dir, &spec); + let bin = ensure_hi_local_binary_with_progress(backend, Some(&progress)).await?; + let host = "127.0.0.1"; + let port = pick_free_port(); + let args = serve_args(&model_path, &spec, host, port); + // Big checkpoints take minutes to load; scale the readiness window to + // the weights instead of failing a healthy-but-still-loading server. + let expected_bytes = model_dir_bytes(&abs_dir); + let deadline = health_deadline_for_model(expected_bytes); + // Spawn first, then wait: publishing the handle lets the UI sample the + // server's memory growth for a real progress bar while weights load. + let process_id = hi_tools::spawn_local_server(&bin, &args)?; + let _ = progress.send(ProvisionPhase::LoadingModel { + deadline_secs: deadline.as_secs(), + server_handle: process_id.clone(), + expected_bytes, + }); + hi_tools::await_local_server_health(&process_id, host, port, deadline) + .await + .with_context(|| { + format!( + "hi-local ({}) failed to become ready (allowed {}s)", + bin.display(), + deadline.as_secs() + ) + })?; + Ok((endpoint_url(host, port), spec.model_id, process_id)) +} + +#[cfg(test)] +mod team_catalog_tests { + use super::*; + + const MLX: Option = Some(LocalBackend::Mlx); + const CUDA: Option = Some(LocalBackend::Cuda); + + #[test] + fn auto_sizing_only_lands_on_serve_verified_models() { + // Live verification (hi team-bench): laguna-s generates ~0 tok/s + // through the current MoE path and qwen3.6's nvfp4 quant can't load, + // so auto must land on the proven dense coder instead. + for ram in [192, 128, 64] { + assert_eq!( + resolve_team_local_model("local", ram, MLX).unwrap().entry.name, + "coder-32b", + "{ram}GB auto-picks the serve-verified dense coder" + ); + } + assert_eq!( + resolve_team_local_model("coder", 24, MLX).unwrap().entry.name, + "coder-14b", + "nemotron-30b scored 1/4 on team-bench — the dense coder owns the 24GB tier" + ); + assert_eq!( + resolve_team_local_model("auto", 16, MLX).unwrap().entry.name, + "coder-7b" + ); + assert_eq!( + resolve_team_local_model("local", 4, MLX).unwrap().entry.name, + "nemotron-4b", + "tiny machines still get a working executor" + ); + } + + #[test] + fn quant_ladders_pick_highest_quality_that_fits_on_explicit_names() { + assert_eq!( + resolve_team_local_model("laguna-s", 192, MLX).unwrap().mlx.unwrap().quant, + "8bit", + "highest quality that fits wins" + ); + assert_eq!( + resolve_team_local_model("laguna-s", 128, MLX).unwrap().mlx.unwrap().quant, + "6bit" + ); + let sixty_four = resolve_team_local_model("laguna-s", 64, MLX).unwrap(); + assert_eq!(sixty_four.mlx.unwrap().quant, "3bit"); + assert_eq!(sixty_four.display(), "laguna-s@3bit"); + } + + #[test] + fn auto_sizing_on_cuda_skips_mlx_only_families() { + assert_eq!( + resolve_team_local_model("local", 128, CUDA).unwrap().entry.name, + "coder-32b", + "the biggest verified-GGUF entry wins on a big CUDA box" + ); + assert_eq!( + resolve_team_local_model("local", 32, CUDA).unwrap().entry.name, + "coder-14b", + "nemotron-30b has no verified GGUF — auto must not pick it on CUDA" + ); + assert_eq!( + resolve_team_local_model("local", 4, CUDA).unwrap().entry.name, + "mini", + "the CUDA fallback is the smallest entry that actually serves on CUDA" + ); + } + + #[test] + fn explicit_quants_and_oversized_picks_are_honored() { + let forced = resolve_team_local_model("laguna-s@2bit", 512, MLX).unwrap(); + assert_eq!(forced.mlx.unwrap().quant, "2bit", "@quant beats auto quality pick"); + assert!( + resolve_team_local_model("laguna-s@5bit", 512, MLX).is_none(), + "unpublished quants don't resolve" + ); + let oversized = resolve_team_local_model("laguna-s", 8, MLX).unwrap(); + assert_eq!( + oversized.mlx.unwrap().quant, + "2bit", + "an explicit pick below every floor falls to the smallest quant" + ); + assert_eq!( + resolve_team_local_model("coder-32b", 8, MLX).unwrap().entry.name, + "coder-32b", + "an explicit pick is honored even below the sizing hint" + ); + assert_eq!( + resolve_team_local_model("glm-5.2-reap50", 8, MLX).unwrap().entry.name, + "glm-5.2-reap50" + ); + assert_eq!( + resolve_team_local_model("coder-7b", 64, MLX).unwrap().display(), + "coder-7b", + "single-form families display without a quant suffix" + ); + assert!(resolve_team_local_model("pipe/glm-4-flash", 64, MLX).is_none()); + assert!(resolve_team_local_model("qwen3-anything", 64, MLX).is_none()); + } + + #[test] + fn specs_map_backend_forms_and_cuda_gaps_are_honest() { + let resolved = resolve_team_local_model("coder-7b", 64, MLX).unwrap(); + let mlx = team_model_spec(resolved, LocalBackend::Mlx).unwrap(); + assert!(mlx.repo.starts_with("mlx-community/")); + assert_eq!(mlx.gguf_file, None); + let cuda = team_model_spec(resolved, LocalBackend::Cuda).unwrap(); + assert!(cuda.gguf_file.is_some()); + + let laguna = resolve_team_local_model("laguna-s", 128, MLX).unwrap(); + let spec = team_model_spec(laguna, LocalBackend::Mlx).unwrap(); + assert_eq!( + spec.repo, "pipenetwork/Laguna-S-2.1-MLX-6bit", + "the spec serves exactly the quant the resolution chose" + ); + assert!( + team_model_spec(laguna, LocalBackend::Cuda).is_err(), + "no guessed GGUF repos: CUDA gap is an honest error" + ); + } + + #[test] + fn model_present_rejects_partial_downloads() { + let dir = std::env::temp_dir().join(format!("hi-present-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&dir); + std::fs::create_dir_all(&dir).unwrap(); + let spec = LocalModelSpec { + repo: "x/y".into(), + model_id: "y".into(), + gguf_file: None, + backend: LocalBackend::Mlx, + }; + assert!(!model_present(&dir, &spec), "empty dir"); + std::fs::write(dir.join("config.json"), "{}").unwrap(); + assert!(!model_present(&dir, &spec), "config alone is not a model"); + std::fs::write( + dir.join("model.safetensors.index.json"), + r#"{"weight_map": {"a": "model-00001-of-00002.safetensors", "b": "model-00002-of-00002.safetensors"}}"#, + ) + .unwrap(); + std::fs::write(dir.join("model-00001-of-00002.safetensors"), "w").unwrap(); + assert!(!model_present(&dir, &spec), "a shard the index names is missing"); + std::fs::write(dir.join("model-00002-of-00002.safetensors"), "w").unwrap(); + assert!(model_present(&dir, &spec), "all shards present"); + std::fs::write(dir.join("model-00002-of-00002.safetensors.aria2"), "ctl").unwrap(); + assert!( + !model_present(&dir, &spec), + "an aria2 control file means the download isn't done" + ); + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn meminfo_parses() { + assert_eq!( + parse_meminfo_total_kb("MemTotal: 16332852 kB\nMemFree: 1 kB"), + Some(16_332_852) + ); + assert_eq!(parse_meminfo_total_kb("nope"), None); + } +} + +/// Health deadline scaled to the model on disk: loading weights into memory +/// dominates startup, so allow ~15s per GiB on top of a 60s floor, capped at +/// ten minutes. A 3B review model waits ~90s; a 19 GB 32B executor gets the +/// minutes it genuinely needs instead of failing at a flat 15s. +pub fn health_deadline_for_model(model_bytes: u64) -> std::time::Duration { + const FLOOR_SECS: u64 = 60; + const PER_GIB_SECS: u64 = 15; + const CAP_SECS: u64 = 600; + let gib = model_bytes / (1024 * 1024 * 1024); + std::time::Duration::from_secs((FLOOR_SECS + gib * PER_GIB_SECS).min(CAP_SECS)) +} + +/// Total bytes directly inside `dir` (model repos download flat). +fn model_dir_bytes(dir: &Path) -> u64 { + std::fs::read_dir(dir) + .map(|entries| { + entries + .flatten() + .filter_map(|entry| entry.metadata().ok()) + .filter(|meta| meta.is_file()) + .map(|meta| meta.len()) + .sum() + }) + .unwrap_or(0) +} + +/// Whether `name` resolves to an executable on PATH. +fn binary_on_path(name: &str) -> bool { + let Some(paths) = std::env::var_os("PATH") else { + return false; + }; + std::env::split_paths(&paths).any(|dir| { + let candidate = dir.join(name); + candidate.is_file() + }) +} + +/// Locate the `hi-local` serving binary, building it when running from a +/// development checkout. Full automation means nobody gets told to go run +/// cargo by hand: if the binary is missing but the current executable lives +/// under a cargo `target/` directory, build `hi-local` for that same profile +/// (quietly — output captured) and use the fresh sibling. +pub async fn ensure_hi_local_binary(backend: LocalBackend) -> Result { + ensure_hi_local_binary_with_progress(backend, None).await +} + +/// The cargo invocation that builds `hi-local` with `backend` compiled in. +/// Pure for testing — getting this wrong shipped a binary that couldn't +/// serve MLX at all. +pub fn hi_local_build_args(profile: &str, backend: LocalBackend) -> Vec { + let mut args = vec![ + "build".to_string(), + "-p".to_string(), + "hi-local".to_string(), + "--features".to_string(), + backend.cargo_feature().to_string(), + ]; + if profile == "release" { + args.push("--release".to_string()); + } + args +} + +/// [`ensure_hi_local_binary`], announcing a dev-checkout compile on +/// `progress`. In a development checkout (`target//hi`) this +/// ALWAYS runs cargo with the backend feature — a stale or feature-less +/// sibling binary must never be served (cargo is a fast no-op when the +/// build is fresh). Installed layouts trust the sibling/PATH binary. +pub async fn ensure_hi_local_binary_with_progress( + backend: LocalBackend, + progress: Option<&tokio::sync::watch::Sender>, +) -> Result { + // Development checkout: /target//hi → (re)build the + // serving binary with the right backend feature, into the same profile. + if let Ok(current) = std::env::current_exe() + && let Some(profile_dir) = current.parent() + && let Some(target_dir) = profile_dir.parent() + && target_dir.file_name().is_some_and(|name| name == "target") + && let Some(workspace) = target_dir.parent() + { + if let Some(progress) = progress { + let _ = progress.send(ProvisionPhase::BuildingServer); + } + let profile = profile_dir + .file_name() + .and_then(|n| n.to_str()) + .unwrap_or("release"); + let args = hi_local_build_args(profile, backend); + let output = tokio::process::Command::new("cargo") + .args(&args) + .current_dir(workspace) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .output() + .await + .context("building hi-local from the development checkout")?; + if !output.status.success() { + let tail: String = String::from_utf8_lossy(&output.stderr) + .chars() + .rev() + .take(400) + .collect::() + .chars() + .rev() + .collect(); + bail!("building hi-local (--features {}) failed: {tail}", backend.cargo_feature()); + } + let built = profile_dir.join(format!("hi-local{}", std::env::consts::EXE_SUFFIX)); + if built.exists() { + return Ok(built); + } + } + let bin = find_hi_local(); + if bin.components().count() > 1 { + if bin.exists() { + return Ok(bin); + } + } else if binary_on_path(&bin.to_string_lossy()) { + return Ok(bin); + } + bail!( + "the hi-local serving binary isn't available (not beside hi, not on PATH); reinstall hi with its local-serving component" + ) +} + +#[cfg(test)] +mod provisioning_support_tests { + use super::*; + + #[test] + fn health_deadline_scales_with_model_size() { + let gib = 1024 * 1024 * 1024; + assert_eq!(health_deadline_for_model(0).as_secs(), 60); + assert_eq!(health_deadline_for_model(2 * gib).as_secs(), 90); + assert_eq!(health_deadline_for_model(19 * gib).as_secs(), 345); + assert_eq!( + health_deadline_for_model(200 * gib).as_secs(), + 600, + "capped at ten minutes" + ); + } + + #[test] + fn build_args_always_carry_the_backend_feature() { + let release = hi_local_build_args("release", LocalBackend::Mlx); + assert!(release.contains(&"--features".to_string())); + assert!(release.contains(&"mlx".to_string())); + assert!(release.contains(&"--release".to_string())); + let debug = hi_local_build_args("debug", LocalBackend::Cuda); + assert!(debug.contains(&"native-cuda".to_string())); + assert!(!debug.contains(&"--release".to_string())); + } + + #[test] + fn binary_on_path_finds_sh_and_rejects_nonsense() { + assert!(binary_on_path("sh")); + assert!(!binary_on_path("hi-definitely-not-a-real-binary-xyz")); + } +} diff --git a/crates/hi-agent/src/subagent.rs b/crates/hi-agent/src/subagent.rs index dc65912e..e3fb4693 100644 --- a/crates/hi-agent/src/subagent.rs +++ b/crates/hi-agent/src/subagent.rs @@ -9,6 +9,34 @@ use async_trait::async_trait; +/// A per-role model route override (team roles): which model — and optionally +/// which OpenAI-compatible endpoint — a subagent runs on, independent of the +/// driver's route. All-`None` means "inherit the driver". This is what lets a +/// big cloud driver plan and integrate while local models do the execution. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct SubagentRoute { + pub model: Option, + pub base_url: Option, + pub api_key: Option, +} + +impl SubagentRoute { + /// Whether this route overrides anything at all. + pub fn is_inherited(&self) -> bool { + self.model.is_none() && self.base_url.is_none() + } + + /// Short display label for role tables: `model @ endpoint` with inherited + /// parts elided. + pub fn label(&self, driver_model: &str) -> String { + let model = self.model.as_deref().unwrap_or(driver_model); + match self.base_url.as_deref() { + Some(url) => format!("{model} @ {url}"), + None => format!("{model} (driver route)"), + } + } +} + /// Outcome of one `delegate` write-subagent run. pub struct DelegateOutcome { /// Authoritative result of the isolated delegate run. A rolled-back, @@ -52,4 +80,19 @@ pub trait DelegateRunner: Send + Sync { } self.run(task, verify).await } + + /// Route-aware delegate execution (team roles). The default ignores the + /// route and preserves existing runner behavior; runners that spawn child + /// processes should apply `route` so delegate work can run on a different + /// model/endpoint than the driver. + async fn run_routed( + &self, + task: &str, + verify: Option<&str>, + route: &SubagentRoute, + cancellation: crate::TurnCancellation, + ) -> DelegateOutcome { + let _ = route; + self.run_cancellable(task, verify, cancellation).await + } } diff --git a/crates/hi-agent/src/tests/explore.rs b/crates/hi-agent/src/tests/explore.rs index 84038d9e..d99e9609 100644 --- a/crates/hi-agent/src/tests/explore.rs +++ b/crates/hi-agent/src/tests/explore.rs @@ -90,6 +90,34 @@ async fn explore_respects_turn_budget() { assert_eq!(agent.subagents.explore_turn_used, 0); } +#[tokio::test] +async fn run_turn_refills_the_explore_budget_each_turn() { + // Regression (user-reported on an old build): "explore budget exhausted + // (8 subagents this session)" — the budget was session-scoped, so a long + // session permanently lost the explore tool. The cap is a per-turn + // runaway guard: `run_turn` itself must refill it at turn start, not + // just `SubagentSessionState::begin_turn` in isolation. + let mut agent = agent( + vec![completion(vec![Content::Text("ok".into())], 1, 1)], + explore_config(), + ); + agent.subagents.explore_turn_used = crate::agent::MAX_EXPLORE_SUBAGENTS_PER_TURN; + agent.subagents.explore_subagents_used = crate::agent::MAX_EXPLORE_SUBAGENTS_PER_TURN; + let mut ui = NullUi; + + agent.run_turn("say ok", &mut ui).await.unwrap(); + + assert_eq!( + agent.subagents.explore_turn_used, 0, + "a new turn starts with a full explore budget" + ); + assert_eq!( + agent.subagents.explore_subagents_used, + crate::agent::MAX_EXPLORE_SUBAGENTS_PER_TURN, + "lifetime slot numbering is preserved across turns" + ); +} + #[tokio::test] async fn explore_runs_child_and_returns_answer() { // Parent and child share the same canned provider (Arc), popping in exactly @@ -345,3 +373,50 @@ async fn explore_batched_failed_offset_reads_are_bounded_before_chat_only_answer ui.statuses ); } + +#[test] +fn team_roles_table_and_route_setters_round_trip() { + let mut agent = agent(Vec::new(), explore_config()); + let roles: Vec<&str> = agent.team_roles().iter().map(|r| r.role).collect(); + assert_eq!( + roles, + vec!["driver", "explore", "delegate", "skeptic", "planner"] + ); + assert!( + agent.team_roles().iter().skip(1).all(|r| r.inherited), + "all worker roles inherit the driver by default" + ); + + agent.set_delegate_route( + Some("qwen3-coder".into()), + Some("http://127.0.0.1:18080/v1".into()), + None, + ); + agent.set_explore_route(Some("qwen3-4b".into()), None, None); + + let table = agent.team_roles(); + let delegate = table.iter().find(|r| r.role == "delegate").unwrap(); + assert_eq!(delegate.model, "qwen3-coder"); + assert_eq!(delegate.route, "http://127.0.0.1:18080/v1"); + assert!(!delegate.inherited); + let explore = table.iter().find(|r| r.role == "explore").unwrap(); + assert_eq!(explore.model, "qwen3-4b"); + assert!(!explore.inherited); + + // The delegate job route the executors will actually receive. + let route = agent.delegate_route(); + assert_eq!(route.model.as_deref(), Some("qwen3-coder")); + assert_eq!(route.base_url.as_deref(), Some("http://127.0.0.1:18080/v1")); + + // Blank/off clears back to inheritance. + agent.set_delegate_route(None, None, None); + agent.set_explore_route(Some(" ".into()), None, None); + assert!( + agent + .team_roles() + .iter() + .filter(|r| r.role == "delegate" || r.role == "explore") + .all(|r| r.inherited), + "cleared routes inherit the driver again" + ); +} diff --git a/crates/hi-agent/src/tests/local_skeptic.rs b/crates/hi-agent/src/tests/local_skeptic.rs index f72a3a48..d0b7ad7c 100644 --- a/crates/hi-agent/src/tests/local_skeptic.rs +++ b/crates/hi-agent/src/tests/local_skeptic.rs @@ -54,7 +54,12 @@ fn model_present_checks_config_json_for_mlx() { let spec = default_model(LocalBackend::Mlx); assert!(!model_present(&dir, &spec), "empty dir is not present"); std::fs::write(dir.join("config.json"), "{}").unwrap(); - assert!(model_present(&dir, &spec), "config.json marks MLX present"); + assert!( + !model_present(&dir, &spec), + "config.json alone is a partial download, not a model" + ); + std::fs::write(dir.join("model.safetensors"), b"w").unwrap(); + assert!(model_present(&dir, &spec), "config + weights mark MLX present"); let _ = std::fs::remove_dir_all(&dir); } diff --git a/crates/hi-cli/Cargo.toml b/crates/hi-cli/Cargo.toml index 8e5a7a76..3ea92f12 100644 --- a/crates/hi-cli/Cargo.toml +++ b/crates/hi-cli/Cargo.toml @@ -36,6 +36,7 @@ flate2.workspace = true ignore.workspace = true tar.workspace = true rusqlite.workspace = true +hi-sqlite-journal.workspace = true sha2.workspace = true uuid.workspace = true diff --git a/crates/hi-cli/src/agent_build.rs b/crates/hi-cli/src/agent_build.rs index fa130827..80459e4f 100644 --- a/crates/hi-cli/src/agent_build.rs +++ b/crates/hi-cli/src/agent_build.rs @@ -110,6 +110,14 @@ pub(crate) fn build_agent( skeptic_endpoint_key: std::env::var("HI_SKEPTIC_ENDPOINT_KEY") .ok() .filter(|s| !s.trim().is_empty()), + // Team-role routes for executors. Env vars seed them at startup + // (mirroring the skeptic knobs); `/team` adjusts them live. + delegate_model: env_route("HI_DELEGATE_MODEL"), + delegate_endpoint: env_route("HI_DELEGATE_ENDPOINT"), + delegate_endpoint_key: env_route("HI_DELEGATE_ENDPOINT_KEY"), + explore_model: env_route("HI_EXPLORE_MODEL"), + explore_endpoint: env_route("HI_EXPLORE_ENDPOINT"), + explore_endpoint_key: env_route("HI_EXPLORE_ENDPOINT_KEY"), // `/goal` is a core CLI contract, not a provider-specific feature. // Delegate children receive bounded tasks and therefore keep it off. long_horizon: goal_drive::long_horizon_enabled(cli.subagent), @@ -144,3 +152,8 @@ pub(crate) fn build_agent( resume_summary, }) } + +/// Read an optional team-role route env var (empty = unset). +fn env_route(name: &str) -> Option { + std::env::var(name).ok().filter(|s| !s.trim().is_empty()) +} diff --git a/crates/hi-cli/src/commands.rs b/crates/hi-cli/src/commands.rs index 45d35c22..2e71a752 100644 --- a/crates/hi-cli/src/commands.rs +++ b/crates/hi-cli/src/commands.rs @@ -191,6 +191,202 @@ pub(crate) fn handle_command(agent: &mut Agent, command: hi_agent::Command) -> b Err(err) => eprintln!("\x1b[33mclear failed: {err}\x1b[0m"), } } + Command::Team(arg) => { + let parts: Vec<&str> = arg.split_whitespace().collect(); + match parts.as_slice() { + [] => { + for row in agent.team_roles() { + let suffix = if row.inherited { " (driver)" } else { "" }; + println!( + "\x1b[2m {:<9} {} @ {}{}\x1b[0m", + row.role, row.model, row.route, suffix + ); + } + println!( + "\x1b[2m /team · /team planner \x1b[0m" + ); + let supported = hi_agent::local_skeptic::SUPPORTED_LOCAL_MODELS + .iter() + .map(|entry| entry.name) + .collect::>() + .join(", "); + println!( + "\x1b[2m local models (auto-download + serve): local (auto-size), {supported}\x1b[0m" + ); + } + ["driver", ..] => { + println!( + "\x1b[2mthe driver is the session model — switch it with /model or /provider\x1b[0m" + ); + } + ["skeptic", ..] => { + println!( + "\x1b[2mskeptic routing has dedicated commands: /config skeptic-local on|off or HI_SKEPTIC_ENDPOINT\x1b[0m" + ); + } + ["planner", "off"] => { + agent.set_planner_model(None); + println!("\x1b[2mplanner → driver model\x1b[0m"); + } + ["planner", model] => { + agent.set_planner_model(Some((*model).to_string())); + println!("\x1b[2mplanner → {model}\x1b[0m"); + } + [role @ ("explore" | "delegate"), "off"] => { + if *role == "delegate" { + agent.set_delegate_route(None, None, None); + } else { + agent.set_explore_route(None, None, None); + } + println!("\x1b[2m{role} → driver route (applies to new {role} runs)\x1b[0m"); + } + [role @ ("explore" | "delegate"), model, rest @ ..] => { + let explicit_endpoint = rest + .first() + .filter(|value| value.starts_with("http")) + .map(|s| (*s).to_string()); + if let Some(endpoint) = explicit_endpoint { + let key = rest.get(1).map(|s| (*s).to_string()); + if *role == "delegate" { + agent.set_delegate_route( + Some((*model).to_string()), + Some(endpoint.clone()), + key, + ); + } else { + agent.set_explore_route( + Some((*model).to_string()), + Some(endpoint.clone()), + key, + ); + } + println!("\x1b[2m{role} → {model} @ {endpoint}\x1b[0m"); + } else if let Some(resolved) = + hi_agent::local_skeptic::resolve_team_local_model( + model, + hi_agent::local_skeptic::system_ram_gb(), + hi_agent::local_skeptic::detect_backend(), + ) + { + // Supported local model: reuse a running server or + // provision inline (plain mode has no background UI). + let reuse = resolved + .mlx + .and_then(|quant| agent.running_local_model_server(quant.model_id)) + .or_else(|| { + resolved.entry.cuda.and_then(|cuda| { + agent.running_local_model_server(cuda.model_id) + }) + }); + let provisioned = match reuse { + Some((endpoint, model_id)) => { + Ok((endpoint, model_id, String::new())) + } + None => { + println!( + "\x1b[2msetting up {} locally (first run may download weights — this can take a while)…\x1b[0m", + resolved.display() + ); + // Plain mode has no background UI; block the + // (synchronous) command loop on the runtime + // and print phase transitions as they land. + let (phase_tx, mut phase_rx) = tokio::sync::watch::channel( + hi_agent::local_skeptic::ProvisionPhase::Resolving, + ); + let printer = tokio::spawn(async move { + while phase_rx.changed().await.is_ok() { + let phase = phase_rx.borrow().clone(); + println!("\x1b[2m {phase:?}\x1b[0m"); + } + }); + let result = tokio::task::block_in_place(|| { + tokio::runtime::Handle::current().block_on( + hi_agent::local_skeptic::provision_team_local_model( + resolved, phase_tx, + ), + ) + }); + printer.abort(); + result + } + }; + match provisioned { + Ok((endpoint, model_id, process_id)) => { + if !process_id.is_empty() { + agent.register_team_local_server( + endpoint.clone(), + model_id.clone(), + process_id, + ); + } + if *role == "delegate" { + agent.set_delegate_route( + Some(model_id.clone()), + Some(endpoint), + None, + ); + } else { + agent.set_explore_route( + Some(model_id.clone()), + Some(endpoint), + None, + ); + } + println!( + "\x1b[2m✓ {role} → {model_id} @ local (applies to new {role} runs)\x1b[0m" + ); + } + Err(error) => { + let reason: String = + format!("{error:#}").chars().take(140).collect(); + println!( + "\x1b[2mcouldn't set up {} locally ({reason}); {role} stays on the driver\x1b[0m", + resolved.display() + ); + } + } + } else { + if *role == "delegate" { + agent.set_delegate_route(Some((*model).to_string()), None, None); + } else { + agent.set_explore_route(Some((*model).to_string()), None, None); + } + println!( + "\x1b[2m{role} → {model} (driver route; applies to new {role} runs)\x1b[0m" + ); + } + } + [role @ ("explore" | "delegate")] => { + // Plain mode has no dropdown; print the catalog with the + // same per-machine sizing the TUI picker shows. + let ram = hi_agent::local_skeptic::system_ram_gb(); + println!( + "\x1b[2mpick a local model for {role}: /team {role} (or `auto`, `name@quant`, a cloud model id, or an explicit endpoint URL)\x1b[0m" + ); + for entry in hi_agent::local_skeptic::SUPPORTED_LOCAL_MODELS { + let fit = match entry.pick_mlx(ram) { + Some(quant) if entry.mlx.len() > 1 => format!( + "{} fits (needs {}GB RAM) · quants {}", + quant.quant, + quant.min_ram_gb, + entry.quant_summary() + ), + Some(quant) => format!("needs {}GB RAM · fits", quant.min_ram_gb), + None => format!( + "needs {}GB+ RAM · too big for this machine", + entry.min_ram_gb(None) + ), + }; + println!("\x1b[2m {:<14} {} · {fit}\x1b[0m", entry.name, entry.label); + } + } + [role, ..] => { + println!( + "\x1b[2munknown role '{role}' — roles: driver, explore, delegate, skeptic, planner\x1b[0m" + ); + } + } + } Command::Config(arg) => { use hi_agent::command::{ConfigArg, parse_config_arg}; match parse_config_arg(&arg) { diff --git a/crates/hi-cli/src/delegate.rs b/crates/hi-cli/src/delegate.rs index f740189a..6bdd4963 100644 --- a/crates/hi-cli/src/delegate.rs +++ b/crates/hi-cli/src/delegate.rs @@ -163,7 +163,57 @@ impl DelegateRunner for CliDelegateRunner { self.run(task, verify).await } + async fn run_routed( + &self, + task: &str, + verify: Option<&str>, + route: &hi_agent::SubagentRoute, + cancellation: hi_agent::TurnCancellation, + ) -> DelegateOutcome { + if cancellation.is_cancelled() { + return outcome(ToolStatus::Cancelled, "delegate cancelled before setup"); + } + self.run_with_route(task, verify, route).await + } + async fn run(&self, task: &str, verify: Option<&str>) -> DelegateOutcome { + self.run_with_route(task, verify, &hi_agent::SubagentRoute::default()) + .await + } +} + +impl CliDelegateRunner { + /// Resolve the child's provider route: team-role overrides win over the + /// runner's defaults. An endpoint override implies the generic + /// OpenAI-compatible provider — local servers (MLX, Ollama, llama.cpp) + /// all speak it. + pub(crate) fn effective_route( + &self, + route: &hi_agent::SubagentRoute, + ) -> (String, String, String, String) { + let model = route.model.clone().unwrap_or_else(|| self.model.clone()); + match route.base_url.as_deref() { + Some(url) => ( + "openai".to_string(), + model, + url.to_string(), + route.api_key.clone().unwrap_or_default(), + ), + None => ( + self.provider.clone(), + model, + self.base_url.clone(), + self.api_key.clone(), + ), + } + } + + async fn run_with_route( + &self, + task: &str, + verify: Option<&str>, + route: &hi_agent::SubagentRoute, + ) -> DelegateOutcome { let Some(verify_cmd) = verify .map(str::to_string) .or_else(|| self.default_verify.clone()) @@ -220,10 +270,7 @@ impl DelegateRunner for CliDelegateRunner { let idx = self.counter.fetch_add(1, Ordering::Relaxed); let exe = self.exe.clone(); - let provider = self.provider.clone(); - let model = self.model.clone(); - let base_url = self.base_url.clone(); - let api_key = self.api_key.clone(); + let (provider, model, base_url, api_key) = self.effective_route(route); let max_steps = self.max_steps; let max_verify = self.max_verify; let task = task.to_string(); diff --git a/crates/hi-cli/src/delegate_tests.rs b/crates/hi-cli/src/delegate_tests.rs index ad048e9d..e3df0ee3 100644 --- a/crates/hi-cli/src/delegate_tests.rs +++ b/crates/hi-cli/src/delegate_tests.rs @@ -300,3 +300,66 @@ fn git_stdout(root: &Path, args: &[&str]) -> String { assert!(output.status.success()); String::from_utf8(output.stdout).unwrap().trim().to_string() } + +#[test] +fn delegate_route_overrides_switch_the_child_to_the_openai_compat_route() { + let id = TEST_ID.fetch_add(1, Ordering::Relaxed); + let base = std::env::temp_dir().join(format!("hi-delegate-route-{}-{id}", std::process::id())); + let workspace = base.join("ws"); + let state = base.join("state"); + std::fs::create_dir_all(&workspace).unwrap(); + std::fs::create_dir_all(&state).unwrap(); + let runner = crate::delegate::CliDelegateRunner::new( + PathBuf::from("hi"), + "pipenetwork".into(), + "pipe/glm-5.2".into(), + "https://api.pipenetwork.ai/v1".into(), + "cloud-key".into(), + Some("true".into()), + None, + 1, + workspace, + state, + ) + .unwrap(); + + // No override → the driver's route, untouched. + let inherited = runner.effective_route(&hi_agent::SubagentRoute::default()); + assert_eq!( + inherited, + ( + "pipenetwork".into(), + "pipe/glm-5.2".into(), + "https://api.pipenetwork.ai/v1".into(), + "cloud-key".into() + ) + ); + + // Model-only override stays on the driver's provider. + let model_only = runner.effective_route(&hi_agent::SubagentRoute { + model: Some("pipe/glm-4-flash".into()), + base_url: None, + api_key: None, + }); + assert_eq!(model_only.0, "pipenetwork"); + assert_eq!(model_only.1, "pipe/glm-4-flash"); + + // Endpoint override moves the child to the generic OpenAI-compatible + // provider (local MLX/Ollama/llama.cpp servers all speak it), and the + // cloud key must NOT leak to the local endpoint. + let local = runner.effective_route(&hi_agent::SubagentRoute { + model: Some("qwen3-coder".into()), + base_url: Some("http://127.0.0.1:18080/v1".into()), + api_key: None, + }); + assert_eq!( + local, + ( + "openai".into(), + "qwen3-coder".into(), + "http://127.0.0.1:18080/v1".into(), + String::new() + ) + ); + let _ = std::fs::remove_dir_all(base); +} diff --git a/crates/hi-cli/src/main.rs b/crates/hi-cli/src/main.rs index ffb64dca..65cdf29e 100644 --- a/crates/hi-cli/src/main.rs +++ b/crates/hi-cli/src/main.rs @@ -27,6 +27,7 @@ mod rsi_bootstrap; mod rsi_observation; mod rsi_policy; mod rsi_remote; +mod team_bench; mod tuning_report; // Wired by the managed RSI entry once descriptor-driven workflow launch lands; // composition and contracts are complete and tested. @@ -138,6 +139,9 @@ async fn run() -> Result<()> { if raw_args.get(1).map(String::as_str) == Some("bench") { return bench::run_bench_cli(&raw_args[2..]).await; } + if raw_args.get(1).map(String::as_str) == Some("team-bench") { + return team_bench::run_team_bench_cli(&raw_args[2..]).await; + } if raw_args.get(1).map(String::as_str) == Some("metrics") { let (_, state_root) = resolve_runtime_roots()?; orchestration_metrics::print_dashboard(&state_root); diff --git a/crates/hi-cli/src/sync.rs b/crates/hi-cli/src/sync.rs index 286f2290..5a5306b8 100644 --- a/crates/hi-cli/src/sync.rs +++ b/crates/hi-cli/src/sync.rs @@ -955,13 +955,12 @@ pub struct SyncSession { impl SyncSession { pub fn new(local: crate::session::JsonlSession, remote: RemoteSessionSink) -> Self { - remote - .reconcile_jsonl(local.path()) - .expect("reconciling durable portal outbox"); - Self { + let session = Self { local, remote: std::sync::Arc::new(remote), - } + }; + session.reconcile_best_effort(); + session } /// Get a handle to the remote sink for flushing / ending the session. @@ -969,18 +968,31 @@ impl SyncSession { pub fn remote_handle(&self) -> std::sync::Arc { self.remote.clone() } + + /// Mirror new JSONL lines into the durable outbox, swallowing failures. + /// The mirror is offset-tracked and idempotent, so an error (a peer + /// holding the SQLite write lock past its busy_timeout, a broken store) + /// only defers those lines to the next reconcile. The local JSONL is the + /// session's source of truth and has already been written — failing the + /// caller's turn over the mirror would surface an error the user can't + /// act on. + fn reconcile_best_effort(&self) { + let _ = self.remote.reconcile_jsonl(self.local.path()); + } } impl SessionSink for SyncSession { fn record(&mut self, messages: &[Message], usage: Usage) -> Result<()> { self.local.record(messages, usage)?; self.remote.observe_messages(messages); - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn record_compaction(&mut self, messages: &[Message]) -> Result<()> { self.local.record_compaction(messages)?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn record_state_replacement( @@ -992,37 +1004,44 @@ impl SessionSink for SyncSession { ) -> Result<()> { self.local .record_state_replacement(messages, goal, decisions, plan)?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn record_checkpoints(&mut self, refs: &[String]) -> Result<()> { self.local.record_checkpoints(refs)?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn record_goal(&mut self, goal: &hi_agent::Goal) -> Result<()> { self.local.record_goal(goal)?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn clear_goal(&mut self) -> Result<()> { self.local.clear_goal()?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn record_plan(&mut self, plan: &[hi_agent::PlanStep]) -> Result<()> { self.local.record_plan(plan)?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn clear_plan(&mut self) -> Result<()> { self.local.clear_plan()?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } fn record_decisions(&mut self, decisions: &hi_agent::DecisionLog) -> Result<()> { self.local.record_decisions(decisions)?; - self.remote.reconcile_jsonl(self.local.path()) + self.reconcile_best_effort(); + Ok(()) } } @@ -2750,6 +2769,47 @@ mod tests { ); } + /// A locked or broken outbox store must never fail the turn: the local + /// JSONL is the source of truth, and the outbox mirror is offset-tracked + /// and idempotent, so a failed reconcile defers to the next record. The + /// old behavior surfaced "database is locked" as a failed turn. + #[test] + fn turn_record_survives_broken_outbox_store() { + let dir = std::env::temp_dir().join(format!( + "hi-sync-broken-store-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&dir).unwrap(); + let store_path = dir.join("outbox.sqlite3"); + let store = Arc::new(crate::sync_store::SyncStore::open_at(store_path.clone()).unwrap()); + let sink = RemoteSessionSink::with_store( + unreachable_config(), + "broken-store".to_string(), + None, + remote_session_http_client(), + store, + ); + // Break the store out from under the sink: every reconcile now errors. + rusqlite::Connection::open(&store_path) + .unwrap() + .execute_batch("DROP TABLE session_sync; DROP TABLE record_outbox;") + .unwrap(); + + let jsonl_path = dir.join("session.jsonl"); + let mut sync = SyncSession::new(crate::session::JsonlSession::new(jsonl_path.clone()), sink); + sync.record(&[Message::user("hello")], Usage::default()) + .expect("turn recording must not fail on outbox errors"); + assert!( + std::fs::metadata(&jsonl_path).unwrap().len() > 0, + "local JSONL still records the turn" + ); + let _ = std::fs::remove_dir_all(&dir); + } + /// The `--session-file` collision bug: session ids derive from the file /// stem, so a second session at a same-named path inherits the first /// session's byte offset into a different file. Reconcile must reset the diff --git a/crates/hi-cli/src/sync_store.rs b/crates/hi-cli/src/sync_store.rs index 5d61f1ca..e3c43df7 100644 --- a/crates/hi-cli/src/sync_store.rs +++ b/crates/hi-cli/src/sync_store.rs @@ -6,7 +6,7 @@ use std::sync::OnceLock; use std::time::{SystemTime, UNIX_EPOCH}; use anyhow::{Context, Result}; -use rusqlite::{Connection, OptionalExtension, params}; +use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params}; use serde::Serialize; use sha2::{Digest, Sha256}; @@ -87,6 +87,20 @@ fn live_event_drop_delta(before: i64, after: i64) -> i64 { before.saturating_sub(after).max(0) } +/// True when any cause in the chain is SQLite's "database is locked" +/// (SQLITE_BUSY / SQLITE_LOCKED) — transient peer contention, safe to retry. +fn is_busy_error(error: &anyhow::Error) -> bool { + error.chain().any(|cause| { + matches!( + cause.downcast_ref::(), + Some(rusqlite::Error::SqliteFailure(failure, _)) if matches!( + failure.code, + rusqlite::ErrorCode::DatabaseBusy | rusqlite::ErrorCode::DatabaseLocked + ) + ) + }) +} + fn hex_sha256(parts: &[&[u8]]) -> String { let mut hash = Sha256::new(); for part in parts { @@ -104,16 +118,36 @@ impl SyncStore { } pub fn open_at(path: PathBuf) -> Result { - let connection = Connection::open(&path) + // Peer hi processes (TUI, daemon, background jobs) share this file + // and hold short write locks. busy_timeout in try_open_at absorbs + // ordinary contention; this bounded retry absorbs a peer that holds + // the lock across an entire timeout window. + let mut attempts = 0u64; + loop { + match Self::try_open_at(&path) { + Err(error) if attempts < 3 && is_busy_error(&error) => { + attempts += 1; + std::thread::sleep(std::time::Duration::from_millis(100 * attempts)); + } + result => return result, + } + } + } + + fn try_open_at(path: &std::path::Path) -> Result { + // hi-sqlite-journal owns the lock-safe open: busy_timeout before any + // lock-taking statement, a poll loop for the journal-mode switch + // (which SQLite refuses to apply busy_timeout to), and a rollback + // journal instead of WAL on network filesystems. + let connection = hi_sqlite_journal::JournalMode::for_db_path(path) + .open(path) .with_context(|| format!("opening portal sync database {}", path.display()))?; #[cfg(unix)] { use std::os::unix::fs::PermissionsExt; - let _ = std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600)); + let _ = std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600)); } - connection.pragma_update(None, "journal_mode", "WAL")?; connection.pragma_update(None, "synchronous", "FULL")?; - connection.busy_timeout(std::time::Duration::from_secs(5))?; connection.execute_batch( "BEGIN IMMEDIATE; CREATE TABLE IF NOT EXISTS sync_settings ( @@ -296,7 +330,10 @@ impl SyncStore { return Ok(()); } let mut connection = self.connection.lock().unwrap(); - let transaction = connection.transaction()?; + // Immediate (here and below): take the write lock at BEGIN, where + // busy_timeout applies. A deferred BEGIN that reads before writing + // can hit SQLITE_BUSY on lock upgrade with no timeout at all. + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; transaction.execute( "INSERT INTO record_outbox(session_id,client_record_id,record_type,payload_json,created_at_unix) VALUES(?1,'pending',?2,?3,?4)", @@ -418,7 +455,7 @@ impl SyncStore { pub fn acknowledge_records(&self, session_id: &str, ids: &[i64], cursor: u64) -> Result<()> { let mut connection = self.connection.lock().unwrap(); - let transaction = connection.transaction()?; + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; for id in ids { transaction.execute("DELETE FROM record_outbox WHERE id=?1", [id])?; } @@ -465,7 +502,7 @@ impl SyncStore { return Ok(()); } let mut connection = self.connection.lock().unwrap(); - let transaction = connection.transaction()?; + let transaction = connection.transaction_with_behavior(TransactionBehavior::Immediate)?; transaction.execute( "INSERT INTO live_event_queue(session_id,event_json,created_at_unix) VALUES(?1,?2,?3)", params![session_id, event_json, now()], @@ -680,6 +717,32 @@ mod tests { let _ = std::fs::remove_file(path); } + /// A peer holding the write lock while another process opens the store + /// must delay the open, not fail it — the "database is locked" turn + /// failures came from lock-taking setup running before any busy_timeout + /// was configured. + #[test] + fn open_waits_for_peer_write_lock_instead_of_failing() { + let path = temp_store_path("open-contended"); + // Seed a pre-WAL database so open_at's journal_mode switch needs an + // exclusive lock, then hold the write lock from a peer connection. + let peer = Connection::open(&path).unwrap(); + peer.execute_batch("CREATE TABLE seed(x); BEGIN IMMEDIATE; INSERT INTO seed VALUES(1);") + .unwrap(); + let opener = { + let path = path.clone(); + std::thread::spawn(move || SyncStore::open_at(path).map(|_| ())) + }; + std::thread::sleep(std::time::Duration::from_millis(300)); + peer.execute_batch("COMMIT;").unwrap(); + opener + .join() + .unwrap() + .expect("open should wait out the peer's short write lock"); + drop(peer); + let _ = std::fs::remove_file(path); + } + #[test] fn track_jsonl_identifies_files_canonically() { let store_path = temp_store_path("track-canonical"); diff --git a/crates/hi-cli/src/team_bench.rs b/crates/hi-cli/src/team_bench.rs new file mode 100644 index 00000000..0ed8584b --- /dev/null +++ b/crates/hi-cli/src/team_bench.rs @@ -0,0 +1,681 @@ +//! `hi team-bench` — a small, honest coding benchmark for the supported +//! local models, answering "which of these is better at what?" with +//! machine-checked results instead of vibes. +//! +//! Each model is provisioned through the exact `/team` path (download → +//! `hi-local` serve → health), then given four verifiable coding tasks: +//! +//! - `codegen` — write a function from a spec; compiled and run against asserts +//! - `bugfix` — repair a buggy function; compiled and run against asserts +//! - `edit` — a precise mechanical edit; textual invariants + compiles +//! - `json` — tool-call fidelity; strict JSON schema + exact values +//! +//! Scores are pass/fail per task plus measured generation speed. Servers are +//! started one at a time and stopped after each model so a 60GB ladder never +//! stacks in RAM. Results print as a table and are saved as JSON under +//! `~/.hi/bench/`. + +use anyhow::{Context, Result, bail}; +use hi_agent::local_skeptic::{ + LocalBackend, ProvisionPhase, ResolvedLocalModel, SUPPORTED_LOCAL_MODELS, + detect_backend_offload, provision_team_local_model, resolve_team_local_model, system_ram_gb, + team_model_spec, +}; +use hi_ai::{ChatRequest, Content, Message, OpenAiProvider, Provider, RequestProfile}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// One benchmark task: a prompt and a machine validator for the reply. +struct BenchTask { + name: &'static str, + prompt: String, + check: fn(&str) -> Result<()>, +} + +#[derive(serde::Serialize, Clone)] +struct TaskResult { + task: &'static str, + pass: bool, + /// Why the task failed (compile error, wrong value, bad JSON…); None on pass. + error: Option, + latency_secs: f64, + output_tokens: u64, + tokens_per_sec: f64, +} + +#[derive(serde::Serialize)] +struct ModelReport { + model: String, + model_id: String, + setup_secs: f64, + results: Vec, +} + +impl ModelReport { + fn passed(&self) -> usize { + self.results.iter().filter(|r| r.pass).count() + } + fn avg_tokens_per_sec(&self) -> f64 { + let generating: Vec<&TaskResult> = + self.results.iter().filter(|r| r.output_tokens > 0).collect(); + if generating.is_empty() { + return 0.0; + } + generating.iter().map(|r| r.tokens_per_sec).sum::() / generating.len() as f64 + } +} + +pub(crate) async fn run_team_bench_cli(args: &[String]) -> Result<()> { + if args.first().map(String::as_str) == Some("--help") { + println!( + "usage: hi team-bench [--list] [model…]\n\n\ + Benchmark supported local models on verifiable coding tasks.\n\ + With no arguments, every catalog model already downloaded for this\n\ + machine's backend (and that fits it) is benchmarked. Naming models\n\ + (e.g. `laguna-s@3bit coder-32b`) benchmarks exactly those,\n\ + downloading weights if needed. --list shows the selection and exits." + ); + return Ok(()); + } + let list_only = args.first().map(String::as_str) == Some("--list"); + let args = if list_only { &args[1..] } else { args }; + let ram = system_ram_gb(); + let Some(backend) = detect_backend_offload().await else { + bail!("no local-inference backend detected (needs Apple Silicon MLX or an NVIDIA CUDA runtime)"); + }; + let selections = if args.is_empty() { + downloaded_selections(ram, backend) + } else { + let mut picked = Vec::new(); + for name in args { + let Some(resolved) = resolve_team_local_model(name, ram, Some(backend)) else { + bail!("'{name}' is not a supported local model — see /team for the catalog"); + }; + picked.push(resolved); + } + picked + }; + if list_only { + for resolved in &selections { + println!("{}", resolved.display()); + } + return Ok(()); + } + if selections.is_empty() { + println!( + "no local models downloaded yet — name one to fetch and benchmark it, e.g.:\n hi team-bench local" + ); + return Ok(()); + } + if !rustc_available() { + bail!("team-bench compiles model output with `rustc`, which isn't on PATH"); + } + + let tasks = bench_tasks(); + println!( + "benchmarking {} model(s) on {} coding task(s); servers run one at a time\n", + selections.len(), + tasks.len() + ); + let mut reports = Vec::new(); + for resolved in selections { + match bench_model(resolved, &tasks).await { + Ok(report) => reports.push(report), + Err(error) => { + // One broken model must not sink the comparison run. + println!(" ✗ {}: {error:#}\n", resolved.display()); + } + } + } + if reports.is_empty() { + bail!("no model completed the benchmark"); + } + print_summary(&reports, &tasks); + if let Some(path) = save_report(&reports) { + println!("\nfull report: {}", path.display()); + } + Ok(()) +} + +/// Every catalog selection whose weights are already on disk for `backend` +/// AND fit this machine — for MLX ladders, the highest-quality downloaded +/// quant of each family. A downloaded-but-oversized model (someone's 390GB +/// GLM copy on a 64GB Mac) must be named explicitly to be benched. +fn downloaded_selections(ram: u64, backend: LocalBackend) -> Vec { + let mut picked = Vec::new(); + for entry in SUPPORTED_LOCAL_MODELS { + let candidates: Vec = match backend { + LocalBackend::Mlx => entry + .mlx + .iter() + .filter(|quant| ram >= quant.min_ram_gb) + .map(|quant| ResolvedLocalModel { entry, mlx: Some(quant) }) + .collect(), + LocalBackend::Cuda => entry + .cuda + .filter(|cuda| ram >= cuda.min_ram_gb) + .map(|_| ResolvedLocalModel { entry, mlx: entry.pick_mlx(ram) }) + .into_iter() + .collect(), + }; + if let Some(downloaded) = candidates.into_iter().find(|resolved| { + team_model_spec(*resolved, backend).is_ok_and(|spec| { + let dir = hi_tools::skeptic_model_dir(&spec.repo); + hi_agent::local_skeptic::model_present(&dir, &spec) + }) + }) { + picked.push(downloaded); + } + } + picked +} + +/// Provision one model through the real `/team` path, run every task against +/// its server, and stop the server before returning. +async fn bench_model(resolved: ResolvedLocalModel, tasks: &[BenchTask]) -> Result { + let display = resolved.display(); + println!("— {display} —"); + let (phase_tx, mut phase_rx) = tokio::sync::watch::channel(ProvisionPhase::Resolving); + let narrator = tokio::spawn(async move { + while phase_rx.changed().await.is_ok() { + let line = match &*phase_rx.borrow() { + ProvisionPhase::Resolving => "resolving".to_string(), + ProvisionPhase::Downloading => "downloading weights (quiet)…".to_string(), + ProvisionPhase::BuildingServer => "building hi-local…".to_string(), + ProvisionPhase::LoadingModel { deadline_secs, .. } => { + format!("server started — loading weights (up to {deadline_secs}s)…") + } + }; + println!(" {line}"); + } + }); + let setup_started = Instant::now(); + let provisioned = provision_team_local_model(resolved, phase_tx).await; + narrator.abort(); + let (endpoint, model_id, process_id) = + provisioned.with_context(|| format!("setting up {display}"))?; + let setup_secs = setup_started.elapsed().as_secs_f64(); + println!(" ready in {setup_secs:.0}s at {endpoint}"); + + let provider = OpenAiProvider::new(endpoint, "local".to_string()); + let mut results = Vec::new(); + for task in tasks { + let result = run_task(&provider, &model_id, task).await; + let mark = if result.pass { "PASS" } else { "FAIL" }; + let detail = result + .error + .as_deref() + .map(|e| format!(" — {}", e.lines().next().unwrap_or_default())) + .unwrap_or_default(); + println!( + " {:<8} {mark} {:>5.1}s {:>6.1} tok/s{detail}", + task.name, result.latency_secs, result.tokens_per_sec + ); + results.push(result); + } + hi_tools::stop_local_server(&process_id); + println!(); + Ok(ModelReport { model: display, model_id, setup_secs, results }) +} + +/// One prompt → validate round. Task failures (bad output) are results, not +/// errors; only transport-level problems surface as FAIL with the cause. +async fn run_task(provider: &OpenAiProvider, model_id: &str, task: &BenchTask) -> TaskResult { + let request = ChatRequest { + model: model_id.to_string(), + request_id: None, + retry_attempt: 0, + user_turn: false, + canonical_objective: None, + messages: Arc::new(vec![Message::user(task.prompt.clone())]), + tools: Arc::new([]), + // Room for models that reason inline before the code; hi-local's + // default per-request ceiling is 8192. + max_tokens: 3500, + temperature: Some(0.0), + top_p: None, + frequency_penalty: None, + thinking_budget: None, + reasoning_effort: None, + profile: RequestProfile::default(), + }; + let started = Instant::now(); + let completion = tokio::time::timeout( + Duration::from_secs(600), + provider.stream(request, &mut |_event| {}), + ) + .await; + let latency_secs = started.elapsed().as_secs_f64(); + let (text, output_tokens) = match completion { + Ok(Ok(completion)) => { + let text: String = completion + .content + .iter() + .filter_map(|content| match content { + Content::Text(text) => Some(text.as_str()), + _ => None, + }) + .collect(); + (text, completion.usage.output_tokens) + } + Ok(Err(error)) => { + return TaskResult { + task: task.name, + pass: false, + error: Some(format!("request failed: {error:#}")), + latency_secs, + output_tokens: 0, + tokens_per_sec: 0.0, + }; + } + Err(_) => { + return TaskResult { + task: task.name, + pass: false, + error: Some("timed out after 600s".to_string()), + latency_secs, + output_tokens: 0, + tokens_per_sec: 0.0, + }; + } + }; + let tokens_per_sec = if latency_secs > 0.0 { + output_tokens as f64 / latency_secs + } else { + 0.0 + }; + match (task.check)(&text) { + Ok(()) => TaskResult { + task: task.name, + pass: true, + error: None, + latency_secs, + output_tokens, + tokens_per_sec, + }, + Err(error) => TaskResult { + task: task.name, + pass: false, + error: Some(format!("{error:#}")), + latency_secs, + output_tokens, + tokens_per_sec, + }, + } +} + +// ---- Tasks ----------------------------------------------------------------- + +const BUGGY_WINDOW_SUM: &str = r#"/// Largest sum over any k consecutive elements; 0 when the slice has +/// fewer than k elements or k is 0. +pub fn max_window_sum(v: &[i64], k: usize) -> i64 { + if k == 0 || v.len() < k { + return 0; + } + let mut best = i64::MIN; + for start in 0..v.len() - k { + let sum: i64 = v[start..start + k].iter().sum(); + if sum > best { + best = sum; + } + } + best +}"#; + +const EDIT_SOURCE: &str = r#"/// Fetch a user record by id. +pub fn fetch_user(id: u64) -> Option { + if id == 0 { + return None; + } + Some(format!("user-{id}")) +} + +pub fn describe(id: u64) -> String { + match fetch_user(id) { + Some(user) => user, + None => "anonymous".to_string(), + } +}"#; + +fn bench_tasks() -> Vec { + vec![ + BenchTask { + name: "codegen", + prompt: "Write a Rust function `pub fn run_length_encode(s: &str) -> Vec<(char, u32)>` \ + that collapses consecutive repeated characters into (character, count) pairs \ + in order of appearance. For example \"aab\" becomes [('a', 2), ('b', 1)]. \ + Return only the function code — no main, no tests, no explanation." + .to_string(), + check: check_codegen, + }, + BenchTask { + name: "bugfix", + prompt: format!( + "This Rust function has a bug: it can miss a window and can return i64::MIN. \ + For example max_window_sum(&[1, 2, 3, 4], 2) must be 7. Fix it and return only \ + the corrected complete function — no explanation.\n\n{BUGGY_WINDOW_SUM}" + ), + check: check_bugfix, + }, + BenchTask { + name: "edit", + prompt: format!( + "In the Rust file below, rename the function `fetch_user` to `load_user` \ + (including every call site) and add a `#[must_use]` attribute on the line \ + directly above `pub fn load_user`. Change nothing else. Return the complete \ + updated file only — no explanation.\n\n{EDIT_SOURCE}" + ), + check: check_edit, + }, + BenchTask { + name: "json", + prompt: "Respond with ONLY a JSON object (no prose, no code fences) with exactly \ + these fields: \"cmd\" (string), \"args\" (array of strings), \ + \"timeout_secs\" (number). It must describe running the shell command \ + `cargo nextest run` with a 300 second timeout, where cmd is the executable \ + and args are its arguments." + .to_string(), + check: check_json, + }, + ] +} + +fn check_codegen(reply: &str) -> Result<()> { + let code = extract_code(reply); + if !code.contains("fn run_length_encode") { + bail!("reply has no run_length_encode function"); + } + let harness = r#" +fn main() { + assert_eq!(run_length_encode("aaabccd"), vec![('a', 3u32), ('b', 1), ('c', 2), ('d', 1)]); + assert_eq!(run_length_encode(""), Vec::<(char, u32)>::new()); + assert_eq!(run_length_encode("xx"), vec![('x', 2u32)]); + println!("ok"); +} +"#; + compile_and_run("codegen", &format!("{code}\n{harness}")) +} + +fn check_bugfix(reply: &str) -> Result<()> { + let code = extract_code(reply); + if !code.contains("fn max_window_sum") { + bail!("reply has no max_window_sum function"); + } + let harness = r#" +fn main() { + assert_eq!(max_window_sum(&[1, 2, 3, 4], 2), 7); + assert_eq!(max_window_sum(&[5], 1), 5); + assert_eq!(max_window_sum(&[2, -1, 2, 3, -9], 3), 4); + assert_eq!(max_window_sum(&[1, 2], 5), 0); + assert_eq!(max_window_sum(&[], 0), 0); + println!("ok"); +} +"#; + compile_and_run("bugfix", &format!("{code}\n{harness}")) +} + +fn check_edit(reply: &str) -> Result<()> { + let code = extract_code(reply); + if !code.contains("#[must_use]") { + bail!("missing #[must_use] attribute"); + } + if !code.contains("pub fn load_user") { + bail!("fetch_user was not renamed to load_user"); + } + if code.contains("fetch_user") { + bail!("a fetch_user reference survived the rename"); + } + if !code.contains("pub fn describe") { + bail!("unrelated code was dropped from the file"); + } + compile_lib("edit", &code) +} + +fn check_json(reply: &str) -> Result<()> { + let value = extract_json(reply)?; + let cmd = value.get("cmd").and_then(|v| v.as_str()).unwrap_or_default(); + if cmd != "cargo" { + bail!("cmd is {cmd:?}, expected \"cargo\""); + } + let args: Vec<&str> = value + .get("args") + .and_then(|v| v.as_array()) + .map(|a| a.iter().filter_map(|v| v.as_str()).collect()) + .unwrap_or_default(); + if args != ["nextest", "run"] { + bail!("args are {args:?}, expected [\"nextest\", \"run\"]"); + } + let timeout = value.get("timeout_secs").and_then(|v| v.as_f64()).unwrap_or_default(); + if timeout != 300.0 { + bail!("timeout_secs is {timeout}, expected 300"); + } + Ok(()) +} + +// ---- Validation helpers ------------------------------------------------------ + +/// The model's code answer: thinking stripped, then the longest fenced block, +/// or the whole reply when nothing is fenced. +fn extract_code(reply: &str) -> String { + let reply = strip_thinking(reply); + let mut best: Option = None; + for (index, segment) in reply.split("```").enumerate() { + // Odd segments sit between fences. + if index % 2 == 1 { + let body = match segment.split_once('\n') { + // Drop a language tag line ("rust", "rs", possibly padded). + Some((first, rest)) if first.trim().len() <= 12 && !first.trim().contains(' ') => rest, + _ => segment, + }; + if best.as_ref().is_none_or(|b| body.len() > b.len()) { + best = Some(body.to_string()); + } + } + } + best.unwrap_or_else(|| reply.trim().to_string()) +} + +/// The first balanced JSON object in the reply (fences and prose tolerated). +fn extract_json(reply: &str) -> Result { + let reply = strip_thinking(reply); + let start = reply.find('{').context("no JSON object in reply")?; + let mut depth = 0usize; + let mut in_string = false; + let mut escaped = false; + for (offset, ch) in reply[start..].char_indices() { + if escaped { + escaped = false; + continue; + } + match ch { + '\\' if in_string => escaped = true, + '"' => in_string = !in_string, + '{' if !in_string => depth += 1, + '}' if !in_string => { + depth -= 1; + if depth == 0 { + let candidate = &reply[start..start + offset + ch.len_utf8()]; + return serde_json::from_str(candidate) + .with_context(|| format!("reply is not valid JSON: {candidate}")); + } + } + _ => {} + } + } + bail!("unbalanced JSON object in reply"); +} + +/// Remove `…` reasoning that some local models emit inline. +fn strip_thinking(reply: &str) -> &str { + match (reply.find(""), reply.find("")) { + (Some(open), Some(close)) if open < close => reply[close + "".len()..].trim(), + _ => reply.trim(), + } +} + +fn bench_scratch(task: &str) -> std::path::PathBuf { + std::env::temp_dir().join(format!("hi-team-bench-{}-{task}", std::process::id())) +} + +fn rustc_available() -> bool { + std::process::Command::new("rustc") + .arg("--version") + .output() + .is_ok_and(|out| out.status.success()) +} + +/// Compile `source` as a binary and run it; pass = clean exit. +fn compile_and_run(task: &str, source: &str) -> Result<()> { + let dir = bench_scratch(task); + std::fs::create_dir_all(&dir)?; + let main_rs = dir.join("main.rs"); + std::fs::write(&main_rs, source)?; + let binary = dir.join("bench-bin"); + let compile = std::process::Command::new("rustc") + .args(["--edition", "2021", "-o"]) + .arg(&binary) + .arg(&main_rs) + .output()?; + if !compile.status.success() { + bail!("doesn't compile: {}", first_error_line(&compile.stderr)); + } + let run = std::process::Command::new(&binary).output()?; + if !run.status.success() { + bail!("wrong behavior: {}", first_error_line(&run.stderr)); + } + Ok(()) +} + +/// Compile `source` as a library; pass = compiles cleanly. +fn compile_lib(task: &str, source: &str) -> Result<()> { + let dir = bench_scratch(task); + std::fs::create_dir_all(&dir)?; + let lib_rs = dir.join("lib.rs"); + std::fs::write(&lib_rs, source)?; + let compile = std::process::Command::new("rustc") + .args(["--edition", "2021", "--crate-type", "lib", "--out-dir"]) + .arg(&dir) + .arg(&lib_rs) + .output()?; + if !compile.status.success() { + bail!("doesn't compile: {}", first_error_line(&compile.stderr)); + } + Ok(()) +} + +fn first_error_line(stderr: &[u8]) -> String { + let text = String::from_utf8_lossy(stderr); + text.lines() + .find(|line| line.contains("error") || line.contains("panicked")) + .unwrap_or_else(|| text.lines().next().unwrap_or("unknown error")) + .trim() + .to_string() +} + +// ---- Reporting --------------------------------------------------------------- + +fn print_summary(reports: &[ModelReport], tasks: &[BenchTask]) { + println!("results ({} tasks, compiled + asserted locally):", tasks.len()); + let name_width = reports.iter().map(|r| r.model.len()).max().unwrap_or(8).max(8); + let mut header = format!("{:6.1} {:>4.0}s", + report.passed(), + report.results.len(), + report.avg_tokens_per_sec(), + report.setup_secs, + )); + println!("{row}"); + } + // Per-task winners: every model that passed, fastest generation first. + for task in tasks { + let mut passed: Vec<(&str, f64)> = reports + .iter() + .filter_map(|report| { + report + .results + .iter() + .find(|r| r.task == task.name && r.pass) + .map(|r| (report.model.as_str(), r.latency_secs)) + }) + .collect(); + passed.sort_by(|a, b| a.1.total_cmp(&b.1)); + let line = if passed.is_empty() { + "nobody passed".to_string() + } else { + passed + .iter() + .map(|(model, secs)| format!("{model} ({secs:.1}s)")) + .collect::>() + .join(", ") + }; + println!(" {:<8} → {line}", task.name); + } +} + +fn save_report(reports: &[ModelReport]) -> Option { + let stamp = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .ok()? + .as_secs(); + let dir = std::env::var_os("HOME").map(std::path::PathBuf::from)?.join(".hi").join("bench"); + std::fs::create_dir_all(&dir).ok()?; + let path = dir.join(format!("team-bench-{stamp}.json")); + let body = serde_json::to_string_pretty(&serde_json::json!({ + "ran_at_unix_secs": stamp, + "models": reports, + })) + .ok()?; + std::fs::write(&path, body).ok()?; + Some(path) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extract_code_prefers_the_largest_fenced_block_and_strips_thinking() { + let reply = "plan planHere you go:\n```rust\npub fn a() {}\n```\nand\n```rust\npub fn bigger() { let x = 1; }\n```"; + assert_eq!(extract_code(reply), "pub fn bigger() { let x = 1; }\n"); + assert_eq!(extract_code("pub fn plain() {}"), "pub fn plain() {}"); + } + + #[test] + fn extract_json_tolerates_prose_and_fences() { + let value = extract_json("Sure!\n```json\n{\"cmd\": \"cargo\", \"args\": [\"nextest\", \"run\"], \"timeout_secs\": 300}\n```").unwrap(); + assert_eq!(value["cmd"], "cargo"); + assert!(extract_json("no json here").is_err()); + } + + #[test] + fn json_check_is_strict_about_values() { + assert!(check_json("{\"cmd\": \"cargo\", \"args\": [\"nextest\", \"run\"], \"timeout_secs\": 300}").is_ok()); + assert!(check_json("{\"cmd\": \"cargo nextest run\", \"args\": [], \"timeout_secs\": 300}").is_err()); + assert!(check_json("{\"cmd\": \"cargo\", \"args\": [\"nextest\", \"run\"], \"timeout_secs\": 30}").is_err()); + } + + #[test] + fn edit_check_requires_a_complete_faithful_rename() { + let good = "#[must_use]\npub fn load_user(id: u64) -> Option { if id == 0 { return None; } Some(format!(\"user-{id}\")) }\npub fn describe(id: u64) -> String { match load_user(id) { Some(user) => user, None => \"anonymous\".to_string() } }"; + assert!(check_edit(good).is_ok()); + let stale_call_site = good.replace("match load_user", "match fetch_user"); + assert!(check_edit(&stale_call_site).is_err()); + } +} diff --git a/crates/hi-local-core/src/model.rs b/crates/hi-local-core/src/model.rs index 2b9061e6..78e7f910 100644 --- a/crates/hi-local-core/src/model.rs +++ b/crates/hi-local-core/src/model.rs @@ -2,7 +2,11 @@ use std::path::PathBuf; use serde::{Deserialize, Serialize}; -pub const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 2048; +/// Per-request output ceiling when the model's generation_config doesn't set +/// `max_new_tokens`. Modern coding models reason inline before answering and +/// team executors emit whole diffs, so 2048 truncated real work mid-function; +/// the context window still bounds the total. +pub const DEFAULT_MAX_OUTPUT_TOKENS: u32 = 8192; #[derive(Clone, Debug, Serialize, Deserialize)] pub struct ModelInfo { diff --git a/crates/hi-sqlite-journal/src/lib.rs b/crates/hi-sqlite-journal/src/lib.rs index 2d1013cb..0eda01aa 100644 --- a/crates/hi-sqlite-journal/src/lib.rs +++ b/crates/hi-sqlite-journal/src/lib.rs @@ -26,6 +26,22 @@ use rusqlite::Connection; /// Wait for peers' locks instead of failing instantly. const BUSY_TIMEOUT_MS: u32 = 5000; +/// Poll cadence and cap for the journal-mode switch, which SQLite refuses to +/// apply busy_timeout to (see [`JournalMode::apply`]). +const MODE_SWITCH_POLL: std::time::Duration = std::time::Duration::from_millis(50); +const MODE_SWITCH_WAIT_MAX: std::time::Duration = std::time::Duration::from_secs(5); + +/// SQLITE_BUSY / SQLITE_LOCKED — transient peer contention, safe to retry. +fn is_busy(error: &rusqlite::Error) -> bool { + matches!( + error, + rusqlite::Error::SqliteFailure(failure, _) if matches!( + failure.code, + rusqlite::ErrorCode::DatabaseBusy | rusqlite::ErrorCode::DatabaseLocked + ) + ) +} + /// Journal mode chosen for a SQLite database based on where it lives. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum JournalMode { @@ -70,12 +86,27 @@ impl JournalMode { /// Apply the journal mode to an existing connection. pub fn apply(&self, conn: &Connection) -> Result<()> { - let mode = self.as_str(); - conn.pragma_update(None, "journal_mode", mode) - .with_context(|| format!("setting journal_mode to {mode}"))?; conn.pragma_update(None, "busy_timeout", BUSY_TIMEOUT_MS) .context("setting busy_timeout")?; - Ok(()) + let mode = self.as_str(); + // Switching journal modes promotes to an exclusive lock, and SQLite + // skips the busy handler on that promotion when a peer holds a write + // lock (waiting could deadlock) — so the pragma returns SQLITE_BUSY + // instantly despite busy_timeout. Poll instead. Steady-state opens + // never loop: the mode is persisted in the file, so re-applying it + // is a lock-free no-op. + let mut waited = std::time::Duration::ZERO; + loop { + match conn.pragma_update(None, "journal_mode", mode) { + Err(error) if is_busy(&error) && waited < MODE_SWITCH_WAIT_MAX => { + std::thread::sleep(MODE_SWITCH_POLL); + waited += MODE_SWITCH_POLL; + } + result => { + return result.with_context(|| format!("setting journal_mode to {mode}")); + } + } + } } /// Open a connection at `db_path` with the appropriate journal mode. @@ -289,6 +320,29 @@ mod tests { assert!(conn.execute("INSERT INTO t VALUES ('nope')", []).is_err()); } + /// busy_timeout must be configured before the journal-mode switch: the + /// switch needs an exclusive lock, and with the default timeout of 0 a + /// peer's write lock fails the open instantly instead of delaying it. + #[test] + fn open_waits_for_peer_write_lock() { + let tmp = tempfile::tempdir().unwrap(); + let db = tmp.path().join("contended.sqlite"); + // A pre-WAL database forces the exclusive-lock path in apply(). + let peer = Connection::open(&db).unwrap(); + peer.execute_batch("CREATE TABLE t(x); BEGIN IMMEDIATE; INSERT INTO t VALUES(1);") + .unwrap(); + let opener = { + let db = db.clone(); + std::thread::spawn(move || JournalMode::Wal.open(&db).map(|_| ())) + }; + std::thread::sleep(std::time::Duration::from_millis(300)); + peer.execute_batch("COMMIT;").unwrap(); + opener + .join() + .unwrap() + .expect("open should wait for the peer's lock instead of failing instantly"); + } + #[test] fn as_str_returns_correct_value() { assert_eq!(JournalMode::Wal.as_str(), "wal"); diff --git a/crates/hi-tools/src/background.rs b/crates/hi-tools/src/background.rs index 5cb9af2d..91ac2709 100644 --- a/crates/hi-tools/src/background.rs +++ b/crates/hi-tools/src/background.rs @@ -392,6 +392,14 @@ impl BackgroundRegistry { kill_all_from(self) } + /// The OS process id (process-group leader) behind a handle, when known. + /// Lets callers sample live resource usage (e.g. RSS while a model + /// server loads weights) for progress display. + pub fn os_pid(&self, id: &str) -> Option { + let processes = self.processes.lock().unwrap(); + processes.get(id).and_then(|proc| proc.pgid) + } + pub fn ids(&self) -> Vec { ids_from(self) } diff --git a/crates/hi-tools/src/hf.rs b/crates/hi-tools/src/hf.rs index 6093fc9e..d71c66d3 100644 --- a/crates/hi-tools/src/hf.rs +++ b/crates/hi-tools/src/hf.rs @@ -233,6 +233,47 @@ pub async fn handle_hf_command_result( } } +/// [`download_repo_keep_foreground`], but with every byte of downloader +/// output captured instead of written to the terminal. For frontends that +/// own the screen (the TUI runs in raw mode on the alternate screen): a +/// background model fetch must never paint over the UI. On failure the +/// captured tail is folded into the error for diagnosis. +pub async fn download_repo_keep_quiet( + repo_source: &str, + output_dir: impl AsRef, +) -> Result<()> { + let client = hi_ai::HuggingFaceHubClient::from_env(); + let repo = hi_ai::HfRepoRef::parse(repo_source)?; + let files = client.list_files(&repo).await?; + if files.is_empty() { + bail!("no files found in {}@{}", repo.repo_id, repo.revision); + } + let output_dir = output_dir.as_ref(); + let command = all_download_command(&client, &repo, &files, output_dir, WholeRepoMode::Keep)?; + let output = tokio::process::Command::new("sh") + .arg("-c") + .arg(&command) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await?; + if !output.status.success() { + let mut tail = String::from_utf8_lossy(&output.stderr).trim().to_string(); + if tail.is_empty() { + tail = String::from_utf8_lossy(&output.stdout).trim().to_string(); + } + let tail: String = tail.chars().rev().take(600).collect::().chars().rev().collect(); + bail!( + "download failed for {}@{} into {} ({}): {tail}", + repo.repo_id, + repo.revision, + output_dir.display(), + output.status + ); + } + Ok(()) +} + pub async fn download_repo_keep_foreground( repo_source: &str, output_dir: impl AsRef, diff --git a/crates/hi-tools/src/lib.rs b/crates/hi-tools/src/lib.rs index 6415cdf0..0bfbe1ff 100644 --- a/crates/hi-tools/src/lib.rs +++ b/crates/hi-tools/src/lib.rs @@ -87,11 +87,11 @@ pub mod infra { run_affected_polyglot_checks, run_affected_polyglot_tests, rust_source_paths, }; pub use crate::hf::{ - HfCommandResult, HfCommandState, HfMlxRun, download_repo_keep_foreground, + HfCommandResult, HfCommandState, HfMlxRun, download_repo_keep_foreground, download_repo_keep_quiet, handle_hf_command, handle_hf_command_result, }; pub use crate::local_server::{ - LocalServerHandle, skeptic_model_dir, start_local_server, stop_all_local_servers, + LocalServerHandle, await_local_server_health, local_server_os_pid, skeptic_model_dir, spawn_local_server, start_local_server, start_local_server_with_deadline, stop_all_local_servers, stop_local_server, }; pub use crate::lsp::lsp_status_report_for; @@ -145,11 +145,11 @@ pub use fast_feedback::{ run_affected_polyglot_checks, run_affected_polyglot_tests, rust_source_paths, }; pub use hf::{ - HfCommandResult, HfCommandState, HfMlxRun, download_repo_keep_foreground, handle_hf_command, + HfCommandResult, HfCommandState, HfMlxRun, download_repo_keep_foreground, download_repo_keep_quiet, handle_hf_command, handle_hf_command_result, }; pub use local_server::{ - LocalServerHandle, skeptic_model_dir, start_local_server, stop_all_local_servers, + LocalServerHandle, await_local_server_health, local_server_os_pid, skeptic_model_dir, spawn_local_server, start_local_server, start_local_server_with_deadline, stop_all_local_servers, stop_local_server, }; pub use lsp::lsp_status_report_for; diff --git a/crates/hi-tools/src/local_server.rs b/crates/hi-tools/src/local_server.rs index bd968416..216784f2 100644 --- a/crates/hi-tools/src/local_server.rs +++ b/crates/hi-tools/src/local_server.rs @@ -25,15 +25,25 @@ pub struct LocalServerHandle { pub endpoint: String, } -/// Cache directory for a downloaded local model, matching `/hf run --mlx`'s -/// layout (`$HI_MLX_MODELS_DIR` or `./.hi/models`, repo id sanitized) so a model -/// fetched by either path is reused rather than downloaded twice. Uses the -/// `main` revision (no `@rev` suffix). +/// Cache directory for a downloaded local model (`$HI_MLX_MODELS_DIR`, else +/// `~/.hi/models`, repo id sanitized) so a model fetched by any project is +/// reused rather than downloaded twice — a 50 GB checkpoint must never be +/// duplicated per working directory. Downloads that predate the shared root +/// (in the old cwd-relative `./.hi/models`) keep working via a per-repo +/// fallback. Uses the `main` revision (no `@rev` suffix). pub fn skeptic_model_dir(repo_id: &str) -> PathBuf { - let root = std::env::var_os("HI_MLX_MODELS_DIR") - .map(PathBuf::from) - .unwrap_or_else(|| PathBuf::from(".hi").join("models")); - root.join(crate::hf::safe_path(repo_id)) + let safe = crate::hf::safe_path(repo_id); + if let Some(root) = std::env::var_os("HI_MLX_MODELS_DIR") { + return PathBuf::from(root).join(safe); + } + let legacy = PathBuf::from(".hi").join("models").join(&safe); + let shared = std::env::var_os("HOME") + .map(|home| PathBuf::from(home).join(".hi").join("models").join(&safe)); + match shared { + Some(shared) if !shared.is_dir() && legacy.is_dir() => legacy, + Some(shared) => shared, + None => legacy, + } } /// Spawn `bin serve ` in the background and wait for `/health` to report @@ -46,21 +56,78 @@ pub async fn start_local_server( host: &str, port: u16, ) -> Result { + start_local_server_with_deadline(bin, serve_args, host, port, Duration::from_secs(15)).await +} + +/// [`start_local_server`] with an explicit health deadline. Large models +/// need it: loading a 19 GB 32B checkpoint into memory takes minutes, not +/// the 15 seconds that suit a small review model. +pub async fn start_local_server_with_deadline( + bin: &Path, + serve_args: &[String], + host: &str, + port: u16, + health_deadline: Duration, +) -> Result { + let process_id = spawn_local_server(bin, serve_args)?; + await_local_server_health(&process_id, host, port, health_deadline).await?; + Ok(LocalServerHandle { + process_id, + endpoint: format!("http://{host}:{port}/v1"), + }) +} + +/// Spawn the server process without waiting for readiness. Callers that want +/// live load progress use this, then poll [`await_local_server_health`] — +/// the handle is available immediately, so its RSS can be sampled while the +/// model loads. +pub fn spawn_local_server(bin: &Path, serve_args: &[String]) -> Result { let mut command = crate::web::shell_quote(&bin.to_string_lossy()); for arg in serve_args { command.push(' '); command.push_str(&crate::web::shell_quote(arg)); } let runner = crate::ProcessRunner::new(std::env::current_dir()?)?; - let process_id = LOCAL_SERVERS.spawn(&runner, &command)?; - match wait_for_health(host, port).await { - Ok(()) => Ok(LocalServerHandle { - process_id, - endpoint: format!("http://{host}:{port}/v1"), - }), + LOCAL_SERVERS.spawn(&runner, &command) +} + +/// The OS pid of a spawned local server (for RSS-based load progress). +pub fn local_server_os_pid(process_id: &str) -> Option { + LOCAL_SERVERS.os_pid(process_id) +} + +/// Wait for a spawned server to become healthy. Fails FAST when the process +/// exits during startup — a crashed server (wrong backend feature, bad +/// weights) must error in seconds with its output, not after staring at a +/// dead port for the whole multi-minute deadline. +pub async fn await_local_server_health( + process_id: &str, + host: &str, + port: u16, + health_deadline: Duration, +) -> Result<()> { + let deadline = tokio::time::Instant::now() + health_deadline; + loop { + if let Ok(outcome) = LOCAL_SERVERS.outcome(process_id) + && outcome.state != crate::BackgroundState::Running + { + let output = LOCAL_SERVERS.poll(process_id).unwrap_or_default(); + let tail: String = output.chars().rev().take(500).collect::().chars().rev().collect(); + anyhow::bail!("the local model server exited during startup: {tail}"); + } + if try_health_once(host, port).await { + return Ok(()); + } + if tokio::time::Instant::now() >= deadline { + break; + } + tokio::time::sleep(Duration::from_millis(400)).await; + } + match wait_for_health(host, port, Duration::from_millis(1)).await { + Ok(()) => Ok(()), Err(err) => { - let output = LOCAL_SERVERS.poll(&process_id).unwrap_or_default(); - let _ = LOCAL_SERVERS.kill(&process_id); + let output = LOCAL_SERVERS.poll(process_id).unwrap_or_default(); + let _ = LOCAL_SERVERS.kill(process_id); bail!("hi-local did not become healthy at http://{host}:{port}: {err}\n{output}"); } } @@ -78,14 +145,14 @@ pub fn stop_all_local_servers() { LOCAL_SERVERS.kill_all(); } -async fn wait_for_health(host: &str, port: u16) -> Result<()> { +async fn wait_for_health(host: &str, port: u16, health_deadline: Duration) -> Result<()> { let url = format!("http://{host}:{port}/health"); let client = reqwest::Client::builder() .connect_timeout(Duration::from_secs(1)) .timeout(Duration::from_secs(2)) .build() .unwrap_or_else(|_| hi_ai::timed_http_client_fallback(1, 2)); - let deadline = tokio::time::Instant::now() + Duration::from_secs(15); + let deadline = tokio::time::Instant::now() + health_deadline; let mut last_error = None; while tokio::time::Instant::now() < deadline { match client.get(&url).send().await { @@ -102,6 +169,60 @@ async fn wait_for_health(host: &str, port: u16) -> Result<()> { Err(last_error.unwrap_or_else(|| anyhow!("health check timed out"))) } +/// One health probe attempt: true when the server answers ready. +async fn try_health_once(host: &str, port: u16) -> bool { + let url = format!("http://{host}:{port}/health"); + let client = reqwest::Client::builder() + .connect_timeout(Duration::from_secs(1)) + .timeout(Duration::from_secs(2)) + .build() + .unwrap_or_else(|_| hi_ai::timed_http_client_fallback(1, 2)); + match client.get(&url).send().await { + Ok(response) if response.status().is_success() => match response.json().await { + Ok(body) => health_ready(&body), + Err(_) => false, + }, + _ => false, + } +} + fn health_ready(body: &serde_json::Value) -> bool { body.get("ready").and_then(serde_json::Value::as_bool) == Some(true) } + +#[cfg(test)] +mod model_dir_tests { + // SAFETY: nextest runs each test in its own process, so env/cwd mutation + // can't race other tests. + #[test] + fn model_dir_prefers_env_then_shared_home_with_legacy_fallback() { + let scratch = std::env::temp_dir().join(format!("hi-modeldir-{}", std::process::id())); + std::fs::create_dir_all(&scratch).unwrap(); + std::env::set_current_dir(&scratch).unwrap(); + let home = scratch.join("home"); + std::fs::create_dir_all(&home).unwrap(); + unsafe { std::env::set_var("HOME", &home) }; + + unsafe { std::env::set_var("HI_MLX_MODELS_DIR", scratch.join("override")) }; + assert!(super::skeptic_model_dir("a/b").starts_with(scratch.join("override"))); + unsafe { std::env::remove_var("HI_MLX_MODELS_DIR") }; + + let shared = home.join(".hi").join("models").join("a_b"); + assert_eq!(super::skeptic_model_dir("a/b"), shared, "default is the shared home root"); + + std::fs::create_dir_all(scratch.join(".hi").join("models").join("a_b")).unwrap(); + assert_eq!( + super::skeptic_model_dir("a/b"), + std::path::PathBuf::from(".hi").join("models").join("a_b"), + "a pre-existing cwd-local download keeps working (cwd-relative, as always)" + ); + + std::fs::create_dir_all(&shared).unwrap(); + assert_eq!( + super::skeptic_model_dir("a/b"), + shared, + "once the shared copy exists it wins over the legacy one" + ); + let _ = std::fs::remove_dir_all(&scratch); + } +} diff --git a/crates/hi-tui/src/app/commands.rs b/crates/hi-tui/src/app/commands.rs index 5ab1c7a4..30ffd718 100644 --- a/crates/hi-tui/src/app/commands.rs +++ b/crates/hi-tui/src/app/commands.rs @@ -1631,6 +1631,9 @@ impl crate::App { }; self.push(Line::styled(msg, dim())); } + Command::Team(arg) => { + self.handle_team_command(agent, &arg); + } Command::Config(arg) => { use hi_agent::command::{ConfigArg, parse_config_arg}; match parse_config_arg(&arg) { diff --git a/crates/hi-tui/src/app/lifecycle.rs b/crates/hi-tui/src/app/lifecycle.rs index 7c0c5c37..e0d0d239 100644 --- a/crates/hi-tui/src/app/lifecycle.rs +++ b/crates/hi-tui/src/app/lifecycle.rs @@ -171,6 +171,8 @@ impl crate::App { session_renamer: None, session_host: None, pending_host_enable: None, + pending_team_provision: None, + team_picker_role: None, sync_control: None, remote_event_tap: None, sync_remote_ui: None, diff --git a/crates/hi-tui/src/app/models.rs b/crates/hi-tui/src/app/models.rs index 943288e4..a0e99960 100644 --- a/crates/hi-tui/src/app/models.rs +++ b/crates/hi-tui/src/app/models.rs @@ -75,13 +75,32 @@ impl crate::App { (window > 0).then(|| (self.context_used * 100 / window).min(100)) } - /// Apply the picker's current selection as the model, then close it. + /// Apply the picker's current selection, then close it. A picker opened + /// by `/team ` assigns the chosen supported local model to that + /// role; otherwise the selection switches the driver model as always. pub(crate) fn pick_model(&mut self, agent: &mut Agent) { let id = self .picker .as_ref() .and_then(|p| p.current()) .map(str::to_string); + if let Some(role) = self.team_picker_role.take() { + self.picker = None; + if let Some(id) = id { + // Rows are rendered as "name — label · fit"; the leading + // token is the catalog name. + let name = id.split_whitespace().next().unwrap_or_default().to_string(); + if let Some(resolved) = hi_agent::local_skeptic::resolve_team_local_model( + &name, + hi_agent::local_skeptic::system_ram_gb(), + hi_agent::local_skeptic::detect_backend(), + ) { + self.assign_supported_local_model(agent, &role, resolved); + } + } + self.follow(); + return; + } if let Some(id) = id { self.select_model(agent, &id); } diff --git a/crates/hi-tui/src/app/run/mod.rs b/crates/hi-tui/src/app/run/mod.rs index f89f09d3..3d5b08ed 100644 --- a/crates/hi-tui/src/app/run/mod.rs +++ b/crates/hi-tui/src/app/run/mod.rs @@ -371,6 +371,9 @@ pub async fn run(agent: &mut Agent, options: crate::RunOptions) -> Result<()> { // Startup host-enable runs in the background; // apply its outcome once it lands. app.poll_pending_host_enable().await; + // `/team` local-model provisioning finishes in + // the background; wire the role when ready. + app.poll_pending_team_provision(agent).await; // Host mode: pull any attach prompts into the // turn queue without a separate daemon process. if app.drain_remote_input() { diff --git a/crates/hi-tui/src/app/sync_commands.rs b/crates/hi-tui/src/app/sync_commands.rs index 43904910..9d2eca02 100644 --- a/crates/hi-tui/src/app/sync_commands.rs +++ b/crates/hi-tui/src/app/sync_commands.rs @@ -884,6 +884,335 @@ impl crate::App { } } + /// `/team [role] [model|local|off] [base-url] [api-key]` — the team-role + /// table: the driver plans on the big model while explore/delegate + /// executors can run elsewhere (typically a local model, so execution + /// rounds cost nothing). Route changes apply to children started after + /// the command. + pub(crate) fn handle_team_command(&mut self, agent: &mut hi_agent::Agent, arg: &str) { + let parts: Vec<&str> = arg.split_whitespace().collect(); + if parts.is_empty() { + for row in agent.team_roles() { + let suffix = if row.inherited { " (driver)" } else { "" }; + self.push(Line::styled( + format!(" {:<9} {} @ {}{}", row.role, row.model, row.route, suffix), + dim(), + )); + } + if let Some(pending) = &self.pending_team_provision { + let phase = pending.phase_rx.borrow().clone(); + self.push(Line::styled( + format!( + " {:<9} setting up {} — {}", + pending.role, + pending.display, + provision_phase_line(&pending.display, &phase) + .trim_start_matches("⟳ ") + .trim_start_matches(&format!("{}: ", pending.display)) + ), + dim(), + )); + } + self.push(Line::styled( + " /team — pick from a list · /team · /team planner ", + dim(), + )); + let supported = hi_agent::local_skeptic::SUPPORTED_LOCAL_MODELS + .iter() + .map(|entry| entry.name) + .collect::>() + .join(", "); + self.push(Line::styled( + format!(" local models (auto-download + serve): local (auto-size), {supported}"), + dim(), + )); + self.push(Line::styled( + " auto picks the best quant for this machine; force one with name@quant (e.g. laguna-s@4bit)", + dim(), + )); + self.follow(); + return; + } + let role = parts[0]; + let value = parts.get(1).copied(); + match (role, value) { + ("driver", _) => { + self.push(Line::styled( + "the driver is the session model — switch it with /model or /provider", + dim(), + )); + } + ("skeptic", _) => { + self.push(Line::styled( + "skeptic routing has dedicated commands: /config skeptic-local on|off (auto-managed) or HI_SKEPTIC_ENDPOINT", + dim(), + )); + } + ("planner", Some("off")) => { + agent.set_planner_model(None); + self.push(Line::styled("planner → driver model", dim())); + } + ("planner", Some(model)) => { + agent.set_planner_model(Some(model.to_string())); + self.push(Line::styled(format!("planner → {model}"), dim())); + } + ("explore" | "delegate", None) => { + self.open_team_model_picker(role); + } + ("explore" | "delegate", Some("off")) => { + if role == "delegate" { + agent.set_delegate_route(None, None, None); + } else { + agent.set_explore_route(None, None, None); + } + self.push(Line::styled( + format!("{role} → driver route (applies to new {role} runs)"), + dim(), + )); + } + ("explore" | "delegate", Some(model)) => { + // Power users may still pass an explicit endpoint; everyone + // else picks a name and hi does the rest. + let explicit_endpoint = parts + .get(2) + .filter(|value| value.starts_with("http")) + .map(|s| s.to_string()); + if let Some(endpoint) = explicit_endpoint { + let key = parts.get(3).map(|s| s.to_string()); + if role == "delegate" { + agent.set_delegate_route(Some(model.to_string()), Some(endpoint.clone()), key); + } else { + agent.set_explore_route(Some(model.to_string()), Some(endpoint.clone()), key); + } + self.push(Line::styled( + format!("{role} → {model} @ {endpoint} (applies to new {role} runs)"), + dim(), + )); + } else if let Some(resolved) = hi_agent::local_skeptic::resolve_team_local_model( + model, + hi_agent::local_skeptic::system_ram_gb(), + hi_agent::local_skeptic::detect_backend(), + ) { + self.assign_supported_local_model(agent, role, resolved); + } else { + // Not a supported local name → a model id on the driver's + // provider (e.g. a cheaper cloud model for recon). + if role == "delegate" { + agent.set_delegate_route(Some(model.to_string()), None, None); + } else { + agent.set_explore_route(Some(model.to_string()), None, None); + } + self.push(Line::styled( + format!("{role} → {model} (driver route; applies to new {role} runs)"), + dim(), + )); + } + } + (other, _) => { + self.push(Line::styled( + format!("unknown role '{other}' — roles: driver, explore, delegate, skeptic, planner"), + dim(), + )); + } + } + self.follow(); + } + + /// Wire a supported local model to a role: reuse a running managed server + /// when one already serves it, otherwise provision (download + spawn) on + /// a background task and wire the role when it completes. + pub(crate) fn assign_supported_local_model( + &mut self, + agent: &mut hi_agent::Agent, + role: &str, + resolved: hi_agent::local_skeptic::ResolvedLocalModel, + ) { + let reuse = resolved + .mlx + .and_then(|quant| agent.running_local_model_server(quant.model_id)) + .or_else(|| { + resolved + .entry + .cuda + .and_then(|cuda| agent.running_local_model_server(cuda.model_id)) + }); + if let Some((endpoint, model_id)) = reuse { + if role == "delegate" { + agent.set_delegate_route(Some(model_id.clone()), Some(endpoint.clone()), None); + } else { + agent.set_explore_route(Some(model_id.clone()), Some(endpoint.clone()), None); + } + self.push(Line::styled( + format!("{role} → {model_id} @ local (reusing the running server; applies to new {role} runs)"), + dim(), + )); + return; + } + if let Some(pending) = &self.pending_team_provision { + self.push(Line::styled( + format!( + "already setting up {} for {} — one local setup at a time; retry when it finishes", + pending.display, pending.role + ), + dim(), + )); + return; + } + let Some(backend) = hi_agent::local_skeptic::detect_backend() else { + self.push(Line::styled( + "no local-inference backend on this machine (needs Apple Silicon or an NVIDIA runtime); the role stays on the driver", + dim(), + )); + return; + }; + let spec = match hi_agent::local_skeptic::team_model_spec(resolved, backend) { + Ok(spec) => spec, + Err(error) => { + self.push(Line::styled(format!("{error:#}"), dim())); + self.follow(); + return; + } + }; + let display = resolved.display(); + let model_dir = hi_tools::skeptic_model_dir(&spec.repo); + let (phase_tx, phase_rx) = + tokio::sync::watch::channel(hi_agent::local_skeptic::ProvisionPhase::Resolving); + let task = tokio::spawn(async move { + hi_agent::local_skeptic::provision_team_local_model(resolved, phase_tx).await + }); + self.pending_team_provision = Some(crate::PendingTeamProvision { + role: role.to_string(), + display: display.clone(), + task, + phase_rx, + announced_phase: hi_agent::local_skeptic::ProvisionPhase::Resolving, + phase_started: std::time::Instant::now(), + model_dir, + ticks_since_report: 0, + last_reported_bytes: 0, + progress_entry_index: None, + }); + self.push(Line::styled( + format!( + "⟳ setting up {display} locally for {role} — the download and server start run in the background; the role wires itself when ready" + ), + dim(), + )); + } + + /// Non-blocking: apply a finished `/team` local-model provisioning, if + /// any, and surface quiet download progress as an occasional dim line + /// (roughly every 30s of ticker time) while one is running. + pub(crate) async fn poll_pending_team_provision(&mut self, agent: &mut hi_agent::Agent) { + let finished = self + .pending_team_provision + .as_ref() + .is_some_and(|pending| pending.task.is_finished()); + if !finished { + let mut transition = None; + let mut heartbeat = None; + let mut index = None; + if let Some(pending) = &mut self.pending_team_provision { + // Phase transition: announce immediately with a permanent line + // and start a fresh in-place progress line for the new phase. + let current = pending.phase_rx.borrow().clone(); + if current != pending.announced_phase { + pending.announced_phase = current.clone(); + pending.phase_started = std::time::Instant::now(); + pending.ticks_since_report = 0; + pending.progress_entry_index = None; + transition = Some(provision_phase_line(&pending.display, ¤t)); + } + // Heartbeat within the phase: a single transcript line that + // updates in place (bar/percent/elapsed), not a spam stream. + pending.ticks_since_report = pending.ticks_since_report.saturating_add(1); + let cadence = provision_heartbeat_ticks(&pending.announced_phase); + if pending.ticks_since_report >= cadence { + pending.ticks_since_report = 0; + let bytes = dir_size_shallow(&pending.model_dir); + if let Some(line) = provision_heartbeat_line( + &pending.display, + &pending.announced_phase, + pending.phase_started.elapsed(), + bytes, + pending.last_reported_bytes, + ) { + pending.last_reported_bytes = bytes; + heartbeat = Some(line); + index = Some(pending.progress_entry_index); + } + } + } + let mut redraw = false; + if let Some(line) = transition { + self.push(Line::styled(line, dim())); + redraw = true; + } + if let Some(line) = heartbeat { + let mut slot = index.flatten(); + self.push_or_replace_progress(&mut slot, "⟳", Line::styled(line, dim())); + if let Some(pending) = &mut self.pending_team_provision { + pending.progress_entry_index = slot; + } + redraw = true; + } + if redraw { + self.follow(); + } + return; + } + let Some(pending) = self.pending_team_provision.take() else { + return; + }; + let result = match pending.task.await { + Ok(result) => result, + Err(join_error) => Err(anyhow::anyhow!("local setup task failed: {join_error}")), + }; + self.apply_team_provision_result(agent, &pending.role, &pending.display, result); + } + + /// Apply a provisioning outcome: success wires the role and registers the + /// server for reuse/teardown; failure leaves the role on the driver with + /// a calm note (never a raw error dump, never a broken workflow). + pub(crate) fn apply_team_provision_result( + &mut self, + agent: &mut hi_agent::Agent, + role: &str, + display: &str, + result: anyhow::Result<(String, String, String)>, + ) { + match result { + Ok((endpoint, model_id, process_id)) => { + agent.register_team_local_server( + endpoint.clone(), + model_id.clone(), + process_id, + ); + if role == "delegate" { + agent.set_delegate_route(Some(model_id.clone()), Some(endpoint), None); + } else { + agent.set_explore_route(Some(model_id.clone()), Some(endpoint), None); + } + self.push(Line::styled( + format!( + "✓ {role} → {model_id} @ local (ready — applies to new {role} runs)" + ), + Style::default().fg(crate::theme::theme().accent_success), + )); + } + Err(error) => { + let reason: String = format!("{error:#}").chars().take(140).collect(); + self.push(Line::styled( + format!( + "couldn't set up {display} locally ({reason}); {role} stays on the driver" + ), + dim(), + )); + } + } + self.follow(); + } + /// Kick off hosted-mode enablement without blocking the UI. The /// controller's network work (portal registration) runs on a background /// task; [`Self::poll_pending_host_enable`] applies the outcome from the @@ -1285,3 +1614,339 @@ impl crate::App { Ok(sessions) } } + +/// Total size of the files directly inside `dir` (model repos download flat). +/// Best-effort: unreadable entries count as zero. +pub(crate) fn dir_size_shallow(dir: &std::path::Path) -> u64 { + std::fs::read_dir(dir) + .map(|entries| { + entries + .flatten() + .filter_map(|entry| entry.metadata().ok()) + .filter(|meta| meta.is_file()) + .map(|meta| meta.len()) + .sum() + }) + .unwrap_or(0) +} + +/// The transcript line announcing a provisioning phase transition. +pub(crate) fn provision_phase_line( + display: &str, + phase: &hi_agent::local_skeptic::ProvisionPhase, +) -> String { + use hi_agent::local_skeptic::ProvisionPhase; + match phase { + ProvisionPhase::Resolving => format!("⟳ {display}: checking hardware and cached weights…"), + ProvisionPhase::Downloading => { + format!("⟳ {display}: downloading weights (quiet, in the background)…") + } + ProvisionPhase::BuildingServer => { + format!("⟳ {display}: compiling the hi-local serving binary (first run)…") + } + ProvisionPhase::LoadingModel { deadline_secs, .. } => format!( + "⟳ {display}: server started — loading weights into memory (can take up to {})…", + format_secs(*deadline_secs) + ), + } +} + +/// Heartbeat cadence per phase, in ~120ms ticker calls. Heartbeats update a +/// single transcript line IN PLACE, so the slow phases can refresh every +/// second or two without spamming: a live bar while weights load, a growing +/// GiB counter while downloading. +pub(crate) fn provision_heartbeat_ticks( + phase: &hi_agent::local_skeptic::ProvisionPhase, +) -> u32 { + use hi_agent::local_skeptic::ProvisionPhase; + match phase { + ProvisionPhase::Downloading => 16, + ProvisionPhase::LoadingModel { .. } => 8, + ProvisionPhase::Resolving | ProvisionPhase::BuildingServer => 40, + } +} + +/// The within-phase heartbeat line, when the phase warrants one. +pub(crate) fn provision_heartbeat_line( + display: &str, + phase: &hi_agent::local_skeptic::ProvisionPhase, + in_phase: std::time::Duration, + bytes_on_disk: u64, + last_reported_bytes: u64, +) -> Option { + use hi_agent::local_skeptic::ProvisionPhase; + match phase { + ProvisionPhase::Downloading => { + if bytes_on_disk > last_reported_bytes { + Some(format!( + "⟳ {display}: downloading — {:.1} GiB on disk…", + bytes_on_disk as f64 / (1024.0 * 1024.0 * 1024.0) + )) + } else { + Some(format!( + "⟳ {display}: still downloading ({} in)…", + format_secs(in_phase.as_secs()) + )) + } + } + ProvisionPhase::LoadingModel { + deadline_secs, + server_handle, + expected_bytes, + } => { + let rss = hi_tools::local_server_os_pid(server_handle).and_then(rss_bytes); + Some(loading_bar_line( + display, + rss, + *expected_bytes, + in_phase, + *deadline_secs, + )) + } + ProvisionPhase::BuildingServer => Some(format!( + "⟳ {display}: still compiling hi-local ({} in)…", + format_secs(in_phase.as_secs()) + )), + ProvisionPhase::Resolving => None, + } +} + +/// `95` → `1m35s`, `40` → `40s`. +pub(crate) fn format_secs(total: u64) -> String { + if total >= 60 { + format!("{}m{:02}s", total / 60, total % 60) + } else { + format!("{total}s") + } +} + +#[cfg(test)] +mod provision_narration_tests { + use super::*; + use hi_agent::local_skeptic::ProvisionPhase; + + fn loading_phase() -> ProvisionPhase { + ProvisionPhase::LoadingModel { + deadline_secs: 345, + server_handle: "bg_none".into(), + expected_bytes: 19 * 1024 * 1024 * 1024, + } + } + + #[test] + fn phase_lines_and_heartbeats_narrate_the_slow_parts() { + assert!(provision_phase_line("coder-32b", &loading_phase()).contains("up to 5m45s")); + // Unknown server pid → honest elapsed line, no fake bar. + let hb = provision_heartbeat_line( + "coder-32b", + &loading_phase(), + std::time::Duration::from_secs(95), + 0, + 0, + ) + .unwrap(); + assert!(hb.contains("1m35s elapsed"), "{hb}"); + let dl = provision_heartbeat_line( + "coder-32b", + &ProvisionPhase::Downloading, + std::time::Duration::from_secs(30), + 3 * 1024 * 1024 * 1024, + 1024, + ) + .unwrap(); + assert!(dl.contains("3.0 GiB on disk"), "{dl}"); + assert!( + provision_heartbeat_line("x", &ProvisionPhase::Resolving, Default::default(), 0, 0) + .is_none(), + "resolving is instant; no heartbeat spam" + ); + assert!( + provision_heartbeat_ticks(&loading_phase()) + < provision_heartbeat_ticks(&ProvisionPhase::Downloading), + "loading refreshes fastest — that's the phase mistaken for a hang" + ); + } + + #[test] + fn loading_bar_reflects_memory_growth_and_clamps() { + let gib = 1024u64 * 1024 * 1024; + let half = loading_bar_line( + "coder-32b", + Some(9 * gib + gib / 2), + 19 * gib, + std::time::Duration::from_secs(70), + 345, + ); + assert!(half.contains("50%"), "{half}"); + assert!(half.contains("9.5/19.0 GiB"), "{half}"); + assert!(half.contains('▰') && half.contains('▱'), "{half}"); + let over = loading_bar_line( + "coder-32b", + Some(25 * gib), + 19 * gib, + std::time::Duration::from_secs(200), + 345, + ); + assert!(over.contains("99%"), "clamps below done: {over}"); + assert_eq!(render_bar(0.5, 10), "▰▰▰▰▰▱▱▱▱▱"); + assert_eq!(render_bar(2.0, 4), "▰▰▰▰"); + } +} + +/// Resident memory of a process in bytes (`ps -o rss=` reports KiB). +pub(crate) fn rss_bytes(pid: i32) -> Option { + let output = std::process::Command::new("ps") + .args(["-o", "rss=", "-p", &pid.to_string()]) + .output() + .ok()?; + String::from_utf8(output.stdout) + .ok()? + .trim() + .parse::() + .ok() + .map(|kib| kib * 1024) +} + +/// The live weights-loading line: a bar of memory growth toward the model's +/// on-disk size. RSS slightly overshoots the weights, so the bar clamps at +/// 99% until the health check flips it to the ✓ line. +pub(crate) fn loading_bar_line( + display: &str, + rss_bytes: Option, + expected_bytes: u64, + in_phase: std::time::Duration, + deadline_secs: u64, +) -> String { + let gib = |bytes: u64| bytes as f64 / (1024.0 * 1024.0 * 1024.0); + match rss_bytes { + Some(rss) if expected_bytes > 0 => { + let frac = (rss as f64 / expected_bytes as f64).clamp(0.0, 0.99); + format!( + "⟳ {display}: loading weights {} {:>2.0}% · {:.1}/{:.1} GiB · {}", + render_bar(frac, 18), + frac * 100.0, + gib(rss), + gib(expected_bytes), + format_secs(in_phase.as_secs()), + ) + } + _ => format!( + "⟳ {display}: loading weights — {} elapsed (allow up to {})…", + format_secs(in_phase.as_secs()), + format_secs(deadline_secs) + ), + } +} + +/// `render_bar(0.5, 10)` → `▰▰▰▰▰▱▱▱▱▱`. +pub(crate) fn render_bar(frac: f64, width: usize) -> String { + let filled = ((frac.clamp(0.0, 1.0)) * width as f64).round() as usize; + let mut bar = String::with_capacity(width * 3); + for i in 0..width { + bar.push(if i < filled { '▰' } else { '▱' }); + } + bar +} + +impl crate::App { + /// `/team ` with no model: open the picker over the supported + /// catalog, largest first, annotated with what fits this machine and + /// what's already downloaded. Enter assigns the selection to the role. + pub(crate) fn open_team_model_picker(&mut self, role: &str) { + let ram = hi_agent::local_skeptic::system_ram_gb(); + let backend = hi_agent::local_skeptic::detect_backend(); + let mut entries: Vec<&'static hi_agent::local_skeptic::SupportedLocalModel> = + hi_agent::local_skeptic::SUPPORTED_LOCAL_MODELS.iter().collect(); + // Largest first, sized by the quant this machine would actually get + // (the ladder means a family's effective size is per-machine). + entries.sort_by_key(|entry| { + std::cmp::Reverse( + entry + .pick_mlx(ram) + .or_else(|| entry.smallest_mlx()) + .map(|quant| quant.min_ram_gb) + .unwrap_or_else(|| entry.min_ram_gb(backend)), + ) + }); + let rows: Vec = entries + .iter() + .map(|entry| team_picker_row(entry, ram, backend)) + .collect(); + let current = rows + .iter() + .find(|row| { + hi_agent::local_skeptic::resolve_team_local_model("local", ram, backend) + .is_some_and(|auto| row.starts_with(auto.entry.name)) + }) + .cloned() + .unwrap_or_default(); + self.team_picker_role = Some(role.to_string()); + self.picker = Some(ModelPicker::new( + rows, + ¤t, + std::collections::HashMap::new(), + &std::collections::HashMap::new(), + )); + self.push(Line::styled( + format!("pick a local model for {role} — ↑↓ to choose, Enter to set up, Esc to cancel"), + dim(), + )); + self.follow(); + } +} + +/// One picker row: `name — label · [· quants …] [· downloaded]`. +/// The fit note names the quant this machine would get, so a ladder like +/// Laguna's reads honestly: `3bit fits (needs 64GB RAM)` on a 64GB Mac. +pub(crate) fn team_picker_row( + entry: &'static hi_agent::local_skeptic::SupportedLocalModel, + ram_gb: u64, + backend: Option, +) -> String { + let chosen = entry.pick_mlx(ram_gb); + let fit = match (backend, chosen, entry.smallest_mlx()) { + (Some(hi_agent::local_skeptic::LocalBackend::Cuda), _, _) => { + match entry.cuda { + Some(cuda) if ram_gb >= cuda.min_ram_gb => { + format!("needs {}GB RAM · fits", cuda.min_ram_gb) + } + Some(cuda) => { + format!("needs {}GB RAM · too big for this machine", cuda.min_ram_gb) + } + None => "MLX-only — not packaged for CUDA yet".to_string(), + } + } + (_, Some(quant), _) if entry.mlx.len() > 1 => { + format!("{} fits (needs {}GB RAM)", quant.quant, quant.min_ram_gb) + } + (_, Some(quant), _) => format!("needs {}GB RAM · fits", quant.min_ram_gb), + (_, None, Some(smallest)) => { + format!( + "needs {}GB+ RAM · too big for this machine", + smallest.min_ram_gb + ) + } + (_, None, None) => "unavailable".to_string(), + }; + let downloaded = backend + .and_then(|backend| { + let resolved = hi_agent::local_skeptic::ResolvedLocalModel { + entry, + mlx: chosen.or_else(|| entry.smallest_mlx()), + }; + hi_agent::local_skeptic::team_model_spec(resolved, backend).ok() + }) + .map(|spec| { + let dir = hi_tools::skeptic_model_dir(&spec.repo); + hi_agent::local_skeptic::model_present(&dir, &spec) + }) + .unwrap_or(false); + let mut row = format!("{} — {} · {fit}", entry.name, entry.label); + if entry.mlx.len() > 1 { + row.push_str(&format!(" · quants {}", entry.quant_summary())); + } + if downloaded { + row.push_str(" · downloaded"); + } + row +} diff --git a/crates/hi-tui/src/app/transcript.rs b/crates/hi-tui/src/app/transcript.rs index f6d2f832..82861b61 100644 --- a/crates/hi-tui/src/app/transcript.rs +++ b/crates/hi-tui/src/app/transcript.rs @@ -55,6 +55,33 @@ impl crate::App { self.cap_transcript(); } + /// Push a live-updating progress line: while it remains the LAST + /// transcript entry (tracked via `index`) and still looks like one of our + /// progress lines (`marker` guard against index drift after capping), + /// the line is replaced in place — one smoothly-updating bar instead of + /// a line of spam per second. + pub(crate) fn push_or_replace_progress( + &mut self, + index: &mut Option, + marker: &str, + line: Line<'static>, + ) { + if let Some(at) = *index + && at + 1 == self.transcript.len() + && matches!( + self.transcript.get(at), + Some(TranscriptEntry::Line(existing)) + if existing.spans.iter().any(|span| span.content.contains(marker)) + ) + { + self.transcript[at] = TranscriptEntry::Line(line); + self.bump_transcript(); + return; + } + self.push(line); + *index = Some(self.transcript.len().saturating_sub(1)); + } + /// Push a user-prompt echo as a structurally-distinct entry so the render /// pass can pin it as a sticky header when scrolled past. pub(crate) fn push_user_prompt(&mut self, line: Line<'static>) { diff --git a/crates/hi-tui/src/lib.rs b/crates/hi-tui/src/lib.rs index bc99bdd7..9089bf0b 100644 --- a/crates/hi-tui/src/lib.rs +++ b/crates/hi-tui/src/lib.rs @@ -188,6 +188,30 @@ pub type SessionHostController = Box< + Sync, >; +/// An in-flight `/team` local-model provisioning task. +pub(crate) struct PendingTeamProvision { + pub(crate) role: String, + pub(crate) display: String, + pub(crate) task: tokio::task::JoinHandle>, + /// Live phase reported by the provisioning task (download → build → + /// load), so the transcript narrates what is actually happening. + pub(crate) phase_rx: tokio::sync::watch::Receiver, + /// The last phase already announced in the transcript. + pub(crate) announced_phase: hi_agent::local_skeptic::ProvisionPhase, + /// When the current phase began (drives "Ns elapsed" heartbeats). + pub(crate) phase_started: std::time::Instant, + /// Where the weights land — polled for size so download heartbeats can + /// say how much is on disk (the downloader itself is fully quiet; raw + /// aria2c output once painted over the alternate screen). + pub(crate) model_dir: std::path::PathBuf, + /// Ticker calls since the last heartbeat line. + pub(crate) ticks_since_report: u32, + /// Bytes on disk at the last heartbeat. + pub(crate) last_reported_bytes: u64, + /// Transcript index of the in-place progress line for the current phase. + pub(crate) progress_entry_index: Option, +} + /// A session cached on this machine, merged into the `/sessions` list view. #[derive(Clone, Debug)] pub struct LocalSessionInfo { @@ -1006,6 +1030,14 @@ pub(crate) struct App { pub(crate) session_renamer: Option, /// Enables/disables remote-input host mode for the active session. pub(crate) session_host: Option, + /// When set, the open model picker assigns its selection to this team + /// role (`/team delegate` with no argument) instead of switching the + /// driver model. + pub(crate) team_picker_role: Option, + /// In-flight `/team` local-model provisioning (download + server spawn on + /// a background task). The event loop applies the outcome when it lands; + /// a 15 GB model fetch must never block the UI. + pub(crate) pending_team_provision: Option, /// In-flight background host-enable (startup auto-host). The controller's /// network work (portal registration) runs off the UI path; the event /// loop applies the outcome when it completes. A dead portal must never diff --git a/crates/hi-tui/src/tests.rs b/crates/hi-tui/src/tests.rs index 5ee3fd19..d8b3b3d4 100644 --- a/crates/hi-tui/src/tests.rs +++ b/crates/hi-tui/src/tests.rs @@ -435,6 +435,253 @@ fn mouse_wheel_scrolls_and_repins_the_transcript() { assert!(app.following, "wheel-down at the bottom should re-pin"); } +#[tokio::test] +async fn team_command_routes_executors_and_clears_back_to_driver() { + // Deterministic backend + isolated weights dir: CI has no MLX/CUDA, and + // an aborted provisioning task must not litter the repo's .hi/models. + // SAFETY: nextest isolates each test in its own process. + unsafe { std::env::set_var("HI_LOCAL_BACKEND", "mlx") }; + unsafe { + std::env::set_var( + "HI_MLX_MODELS_DIR", + std::env::temp_dir().join(format!("hi-team-test-{}", std::process::id())), + ) + }; + let provider = std::sync::Arc::new(hi_ai::OpenAiProvider::new( + "http://127.0.0.1:1/v1".into(), + "test".into(), + )); + let mut agent = hi_agent::Agent::new(provider, hi_agent::AgentConfig::default()).unwrap(); + let mut app = test_app("openai", "gpt-4o"); + + // Bare /team renders the role table. + app.handle_command(&mut agent, hi_agent::Command::Team(String::new())) + .await; + let text = app.transcript_text(); + assert!(text.contains("driver"), "role table shown: {text}"); + assert!(text.contains("delegate"), "role table shown: {text}"); + + // Route the executors: delegate to a local endpoint, explore by model. + app.handle_command( + &mut agent, + hi_agent::Command::Team("delegate qwen3-coder http://127.0.0.1:18080/v1".into()), + ) + .await; + assert!( + app.transcript_text() + .contains("delegate → qwen3-coder @ http://127.0.0.1:18080/v1"), + "{}", + app.transcript_text() + ); + let delegate = agent + .team_roles() + .into_iter() + .find(|r| r.role == "delegate") + .unwrap(); + assert!(!delegate.inherited); + assert_eq!(delegate.model, "qwen3-coder"); + + // `local` is fully automated: it resolves a supported model for this + // hardware and starts the download/server setup in the background. + app.handle_command(&mut agent, hi_agent::Command::Team("explore local".into())) + .await; + assert!( + app.transcript_text().contains("locally for explore"), + "{}", + app.transcript_text() + ); + if let Some(pending) = app.pending_team_provision.take() { + pending.task.abort(); + } + + // `off` returns the role to the driver. + app.handle_command(&mut agent, hi_agent::Command::Team("delegate off".into())) + .await; + assert!( + agent + .team_roles() + .into_iter() + .find(|r| r.role == "delegate") + .unwrap() + .inherited, + "delegate returned to the driver route" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn team_supported_model_provisions_in_background_and_wires_on_success() { + // Deterministic backend + isolated weights dir: CI has no MLX/CUDA, and + // an aborted provisioning task must not litter the repo's .hi/models. + // SAFETY: nextest isolates each test in its own process. + unsafe { std::env::set_var("HI_LOCAL_BACKEND", "mlx") }; + unsafe { + std::env::set_var( + "HI_MLX_MODELS_DIR", + std::env::temp_dir().join(format!("hi-team-test-{}", std::process::id())), + ) + }; + let provider = std::sync::Arc::new(hi_ai::OpenAiProvider::new( + "http://127.0.0.1:1/v1".into(), + "test".into(), + )); + let mut agent = hi_agent::Agent::new(provider, hi_agent::AgentConfig::default()).unwrap(); + let mut app = test_app("openai", "gpt-4o"); + + // A supported name starts a background setup instead of demanding a URL. + app.handle_command( + &mut agent, + hi_agent::Command::Team("delegate coder-7b".into()), + ) + .await; + assert!( + app.transcript_text().contains("setting up coder-7b locally"), + "{}", + app.transcript_text() + ); + assert!(app.pending_team_provision.is_some(), "background task runs"); + assert!( + agent + .team_roles() + .into_iter() + .find(|r| r.role == "delegate") + .unwrap() + .inherited, + "role stays on the driver until the setup lands" + ); + + // Second request while one is in flight: calm, no second task. + app.handle_command( + &mut agent, + hi_agent::Command::Team("explore coder-7b".into()), + ) + .await; + assert!( + app.transcript_text().contains("one local setup at a time"), + "{}", + app.transcript_text() + ); + + // Simulate the provisioning outcome landing (the real task would need a + // model download; the apply path is what must be correct). + if let Some(pending) = app.pending_team_provision.take() { + pending.task.abort(); + } + app.apply_team_provision_result( + &mut agent, + "delegate", + "coder-7b", + Ok(( + "http://127.0.0.1:18080/v1".into(), + "Qwen2.5-Coder-7B-Instruct-4bit".into(), + "bg_local_1".into(), + )), + ); + let delegate = agent + .team_roles() + .into_iter() + .find(|r| r.role == "delegate") + .unwrap(); + assert_eq!(delegate.model, "Qwen2.5-Coder-7B-Instruct-4bit"); + assert_eq!(delegate.route, "http://127.0.0.1:18080/v1"); + assert!( + app.transcript_text().contains("✓ delegate → Qwen2.5-Coder-7B-Instruct-4bit @ local"), + "{}", + app.transcript_text() + ); + + // The registered server is now reused instantly for another role. + app.handle_command( + &mut agent, + hi_agent::Command::Team("explore coder-7b".into()), + ) + .await; + assert!( + app.transcript_text().contains("reusing the running server"), + "{}", + app.transcript_text() + ); + assert!(app.pending_team_provision.is_none(), "no new task needed"); + + // Failures stay calm: role unchanged, no raw error dump. + app.apply_team_provision_result( + &mut agent, + "explore", + "coder-32b", + Err(anyhow::anyhow!("no local-inference backend detected")), + ); + assert!( + app.transcript_text() + .contains("couldn't set up coder-32b locally"), + "{}", + app.transcript_text() + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn bare_team_role_opens_picker_and_selection_starts_setup() { + // SAFETY: nextest isolates each test in its own process. + unsafe { std::env::set_var("HI_LOCAL_BACKEND", "mlx") }; + unsafe { + std::env::set_var( + "HI_MLX_MODELS_DIR", + std::env::temp_dir().join(format!("hi-team-picker-{}", std::process::id())), + ) + }; + let provider = std::sync::Arc::new(hi_ai::OpenAiProvider::new( + "http://127.0.0.1:1/v1".into(), + "test".into(), + )); + let mut agent = hi_agent::Agent::new(provider, hi_agent::AgentConfig::default()).unwrap(); + let mut app = test_app("openai", "gpt-4o"); + + // `/team delegate` with no model opens the catalog picker. + app.handle_command(&mut agent, hi_agent::Command::Team("delegate".into())) + .await; + assert!(app.picker.is_some(), "picker opens"); + assert_eq!(app.team_picker_role.as_deref(), Some("delegate")); + let rows = app.picker.as_ref().unwrap().all.clone(); + assert!( + rows.iter().any(|row| row.starts_with("laguna-s")), + "laguna listed: {rows:?}" + ); + assert!( + rows.iter().any(|row| row.starts_with("qwen3.6-35b")), + "qwen3.6 listed: {rows:?}" + ); + assert!( + rows.iter() + .any(|row| row.starts_with("glm-5.2") && row.contains("too big")), + "oversized entries are annotated honestly: {rows:?}" + ); + + // Selecting a row routes to team setup, not a driver-model switch. + let coder_row_index = app.picker.as_ref().unwrap().all.iter() + .position(|row| row.starts_with("coder-7b")) + .expect("coder-7b row"); + if let Some(picker) = app.picker.as_mut() { + picker.selected = picker + .matches + .iter() + .position(|&i| i == coder_row_index) + .unwrap_or(0); + } + app.pick_model(&mut agent); + assert!(app.picker.is_none(), "picker closes"); + assert!(app.team_picker_role.is_none()); + assert_eq!( + "gpt-4o", app.model, + "the driver model must NOT change from a team pick" + ); + assert!( + app.transcript_text().contains("setting up coder-7b locally"), + "{}", + app.transcript_text() + ); + if let Some(pending) = app.pending_team_provision.take() { + pending.task.abort(); + } +} + #[tokio::test] async fn config_command_sets_disables_and_restores_automatic_step_limit() { let provider = std::sync::Arc::new(hi_ai::OpenAiProvider::new(