diff --git a/Cargo.lock b/Cargo.lock index 022e0d5..1a44e43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1036,6 +1036,9 @@ dependencies = [ "clap", "mofa-engine-core", "mofa-engine-sdk", + "mofa-kernel", + "serde", + "serde_json", "tokio", "tracing", "tracing-subscriber", @@ -1072,6 +1075,7 @@ dependencies = [ "chrono", "mofa-engine-core", "mofa-kernel", + "reqwest", "serde", "serde_json", "thiserror", @@ -1080,6 +1084,7 @@ dependencies = [ "tower", "tower-http", "tracing", + "uuid", ] [[package]] @@ -1090,6 +1095,7 @@ dependencies = [ "serde", "serde_json", "thiserror", + "tokio", "uuid", ] diff --git a/config.example.toml b/config.example.toml new file mode 100644 index 0000000..b132a08 --- /dev/null +++ b/config.example.toml @@ -0,0 +1,91 @@ +# MoFA Engine example configuration. +# +# Copy to `config.toml` (current directory) or +# `~/.config/mofa-engine/config.toml`. Any section may be omitted; defaults +# shown below are applied. API keys support `${ENV_VAR}` substitution. + +[listen] +host = "127.0.0.1" # loopback by default; remote access is opt-in +port = 8420 + +[memory] +# Omit budget_mb to auto-detect (70% of system RAM). +budget_mb = 12288 +idle_timeout_secs = 120 # unload a resident model after this many idle seconds + +[timeouts] +request_secs = 240 # overall budget for one /v1/invoke, across all candidates +queue_secs = 30 # max wait for a concurrency permit +load_secs = 60 # max time to load/warm one model +inference_secs = 180 # max time for one provider inference call +discovery_secs = 8 # per-provider discovery probe +health_secs = 5 # per-provider health probe + +[preflight] +enabled = true # master switch for predictive warming +history_learning = true # learn capability transition history +speculative_warming = true # warm models from hints/history/subscriptions +confidence_threshold = 0.6 # min confidence (0.0-1.0) before a history prediction warms +min_samples = 3 # min observed transitions before history is trusted +subscription_ttl_secs = 3600 + +[artifacts] +# dir = "/tmp" # where generated audio/artifacts are scanned (default: system temp) +retention_secs = 3600 # delete engine-generated artifacts older than this; 0 disables + +[security] +# Allowlist of directory roots that request `input_file` paths must fall under. +# Empty (default) allows any local path — appropriate only on a trusted machine. +# input_roots = ["/Users/me/mofa-workspace"] +input_roots = [] + +# Local Ollama (auto-discovered models). Priority is lower = preferred. +[[providers]] +name = "ollama" +kind = "ollama" +base_url = "http://127.0.0.1:11434" +priority = 1 +cost_tier = "free" + +# Cloud OpenAI-compatible backend. Models are declared explicitly. +[[providers]] +name = "openai" +kind = "openai_compatible" +base_url = "https://api.openai.com/v1" +api_key = "${OPENAI_API_KEY}" +priority = 10 +cost_tier = "high" + + [[providers.models]] + name = "gpt-4o" + capability = "chat" + context_window = 128000 + + [[providers.models]] + name = "tts-1" + capability = "tts" + + [[providers.models]] + name = "whisper-1" + capability = "asr" + +# Local TTS process-adapter backend. Runs a local command (an MLX/Kokoro or +# Piper-style CLI) once per synthesis and returns the produced audio file. +# Enable and point `command`/`args` at your installed runtime. Placeholders: +# {text} input text passed as an argument +# {text_file} path to a temp file holding the input text +# {output} path the command must write audio to +[[providers]] +name = "local-tts" +kind = "local_tts" +base_url = "" # unused for local_tts, but the field is required +command = "kokoro-tts" +args = ["--text-file", "{text_file}", "--output", "{output}"] +output_format = "wav" +priority = 2 +enabled = false # opt-in: set true once the command is installed + + [[providers.models]] + name = "kokoro" + capability = "tts" + memory_mb = 1024 diff --git a/docs/acceptance.md b/docs/acceptance.md index 5766815..156905a 100644 --- a/docs/acceptance.md +++ b/docs/acceptance.md @@ -14,15 +14,61 @@ This file converts the RFC acceptance checklist into implementation evidence. - [x] Ambiguous short model names are rejected. - [x] Unsupported provider operations return typed errors. -## MVP Acceptance Still Pending - -- [ ] Auto-discover local MLX models. -- [ ] Local Qwen via Ollama handles type-level LLM calls. -- [ ] Named GPT request routes to cloud exactly. -- [ ] Local TTS returns a managed audio artifact. -- [ ] Hint Preflight starts TTS loading before LLM completion. -- [ ] Idle models unload and release accounted memory. -- [ ] Local TTS failure follows documented fallback policy. -- [ ] HTTP, Rust, and Python interfaces complete equivalent flows. -- [ ] mofa-fm completes English article to Chinese podcast end to end. +## Stage 3 — Routing, Concurrency, and Failure Handling + +- [x] Routing applies hard constraints, then ranks candidates and returns a + reason (`router::route_ranked`). +- [x] Static memory feasibility filter drops models larger than the budget. +- [x] Primary selection and failover walk the *same* ranked candidate plan. +- [x] Capability routing prefers local models under the default policy. +- [x] Named routing resolves exactly; ambiguous short names are rejected. +- [x] A retryable local failure falls back to the next valid candidate. +- [x] Invalid input and unsupported operations do **not** trigger fallback. +- [x] Named routing is strict by default; `allow_named` opts into fallback. +- [x] Per-model concurrency admission via semaphores with a queue timeout. +- [x] Ten concurrent requests respect the model's concurrency limit (no over-admission). +- [x] Circuit breaker admits exactly one half-open probe. +- [x] Overall, queue, load, inference, discovery, and health timeouts are configurable. + +## Stage 4 — Memory Manager and Model Lifecycle + +- [x] Atomic reservation before loading (no time-of-check/time-of-use overcommit). +- [x] Reserved vs observed memory tracked separately; reconciled after load. +- [x] Active leases prevent eviction of in-flight models. +- [x] LRU eviction with incoming/subscription protection. +- [x] Reservation rolled back on load failure or timeout. +- [x] Supervised idle-timeout task unloads stale models and releases memory. +- [x] Lifecycle history and current allocations exposed via `/v1/lifecycle`, `/v1/memory`. +- [x] Memory pressure returns a structured error when nothing can be freed. + +## Stage 5 — Preflight v1 + +- [x] Explicit hints warm the next model concurrently with the current request. +- [x] Warm tasks are deduplicated and cancellable; speculative loads go through + normal memory admission. +- [x] Capability subscriptions (app/session-owned, TTL or explicit removal) warm + and protect their models. +- [x] History is keyed by app/session with a global fallback; one app's history + does not pollute another's predictions. +- [x] Predictions are gated by minimum samples and a confidence threshold. +- [x] Memory-unsafe predictions do not load a model. +- [x] Preflight effectiveness is visible via events and `/v1/preflight`. +- [x] History learning and speculative warming can be disabled by config. + +## MVP Acceptance Still Pending (needs live backends / later stages) + +- [ ] Auto-discover local MLX models. *(Stage 6.)* +- [x] Engine routes type-level LLM calls to a local model. *(Logic + tests done; + live demo needs a running Ollama with Qwen.)* +- [x] Named request routes to the cloud model exactly. *(Logic + tests done; + live demo needs a configured cloud key.)* +- [ ] Local TTS returns a managed audio artifact. *(Cloud TTS works today; local + MLX/Kokoro is Stage 6.)* +- [x] Hint Preflight starts the next model loading before the current completes. +- [x] Idle models unload and release accounted memory. +- [x] A retryable failure follows the documented fallback policy (engine logic); + local-TTS-to-cloud demo needs the Stage 6 local backend. +- [ ] Python UniFFI interface completes equivalent flows. *(Stage 7; HTTP done.)* +- [ ] mofa-fm completes English article to Chinese podcast end to end. *(Needs + live Ollama + TTS; engine-side flow is in place.)* diff --git a/docs/decisions/0002-stage-3-5-routing-memory-preflight.md b/docs/decisions/0002-stage-3-5-routing-memory-preflight.md new file mode 100644 index 0000000..892ca54 --- /dev/null +++ b/docs/decisions/0002-stage-3-5-routing-memory-preflight.md @@ -0,0 +1,80 @@ +# ADR 0002: Stage 3-5 — Routing, Memory Lifecycle, and Preflight + +## Status + +Accepted for the Stage 3-5 implementation. + +## Context + +ADR 0001 froze the kernel contracts and registry. Stages 3-5 turn that baseline +into a system that selects models explainably, manages constrained memory +safely, and reduces cold-start latency without destabilising scheduling. The +prototype had a single selection path, no real concurrency control, no memory +reservation, no idle eviction, and a global (cross-application) preflight chain. + +## Decisions + +### Routing & failure handling (Stage 3) + +- The router produces one **ranked candidate plan** (`route_ranked`). The engine + walks it in order; the first entry is the primary, each subsequent one a + failover. Primary and failover therefore never diverge. +- A **static memory-feasibility** filter removes local models whose estimate + exceeds the whole budget; remote (cloud) models are exempt. +- Fallback advances **only on retryable errors**. `InvalidRequest` and + `UnsupportedOperation` fail immediately. This satisfies "invalid input does + not trigger fallback." +- **Named routing is strict by default** (only the named model is attempted). + `fallback_policy = allow_named` appends capability candidates after the named + model. Capability requests fail over across all candidates unless + `fallback_policy = disabled`. +- **Concurrency admission** is a per-model `tokio::Semaphore` sized by the + model's `max_concurrency`, acquired with a configurable queue timeout. Busy + models are no longer hard-filtered from routing; they queue. A queue-wait + timeout is retryable, so it can fail over to another candidate. +- The circuit breaker admits **exactly one half-open probe**; concurrent callers + fail fast until the probe's outcome is recorded. +- All phase timeouts (overall request, queue, load, inference, discovery, + health) are configurable under `[timeouts]`. + +### Memory & lifecycle (Stage 4) + +- Memory is **reserved atomically before loading**. The evict-then-reserve + critical section runs under a `load_gate`; the slow backend load runs outside + it so independent models can load in parallel without overcommitting. +- The manager tracks **reserved** (accounting) and **observed** (reported by the + backend) bytes separately, reconciling reserved upward to observed. +- **Leases** protect in-flight models from eviction; eviction and the idle sweep + both skip leased and subscription-protected models. +- A **supervised idle-timeout task** (holding a `Weak`, aborted on drop) + unloads models idle past `memory.idle_timeout_secs` using monotonic time. +- Failed/timed-out loads **roll back** the reservation. When nothing can be + freed, admission returns a structured `MemoryPressure` error. +- Lifecycle history (bounded ring buffer) and current allocations are exposed via + `/v1/lifecycle` and `/v1/memory`. + +### Preflight (Stage 5) + +- Signal priority follows the RFC feedback: **hint → subscription → history**. +- History is a **per-scope Markov chain** keyed by `session_id`, else `app_id`, + else a shared global scope. A scope with too little data falls back to global; + one application's ordering never leaks into another's predictions. +- Predictions are gated by `min_samples` and `confidence_threshold`. +- Warm tasks are **deduplicated per model** and **cancellable**, and route + speculative loads through the same reservation/eviction admission as real + requests, so a memory-unsafe prediction simply fails to load. +- Subscriptions are app/session-owned with a TTL or explicit removal; subscribed + resident models are eviction-protected. +- Effectiveness counters (warms, predictions, hits, misses) are exposed via + `/v1/preflight`. + +## Consequences + +- Subscription protection can, by design, make memory pressure unsatisfiable + rather than break a keep-alive policy; the engine returns `MemoryPressure` in + that case. Revisit if subscription/pressure conflict needs a softer policy + (RFC open question 8). +- History and lifecycle state are in-memory only; persistence across restarts is + deferred (RFC open question 11) behind clean seams. +- The local TTS backend (Stage 6) and Python UniFFI bindings (Stage 7) remain + out of scope; cloud TTS and the HTTP API cover the demo path meanwhile. diff --git a/docs/openapi.yaml b/docs/openapi.yaml index d42476d..78b2abd 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -2,7 +2,10 @@ openapi: 3.1.0 info: title: MoFA Engine API version: 0.1.0 - description: Compatibility API for the Stage 0-2 engine baseline. + description: >- + MoFA Engine HTTP API. Covers the Stage 0-2 baseline plus the Stage 3-5 + additions: deterministic routing with failover, reservation-based memory and + lifecycle, and Preflight (hint/subscription/history) warming. paths: /health: get: @@ -74,6 +77,185 @@ paths: application/json: schema: $ref: "#/components/schemas/ErrorInfo" + /v1/invoke/stream: + post: + summary: Invoke a model and stream typed output as Server-Sent Events + description: > + Streams the response as SSE. Each `data:` line is a JSON StreamChunk: a + `started` event, then zero or more `text` deltas, then a terminal + `completed` or `error`. Errors after the stream opens are delivered + in-band as an `error` chunk (HTTP status is 200 once streaming begins), + so clients handle success and failure through one channel. The interface + is versioned and supports a non-streaming compatibility mode, in which a + backend emits `started`, a single full-text `text` chunk, then + `completed`. + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/InferenceRequest" + responses: + "200": + description: A Server-Sent Events stream of StreamChunk objects. + content: + text/event-stream: + schema: + $ref: "#/components/schemas/StreamChunk" + /v1/memory: + get: + summary: Current memory accounting (budget, reservations, leases) + responses: + "200": + description: Memory report. + content: + application/json: + schema: + $ref: "#/components/schemas/MemoryReport" + /v1/lifecycle: + get: + summary: Rolling model lifecycle history (load/unload/evict) + responses: + "200": + description: Lifecycle records, oldest first. + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/LifecycleRecord" + /v1/preflight: + get: + summary: Preflight effectiveness counters + responses: + "200": + description: Preflight statistics. + content: + application/json: + schema: + $ref: "#/components/schemas/PreflightStats" + /v1/subscriptions: + get: + summary: List active capability subscriptions + responses: + "200": + description: Active subscriptions. + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/SubscriptionInfo" + post: + summary: Create a capability subscription (kept warm + eviction-protected) + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/SubscribeRequest" + responses: + "200": + description: Subscription created. + content: + application/json: + schema: + type: object + properties: + id: + type: integer + "400": + description: Empty or invalid capability list. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + /v1/subscriptions/{id}: + delete: + summary: Remove a capability subscription + parameters: + - name: id + in: path + required: true + schema: + type: integer + responses: + "200": + description: Subscription removed. + content: + application/json: + schema: + type: object + properties: + removed: + type: boolean + "404": + description: No such subscription. + content: + application/json: + schema: + type: object + properties: + removed: + type: boolean + /v1/models/load: + post: + summary: Manually load (warm) a model + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ModelActionRequest" + responses: + "200": + description: Model loaded. + content: + application/json: + schema: + $ref: "#/components/schemas/ModelActionResponse" + "400": + description: Unknown or busy model. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + /v1/models/unload: + post: + summary: Manually unload a model + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ModelActionRequest" + responses: + "200": + description: Model unloaded (or was not known). + content: + application/json: + schema: + $ref: "#/components/schemas/ModelActionResponse" + "400": + description: Model is busy and cannot be unloaded. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + /metrics: + get: + summary: Prometheus text-exposition metrics + description: > + Bounded-cardinality engine metrics (requests, failures, fallback, + load/unload/eviction counters, a request-latency histogram, and memory / + model / provider gauges). Public — not behind the `/v1` auth gate. + responses: + "200": + description: Prometheus exposition text. + content: + text/plain: + schema: + type: string components: schemas: InferenceRequest: @@ -135,6 +317,55 @@ components: type: boolean routing_reason: type: [string, "null"] + StreamChunk: + description: > + One event in a streaming response, discriminated by `type`. Serialized + as internally-tagged JSON. + oneOf: + - type: object + required: [type, request_id, model_used, provider] + properties: + type: { const: started } + request_id: { type: string } + model_used: { type: string } + provider: { type: string } + - type: object + required: [type, delta] + properties: + type: { const: text } + delta: { type: string } + - type: object + required: [type, duration_ms, fallback_used] + properties: + type: { const: completed } + duration_ms: { type: integer } + tokens_used: { type: [integer, "null"] } + file: { type: [string, "null"] } + fallback_used: { type: boolean } + routing_reason: { type: [string, "null"] } + - type: object + required: [type, code, message, retryable] + properties: + type: { const: error } + code: { type: string } + message: { type: string } + retryable: { type: boolean } + source: { type: [string, "null"] } + ModelActionRequest: + type: object + required: [model_id] + properties: + model_id: + type: string + description: Canonical model id (`provider/model`). + ModelActionResponse: + type: object + required: [model_id, changed] + properties: + model_id: + type: string + changed: + type: boolean ErrorInfo: type: object required: [code, message, retryable] @@ -147,3 +378,91 @@ components: type: boolean source: type: string + AllocationSnapshot: + type: object + properties: + model_id: + type: string + reserved_bytes: + type: integer + observed_bytes: + type: [integer, "null"] + leases: + type: integer + MemoryReport: + type: object + properties: + used_bytes: + type: integer + budget_bytes: + type: integer + available_bytes: + type: integer + allocations: + type: array + items: + $ref: "#/components/schemas/AllocationSnapshot" + LifecycleRecord: + type: object + properties: + seq: + type: integer + at_ms: + type: integer + description: Milliseconds since engine start (monotonic). + model_id: + type: string + event: + type: string + description: load, unload, evict, idle_unload, load_failed, etc. + detail: + type: [string, "null"] + PreflightStats: + type: object + properties: + warms_started: + type: integer + warms_completed: + type: integer + warms_failed: + type: integer + warms_skipped: + type: integer + predictions: + type: integer + hits: + type: integer + misses: + type: integer + SubscriptionInfo: + type: object + properties: + id: + type: integer + app_id: + type: [string, "null"] + session_id: + type: [string, "null"] + capabilities: + type: array + items: + type: string + enum: [chat, tts, asr, image_gen, video_gen, vlm, embedding] + expires_in_secs: + type: [integer, "null"] + SubscribeRequest: + type: object + required: [capabilities] + properties: + app_id: + type: string + session_id: + type: string + capabilities: + type: array + items: + type: string + enum: [chat, tts, asr, image_gen, video_gen, vlm, embedding] + ttl_secs: + type: integer + description: Optional lifetime in seconds; omitted means until removed. diff --git a/mofa-engine-app/Cargo.toml b/mofa-engine-app/Cargo.toml index 6e7bc16..66ad1a9 100644 --- a/mofa-engine-app/Cargo.toml +++ b/mofa-engine-app/Cargo.toml @@ -6,9 +6,12 @@ license.workspace = true description = "MoFA Engine binary entry point" [dependencies] +mofa-kernel = { workspace = true } mofa-engine-core = { workspace = true } mofa-engine-sdk = { workspace = true } tokio = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } clap = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/mofa-engine-app/src/main.rs b/mofa-engine-app/src/main.rs index 2f163c8..891be7b 100644 --- a/mofa-engine-app/src/main.rs +++ b/mofa-engine-app/src/main.rs @@ -1,59 +1,189 @@ //! MoFA Engine — multimodal AI model orchestration engine. //! -//! Binary entry point. Parses CLI arguments, loads configuration, -//! initialises the engine, and starts the HTTP server. +//! Binary entry point. Runs the HTTP server by default, and offers CLI +//! subcommands (`status`, `capabilities`, `invoke`, `refresh`, +//! `validate-config`) that talk to a running daemon over the `/v1` API. -use clap::Parser; +use clap::{Parser, Subcommand}; use mofa_engine_core::{Engine, EngineConfig}; -use mofa_engine_sdk::start_server; +use mofa_engine_sdk::{DaemonClient, start_server}; +use mofa_kernel::{Capability, InferenceRequest, Message}; use std::path::PathBuf; -use tracing_subscriber::{EnvFilter, fmt}; +use tracing_subscriber::{EnvFilter, prelude::*}; /// MoFA Engine — multimodal AI model orchestration #[derive(Parser, Debug)] #[command(name = "mofa-engine", version, about)] struct Cli { /// Path to config.toml (default: auto-detect) - #[arg(short, long)] + #[arg(short, long, global = true)] config: Option, - /// Override listen port + /// Override listen port (serve mode) #[arg(short, long)] port: Option, + + #[command(subcommand)] + command: Option, +} + +/// Base URL argument shared by the daemon-facing subcommands. +#[derive(clap::Args, Debug)] +struct DaemonArgs { + /// Base URL of the running engine daemon. + #[arg(long, default_value = "http://127.0.0.1:8420")] + url: String, + /// Bearer token for a secured daemon (defaults to $MOFA_API_TOKEN). + #[arg(long)] + token: Option, +} + +/// Build a daemon client from the shared args, taking the bearer token from +/// `--token` or, failing that, `$MOFA_API_TOKEN`. +fn daemon_client(args: DaemonArgs) -> DaemonClient { + let token = args + .token + .or_else(|| std::env::var("MOFA_API_TOKEN").ok()) + .filter(|t| !t.is_empty()); + let client = DaemonClient::new(args.url); + match token { + Some(t) => client.with_token(t), + None => client, + } +} + +#[derive(Subcommand, Debug)] +enum Command { + /// Run the HTTP server (default). + Serve, + /// Validate the configuration and exit. + ValidateConfig, + /// List a running daemon's capabilities. + Capabilities(DaemonArgs), + /// Show a running daemon's status. + Status(DaemonArgs), + /// Re-run discovery on a running daemon. + Refresh(DaemonArgs), + /// Invoke a model on a running daemon. + Invoke { + #[command(flatten)] + daemon: DaemonArgs, + /// Capability to request (e.g. chat, tts). + #[arg(long)] + capability: Option, + /// Specific model name to request. + #[arg(long)] + model: Option, + /// Text input. + #[arg(long)] + text: String, + }, +} + +/// Initialise tracing. Set `MOFA_LOG_FORMAT=json` for structured JSON logs +/// (suitable for aggregation); otherwise human-readable text. Levels honour +/// `RUST_LOG`. +fn init_tracing() { + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); + let json = std::env::var("MOFA_LOG_FORMAT") + .map(|v| v.eq_ignore_ascii_case("json")) + .unwrap_or(false); + + let registry = tracing_subscriber::registry().with(filter); + if json { + registry + .with( + tracing_subscriber::fmt::layer() + .json() + .with_current_span(true) + .with_span_list(false), + ) + .init(); + } else { + registry + .with(tracing_subscriber::fmt::layer().with_target(false)) + .init(); + } } #[tokio::main] async fn main() -> anyhow::Result<()> { - // Initialise tracing - fmt() - .with_env_filter( - EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")), - ) - .with_target(false) - .init(); - + init_tracing(); let cli = Cli::parse(); - // Load and validate configuration. - let mut config = EngineConfig::load_checked(cli.config.as_deref())?; + match cli.command { + None | Some(Command::Serve) => serve(cli.config, cli.port).await, + Some(Command::ValidateConfig) => validate_config(cli.config), + Some(Command::Capabilities(d)) => print_json(daemon_client(d).capabilities().await), + Some(Command::Status(d)) => print_json(daemon_client(d).status().await), + Some(Command::Refresh(d)) => print_json(daemon_client(d).refresh().await), + Some(Command::Invoke { + daemon, + capability, + model, + text, + }) => { + let request = InferenceRequest { + capability: capability.as_deref().and_then(Capability::from_str_loose), + model, + app_id: None, + session_id: None, + fallback_policy: Default::default(), + messages: vec![Message { + role: "user".into(), + content: text, + }], + input_file: None, + params: serde_json::Value::Null, + hint_next: None, + request_id: String::new(), + }; + print_json(daemon_client(daemon).invoke(&request).await) + } + } +} - // CLI port override - if let Some(port) = cli.port { +/// Run the daemon. +async fn serve(config_path: Option, port: Option) -> anyhow::Result<()> { + let mut config = EngineConfig::load_checked(config_path.as_deref())?; + if let Some(port) = port { config.listen.port = port; } - let host = config.listen.host.clone(); let port = config.listen.port; tracing::info!("MoFA Engine v{} starting", env!("CARGO_PKG_VERSION")); - - // Create engine let engine = Engine::try_new(config).await?; - - // Start server start_server(engine, &host, port) .await - .map_err(|e| anyhow::anyhow!("{e}"))?; + .map_err(|e| anyhow::anyhow!("{e}")) +} + +/// Validate configuration without starting the engine. +fn validate_config(config_path: Option) -> anyhow::Result<()> { + match EngineConfig::load_checked(config_path.as_deref()) { + Ok(cfg) => { + println!( + "configuration is valid: {} provider(s), listen {}:{}", + cfg.providers.len(), + cfg.listen.host, + cfg.listen.port + ); + Ok(()) + } + Err(e) => Err(anyhow::anyhow!("invalid configuration: {e}")), + } +} - Ok(()) +/// Pretty-print a client result as JSON, or fail with the error. +fn print_json( + result: Result, +) -> anyhow::Result<()> { + match result { + Ok(value) => { + println!("{}", serde_json::to_string_pretty(&value)?); + Ok(()) + } + Err(e) => Err(anyhow::anyhow!("{e}")), + } } diff --git a/mofa-engine-core/src/artifacts.rs b/mofa-engine-core/src/artifacts.rs new file mode 100644 index 0000000..84a64ab --- /dev/null +++ b/mofa-engine-core/src/artifacts.rs @@ -0,0 +1,115 @@ +//! Retention and cleanup for engine-generated artifacts. +//! +//! Backends that produce files (TTS audio, and later image/video output) write +//! them with a `mofa_` name prefix into a shared directory. Left alone these +//! accumulate, so the engine runs a periodic sweep that deletes artifacts older +//! than a configured retention. Only files the engine created (matched by +//! prefix) are ever removed, so pointing the sweeper at a shared temp dir cannot +//! touch unrelated files. + +use std::path::PathBuf; +use std::time::{Duration, SystemTime}; + +/// Name prefix stamped on every engine-generated artifact. +pub const ARTIFACT_PREFIX: &str = "mofa_"; + +/// Deletes stale engine artifacts from a directory. +#[derive(Debug, Clone)] +pub struct ArtifactSweeper { + dir: PathBuf, + retention: Duration, +} + +impl ArtifactSweeper { + /// Create a sweeper for `dir` (defaulting to the system temp dir) that + /// removes engine artifacts older than `retention`. + pub fn new(dir: Option, retention: Duration) -> Self { + Self { + dir: dir.unwrap_or_else(std::env::temp_dir), + retention, + } + } + + /// The directory this sweeper scans. + pub fn dir(&self) -> &PathBuf { + &self.dir + } + + /// The retention window; artifacts at least this old are removed. + pub fn retention(&self) -> Duration { + self.retention + } + + /// Delete engine artifacts older than the retention. Returns how many files + /// were removed. Non-engine files and unreadable entries are left untouched. + pub fn sweep(&self) -> usize { + let now = SystemTime::now(); + let Ok(entries) = std::fs::read_dir(&self.dir) else { + return 0; + }; + let mut removed = 0; + for entry in entries.flatten() { + let name = entry.file_name(); + let Some(name) = name.to_str() else { continue }; + if !name.starts_with(ARTIFACT_PREFIX) { + continue; + } + let is_stale = entry + .metadata() + .and_then(|m| m.modified()) + .map(|modified| { + now.duration_since(modified) + .map(|age| age >= self.retention) + .unwrap_or(false) + }) + .unwrap_or(false); + if is_stale && std::fs::remove_file(entry.path()).is_ok() { + removed += 1; + } + } + removed + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn sweeps_only_stale_engine_files() { + let dir = tempfile::tempdir().unwrap(); + let dir_path = dir.path().to_path_buf(); + + let artifact = dir_path.join("mofa_tts_abc.wav"); + let unrelated = dir_path.join("important.txt"); + std::fs::write(&artifact, b"audio").unwrap(); + std::fs::write(&unrelated, b"keep me").unwrap(); + + // Retention 0 → every engine artifact is stale and removed; the + // non-engine file is never eligible. + let sweeper = ArtifactSweeper::new(Some(dir_path.clone()), Duration::ZERO); + assert_eq!(sweeper.sweep(), 1); + assert!(!artifact.exists()); + assert!(unrelated.exists()); + } + + #[test] + fn keeps_fresh_artifacts() { + let dir = tempfile::tempdir().unwrap(); + let artifact = dir.path().join("mofa_tts_fresh.wav"); + std::fs::write(&artifact, b"audio").unwrap(); + + // A one-hour retention keeps a just-written file. + let sweeper = + ArtifactSweeper::new(Some(dir.path().to_path_buf()), Duration::from_secs(3600)); + assert_eq!(sweeper.sweep(), 0); + assert!(artifact.exists()); + } + + #[test] + fn missing_directory_is_a_noop() { + let sweeper = + ArtifactSweeper::new(Some(PathBuf::from("/nonexistent/mofa/dir")), Duration::ZERO); + assert_eq!(sweeper.sweep(), 0); + } +} diff --git a/mofa-engine-core/src/backends.rs b/mofa-engine-core/src/backends.rs index b4332fa..7783c49 100644 --- a/mofa-engine-core/src/backends.rs +++ b/mofa-engine-core/src/backends.rs @@ -1,7 +1,9 @@ //! Provider backend implementations. +pub mod local_tts; pub mod ollama; pub mod openai_compat; +pub use local_tts::LocalTtsProvider; pub use ollama::OllamaProvider; pub use openai_compat::OpenAiCompatProvider; diff --git a/mofa-engine-core/src/backends/local_tts.rs b/mofa-engine-core/src/backends/local_tts.rs new file mode 100644 index 0000000..4959866 --- /dev/null +++ b/mofa-engine-core/src/backends/local_tts.rs @@ -0,0 +1,479 @@ +//! Local TTS process-adapter backend. +//! +//! Runs a configured local command — an MLX/Kokoro or Piper-style TTS CLI — to +//! synthesize speech, returning the produced audio as a managed artifact. +//! +//! Runtime- and device-specific concerns stay behind this `Provider` boundary, +//! so the engine treats a local TTS model like any other backend: it is +//! discovered, admitted through memory reservation, warmed, idle-unloaded, and +//! can fail over to a cloud TTS backend when the local one is unavailable. +//! +//! ## Lifecycle model +//! +//! This adapter spawns the command once per synthesis (a cold, stateless +//! process). `load` performs a cheap readiness probe (the program resolves and +//! is executable) and reports the model as resident with its configured memory +//! estimate; that conservative reservation lets the engine's coexistence and +//! idle-eviction logic apply to local TTS exactly as it does to Ollama. A +//! long-running server variant can later make `load` start the process without +//! changing the engine contract. +//! +//! ## Cancellation +//! +//! Child processes are spawned with `kill_on_drop(true)`, so when the engine's +//! inference timeout fires and drops the invocation future, the underlying +//! synthesis process is terminated rather than leaked. + +use std::path::{Path, PathBuf}; + +use async_trait::async_trait; +use mofa_kernel::{ + BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelResidency, Provider, + ProviderKind, canonical_model_id, model_id_name, +}; +use tokio::process::Command; + +use crate::config::ModelDef; + +/// A process-adapter provider that shells out to a local TTS command. +pub struct LocalTtsProvider { + /// Display name. + name: String, + /// Program to execute per synthesis. + command: String, + /// Argument template with `{text}`, `{text_file}`, and `{output}` placeholders. + args: Vec, + /// Output audio extension/container (e.g. `wav`, `mp3`). + output_format: String, + /// Directory for generated artifacts. + output_dir: PathBuf, + /// Configured models this backend serves. + models: Vec, +} + +impl LocalTtsProvider { + /// Create a new local TTS process adapter. + pub fn new( + name: impl Into, + command: impl Into, + args: Vec, + output_format: Option, + output_dir: Option, + models: Vec, + ) -> Self { + Self { + name: name.into(), + command: command.into(), + args, + output_format: output_format + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "wav".into()), + output_dir: output_dir + .filter(|s| !s.is_empty()) + .map(PathBuf::from) + .unwrap_or_else(std::env::temp_dir), + models, + } + } + + /// Resolve the configured program to an executable path, searching `PATH` + /// for a bare name. Returns `None` when it cannot be found. + fn resolve_program(&self) -> Option { + let program = Path::new(&self.command); + if program.is_absolute() || self.command.contains(std::path::MAIN_SEPARATOR) { + return program.is_file().then(|| program.to_path_buf()); + } + let paths = std::env::var_os("PATH")?; + std::env::split_paths(&paths) + .map(|dir| dir.join(&self.command)) + .find(|candidate| candidate.is_file()) + } + + /// Whether a configured model serves the given capability. + fn model_supports(&self, model_name: &str, capability: Capability) -> bool { + self.models.iter().any(|m| { + m.name == model_name && Capability::from_str_loose(&m.capability) == Some(capability) + }) + } + + async fn synthesize( + &self, + model_name: &str, + request: &InferenceRequest, + start: std::time::Instant, + ) -> Result { + let text = request + .messages + .first() + .map(|m| m.content.clone()) + .or_else(|| { + request + .params + .get("input") + .and_then(|v| v.as_str()) + .map(String::from) + }) + .filter(|t| !t.trim().is_empty()) + .ok_or_else(|| EngineError::InvalidRequest("TTS requires text input".into()))?; + + let output_path = self.output_dir.join(format!( + "mofa_tts_{}.{}", + uuid::Uuid::new_v4(), + self.output_format + )); + + // Materialize a temp text file only if the command template references it. + let needs_text_file = self.args.iter().any(|a| a.contains("{text_file}")); + let text_file = if needs_text_file { + let path = self + .output_dir + .join(format!("mofa_tts_{}.txt", uuid::Uuid::new_v4())); + tokio::fs::write(&path, &text) + .await + .map_err(|e| EngineError::Internal(format!("cannot write TTS input file: {e}")))?; + Some(path) + } else { + None + }; + + let output_str = output_path.to_string_lossy().to_string(); + let text_file_str = text_file + .as_ref() + .map(|p| p.to_string_lossy().to_string()) + .unwrap_or_default(); + + let resolved_args: Vec = self + .args + .iter() + .map(|arg| { + arg.replace("{text}", &text) + .replace("{text_file}", &text_file_str) + .replace("{output}", &output_str) + }) + .collect(); + + let result = Command::new(&self.command) + .args(&resolved_args) + .kill_on_drop(true) + .output() + .await; + + // Best-effort cleanup of the transient input file regardless of outcome. + if let Some(path) = &text_file { + let _ = tokio::fs::remove_file(path).await; + } + + let output = result.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("failed to spawn '{}': {e}", self.command), + })?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + let _ = tokio::fs::remove_file(&output_path).await; + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: format!( + "TTS command exited with {}: {}", + output.status, + stderr.trim() + ), + }); + } + + // A successful exit must have produced a non-empty artifact. + match tokio::fs::metadata(&output_path).await { + Ok(meta) if meta.len() > 0 => {} + _ => { + // Remove any empty/partial file the command may have created. + let _ = tokio::fs::remove_file(&output_path).await; + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "TTS command produced no audio output".into(), + }); + } + } + + Ok(InferenceResponse { + text: None, + file: Some(output_str), + model_used: model_name.to_string(), + provider: self.name.clone(), + duration_ms: start.elapsed().as_millis() as u64, + request_id: request.request_id.clone(), + tokens_used: None, + fallback_used: false, + routing_reason: None, + }) + } +} + +#[async_trait] +impl Provider for LocalTtsProvider { + fn name(&self) -> &str { + &self.name + } + + fn kind(&self) -> ProviderKind { + ProviderKind::LocalTts + } + + fn features(&self) -> Vec { + vec![ + BackendFeature::Discovery, + BackendFeature::Load, + BackendFeature::Unload, + BackendFeature::MemoryReporting, + ] + } + + async fn discover(&self) -> Result, EngineError> { + let cards = self + .models + .iter() + .filter_map(|m| { + let cap = Capability::from_str_loose(&m.capability)?; + let mut card = + ModelCard::new(self.name.clone(), m.name.clone(), cap, CostTier::Free); + card.id = canonical_model_id(&self.name, &m.name); + card.availability = ModelAvailability::Configured; + card.residency = ModelResidency::Unloaded; + card.memory_estimate_bytes = m.memory_mb.unwrap_or(0) * 1024 * 1024; + card.refresh_status(); + Some(card) + }) + .collect(); + Ok(cards) + } + + async fn health(&self) -> Result { + if self.resolve_program().is_some() { + Ok(BackendHealth::Healthy) + } else { + Ok(BackendHealth::Unavailable) + } + } + + async fn load(&self, model_id: &str) -> Result { + // Readiness probe: the command must resolve before we claim residency. + if self.resolve_program().is_none() { + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("TTS command '{}' not found", self.command), + }); + } + let model_name = model_id_name(model_id); + let estimate = self + .models + .iter() + .find(|m| m.name == model_name) + .and_then(|m| m.memory_mb) + .map(|mb| mb * 1024 * 1024); + Ok(LifecycleResult { + model_id: canonical_model_id(&self.name, model_name), + residency: ModelResidency::Loaded, + memory_bytes: estimate, + changed: true, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: canonical_model_id(&self.name, model_id_name(model_id)), + residency: ModelResidency::Unloaded, + memory_bytes: Some(0), + changed: true, + }) + } + + async fn invoke( + &self, + model_id: &str, + request: &InferenceRequest, + ) -> Result { + let model_name = model_id_name(model_id); + let capability = request.capability.unwrap_or(Capability::Tts); + + if capability != Capability::Tts { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' only supports tts, not {capability}", + self.name + ))); + } + if !self.model_supports(model_name, Capability::Tts) { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' model '{model_name}' does not support tts", + self.name + ))); + } + + self.synthesize(model_name, request, std::time::Instant::now()) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mofa_kernel::Message; + + fn tts_models() -> Vec { + vec![ModelDef { + name: "fixture".into(), + capability: "tts".into(), + context_window: None, + memory_mb: Some(64), + }] + } + + fn tts_request(text: &str) -> InferenceRequest { + InferenceRequest { + capability: Some(Capability::Tts), + model: None, + app_id: None, + session_id: None, + fallback_policy: Default::default(), + messages: vec![Message { + role: "user".into(), + content: text.into(), + }], + input_file: None, + params: serde_json::Value::Null, + hint_next: None, + request_id: "test".into(), + } + } + + #[test] + fn metadata_and_defaults() { + let p = LocalTtsProvider::new("local-tts", "sh", vec![], None, None, tts_models()); + assert_eq!(p.kind(), ProviderKind::LocalTts); + assert_eq!(p.name(), "local-tts"); + assert_eq!(p.output_format, "wav"); + assert!(p.features().contains(&BackendFeature::Load)); + } + + #[tokio::test] + async fn discover_reports_configured_tts_model() { + let p = LocalTtsProvider::new("local-tts", "sh", vec![], None, None, tts_models()); + let cards = p.discover().await.unwrap(); + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "local-tts/fixture"); + assert_eq!(cards[0].capability, Capability::Tts); + assert_eq!(cards[0].residency, ModelResidency::Unloaded); + assert_eq!(cards[0].memory_estimate_bytes, 64 * 1024 * 1024); + } + + #[tokio::test] + async fn health_reflects_command_availability() { + // `sh` resolves on any POSIX PATH. + let ok = LocalTtsProvider::new("local-tts", "sh", vec![], None, None, tts_models()); + assert_eq!(ok.health().await.unwrap(), BackendHealth::Healthy); + + let missing = LocalTtsProvider::new( + "local-tts", + "definitely-not-a-real-binary-xyz", + vec![], + None, + None, + tts_models(), + ); + assert_eq!(missing.health().await.unwrap(), BackendHealth::Unavailable); + } + + #[tokio::test] + async fn synthesizes_audio_artifact_via_fixture_command() { + let dir = tempfile::tempdir().unwrap(); + // A deterministic "TTS" fixture: write bytes to the requested output path. + let p = LocalTtsProvider::new( + "local-tts", + "sh", + vec![ + "-c".into(), + "printf 'RIFF-fake-audio' > \"$1\"".into(), + "sh".into(), + "{output}".into(), + ], + Some("wav".into()), + Some(dir.path().to_string_lossy().to_string()), + tts_models(), + ); + + let resp = p + .invoke("local-tts/fixture", &tts_request("hello world")) + .await + .unwrap(); + let file = resp.file.expect("a TTS artifact path"); + let bytes = std::fs::read(&file).unwrap(); + assert_eq!(bytes, b"RIFF-fake-audio"); + assert!(file.ends_with(".wav")); + } + + #[tokio::test] + async fn passes_text_through_a_temp_file_when_templated() { + let dir = tempfile::tempdir().unwrap(); + // Copy the input text file to the output so we can assert it round-trips. + let p = LocalTtsProvider::new( + "local-tts", + "sh", + vec![ + "-c".into(), + "cat \"$1\" > \"$2\"".into(), + "sh".into(), + "{text_file}".into(), + "{output}".into(), + ], + Some("wav".into()), + Some(dir.path().to_string_lossy().to_string()), + tts_models(), + ); + + let resp = p + .invoke("local-tts/fixture", &tts_request("spoken words")) + .await + .unwrap(); + let produced = std::fs::read_to_string(resp.file.unwrap()).unwrap(); + assert_eq!(produced, "spoken words"); + } + + #[tokio::test] + async fn command_failure_is_a_retryable_provider_error() { + let dir = tempfile::tempdir().unwrap(); + let p = LocalTtsProvider::new( + "local-tts", + "sh", + vec!["-c".into(), "echo boom >&2; exit 3".into()], + Some("wav".into()), + Some(dir.path().to_string_lossy().to_string()), + tts_models(), + ); + + let err = p + .invoke("local-tts/fixture", &tts_request("hi")) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + // Retryable so the engine can fail over to a cloud TTS backend. + assert!(err.retryable()); + } + + #[tokio::test] + async fn empty_text_is_an_invalid_request() { + let p = LocalTtsProvider::new("local-tts", "sh", vec![], None, None, tts_models()); + let err = p + .invoke("local-tts/fixture", &tts_request(" ")) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + } + + #[tokio::test] + async fn rejects_unsupported_capability() { + let p = LocalTtsProvider::new("local-tts", "sh", vec![], None, None, tts_models()); + let mut req = tts_request("hi"); + req.capability = Some(Capability::Chat); + let err = p.invoke("local-tts/fixture", &req).await.unwrap_err(); + assert!(matches!(err, EngineError::UnsupportedOperation(_))); + } +} diff --git a/mofa-engine-core/src/backends/ollama.rs b/mofa-engine-core/src/backends/ollama.rs index aa0a42f..71e7115 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -25,19 +25,23 @@ pub struct OllamaProvider { impl OllamaProvider { /// Create a new Ollama provider. - pub fn new(name: impl Into, base_url: impl Into) -> Self { + /// + /// Fails (rather than panicking or silently dropping the configured + /// timeouts/`no_proxy`) if the system TLS/HTTP stack cannot build a client, + /// so the engine surfaces a clean startup error instead of crashing. + pub fn new(name: impl Into, base_url: impl Into) -> Result { let client = Client::builder() .no_proxy() .connect_timeout(Duration::from_secs(5)) .timeout(Duration::from_secs(180)) .build() - .unwrap_or_default(); + .map_err(|e| EngineError::Config(format!("failed to build Ollama HTTP client: {e}")))?; - Self { + Ok(Self { name: name.into(), base_url: base_url.into(), client, - } + }) } async fn loaded_models(&self) -> HashSet { @@ -362,7 +366,7 @@ mod tests { #[test] fn provider_metadata() { - let p = OllamaProvider::new("test-ollama", "http://localhost:11434"); + let p = OllamaProvider::new("test-ollama", "http://localhost:11434").unwrap(); assert_eq!(p.kind(), ProviderKind::Ollama); assert_eq!(p.name(), "test-ollama"); assert!(p.features().contains(&BackendFeature::Discovery)); diff --git a/mofa-engine-core/src/backends/openai_compat.rs b/mofa-engine-core/src/backends/openai_compat.rs index 60a1ba9..45713a4 100644 --- a/mofa-engine-core/src/backends/openai_compat.rs +++ b/mofa-engine-core/src/backends/openai_compat.rs @@ -26,6 +26,8 @@ pub struct OpenAiCompatProvider { models: Vec, /// Cost tier for all models from this provider. cost_tier: CostTier, + /// Directory for generated TTS artifacts. + output_dir: std::path::PathBuf, /// HTTP client. client: Client, } @@ -38,21 +40,45 @@ impl OpenAiCompatProvider { api_key: impl Into, models: Vec, cost_tier: CostTier, - ) -> Self { + ) -> Result { + Self::with_output_dir(name, base_url, api_key, models, cost_tier, None) + } + + /// Create a provider, writing TTS artifacts into `output_dir` (or the system + /// temp dir when `None`) so they land where the artifact sweeper looks. + /// + /// Fails (rather than panicking or silently dropping the configured + /// timeouts) if the system TLS/HTTP stack cannot build a client. + pub fn with_output_dir( + name: impl Into, + base_url: impl Into, + api_key: impl Into, + models: Vec, + cost_tier: CostTier, + output_dir: Option, + ) -> Result { let client = Client::builder() .connect_timeout(Duration::from_secs(10)) .timeout(Duration::from_secs(120)) .build() - .unwrap_or_default(); + .map_err(|e| { + EngineError::Config(format!( + "failed to build OpenAI-compatible HTTP client: {e}" + )) + })?; - Self { + Ok(Self { name: name.into(), base_url: base_url.into(), api_key: api_key.into(), models, cost_tier, + output_dir: output_dir + .filter(|s| !s.is_empty()) + .map(std::path::PathBuf::from) + .unwrap_or_else(std::env::temp_dir), client, - } + }) } } @@ -385,8 +411,11 @@ impl OpenAiCompatProvider { detail: format!("TTS read error: {e}"), })?; - let path = std::env::temp_dir().join(format!("mofa_tts_{}.mp3", uuid::Uuid::new_v4())); - std::fs::write(&path, &bytes) + let path = self + .output_dir + .join(format!("mofa_tts_{}.mp3", uuid::Uuid::new_v4())); + tokio::fs::write(&path, &bytes) + .await .map_err(|e| EngineError::Internal(format!("write error: {e}")))?; let path = path.to_string_lossy().to_string(); @@ -415,7 +444,7 @@ impl OpenAiCompatProvider { .as_deref() .ok_or_else(|| EngineError::InvalidRequest("ASR requires input_file".into()))?; - let file_bytes = std::fs::read(file_path).map_err(|e| { + let file_bytes = tokio::fs::read(file_path).await.map_err(|e| { EngineError::InvalidRequest(format!("cannot read file '{file_path}': {e}")) })?; @@ -505,7 +534,8 @@ mod tests { }, ], CostTier::Medium, - ); + ) + .unwrap(); let cards = provider.discover().await.unwrap(); assert_eq!(cards.len(), 2); @@ -518,7 +548,8 @@ mod tests { #[test] fn kind_is_openai_compat() { - let p = OpenAiCompatProvider::new("x", "https://example.com", "key", vec![], CostTier::Low); + let p = OpenAiCompatProvider::new("x", "https://example.com", "key", vec![], CostTier::Low) + .unwrap(); assert_eq!(p.kind(), ProviderKind::OpenAiCompatible); assert_eq!(p.name(), "x"); } diff --git a/mofa-engine-core/src/circuit_breaker.rs b/mofa-engine-core/src/circuit_breaker.rs index d62f650..a2ec7c7 100644 --- a/mofa-engine-core/src/circuit_breaker.rs +++ b/mofa-engine-core/src/circuit_breaker.rs @@ -50,6 +50,12 @@ struct ProviderBreaker { state: CircuitState, failure_count: u32, last_failure: Option, + /// Whether a half-open probe has been admitted and is awaiting its outcome. + /// Guarantees that exactly one probe is in flight while half-open. + probe_in_flight: bool, + /// When the in-flight half-open probe was admitted, used to self-heal if the + /// caller never records the probe's outcome. + probe_started: Option, config: CircuitBreakerConfig, } @@ -59,27 +65,60 @@ impl ProviderBreaker { state: CircuitState::Closed, failure_count: 0, last_failure: None, + probe_in_flight: false, + probe_started: None, config, } } + /// How long to wait for a half-open probe to report an outcome before + /// assuming it was dropped and admitting a fresh one. Floored at 1s so a + /// zero cool-down (used in tests) still admits exactly one probe at a time. + fn probe_timeout(&self) -> Duration { + Duration::from_secs(self.config.cool_down_secs.max(1)) + } + /// Check if a request is allowed through. + /// + /// In `HalfOpen`, admits exactly one probe at a time: the first caller after + /// the cool-down wins, and concurrent callers fail fast until that probe's + /// outcome is recorded (which moves the circuit to `Closed` or `Open`). If an + /// admitted probe never records an outcome (e.g. its task was dropped), the + /// slot self-heals after `probe_timeout` so the provider is not black-holed + /// forever. fn allow_request(&mut self) -> bool { match self.state { CircuitState::Closed => true, CircuitState::Open => { - // Check if cool-down period has elapsed + // Check if cool-down period has elapsed. if let Some(last) = self.last_failure && last.elapsed() >= Duration::from_secs(self.config.cool_down_secs) { self.state = CircuitState::HalfOpen; - tracing::info!("circuit breaker → half_open (cool-down elapsed)"); + self.probe_in_flight = true; + self.probe_started = Some(Instant::now()); + tracing::info!("circuit breaker → half_open (admitting one probe)"); return true; } false } CircuitState::HalfOpen => { - // Allow exactly one probe request + if self.probe_in_flight { + // A probe is being tested; reject everyone else — unless the + // probe has gone silent past its timeout, in which case admit + // a fresh one so a dropped probe cannot wedge the circuit. + let stale = self + .probe_started + .map(|t| t.elapsed() >= self.probe_timeout()) + .unwrap_or(true); + if stale { + self.probe_started = Some(Instant::now()); + return true; + } + return false; + } + self.probe_in_flight = true; + self.probe_started = Some(Instant::now()); true } } @@ -87,8 +126,21 @@ impl ProviderBreaker { /// Record a successful request. fn record_success(&mut self) { - self.failure_count = 0; - self.state = CircuitState::Closed; + match self.state { + // The authorized probe succeeded → recover. + CircuitState::HalfOpen => { + self.state = CircuitState::Closed; + self.failure_count = 0; + self.probe_in_flight = false; + self.probe_started = None; + } + CircuitState::Closed => { + self.failure_count = 0; + } + // A late success from a request dispatched before the circuit opened + // must not reopen the gates; the cool-down governs recovery. + CircuitState::Open => {} + } } /// Record a failed request. @@ -104,12 +156,14 @@ impl ProviderBreaker { } } CircuitState::HalfOpen => { - // Probe failed — back to Open + // Probe failed — back to Open and release the probe slot. self.state = CircuitState::Open; + self.probe_in_flight = false; + self.probe_started = None; tracing::warn!("circuit breaker → open (probe failed)"); } CircuitState::Open => { - // Already open + // Already open. } } } @@ -228,6 +282,90 @@ mod tests { assert_eq!(reg.state("p"), CircuitState::Closed); } + #[test] + fn half_open_admits_exactly_one_probe() { + let cfg = CircuitBreakerConfig { + failure_threshold: 1, + cool_down_secs: 0, + }; + let reg = CircuitBreakerRegistry::new(cfg); + + reg.record_failure("p"); + assert_eq!(reg.state("p"), CircuitState::Open); + std::thread::sleep(Duration::from_millis(10)); + + // First caller after cool-down wins the single probe slot. + assert!(reg.allow_request("p")); + assert_eq!(reg.state("p"), CircuitState::HalfOpen); + + // Concurrent callers while the probe is in flight are rejected. + assert!(!reg.allow_request("p")); + assert!(!reg.allow_request("p")); + + // Once the probe succeeds, the circuit closes and admits traffic again. + reg.record_success("p"); + assert_eq!(reg.state("p"), CircuitState::Closed); + assert!(reg.allow_request("p")); + } + + #[test] + fn half_open_probe_slot_is_released_on_reopen() { + let cfg = CircuitBreakerConfig { + failure_threshold: 1, + cool_down_secs: 0, + }; + let reg = CircuitBreakerRegistry::new(cfg); + + reg.record_failure("p"); + std::thread::sleep(Duration::from_millis(10)); + assert!(reg.allow_request("p")); // → HalfOpen, probe in flight + assert!(!reg.allow_request("p")); // rejected while in flight + reg.record_failure("p"); // probe fails → Open, slot released + assert_eq!(reg.state("p"), CircuitState::Open); + + // After another cool-down, a fresh single probe is admitted. + std::thread::sleep(Duration::from_millis(10)); + assert!(reg.allow_request("p")); + assert!(!reg.allow_request("p")); + } + + #[test] + fn late_success_does_not_reopen_an_open_circuit() { + // A success recorded for a request that was in flight before the circuit + // opened must not flip it back to Closed. + let reg = CircuitBreakerRegistry::new(test_config()); + reg.record_failure("p"); + reg.record_failure("p"); + reg.record_failure("p"); + assert_eq!(reg.state("p"), CircuitState::Open); + + reg.record_success("p"); + assert_eq!(reg.state("p"), CircuitState::Open); + assert!(!reg.allow_request("p")); + } + + #[test] + fn stuck_half_open_probe_self_heals() { + // A probe that never records an outcome must not black-hole the provider. + let cfg = CircuitBreakerConfig { + failure_threshold: 1, + cool_down_secs: 1, // probe_timeout floors at 1s + }; + let reg = CircuitBreakerRegistry::new(cfg); + reg.record_failure("p"); + std::thread::sleep(Duration::from_millis(1100)); + + // First probe admitted, then it goes silent (no success/failure recorded). + assert!(reg.allow_request("p")); + assert_eq!(reg.state("p"), CircuitState::HalfOpen); + assert!(!reg.allow_request("p")); // still within probe timeout + + // After the probe timeout, a fresh probe is admitted rather than the + // circuit staying wedged forever. + std::thread::sleep(Duration::from_millis(1100)); + assert!(reg.allow_request("p")); + } + #[test] fn failure_in_half_open_reopens_circuit() { let cfg = CircuitBreakerConfig { diff --git a/mofa-engine-core/src/config.rs b/mofa-engine-core/src/config.rs index a31c82c..f39927c 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; /// Top-level engine configuration. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct EngineConfig { /// Network listen settings #[serde(default)] @@ -19,11 +19,58 @@ pub struct EngineConfig { /// Memory management settings #[serde(default)] pub memory: MemoryConfig, + /// Operation timeout settings + #[serde(default)] + pub timeouts: TimeoutConfig, + /// Preflight (predictive warming) settings + #[serde(default)] + pub preflight: PreflightConfig, + /// Generated-artifact retention settings + #[serde(default)] + pub artifacts: ArtifactConfig, + /// Security and file-access settings + #[serde(default)] + pub security: SecurityConfig, /// Provider definitions #[serde(default)] pub providers: Vec, } +/// Retention policy for engine-generated artifacts (e.g. TTS audio files). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ArtifactConfig { + /// Directory scanned for cleanup (default: the system temp dir). + #[serde(default)] + pub dir: Option, + /// Delete engine artifacts older than this many seconds. `0` disables the + /// background sweep (artifacts are then the caller's responsibility). + #[serde(default = "default_artifact_retention")] + pub retention_secs: u64, +} + +impl Default for ArtifactConfig { + fn default() -> Self { + Self { + dir: None, + retention_secs: default_artifact_retention(), + } + } +} + +fn default_artifact_retention() -> u64 { + 3600 +} + +/// Security and file-access constraints. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct SecurityConfig { + /// Allowlist of directory roots that request `input_file` paths must fall + /// under. Empty (the default) allows any local path — appropriate only for a + /// trusted single-user machine. + #[serde(default)] + pub input_roots: Vec, +} + /// Network listen configuration. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ListenConfig { @@ -75,6 +122,142 @@ fn default_idle_timeout() -> u64 { 120 } +/// Timeouts for the distinct phases of a request, all in seconds. +/// +/// These bound the time spent waiting in the admission queue, loading a model, +/// running inference, and the request as a whole, plus the per-provider +/// discovery and health probes. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TimeoutConfig { + /// Overall budget for a single `invoke`, spanning queueing, loading, and + /// inference across every candidate attempted. + #[serde(default = "default_request_timeout")] + pub request_secs: u64, + /// Maximum time a request may wait for a concurrency permit before failing. + #[serde(default = "default_queue_timeout")] + pub queue_secs: u64, + /// Maximum time to load/warm a single model. + #[serde(default = "default_load_timeout")] + pub load_secs: u64, + /// Maximum time for a single inference call to a provider. + #[serde(default = "default_inference_timeout")] + pub inference_secs: u64, + /// Maximum time for a single provider discovery probe. + #[serde(default = "default_discovery_timeout")] + pub discovery_secs: u64, + /// Maximum time for a single provider health probe. + #[serde(default = "default_health_timeout")] + pub health_secs: u64, +} + +impl Default for TimeoutConfig { + fn default() -> Self { + Self { + request_secs: default_request_timeout(), + queue_secs: default_queue_timeout(), + load_secs: default_load_timeout(), + inference_secs: default_inference_timeout(), + discovery_secs: default_discovery_timeout(), + health_secs: default_health_timeout(), + } + } +} + +impl TimeoutConfig { + /// Overall request budget as a `Duration`. + pub fn request(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.request_secs) + } + /// Queue-wait budget as a `Duration`. + pub fn queue(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.queue_secs) + } + /// Load budget as a `Duration`. + pub fn load(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.load_secs) + } + /// Inference budget as a `Duration`. + pub fn inference(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.inference_secs) + } + /// Discovery probe budget as a `Duration`. + pub fn discovery(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.discovery_secs) + } + /// Health probe budget as a `Duration`. + pub fn health(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.health_secs) + } +} + +fn default_request_timeout() -> u64 { + 240 +} +fn default_queue_timeout() -> u64 { + 30 +} +fn default_load_timeout() -> u64 { + 60 +} +fn default_inference_timeout() -> u64 { + 180 +} +fn default_discovery_timeout() -> u64 { + 8 +} +fn default_health_timeout() -> u64 { + 5 +} + +/// Preflight (predictive warming) configuration. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct PreflightConfig { + /// Master switch for all predictive warming. + #[serde(default = "default_true")] + pub enabled: bool, + /// Whether the engine learns capability transition history. + #[serde(default = "default_true")] + pub history_learning: bool, + /// Whether the engine speculatively warms models from hints/history/subscriptions. + #[serde(default = "default_true")] + pub speculative_warming: bool, + /// Minimum confidence (0.0–1.0) a history prediction needs before it warms a model. + #[serde(default = "default_confidence_threshold")] + pub confidence_threshold: f64, + /// Minimum number of observed transitions before history is trusted. + #[serde(default = "default_min_samples")] + pub min_samples: u64, + /// Default lifetime of a capability subscription, in seconds. + #[serde(default = "default_subscription_ttl")] + pub subscription_ttl_secs: u64, +} + +impl Default for PreflightConfig { + fn default() -> Self { + Self { + enabled: default_true(), + history_learning: default_true(), + speculative_warming: default_true(), + confidence_threshold: default_confidence_threshold(), + min_samples: default_min_samples(), + subscription_ttl_secs: default_subscription_ttl(), + } + } +} + +fn default_true() -> bool { + true +} +fn default_confidence_threshold() -> f64 { + 0.6 +} +fn default_min_samples() -> u64 { + 3 +} +fn default_subscription_ttl() -> u64 { + 3600 +} + /// Configuration for a single provider. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ProviderConfig { @@ -98,6 +281,43 @@ pub struct ProviderConfig { /// Whether this provider is enabled #[serde(default = "default_enabled")] pub enabled: bool, + /// Program to execute for a `local_tts` process-adapter backend. + /// + /// Only used when `kind = "local_tts"`. The engine runs this command once + /// per synthesis, substituting placeholders in `args`. + #[serde(default)] + pub command: Option, + /// Argument template for a `local_tts` command. Each element may contain the + /// placeholders `{text}` (input text), `{text_file}` (path to a temp file + /// holding the input text), and `{output}` (path the command must write + /// audio to). + #[serde(default)] + pub args: Vec, + /// Output audio container/extension for `local_tts` artifacts (default `wav`). + #[serde(default)] + pub output_format: Option, + /// Directory for `local_tts` audio artifacts (default: the system temp dir). + #[serde(default)] + pub output_dir: Option, +} + +impl Default for ProviderConfig { + fn default() -> Self { + Self { + name: String::new(), + kind: String::new(), + base_url: String::new(), + api_key: None, + priority: default_priority(), + cost_tier: default_cost_tier(), + models: Vec::new(), + enabled: default_enabled(), + command: None, + args: Vec::new(), + output_format: None, + output_dir: None, + } + } } fn default_priority() -> u8 { @@ -192,6 +412,22 @@ impl EngineConfig { /// Validate configuration before constructing the engine. pub fn validate(&self) -> Result<(), EngineError> { + if !(0.0..=1.0).contains(&self.preflight.confidence_threshold) { + return Err(EngineError::Config(format!( + "preflight.confidence_threshold must be within 0.0..=1.0, got {}", + self.preflight.confidence_threshold + ))); + } + // Each configured allowlist root must resolve. Otherwise a typo would be + // silently dropped, leaving an *empty* allowlist that permits every path + // — a security downgrade rather than the intended restriction. + for root in &self.security.input_roots { + if std::fs::canonicalize(root).is_err() { + return Err(EngineError::Config(format!( + "security.input_roots entry '{root}' does not exist or cannot be resolved" + ))); + } + } for provider in &self.providers { if !provider.enabled { continue; @@ -205,6 +441,20 @@ impl EngineConfig { provider.name ))); } + if provider.kind == "local_tts" { + if provider.command.as_deref().unwrap_or_default().is_empty() { + return Err(EngineError::Config(format!( + "provider '{}' (local_tts) requires a non-empty command", + provider.name + ))); + } + if provider.models.is_empty() { + return Err(EngineError::Config(format!( + "provider '{}' (local_tts) must declare at least one model", + provider.name + ))); + } + } for model in &provider.models { if mofa_kernel::Capability::from_str_loose(&model.capability).is_none() { return Err(EngineError::Config(format!( @@ -231,6 +481,7 @@ impl EngineConfig { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }); // OpenAI @@ -275,6 +526,7 @@ impl EngineConfig { }, ], enabled: true, + ..Default::default() }); } @@ -302,6 +554,7 @@ impl EngineConfig { }, ], enabled: true, + ..Default::default() }); } @@ -335,6 +588,7 @@ impl EngineConfig { }, ], enabled: true, + ..Default::default() }); } @@ -354,6 +608,7 @@ impl EngineConfig { memory_mb: None, }], enabled: true, + ..Default::default() }); } @@ -373,6 +628,7 @@ impl EngineConfig { memory_mb: None, }], enabled: true, + ..Default::default() }); } @@ -392,12 +648,17 @@ impl EngineConfig { memory_mb: None, }], enabled: true, + ..Default::default() }); } EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers, } } @@ -421,6 +682,7 @@ impl ProviderConfig { match self.kind.as_str() { "ollama" => Ok(ProviderKind::Ollama), "openai_compatible" => Ok(ProviderKind::OpenAiCompatible), + "local_tts" => Ok(ProviderKind::LocalTts), other => Err(EngineError::Config(format!( "unknown provider kind '{}' for provider '{}'", other, self.name @@ -472,6 +734,10 @@ mod tests { let cfg = EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers: vec![ProviderConfig { name: "test".into(), kind: "ollama".into(), @@ -481,6 +747,7 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }], }; let toml_str = toml::to_string_pretty(&cfg).unwrap(); @@ -489,11 +756,46 @@ mod tests { assert_eq!(parsed.providers[0].name, "test"); } + #[test] + fn example_config_is_valid() { + // The shipped example must always parse and validate so docs cannot drift. + let path = concat!(env!("CARGO_MANIFEST_DIR"), "/../config.example.toml"); + let content = std::fs::read_to_string(path).expect("config.example.toml should exist"); + let cfg: EngineConfig = toml::from_str(&content).expect("example config should parse"); + cfg.validate().expect("example config should validate"); + assert!(cfg.providers.iter().any(|p| p.name == "ollama")); + assert_eq!(cfg.timeouts.queue_secs, 30); + assert!(cfg.preflight.enabled); + } + + #[test] + fn timeout_and_preflight_defaults_apply_from_partial_toml() { + // A config that omits [timeouts] and [preflight] must still parse with defaults. + let cfg: EngineConfig = toml::from_str("[memory]\nbudget_mb = 1024\n").unwrap(); + assert_eq!(cfg.timeouts.queue_secs, 30); + assert_eq!(cfg.timeouts.inference_secs, 180); + assert!(cfg.preflight.enabled); + assert!(cfg.preflight.speculative_warming); + assert_eq!(cfg.preflight.min_samples, 3); + cfg.validate().unwrap(); + } + + #[test] + fn validate_rejects_out_of_range_confidence() { + let mut cfg = EngineConfig::from_env(); + cfg.preflight.confidence_threshold = 1.5; + assert!(cfg.validate().is_err()); + } + #[test] fn validate_rejects_unknown_provider_kind() { let cfg = EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers: vec![ProviderConfig { name: "bad".into(), kind: "bad_kind".into(), @@ -503,6 +805,7 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }], }; assert!(cfg.validate().is_err()); diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index e656d6c..ded37b3 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -1,26 +1,38 @@ //! The main engine orchestrator. //! //! `Engine` ties together providers, routing, discovery, health, lifecycle, -//! circuit breaking, memory accounting, and observability. +//! circuit breaking, reservation-based memory admission, concurrency control, +//! idle eviction, and observability. -use std::collections::HashSet; -use std::sync::Arc; +use std::collections::{HashSet, VecDeque}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock, Weak}; use std::time::{Duration, Instant}; use dashmap::DashMap; use mofa_kernel::{ - BackendHealth, BackendStatus, CostTier, EngineError, EngineEvent, EngineStatus, FallbackPolicy, - InferenceRequest, InferenceResponse, ModelCard, ModelResidency, Provider, ProviderHealth, - ProviderKind, + BackendHealth, BackendStatus, Capability, CostTier, EngineError, EngineEvent, EngineStatus, + FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, ModelResidency, Provider, + ProviderHealth, ProviderKind, StreamChunk, model_id_name, }; -use tokio::sync::broadcast; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex as AsyncMutex, Semaphore, broadcast, mpsc}; +use tokio::task::{AbortHandle, JoinHandle}; -use crate::backends::{OllamaProvider, OpenAiCompatProvider}; +use crate::backends::{LocalTtsProvider, OllamaProvider, OpenAiCompatProvider}; use crate::circuit_breaker::{CircuitBreakerConfig, CircuitBreakerRegistry, CircuitState}; -use crate::config::EngineConfig; -use crate::memory::MemoryManager; -use crate::preflight::PreflightPredictor; -use crate::router::{Router, RoutingProvider}; +use crate::config::{EngineConfig, PreflightConfig, TimeoutConfig}; +use crate::memory::{AllocationSnapshot, MemoryManager}; +use crate::metrics::{EngineMetrics, MetricsGauges}; +use crate::preflight::{GLOBAL_SCOPE, PreflightMetrics, PreflightPredictor, PreflightStats}; +use crate::router::{RouteDecision, Router, RoutingProvider}; +use crate::subscription::{SubscriptionInfo, SubscriptionRegistry}; + +/// Maximum number of lifecycle records retained in the rolling history. +const LIFECYCLE_CAPACITY: usize = 256; +/// Maximum number of in-flight predictions tracked for hit/miss accounting. +/// Bounds memory against an unbounded stream of unique scope identifiers. +const MAX_PENDING_PREDICTIONS: usize = 4096; #[derive(Clone)] struct RegisteredProvider { @@ -30,6 +42,34 @@ struct RegisteredProvider { provider: Arc, } +/// A single entry in the model lifecycle history. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LifecycleRecord { + /// Monotonic sequence number. + pub seq: u64, + /// Milliseconds since engine start (monotonic). + pub at_ms: u64, + /// Affected model. + pub model_id: String, + /// Event kind: `load`, `unload`, `evict`, `idle_unload`, `load_failed`, etc. + pub event: String, + /// Optional human-readable detail. + pub detail: Option, +} + +/// A snapshot of the engine's memory accounting. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemoryReport { + /// Bytes reserved across all models. + pub used_bytes: u64, + /// Total budget. + pub budget_bytes: u64, + /// Bytes still free within the budget. + pub available_bytes: u64, + /// Per-model allocations. + pub allocations: Vec, +} + /// The main MoFA Engine orchestrator. pub struct Engine { /// Named providers. @@ -38,18 +78,56 @@ pub struct Engine { models: DashMap, /// Latest backend health by provider. backend_health: DashMap, - /// Memory manager. + /// Reservation-based memory manager. memory: MemoryManager, /// Circuit breaker registry. circuit_breakers: CircuitBreakerRegistry, - /// Preflight predictor. + /// Preflight predictor (app/session-scoped transition history). preflight: PreflightPredictor, + /// Preflight effectiveness counters. + preflight_metrics: PreflightMetrics, + /// Active capability subscriptions. + subscriptions: SubscriptionRegistry, + /// In-flight speculative warm tasks, keyed by model id (dedup + cancel). + warming: DashMap, + /// Per-scope pending prediction awaiting confirmation by the next request. + pending_predictions: DashMap, + /// Operation timeouts. + timeouts: TimeoutConfig, + /// Preflight configuration. + preflight_config: PreflightConfig, + /// Idle timeout before a resident model is auto-unloaded. + idle_timeout: Duration, + /// Per-model concurrency admission. + semaphores: DashMap>, + /// Serializes the evict-then-reserve admission critical section. + load_gate: AsyncMutex<()>, + /// Rolling lifecycle history. + lifecycle: Mutex>, + /// Lifecycle sequence counter. + lifecycle_seq: AtomicU64, + /// Background idle-eviction task; aborted on drop. + idle_task: Mutex>>, + /// Background artifact-cleanup task; aborted on drop. + artifact_task: Mutex>>, + /// Weak self-reference for background tasks. + weak_self: OnceLock>, /// Event broadcast channel. event_tx: broadcast::Sender, + /// Bounded-cardinality metrics counters. + metrics: EngineMetrics, + /// Canonicalized allowlist of roots for request `input_file` paths; empty + /// means any local path is accepted. + input_roots: Vec, /// Engine start time. started_at: Instant, } +/// Remaining time until `deadline`, saturating at zero. +fn remaining(deadline: Instant) -> Duration { + deadline.saturating_duration_since(Instant::now()) +} + impl Engine { /// Create and initialize a new engine from configuration. /// @@ -69,6 +147,7 @@ impl Engine { let memory = MemoryManager::new(config.memory.budget_mb); let circuit_breakers = CircuitBreakerRegistry::new(CircuitBreakerConfig::default()); let preflight = PreflightPredictor::new(); + let idle_timeout = Duration::from_secs(config.memory.idle_timeout_secs); let mut providers = Vec::new(); for pc in &config.providers { @@ -80,14 +159,31 @@ impl Engine { let kind = pc.provider_kind()?; let cost_tier = CostTier::from_str_loose(&pc.cost_tier); let provider: Arc = match kind { - ProviderKind::Ollama => Arc::new(OllamaProvider::new(&pc.name, &pc.base_url)), - ProviderKind::OpenAiCompatible => Arc::new(OpenAiCompatProvider::new( + ProviderKind::Ollama => Arc::new(OllamaProvider::new(&pc.name, &pc.base_url)?), + ProviderKind::OpenAiCompatible => Arc::new(OpenAiCompatProvider::with_output_dir( &pc.name, &pc.base_url, pc.api_key.clone().unwrap_or_default(), pc.models.clone(), cost_tier, - )), + config.artifacts.dir.clone(), + )?), + ProviderKind::LocalTts => { + let command = pc.command.clone().ok_or_else(|| { + EngineError::Config(format!( + "provider '{}' (local_tts) requires a command", + pc.name + )) + })?; + Arc::new(LocalTtsProvider::new( + &pc.name, + command, + pc.args.clone(), + pc.output_format.clone(), + pc.output_dir.clone(), + pc.models.clone(), + )) + } _ => { return Err(EngineError::Config(format!( "provider '{}' uses unsupported provider kind", @@ -104,6 +200,19 @@ impl Engine { }); } + // Canonicalize the input-path allowlist up front; unresolvable roots are + // dropped so a typo cannot silently widen access. + let input_roots = config + .security + .input_roots + .iter() + .filter_map(|r| std::fs::canonicalize(r).ok()) + .collect::>(); + let artifact_sweeper = crate::artifacts::ArtifactSweeper::new( + config.artifacts.dir.clone().map(std::path::PathBuf::from), + Duration::from_secs(config.artifacts.retention_secs), + ); + let engine = Arc::new(Self { providers, models: DashMap::new(), @@ -111,11 +220,30 @@ impl Engine { memory, circuit_breakers, preflight, + preflight_metrics: PreflightMetrics::default(), + subscriptions: SubscriptionRegistry::new(), + warming: DashMap::new(), + pending_predictions: DashMap::new(), + timeouts: config.timeouts.clone(), + preflight_config: config.preflight.clone(), + idle_timeout, + semaphores: DashMap::new(), + load_gate: AsyncMutex::new(()), + lifecycle: Mutex::new(VecDeque::with_capacity(LIFECYCLE_CAPACITY)), + lifecycle_seq: AtomicU64::new(0), + idle_task: Mutex::new(None), + artifact_task: Mutex::new(None), + weak_self: OnceLock::new(), event_tx, + metrics: EngineMetrics::default(), + input_roots, started_at: Instant::now(), }); + let _ = engine.weak_self.set(Arc::downgrade(&engine)); engine.refresh_resources().await; + engine.spawn_idle_eviction(); + engine.spawn_artifact_sweep(artifact_sweeper); Ok(engine) } @@ -127,6 +255,7 @@ impl Engine { /// Discover models from all providers with bounded per-provider timeouts. async fn discover_all(&self) { + let discovery_timeout = self.timeouts.discovery(); let handles = self .providers .iter() @@ -134,7 +263,7 @@ impl Engine { let name = registered.name.clone(); let provider = Arc::clone(®istered.provider); tokio::spawn(async move { - let result = tokio::time::timeout(Duration::from_secs(8), provider.discover()) + let result = tokio::time::timeout(discovery_timeout, provider.discover()) .await .map_err(|_| EngineError::Timeout(format!("discovery timeout for {name}"))) .and_then(|inner| inner); @@ -163,19 +292,40 @@ impl Engine { for id in stale { self.models.remove(&id); self.memory.deallocate(&id); + self.semaphores.remove(&id); + self.warming.remove(&id); } if had_stale { - let _ = self.event_tx.send(EngineEvent::MemoryChanged { - used_bytes: self.memory.used_bytes(), - total_bytes: self.memory.budget_bytes(), - }); + self.emit_memory_changed(); } let count = cards.len(); + let mut freed = false; for mut card in cards { + // Reconcile with the engine's current view. `.map` drops the + // shard guard immediately so the later insert cannot deadlock. + let previous = self.models.get(&card.id).map(|c| c.residency); + match previous { + // A load is in flight; don't let discovery regress it. + Some(ModelResidency::Loading) => { + card.residency = ModelResidency::Loading; + } + // The backend reports the model is no longer resident, so + // release the reservation we still held for it. + Some(ModelResidency::Loaded) + if !matches!(card.residency, ModelResidency::Loaded) => + { + self.memory.deallocate(&card.id); + freed = true; + } + _ => {} + } card.refresh_status(); self.models.insert(card.id.clone(), card); } + if freed { + self.emit_memory_changed(); + } let _ = self.event_tx.send(EngineEvent::DiscoveryCompleted { provider: name.clone(), models: count, @@ -196,6 +346,7 @@ impl Engine { } async fn refresh_health(&self) { + let health_timeout = self.timeouts.health(); let handles = self .providers .iter() @@ -203,7 +354,7 @@ impl Engine { let name = registered.name.clone(); let provider = Arc::clone(®istered.provider); tokio::spawn(async move { - let health = tokio::time::timeout(Duration::from_secs(5), provider.health()) + let health = tokio::time::timeout(health_timeout, provider.health()) .await .map_err(|_| EngineError::Timeout(format!("health timeout for {name}"))) .and_then(|inner| inner) @@ -229,6 +380,11 @@ impl Engine { /// Return all known model cards. pub async fn capabilities(&self) -> Vec { + self.models_snapshot() + } + + /// Synchronous snapshot of all model cards, sorted by id. + fn models_snapshot(&self) -> Vec { let mut models = self .models .iter() @@ -238,53 +394,125 @@ impl Engine { models } - /// Run inference: route → circuit-check → load → invoke → optional fallback. + /// Run inference, recording request-level metrics around the attempt. pub async fn invoke(&self, req: InferenceRequest) -> Result { + let started = Instant::now(); + let result = self.invoke_inner(req).await; + let duration_ms = started.elapsed().as_millis() as u64; + match &result { + Ok(resp) => self + .metrics + .record_request(true, duration_ms, resp.fallback_used), + Err(_) => self.metrics.record_request(false, duration_ms, false), + } + result + } + + /// Run inference. + /// + /// Builds a single ranked candidate plan and walks it in order: the first + /// candidate is the primary selection, and each subsequent one is a failover. + /// Only *retryable* failures advance to the next candidate, so malformed or + /// unsupported requests fail immediately rather than masquerading as + /// transient errors. The candidate list itself encodes the fallback policy. + async fn invoke_inner(&self, req: InferenceRequest) -> Result { self.reject_ambiguous_short_name(&req)?; + self.check_input_path(&req)?; + let overall_deadline = Instant::now() + self.timeouts.request(); + + // Confirm the previous request's prediction for this scope before a new + // one is formed during this request's admission. + self.confirm_prediction(&req); let all_models = self.capabilities().await; let providers = self.routing_providers(); - let selected = Router::route(&all_models, &req, &providers).ok_or_else(|| { - let cap_str = req - .capability - .map(|c| c.to_string()) - .unwrap_or_else(|| "any".into()); - EngineError::NoCapableModel(cap_str) - })?; + let candidates = self.build_candidates(&all_models, &req, &providers); - let model_id = selected.model.id.clone(); - let provider_name = selected.model.provider.clone(); - let routing_reason = selected.reason.clone(); + if candidates.is_empty() { + return Err(EngineError::NoCapableModel(Self::requested_capability( + &req, + ))); + } - match self - .try_invoke(&model_id, &provider_name, &req, Some(routing_reason)) - .await - { - Ok(resp) => Ok(resp), - Err(primary_err) if self.can_fallback(&req) => { - tracing::warn!("primary model '{model_id}' failed: {primary_err}, trying fallback"); - let mut fallback_req = req.clone(); - fallback_req.model = None; - let fallback_models = self - .capabilities() - .await - .into_iter() - .filter(|m| m.id != model_id && m.provider != provider_name) - .collect::>(); - let Some(fallback) = Router::route(&fallback_models, &fallback_req, &providers) - else { - return Err(primary_err); - }; - let fb_id = fallback.model.id.clone(); - let fb_provider = fallback.model.provider.clone(); - let mut resp = self - .try_invoke(&fb_id, &fb_provider, &fallback_req, Some(fallback.reason)) - .await - .map_err(|_| primary_err)?; - resp.fallback_used = true; - Ok(resp) + let mut last_err: Option = None; + for (idx, decision) in candidates.iter().enumerate() { + if remaining(overall_deadline).is_zero() { + last_err = Some(EngineError::Timeout( + "overall request deadline exceeded".into(), + )); + break; + } + + let model_id = decision.model.id.clone(); + let provider_name = decision.model.provider.clone(); + let max_concurrency = decision.model.execution.max_concurrency; + + match self + .try_invoke( + &model_id, + &provider_name, + max_concurrency, + &req, + decision.reason.clone(), + overall_deadline, + ) + .await + { + Ok(mut resp) => { + resp.fallback_used = idx > 0; + return Ok(resp); + } + Err(e) => { + if !e.retryable() { + // Invalid request, unsupported operation, etc. — never fall over. + return Err(e); + } + tracing::warn!("candidate '{model_id}' failed (retryable): {e}"); + last_err = Some(e); + } + } + } + + Err(last_err + .unwrap_or_else(|| EngineError::NoCapableModel(Self::requested_capability(&req)))) + } + + fn requested_capability(req: &InferenceRequest) -> String { + req.capability + .map(|c| c.to_string()) + .unwrap_or_else(|| "any".into()) + } + + /// Build the ordered candidate plan, applying the request's fallback policy. + fn build_candidates<'a>( + &self, + all_models: &'a [ModelCard], + req: &InferenceRequest, + providers: &[RoutingProvider], + ) -> Vec> { + let budget = Some(self.memory.budget_bytes()); + let primary = Router::route_ranked(all_models, req, providers, budget); + + match (req.model.is_some(), req.fallback_policy) { + // Named request, strict (default) or fallback disabled: only the named model. + (true, FallbackPolicy::CapabilityOnly | FallbackPolicy::Disabled) => primary, + // Named request with explicit opt-in: the named model first, then any + // capability-compatible candidate. + (true, FallbackPolicy::AllowNamed) => { + let mut cands = primary; + let mut cap_req = req.clone(); + cap_req.model = None; + for d in Router::route_ranked(all_models, &cap_req, providers, budget) { + if !cands.iter().any(|c| c.model.id == d.model.id) { + cands.push(d); + } + } + cands } - Err(primary_err) => Err(primary_err), + // Capability request, fallback disabled: only the single best. + (false, FallbackPolicy::Disabled) => primary.into_iter().take(1).collect(), + // Capability request, default/allow: the full ranked list. + (false, _) => primary, } } @@ -317,22 +545,16 @@ impl Engine { Ok(()) } - fn can_fallback(&self, req: &InferenceRequest) -> bool { - match (req.model.is_some(), req.fallback_policy) { - (_, FallbackPolicy::Disabled) => false, - (true, FallbackPolicy::CapabilityOnly) => false, - (true, FallbackPolicy::AllowNamed) => true, - (false, _) => true, - } - } - - /// Attempt to invoke a specific model, handling circuit breaker and loading. + /// Attempt one candidate: circuit check → ensure loaded → admit concurrency + /// → invoke, with phase-bounded timeouts that respect the overall deadline. async fn try_invoke( &self, model_id: &str, provider_name: &str, + max_concurrency: u32, req: &InferenceRequest, - routing_reason: Option, + routing_reason: String, + overall_deadline: Instant, ) -> Result { if !self.circuit_breakers.allow_request(provider_name) { return Err(EngineError::CircuitOpen(provider_name.into())); @@ -349,267 +571,914 @@ impl Engine { }); if let Err(e) = self.ensure_loaded(model_id, provider_name, &provider).await { - let _ = self.event_tx.send(EngineEvent::RequestCompleted { - request_id: req.request_id.clone(), - duration_ms: 0, - success: false, - }); + self.emit_request_completed(&req.request_id, 0, false); return Err(e); } + + // Concurrency admission with a bounded queue wait. + let sem = self.semaphore_for(model_id, max_concurrency); + let queue_budget = remaining(overall_deadline).min(self.timeouts.queue()); + let permit = match tokio::time::timeout(queue_budget, sem.acquire_owned()).await { + Ok(Ok(p)) => p, + Ok(Err(_)) => { + self.emit_request_completed(&req.request_id, 0, false); + return Err(EngineError::Internal("concurrency semaphore closed".into())); + } + Err(_) => { + self.emit_request_completed(&req.request_id, 0, false); + return Err(EngineError::Timeout(format!( + "queue wait exceeded for '{model_id}'" + ))); + } + }; + self.begin_execution(model_id); + // Kick off predictive warming of the *next* model concurrently with this + // inference, so a hinted/predicted model is hot by the time it is needed. + self.trigger_preflight(req); + + let inference_budget = remaining(overall_deadline).min(self.timeouts.inference()); let start = Instant::now(); - let invoke_timeout = Duration::from_secs(180); - let result = tokio::time::timeout(invoke_timeout, provider.invoke(model_id, req)).await; + let result = tokio::time::timeout(inference_budget, provider.invoke(model_id, req)).await; let duration_ms = start.elapsed().as_millis() as u64; self.end_execution(model_id); + drop(permit); match result { Ok(Ok(mut resp)) => { - resp.routing_reason = routing_reason; + resp.routing_reason = Some(routing_reason); self.circuit_breakers.record_success(provider_name); - if let Some(cap) = req.capability { - self.preflight.record(cap); - } - let _ = self.event_tx.send(EngineEvent::RequestCompleted { - request_id: req.request_id.clone(), - duration_ms, - success: true, - }); + self.record_transition(req); + self.emit_request_completed(&req.request_id, duration_ms, true); Ok(resp) } Ok(Err(e)) => { self.circuit_breakers.record_failure(provider_name); - let _ = self.event_tx.send(EngineEvent::RequestCompleted { - request_id: req.request_id.clone(), - duration_ms, - success: false, - }); + self.emit_request_completed(&req.request_id, duration_ms, false); Err(e) } Err(_) => { self.circuit_breakers.record_failure(provider_name); - let _ = self.event_tx.send(EngineEvent::RequestCompleted { - request_id: req.request_id.clone(), - duration_ms: invoke_timeout.as_millis() as u64, - success: false, - }); + self.emit_request_completed( + &req.request_id, + inference_budget.as_millis() as u64, + false, + ); Err(EngineError::Timeout(format!( - "provider '{}' did not respond within {}s", - provider_name, - invoke_timeout.as_secs() + "provider '{provider_name}' did not respond within {}s", + inference_budget.as_secs() ))) } } } + /// Run inference and stream typed output chunks. + /// + /// Returns a receiver that yields a `Started` chunk, then `Text` deltas as + /// the backend produces them, then a terminal `Completed` or `Error`. Errors + /// are delivered in-band as [`StreamChunk::Error`] rather than as a `Result`, + /// so a consumer only has to read one channel. + /// + /// Streaming targets the single best candidate. Unlike [`invoke`](Self::invoke) + /// it does not fail over once output has begun, since partial output cannot + /// be un-sent; a failure before the first token surfaces as `Error`. + pub fn invoke_stream(&self, req: InferenceRequest) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(64); + if let Some(weak) = self.weak() { + tokio::spawn(async move { + if let Some(engine) = weak.upgrade() { + engine.run_stream(req, tx).await; + } + }); + } + rx + } + + /// Drive one streaming request end to end, emitting chunks to `out`. + async fn run_stream(&self, req: InferenceRequest, out: mpsc::Sender) { + if let Err(e) = self + .reject_ambiguous_short_name(&req) + .and_then(|()| self.check_input_path(&req)) + { + let _ = out.send(StreamChunk::Error(e.info())).await; + return; + } + self.confirm_prediction(&req); + + let all_models = self.capabilities().await; + let providers = self.routing_providers(); + let candidates = self.build_candidates(&all_models, &req, &providers); + let Some(decision) = candidates.first() else { + let err = EngineError::NoCapableModel(Self::requested_capability(&req)); + let _ = out.send(StreamChunk::Error(err.info())).await; + return; + }; + + let model_id = decision.model.id.clone(); + let provider_name = decision.model.provider.clone(); + let max_concurrency = decision.model.execution.max_concurrency; + let routing_reason = decision.reason.clone(); + let overall_deadline = Instant::now() + self.timeouts.request(); + + if !self.circuit_breakers.allow_request(&provider_name) { + let err = EngineError::CircuitOpen(provider_name.clone()); + let _ = out.send(StreamChunk::Error(err.info())).await; + return; + } + let Some(provider) = self.find_provider(&provider_name) else { + let err = EngineError::Internal(format!("provider '{provider_name}' not found")); + let _ = out.send(StreamChunk::Error(err.info())).await; + return; + }; + + let _ = self.event_tx.send(EngineEvent::RequestStarted { + request_id: req.request_id.clone(), + capability: req.capability, + model_id: model_id.clone(), + }); + + if let Err(e) = self + .ensure_loaded(&model_id, &provider_name, &provider) + .await + { + self.emit_request_completed(&req.request_id, 0, false); + let _ = out.send(StreamChunk::Error(e.info())).await; + return; + } + + let sem = self.semaphore_for(&model_id, max_concurrency); + let queue_budget = remaining(overall_deadline).min(self.timeouts.queue()); + let permit = match tokio::time::timeout(queue_budget, sem.acquire_owned()).await { + Ok(Ok(p)) => p, + Ok(Err(_)) => { + self.emit_request_completed(&req.request_id, 0, false); + let err = EngineError::Internal("concurrency semaphore closed".into()); + let _ = out.send(StreamChunk::Error(err.info())).await; + return; + } + Err(_) => { + self.emit_request_completed(&req.request_id, 0, false); + let err = EngineError::Timeout(format!("queue wait exceeded for '{model_id}'")); + let _ = out.send(StreamChunk::Error(err.info())).await; + return; + } + }; + + self.begin_execution(&model_id); + self.trigger_preflight(&req); + + let _ = out + .send(StreamChunk::Started { + request_id: req.request_id.clone(), + model_used: model_id_name(&model_id).to_string(), + provider: provider_name.clone(), + }) + .await; + + // Forward the provider's text deltas into the output stream as they + // arrive. The provider drops its sink when it returns, ending the + // forwarder, so all deltas are flushed before the terminal chunk. + let (delta_tx, mut delta_rx) = mpsc::channel::(64); + let forward_out = out.clone(); + let forwarder = tokio::spawn(async move { + while let Some(delta) = delta_rx.recv().await { + if forward_out.send(StreamChunk::Text { delta }).await.is_err() { + break; + } + } + }); + + let inference_budget = remaining(overall_deadline).min(self.timeouts.inference()); + let start = Instant::now(); + let result = + tokio::time::timeout(inference_budget, provider.stream(&model_id, &req, delta_tx)) + .await; + let duration_ms = start.elapsed().as_millis() as u64; + + let _ = forwarder.await; + self.end_execution(&model_id); + drop(permit); + + let terminal = match result { + Ok(Ok(resp)) => { + self.circuit_breakers.record_success(&provider_name); + self.record_transition(&req); + self.emit_request_completed(&req.request_id, duration_ms, true); + self.metrics.record_request(true, duration_ms, false); + StreamChunk::Completed { + duration_ms, + tokens_used: resp.tokens_used, + file: resp.file, + fallback_used: false, + routing_reason: Some(routing_reason), + } + } + Ok(Err(e)) => { + self.circuit_breakers.record_failure(&provider_name); + self.emit_request_completed(&req.request_id, duration_ms, false); + self.metrics.record_request(false, duration_ms, false); + StreamChunk::Error(e.info()) + } + Err(_) => { + self.circuit_breakers.record_failure(&provider_name); + let budget_ms = inference_budget.as_millis() as u64; + self.emit_request_completed(&req.request_id, budget_ms, false); + self.metrics.record_request(false, budget_ms, false); + let err = EngineError::Timeout(format!( + "provider '{provider_name}' did not respond within {}s", + inference_budget.as_secs() + )); + StreamChunk::Error(err.info()) + } + }; + let _ = out.send(terminal).await; + } + + /// Ensure a model is resident, reserving memory before loading. + /// + /// The evict-then-reserve step runs under `load_gate` so concurrent loads + /// cannot jointly overcommit memory; the slow backend load then runs outside + /// the gate so independent models can load in parallel. async fn ensure_loaded( &self, model_id: &str, provider_name: &str, provider: &Arc, ) -> Result<(), EngineError> { - let residency = self - .models - .get(model_id) - .map(|m| m.residency) - .unwrap_or(ModelResidency::Unknown); - if matches!(residency, ModelResidency::Loaded | ModelResidency::Remote) { + if self.is_resident(model_id) { + self.memory.touch(model_id); return Ok(()); } - self.set_model_residency(model_id, ModelResidency::Loading); - let load_result = tokio::time::timeout(Duration::from_secs(30), provider.load(model_id)) - .await - .map_err(|_| { - self.circuit_breakers.record_failure(provider_name); - self.set_model_residency(model_id, ModelResidency::Unloaded); - EngineError::Timeout(format!("load timeout for {provider_name}/{model_id}")) - })? - .inspect_err(|_| { - self.circuit_breakers.record_failure(provider_name); - self.set_model_residency(model_id, ModelResidency::Unloaded); - })?; + let estimate = self + .models + .get(model_id) + .map(|m| m.memory_estimate_bytes) + .unwrap_or(0); - self.set_model_residency(model_id, load_result.residency); - if matches!(load_result.residency, ModelResidency::Loaded) - && let Some(card) = self.models.get(model_id) { - let bytes = load_result - .memory_bytes - .unwrap_or(card.memory_estimate_bytes); - if bytes > 0 { - if !self.memory.can_fit(bytes) { - self.evict_to_fit(model_id, bytes).await; + let _gate = self.load_gate.lock().await; + // Another task may have loaded this model while we waited for the gate. + if self.is_resident(model_id) { + self.memory.touch(model_id); + return Ok(()); + } + if estimate > 0 { + self.admit_memory(model_id, estimate).await?; + } + self.set_model_residency(model_id, ModelResidency::Loading); + } + + let load_result = + match tokio::time::timeout(self.timeouts.load(), provider.load(model_id)).await { + Ok(Ok(r)) => r, + Ok(Err(e)) => { + self.circuit_breakers.record_failure(provider_name); + self.rollback_reservation(model_id, estimate); + self.record_lifecycle(model_id, "load_failed", Some(&e.to_string())); + return Err(e); } - if !self.memory.can_fit(bytes) { - let _ = provider.unload(model_id).await; - self.set_model_residency(model_id, ModelResidency::Unloaded); - return Err(EngineError::MemoryPressure { - need: bytes, - available: self.memory.available_bytes(), - }); + Err(_) => { + self.circuit_breakers.record_failure(provider_name); + self.rollback_reservation(model_id, estimate); + self.record_lifecycle(model_id, "load_timeout", None); + return Err(EngineError::Timeout(format!( + "load timeout for {provider_name}/{model_id}" + ))); } - self.memory.allocate(model_id, bytes); - let _ = self.event_tx.send(EngineEvent::MemoryChanged { - used_bytes: self.memory.used_bytes(), - total_bytes: self.memory.budget_bytes(), - }); + }; + + self.set_model_residency(model_id, load_result.residency); + if matches!(load_result.residency, ModelResidency::Loaded) { + if let Some(observed) = load_result.memory_bytes + && observed > 0 + { + self.memory.reconcile(model_id, observed); } + // Mark freshly loaded so it is not the immediate LRU victim before its + // first lease is taken. + self.memory.touch(model_id); + self.emit_memory_changed(); + self.record_lifecycle(model_id, "load", None); } Ok(()) } - /// Evict least-recently-used idle models until `bytes` can fit, or until - /// nothing more can be freed. Models with in-flight requests and the - /// incoming model itself are never evicted. - async fn evict_to_fit(&self, incoming_id: &str, bytes: u64) { - let mut last_victim: Option = None; - while !self.memory.can_fit(bytes) { - let protected = self.protected_model_ids(incoming_id); + /// Reserve `bytes` for `model_id`, evicting idle models until it fits. + /// Caller must hold `load_gate`. + async fn admit_memory(&self, model_id: &str, bytes: u64) -> Result<(), EngineError> { + if self.memory.try_reserve(model_id, bytes) { + return Ok(()); + } + loop { + let mut protected = self.eviction_protected_ids(); + protected.push(model_id.to_string()); let Some(victim) = self.memory.lru_candidate(&protected) else { break; }; - // Defensive: if the previous eviction freed nothing, stop instead of spinning. - if last_victim.as_deref() == Some(victim.as_str()) { - break; - } - - let provider = self - .models - .get(&victim) - .and_then(|card| self.find_provider(&card.provider)); - if let Some(provider) = provider { - let _ = provider.unload(&victim).await; + self.unload_model(&victim, "memory_pressure").await; + if self.memory.try_reserve(model_id, bytes) { + return Ok(()); } - // Frees the allocation and emits ModelResidencyChanged/MemoryChanged. - self.set_model_residency(&victim, ModelResidency::Unloaded); - tracing::info!("evicted '{victim}' to admit {bytes} bytes"); - last_victim = Some(victim); } + Err(EngineError::MemoryPressure { + need: bytes, + available: self.memory.available_bytes(), + }) } - /// Model IDs that must never be evicted: the incoming model plus any model - /// with in-flight requests. - fn protected_model_ids(&self, incoming_id: &str) -> Vec { - let mut protected = vec![incoming_id.to_string()]; - protected.extend( - self.models - .iter() - .filter(|card| card.execution.active_requests > 0) - .map(|card| card.key().clone()), - ); - protected + fn rollback_reservation(&self, model_id: &str, estimate: u64) { + if estimate > 0 { + self.memory.deallocate(model_id); + } + self.set_model_residency(model_id, ModelResidency::Unloaded); + self.emit_memory_changed(); } - fn find_provider(&self, name: &str) -> Option> { - self.providers - .iter() - .find(|registered| registered.name == name) - .map(|registered| Arc::clone(®istered.provider)) + /// Unload a model from a backend and release its accounting. + async fn unload_model(&self, model_id: &str, reason: &str) { + let provider = self + .models + .get(model_id) + .and_then(|card| self.find_provider(&card.provider)); + if let Some(provider) = provider { + let _ = provider.unload(model_id).await; + } + // Frees the reservation and emits residency/memory events. + self.set_model_residency(model_id, ModelResidency::Unloaded); + // Distinguish the cause so metrics count each correctly: an idle sweep is + // an idle-unload, an operator action is a plain unload, and everything + // else (memory pressure) is an eviction. + let event = match reason { + "idle_timeout" => "idle_unload", + "manual" => "unload", + _ => "evict", + }; + self.record_lifecycle(model_id, event, Some(reason)); + let _ = self.event_tx.send(EngineEvent::ModelEvicted { + model_id: model_id.to_string(), + reason: reason.to_string(), + }); + tracing::info!("unloaded '{model_id}' ({reason})"); } - fn routing_providers(&self) -> Vec { - self.providers + /// Model IDs that must never be evicted by memory pressure or the idle sweep. + /// + /// Leased (in-flight) models are already excluded inside the memory manager; + /// this additionally protects: + /// - models whose load is currently in flight (`Loading`), so a concurrent + /// admission cannot unload a model out from under an active load and leave + /// its reservation orphaned; and + /// - resident models whose capability is covered by an active subscription + /// (the RFC "resident keep-alive" policy). + /// + /// If every candidate is protected, admission returns structured memory + /// pressure rather than breaking a subscription or a load. + fn eviction_protected_ids(&self) -> Vec { + let subscribed = self.subscriptions.active_capabilities(); + self.models .iter() - .map(|registered| { - let health = self - .backend_health - .get(®istered.name) - .map(|h| *h) - .unwrap_or(BackendHealth::Unknown); - RoutingProvider { - name: registered.name.clone(), - kind: registered.kind, - priority: registered.priority, - health, - circuit_open: self.circuit_breakers.state(®istered.name) - == CircuitState::Open, - } + .filter(|card| { + matches!(card.residency, ModelResidency::Loading) + || (matches!(card.residency, ModelResidency::Loaded) + && !subscribed.is_empty() + && card.capabilities.iter().any(|cap| subscribed.contains(cap))) }) + .map(|card| card.key().clone()) .collect() } - fn begin_execution(&self, model_id: &str) { - if let Some(mut card) = self.models.get_mut(model_id) { - let old_status = card.status; - card.execution.active_requests = card.execution.active_requests.saturating_add(1); - card.refresh_status(); - if old_status != card.status { - let _ = self.event_tx.send(EngineEvent::ModelStatusChanged { - model_id: model_id.to_string(), - old: old_status, - new: card.status, - }); + // ----- Preflight / subscriptions ------------------------------------------------- + + /// Scope key for history and predictions: prefer session, then app, else global. + fn scope_key(req: &InferenceRequest) -> String { + req.session_id + .clone() + .or_else(|| req.app_id.clone()) + .unwrap_or_else(|| GLOBAL_SCOPE.to_string()) + } + + fn weak(&self) -> Option> { + self.weak_self.get().cloned() + } + + /// Confirm (hit/miss) the prediction made for this scope by the *previous* + /// request, against the capability the current request actually asked for. + /// Runs once per request, before a new prediction is made. + fn confirm_prediction(&self, req: &InferenceRequest) { + let Some(cap) = req.capability else { + return; + }; + let scope = Self::scope_key(req); + if let Some((_, predicted)) = self.pending_predictions.remove(&scope) { + if predicted == cap { + self.preflight_metrics.hit(); + } else { + self.preflight_metrics.miss(); } } - self.memory.touch(model_id); } - fn end_execution(&self, model_id: &str) { - if let Some(mut card) = self.models.get_mut(model_id) { - let old_status = card.status; - card.execution.active_requests = card.execution.active_requests.saturating_sub(1); - card.refresh_status(); - if old_status != card.status { - let _ = self.event_tx.send(EngineEvent::ModelStatusChanged { - model_id: model_id.to_string(), - old: old_status, - new: card.status, - }); - } + /// Remember a prediction for later hit/miss confirmation, bounding the map + /// so caller-supplied scope keys cannot grow it without limit. + fn remember_prediction(&self, scope: String, capability: Capability) { + if self.pending_predictions.len() >= MAX_PENDING_PREDICTIONS + && !self.pending_predictions.contains_key(&scope) + && let Some(victim) = self + .pending_predictions + .iter() + .next() + .map(|e| e.key().clone()) + { + self.pending_predictions.remove(&victim); } - self.memory.touch(model_id); + self.pending_predictions.insert(scope, capability); } - fn set_model_residency(&self, model_id: &str, new_residency: ModelResidency) { - if let Some(mut card) = self.models.get_mut(model_id) { - let old_residency = card.residency; - let old_status = card.status; - if old_residency != new_residency { - card.residency = new_residency; - card.refresh_status(); - let new_status = card.status; - // Release the shard guard before touching the memory manager / event bus. - drop(card); - let _ = self.event_tx.send(EngineEvent::ModelResidencyChanged { - model_id: model_id.to_string(), - old: old_residency, - new: new_residency, - }); - if old_status != new_status { - let _ = self.event_tx.send(EngineEvent::ModelStatusChanged { - model_id: model_id.to_string(), - old: old_status, - new: new_status, - }); - } - // A model that stops being resident releases its memory budget. - if old_residency == ModelResidency::Loaded - && new_residency != ModelResidency::Loaded - { - self.memory.deallocate(model_id); - let _ = self.event_tx.send(EngineEvent::MemoryChanged { - used_bytes: self.memory.used_bytes(), - total_bytes: self.memory.budget_bytes(), - }); - } + /// Record a completed capability into the scope's transition history. + fn record_transition(&self, req: &InferenceRequest) { + if !self.preflight_config.history_learning { + return; + } + if let Some(cap) = req.capability { + self.preflight.record(&Self::scope_key(req), cap); + } + } + + /// Choose which capability (if any) to speculatively warm for this request, + /// in priority order: explicit hint, then subscription, then learned history. + fn select_warm_capability(&self, req: &InferenceRequest) -> Option<(Capability, &'static str)> { + if let Some(hint) = req + .hint_next + .as_deref() + .and_then(Capability::from_str_loose) + { + return Some((hint, "hint")); + } + if let Some(cap) = self.first_subscribed_without_resident() { + return Some((cap, "subscription")); + } + if self.preflight_config.history_learning + && let Some(current) = req.capability + { + let scope = Self::scope_key(req); + if let Some(pred) = self.preflight.predict( + &scope, + current, + self.preflight_config.min_samples, + self.preflight_config.confidence_threshold, + ) { + self.preflight_metrics.prediction(); + self.remember_prediction(scope, pred.capability); + return Some((pred.capability, "history")); } } + None } - /// Get a snapshot of the engine status. - pub async fn status(&self) -> EngineStatus { - let all_models = self.capabilities().await; - let total_models = all_models.len(); - let loaded_models = all_models - .iter() - .filter(|m| matches!(m.residency, ModelResidency::Loaded | ModelResidency::Remote)) - .count(); + /// A subscribed capability that currently has no resident model, if any. + fn first_subscribed_without_resident(&self) -> Option { + let mut subscribed: Vec = self + .subscriptions + .active_capabilities() + .into_iter() + .collect(); + // Deterministic order so warming is reproducible. + subscribed.sort_by_key(|c| c.to_string()); + subscribed + .into_iter() + .find(|cap| !self.has_resident_for(*cap)) + } - let backends = self.backend_statuses(); - let provider_health = backends + fn has_resident_for(&self, cap: Capability) -> bool { + self.models.iter().any(|card| { + card.supports(cap) + && matches!( + card.residency, + ModelResidency::Loaded | ModelResidency::Remote + ) + }) + } + + /// Decide and launch speculative warming for the request, if enabled. + fn trigger_preflight(&self, req: &InferenceRequest) { + if !self.preflight_config.enabled || !self.preflight_config.speculative_warming { + return; + } + if let Some((cap, source)) = self.select_warm_capability(req) { + self.warm_capability(cap, source, req.app_id.clone(), req.session_id.clone()); + } + } + + /// Route a capability to its best model and warm it on a background task. + /// + /// Speculative loads go through the same reservation/eviction admission as + /// regular requests (via `ensure_loaded`), are deduplicated per model, and + /// are skipped when the model is already resident. + fn warm_capability( + &self, + cap: Capability, + source: &str, + app_id: Option, + session_id: Option, + ) { + if !self.preflight_config.enabled || !self.preflight_config.speculative_warming { + return; + } + + let all = self.models_snapshot(); + let providers = self.routing_providers(); + let route_req = InferenceRequest { + capability: Some(cap), + model: None, + app_id, + session_id, + fallback_policy: FallbackPolicy::CapabilityOnly, + messages: Vec::new(), + input_file: None, + params: serde_json::Value::Null, + hint_next: None, + request_id: String::new(), + }; + let budget = Some(self.memory.budget_bytes()); + let Some(best) = Router::route_ranked(&all, &route_req, &providers, budget) + .into_iter() + .next() + else { + return; + }; + + let model_id = best.model.id.clone(); + // Already hot (loaded locally or cloud-backed) → nothing to warm. + if self.is_resident(&model_id) { + return; + } + // A load is already in flight for this model (driven by a request or an + // earlier warm) → don't pile on a redundant concurrent load. + if self + .models + .get(&model_id) + .is_some_and(|c| c.residency == ModelResidency::Loading) + { + self.preflight_metrics.warm_skipped(); + return; + } + // Deduplicate: skip if a warm for this model is still running. We key on + // `is_finished()` rather than mere presence so a completed task's handle + // (left in the map intentionally) never permanently blocks future warms. + if self + .warming + .get(&model_id) + .is_some_and(|h| !h.is_finished()) + { + self.preflight_metrics.warm_skipped(); + return; + } + let provider_name = best.model.provider.clone(); + let Some(provider) = self.find_provider(&provider_name) else { + return; + }; + let Some(weak) = self.weak() else { + return; + }; + + self.preflight_metrics.warm_started(); + let _ = self.event_tx.send(EngineEvent::PreflightWarmStarted { + model_id: model_id.clone(), + source: source.to_string(), + }); + + let warm_id = model_id.clone(); + let handle = tokio::spawn(async move { + let Some(engine) = weak.upgrade() else { + return; + }; + let ok = engine + .ensure_loaded(&warm_id, &provider_name, &provider) + .await + .is_ok(); + // The handle is intentionally left in `warming`; dedup checks + // `is_finished()`, and the next warm of this model overwrites it. + if ok { + engine.preflight_metrics.warm_completed(); + } else { + engine.preflight_metrics.warm_failed(); + } + let _ = engine.event_tx.send(EngineEvent::PreflightWarmCompleted { + model_id: warm_id, + success: ok, + }); + }); + self.warming.insert(model_id, handle.abort_handle()); + } + + /// Register a capability subscription and immediately warm its models. + /// + /// Returns the new subscription id. + pub fn subscribe( + &self, + app_id: Option, + session_id: Option, + capabilities: Vec, + ttl: Option, + ) -> u64 { + let id = self.subscriptions.subscribe( + app_id.clone(), + session_id.clone(), + capabilities.clone(), + ttl, + ); + for cap in capabilities { + self.warm_capability(cap, "subscription", app_id.clone(), session_id.clone()); + } + id + } + + /// Remove a subscription by id. Returns whether one existed. + pub fn unsubscribe(&self, id: u64) -> bool { + self.subscriptions.unsubscribe(id) + } + + /// List all active subscriptions. + pub fn subscriptions(&self) -> Vec { + self.subscriptions.list() + } + + /// Snapshot of Preflight effectiveness counters. + pub fn preflight_stats(&self) -> PreflightStats { + self.preflight_metrics.snapshot() + } + + /// Background task: unload models idle beyond the configured timeout. + fn spawn_idle_eviction(&self) { + if self.idle_timeout.is_zero() { + return; + } + let Some(weak) = self.weak_self.get().cloned() else { + return; + }; + // Tick often enough to act soon after the timeout without busy-looping. + let tick = Duration::from_secs((self.idle_timeout.as_secs().max(2) / 2).clamp(1, 15)); + let handle = tokio::spawn(async move { + let mut ticker = tokio::time::interval(tick); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + ticker.tick().await; + let Some(engine) = weak.upgrade() else { + break; + }; + engine.run_idle_sweep().await; + } + }); + *self.idle_task.lock().unwrap_or_else(|e| e.into_inner()) = Some(handle); + } + + /// Background task: periodically delete stale engine artifacts. + fn spawn_artifact_sweep(&self, sweeper: crate::artifacts::ArtifactSweeper) { + let retention = sweeper.retention(); + if retention.is_zero() { + return; // retention 0 disables cleanup + } + let Some(weak) = self.weak_self.get().cloned() else { + return; + }; + // Sweep on roughly half the retention interval, bounded so tests and + // long-lived daemons both behave. + let tick = Duration::from_secs((retention.as_secs().max(2) / 2).clamp(1, 300)); + let handle = tokio::spawn(async move { + let mut ticker = tokio::time::interval(tick); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + ticker.tick().await; + if weak.upgrade().is_none() { + break; + } + let removed = sweeper.sweep(); + if removed > 0 { + tracing::debug!("artifact sweep removed {removed} stale file(s)"); + } + } + }); + *self.artifact_task.lock().unwrap_or_else(|e| e.into_inner()) = Some(handle); + } + + /// Reject a request whose `input_file` falls outside the configured + /// allowlist. A no-op when no roots are configured. + fn check_input_path(&self, req: &InferenceRequest) -> Result<(), EngineError> { + if self.input_roots.is_empty() { + return Ok(()); + } + let Some(path) = req.input_file.as_deref() else { + return Ok(()); + }; + let canonical = std::fs::canonicalize(path).map_err(|_| { + EngineError::InvalidRequest(format!("input_file '{path}' cannot be resolved")) + })?; + if self + .input_roots + .iter() + .any(|root| canonical.starts_with(root)) + { + Ok(()) + } else { + Err(EngineError::InvalidRequest(format!( + "input_file '{path}' is outside the allowed roots" + ))) + } + } + + /// Unload every model that has been idle past the timeout. + async fn run_idle_sweep(&self) { + let protected = self.eviction_protected_ids(); + let victims = self.memory.idle_candidates(self.idle_timeout, &protected); + for victim in victims { + // Re-check under no lock: only unload models that are genuinely + // resident and not leased right now. + let resident = self + .models + .get(&victim) + .map(|c| matches!(c.residency, ModelResidency::Loaded)) + .unwrap_or(false); + if !resident || self.memory.lease_count(&victim) > 0 { + continue; + } + self.unload_model(&victim, "idle_timeout").await; + } + } + + fn is_resident(&self, model_id: &str) -> bool { + self.models + .get(model_id) + .map(|m| matches!(m.residency, ModelResidency::Loaded | ModelResidency::Remote)) + .unwrap_or(false) + } + + fn semaphore_for(&self, model_id: &str, max_concurrency: u32) -> Arc { + let permits = max_concurrency.max(1) as usize; + self.semaphores + .entry(model_id.to_string()) + .or_insert_with(|| Arc::new(Semaphore::new(permits))) + .clone() + } + + fn find_provider(&self, name: &str) -> Option> { + self.providers + .iter() + .find(|registered| registered.name == name) + .map(|registered| Arc::clone(®istered.provider)) + } + + fn routing_providers(&self) -> Vec { + self.providers + .iter() + .map(|registered| { + let health = self + .backend_health + .get(®istered.name) + .map(|h| *h) + .unwrap_or(BackendHealth::Unknown); + RoutingProvider { + name: registered.name.clone(), + kind: registered.kind, + priority: registered.priority, + health, + circuit_open: self.circuit_breakers.state(®istered.name) + == CircuitState::Open, + } + }) + .collect() + } + + fn begin_execution(&self, model_id: &str) { + if let Some(mut card) = self.models.get_mut(model_id) { + let old_status = card.status; + card.execution.active_requests = card.execution.active_requests.saturating_add(1); + card.refresh_status(); + if old_status != card.status { + let _ = self.event_tx.send(EngineEvent::ModelStatusChanged { + model_id: model_id.to_string(), + old: old_status, + new: card.status, + }); + } + } + // The memory lease protects this model from eviction for the duration. + self.memory.lease(model_id); + } + + fn end_execution(&self, model_id: &str) { + if let Some(mut card) = self.models.get_mut(model_id) { + let old_status = card.status; + card.execution.active_requests = card.execution.active_requests.saturating_sub(1); + card.refresh_status(); + if old_status != card.status { + let _ = self.event_tx.send(EngineEvent::ModelStatusChanged { + model_id: model_id.to_string(), + old: old_status, + new: card.status, + }); + } + } + self.memory.release_lease(model_id); + } + + fn set_model_residency(&self, model_id: &str, new_residency: ModelResidency) { + if let Some(mut card) = self.models.get_mut(model_id) { + let old_residency = card.residency; + let old_status = card.status; + if old_residency != new_residency { + card.residency = new_residency; + card.refresh_status(); + let new_status = card.status; + // Release the shard guard before touching the memory manager / event bus. + drop(card); + let _ = self.event_tx.send(EngineEvent::ModelResidencyChanged { + model_id: model_id.to_string(), + old: old_residency, + new: new_residency, + }); + if old_status != new_status { + let _ = self.event_tx.send(EngineEvent::ModelStatusChanged { + model_id: model_id.to_string(), + old: old_status, + new: new_status, + }); + } + // A model that stops being resident releases its memory budget. + if old_residency == ModelResidency::Loaded + && new_residency != ModelResidency::Loaded + { + self.memory.deallocate(model_id); + self.emit_memory_changed(); + } + } + } + } + + fn emit_memory_changed(&self) { + let _ = self.event_tx.send(EngineEvent::MemoryChanged { + used_bytes: self.memory.used_bytes(), + total_bytes: self.memory.budget_bytes(), + }); + } + + fn emit_request_completed(&self, request_id: &str, duration_ms: u64, success: bool) { + let _ = self.event_tx.send(EngineEvent::RequestCompleted { + request_id: request_id.to_string(), + duration_ms, + success, + }); + } + + fn record_lifecycle(&self, model_id: &str, event: &str, detail: Option<&str>) { + self.metrics.record_lifecycle(event); + let record = LifecycleRecord { + seq: self.lifecycle_seq.fetch_add(1, Ordering::Relaxed), + at_ms: self.started_at.elapsed().as_millis() as u64, + model_id: model_id.to_string(), + event: event.to_string(), + detail: detail.map(ToString::to_string), + }; + let mut hist = self.lifecycle.lock().unwrap_or_else(|e| e.into_inner()); + if hist.len() >= LIFECYCLE_CAPACITY { + hist.pop_front(); + } + hist.push_back(record); + } + + /// Return the rolling lifecycle history, oldest first. + pub fn lifecycle_history(&self) -> Vec { + self.lifecycle + .lock() + .unwrap_or_else(|e| e.into_inner()) + .iter() + .cloned() + .collect() + } + + /// Return a snapshot of the engine's memory accounting. + pub fn memory_report(&self) -> MemoryReport { + MemoryReport { + used_bytes: self.memory.used_bytes(), + budget_bytes: self.memory.budget_bytes(), + available_bytes: self.memory.available_bytes(), + allocations: self.memory.snapshot(), + } + } + + /// Get a snapshot of the engine status. + pub async fn status(&self) -> EngineStatus { + let all_models = self.capabilities().await; + let total_models = all_models.len(); + let loaded_models = all_models + .iter() + .filter(|m| matches!(m.residency, ModelResidency::Loaded | ModelResidency::Remote)) + .count(); + + let backends = self.backend_statuses(); + let provider_health = backends .iter() .map(|backend| ProviderHealth { name: backend.name.clone(), @@ -654,12 +1523,100 @@ impl Engine { self.event_tx.subscribe() } + /// Render the current metrics in Prometheus text-exposition format. + /// + /// Combines the process-wide counters with point-in-time gauges sampled + /// from the registry, memory manager, preflight counters, and per-provider + /// health. Label cardinality is bounded: only the fixed provider set adds + /// labels. + pub fn metrics_prometheus(&self) -> String { + let models = self.models_snapshot(); + let loaded = models + .iter() + .filter(|m| matches!(m.residency, ModelResidency::Loaded | ModelResidency::Remote)) + .count() as u64; + let preflight = self.preflight_metrics.snapshot(); + let provider_up = self + .providers + .iter() + .map(|p| { + let health = self + .backend_health + .get(&p.name) + .map(|h| *h) + .unwrap_or(BackendHealth::Unknown); + let up = health.is_routable() + && self.circuit_breakers.state(&p.name) != CircuitState::Open; + (p.name.clone(), up) + }) + .collect(); + + let gauges = MetricsGauges { + models_total: models.len() as u64, + models_loaded: loaded, + memory_used_bytes: self.memory.used_bytes(), + memory_budget_bytes: self.memory.budget_bytes(), + preflight_warms_started: preflight.warms_started, + preflight_hits: preflight.hits, + provider_up, + }; + self.metrics.render_prometheus(&gauges) + } + + /// Manually load (warm) a model by id. Used by the management interface. + /// + /// Goes through the same reservation-based admission and lifecycle path as + /// an on-demand load, so memory accounting and eviction rules still hold. + pub async fn load_model(&self, model_id: &str) -> Result<(), EngineError> { + let provider_name = self + .models + .get(model_id) + .map(|c| c.provider.clone()) + .ok_or_else(|| EngineError::InvalidRequest(format!("unknown model '{model_id}'")))?; + let provider = self.find_provider(&provider_name).ok_or_else(|| { + EngineError::Internal(format!("provider '{provider_name}' not found")) + })?; + self.ensure_loaded(model_id, &provider_name, &provider) + .await + } + + /// Manually unload a model by id. Returns whether the model was known. + /// + /// Refuses to unload a model with active leases (in-flight inference) so a + /// manual action cannot corrupt a running request. + pub async fn unload_model_manual(&self, model_id: &str) -> Result { + if !self.models.contains_key(model_id) { + return Ok(false); + } + if self.memory.lease_count(model_id) > 0 { + return Err(EngineError::InvalidRequest(format!( + "model '{model_id}' is busy and cannot be unloaded" + ))); + } + self.unload_model(model_id, "manual").await; + Ok(true) + } + /// Engine uptime in seconds. pub fn uptime_secs(&self) -> u64 { self.started_at.elapsed().as_secs() } } +impl Drop for Engine { + fn drop(&mut self) { + for task in [&self.idle_task, &self.artifact_task] { + if let Some(handle) = task.lock().unwrap_or_else(|e| e.into_inner()).take() { + handle.abort(); + } + } + // Cancel any in-flight speculative warm tasks. + for entry in self.warming.iter() { + entry.value().abort(); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -668,9 +1625,14 @@ mod tests { BackendFeature, Capability, ExecutionState, LifecycleResult, Message, ModelAvailability, canonical_model_id, }; - use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::atomic::AtomicUsize; + + use crate::config::{ + EngineConfig, ListenConfig, MemoryConfig, PreflightConfig, ProviderConfig, SecurityConfig, + TimeoutConfig, + }; - use crate::config::{EngineConfig, ListenConfig, MemoryConfig, ProviderConfig}; + const MB: u64 = 1024 * 1024; fn minimal_config() -> EngineConfig { EngineConfig { @@ -679,51 +1641,104 @@ mod tests { budget_mb: Some(100), idle_timeout_secs: 60, }, + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), providers: vec![], } } - #[derive(Default)] - struct MockProvider { + /// Behavior knobs for the configurable test provider. + #[derive(Clone, Copy)] + enum InvokeBehavior { + Ok, + SlowOkMs(u64), + FailRetryable, + FailInvalid, + } + + struct TestProvider { name: String, - discover_calls: AtomicUsize, + kind: ProviderKind, + /// (model_name, capability, residency, mem_bytes, max_concurrency) + template: Vec<(String, Capability, ModelResidency, u64, u32)>, + behavior: InvokeBehavior, + load_should_fail: bool, + in_flight: Arc, + max_in_flight: Arc, + invoke_calls: Arc, } - impl MockProvider { - fn new(name: &str) -> Self { + impl TestProvider { + fn new(name: &str, kind: ProviderKind) -> Self { Self { name: name.into(), - discover_calls: AtomicUsize::new(0), + kind, + template: Vec::new(), + behavior: InvokeBehavior::Ok, + load_should_fail: false, + in_flight: Arc::new(AtomicUsize::new(0)), + max_in_flight: Arc::new(AtomicUsize::new(0)), + invoke_calls: Arc::new(AtomicUsize::new(0)), } } + + fn with_model( + mut self, + name: &str, + cap: Capability, + residency: ModelResidency, + mem_bytes: u64, + max_concurrency: u32, + ) -> Self { + self.template + .push((name.into(), cap, residency, mem_bytes, max_concurrency)); + self + } + + fn behavior(mut self, behavior: InvokeBehavior) -> Self { + self.behavior = behavior; + self + } + + fn failing_load(mut self) -> Self { + self.load_should_fail = true; + self + } } #[async_trait] - impl Provider for MockProvider { + impl Provider for TestProvider { fn name(&self) -> &str { &self.name } - fn kind(&self) -> ProviderKind { - ProviderKind::OpenAiCompatible + self.kind } - fn features(&self) -> Vec { - vec![BackendFeature::Discovery] + vec![BackendFeature::Discovery, BackendFeature::Load] } async fn discover(&self) -> Result, EngineError> { - self.discover_calls.fetch_add(1, Ordering::SeqCst); - let mut card = ModelCard::new(&self.name, "mock-chat", Capability::Chat, CostTier::Low); - card.id = canonical_model_id(&self.name, "mock-chat"); - card.availability = ModelAvailability::Configured; - card.residency = ModelResidency::Remote; - card.execution = ExecutionState { - active_requests: 0, - max_concurrency: 4, - }; - card.refresh_status(); - Ok(vec![card]) + let cards = self + .template + .iter() + .map(|(name, cap, residency, mem, conc)| { + let mut card = ModelCard::new(&self.name, name, *cap, CostTier::Low); + card.id = canonical_model_id(&self.name, name); + card.availability = ModelAvailability::Discovered; + card.residency = *residency; + card.memory_estimate_bytes = *mem; + card.execution = ExecutionState { + active_requests: 0, + max_concurrency: *conc, + }; + card.refresh_status(); + card + }) + .collect(); + Ok(cards) } async fn health(&self) -> Result { @@ -731,20 +1746,36 @@ mod tests { } async fn load(&self, model_id: &str) -> Result { + if self.load_should_fail { + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "load failed (test)".into(), + }); + } + let residency = if self.kind == ProviderKind::Ollama { + ModelResidency::Loaded + } else { + ModelResidency::Remote + }; + let mem = self + .template + .iter() + .find(|(n, ..)| canonical_model_id(&self.name, n) == model_id) + .map(|(.., mem, _)| *mem); Ok(LifecycleResult { model_id: model_id.into(), - residency: ModelResidency::Remote, - memory_bytes: Some(0), - changed: false, + residency, + memory_bytes: mem, + changed: true, }) } async fn unload(&self, model_id: &str) -> Result { Ok(LifecycleResult { model_id: model_id.into(), - residency: ModelResidency::Remote, + residency: ModelResidency::Unloaded, memory_bytes: Some(0), - changed: false, + changed: true, }) } @@ -753,6 +1784,26 @@ mod tests { model_id: &str, request: &InferenceRequest, ) -> Result { + self.invoke_calls.fetch_add(1, Ordering::SeqCst); + let now = self.in_flight.fetch_add(1, Ordering::SeqCst) + 1; + self.max_in_flight.fetch_max(now, Ordering::SeqCst); + + let behavior = self.behavior; + let result = match behavior { + InvokeBehavior::Ok => Ok(()), + InvokeBehavior::SlowOkMs(ms) => { + tokio::time::sleep(Duration::from_millis(ms)).await; + Ok(()) + } + InvokeBehavior::FailRetryable => Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "boom".into(), + }), + InvokeBehavior::FailInvalid => Err(EngineError::InvalidRequest("bad input".into())), + }; + self.in_flight.fetch_sub(1, Ordering::SeqCst); + result?; + Ok(InferenceResponse { text: Some("ok".into()), file: None, @@ -767,32 +1818,71 @@ mod tests { } } - fn engine_with_provider(provider: Arc) -> Arc { + fn build_engine( + providers: Vec>, + budget_mb: u64, + idle_secs: u64, + ) -> Arc { let (event_tx, _) = broadcast::channel(256); - Arc::new(Engine { - providers: vec![RegisteredProvider { + let registered = providers + .into_iter() + .map(|provider| RegisteredProvider { name: provider.name().into(), kind: provider.kind(), - priority: 1, + priority: if provider.kind() == ProviderKind::Ollama { + 1 + } else { + 10 + }, provider, - }], + }) + .collect(); + let engine = Arc::new(Engine { + providers: registered, models: DashMap::new(), backend_health: DashMap::new(), - memory: MemoryManager::new(Some(100)), + memory: MemoryManager::new(Some(budget_mb)), circuit_breakers: CircuitBreakerRegistry::new(CircuitBreakerConfig::default()), preflight: PreflightPredictor::new(), + preflight_metrics: PreflightMetrics::default(), + subscriptions: SubscriptionRegistry::new(), + warming: DashMap::new(), + pending_predictions: DashMap::new(), + timeouts: TimeoutConfig::default(), + preflight_config: PreflightConfig::default(), + idle_timeout: Duration::from_secs(idle_secs), + semaphores: DashMap::new(), + load_gate: AsyncMutex::new(()), + lifecycle: Mutex::new(VecDeque::new()), + lifecycle_seq: AtomicU64::new(0), + idle_task: Mutex::new(None), + artifact_task: Mutex::new(None), + weak_self: OnceLock::new(), event_tx, + metrics: EngineMetrics::default(), + input_roots: Vec::new(), started_at: Instant::now(), - }) + }); + let _ = engine.weak_self.set(Arc::downgrade(&engine)); + engine.spawn_idle_eviction(); + engine } fn chat_request(model: Option<&str>) -> InferenceRequest { + request(Capability::Chat, model, FallbackPolicy::default()) + } + + fn request( + capability: Capability, + model: Option<&str>, + fallback_policy: FallbackPolicy, + ) -> InferenceRequest { InferenceRequest { - capability: Some(Capability::Chat), + capability: Some(capability), model: model.map(str::to_owned), app_id: None, session_id: None, - fallback_policy: FallbackPolicy::default(), + fallback_policy, messages: vec![Message { role: "user".into(), content: "hello".into(), @@ -807,8 +1897,7 @@ mod tests { #[tokio::test] async fn engine_starts_with_empty_config() { let engine = Engine::new(minimal_config()).await; - let caps = engine.capabilities().await; - assert!(caps.is_empty()); + assert!(engine.capabilities().await.is_empty()); } #[tokio::test] @@ -830,14 +1919,11 @@ mod tests { async fn event_subscription_works() { let engine = Engine::new(minimal_config()).await; let mut rx = engine.subscribe_events(); - let _ = engine.event_tx.send(EngineEvent::ProviderHealthChanged { provider: "test".into(), health: BackendHealth::Healthy, }); - - let event = rx.recv().await.unwrap(); - match event { + match rx.recv().await.unwrap() { EngineEvent::ProviderHealthChanged { provider, health } => { assert_eq!(provider, "test"); assert_eq!(health, BackendHealth::Healthy); @@ -854,6 +1940,10 @@ mod tests { budget_mb: Some(100), idle_timeout_secs: 60, }, + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), providers: vec![ProviderConfig { name: "disabled-ollama".into(), kind: "ollama".into(), @@ -863,99 +1953,781 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: false, + ..Default::default() }], }; let engine = Engine::new(config).await; - let status = engine.status().await; - assert_eq!(status.providers, 0); + assert_eq!(engine.status().await.providers, 0); } #[tokio::test] - async fn mock_provider_refresh_and_invoke_are_deterministic() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); + async fn discovers_and_invokes_deterministically() { + let provider = Arc::new( + TestProvider::new("cloud", ProviderKind::OpenAiCompatible).with_model( + "mock-chat", + Capability::Chat, + ModelResidency::Remote, + 0, + 4, + ), + ); + let engine = build_engine(vec![provider], 100, 0); engine.refresh_resources().await; let caps = engine.capabilities().await; assert_eq!(caps.len(), 1); - assert_eq!(caps[0].id, "mock/mock-chat"); + assert_eq!(caps[0].id, "cloud/mock-chat"); let resp = engine.invoke(chat_request(None)).await.unwrap(); assert_eq!(resp.text.as_deref(), Some("ok")); - assert_eq!(resp.provider, "mock"); + assert_eq!(resp.provider, "cloud"); assert!(resp.routing_reason.is_some()); + assert!(!resp.fallback_used); } #[tokio::test] async fn ambiguous_short_model_name_is_rejected() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); - let mut a = ModelCard::new("a", "same", Capability::Chat, CostTier::Free); - a.residency = ModelResidency::Remote; - a.refresh_status(); - let mut b = ModelCard::new("b", "same", Capability::Chat, CostTier::Free); - b.residency = ModelResidency::Remote; - b.refresh_status(); - engine.models.insert(a.id.clone(), a); - engine.models.insert(b.id.clone(), b); - + let engine = build_engine(vec![], 100, 0); + for prov in ["a", "b"] { + let mut card = ModelCard::new(prov, "same", Capability::Chat, CostTier::Free); + card.residency = ModelResidency::Remote; + card.refresh_status(); + engine.models.insert(card.id.clone(), card); + } let err = engine.invoke(chat_request(Some("same"))).await.unwrap_err(); assert!(matches!(err, EngineError::InvalidRequest(_))); } + #[tokio::test] + async fn capability_routing_prefers_local() { + let local = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible).with_model( + "gpt", + Capability::Chat, + ModelResidency::Remote, + 0, + 32, + ), + ); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + let resp = engine.invoke(chat_request(None)).await.unwrap(); + assert_eq!(resp.model_used, "qwen"); + assert!(!resp.fallback_used); + } + + #[tokio::test] + async fn named_request_routes_exactly() { + let local = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Loaded, + 10 * MB, + 1, + ), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible).with_model( + "gpt-4o", + Capability::Chat, + ModelResidency::Remote, + 0, + 32, + ), + ); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + let resp = engine.invoke(chat_request(Some("gpt-4o"))).await.unwrap(); + assert_eq!(resp.model_used, "gpt-4o"); + assert_eq!(resp.provider, "openai"); + } + + #[tokio::test] + async fn retryable_failure_falls_back_to_next_candidate() { + let local = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model("qwen", Capability::Chat, ModelResidency::Loaded, 10 * MB, 1) + .behavior(InvokeBehavior::FailRetryable), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible) + .with_model("gpt", Capability::Chat, ModelResidency::Remote, 0, 32) + .behavior(InvokeBehavior::Ok), + ); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + let resp = engine.invoke(chat_request(None)).await.unwrap(); + assert_eq!(resp.provider, "openai"); + assert!(resp.fallback_used); + } + + #[tokio::test] + async fn local_tts_failure_falls_back_to_cloud_tts() { + // A local TTS backend (preferred by locality) crashes mid-synthesis; the + // engine must fail over to the configured cloud TTS model. This is RFC + // acceptance step 8 exercised deterministically without a real backend. + let local_tts = Arc::new( + TestProvider::new("local-tts", ProviderKind::LocalTts) + .with_model( + "kokoro", + Capability::Tts, + ModelResidency::Loaded, + 10 * MB, + 1, + ) + .behavior(InvokeBehavior::FailRetryable), + ); + let cloud_tts = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible) + .with_model("tts-1", Capability::Tts, ModelResidency::Remote, 0, 32) + .behavior(InvokeBehavior::Ok), + ); + let local_calls = local_tts.invoke_calls.clone(); + let engine = build_engine(vec![local_tts, cloud_tts], 1000, 0); + engine.refresh_resources().await; + + let resp = engine + .invoke(request(Capability::Tts, None, FallbackPolicy::default())) + .await + .unwrap(); + // The local backend was tried first, then the engine fell over to cloud. + assert_eq!(local_calls.load(Ordering::SeqCst), 1); + assert_eq!(resp.provider, "openai"); + assert_eq!(resp.model_used, "tts-1"); + assert!(resp.fallback_used); + } + + #[tokio::test] + async fn invalid_request_does_not_fall_back() { + let local = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model("qwen", Capability::Chat, ModelResidency::Loaded, 10 * MB, 1) + .behavior(InvokeBehavior::FailInvalid), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible) + .with_model("gpt", Capability::Chat, ModelResidency::Remote, 0, 32) + .behavior(InvokeBehavior::Ok), + ); + let cloud_calls = cloud.invoke_calls.clone(); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + let err = engine.invoke(chat_request(None)).await.unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + // The cloud candidate must never have been tried. + assert_eq!(cloud_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn named_strict_does_not_fall_back() { + let local = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model("qwen", Capability::Chat, ModelResidency::Loaded, 10 * MB, 1) + .behavior(InvokeBehavior::FailRetryable), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible) + .with_model("gpt", Capability::Chat, ModelResidency::Remote, 0, 32) + .behavior(InvokeBehavior::Ok), + ); + let cloud_calls = cloud.invoke_calls.clone(); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + // Default policy is strict for named requests. + let err = engine.invoke(chat_request(Some("qwen"))).await.unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + assert_eq!(cloud_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn ten_concurrent_requests_respect_concurrency_limit() { + let provider = Arc::new( + TestProvider::new("cloud", ProviderKind::OpenAiCompatible) + .with_model("m", Capability::Chat, ModelResidency::Remote, 0, 4) + .behavior(InvokeBehavior::SlowOkMs(40)), + ); + let max_seen = provider.max_in_flight.clone(); + let engine = build_engine(vec![provider], 1000, 0); + engine.refresh_resources().await; + + let mut handles = Vec::new(); + for _ in 0..10 { + let engine = Arc::clone(&engine); + handles.push(tokio::spawn(async move { + engine.invoke(chat_request(None)).await + })); + } + for h in handles { + assert!(h.await.unwrap().is_ok()); + } + // The semaphore must have capped concurrency at the model's limit. + assert!( + max_seen.load(Ordering::SeqCst) <= 4, + "over-admitted requests" + ); + } + #[tokio::test] async fn memory_pressure_evicts_lru_idle_model() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); // 100 MB budget + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 0); // 100 MB budget - // Two resident, idle models at 40 MB each (80 MB used). for name in ["old", "new"] { - let mut card = ModelCard::new("mock", name, Capability::Chat, CostTier::Low); + let mut card = ModelCard::new("ollama", name, Capability::Chat, CostTier::Low); card.residency = ModelResidency::Loaded; card.refresh_status(); let id = card.id.clone(); engine.models.insert(id.clone(), card); - engine.memory.allocate(&id, 40 * 1024 * 1024); + engine.memory.allocate(&id, 40 * MB); } - // Touch "new" so "old" is the least-recently-used victim. - engine.memory.touch("mock/new"); - assert_eq!(engine.memory.used_bytes(), 80 * 1024 * 1024); + engine.memory.touch("ollama/new"); // make "old" the LRU victim + assert_eq!(engine.memory.used_bytes(), 80 * MB); - // Admitting another 40 MB needs 20 MB freed → evict exactly the LRU model. - engine.evict_to_fit("mock/incoming", 40 * 1024 * 1024).await; + // Admitting 40 MB needs 20 MB freed → evict exactly the LRU model. + engine + .admit_memory("ollama/incoming", 40 * MB) + .await + .unwrap(); - assert!(engine.memory.can_fit(40 * 1024 * 1024)); assert_eq!( - engine.models.get("mock/old").unwrap().residency, + engine.models.get("ollama/old").unwrap().residency, ModelResidency::Unloaded ); assert_eq!( - engine.models.get("mock/new").unwrap().residency, + engine.models.get("ollama/new").unwrap().residency, ModelResidency::Loaded ); + assert!(engine.memory.lease_count("ollama/incoming") == 0); } #[tokio::test] - async fn busy_models_are_never_evicted() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); // 100 MB budget + async fn leased_models_are_never_evicted() { + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 0); - // A single resident model with an in-flight request, filling the budget. - let mut card = ModelCard::new("mock", "busy", Capability::Chat, CostTier::Low); + let mut card = ModelCard::new("ollama", "busy", Capability::Chat, CostTier::Low); card.residency = ModelResidency::Loaded; - card.execution.active_requests = 1; card.refresh_status(); let id = card.id.clone(); engine.models.insert(id.clone(), card); - engine.memory.allocate(&id, 90 * 1024 * 1024); + engine.memory.allocate(&id, 90 * MB); + engine.memory.lease(&id); // in-flight inference + + // Cannot evict the only (leased) model → memory pressure. + let err = engine + .admit_memory("ollama/incoming", 40 * MB) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::MemoryPressure { .. })); + assert_eq!( + engine.models.get("ollama/busy").unwrap().residency, + ModelResidency::Loaded + ); + } + + #[tokio::test] + async fn load_failure_rolls_back_reservation() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 10 * MB, + 1, + ) + .failing_load(), + ); + let engine = build_engine(vec![provider], 100, 0); + engine.refresh_resources().await; + + let err = engine.invoke(chat_request(None)).await.unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + // The reservation must have been rolled back. + assert_eq!(engine.memory.used_bytes(), 0); + assert_eq!( + engine.models.get("ollama/qwen").unwrap().residency, + ModelResidency::Unloaded + ); + } - // It cannot be evicted, so the budget stays full. - engine.evict_to_fit("mock/incoming", 40 * 1024 * 1024).await; + #[tokio::test] + async fn successful_load_reserves_and_records_memory() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 20 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 100, 0); + engine.refresh_resources().await; - assert!(!engine.memory.can_fit(40 * 1024 * 1024)); + engine.invoke(chat_request(None)).await.unwrap(); + assert_eq!(engine.memory.used_bytes(), 20 * MB); assert_eq!( - engine.models.get("mock/busy").unwrap().residency, + engine.models.get("ollama/qwen").unwrap().residency, ModelResidency::Loaded ); + // A load event must be in the lifecycle history. + assert!( + engine + .lifecycle_history() + .iter() + .any(|r| r.model_id == "ollama/qwen" && r.event == "load") + ); + } + + #[tokio::test] + async fn loading_model_is_not_evicted() { + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 0); + + // A model whose load is in flight, holding most of the budget. + let mut card = ModelCard::new("ollama", "loading", Capability::Chat, CostTier::Low); + card.residency = ModelResidency::Loading; + card.refresh_status(); + let id = card.id.clone(); + engine.models.insert(id.clone(), card); + engine.memory.allocate(&id, 90 * MB); + + // It must be protected, so admission cannot satisfy the request. + assert!(engine.eviction_protected_ids().contains(&id)); + let err = engine + .admit_memory("ollama/incoming", 40 * MB) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::MemoryPressure { .. })); + assert_eq!( + engine.models.get(&id).unwrap().residency, + ModelResidency::Loading + ); + } + + #[tokio::test] + async fn rediscovery_releases_reservation_for_unloaded_model() { + // The provider reports the model as Unloaded on discovery. + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 100, 0); + + // The engine currently believes the model is loaded and reserved. + let mut card = ModelCard::new("ollama", "qwen", Capability::Chat, CostTier::Low); + card.residency = ModelResidency::Loaded; + card.refresh_status(); + engine.models.insert(card.id.clone(), card); + engine.memory.allocate("ollama/qwen", 10 * MB); + assert_eq!(engine.memory.used_bytes(), 10 * MB); + + // Rediscovery reconciles: the backend says Unloaded, so the reservation + // is released to match reality. + engine.refresh_resources().await; + assert_eq!( + engine.models.get("ollama/qwen").unwrap().residency, + ModelResidency::Unloaded + ); + assert_eq!(engine.memory.used_bytes(), 0); + } + + async fn wait_until(mut cond: impl FnMut() -> bool) -> bool { + // Up to ~4s, returning as soon as the condition holds. + for _ in 0..400 { + if cond() { + return true; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + cond() + } + + fn scoped_request(cap: Capability, app: &str, session: &str) -> InferenceRequest { + let mut r = request(cap, None, FallbackPolicy::default()); + r.app_id = Some(app.into()); + r.session_id = Some(session.into()); + r + } + + #[tokio::test] + async fn hint_warms_next_model_concurrently() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model("qwen", Capability::Chat, ModelResidency::Loaded, 10 * MB, 1) + .with_model( + "kokoro", + Capability::Tts, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 1000, 0); + engine.refresh_resources().await; + + let mut req = chat_request(None); + req.hint_next = Some("tts".into()); + engine.invoke(req).await.unwrap(); + + let engine_ref = Arc::clone(&engine); + let warmed = wait_until(move || engine_ref.is_resident("ollama/kokoro")).await; + assert!(warmed, "the hinted TTS model should have been warmed"); + assert!(engine.preflight_stats().warms_started >= 1); + } + + #[tokio::test] + async fn subscription_warms_and_protects_model() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "kokoro", + Capability::Tts, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 1000, 1); + engine.refresh_resources().await; + + engine.subscribe(Some("mofa-fm".into()), None, vec![Capability::Tts], None); + + let engine_ref = Arc::clone(&engine); + let warmed = wait_until(move || engine_ref.is_resident("ollama/kokoro")).await; + assert!(warmed, "subscription should warm its capability's model"); + // A subscribed resident model is protected from eviction. + assert!( + engine + .eviction_protected_ids() + .contains(&"ollama/kokoro".to_string()) + ); + } + + #[tokio::test] + async fn history_predicts_and_confirms_hits() { + let chat = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Loaded, + 10 * MB, + 4, + ), + ); + let tts = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible).with_model( + "tts-1", + Capability::Tts, + ModelResidency::Remote, + 0, + 32, + ), + ); + let engine = build_engine(vec![chat, tts], 1000, 0); + engine.refresh_resources().await; + + // Build a Chat → Tts habit within one session. + for _ in 0..5 { + engine + .invoke(scoped_request(Capability::Chat, "fm", "s1")) + .await + .unwrap(); + engine + .invoke(scoped_request(Capability::Tts, "fm", "s1")) + .await + .unwrap(); + } + + let stats = engine.preflight_stats(); + assert!(stats.predictions > 0, "history should produce predictions"); + assert!(stats.hits > 0, "predictions should be confirmed as hits"); + } + + /// The headline article-to-podcast flow across Stages 3-5: a hinted chat + /// call warms TTS concurrently, the follow-up TTS call is served by the + /// already-warm model, and both then unload on idle. + #[tokio::test] + async fn article_to_podcast_flow() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 20 * MB, + 1, + ) + .with_model( + "kokoro", + Capability::Tts, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 1000, 1); + engine.refresh_resources().await; + + // Step 1: LLM translation, hinting that TTS is next. + let mut chat = scoped_request(Capability::Chat, "mofa-fm", "s1"); + chat.hint_next = Some("tts".into()); + let chat_resp = engine.invoke(chat).await.unwrap(); + assert_eq!(chat_resp.model_used, "qwen"); + + // The hint should have warmed kokoro concurrently. + let engine_ref = Arc::clone(&engine); + assert!( + wait_until(move || engine_ref.is_resident("ollama/kokoro")).await, + "TTS should be warmed by the hint before it is requested" + ); + + // Step 2: TTS synthesis — served by the already-resident model. + let tts_resp = engine + .invoke(scoped_request(Capability::Tts, "mofa-fm", "s1")) + .await + .unwrap(); + assert_eq!(tts_resp.model_used, "kokoro"); + assert!(!tts_resp.fallback_used); + assert!(engine.preflight_stats().warms_started >= 1); + + // Step 3: both models unload once idle, releasing memory. + let engine_ref = Arc::clone(&engine); + assert!( + wait_until(move || engine_ref.memory.used_bytes() == 0).await, + "idle models should unload and release memory" + ); + } + + #[tokio::test] + async fn idle_sweep_unloads_stale_models() { + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 1); // 1s idle timeout + + let mut card = ModelCard::new("ollama", "qwen", Capability::Chat, CostTier::Low); + card.residency = ModelResidency::Loaded; + card.refresh_status(); + let id = card.id.clone(); + engine.models.insert(id.clone(), card); + engine.memory.allocate(&id, 30 * MB); + + // The idle timeout uses monotonic last-access; wait past it then sweep. + tokio::time::sleep(Duration::from_millis(1100)).await; + engine.run_idle_sweep().await; + + assert_eq!( + engine.models.get(&id).unwrap().residency, + ModelResidency::Unloaded + ); + assert_eq!(engine.memory.used_bytes(), 0); + assert!( + engine + .lifecycle_history() + .iter() + .any(|r| r.event == "idle_unload") + ); + } + + async fn collect_stream(mut rx: mpsc::Receiver) -> Vec { + let mut chunks = Vec::new(); + while let Some(c) = rx.recv().await { + chunks.push(c); + } + chunks + } + + /// A provider that emits several real text deltas, to prove the engine + /// forwards incremental output in order (not just the compatibility path). + struct MultiDeltaProvider { + name: String, + deltas: Vec, + } + + #[async_trait] + impl Provider for MultiDeltaProvider { + fn name(&self) -> &str { + &self.name + } + fn kind(&self) -> ProviderKind { + ProviderKind::Ollama + } + fn features(&self) -> Vec { + vec![BackendFeature::Discovery, BackendFeature::Streaming] + } + async fn discover(&self) -> Result, EngineError> { + let mut card = ModelCard::new(&self.name, "streamer", Capability::Chat, CostTier::Free); + card.residency = ModelResidency::Loaded; + card.refresh_status(); + Ok(vec![card]) + } + async fn health(&self) -> Result { + Ok(BackendHealth::Healthy) + } + async fn load(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: model_id.into(), + residency: ModelResidency::Loaded, + memory_bytes: Some(0), + changed: true, + }) + } + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: model_id.into(), + residency: ModelResidency::Unloaded, + memory_bytes: Some(0), + changed: true, + }) + } + async fn invoke( + &self, + model_id: &str, + request: &InferenceRequest, + ) -> Result { + Ok(InferenceResponse { + text: Some(self.deltas.concat()), + file: None, + model_used: mofa_kernel::model_id_name(model_id).into(), + provider: self.name.clone(), + duration_ms: 1, + request_id: request.request_id.clone(), + tokens_used: Some(self.deltas.len() as u32), + fallback_used: false, + routing_reason: None, + }) + } + async fn stream( + &self, + model_id: &str, + request: &InferenceRequest, + sink: mofa_kernel::StreamSink, + ) -> Result { + for d in &self.deltas { + let _ = sink.send(d.clone()).await; + } + self.invoke(model_id, request).await + } + } + + #[tokio::test] + async fn stream_emits_started_text_completed_in_order() { + // Default (compat) streaming: one Text chunk carrying the full output. + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Loaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 1000, 0); + engine.refresh_resources().await; + + let chunks = collect_stream(engine.invoke_stream(chat_request(None))).await; + assert!(matches!(chunks.first(), Some(StreamChunk::Started { .. }))); + assert!(matches!(chunks.last(), Some(StreamChunk::Completed { .. }))); + let text: String = chunks + .iter() + .filter_map(|c| match c { + StreamChunk::Text { delta } => Some(delta.clone()), + _ => None, + }) + .collect(); + assert_eq!(text, "ok"); + } + + #[tokio::test] + async fn stream_forwards_incremental_deltas_in_order() { + let provider = Arc::new(MultiDeltaProvider { + name: "ollama".into(), + deltas: vec!["Hello".into(), ", ".into(), "world".into()], + }); + let engine = build_engine(vec![provider], 1000, 0); + engine.refresh_resources().await; + + let chunks = collect_stream(engine.invoke_stream(chat_request(None))).await; + let deltas: Vec = chunks + .iter() + .filter_map(|c| match c { + StreamChunk::Text { delta } => Some(delta.clone()), + _ => None, + }) + .collect(); + assert_eq!(deltas, vec!["Hello", ", ", "world"]); + // Completed reports the token count the provider aggregated. + match chunks.last() { + Some(StreamChunk::Completed { tokens_used, .. }) => { + assert_eq!(*tokens_used, Some(3)); + } + other => panic!("expected Completed, got {other:?}"), + } + } + + #[tokio::test] + async fn input_path_allowlist_blocks_paths_outside_roots() { + let root = tempfile::tempdir().unwrap(); + let allowed = root.path().join("clip.wav"); + std::fs::write(&allowed, b"x").unwrap(); + let outside = std::env::temp_dir().join("mofa_outside_allowlist.wav"); + std::fs::write(&outside, b"x").unwrap(); + + let config = EngineConfig { + security: SecurityConfig { + input_roots: vec![root.path().to_string_lossy().to_string()], + }, + ..minimal_config() + }; + let engine = Engine::new(config).await; + + // A path outside the allowlist is rejected before any routing happens. + let mut req = request(Capability::Asr, None, FallbackPolicy::default()); + req.input_file = Some(outside.to_string_lossy().to_string()); + assert!(matches!( + engine.invoke(req).await, + Err(EngineError::InvalidRequest(_)) + )); + + // A path inside the allowlist passes the check (then fails later only + // because no ASR model is registered). + let mut ok = request(Capability::Asr, None, FallbackPolicy::default()); + ok.input_file = Some(allowed.to_string_lossy().to_string()); + assert!(matches!( + engine.invoke(ok).await, + Err(EngineError::NoCapableModel(_)) + )); + } + + #[tokio::test] + async fn stream_reports_no_capable_model_as_error_chunk() { + let engine = build_engine(vec![], 100, 0); + let chunks = collect_stream(engine.invoke_stream(chat_request(None))).await; + assert_eq!(chunks.len(), 1); + match &chunks[0] { + StreamChunk::Error(info) => { + assert_eq!(info.code, mofa_kernel::ErrorCode::NoCapableModel); + } + other => panic!("expected Error chunk, got {other:?}"), + } } } diff --git a/mofa-engine-core/src/lib.rs b/mofa-engine-core/src/lib.rs index f747665..9943486 100644 --- a/mofa-engine-core/src/lib.rs +++ b/mofa-engine-core/src/lib.rs @@ -4,13 +4,16 @@ //! memory management, circuit breaker, preflight prediction, //! and the main `Engine` orchestrator. +pub mod artifacts; pub mod backends; pub mod circuit_breaker; pub mod config; pub mod engine; pub mod memory; +pub mod metrics; pub mod preflight; pub mod router; +pub mod subscription; pub use config::EngineConfig; pub use engine::Engine; diff --git a/mofa-engine-core/src/memory.rs b/mofa-engine-core/src/memory.rs index 6abf8da..c874fe3 100644 --- a/mofa-engine-core/src/memory.rs +++ b/mofa-engine-core/src/memory.rs @@ -1,23 +1,63 @@ -//! Memory manager with LRU eviction. +//! Reservation-based memory manager with LRU eviction. //! -//! Tracks per-model memory allocations, detects real system memory, -//! and evicts least-recently-used models when under pressure. +//! Memory is admitted *before* a model is loaded: callers atomically reserve +//! capacity, evict idle models if needed, and only then ask the backend to +//! load. This prevents the time-of-check/time-of-use race where two concurrent +//! loads each observe enough free memory and together overcommit the device. +//! +//! Three quantities are tracked per model: +//! - **reserved** — the bytes admitted against the budget (a conservative +//! estimate, later reconciled toward observed usage); +//! - **observed** — actual usage reported by the backend, when available; +//! - **leases** — in-flight inferences; a model with active leases is never +//! evicted, even under pressure. +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Mutex; -use std::time::Instant; +use std::time::{Duration, Instant}; -/// Tracks memory allocations and enforces a budget. +/// Tracks per-model memory reservations and enforces a global budget. pub struct MemoryManager { - /// Total budget in bytes + /// Total budget in bytes. budget_bytes: u64, - /// Per-model allocations: model_id → (bytes, last_access) + /// Per-model allocations. allocations: Mutex>, } struct Allocation { - bytes: u64, + /// Bytes admitted against the budget. Accounting always uses this figure. + reserved_bytes: u64, + /// Actual usage reported by the backend, when known. + observed_bytes: Option, + /// Monotonic time of the last access, for LRU and idle-timeout decisions. last_access: Instant, + /// Active inference leases. A model with `leases > 0` is never evicted. + leases: u32, +} + +impl Allocation { + fn new(reserved_bytes: u64) -> Self { + Self { + reserved_bytes, + observed_bytes: None, + last_access: Instant::now(), + leases: 0, + } + } +} + +/// A point-in-time view of one model's memory accounting. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AllocationSnapshot { + /// Model identifier. + pub model_id: String, + /// Bytes reserved against the budget. + pub reserved_bytes: u64, + /// Bytes observed by the backend, if reported. + pub observed_bytes: Option, + /// Active inference leases. + pub leases: u32, } impl MemoryManager { @@ -26,7 +66,7 @@ impl MemoryManager { /// If `budget_mb` is `None`, auto-detect from system RAM (use 70% of total). pub fn new(budget_mb: Option) -> Self { let budget_bytes = match budget_mb { - Some(mb) => mb * 1024 * 1024, + Some(mb) => mb.saturating_mul(1024 * 1024), None => Self::detect_system_memory(), }; @@ -45,69 +85,164 @@ impl MemoryManager { (total as f64 * 0.7) as u64 } + fn lock(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.allocations.lock().unwrap_or_else(|e| e.into_inner()) + } + /// Total budget in bytes. pub fn budget_bytes(&self) -> u64 { self.budget_bytes } - /// Currently used bytes. + /// Currently reserved bytes across all models. pub fn used_bytes(&self) -> u64 { - let allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); - allocs.values().map(|a| a.bytes).sum() + self.lock().values().map(|a| a.reserved_bytes).sum() } - /// Available bytes. + /// Bytes still available within the budget. pub fn available_bytes(&self) -> u64 { self.budget_bytes.saturating_sub(self.used_bytes()) } - /// Record that a model is using memory. Updates last-access time. + /// Whether `needed_bytes` currently fits within the budget without eviction. + pub fn can_fit(&self, needed_bytes: u64) -> bool { + self.available_bytes() >= needed_bytes + } + + /// Atomically reserve `bytes` for `model_id` if it fits within the budget. + /// + /// The check and the insert happen under a single lock, so concurrent + /// reservations can never jointly exceed the budget. Re-reserving a model + /// that already holds an allocation accounts only for the *delta*, and + /// preserves its lease count. Returns `true` if the reservation succeeded. + pub fn try_reserve(&self, model_id: &str, bytes: u64) -> bool { + let mut allocs = self.lock(); + let used_by_others: u64 = allocs + .iter() + .filter(|(id, _)| id.as_str() != model_id) + .map(|(_, a)| a.reserved_bytes) + .sum(); + if used_by_others.saturating_add(bytes) > self.budget_bytes { + return false; + } + match allocs.get_mut(model_id) { + Some(existing) => { + existing.reserved_bytes = bytes; + existing.last_access = Instant::now(); + } + None => { + allocs.insert(model_id.to_string(), Allocation::new(bytes)); + } + } + true + } + + /// Unconditionally record a reservation, ignoring the budget. + /// + /// Intended for tests and for accounting state the backend reports as already + /// resident. Prefer [`try_reserve`](Self::try_reserve) for admission. pub fn allocate(&self, model_id: &str, bytes: u64) { - let mut allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); - allocs.insert( - model_id.to_string(), - Allocation { - bytes, - last_access: Instant::now(), - }, - ); + let mut allocs = self.lock(); + match allocs.get_mut(model_id) { + Some(existing) => { + existing.reserved_bytes = bytes; + existing.last_access = Instant::now(); + } + None => { + allocs.insert(model_id.to_string(), Allocation::new(bytes)); + } + } } - /// Remove a model's allocation. + /// Release a model's reservation entirely. pub fn deallocate(&self, model_id: &str) { - let mut allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); - allocs.remove(model_id); + self.lock().remove(model_id); } - /// Touch a model (update its last-access time). - pub fn touch(&self, model_id: &str) { - let mut allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); + /// Reconcile a reservation with memory the backend actually reports. + /// + /// The observed value is recorded for reporting, and the reservation is + /// raised to it when the model uses more than estimated, so accounting never + /// undercounts real usage. + pub fn reconcile(&self, model_id: &str, observed_bytes: u64) { + let mut allocs = self.lock(); if let Some(a) = allocs.get_mut(model_id) { + a.observed_bytes = Some(observed_bytes); + a.reserved_bytes = a.reserved_bytes.max(observed_bytes); + } + } + + /// Update a model's last-access time (for LRU and idle-timeout decisions). + pub fn touch(&self, model_id: &str) { + if let Some(a) = self.lock().get_mut(model_id) { a.last_access = Instant::now(); } } - /// Check if there is enough room for `needed_bytes`. - pub fn can_fit(&self, needed_bytes: u64) -> bool { - self.available_bytes() >= needed_bytes + /// Acquire an inference lease, protecting the model from eviction. + /// + /// No-op for models without an allocation (e.g. cloud-backed models, which + /// hold no local memory and are never eviction candidates). + pub fn lease(&self, model_id: &str) { + if let Some(a) = self.lock().get_mut(model_id) { + a.leases = a.leases.saturating_add(1); + a.last_access = Instant::now(); + } + } + + /// Release a previously acquired inference lease. + pub fn release_lease(&self, model_id: &str) { + if let Some(a) = self.lock().get_mut(model_id) { + a.leases = a.leases.saturating_sub(1); + a.last_access = Instant::now(); + } } - /// Find the least-recently-used model that can be evicted. + /// Number of active leases on a model. + pub fn lease_count(&self, model_id: &str) -> u32 { + self.lock().get(model_id).map(|a| a.leases).unwrap_or(0) + } + + /// Find the least-recently-used evictable model. /// - /// `protected` is a set of model IDs that must not be evicted (e.g. resident models). + /// Models in `protected`, and any model holding an active lease, are never + /// returned. pub fn lru_candidate(&self, protected: &[String]) -> Option { - let allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); - allocs + self.lock() .iter() - .filter(|(id, _)| !protected.contains(id)) + .filter(|(id, a)| a.leases == 0 && !protected.iter().any(|p| p == *id)) .min_by_key(|(_, a)| a.last_access) .map(|(id, _)| id.clone()) } - /// Return a snapshot of all allocations (model_id, bytes). - pub fn snapshot(&self) -> Vec<(String, u64)> { - let allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); - allocs.iter().map(|(id, a)| (id.clone(), a.bytes)).collect() + /// Return every model whose idle time meets or exceeds `idle`, excluding + /// leased and `protected` models. Used by the idle-timeout sweep. + pub fn idle_candidates(&self, idle: Duration, protected: &[String]) -> Vec { + self.lock() + .iter() + .filter(|(id, a)| { + a.leases == 0 + && !protected.iter().any(|p| p == *id) + && a.last_access.elapsed() >= idle + }) + .map(|(id, _)| id.clone()) + .collect() + } + + /// Snapshot of all current allocations, sorted by model id for determinism. + pub fn snapshot(&self) -> Vec { + let mut out: Vec = self + .lock() + .iter() + .map(|(id, a)| AllocationSnapshot { + model_id: id.clone(), + reserved_bytes: a.reserved_bytes, + observed_bytes: a.observed_bytes, + leases: a.leases, + }) + .collect(); + out.sort_by(|a, b| a.model_id.cmp(&b.model_id)); + out } } @@ -115,65 +250,130 @@ impl MemoryManager { mod tests { use super::*; + const MB: u64 = 1024 * 1024; + #[test] fn allocate_and_track() { - let mm = MemoryManager::new(Some(100)); // 100 MB - assert_eq!(mm.budget_bytes(), 100 * 1024 * 1024); + let mm = MemoryManager::new(Some(100)); + assert_eq!(mm.budget_bytes(), 100 * MB); assert_eq!(mm.used_bytes(), 0); - mm.allocate("model-a", 10 * 1024 * 1024); - assert_eq!(mm.used_bytes(), 10 * 1024 * 1024); + mm.allocate("model-a", 10 * MB); + assert_eq!(mm.used_bytes(), 10 * MB); - mm.allocate("model-b", 20 * 1024 * 1024); - assert_eq!(mm.used_bytes(), 30 * 1024 * 1024); + mm.allocate("model-b", 20 * MB); + assert_eq!(mm.used_bytes(), 30 * MB); mm.deallocate("model-a"); - assert_eq!(mm.used_bytes(), 20 * 1024 * 1024); + assert_eq!(mm.used_bytes(), 20 * MB); + } + + #[test] + fn try_reserve_respects_budget() { + let mm = MemoryManager::new(Some(100)); + assert!(mm.try_reserve("a", 60 * MB)); + // 60 used, only 40 left → a 50 MB reservation must fail. + assert!(!mm.try_reserve("b", 50 * MB)); + assert!(mm.try_reserve("b", 40 * MB)); + assert_eq!(mm.used_bytes(), 100 * MB); + } + + #[test] + fn try_reserve_same_model_accounts_delta_not_sum() { + let mm = MemoryManager::new(Some(100)); + assert!(mm.try_reserve("a", 60 * MB)); + // Re-reserving the same model at 80 MB replaces, not adds → fits. + assert!(mm.try_reserve("a", 80 * MB)); + assert_eq!(mm.used_bytes(), 80 * MB); + } + + #[test] + fn reconcile_raises_reservation_to_observed() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("a", 10 * MB); + mm.reconcile("a", 25 * MB); + assert_eq!(mm.used_bytes(), 25 * MB); + // A lower observed value does not shrink the conservative reservation. + mm.reconcile("a", 5 * MB); + assert_eq!(mm.used_bytes(), 25 * MB); } #[test] fn can_fit_check() { let mm = MemoryManager::new(Some(100)); - mm.allocate("m1", 90 * 1024 * 1024); - assert!(!mm.can_fit(20 * 1024 * 1024)); - assert!(mm.can_fit(10 * 1024 * 1024)); + mm.allocate("m1", 90 * MB); + assert!(!mm.can_fit(20 * MB)); + assert!(mm.can_fit(10 * MB)); } #[test] fn lru_eviction_order() { let mm = MemoryManager::new(Some(100)); - mm.allocate("old", 10 * 1024 * 1024); - // tiny delay so timestamps differ - std::thread::sleep(std::time::Duration::from_millis(10)); - mm.allocate("new", 10 * 1024 * 1024); - - let candidate = mm.lru_candidate(&[]).unwrap(); - assert_eq!(candidate, "old"); + mm.allocate("old", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("new", 10 * MB); + assert_eq!(mm.lru_candidate(&[]).unwrap(), "old"); } #[test] fn lru_respects_protected() { let mm = MemoryManager::new(Some(100)); - mm.allocate("protected", 10 * 1024 * 1024); - std::thread::sleep(std::time::Duration::from_millis(10)); - mm.allocate("evictable", 10 * 1024 * 1024); + mm.allocate("protected", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("evictable", 10 * MB); + assert_eq!( + mm.lru_candidate(&["protected".into()]).unwrap(), + "evictable" + ); + } - let candidate = mm.lru_candidate(&["protected".into()]).unwrap(); - assert_eq!(candidate, "evictable"); + #[test] + fn leased_models_are_never_lru_candidates() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("leased", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("free", 10 * MB); + mm.lease("leased"); + // "leased" is older but is protected by its lease. + assert_eq!(mm.lru_candidate(&[]).unwrap(), "free"); + mm.release_lease("leased"); + assert_eq!(mm.lease_count("leased"), 0); } #[test] fn touch_updates_access_time() { let mm = MemoryManager::new(Some(100)); - mm.allocate("a", 10 * 1024 * 1024); - std::thread::sleep(std::time::Duration::from_millis(10)); - mm.allocate("b", 10 * 1024 * 1024); - - // 'a' is older, but touch it to make it newer - std::thread::sleep(std::time::Duration::from_millis(10)); + mm.allocate("a", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("b", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); mm.touch("a"); + assert_eq!(mm.lru_candidate(&[]).unwrap(), "b"); + } + + #[test] + fn idle_candidates_reports_only_stale_unleased() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("stale", 10 * MB); + mm.allocate("leased_stale", 10 * MB); + mm.lease("leased_stale"); + std::thread::sleep(Duration::from_millis(20)); + mm.allocate("fresh", 10 * MB); - let candidate = mm.lru_candidate(&[]).unwrap(); - assert_eq!(candidate, "b"); + let idle = mm.idle_candidates(Duration::from_millis(10), &[]); + assert_eq!(idle, vec!["stale".to_string()]); + } + + #[test] + fn snapshot_is_sorted_and_complete() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("b", 20 * MB); + mm.allocate("a", 10 * MB); + mm.reconcile("a", 15 * MB); + let snap = mm.snapshot(); + assert_eq!(snap.len(), 2); + assert_eq!(snap[0].model_id, "a"); + assert_eq!(snap[0].observed_bytes, Some(15 * MB)); + assert_eq!(snap[1].model_id, "b"); } } diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs new file mode 100644 index 0000000..3eda177 --- /dev/null +++ b/mofa-engine-core/src/metrics.rs @@ -0,0 +1,300 @@ +//! Bounded-cardinality engine metrics and Prometheus text rendering. +//! +//! These counters are the raw signal behind the `/metrics` endpoint. They are +//! deliberately label-free (or labelled only by provider, a small fixed set) so +//! metric cardinality stays bounded regardless of request volume — no per-model, +//! per-request, or user-supplied labels leak into the series. + +use std::sync::atomic::{AtomicU64, Ordering}; + +/// Upper bounds (in milliseconds) for the request-latency histogram buckets. +/// The final `+Inf` bucket is implicit. +const LATENCY_BUCKETS_MS: [u64; 8] = [50, 100, 250, 500, 1000, 2500, 5000, 10000]; + +/// Process-wide engine counters. +#[derive(Default)] +pub struct EngineMetrics { + requests_total: AtomicU64, + requests_failed: AtomicU64, + requests_fallback: AtomicU64, + request_duration_ms_sum: AtomicU64, + /// Cumulative-per-bucket counts plus a trailing `+Inf` bucket. + latency_buckets: [AtomicU64; LATENCY_BUCKETS_MS.len() + 1], + loads_total: AtomicU64, + unloads_total: AtomicU64, + evictions_total: AtomicU64, + idle_unloads_total: AtomicU64, + load_failures_total: AtomicU64, +} + +impl EngineMetrics { + /// Record the outcome of one logical inference request. + pub fn record_request(&self, success: bool, duration_ms: u64, fallback_used: bool) { + self.requests_total.fetch_add(1, Ordering::Relaxed); + if !success { + self.requests_failed.fetch_add(1, Ordering::Relaxed); + } + if fallback_used { + self.requests_fallback.fetch_add(1, Ordering::Relaxed); + } + self.request_duration_ms_sum + .fetch_add(duration_ms, Ordering::Relaxed); + // Increment the first bucket whose upper bound covers this observation. + let mut idx = LATENCY_BUCKETS_MS.len(); + for (i, bound) in LATENCY_BUCKETS_MS.iter().enumerate() { + if duration_ms <= *bound { + idx = i; + break; + } + } + self.latency_buckets[idx].fetch_add(1, Ordering::Relaxed); + } + + /// Fold a lifecycle event (as recorded in the lifecycle history) into the + /// relevant lifecycle counter. + pub fn record_lifecycle(&self, event: &str) { + match event { + "load" => &self.loads_total, + "unload" => &self.unloads_total, + "evict" => &self.evictions_total, + "idle_unload" => &self.idle_unloads_total, + "load_failed" | "load_timeout" => &self.load_failures_total, + _ => return, + } + .fetch_add(1, Ordering::Relaxed); + } + + /// Total requests observed. + pub fn requests_total(&self) -> u64 { + self.requests_total.load(Ordering::Relaxed) + } + + /// Render the Prometheus text-exposition body for these counters. + /// + /// `gauges` are point-in-time values pulled from the engine (memory, model + /// counts, preflight, and per-provider routability) sampled at scrape time. + pub fn render_prometheus(&self, gauges: &MetricsGauges) -> String { + let mut out = String::with_capacity(2048); + + let counter = |out: &mut String, name: &str, help: &str, value: u64| { + out.push_str(&format!( + "# HELP {name} {help}\n# TYPE {name} counter\n{name} {value}\n" + )); + }; + let gauge = |out: &mut String, name: &str, help: &str, value: u64| { + out.push_str(&format!( + "# HELP {name} {help}\n# TYPE {name} gauge\n{name} {value}\n" + )); + }; + + counter( + &mut out, + "mofa_requests_total", + "Total inference requests handled.", + self.requests_total.load(Ordering::Relaxed), + ); + counter( + &mut out, + "mofa_requests_failed_total", + "Inference requests that returned an error.", + self.requests_failed.load(Ordering::Relaxed), + ); + counter( + &mut out, + "mofa_requests_fallback_total", + "Requests served by a fallback candidate.", + self.requests_fallback.load(Ordering::Relaxed), + ); + counter( + &mut out, + "mofa_model_loads_total", + "Model load operations.", + self.loads_total.load(Ordering::Relaxed), + ); + counter( + &mut out, + "mofa_model_unloads_total", + "Explicit model unload operations.", + self.unloads_total.load(Ordering::Relaxed), + ); + counter( + &mut out, + "mofa_model_evictions_total", + "Models evicted under memory pressure.", + self.evictions_total.load(Ordering::Relaxed), + ); + counter( + &mut out, + "mofa_model_idle_unloads_total", + "Models unloaded by the idle sweep.", + self.idle_unloads_total.load(Ordering::Relaxed), + ); + counter( + &mut out, + "mofa_model_load_failures_total", + "Failed or timed-out model loads.", + self.load_failures_total.load(Ordering::Relaxed), + ); + + // Latency histogram (cumulative buckets, per Prometheus convention). + out.push_str( + "# HELP mofa_request_duration_ms Request latency in milliseconds.\n\ + # TYPE mofa_request_duration_ms histogram\n", + ); + let mut cumulative = 0u64; + for (i, bound) in LATENCY_BUCKETS_MS.iter().enumerate() { + cumulative += self.latency_buckets[i].load(Ordering::Relaxed); + out.push_str(&format!( + "mofa_request_duration_ms_bucket{{le=\"{bound}\"}} {cumulative}\n" + )); + } + cumulative += self.latency_buckets[LATENCY_BUCKETS_MS.len()].load(Ordering::Relaxed); + out.push_str(&format!( + "mofa_request_duration_ms_bucket{{le=\"+Inf\"}} {cumulative}\n" + )); + out.push_str(&format!( + "mofa_request_duration_ms_sum {}\n", + self.request_duration_ms_sum.load(Ordering::Relaxed) + )); + out.push_str(&format!("mofa_request_duration_ms_count {cumulative}\n")); + + // Point-in-time gauges from the engine. + gauge( + &mut out, + "mofa_models_total", + "Known models in the registry.", + gauges.models_total, + ); + gauge( + &mut out, + "mofa_models_loaded", + "Models currently resident.", + gauges.models_loaded, + ); + gauge( + &mut out, + "mofa_memory_used_bytes", + "Reserved model memory in bytes.", + gauges.memory_used_bytes, + ); + gauge( + &mut out, + "mofa_memory_budget_bytes", + "Total model memory budget in bytes.", + gauges.memory_budget_bytes, + ); + gauge( + &mut out, + "mofa_preflight_warms_total", + "Speculative warm tasks started.", + gauges.preflight_warms_started, + ); + gauge( + &mut out, + "mofa_preflight_hits_total", + "Confirmed predictive-warm hits.", + gauges.preflight_hits, + ); + + // One bounded-cardinality gauge family labelled by provider. + out.push_str( + "# HELP mofa_provider_up Provider routability (1 = routable).\n\ + # TYPE mofa_provider_up gauge\n", + ); + for (provider, up) in &gauges.provider_up { + out.push_str(&format!( + "mofa_provider_up{{provider=\"{}\"}} {}\n", + escape_label_value(provider), + u8::from(*up) + )); + } + + out + } +} + +/// Escape a Prometheus label value per the exposition format: backslash, +/// double-quote, and newline must be escaped so a provider name cannot produce +/// malformed output. +fn escape_label_value(value: &str) -> String { + value + .replace('\\', "\\\\") + .replace('"', "\\\"") + .replace('\n', "\\n") +} + +/// Point-in-time gauge values sampled from the engine at scrape time. +#[derive(Debug, Default, Clone)] +pub struct MetricsGauges { + /// Known models in the registry. + pub models_total: u64, + /// Models currently resident. + pub models_loaded: u64, + /// Reserved model memory in bytes. + pub memory_used_bytes: u64, + /// Total model memory budget in bytes. + pub memory_budget_bytes: u64, + /// Speculative warm tasks started. + pub preflight_warms_started: u64, + /// Confirmed predictive-warm hits. + pub preflight_hits: u64, + /// Per-provider routability, `(provider_name, is_routable)`. + pub provider_up: Vec<(String, bool)>, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn histogram_and_counters_render() { + let m = EngineMetrics::default(); + m.record_request(true, 30, false); // → le=50 bucket + m.record_request(false, 1500, true); // → le=2500 bucket, failed + fallback + m.record_lifecycle("load"); + m.record_lifecycle("idle_unload"); + + let gauges = MetricsGauges { + models_total: 3, + models_loaded: 1, + memory_budget_bytes: 1000, + provider_up: vec![("ollama".into(), true)], + ..Default::default() + }; + let text = m.render_prometheus(&gauges); + + assert!(text.contains("mofa_requests_total 2")); + assert!(text.contains("mofa_requests_failed_total 1")); + assert!(text.contains("mofa_requests_fallback_total 1")); + assert!(text.contains("mofa_model_loads_total 1")); + assert!(text.contains("mofa_model_idle_unloads_total 1")); + // Cumulative histogram: both observations are <= +Inf. + assert!(text.contains("mofa_request_duration_ms_bucket{le=\"+Inf\"} 2")); + assert!(text.contains("mofa_request_duration_ms_count 2")); + assert!(text.contains("mofa_request_duration_ms_sum 1530")); + assert!(text.contains("mofa_provider_up{provider=\"ollama\"} 1")); + } + + #[test] + fn provider_label_values_are_escaped() { + let m = EngineMetrics::default(); + let gauges = MetricsGauges { + provider_up: vec![("we\"ird\\name".into(), true)], + ..Default::default() + }; + let text = m.render_prometheus(&gauges); + // The quote and backslash must be escaped so the exposition stays valid. + assert!(text.contains(r#"mofa_provider_up{provider="we\"ird\\name"} 1"#)); + } + + #[test] + fn latency_buckets_are_cumulative() { + let m = EngineMetrics::default(); + m.record_request(true, 10, false); // le=50 + m.record_request(true, 80, false); // le=100 + let text = m.render_prometheus(&MetricsGauges::default()); + // le=50 covers 1, le=100 covers both. + assert!(text.contains("mofa_request_duration_ms_bucket{le=\"50\"} 1")); + assert!(text.contains("mofa_request_duration_ms_bucket{le=\"100\"} 2")); + } +} diff --git a/mofa-engine-core/src/preflight.rs b/mofa-engine-core/src/preflight.rs index 4a4a4c3..70012b4 100644 --- a/mofa-engine-core/src/preflight.rs +++ b/mofa-engine-core/src/preflight.rs @@ -1,39 +1,113 @@ -//! Preflight prediction via Markov chain. +//! Preflight prediction via per-scope Markov chains. //! -//! Tracks transitions between capabilities (e.g. Chat → Tts) and -//! predicts the next capability so the engine can pre-warm models. +//! Tracks transitions between capabilities (e.g. `Chat → Tts`) so the engine can +//! pre-warm the next model before it is requested. History is keyed by an +//! application/session **scope** so one application's behavior never pollutes +//! another's predictions, with a global chain as a fallback when a scope has not +//! yet accumulated enough samples. use mofa_kernel::Capability; +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Instant; /// Decay factor applied every `DECAY_INTERVAL` transitions. const DECAY_FACTOR: f64 = 0.95; /// Number of transitions between decay applications. const DECAY_INTERVAL: u64 = 100; -/// Minimum observations before making a prediction. -const MIN_OBSERVATIONS: u64 = 3; +/// Maximum number of distinct (non-global) scopes retained. Scope keys come from +/// caller-supplied app/session IDs, so this bounds memory against an unbounded +/// stream of unique identifiers by evicting the least-recently-used scope. +const MAX_SCOPES: usize = 2048; +/// Scope key under which cross-application history is aggregated. +pub const GLOBAL_SCOPE: &str = "__global__"; -/// Markov chain transition tracker with decay. +/// Per-scope Markov chain transition tracker with decay. pub struct PreflightPredictor { - inner: Mutex, + inner: Mutex>, } -struct PreflightInner { - /// Transition counts: from_capability → (to_capability → count) +struct ScopeChain { + /// Transition counts: from_capability → (to_capability → weight). transitions: HashMap>, - /// Total transitions observed + /// Total transitions observed in this scope. total_transitions: u64, - /// Last capability observed + /// Last capability observed in this scope. last_capability: Option, + /// Last time this scope was recorded into or predicted from (for LRU). + last_access: Instant, +} + +impl ScopeChain { + fn new() -> Self { + Self { + transitions: HashMap::new(), + total_transitions: 0, + last_capability: None, + last_access: Instant::now(), + } + } + + /// Observe a capability used within this scope, recording the `prev → this` + /// transition against the scope's own running order. + fn observe(&mut self, capability: Capability) { + self.last_access = Instant::now(); + if let Some(prev) = self.last_capability { + self.record_transition(prev, capability); + } + self.last_capability = Some(capability); + } + + /// Add a single `from → to` transition to this chain's counts, applying + /// decay on the interval. Unlike [`observe`](Self::observe) it does not + /// touch `last_capability`, so it can mirror a transition that already + /// happened elsewhere (e.g. into the global aggregate) without inventing a + /// spurious edge from whatever this chain saw last. + fn record_transition(&mut self, from: Capability, to: Capability) { + self.last_access = Instant::now(); + *self + .transitions + .entry(from) + .or_default() + .entry(to) + .or_insert(0.0) += 1.0; + self.total_transitions += 1; + + if self.total_transitions.is_multiple_of(DECAY_INTERVAL) { + for counts in self.transitions.values_mut() { + for count in counts.values_mut() { + *count *= DECAY_FACTOR; + } + } + } + } + + fn predict(&self, current: Capability, min_samples: u64) -> Option { + let counts = self.transitions.get(¤t)?; + let total: f64 = counts.values().sum(); + // Compare in f64: `total` may be fractional after decay, and a lossy + // cast to u64 would be an unnecessary truncation. + if total < min_samples as f64 { + return None; + } + counts + .iter() + .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(cap, count)| Prediction { + capability: *cap, + confidence: count / total, + }) + } } /// Prediction result. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq)] pub struct Prediction { - /// Predicted next capability + /// Predicted next capability. pub capability: Capability, - /// Confidence score (0.0 - 1.0) + /// Confidence score (0.0 – 1.0). pub confidence: f64, } @@ -41,63 +115,100 @@ impl PreflightPredictor { /// Create a new predictor. pub fn new() -> Self { Self { - inner: Mutex::new(PreflightInner { - transitions: HashMap::new(), - total_transitions: 0, - last_capability: None, - }), - } - } - - /// Record that a capability was just used. - /// This updates the transition counts from the previous capability. - pub fn record(&self, capability: Capability) { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - - if let Some(prev) = inner.last_capability { - let entry = inner - .transitions - .entry(prev) - .or_default() - .entry(capability) - .or_insert(0.0); - *entry += 1.0; - - inner.total_transitions += 1; - - // Apply decay periodically - if inner.total_transitions.is_multiple_of(DECAY_INTERVAL) { - for counts in inner.transitions.values_mut() { - for count in counts.values_mut() { - *count *= DECAY_FACTOR; - } - } - } + inner: Mutex::new(HashMap::new()), } + } - inner.last_capability = Some(capability); + fn lock(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.inner.lock().unwrap_or_else(|e| e.into_inner()) } - /// Predict the most likely next capability given the current one. + /// Record that a capability was used within `scope`. /// - /// Returns `None` if there are not enough observations. - pub fn predict(&self, current: Capability) -> Option { - let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + /// The observation updates the scope's own chain, and the *same* per-scope + /// transition (if any) is mirrored into the global aggregate. The global + /// chain therefore only ever accumulates transitions that genuinely occurred + /// within some scope — an edge that spanned two different scopes (because + /// their requests interleaved) is never invented. + pub fn record(&self, scope: &str, capability: Capability) { + let mut scopes = self.lock(); + Self::evict_if_full(&mut scopes, scope); - let counts = inner.transitions.get(¤t)?; - let total: f64 = counts.values().sum(); + // Observe into the scope's own chain, capturing the transition edge it + // just formed (its previous capability, if any). + let edge_from = { + let chain = scopes + .entry(scope.to_string()) + .or_insert_with(ScopeChain::new); + let prev = chain.last_capability; + chain.observe(capability); + prev + }; - if (total as u64) < MIN_OBSERVATIONS { - return None; + // Mirror only that real edge into the global aggregate. + if scope != GLOBAL_SCOPE + && let Some(from) = edge_from + { + scopes + .entry(GLOBAL_SCOPE.to_string()) + .or_insert_with(ScopeChain::new) + .record_transition(from, capability); } + } - counts + /// Evict the least-recently-used non-global scope when at capacity and the + /// incoming scope is new. The global scope is never evicted. + fn evict_if_full(scopes: &mut HashMap, incoming: &str) { + if scopes.len() < MAX_SCOPES || scopes.contains_key(incoming) { + return; + } + if let Some(victim) = scopes .iter() - .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(cap, count)| Prediction { - capability: *cap, - confidence: count / total, - }) + .filter(|(k, _)| k.as_str() != GLOBAL_SCOPE) + .min_by_key(|(_, c)| c.last_access) + .map(|(k, _)| k.clone()) + { + scopes.remove(&victim); + } + } + + /// Predict the next capability for `scope` given `current`. + /// + /// Tries the scope's own chain first; if it lacks enough samples or + /// confidence, falls back to the global chain. Returns `None` when neither + /// clears `min_samples` and `confidence_threshold`. + pub fn predict( + &self, + scope: &str, + current: Capability, + min_samples: u64, + confidence_threshold: f64, + ) -> Option { + let mut scopes = self.lock(); + // Touch the scope so an actively-predicted-from scope is not evicted. + if let Some(chain) = scopes.get_mut(scope) { + chain.last_access = Instant::now(); + } + let scoped = scopes + .get(scope) + .and_then(|chain| chain.predict(current, min_samples)) + .filter(|p| p.confidence >= confidence_threshold); + if scoped.is_some() { + return scoped; + } + if scope == GLOBAL_SCOPE { + return None; + } + scopes + .get(GLOBAL_SCOPE) + .and_then(|chain| chain.predict(current, min_samples)) + .filter(|p| p.confidence >= confidence_threshold) + } + + /// Number of scopes currently retained (test/diagnostic use). + #[cfg(test)] + fn scope_count(&self) -> usize { + self.lock().len() } } @@ -107,54 +218,220 @@ impl Default for PreflightPredictor { } } +/// Live counters describing Preflight effectiveness. +#[derive(Debug, Default)] +pub struct PreflightMetrics { + warms_started: AtomicU64, + warms_completed: AtomicU64, + warms_failed: AtomicU64, + warms_skipped: AtomicU64, + predictions: AtomicU64, + hits: AtomicU64, + misses: AtomicU64, +} + +/// Serializable snapshot of [`PreflightMetrics`]. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct PreflightStats { + /// Warm tasks spawned. + pub warms_started: u64, + /// Warm tasks that completed successfully. + pub warms_completed: u64, + /// Warm tasks that failed. + pub warms_failed: u64, + /// Warm requests skipped (deduplicated, disabled, or memory-unsafe). + pub warms_skipped: u64, + /// History/hint predictions produced. + pub predictions: u64, + /// Predictions later confirmed by the next request. + pub hits: u64, + /// Predictions the next request did not match. + pub misses: u64, +} + +impl PreflightMetrics { + /// A warm task was spawned. + pub fn warm_started(&self) { + self.warms_started.fetch_add(1, Ordering::Relaxed); + } + /// A warm task completed successfully. + pub fn warm_completed(&self) { + self.warms_completed.fetch_add(1, Ordering::Relaxed); + } + /// A warm task failed. + pub fn warm_failed(&self) { + self.warms_failed.fetch_add(1, Ordering::Relaxed); + } + /// A warm was skipped (deduplicated, disabled, or unsafe). + pub fn warm_skipped(&self) { + self.warms_skipped.fetch_add(1, Ordering::Relaxed); + } + /// A prediction was produced. + pub fn prediction(&self) { + self.predictions.fetch_add(1, Ordering::Relaxed); + } + /// A prior prediction matched the next request. + pub fn hit(&self) { + self.hits.fetch_add(1, Ordering::Relaxed); + } + /// A prior prediction did not match the next request. + pub fn miss(&self) { + self.misses.fetch_add(1, Ordering::Relaxed); + } + + /// Read all counters into a `PreflightStats`. + /// + /// Each counter is read independently with relaxed ordering, so under heavy + /// concurrent updates the fields may reflect slightly different instants; + /// they are monotonic and eventually consistent, which is what the metrics + /// consumers need. + pub fn snapshot(&self) -> PreflightStats { + PreflightStats { + warms_started: self.warms_started.load(Ordering::Relaxed), + warms_completed: self.warms_completed.load(Ordering::Relaxed), + warms_failed: self.warms_failed.load(Ordering::Relaxed), + warms_skipped: self.warms_skipped.load(Ordering::Relaxed), + predictions: self.predictions.load(Ordering::Relaxed), + hits: self.hits.load(Ordering::Relaxed), + misses: self.misses.load(Ordering::Relaxed), + } + } +} + #[cfg(test)] mod tests { use super::*; + const SAMPLES: u64 = 3; + const THRESHOLD: f64 = 0.6; + #[test] fn no_prediction_without_observations() { let p = PreflightPredictor::new(); - assert!(p.predict(Capability::Chat).is_none()); + assert!( + p.predict("app", Capability::Chat, SAMPLES, THRESHOLD) + .is_none() + ); } #[test] fn no_prediction_below_threshold() { let p = PreflightPredictor::new(); - p.record(Capability::Chat); - p.record(Capability::Tts); - // Only 1 transition Chat→Tts, need 3 - assert!(p.predict(Capability::Chat).is_none()); + p.record("app", Capability::Chat); + p.record("app", Capability::Tts); + // Only one Chat→Tts transition; need three samples. + assert!( + p.predict("app", Capability::Chat, SAMPLES, THRESHOLD) + .is_none() + ); } #[test] fn predicts_after_enough_observations() { let p = PreflightPredictor::new(); - - // Record Chat → Tts four times for _ in 0..4 { - p.record(Capability::Chat); - p.record(Capability::Tts); + p.record("app", Capability::Chat); + p.record("app", Capability::Tts); } - - let pred = p.predict(Capability::Chat).unwrap(); + let pred = p + .predict("app", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap(); assert_eq!(pred.capability, Capability::Tts); assert!(pred.confidence > 0.9); } #[test] - fn handles_multiple_successors() { + fn scopes_do_not_pollute_each_other() { let p = PreflightPredictor::new(); + // App A always does Chat → Tts. + for _ in 0..5 { + p.record("app-a", Capability::Chat); + p.record("app-a", Capability::Tts); + } + // App B always does Chat → Asr. + for _ in 0..5 { + p.record("app-b", Capability::Chat); + p.record("app-b", Capability::Asr); + } + assert_eq!( + p.predict("app-a", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap() + .capability, + Capability::Tts + ); + assert_eq!( + p.predict("app-b", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap() + .capability, + Capability::Asr + ); + } - // Chat → Tts (3x), Chat → Asr (1x) - for _ in 0..3 { - p.record(Capability::Chat); - p.record(Capability::Tts); + #[test] + fn unknown_scope_falls_back_to_global() { + let p = PreflightPredictor::new(); + for _ in 0..5 { + p.record("app-a", Capability::Chat); + p.record("app-a", Capability::Tts); } - p.record(Capability::Chat); - p.record(Capability::Asr); + // A brand-new scope has no history, but global aggregates app-a's. + let pred = p + .predict("fresh-scope", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap(); + assert_eq!(pred.capability, Capability::Tts); + } - let pred = p.predict(Capability::Chat).unwrap(); + #[test] + fn global_chain_ignores_cross_scope_transitions() { + let p = PreflightPredictor::new(); + // Two scopes run concurrently with interleaved requests. App A always + // does Chat → Tts; app B always does Asr → Vlm. If the global chain + // observed the interleaved stream it would learn phantom edges like + // Chat → Asr (A's Chat followed by B's Asr), which never happened in any + // real sequence. + for _ in 0..6 { + p.record("app-a", Capability::Chat); + p.record("app-b", Capability::Asr); + p.record("app-a", Capability::Tts); + p.record("app-b", Capability::Vlm); + } + // A fresh scope falls back to the global aggregate. It must predict the + // real edge Chat → Tts, not the phantom Chat → Asr the old code learned. + let pred = p + .predict("fresh", Capability::Chat, SAMPLES, THRESHOLD) + .expect("global should predict the real transition"); assert_eq!(pred.capability, Capability::Tts); - assert!(pred.confidence > 0.7); + assert!(pred.confidence > 0.9); + } + + #[test] + fn scope_count_is_bounded() { + let p = PreflightPredictor::new(); + // Record far more distinct scopes than the cap allows. + for i in 0..(MAX_SCOPES + 200) { + let scope = format!("session-{i}"); + p.record(&scope, Capability::Chat); + p.record(&scope, Capability::Tts); + } + // Bounded by the cap plus the always-retained global scope. + assert!( + p.scope_count() <= MAX_SCOPES + 1, + "scope map grew to {} entries", + p.scope_count() + ); + } + + #[test] + fn metrics_snapshot_counts() { + let m = PreflightMetrics::default(); + m.warm_started(); + m.warm_started(); + m.warm_completed(); + m.hit(); + let snap = m.snapshot(); + assert_eq!(snap.warms_started, 2); + assert_eq!(snap.warms_completed, 1); + assert_eq!(snap.hits, 1); + assert_eq!(snap.misses, 0); } } diff --git a/mofa-engine-core/src/router.rs b/mofa-engine-core/src/router.rs index 7a8e0c5..0c036b1 100644 --- a/mofa-engine-core/src/router.rs +++ b/mofa-engine-core/src/router.rs @@ -38,13 +38,34 @@ pub struct Router; impl Router { /// Score and rank models, returning the best match and reason. + /// + /// Equivalent to the first entry of [`route_ranked`] with no memory budget. pub fn route<'a>( models: &'a [ModelCard], request: &InferenceRequest, providers: &[RoutingProvider], ) -> Option> { + Self::route_ranked(models, request, providers, None) + .into_iter() + .next() + } + + /// Apply hard constraints, then score and rank every valid candidate, + /// best first. The engine uses the same ordered plan for both primary + /// selection and failover so the two never diverge. + /// + /// `budget_bytes` enables a static memory-feasibility filter: a local model + /// whose estimated footprint exceeds the entire budget can never be admitted + /// even after evicting everything else, so it is dropped here rather than + /// failing later at load time. Pass `None` to skip this filter. + pub fn route_ranked<'a>( + models: &'a [ModelCard], + request: &InferenceRequest, + providers: &[RoutingProvider], + budget_bytes: Option, + ) -> Vec> { let desired_cap = request.capability; - let mut best: Option> = None; + let mut decisions: Vec> = Vec::new(); for model in models { if !Self::matches_model_name(model, request.model.as_deref()) { @@ -69,7 +90,13 @@ impl Router { continue; } - if !model.execution.has_capacity() { + // Static memory feasibility: a local model that cannot fit the budget + // even on an empty device is never a viable candidate. Remote models + // report a zero footprint and are unaffected. + if let Some(budget) = budget_bytes + && model.residency != ModelResidency::Remote + && model.memory_estimate_bytes > budget + { continue; } @@ -79,26 +106,22 @@ impl Router { } let reason = Self::reason(model, provider, score); - let decision = RouteDecision { + decisions.push(RouteDecision { model, score, reason, - }; - - match &best { - None => best = Some(decision), - Some(current) if decision.score > current.score => best = Some(decision), - Some(current) - if decision.score == current.score - && Self::tie_breaks_before(model, current.model) => - { - best = Some(decision); - } - _ => {} - } + }); } - best + // Highest score first; ties resolved toward already-resident models, then + // by canonical id for a fully deterministic, stable ordering. + decisions.sort_by(|a, b| { + b.score + .cmp(&a.score) + .then_with(|| Self::residency_rank(b.model).cmp(&Self::residency_rank(a.model))) + .then_with(|| a.model.id.cmp(&b.model.id)) + }); + decisions } /// Compatibility helper returning just the selected model. @@ -110,6 +133,15 @@ impl Router { Self::route(models, request, providers).map(|d| d.model) } + /// Tie-break rank: resident models (loaded locally or cloud-backed) rank + /// above those that would require a cold start. + fn residency_rank(model: &ModelCard) -> u8 { + u8::from(matches!( + model.residency, + ModelResidency::Loaded | ModelResidency::Remote + )) + } + fn matches_model_name(model: &ModelCard, target: Option<&str>) -> bool { let Some(target) = target else { return true; @@ -172,11 +204,7 @@ impl Router { } fn locality_score(kind: ProviderKind) -> i64 { - match kind { - ProviderKind::Ollama => 100, - ProviderKind::OpenAiCompatible => 0, - _ => 0, - } + if kind.is_local() { 100 } else { 0 } } fn cost_score(tier: CostTier) -> i64 { @@ -198,7 +226,9 @@ impl Router { } fn priority_score(priority: u8) -> i64 { - 100_i64.saturating_sub(priority as i64) + // Lower configured priority is preferred; floor at 0 so a priority above + // 100 cannot contribute a negative term to the score. + (100_i64 - priority as i64).max(0) } fn health_score(health: BackendHealth) -> i64 { @@ -216,18 +246,6 @@ impl Router { .saturating_sub(model.execution.active_requests) as i64 } - fn tie_breaks_before(candidate: &ModelCard, current: &ModelCard) -> bool { - let candidate_loaded = matches!( - candidate.residency, - ModelResidency::Loaded | ModelResidency::Remote - ); - let current_loaded = matches!( - current.residency, - ModelResidency::Loaded | ModelResidency::Remote - ); - candidate_loaded && !current_loaded - } - fn reason(model: &ModelCard, provider: &RoutingProvider, score: i64) -> String { format!( "selected {} via {}: score={score}, kind={:?}, health={:?}, priority={}, residency={:?}, cost={:?}", @@ -290,6 +308,51 @@ mod tests { } } + #[test] + #[ignore = "timing-sensitive perf guard; run explicitly with `--ignored` or as a benchmark"] + fn routing_decision_meets_latency_target() { + // RFC budget: a scheduling decision should take < 1ms (excluding load). + // Wall-clock assertions are sensitive to CI host load, so this is a + // perf guard run on demand rather than a gating unit test. + // Build a realistic pool: three providers, several models each. + let providers = vec![ + provider("ollama", ProviderKind::Ollama, 1), + provider("openai", ProviderKind::OpenAiCompatible, 10), + provider("deepseek", ProviderKind::OpenAiCompatible, 8), + ]; + let mut models = Vec::new(); + for (prov, count) in [("ollama", 8), ("openai", 6), ("deepseek", 6)] { + for i in 0..count { + models.push(make_model( + &format!("m{i}"), + prov, + Capability::Chat, + if i % 2 == 0 { + ModelResidency::Unloaded + } else { + ModelResidency::Loaded + }, + CostTier::Low, + )); + } + } + let req = request(Capability::Chat, None); + let budget = Some(16 * 1024 * 1024 * 1024); + + let iterations = 2000; + let start = std::time::Instant::now(); + for _ in 0..iterations { + let plan = Router::route_ranked(&models, &req, &providers, budget); + std::hint::black_box(&plan); + } + let avg = start.elapsed() / iterations; + assert!( + avg < std::time::Duration::from_millis(1), + "routing averaged {avg:?}/decision over {} models, exceeding the 1ms target", + models.len() + ); + } + #[test] fn prefers_loaded_over_unloaded() { let models = vec![ @@ -393,6 +456,90 @@ mod tests { assert_eq!(selected.name, "b"); } + #[test] + fn route_ranked_orders_local_before_cloud() { + let models = vec![ + make_model( + "cloud", + "openai", + Capability::Chat, + ModelResidency::Remote, + CostTier::High, + ), + make_model( + "local", + "ollama", + Capability::Chat, + ModelResidency::Loaded, + CostTier::Free, + ), + ]; + let providers = vec![ + provider("openai", ProviderKind::OpenAiCompatible, 10), + provider("ollama", ProviderKind::Ollama, 1), + ]; + let ranked = + Router::route_ranked(&models, &request(Capability::Chat, None), &providers, None); + assert_eq!(ranked.len(), 2); + assert_eq!(ranked[0].model.name, "local"); + assert_eq!(ranked[1].model.name, "cloud"); + // Scores are strictly ordered best-first. + assert!(ranked[0].score >= ranked[1].score); + } + + #[test] + fn route_ranked_filters_models_larger_than_budget() { + let mut big = make_model( + "big", + "ollama", + Capability::Chat, + ModelResidency::Unloaded, + CostTier::Free, + ); + big.memory_estimate_bytes = 32 * 1024 * 1024 * 1024; // 32 GB + let small = make_model( + "small", + "ollama", + Capability::Chat, + ModelResidency::Unloaded, + CostTier::Free, + ); + let models = vec![big, small]; + let providers = vec![provider("ollama", ProviderKind::Ollama, 1)]; + + let budget = Some(8 * 1024 * 1024 * 1024); // 8 GB + let ranked = Router::route_ranked( + &models, + &request(Capability::Chat, None), + &providers, + budget, + ); + assert_eq!(ranked.len(), 1); + assert_eq!(ranked[0].model.name, "small"); + } + + #[test] + fn route_ranked_budget_does_not_filter_remote() { + let mut cloud = make_model( + "cloud", + "openai", + Capability::Chat, + ModelResidency::Remote, + CostTier::High, + ); + // A nonsensical estimate must not exclude a cloud model, which holds no local memory. + cloud.memory_estimate_bytes = u64::MAX; + let models = vec![cloud]; + let providers = vec![provider("openai", ProviderKind::OpenAiCompatible, 10)]; + let ranked = Router::route_ranked( + &models, + &request(Capability::Chat, None), + &providers, + Some(1024), + ); + assert_eq!(ranked.len(), 1); + } + #[test] fn skips_unhealthy_provider() { let models = vec![make_model( diff --git a/mofa-engine-core/src/subscription.rs b/mofa-engine-core/src/subscription.rs new file mode 100644 index 0000000..5e30fa1 --- /dev/null +++ b/mofa-engine-core/src/subscription.rs @@ -0,0 +1,197 @@ +//! Capability subscriptions. +//! +//! An application declares at startup which capabilities it will use. The engine +//! keeps subscribed capabilities' models warm and protects them from eviction +//! for the lifetime of the subscription (until it is explicitly removed or its +//! TTL expires). Subscriptions are owned by an `app_id`/`session_id` so they can +//! be reasoned about and cleaned up per application. + +use mofa_kernel::Capability; +use serde::{Deserialize, Serialize}; +use std::collections::HashSet; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +/// Registry of active capability subscriptions. +pub struct SubscriptionRegistry { + inner: Mutex>, + next_id: AtomicU64, +} + +#[derive(Clone)] +struct Subscription { + id: u64, + app_id: Option, + session_id: Option, + capabilities: Vec, + expires_at: Option, +} + +impl Subscription { + fn is_expired(&self, now: Instant) -> bool { + self.expires_at.is_some_and(|e| now >= e) + } +} + +/// Public view of a subscription. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SubscriptionInfo { + /// Unique subscription id. + pub id: u64, + /// Owning application, if provided. + pub app_id: Option, + /// Owning session, if provided. + pub session_id: Option, + /// Subscribed capabilities. + pub capabilities: Vec, + /// Seconds until expiry, or `None` if the subscription does not expire. + pub expires_in_secs: Option, +} + +impl SubscriptionRegistry { + /// Create an empty registry. + pub fn new() -> Self { + Self { + inner: Mutex::new(Vec::new()), + next_id: AtomicU64::new(1), + } + } + + fn lock(&self) -> std::sync::MutexGuard<'_, Vec> { + self.inner.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Register a subscription, returning its id. + /// + /// `ttl` bounds the subscription's lifetime; pass `None` for one that lives + /// until explicitly removed. + pub fn subscribe( + &self, + app_id: Option, + session_id: Option, + capabilities: Vec, + ttl: Option, + ) -> u64 { + let id = self.next_id.fetch_add(1, Ordering::Relaxed); + let expires_at = ttl.map(|d| Instant::now() + d); + self.lock().push(Subscription { + id, + app_id, + session_id, + capabilities, + expires_at, + }); + id + } + + /// Remove a subscription by id. Returns whether one was removed. + pub fn unsubscribe(&self, id: u64) -> bool { + let mut subs = self.lock(); + let before = subs.len(); + subs.retain(|s| s.id != id); + subs.len() != before + } + + /// The set of capabilities currently kept warm by any live subscription. + /// + /// Prunes and reads under a single lock, so the result is a consistent + /// snapshot and expired entries are dropped exactly once. + pub fn active_capabilities(&self) -> HashSet { + let now = Instant::now(); + let mut subs = self.lock(); + subs.retain(|s| !s.is_expired(now)); + subs.iter() + .flat_map(|s| s.capabilities.iter().copied()) + .collect() + } + + /// Whether `capability` is currently subscribed. + pub fn is_subscribed(&self, capability: Capability) -> bool { + self.active_capabilities().contains(&capability) + } + + /// List all live subscriptions. + pub fn list(&self) -> Vec { + let now = Instant::now(); + let mut subs = self.lock(); + subs.retain(|s| !s.is_expired(now)); + subs.iter() + .map(|s| SubscriptionInfo { + id: s.id, + app_id: s.app_id.clone(), + session_id: s.session_id.clone(), + capabilities: s.capabilities.clone(), + expires_in_secs: s + .expires_at + .map(|e| e.saturating_duration_since(now).as_secs()), + }) + .collect() + } +} + +impl Default for SubscriptionRegistry { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn subscribe_and_query() { + let reg = SubscriptionRegistry::new(); + let id = reg.subscribe( + Some("mofa-fm".into()), + None, + vec![Capability::Chat, Capability::Tts], + None, + ); + assert!(reg.is_subscribed(Capability::Chat)); + assert!(reg.is_subscribed(Capability::Tts)); + assert!(!reg.is_subscribed(Capability::Asr)); + assert_eq!(reg.list().len(), 1); + + assert!(reg.unsubscribe(id)); + assert!(!reg.is_subscribed(Capability::Chat)); + assert!(reg.list().is_empty()); + } + + #[test] + fn unsubscribe_unknown_is_false() { + let reg = SubscriptionRegistry::new(); + assert!(!reg.unsubscribe(999)); + } + + #[test] + fn expired_subscriptions_are_pruned() { + let reg = SubscriptionRegistry::new(); + reg.subscribe( + None, + None, + vec![Capability::Tts], + Some(Duration::from_millis(10)), + ); + assert!(reg.is_subscribed(Capability::Tts)); + std::thread::sleep(Duration::from_millis(20)); + assert!(!reg.is_subscribed(Capability::Tts)); + assert!(reg.list().is_empty()); + } + + #[test] + fn list_reports_remaining_ttl() { + let reg = SubscriptionRegistry::new(); + reg.subscribe( + None, + None, + vec![Capability::Chat], + Some(Duration::from_secs(3600)), + ); + let info = reg.list(); + assert_eq!(info.len(), 1); + let ttl = info[0].expires_in_secs.unwrap(); + assert!(ttl > 3590 && ttl <= 3600); + } +} diff --git a/mofa-engine-sdk/Cargo.toml b/mofa-engine-sdk/Cargo.toml index eb785c4..d0bc8e2 100644 --- a/mofa-engine-sdk/Cargo.toml +++ b/mofa-engine-sdk/Cargo.toml @@ -11,6 +11,7 @@ mofa-engine-core = { workspace = true } tokio = { workspace = true } serde = { workspace = true } serde_json = { workspace = true } +reqwest = { workspace = true } axum = { workspace = true } axum-extra = { workspace = true } tower = { workspace = true } @@ -19,3 +20,4 @@ thiserror = { workspace = true } tracing = { workspace = true } chrono = { workspace = true } tokio-stream = { workspace = true } +uuid = { workspace = true } diff --git a/mofa-engine-sdk/src/client.rs b/mofa-engine-sdk/src/client.rs new file mode 100644 index 0000000..1b51f0c --- /dev/null +++ b/mofa-engine-sdk/src/client.rs @@ -0,0 +1,384 @@ +//! Native Rust SDK: embedded and daemon clients. +//! +//! The RFC feedback separates two deployment modes that the prototype conflated: +//! +//! - **Embedded mode** — the engine runs in the caller's process. [`EmbeddedEngine`] +//! is a small *synchronous* facade over an internally managed Tokio runtime, +//! so a non-async caller (including a UniFFI-generated Python binding) can drive +//! the engine without owning a runtime. This is the intended UniFFI target. +//! - **Daemon mode** — the engine runs as a separate process and is reached over +//! the versioned HTTP API. [`DaemonClient`] is a typed client for that surface. +//! +//! Both speak the same domain request/response types from `mofa-kernel`, so a +//! caller can switch modes without changing how it constructs requests. + +use std::sync::Arc; +use std::time::Duration; + +use mofa_engine_core::engine::{LifecycleRecord, MemoryReport}; +use mofa_engine_core::preflight::PreflightStats; +use mofa_engine_core::subscription::SubscriptionInfo; +use mofa_engine_core::{Engine, EngineConfig}; +use mofa_kernel::{ + Capability, EngineError, EngineStatus, ErrorInfo, InferenceRequest, InferenceResponse, + ModelCard, +}; + +/// Errors returned by [`DaemonClient`]. +#[derive(Debug, thiserror::Error)] +pub enum ClientError { + /// Network/transport failure reaching the daemon. + #[error("transport error: {0}")] + Transport(String), + /// The daemon returned a structured engine error. + #[error("engine error: {}", .0.message)] + Engine(ErrorInfo), + /// The daemon returned a non-success status without a structured body + /// (e.g. a 401/404/413 from the framework rather than the handler). + #[error("http {status}: {body}")] + Status { + /// HTTP status code. + status: u16, + /// Raw response body. + body: String, + }, + /// A success response body could not be decoded into the expected type. + #[error("decode error: {0}")] + Decode(String), +} + +impl ClientError { + /// The structured engine error, if this was an engine-level failure. + pub fn engine_error(&self) -> Option<&ErrorInfo> { + match self { + Self::Engine(info) => Some(info), + _ => None, + } + } +} + +/// A synchronous, in-process facade over the engine for embedded/UniFFI use. +/// +/// Owns a multi-threaded Tokio runtime and blocks on each call, so it can be +/// driven from ordinary synchronous code. Background engine tasks (idle +/// eviction, speculative warming) run on the owned runtime for the lifetime of +/// this value. +pub struct EmbeddedEngine { + runtime: tokio::runtime::Runtime, + engine: Arc, +} + +impl EmbeddedEngine { + /// Build and initialize an embedded engine from configuration. + pub fn new(config: EngineConfig) -> Result { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .map_err(|e| EngineError::Internal(format!("failed to build runtime: {e}")))?; + let engine = runtime.block_on(Engine::try_new(config))?; + Ok(Self { runtime, engine }) + } + + /// List all known model cards. + pub fn capabilities(&self) -> Vec { + self.runtime.block_on(self.engine.capabilities()) + } + + /// Run one inference request to completion. + pub fn invoke(&self, request: InferenceRequest) -> Result { + self.runtime.block_on(self.engine.invoke(request)) + } + + /// Snapshot the engine status. + pub fn status(&self) -> EngineStatus { + self.runtime.block_on(self.engine.status()) + } + + /// Re-run discovery and health probes. + pub fn refresh(&self) { + self.runtime.block_on(self.engine.refresh_resources()); + } + + /// Register a capability subscription; returns its id. + pub fn subscribe( + &self, + app_id: Option, + session_id: Option, + capabilities: Vec, + ttl_secs: Option, + ) -> u64 { + // `Engine::subscribe` spawns background warm tasks, so it must run with + // the owned runtime entered on this thread — otherwise `tokio::spawn` + // panics when called from a plain (non-async) caller. + let _guard = self.runtime.enter(); + self.engine.subscribe( + app_id, + session_id, + capabilities, + ttl_secs.map(Duration::from_secs), + ) + } + + /// Remove a subscription by id. + pub fn unsubscribe(&self, id: u64) -> bool { + let _guard = self.runtime.enter(); + self.engine.unsubscribe(id) + } + + /// Access the underlying async engine for advanced, async-native use. + pub fn engine(&self) -> &Arc { + &self.engine + } +} + +/// JSON string facade — the surface a UniFFI binding exports to Python. +/// +/// UniFFI cannot carry `serde_json::Value`, borrowed types, or the full domain +/// structs across the FFI boundary without hand-written type maps. These +/// string-in / string-out methods sidestep that: they are stable, language +/// agnostic, and map onto UniFFI's `string` and `Result` +/// natively. A generated binding is a thin wrapper that forwards JSON. +impl EmbeddedEngine { + /// Capability list as a JSON array. + pub fn capabilities_json(&self) -> String { + serde_json::to_string(&self.capabilities()).unwrap_or_else(|_| "[]".into()) + } + + /// Engine status as a JSON object. + pub fn status_json(&self) -> String { + serde_json::to_string(&self.status()).unwrap_or_else(|_| "{}".into()) + } + + /// Parse an [`InferenceRequest`] from JSON, invoke, and return the response + /// as JSON. On failure the `Err` payload is a JSON [`ErrorInfo`] with a + /// stable code, so a binding can raise a typed exception. + pub fn invoke_json(&self, request_json: &str) -> Result { + let request: InferenceRequest = serde_json::from_str(request_json) + .map_err(|e| err_json(EngineError::InvalidRequest(e.to_string())))?; + self.invoke(request) + .map(|resp| serde_json::to_string(&resp).unwrap_or_default()) + .map_err(err_json) + } +} + +/// Serialize an engine error to a JSON `ErrorInfo` string. +fn err_json(e: EngineError) -> String { + serde_json::to_string(&e.info()).unwrap_or_else(|_| e.to_string()) +} + +/// A typed client for the engine's versioned HTTP API (daemon mode). +pub struct DaemonClient { + base_url: String, + http: reqwest::Client, + token: Option, +} + +impl DaemonClient { + /// Create a client for a daemon at `base_url` (e.g. `http://localhost:8420`). + pub fn new(base_url: impl Into) -> Self { + Self::with_client(base_url, reqwest::Client::new()) + } + + /// Create a client with a caller-provided HTTP client. + pub fn with_client(base_url: impl Into, http: reqwest::Client) -> Self { + let base_url = base_url.into().trim_end_matches('/').to_string(); + Self { + base_url, + http, + token: None, + } + } + + /// Attach a bearer token, sent on every request. Required when the daemon + /// runs with `MOFA_API_TOKEN` set. + pub fn with_token(mut self, token: impl Into) -> Self { + self.token = Some(token.into()); + self + } + + fn url(&self, path: &str) -> String { + format!("{}/{}", self.base_url, path.trim_start_matches('/')) + } + + /// Attach the bearer token to a request builder when one is configured. + fn authed(&self, rb: reqwest::RequestBuilder) -> reqwest::RequestBuilder { + match &self.token { + Some(token) => rb.bearer_auth(token), + None => rb, + } + } + + async fn get_json(&self, path: &str) -> Result { + let resp = self + .authed(self.http.get(self.url(path))) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } + + async fn decode( + resp: reqwest::Response, + ) -> Result { + let status = resp.status(); + if status.is_success() { + return resp + .json::() + .await + .map_err(|e| ClientError::Decode(e.to_string())); + } + // Read the body once, then prefer a structured ErrorInfo; fall back to a + // status-bearing error for framework responses (401/404/413/…) so the + // status and reason are not lost. + let body = resp.text().await.unwrap_or_default(); + match serde_json::from_str::(&body) { + Ok(info) => Err(ClientError::Engine(info)), + Err(_) => Err(ClientError::Status { + status: status.as_u16(), + body, + }), + } + } + + /// `GET /v1/capabilities` + pub async fn capabilities(&self) -> Result, ClientError> { + self.get_json("/v1/capabilities").await + } + + /// `GET /v1/status` + pub async fn status(&self) -> Result { + self.get_json("/v1/status").await + } + + /// `GET /v1/memory` + pub async fn memory(&self) -> Result { + self.get_json("/v1/memory").await + } + + /// `GET /v1/lifecycle` + pub async fn lifecycle(&self) -> Result, ClientError> { + self.get_json("/v1/lifecycle").await + } + + /// `GET /v1/preflight` + pub async fn preflight(&self) -> Result { + self.get_json("/v1/preflight").await + } + + /// `GET /v1/subscriptions` + pub async fn subscriptions(&self) -> Result, ClientError> { + self.get_json("/v1/subscriptions").await + } + + /// `POST /v1/invoke` + pub async fn invoke( + &self, + request: &InferenceRequest, + ) -> Result { + let resp = self + .authed(self.http.post(self.url("/v1/invoke")).json(request)) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } + + /// `POST /v1/discovery/refresh` + pub async fn refresh(&self) -> Result { + let resp = self + .authed(self.http.post(self.url("/v1/discovery/refresh"))) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mofa_engine_core::config::{ListenConfig, MemoryConfig, PreflightConfig, TimeoutConfig}; + + fn empty_config() -> EngineConfig { + EngineConfig { + listen: ListenConfig::default(), + memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), + providers: vec![], + } + } + + #[test] + fn embedded_engine_boots_and_reports_status() { + // A plain (non-async) caller can drive the engine via the facade. + let engine = EmbeddedEngine::new(empty_config()).expect("embedded engine builds"); + assert!(engine.capabilities().is_empty()); + assert_eq!(engine.status().total_models, 0); + } + + #[test] + fn embedded_invoke_without_models_errors() { + let engine = EmbeddedEngine::new(empty_config()).unwrap(); + let req = InferenceRequest { + capability: Some(Capability::Chat), + model: None, + app_id: None, + session_id: None, + fallback_policy: Default::default(), + messages: vec![], + input_file: None, + params: serde_json::Value::Null, + hint_next: None, + request_id: "t".into(), + }; + assert!(matches!( + engine.invoke(req), + Err(EngineError::NoCapableModel(_)) + )); + } + + #[test] + fn embedded_json_facade_round_trips() { + let engine = EmbeddedEngine::new(empty_config()).unwrap(); + assert_eq!(engine.capabilities_json(), "[]"); + assert!(engine.status_json().contains("total_models")); + + // A malformed request body yields a JSON ErrorInfo in the Err arm. + let err = engine.invoke_json("not json").unwrap_err(); + assert!(err.contains("invalid_request")); + + // A well-formed request with no models yields a no_capable_model error. + let err = engine + .invoke_json(r#"{"capability":"chat","messages":[]}"#) + .unwrap_err(); + assert!(err.contains("no_capable_model")); + } + + #[test] + fn daemon_client_normalizes_urls() { + let client = DaemonClient::new("http://localhost:8420/"); + assert_eq!(client.url("/v1/status"), "http://localhost:8420/v1/status"); + assert_eq!(client.url("v1/status"), "http://localhost:8420/v1/status"); + } + + #[test] + fn with_token_is_retained_and_url_unaffected() { + let client = DaemonClient::new("http://localhost:8420").with_token("secret"); + assert_eq!(client.token.as_deref(), Some("secret")); + assert_eq!(client.url("/v1/status"), "http://localhost:8420/v1/status"); + } + + #[test] + fn client_error_status_preserves_code_and_body() { + let err = ClientError::Status { + status: 401, + body: "unauthorized".into(), + }; + assert!(err.to_string().contains("401")); + assert!(err.to_string().contains("unauthorized")); + assert!(err.engine_error().is_none()); + } +} diff --git a/mofa-engine-sdk/src/lib.rs b/mofa-engine-sdk/src/lib.rs index df91a52..8ed550e 100644 --- a/mofa-engine-sdk/src/lib.rs +++ b/mofa-engine-sdk/src/lib.rs @@ -1,9 +1,18 @@ //! # mofa-engine-sdk //! -//! HTTP API server (Axum), SSE event streaming, and embedded dashboard -//! for the MoFA Engine. +//! The unified call layer for the MoFA Engine: +//! +//! - [`server`] — the versioned Axum HTTP API, SSE streaming, and dashboard. +//! - [`client`] — a native Rust SDK with an embedded ([`client::EmbeddedEngine`]) +//! and a daemon ([`client::DaemonClient`]) mode. +//! +//! The embedded facade is synchronous and intended as the UniFFI target for +//! Python bindings; the daemon client speaks the same HTTP surface the server +//! exposes. +pub mod client; pub mod dashboard; pub mod server; +pub use client::{ClientError, DaemonClient, EmbeddedEngine}; pub use server::start_server; diff --git a/mofa-engine-sdk/src/server.rs b/mofa-engine-sdk/src/server.rs index 68c2623..2af2965 100644 --- a/mofa-engine-sdk/src/server.rs +++ b/mofa-engine-sdk/src/server.rs @@ -6,63 +6,177 @@ use std::time::Duration; use axum::{ Json, Router, - extract::State, - http::StatusCode, + extract::{DefaultBodyLimit, Path, Request, State}, + http::{HeaderValue, StatusCode, header}, + middleware::{self, Next}, response::{ - Html, + Html, IntoResponse, Response, sse::{Event, KeepAlive, Sse}, }, - routing::{get, post}, + routing::{delete, get, post}, }; use mofa_engine_core::Engine; -use mofa_kernel::{ErrorInfo, InferenceRequest}; -use serde::Serialize; +use mofa_engine_core::engine::{LifecycleRecord, MemoryReport}; +use mofa_engine_core::preflight::PreflightStats; +use mofa_engine_core::subscription::SubscriptionInfo; +use mofa_kernel::{Capability, EngineError, ErrorInfo, InferenceRequest}; +use serde::{Deserialize, Serialize}; use tokio_stream::StreamExt; -use tokio_stream::wrappers::BroadcastStream; +use tokio_stream::wrappers::{BroadcastStream, ReceiverStream}; use tower_http::cors::CorsLayer; use tower_http::trace::TraceLayer; +use tracing::Instrument; use crate::dashboard; +/// Header carrying the request correlation id. +const REQUEST_ID_HEADER: &str = "x-request-id"; +/// Maximum accepted request body size (16 MiB) — bounds base64/JSON payloads. +const MAX_BODY_BYTES: usize = 16 * 1024 * 1024; + /// Shared application state. #[derive(Clone)] struct AppState { engine: Arc, started_at: std::time::Instant, + /// Optional bearer token; when set, `/v1` routes require it. + api_token: Option>, } /// Start the HTTP server. +/// +/// Binds loopback by default (the caller passes the host). Set `MOFA_API_TOKEN` +/// to require `Authorization: Bearer ` on all `/v1` routes; when unset, +/// the API is open (appropriate only for a trusted local machine). pub async fn start_server( engine: Arc, host: &str, port: u16, ) -> Result<(), Box> { + let api_token = std::env::var("MOFA_API_TOKEN") + .ok() + .filter(|t| !t.is_empty()) + .map(Arc::new); + if api_token.is_some() { + tracing::info!("API authentication enabled (bearer token required on /v1)"); + } + let state = AppState { engine, started_at: std::time::Instant::now(), + api_token, }; - let app = Router::new() - // Dashboard - .route("/", get(dashboard_handler)) - // API routes - .route("/health", get(health_handler)) + let app = build_app(state); + + let addr = format!("{host}:{port}"); + tracing::info!("MoFA Engine listening on http://{addr}"); + + let listener = tokio::net::TcpListener::bind(&addr).await?; + axum::serve(listener, app).await?; + + Ok(()) +} + +/// Assemble the full router: public routes, an auth-gated `/v1` API, and the +/// cross-cutting middleware stack. Extracted so tests can exercise it without +/// binding a socket. +fn build_app(state: AppState) -> Router { + // `/v1` routes sit behind the auth gate; public routes do not. + let api = Router::new() .route("/v1/capabilities", get(capabilities_handler)) .route("/v1/invoke", post(invoke_handler)) + .route("/v1/invoke/stream", post(invoke_stream_handler)) .route("/v1/status", get(status_handler)) + .route("/v1/memory", get(memory_handler)) + .route("/v1/lifecycle", get(lifecycle_handler)) + .route("/v1/preflight", get(preflight_handler)) + .route( + "/v1/subscriptions", + get(list_subscriptions_handler).post(subscribe_handler), + ) + .route("/v1/subscriptions/{id}", delete(unsubscribe_handler)) .route("/v1/events", get(events_handler)) .route("/v1/discovery/refresh", post(refresh_handler)) + .route("/v1/models/load", post(load_model_handler)) + .route("/v1/models/unload", post(unload_model_handler)) + .route_layer(middleware::from_fn_with_state( + state.clone(), + auth_middleware, + )); + + Router::new() + // Public (unauthenticated) routes. + .route("/", get(dashboard_handler)) + .route("/health", get(health_handler)) + .route("/metrics", get(metrics_handler)) + .merge(api) + .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) .layer(CorsLayer::permissive()) .layer(TraceLayer::new_for_http()) - .with_state(state); + // Outermost, so the request-id span encloses the trace layer and every + // response (including framework rejections) carries `x-request-id`. + .layer(middleware::from_fn(correlation_middleware)) + .with_state(state) +} - let addr = format!("{host}:{port}"); - tracing::info!("MoFA Engine listening on http://{addr}"); +/// Maximum accepted length of a client-supplied `x-request-id`. A longer (or +/// empty) value is replaced with a generated id so a caller cannot amplify log +/// or allocation cost through the correlation header. +const MAX_REQUEST_ID_LEN: usize = 128; - let listener = tokio::net::TcpListener::bind(&addr).await?; - axum::serve(listener, app).await?; +/// Attach an `x-request-id` to every request/response and open a tracing span +/// carrying it, so logs can be correlated across a single request. +async fn correlation_middleware(req: Request, next: Next) -> Response { + let request_id = req + .headers() + .get(REQUEST_ID_HEADER) + .and_then(|v| v.to_str().ok()) + .filter(|v| !v.is_empty() && v.len() <= MAX_REQUEST_ID_LEN) + .map(str::to_owned) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - Ok(()) + let span = tracing::info_span!("request", request_id = %request_id); + let mut resp = next.run(req).instrument(span).await; + + if let Ok(value) = HeaderValue::from_str(&request_id) { + resp.headers_mut().insert(REQUEST_ID_HEADER, value); + } + resp +} + +/// Enforce bearer-token auth on `/v1` routes when a token is configured. +async fn auth_middleware(State(state): State, req: Request, next: Next) -> Response { + let Some(expected) = state.api_token.as_ref() else { + return next.run(req).await; + }; + let presented = req + .headers() + .get(header::AUTHORIZATION) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.strip_prefix("Bearer ")); + let authorized = presented + .map(|p| constant_time_eq(p.as_bytes(), expected.as_bytes())) + .unwrap_or(false); + if authorized { + next.run(req).await + } else { + let err = EngineError::InvalidRequest("missing or invalid bearer token".into()); + (StatusCode::UNAUTHORIZED, Json(err.info())).into_response() + } +} + +/// Compare two byte strings in time independent of how many leading bytes match, +/// so token verification does not leak the secret through response timing. +fn constant_time_eq(a: &[u8], b: &[u8]) -> bool { + if a.len() != b.len() { + return false; + } + let mut diff = 0u8; + for (x, y) in a.iter().zip(b.iter()) { + diff |= x ^ y; + } + diff == 0 } /// Health check response. @@ -107,6 +221,33 @@ async fn invoke_handler( } } +/// Stream inference output as Server-Sent Events. +/// +/// Each SSE `data:` line is a JSON [`StreamChunk`](mofa_kernel::StreamChunk): +/// a `started` event, then `text` deltas, then a terminal `completed` or +/// `error`. Errors are delivered in-band as an `error` chunk (HTTP status is +/// always 200 once the stream opens), so clients handle failures uniformly. +async fn invoke_stream_handler( + State(state): State, + Json(req): Json, +) -> Sse>> { + let rx = state.engine.invoke_stream(req); + let stream = ReceiverStream::new(rx).map(|chunk| { + // Every SSE frame must be a valid StreamChunk JSON. Serialization of a + // chunk cannot realistically fail, but if it did, emit a well-formed + // error chunk rather than an empty (invalid) frame. + let data = serde_json::to_string(&chunk).unwrap_or_else(|_| { + r#"{"type":"error","code":"internal","message":"failed to serialize stream chunk","retryable":false,"source":null}"#.to_string() + }); + Ok(Event::default().data(data)) + }); + Sse::new(stream).keep_alive( + KeepAlive::new() + .interval(Duration::from_secs(15)) + .text("ping"), + ) +} + async fn status_handler(State(state): State) -> Json { let status = state.engine.status().await; Json(status) @@ -117,6 +258,133 @@ async fn refresh_handler(State(state): State) -> Json) -> impl IntoResponse { + let body = state.engine.metrics_prometheus(); + ([(header::CONTENT_TYPE, "text/plain; version=0.0.4")], body) +} + +/// Map an engine error to the HTTP status used by the management endpoints. +fn error_status(e: &EngineError) -> StatusCode { + match e { + EngineError::NoCapableModel(_) => StatusCode::NOT_FOUND, + EngineError::InvalidRequest(_) | EngineError::UnsupportedOperation(_) => { + StatusCode::BAD_REQUEST + } + EngineError::CircuitOpen(_) => StatusCode::SERVICE_UNAVAILABLE, + EngineError::Timeout(_) => StatusCode::GATEWAY_TIMEOUT, + _ => StatusCode::INTERNAL_SERVER_ERROR, + } +} + +/// Body for the manual model load/unload management endpoints. +#[derive(Deserialize)] +struct ModelActionRequest { + /// Canonical model id (`provider/model`). + model_id: String, +} + +#[derive(Serialize)] +struct ModelActionResponse { + model_id: String, + changed: bool, +} + +/// `POST /v1/models/load` — manually warm a model. +async fn load_model_handler( + State(state): State, + Json(req): Json, +) -> Result, (StatusCode, Json)> { + match state.engine.load_model(&req.model_id).await { + Ok(()) => Ok(Json(ModelActionResponse { + model_id: req.model_id, + changed: true, + })), + Err(e) => Err((error_status(&e), Json(e.info()))), + } +} + +/// `POST /v1/models/unload` — manually unload a model. +async fn unload_model_handler( + State(state): State, + Json(req): Json, +) -> Result, (StatusCode, Json)> { + match state.engine.unload_model_manual(&req.model_id).await { + Ok(changed) => Ok(Json(ModelActionResponse { + model_id: req.model_id, + changed, + })), + Err(e) => Err((error_status(&e), Json(e.info()))), + } +} + +async fn memory_handler(State(state): State) -> Json { + Json(state.engine.memory_report()) +} + +async fn lifecycle_handler(State(state): State) -> Json> { + Json(state.engine.lifecycle_history()) +} + +async fn preflight_handler(State(state): State) -> Json { + Json(state.engine.preflight_stats()) +} + +async fn list_subscriptions_handler(State(state): State) -> Json> { + Json(state.engine.subscriptions()) +} + +/// Body for `POST /v1/subscriptions`. +#[derive(Deserialize)] +struct SubscribeRequest { + #[serde(default)] + app_id: Option, + #[serde(default)] + session_id: Option, + capabilities: Vec, + /// Optional lifetime in seconds; omitted means it lives until removed. + #[serde(default)] + ttl_secs: Option, +} + +#[derive(Serialize)] +struct SubscribeResponse { + id: u64, +} + +async fn subscribe_handler( + State(state): State, + Json(req): Json, +) -> Result, (StatusCode, Json)> { + if req.capabilities.is_empty() { + let err = EngineError::InvalidRequest("capabilities must not be empty".into()); + return Err((StatusCode::BAD_REQUEST, Json(err.info()))); + } + let ttl = req.ttl_secs.map(Duration::from_secs); + let id = state + .engine + .subscribe(req.app_id, req.session_id, req.capabilities, ttl); + Ok(Json(SubscribeResponse { id })) +} + +#[derive(Serialize)] +struct UnsubscribeResponse { + removed: bool, +} + +async fn unsubscribe_handler( + State(state): State, + Path(id): Path, +) -> (StatusCode, Json) { + let removed = state.engine.unsubscribe(id); + let status = if removed { + StatusCode::OK + } else { + StatusCode::NOT_FOUND + }; + (status, Json(UnsubscribeResponse { removed })) +} + async fn events_handler( State(state): State, ) -> Sse>> { @@ -143,6 +411,11 @@ async fn dashboard_handler() -> Html<&'static str> { #[cfg(test)] mod tests { use super::*; + use axum::body::Body; + use axum::http::Request as HttpRequest; + use mofa_engine_core::EngineConfig; + use mofa_engine_core::config::{ListenConfig, MemoryConfig, PreflightConfig, TimeoutConfig}; + use tower::ServiceExt; // for `oneshot` #[test] fn health_response_serializes() { @@ -154,4 +427,122 @@ mod tests { let json = serde_json::to_string(&resp).unwrap(); assert!(json.contains("\"status\":\"ok\"")); } + + async fn test_state(api_token: Option<&str>) -> AppState { + let config = EngineConfig { + listen: ListenConfig::default(), + memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), + providers: vec![], + }; + AppState { + engine: mofa_engine_core::Engine::new(config).await, + started_at: std::time::Instant::now(), + api_token: api_token.map(|t| Arc::new(t.to_string())), + } + } + + #[tokio::test] + async fn metrics_endpoint_serves_prometheus_text() { + let app = build_app(test_state(None).await); + let resp = app + .oneshot( + HttpRequest::builder() + .uri("/metrics") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + assert!( + resp.headers() + .get(header::CONTENT_TYPE) + .unwrap() + .to_str() + .unwrap() + .starts_with("text/plain") + ); + let body = axum::body::to_bytes(resp.into_body(), 65536).await.unwrap(); + let text = String::from_utf8(body.to_vec()).unwrap(); + assert!(text.contains("mofa_requests_total")); + } + + #[tokio::test] + async fn responses_carry_a_request_id_header() { + let app = build_app(test_state(None).await); + let resp = app + .oneshot( + HttpRequest::builder() + .uri("/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert!(resp.headers().contains_key(REQUEST_ID_HEADER)); + } + + #[tokio::test] + async fn provided_request_id_is_echoed_back() { + let app = build_app(test_state(None).await); + let resp = app + .oneshot( + HttpRequest::builder() + .uri("/health") + .header(REQUEST_ID_HEADER, "abc-123") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.headers().get(REQUEST_ID_HEADER).unwrap(), "abc-123"); + } + + #[tokio::test] + async fn auth_gate_rejects_missing_token_but_allows_public_routes() { + let app = build_app(test_state(Some("secret")).await); + + // A /v1 route without the token is unauthorized. + let resp = app + .clone() + .oneshot( + HttpRequest::builder() + .uri("/v1/status") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::UNAUTHORIZED); + + // The same route with the correct bearer token succeeds. + let resp = app + .clone() + .oneshot( + HttpRequest::builder() + .uri("/v1/status") + .header(header::AUTHORIZATION, "Bearer secret") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Public routes remain reachable without a token. + let resp = app + .oneshot( + HttpRequest::builder() + .uri("/health") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } } diff --git a/mofa-kernel/Cargo.toml b/mofa-kernel/Cargo.toml index e269434..0ebd268 100644 --- a/mofa-kernel/Cargo.toml +++ b/mofa-kernel/Cargo.toml @@ -11,3 +11,4 @@ serde = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } uuid = { workspace = true } +tokio = { version = "1", features = ["sync"] } diff --git a/mofa-kernel/src/traits.rs b/mofa-kernel/src/traits.rs index 31fc521..6767793 100644 --- a/mofa-kernel/src/traits.rs +++ b/mofa-kernel/src/traits.rs @@ -5,7 +5,7 @@ use async_trait::async_trait; use crate::error::EngineError; use crate::types::{ BackendFeature, BackendHealth, InferenceRequest, InferenceResponse, LifecycleResult, ModelCard, - ProviderKind, + ProviderKind, StreamSink, }; /// A model provider backend. @@ -41,6 +41,31 @@ pub trait Provider: Send + Sync { model_id: &str, request: &InferenceRequest, ) -> Result; + + /// Stream inference output as it is produced. + /// + /// Providers push incremental text deltas to `sink` and return the final + /// aggregate response (with the full text/file, tokens, etc.). The engine + /// owns the surrounding `Started`/`Completed`/`Error` envelope. + /// + /// The default implementation provides *non-streaming compatibility*: it + /// runs [`invoke`](Self::invoke) and emits the whole text output as a single + /// delta. Backends that support incremental generation (e.g. Ollama) should + /// override this to stream real tokens. + async fn stream( + &self, + model_id: &str, + request: &InferenceRequest, + sink: StreamSink, + ) -> Result { + let response = self.invoke(model_id, request).await?; + if let Some(text) = &response.text + && !text.is_empty() + { + let _ = sink.send(text.clone()).await; + } + Ok(response) + } } #[cfg(test)] diff --git a/mofa-kernel/src/types.rs b/mofa-kernel/src/types.rs index 196dabe..3eabccb 100644 --- a/mofa-kernel/src/types.rs +++ b/mofa-kernel/src/types.rs @@ -3,6 +3,8 @@ use serde::{Deserialize, Serialize}; use uuid::Uuid; +use crate::error::ErrorInfo; + /// Capabilities that a model can provide. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] @@ -64,6 +66,17 @@ pub enum ProviderKind { Ollama, /// Any OpenAI-compatible API. OpenAiCompatible, + /// Local process-adapter backend (e.g. an MLX/Kokoro or Piper TTS CLI). + LocalTts, +} + +impl ProviderKind { + /// Whether this backend runs on the local machine (as opposed to a remote + /// API). Used by routing to prefer local models and by the memory manager + /// to account for on-device residency. + pub fn is_local(self) -> bool { + matches!(self, Self::Ollama | Self::LocalTts) + } } /// Backend health is independent from model residency. @@ -400,6 +413,52 @@ pub struct InferenceResponse { pub routing_reason: Option, } +/// A single event in a streaming inference response. +/// +/// The streaming interface is versioned and can operate in a *non-streaming +/// compatibility mode*: a backend that cannot emit incremental tokens sends +/// `Started`, one `Text` chunk carrying the full output, then `Completed`. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +#[non_exhaustive] +pub enum StreamChunk { + /// Emitted once before any content; identifies the serving model. + Started { + /// Request correlation ID. + request_id: String, + /// Model actually serving the request. + model_used: String, + /// Provider serving the request. + provider: String, + }, + /// An incremental piece of text output. + Text { + /// Text delta appended to the output so far. + delta: String, + }, + /// Terminal success event with aggregate metadata and any file output. + Completed { + /// Wall-clock duration in milliseconds. + duration_ms: u64, + /// Token usage, if reported. + tokens_used: Option, + /// File output path (for TTS, image gen, etc.). + file: Option, + /// Whether a fallback candidate served the request. + fallback_used: bool, + /// Machine-readable routing reason. + routing_reason: Option, + }, + /// Terminal error event. + Error(ErrorInfo), +} + +/// Channel a provider uses to emit incremental text deltas while streaming. +/// +/// The engine owns the surrounding envelope (`Started`/`Completed`/`Error`); +/// providers push only text deltas here and return the final aggregate response. +pub type StreamSink = tokio::sync::mpsc::Sender; + /// Result of a lifecycle operation. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct LifecycleResult { @@ -476,6 +535,27 @@ pub enum EngineEvent { /// Total budget bytes. total_bytes: u64, }, + /// A model was evicted from local memory. + ModelEvicted { + /// Model identifier. + model_id: String, + /// Why it was evicted, e.g. `memory_pressure` or `idle_timeout`. + reason: String, + }, + /// Predictive (Preflight) warming started for a model. + PreflightWarmStarted { + /// Model identifier being warmed. + model_id: String, + /// What triggered the warm: `hint`, `subscription`, or `history`. + source: String, + }, + /// Predictive (Preflight) warming finished. + PreflightWarmCompleted { + /// Model identifier that was warmed. + model_id: String, + /// Whether the warm succeeded. + success: bool, + }, /// Provider health changed. ProviderHealthChanged { /// Provider name.