From 94ceb4ff0580ccd81177d95080087c8db171205b Mon Sep 17 00:00:00 2001 From: aayank13 Date: Sun, 28 Jun 2026 18:38:22 +0530 Subject: [PATCH 01/22] feat: implement per-scope Markov chain preflight prediction - Refactor PreflightPredictor to track transitions between capabilities per application/session scope. - Introduce global fallback for scopes with insufficient data. - Add metrics for preflight effectiveness and expose them via API. - Implement SubscriptionRegistry for managing capability subscriptions with TTL. - Enhance Router to support memory and lifecycle reporting, and subscription management. - Add new EngineEvent variants for model eviction and preflight warming events. --- config.example.toml | 60 + docs/acceptance.md | 68 +- ...0002-stage-3-5-routing-memory-preflight.md | 80 + docs/openapi.yaml | 189 +- mofa-engine-core/src/circuit_breaker.rs | 72 +- mofa-engine-core/src/config.rs | 185 ++ mofa-engine-core/src/engine.rs | 1712 ++++++++++++++--- mofa-engine-core/src/lib.rs | 1 + mofa-engine-core/src/memory.rs | 338 +++- mofa-engine-core/src/preflight.rs | 372 +++- mofa-engine-core/src/router.rs | 204 +- mofa-engine-core/src/subscription.rs | 199 ++ mofa-engine-sdk/src/server.rs | 86 +- mofa-kernel/src/types.rs | 21 + 14 files changed, 3135 insertions(+), 452 deletions(-) create mode 100644 config.example.toml create mode 100644 docs/decisions/0002-stage-3-5-routing-memory-preflight.md create mode 100644 mofa-engine-core/src/subscription.rs diff --git a/config.example.toml b/config.example.toml new file mode 100644 index 0000000..6c8c3f3 --- /dev/null +++ b/config.example.toml @@ -0,0 +1,60 @@ +# 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 + +# 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" 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..e4b09f4 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,102 @@ paths: application/json: schema: $ref: "#/components/schemas/ErrorInfo" + /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 components: schemas: InferenceRequest: @@ -147,3 +246,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-core/src/circuit_breaker.rs b/mofa-engine-core/src/circuit_breaker.rs index d62f650..97f5802 100644 --- a/mofa-engine-core/src/circuit_breaker.rs +++ b/mofa-engine-core/src/circuit_breaker.rs @@ -50,6 +50,9 @@ 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, config: CircuitBreakerConfig, } @@ -59,27 +62,37 @@ impl ProviderBreaker { state: CircuitState::Closed, failure_count: 0, last_failure: None, + probe_in_flight: false, config, } } /// 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`). 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; + tracing::info!("circuit breaker → half_open (admitting one probe)"); return true; } false } CircuitState::HalfOpen => { - // Allow exactly one probe request + // A probe is already being tested; reject everyone else. + if self.probe_in_flight { + return false; + } + self.probe_in_flight = true; true } } @@ -88,6 +101,7 @@ impl ProviderBreaker { /// Record a successful request. fn record_success(&mut self) { self.failure_count = 0; + self.probe_in_flight = false; self.state = CircuitState::Closed; } @@ -104,12 +118,13 @@ 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; tracing::warn!("circuit breaker → open (probe failed)"); } CircuitState::Open => { - // Already open + // Already open. } } } @@ -228,6 +243,53 @@ 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 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..38da9a5 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -19,6 +19,12 @@ 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, /// Provider definitions #[serde(default)] pub providers: Vec, @@ -75,6 +81,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 { @@ -192,6 +334,12 @@ 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 + ))); + } for provider in &self.providers { if !provider.enabled { continue; @@ -398,6 +546,8 @@ impl EngineConfig { EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), providers, } } @@ -472,6 +622,8 @@ mod tests { let cfg = EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), providers: vec![ProviderConfig { name: "test".into(), kind: "ollama".into(), @@ -489,11 +641,44 @@ 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(), providers: vec![ProviderConfig { name: "bad".into(), kind: "bad_kind".into(), diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index e656d6c..8b77383 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -1,26 +1,37 @@ //! 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, }; -use tokio::sync::broadcast; +use serde::Serialize; +use tokio::sync::{Mutex as AsyncMutex, Semaphore, broadcast}; +use tokio::task::{AbortHandle, JoinHandle}; use crate::backends::{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::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 +41,34 @@ struct RegisteredProvider { provider: Arc, } +/// A single entry in the model lifecycle history. +#[derive(Debug, Clone, Serialize)] +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)] +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 +77,49 @@ 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>>, + /// Weak self-reference for background tasks. + weak_self: OnceLock>, /// Event broadcast channel. event_tx: broadcast::Sender, /// 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 +139,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 { @@ -111,11 +182,26 @@ 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), + weak_self: OnceLock::new(), event_tx, started_at: Instant::now(), }); + let _ = engine.weak_self.set(Arc::downgrade(&engine)); engine.refresh_resources().await; + engine.spawn_idle_eviction(); Ok(engine) } @@ -127,6 +213,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 +221,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 +250,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 +304,7 @@ impl Engine { } async fn refresh_health(&self) { + let health_timeout = self.timeouts.health(); let handles = self .providers .iter() @@ -203,7 +312,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 +338,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 +352,110 @@ impl Engine { models } - /// Run inference: route → circuit-check → load → invoke → optional fallback. + /// 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. pub async fn invoke(&self, req: InferenceRequest) -> Result { self.reject_ambiguous_short_name(&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); + + if candidates.is_empty() { + return Err(EngineError::NoCapableModel(Self::requested_capability( + &req, + ))); + } - let model_id = selected.model.id.clone(); - let provider_name = selected.model.provider.clone(); - let routing_reason = selected.reason.clone(); + 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; + } - 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 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(primary_err) => Err(primary_err), + } + + 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 + } + // 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 +488,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,158 +514,537 @@ 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() ))) } } } + /// 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; + self.unload_model(&victim, "memory_pressure").await; + if self.memory.try_reserve(model_id, bytes) { + return Ok(()); } + } + Err(EngineError::MemoryPressure { + need: bytes, + available: self.memory.available_bytes(), + }) + } - let provider = self - .models - .get(&victim) - .and_then(|card| self.find_provider(&card.provider)); - if let Some(provider) = provider { - let _ = provider.unload(&victim).await; + 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(); + } + + /// 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); + let event = if reason == "idle_timeout" { + "idle_unload" + } else { + "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})"); + } + + /// 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() + .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() + } + + // ----- 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(); } - // 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); } } - /// 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 + /// 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() - .filter(|card| card.execution.active_requests > 0) - .map(|card| card.key().clone()), + .next() + .map(|e| e.key().clone()) + { + self.pending_predictions.remove(&victim); + } + self.pending_predictions.insert(scope, capability); + } + + /// 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 + } + + /// 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)) + } + + 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, ); - protected + 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); + } + + /// 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> { @@ -544,7 +1088,8 @@ impl Engine { }); } } - self.memory.touch(model_id); + // The memory lease protects this model from eviction for the duration. + self.memory.lease(model_id); } fn end_execution(&self, model_id: &str) { @@ -560,7 +1105,7 @@ impl Engine { }); } } - self.memory.touch(model_id); + self.memory.release_lease(model_id); } fn set_model_residency(&self, model_id: &str, new_residency: ModelResidency) { @@ -590,15 +1135,62 @@ impl Engine { && 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(), - }); + 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>) { + 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; @@ -660,6 +1252,23 @@ impl Engine { } } +impl Drop for Engine { + fn drop(&mut self) { + if let Some(handle) = self + .idle_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 +1277,13 @@ 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, TimeoutConfig, + }; - use crate::config::{EngineConfig, ListenConfig, MemoryConfig, ProviderConfig}; + const MB: u64 = 1024 * 1024; fn minimal_config() -> EngineConfig { EngineConfig { @@ -679,51 +1292,102 @@ mod tests { budget_mb: Some(100), idle_timeout_secs: 60, }, + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::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 +1395,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 +1433,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 +1467,68 @@ 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), + weak_self: OnceLock::new(), event_tx, 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 +1543,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 +1565,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 +1586,8 @@ mod tests { budget_mb: Some(100), idle_timeout_secs: 60, }, + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), providers: vec![ProviderConfig { name: "disabled-ollama".into(), kind: "ollama".into(), @@ -866,96 +1600,558 @@ mod tests { }], }; 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 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 - // It cannot be evicted, so the budget stays full. - engine.evict_to_fit("mock/incoming", 40 * 1024 * 1024).await; + // 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 + ); + } - assert!(!engine.memory.can_fit(40 * 1024 * 1024)); + #[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("mock/busy").unwrap().residency, + engine.models.get("ollama/qwen").unwrap().residency, + ModelResidency::Unloaded + ); + } + + #[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; + + engine.invoke(chat_request(None)).await.unwrap(); + assert_eq!(engine.memory.used_bytes(), 20 * MB); + assert_eq!( + 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") + ); } } diff --git a/mofa-engine-core/src/lib.rs b/mofa-engine-core/src/lib.rs index f747665..bbaab94 100644 --- a/mofa-engine-core/src/lib.rs +++ b/mofa-engine-core/src/lib.rs @@ -11,6 +11,7 @@ pub mod engine; pub mod memory; 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..5aa4751 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::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)] +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 { @@ -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 + 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/preflight.rs b/mofa-engine-core/src/preflight.rs index 4a4a4c3..98b1fc7 100644 --- a/mofa-engine-core/src/preflight.rs +++ b/mofa-engine-core/src/preflight.rs @@ -1,96 +1,83 @@ -//! 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::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, } -/// Prediction result. -#[derive(Debug, Clone)] -pub struct Prediction { - /// Predicted next capability - pub capability: Capability, - /// Confidence score (0.0 - 1.0) - pub confidence: f64, -} - -impl PreflightPredictor { - /// Create a new predictor. - pub fn new() -> Self { +impl ScopeChain { + fn new() -> Self { Self { - inner: Mutex::new(PreflightInner { - transitions: HashMap::new(), - total_transitions: 0, - last_capability: None, - }), + transitions: HashMap::new(), + total_transitions: 0, + last_capability: None, + last_access: Instant::now(), } } - /// 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 + fn observe(&mut self, capability: Capability) { + self.last_access = Instant::now(); + if let Some(prev) = self.last_capability { + *self .transitions .entry(prev) .or_default() .entry(capability) - .or_insert(0.0); - *entry += 1.0; - - inner.total_transitions += 1; + .or_insert(0.0) += 1.0; + self.total_transitions += 1; - // Apply decay periodically - if inner.total_transitions.is_multiple_of(DECAY_INTERVAL) { - for counts in inner.transitions.values_mut() { + 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; } } } } - - inner.last_capability = Some(capability); + self.last_capability = Some(capability); } - /// Predict the most likely next capability given the current one. - /// - /// 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()); - - let counts = inner.transitions.get(¤t)?; + fn predict(&self, current: Capability, min_samples: u64) -> Option { + let counts = self.transitions.get(¤t)?; let total: f64 = counts.values().sum(); - - if (total as u64) < MIN_OBSERVATIONS { + if (total as u64) < min_samples { return None; } - counts .iter() .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) @@ -101,60 +88,295 @@ impl PreflightPredictor { } } +/// Prediction result. +#[derive(Debug, Clone, PartialEq)] +pub struct Prediction { + /// Predicted next capability. + pub capability: Capability, + /// Confidence score (0.0 – 1.0). + pub confidence: f64, +} + +impl PreflightPredictor { + /// Create a new predictor. + pub fn new() -> Self { + Self { + inner: Mutex::new(HashMap::new()), + } + } + + fn lock(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.inner.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Record that a capability was used within `scope`. + /// + /// The observation updates both the scope's own chain and the global chain, + /// so global history accumulates without leaking one scope's ordering into + /// another's predictions. + pub fn record(&self, scope: &str, capability: Capability) { + let mut scopes = self.lock(); + Self::evict_if_full(&mut scopes, scope); + scopes + .entry(scope.to_string()) + .or_insert_with(ScopeChain::new) + .observe(capability); + if scope != GLOBAL_SCOPE { + scopes + .entry(GLOBAL_SCOPE.to_string()) + .or_insert_with(ScopeChain::new) + .observe(capability); + } + } + + /// 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() + .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() + } +} + impl Default for PreflightPredictor { fn default() -> Self { Self::new() } } +/// 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, 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); + } + + /// Take a consistent snapshot of all counters. + 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(); - - // Chat → Tts (3x), Chat → Asr (1x) - for _ in 0..3 { - p.record(Capability::Chat); - p.record(Capability::Tts); + // 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); } - p.record(Capability::Chat); - p.record(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 + ); + } - let pred = p.predict(Capability::Chat).unwrap(); + #[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); + } + // 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); - assert!(pred.confidence > 0.7); + } + + #[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..0401367 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; @@ -216,18 +248,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 +310,48 @@ mod tests { } } + #[test] + fn routing_decision_meets_latency_target() { + // RFC budget: a scheduling decision should take < 1ms (excluding load). + // 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 +455,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..0c970e1 --- /dev/null +++ b/mofa-engine-core/src/subscription.rs @@ -0,0 +1,199 @@ +//! 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::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, 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 + } + + /// Drop every expired subscription. + fn prune(&self, now: Instant) { + self.lock().retain(|s| !s.is_expired(now)); + } + + /// The set of capabilities currently kept warm by any live subscription. + pub fn active_capabilities(&self) -> HashSet { + let now = Instant::now(); + self.prune(now); + self.lock() + .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(); + self.prune(now); + self.lock() + .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/src/server.rs b/mofa-engine-sdk/src/server.rs index 68c2623..7361dbd 100644 --- a/mofa-engine-sdk/src/server.rs +++ b/mofa-engine-sdk/src/server.rs @@ -6,17 +6,20 @@ use std::time::Duration; use axum::{ Json, Router, - extract::State, + extract::{Path, State}, http::StatusCode, response::{ Html, 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 tower_http::cors::CorsLayer; @@ -50,6 +53,14 @@ pub async fn start_server( .route("/v1/capabilities", get(capabilities_handler)) .route("/v1/invoke", post(invoke_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)) .layer(CorsLayer::permissive()) @@ -117,6 +128,73 @@ async fn refresh_handler(State(state): State) -> Json) -> 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>> { diff --git a/mofa-kernel/src/types.rs b/mofa-kernel/src/types.rs index 196dabe..9baab0d 100644 --- a/mofa-kernel/src/types.rs +++ b/mofa-kernel/src/types.rs @@ -476,6 +476,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. From f97f3efaebf6d8369b1be65159c2bc5e3e19c228 Mon Sep 17 00:00:00 2001 From: aayank13 Date: Tue, 7 Jul 2026 02:04:31 +0530 Subject: [PATCH 02/22] feat: add streaming support for inference requests - Introduced a new `stream` method in the `Provider` trait to allow providers to emit incremental text deltas during inference. - Added `StreamChunk` enum to represent different types of events in a streaming inference response, including `Started`, `Text`, `Completed`, and `Error`. - Implemented `invoke_stream` in the `DaemonClient` to handle streaming responses over SSE. - Updated the HTTP server to include a new `/v1/invoke/stream` endpoint for streaming inference output. - Enhanced the `Router` to support bearer token authentication for API routes. - Added metrics endpoint for Prometheus scraping. - Improved error handling and request correlation with `x-request-id` headers. - Updated dependencies in `Cargo.toml` for `reqwest` and `tokio`. --- Cargo.lock | 6 + config.example.toml | 31 + docs/openapi.yaml | 132 ++++ mofa-engine-app/Cargo.toml | 3 + mofa-engine-app/src/main.rs | 165 ++++- mofa-engine-core/src/artifacts.rs | 115 ++++ mofa-engine-core/src/backends.rs | 2 + mofa-engine-core/src/backends/local_tts.rs | 479 ++++++++++++++ mofa-engine-core/src/backends/ollama.rs | 4 +- .../src/backends/openai_compat.rs | 32 +- mofa-engine-core/src/circuit_breaker.rs | 86 ++- mofa-engine-core/src/config.rs | 120 +++- mofa-engine-core/src/engine.rs | 619 +++++++++++++++++- mofa-engine-core/src/lib.rs | 2 + mofa-engine-core/src/memory.rs | 8 +- mofa-engine-core/src/metrics.rs | 288 ++++++++ mofa-engine-core/src/preflight.rs | 11 +- mofa-engine-core/src/router.rs | 10 +- mofa-engine-core/src/subscription.rs | 24 +- mofa-engine-sdk/Cargo.toml | 2 + mofa-engine-sdk/src/client.rs | 329 ++++++++++ mofa-engine-sdk/src/lib.rs | 13 +- mofa-engine-sdk/src/server.rs | 314 ++++++++- mofa-kernel/Cargo.toml | 1 + mofa-kernel/src/traits.rs | 27 +- mofa-kernel/src/types.rs | 59 ++ 26 files changed, 2779 insertions(+), 103 deletions(-) create mode 100644 mofa-engine-core/src/artifacts.rs create mode 100644 mofa-engine-core/src/backends/local_tts.rs create mode 100644 mofa-engine-core/src/metrics.rs create mode 100644 mofa-engine-sdk/src/client.rs 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 index 6c8c3f3..b132a08 100644 --- a/config.example.toml +++ b/config.example.toml @@ -29,6 +29,16 @@ confidence_threshold = 0.6 # min confidence (0.0-1.0) before a history predicti 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" @@ -58,3 +68,24 @@ cost_tier = "high" [[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/openapi.yaml b/docs/openapi.yaml index e4b09f4..78b2abd 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -77,6 +77,31 @@ 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) @@ -173,6 +198,64 @@ paths: 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: @@ -234,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] 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..0ed9034 100644 --- a/mofa-engine-app/src/main.rs +++ b/mofa-engine-app/src/main.rs @@ -1,59 +1,172 @@ //! 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, +} + +#[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(DaemonClient::new(d.url).capabilities().await), + Some(Command::Status(d)) => print_json(DaemonClient::new(d.url).status().await), + Some(Command::Refresh(d)) => print_json(DaemonClient::new(d.url).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(DaemonClient::new(daemon.url).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}")) +} - Ok(()) +/// 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}")), + } +} + +/// 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..a2d30b1 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -26,12 +26,14 @@ pub struct OllamaProvider { impl OllamaProvider { /// Create a new Ollama provider. pub fn new(name: impl Into, base_url: impl Into) -> Self { + // A build failure means the system TLS/HTTP stack is unusable; fail + // loudly rather than silently dropping the configured timeouts/no_proxy. let client = Client::builder() .no_proxy() .connect_timeout(Duration::from_secs(5)) .timeout(Duration::from_secs(180)) .build() - .unwrap_or_default(); + .expect("failed to build HTTP client"); Self { name: name.into(), diff --git a/mofa-engine-core/src/backends/openai_compat.rs b/mofa-engine-core/src/backends/openai_compat.rs index 60a1ba9..b5ea556 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, } @@ -39,11 +41,26 @@ impl OpenAiCompatProvider { models: Vec, cost_tier: CostTier, ) -> Self { + 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. + pub fn with_output_dir( + name: impl Into, + base_url: impl Into, + api_key: impl Into, + models: Vec, + cost_tier: CostTier, + output_dir: Option, + ) -> Self { + // A failure here means the system TLS/HTTP stack is unusable, so fail + // loudly rather than falling back to a client without our timeouts. let client = Client::builder() .connect_timeout(Duration::from_secs(10)) .timeout(Duration::from_secs(120)) .build() - .unwrap_or_default(); + .expect("failed to build HTTP client"); Self { name: name.into(), @@ -51,6 +68,10 @@ impl OpenAiCompatProvider { 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 +406,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 +439,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}")) })?; diff --git a/mofa-engine-core/src/circuit_breaker.rs b/mofa-engine-core/src/circuit_breaker.rs index 97f5802..a2ec7c7 100644 --- a/mofa-engine-core/src/circuit_breaker.rs +++ b/mofa-engine-core/src/circuit_breaker.rs @@ -53,6 +53,9 @@ struct ProviderBreaker { /// 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, } @@ -63,15 +66,26 @@ impl ProviderBreaker { 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`). + /// 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, @@ -82,17 +96,29 @@ impl ProviderBreaker { { self.state = CircuitState::HalfOpen; 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 => { - // A probe is already being tested; reject everyone else. 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 } } @@ -100,9 +126,21 @@ impl ProviderBreaker { /// Record a successful request. fn record_success(&mut self) { - self.failure_count = 0; - self.probe_in_flight = false; - 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. @@ -121,6 +159,7 @@ impl ProviderBreaker { // 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 => { @@ -290,6 +329,43 @@ mod tests { 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 38da9a5..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)] @@ -25,11 +25,52 @@ pub struct EngineConfig { /// 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 { @@ -240,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 { @@ -340,6 +418,16 @@ impl EngineConfig { 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; @@ -353,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!( @@ -379,6 +481,7 @@ impl EngineConfig { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }); // OpenAI @@ -423,6 +526,7 @@ impl EngineConfig { }, ], enabled: true, + ..Default::default() }); } @@ -450,6 +554,7 @@ impl EngineConfig { }, ], enabled: true, + ..Default::default() }); } @@ -483,6 +588,7 @@ impl EngineConfig { }, ], enabled: true, + ..Default::default() }); } @@ -502,6 +608,7 @@ impl EngineConfig { memory_mb: None, }], enabled: true, + ..Default::default() }); } @@ -521,6 +628,7 @@ impl EngineConfig { memory_mb: None, }], enabled: true, + ..Default::default() }); } @@ -540,6 +648,7 @@ impl EngineConfig { memory_mb: None, }], enabled: true, + ..Default::default() }); } @@ -548,6 +657,8 @@ impl EngineConfig { memory: MemoryConfig::default(), timeouts: TimeoutConfig::default(), preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers, } } @@ -571,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 @@ -624,6 +736,8 @@ mod tests { memory: MemoryConfig::default(), timeouts: TimeoutConfig::default(), preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers: vec![ProviderConfig { name: "test".into(), kind: "ollama".into(), @@ -633,6 +747,7 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }], }; let toml_str = toml::to_string_pretty(&cfg).unwrap(); @@ -679,6 +794,8 @@ mod tests { 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(), @@ -688,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 8b77383..2352a42 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -13,16 +13,17 @@ use dashmap::DashMap; use mofa_kernel::{ BackendHealth, BackendStatus, Capability, CostTier, EngineError, EngineEvent, EngineStatus, FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, ModelResidency, Provider, - ProviderHealth, ProviderKind, + ProviderHealth, ProviderKind, StreamChunk, model_id_name, }; -use serde::Serialize; -use tokio::sync::{Mutex as AsyncMutex, Semaphore, 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, 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}; @@ -42,7 +43,7 @@ struct RegisteredProvider { } /// A single entry in the model lifecycle history. -#[derive(Debug, Clone, Serialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct LifecycleRecord { /// Monotonic sequence number. pub seq: u64, @@ -57,7 +58,7 @@ pub struct LifecycleRecord { } /// A snapshot of the engine's memory accounting. -#[derive(Debug, Clone, Serialize)] +#[derive(Debug, Clone, Serialize, Deserialize)] pub struct MemoryReport { /// Bytes reserved across all models. pub used_bytes: u64, @@ -107,10 +108,17 @@ pub struct Engine { 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, } @@ -152,13 +160,32 @@ impl Engine { 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( - &pc.name, - &pc.base_url, - pc.api_key.clone().unwrap_or_default(), - pc.models.clone(), - cost_tier, - )), + 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", @@ -175,6 +202,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(), @@ -194,14 +234,18 @@ impl Engine { 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) } @@ -352,6 +396,20 @@ impl Engine { models } + /// 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 @@ -359,8 +417,9 @@ impl Engine { /// 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. - pub async fn invoke(&self, req: InferenceRequest) -> Result { + 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 @@ -577,6 +636,168 @@ impl Engine { } } + /// 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 @@ -1013,6 +1234,60 @@ impl Engine { *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(); @@ -1157,6 +1432,7 @@ impl Engine { } 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, @@ -1246,6 +1522,80 @@ 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() @@ -1254,13 +1604,10 @@ impl Engine { impl Drop for Engine { fn drop(&mut self) { - if let Some(handle) = self - .idle_task - .lock() - .unwrap_or_else(|e| e.into_inner()) - .take() - { - handle.abort(); + 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() { @@ -1280,7 +1627,8 @@ mod tests { use std::sync::atomic::AtomicUsize; use crate::config::{ - EngineConfig, ListenConfig, MemoryConfig, PreflightConfig, ProviderConfig, TimeoutConfig, + EngineConfig, ListenConfig, MemoryConfig, PreflightConfig, ProviderConfig, SecurityConfig, + TimeoutConfig, }; const MB: u64 = 1024 * 1024; @@ -1294,6 +1642,8 @@ mod tests { }, timeouts: TimeoutConfig::default(), preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), providers: vec![], } } @@ -1505,8 +1855,11 @@ mod tests { 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)); @@ -1588,6 +1941,8 @@ mod tests { }, timeouts: TimeoutConfig::default(), preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), providers: vec![ProviderConfig { name: "disabled-ollama".into(), kind: "ollama".into(), @@ -1597,6 +1952,7 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: false, + ..Default::default() }], }; let engine = Engine::new(config).await; @@ -1717,6 +2073,42 @@ mod tests { 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( @@ -2154,4 +2546,187 @@ mod tests { .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 bbaab94..9943486 100644 --- a/mofa-engine-core/src/lib.rs +++ b/mofa-engine-core/src/lib.rs @@ -4,11 +4,13 @@ //! 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; diff --git a/mofa-engine-core/src/memory.rs b/mofa-engine-core/src/memory.rs index 5aa4751..c874fe3 100644 --- a/mofa-engine-core/src/memory.rs +++ b/mofa-engine-core/src/memory.rs @@ -12,7 +12,7 @@ //! - **leases** — in-flight inferences; a model with active leases is never //! evicted, even under pressure. -use serde::Serialize; +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Mutex; use std::time::{Duration, Instant}; @@ -48,7 +48,7 @@ impl Allocation { } /// A point-in-time view of one model's memory accounting. -#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] pub struct AllocationSnapshot { /// Model identifier. pub model_id: String, @@ -66,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(), }; @@ -122,7 +122,7 @@ impl MemoryManager { .filter(|(id, _)| id.as_str() != model_id) .map(|(_, a)| a.reserved_bytes) .sum(); - if used_by_others + bytes > self.budget_bytes { + if used_by_others.saturating_add(bytes) > self.budget_bytes { return false; } match allocs.get_mut(model_id) { diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs new file mode 100644 index 0000000..f92c2c9 --- /dev/null +++ b/mofa-engine-core/src/metrics.rs @@ -0,0 +1,288 @@ +//! 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 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 98b1fc7..1e9e34c 100644 --- a/mofa-engine-core/src/preflight.rs +++ b/mofa-engine-core/src/preflight.rs @@ -7,7 +7,7 @@ //! yet accumulated enough samples. use mofa_kernel::Capability; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Mutex; use std::sync::atomic::{AtomicU64, Ordering}; @@ -204,7 +204,7 @@ pub struct PreflightMetrics { } /// Serializable snapshot of [`PreflightMetrics`]. -#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct PreflightStats { /// Warm tasks spawned. pub warms_started: u64, @@ -252,7 +252,12 @@ impl PreflightMetrics { self.misses.fetch_add(1, Ordering::Relaxed); } - /// Take a consistent snapshot of all counters. + /// 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), diff --git a/mofa-engine-core/src/router.rs b/mofa-engine-core/src/router.rs index 0401367..a478e7b 100644 --- a/mofa-engine-core/src/router.rs +++ b/mofa-engine-core/src/router.rs @@ -204,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 { @@ -230,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 { diff --git a/mofa-engine-core/src/subscription.rs b/mofa-engine-core/src/subscription.rs index 0c970e1..5e30fa1 100644 --- a/mofa-engine-core/src/subscription.rs +++ b/mofa-engine-core/src/subscription.rs @@ -7,7 +7,7 @@ //! be reasoned about and cleaned up per application. use mofa_kernel::Capability; -use serde::Serialize; +use serde::{Deserialize, Serialize}; use std::collections::HashSet; use std::sync::Mutex; use std::sync::atomic::{AtomicU64, Ordering}; @@ -35,7 +35,7 @@ impl Subscription { } /// Public view of a subscription. -#[derive(Debug, Clone, Serialize, PartialEq, Eq)] +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct SubscriptionInfo { /// Unique subscription id. pub id: u64, @@ -93,17 +93,15 @@ impl SubscriptionRegistry { subs.len() != before } - /// Drop every expired subscription. - fn prune(&self, now: Instant) { - self.lock().retain(|s| !s.is_expired(now)); - } - /// 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(); - self.prune(now); - self.lock() - .iter() + let mut subs = self.lock(); + subs.retain(|s| !s.is_expired(now)); + subs.iter() .flat_map(|s| s.capabilities.iter().copied()) .collect() } @@ -116,9 +114,9 @@ impl SubscriptionRegistry { /// List all live subscriptions. pub fn list(&self) -> Vec { let now = Instant::now(); - self.prune(now); - self.lock() - .iter() + 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(), 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..f3bdab0 --- /dev/null +++ b/mofa-engine-sdk/src/client.rs @@ -0,0 +1,329 @@ +//! 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), + /// A response body could not be decoded. + #[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 { + 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 { + 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, +} + +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 } + } + + fn url(&self, path: &str) -> String { + format!("{}/{}", self.base_url, path.trim_start_matches('/')) + } + + async fn get_json(&self, path: &str) -> Result { + let resp = 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 { + if resp.status().is_success() { + resp.json::() + .await + .map_err(|e| ClientError::Decode(e.to_string())) + } else { + // Error responses carry a structured ErrorInfo body. + match resp.json::().await { + Ok(info) => Err(ClientError::Engine(info)), + Err(e) => Err(ClientError::Decode(e.to_string())), + } + } + } + + /// `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 + .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 + .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"); + } +} 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 7361dbd..c211d02 100644 --- a/mofa-engine-sdk/src/server.rs +++ b/mofa-engine-sdk/src/server.rs @@ -6,10 +6,11 @@ use std::time::Duration; use axum::{ Json, Router, - extract::{Path, 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::{delete, get, post}, @@ -21,37 +22,71 @@ 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)) @@ -63,17 +98,61 @@ pub async fn start_server( .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(middleware::from_fn(correlation_middleware)) + .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) .layer(CorsLayer::permissive()) .layer(TraceLayer::new_for_http()) - .with_state(state); + .with_state(state) +} - let addr = format!("{host}:{port}"); - tracing::info!("MoFA Engine listening on http://{addr}"); +/// 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()) + .map(str::to_owned) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); - let listener = tokio::net::TcpListener::bind(&addr).await?; - axum::serve(listener, app).await?; + let span = tracing::info_span!("request", request_id = %request_id); + let mut resp = next.run(req).instrument(span).await; - Ok(()) + 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 ")); + if presented == Some(expected.as_str()) { + next.run(req).await + } else { + let err = EngineError::InvalidRequest("missing or invalid bearer token".into()); + (StatusCode::UNAUTHORIZED, Json(err.info())).into_response() + } } /// Health check response. @@ -118,6 +197,28 @@ 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| { + let data = serde_json::to_string(&chunk).unwrap_or_default(); + 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) @@ -128,6 +229,66 @@ 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()) } @@ -221,6 +382,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() { @@ -232,4 +398,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 9baab0d..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 { From abc57d90f24ecacc3ae4a2a2cfba03f341dc0ffe Mon Sep 17 00:00:00 2001 From: aayank13 Date: Tue, 7 Jul 2026 02:05:53 +0530 Subject: [PATCH 03/22] chore: cargo fmt --- mofa-engine-core/src/engine.rs | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index 2352a42..e92221e 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -160,16 +160,14 @@ impl Engine { 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::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::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!( From 63c20384dd213c7c02e3acf5446313e9d1c55805 Mon Sep 17 00:00:00 2001 From: aayank13 Date: Tue, 7 Jul 2026 19:00:28 +0530 Subject: [PATCH 04/22] feat: add bearer token support for daemon client and improve error handling --- mofa-engine-app/src/main.rs | 25 ++++++-- mofa-engine-core/src/engine.rs | 11 ++-- mofa-engine-core/src/metrics.rs | 12 ++++ mofa-engine-core/src/preflight.rs | 94 +++++++++++++++++++++++-------- mofa-engine-sdk/src/client.rs | 91 ++++++++++++++++++++++++------ mofa-engine-sdk/src/server.rs | 29 +++++++++- 6 files changed, 210 insertions(+), 52 deletions(-) diff --git a/mofa-engine-app/src/main.rs b/mofa-engine-app/src/main.rs index 0ed9034..891be7b 100644 --- a/mofa-engine-app/src/main.rs +++ b/mofa-engine-app/src/main.rs @@ -33,6 +33,23 @@ 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)] @@ -97,9 +114,9 @@ async fn main() -> anyhow::Result<()> { 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(DaemonClient::new(d.url).capabilities().await), - Some(Command::Status(d)) => print_json(DaemonClient::new(d.url).status().await), - Some(Command::Refresh(d)) => print_json(DaemonClient::new(d.url).refresh().await), + 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, @@ -121,7 +138,7 @@ async fn main() -> anyhow::Result<()> { hint_next: None, request_id: String::new(), }; - print_json(DaemonClient::new(daemon.url).invoke(&request).await) + print_json(daemon_client(daemon).invoke(&request).await) } } } diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index e92221e..efa008b 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -908,10 +908,13 @@ impl Engine { } // Frees the reservation and emits residency/memory events. self.set_model_residency(model_id, ModelResidency::Unloaded); - let event = if reason == "idle_timeout" { - "idle_unload" - } else { - "evict" + // 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 { diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs index f92c2c9..3eda177 100644 --- a/mofa-engine-core/src/metrics.rs +++ b/mofa-engine-core/src/metrics.rs @@ -275,6 +275,18 @@ mod tests { 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(); diff --git a/mofa-engine-core/src/preflight.rs b/mofa-engine-core/src/preflight.rs index 1e9e34c..f36ad7d 100644 --- a/mofa-engine-core/src/preflight.rs +++ b/mofa-engine-core/src/preflight.rs @@ -50,26 +50,38 @@ impl ScopeChain { } } + /// 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 - .transitions - .entry(prev) - .or_default() - .entry(capability) - .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; - } + 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; } } } - self.last_capability = Some(capability); } fn predict(&self, current: Capability, min_samples: u64) -> Option { @@ -111,21 +123,34 @@ impl PreflightPredictor { /// Record that a capability was used within `scope`. /// - /// The observation updates both the scope's own chain and the global chain, - /// so global history accumulates without leaking one scope's ordering into - /// another's predictions. + /// 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); - scopes - .entry(scope.to_string()) - .or_insert_with(ScopeChain::new) - .observe(capability); - if scope != GLOBAL_SCOPE { + + // 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 + }; + + // 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) - .observe(capability); + .record_transition(from, capability); } } @@ -354,6 +379,29 @@ mod tests { assert_eq!(pred.capability, Capability::Tts); } + #[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.9); + } + #[test] fn scope_count_is_bounded() { let p = PreflightPredictor::new(); diff --git a/mofa-engine-sdk/src/client.rs b/mofa-engine-sdk/src/client.rs index f3bdab0..1b51f0c 100644 --- a/mofa-engine-sdk/src/client.rs +++ b/mofa-engine-sdk/src/client.rs @@ -33,7 +33,16 @@ pub enum ClientError { /// The daemon returned a structured engine error. #[error("engine error: {}", .0.message)] Engine(ErrorInfo), - /// A response body could not be decoded. + /// 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), } @@ -98,6 +107,10 @@ impl EmbeddedEngine { 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, @@ -108,6 +121,7 @@ impl EmbeddedEngine { /// Remove a subscription by id. pub fn unsubscribe(&self, id: u64) -> bool { + let _guard = self.runtime.enter(); self.engine.unsubscribe(id) } @@ -156,6 +170,7 @@ fn err_json(e: EngineError) -> String { pub struct DaemonClient { base_url: String, http: reqwest::Client, + token: Option, } impl DaemonClient { @@ -167,17 +182,35 @@ impl DaemonClient { /// 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 } + 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 - .http - .get(self.url(path)) + .authed(self.http.get(self.url(path))) .send() .await .map_err(|e| ClientError::Transport(e.to_string()))?; @@ -187,16 +220,23 @@ impl DaemonClient { async fn decode( resp: reqwest::Response, ) -> Result { - if resp.status().is_success() { - resp.json::() + let status = resp.status(); + if status.is_success() { + return resp + .json::() .await - .map_err(|e| ClientError::Decode(e.to_string())) - } else { - // Error responses carry a structured ErrorInfo body. - match resp.json::().await { - Ok(info) => Err(ClientError::Engine(info)), - Err(e) => Err(ClientError::Decode(e.to_string())), - } + .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, + }), } } @@ -236,9 +276,7 @@ impl DaemonClient { request: &InferenceRequest, ) -> Result { let resp = self - .http - .post(self.url("/v1/invoke")) - .json(request) + .authed(self.http.post(self.url("/v1/invoke")).json(request)) .send() .await .map_err(|e| ClientError::Transport(e.to_string()))?; @@ -248,8 +286,7 @@ impl DaemonClient { /// `POST /v1/discovery/refresh` pub async fn refresh(&self) -> Result { let resp = self - .http - .post(self.url("/v1/discovery/refresh")) + .authed(self.http.post(self.url("/v1/discovery/refresh"))) .send() .await .map_err(|e| ClientError::Transport(e.to_string()))?; @@ -326,4 +363,22 @@ mod tests { 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/server.rs b/mofa-engine-sdk/src/server.rs index c211d02..302d7cb 100644 --- a/mofa-engine-sdk/src/server.rs +++ b/mofa-engine-sdk/src/server.rs @@ -111,10 +111,12 @@ fn build_app(state: AppState) -> Router { .route("/health", get(health_handler)) .route("/metrics", get(metrics_handler)) .merge(api) - .layer(middleware::from_fn(correlation_middleware)) .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) .layer(CorsLayer::permissive()) .layer(TraceLayer::new_for_http()) + // 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) } @@ -147,7 +149,10 @@ async fn auth_middleware(State(state): State, req: Request, next: Next .get(header::AUTHORIZATION) .and_then(|v| v.to_str().ok()) .and_then(|v| v.strip_prefix("Bearer ")); - if presented == Some(expected.as_str()) { + 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()); @@ -155,6 +160,19 @@ async fn auth_middleware(State(state): State, req: Request, next: Next } } +/// 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. #[derive(Serialize)] struct HealthResponse { @@ -209,7 +227,12 @@ async fn invoke_stream_handler( ) -> Sse>> { let rx = state.engine.invoke_stream(req); let stream = ReceiverStream::new(rx).map(|chunk| { - let data = serde_json::to_string(&chunk).unwrap_or_default(); + // 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( From 9aaeab2a3c83b2a21c23364619dfc501a09bd07c Mon Sep 17 00:00:00 2001 From: aayank13 Date: Tue, 7 Jul 2026 19:19:17 +0530 Subject: [PATCH 05/22] feat: enhance provider initialization to return Result for better error handling --- mofa-engine-core/src/backends/ollama.rs | 16 ++++++------ .../src/backends/openai_compat.rs | 25 ++++++++++++------- mofa-engine-core/src/engine.rs | 4 +-- mofa-engine-core/src/preflight.rs | 4 ++- mofa-engine-core/src/router.rs | 3 +++ mofa-engine-sdk/src/server.rs | 6 +++++ 6 files changed, 39 insertions(+), 19 deletions(-) diff --git a/mofa-engine-core/src/backends/ollama.rs b/mofa-engine-core/src/backends/ollama.rs index a2d30b1..71e7115 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -25,21 +25,23 @@ pub struct OllamaProvider { impl OllamaProvider { /// Create a new Ollama provider. - pub fn new(name: impl Into, base_url: impl Into) -> Self { - // A build failure means the system TLS/HTTP stack is unusable; fail - // loudly rather than silently dropping the configured timeouts/no_proxy. + /// + /// 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() - .expect("failed to build HTTP client"); + .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 { @@ -364,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 b5ea556..45713a4 100644 --- a/mofa-engine-core/src/backends/openai_compat.rs +++ b/mofa-engine-core/src/backends/openai_compat.rs @@ -40,12 +40,15 @@ 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, @@ -53,16 +56,18 @@ impl OpenAiCompatProvider { models: Vec, cost_tier: CostTier, output_dir: Option, - ) -> Self { - // A failure here means the system TLS/HTTP stack is unusable, so fail - // loudly rather than falling back to a client without our timeouts. + ) -> Result { let client = Client::builder() .connect_timeout(Duration::from_secs(10)) .timeout(Duration::from_secs(120)) .build() - .expect("failed to build HTTP client"); + .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(), @@ -73,7 +78,7 @@ impl OpenAiCompatProvider { .map(std::path::PathBuf::from) .unwrap_or_else(std::env::temp_dir), client, - } + }) } } @@ -529,7 +534,8 @@ mod tests { }, ], CostTier::Medium, - ); + ) + .unwrap(); let cards = provider.discover().await.unwrap(); assert_eq!(cards.len(), 2); @@ -542,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/engine.rs b/mofa-engine-core/src/engine.rs index efa008b..ded37b3 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -159,7 +159,7 @@ 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::Ollama => Arc::new(OllamaProvider::new(&pc.name, &pc.base_url)?), ProviderKind::OpenAiCompatible => Arc::new(OpenAiCompatProvider::with_output_dir( &pc.name, &pc.base_url, @@ -167,7 +167,7 @@ impl Engine { pc.models.clone(), cost_tier, config.artifacts.dir.clone(), - )), + )?), ProviderKind::LocalTts => { let command = pc.command.clone().ok_or_else(|| { EngineError::Config(format!( diff --git a/mofa-engine-core/src/preflight.rs b/mofa-engine-core/src/preflight.rs index f36ad7d..70012b4 100644 --- a/mofa-engine-core/src/preflight.rs +++ b/mofa-engine-core/src/preflight.rs @@ -87,7 +87,9 @@ impl ScopeChain { fn predict(&self, current: Capability, min_samples: u64) -> Option { let counts = self.transitions.get(¤t)?; let total: f64 = counts.values().sum(); - if (total as u64) < min_samples { + // 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 diff --git a/mofa-engine-core/src/router.rs b/mofa-engine-core/src/router.rs index a478e7b..0c036b1 100644 --- a/mofa-engine-core/src/router.rs +++ b/mofa-engine-core/src/router.rs @@ -309,8 +309,11 @@ 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), diff --git a/mofa-engine-sdk/src/server.rs b/mofa-engine-sdk/src/server.rs index 302d7cb..2af2965 100644 --- a/mofa-engine-sdk/src/server.rs +++ b/mofa-engine-sdk/src/server.rs @@ -120,6 +120,11 @@ fn build_app(state: AppState) -> Router { .with_state(state) } +/// 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; + /// 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 { @@ -127,6 +132,7 @@ async fn correlation_middleware(req: Request, next: Next) -> Response { .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()); From 5287b436cd3d49beb03bbf4f516ad1a3ec94fab0 Mon Sep 17 00:00:00 2001 From: aayank13 Date: Sun, 12 Jul 2026 20:31:04 +0530 Subject: [PATCH 06/22] feat: prefer=local routing, token/cost observability, real Ollama streaming --- config.example.toml | 4 + docs/openapi.yaml | 17 ++ mofa-engine-app/src/main.rs | 1 + mofa-engine-core/src/backends/local_tts.rs | 2 + mofa-engine-core/src/backends/ollama.rs | 199 +++++++++++++++++- .../src/backends/openai_compat.rs | 14 +- mofa-engine-core/src/config.rs | 9 + mofa-engine-core/src/engine.rs | 120 ++++++++++- mofa-engine-core/src/metrics.rs | 85 ++++++++ mofa-engine-core/src/router.rs | 102 ++++++++- mofa-engine-sdk/src/client.rs | 1 + mofa-kernel/src/types.rs | 61 +++++- 12 files changed, 598 insertions(+), 17 deletions(-) diff --git a/config.example.toml b/config.example.toml index b132a08..f7de69b 100644 --- a/config.example.toml +++ b/config.example.toml @@ -55,6 +55,10 @@ base_url = "https://api.openai.com/v1" api_key = "${OPENAI_API_KEY}" priority = 10 cost_tier = "high" +# Optional token pricing (USD per 1000 tokens) for cost tracking in /metrics and +# per-response `cost_usd`. Omit or set 0 for free/local models. +price_input_per_1k = 0.0025 +price_output_per_1k = 0.01 [[providers.models]] name = "gpt-4o" diff --git a/docs/openapi.yaml b/docs/openapi.yaml index 78b2abd..c6edc90 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -275,6 +275,13 @@ components: type: string enum: [capability_only, disabled, allow_named] default: capability_only + locality: + type: string + enum: [auto, prefer_local, local_only] + default: auto + description: > + Routing locality guardrail. `local_only` restricts the request to + local backends (privacy); `prefer_local` ranks local ahead of cloud. messages: type: array items: @@ -313,6 +320,13 @@ components: type: string tokens_used: type: [integer, "null"] + prompt_tokens: + type: [integer, "null"] + completion_tokens: + type: [integer, "null"] + cost_usd: + type: [number, "null"] + description: Estimated USD cost, if a price is configured for the provider. fallback_used: type: boolean routing_reason: @@ -340,6 +354,9 @@ components: type: { const: completed } duration_ms: { type: integer } tokens_used: { type: [integer, "null"] } + prompt_tokens: { type: [integer, "null"] } + completion_tokens: { type: [integer, "null"] } + cost_usd: { type: [number, "null"] } file: { type: [string, "null"] } fallback_used: { type: boolean } routing_reason: { type: [string, "null"] } diff --git a/mofa-engine-app/src/main.rs b/mofa-engine-app/src/main.rs index 891be7b..63fcc15 100644 --- a/mofa-engine-app/src/main.rs +++ b/mofa-engine-app/src/main.rs @@ -137,6 +137,7 @@ async fn main() -> anyhow::Result<()> { params: serde_json::Value::Null, hint_next: None, request_id: String::new(), + ..Default::default() }; print_json(daemon_client(daemon).invoke(&request).await) } diff --git a/mofa-engine-core/src/backends/local_tts.rs b/mofa-engine-core/src/backends/local_tts.rs index 4959866..53db6b4 100644 --- a/mofa-engine-core/src/backends/local_tts.rs +++ b/mofa-engine-core/src/backends/local_tts.rs @@ -205,6 +205,7 @@ impl LocalTtsProvider { tokens_used: None, fallback_used: false, routing_reason: None, + ..Default::default() }) } } @@ -342,6 +343,7 @@ mod tests { params: serde_json::Value::Null, hint_next: None, request_id: "test".into(), + ..Default::default() } } diff --git a/mofa-engine-core/src/backends/ollama.rs b/mofa-engine-core/src/backends/ollama.rs index 71e7115..fd19ef2 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -6,12 +6,12 @@ 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, + ProviderKind, StreamSink, canonical_model_id, model_id_name, }; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::collections::HashSet; -use std::time::Duration; +use std::time::{Duration, Instant}; /// Provider for a local Ollama instance. pub struct OllamaProvider { @@ -106,12 +106,41 @@ struct OllamaMessage { #[allow(dead_code)] struct OllamaChatResponse { message: Option, + /// Completion/output tokens. #[serde(default)] eval_count: Option, + /// Prompt/input tokens. + #[serde(default)] + prompt_eval_count: Option, #[serde(default)] total_duration: Option, } +/// One NDJSON object from Ollama's `stream: true` response. +#[derive(Debug, Deserialize)] +struct OllamaStreamChunk { + /// Incremental message with the next token(s) of content. + message: Option, + /// Set on the final object. + #[serde(default)] + done: bool, + /// Completion/output tokens (final object only). + #[serde(default)] + eval_count: Option, + /// Prompt/input tokens (final object only). + #[serde(default)] + prompt_eval_count: Option, +} + +/// Parse one NDJSON line from an Ollama stream. Blank/whitespace lines and +/// unparseable fragments yield `None`. +fn parse_stream_line(line: &[u8]) -> Option { + if line.iter().all(u8::is_ascii_whitespace) { + return None; + } + serde_json::from_slice(line).ok() +} + #[async_trait] impl Provider for OllamaProvider { fn name(&self) -> &str { @@ -345,6 +374,12 @@ impl Provider for OllamaProvider { let duration_ms = start.elapsed().as_millis() as u64; let text = chat_resp.message.map(|m| m.content).unwrap_or_default(); + let prompt_tokens = chat_resp.prompt_eval_count; + let completion_tokens = chat_resp.eval_count; + let tokens_used = match (prompt_tokens, completion_tokens) { + (None, None) => None, + (p, c) => Some(p.unwrap_or(0) + c.unwrap_or(0)), + }; Ok(InferenceResponse { text: Some(text), @@ -353,9 +388,121 @@ impl Provider for OllamaProvider { provider: self.name.clone(), duration_ms, request_id: request.request_id.clone(), - tokens_used: chat_resp.eval_count, + tokens_used, + prompt_tokens, + completion_tokens, + fallback_used: false, + routing_reason: None, + ..Default::default() + }) + } + + async fn stream( + &self, + model_id: &str, + request: &InferenceRequest, + sink: StreamSink, + ) -> Result { + let model_name = model_id_name(model_id); + let messages: Vec = request + .messages + .iter() + .map(|m| OllamaMessage { + role: m.role.clone(), + content: m.content.clone(), + }) + .collect(); + if messages.is_empty() { + return Err(EngineError::InvalidRequest("no messages provided".into())); + } + + let body = OllamaChatRequest { + model: model_name.to_string(), + messages, + stream: true, + keep_alive: None, + }; + let url = format!("{}/api/chat", self.base_url); + let start = Instant::now(); + + let mut resp = self + .client + .post(&url) + .json(&body) + .send() + .await + .map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: e.to_string(), + })?; + + if !resp.status().is_success() { + let status = resp.status(); + let text = resp.text().await.unwrap_or_default(); + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("HTTP {status}: {text}"), + }); + } + + // Ollama streams newline-delimited JSON. Buffer bytes and emit each + // token's content as it completes a line, accumulating the full text. + let mut buf: Vec = Vec::new(); + let mut full = String::new(); + let mut prompt_tokens = None; + let mut completion_tokens = None; + + let mut apply = |chunk: OllamaStreamChunk, full: &mut String| -> Option { + if chunk.done { + prompt_tokens = chunk.prompt_eval_count; + completion_tokens = chunk.eval_count; + } + let content = chunk.message.map(|m| m.content)?; + if content.is_empty() { + return None; + } + full.push_str(&content); + Some(content) + }; + + while let Some(bytes) = resp.chunk().await.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("stream read error: {e}"), + })? { + buf.extend_from_slice(&bytes); + while let Some(pos) = buf.iter().position(|&b| b == b'\n') { + let line: Vec = buf.drain(..=pos).collect(); + if let Some(chunk) = parse_stream_line(&line) + && let Some(delta) = apply(chunk, &mut full) + { + let _ = sink.send(delta).await; + } + } + } + // Any trailing bytes without a final newline. + if let Some(chunk) = parse_stream_line(&buf) + && let Some(delta) = apply(chunk, &mut full) + { + let _ = sink.send(delta).await; + } + + let tokens_used = match (prompt_tokens, completion_tokens) { + (None, None) => None, + (p, c) => Some(p.unwrap_or(0) + c.unwrap_or(0)), + }; + Ok(InferenceResponse { + text: Some(full), + file: None, + 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, + prompt_tokens, + completion_tokens, fallback_used: false, routing_reason: None, + ..Default::default() }) } } @@ -377,4 +524,50 @@ mod tests { assert_eq!(model_id_name("ollama/llama3:8b"), "llama3:8b"); assert_eq!(model_id_name("ollama::llama3:8b"), "llama3:8b"); } + + #[test] + fn ndjson_stream_lines_accumulate_text_and_tokens() { + // A realistic Ollama `stream: true` sequence, including a blank line that + // must be ignored and a final `done` object carrying token counts. + let lines: Vec<&[u8]> = vec![ + br#"{"message":{"role":"assistant","content":"He"},"done":false}"#, + b"", + br#"{"message":{"role":"assistant","content":"llo"},"done":false}"#, + br#"{"message":{"role":"assistant","content":""},"done":false}"#, + br#"{"message":{"role":"assistant","content":"!"},"done":false}"#, + br#"{"done":true,"eval_count":3,"prompt_eval_count":5}"#, + ]; + + let mut full = String::new(); + let mut deltas = Vec::new(); + let mut prompt = None; + let mut completion = None; + for line in lines { + if let Some(chunk) = parse_stream_line(line) { + if chunk.done { + prompt = chunk.prompt_eval_count; + completion = chunk.eval_count; + } + if let Some(content) = chunk.message.map(|m| m.content) + && !content.is_empty() + { + full.push_str(&content); + deltas.push(content); + } + } + } + + // Multiple real deltas (not one lump), empty content skipped. + assert_eq!(deltas, vec!["He", "llo", "!"]); + assert_eq!(full, "Hello!"); + assert_eq!(prompt, Some(5)); + assert_eq!(completion, Some(3)); + } + + #[test] + fn parse_stream_line_ignores_blank_and_garbage() { + assert!(parse_stream_line(b"").is_none()); + assert!(parse_stream_line(b" \n").is_none()); + assert!(parse_stream_line(b"not json").is_none()); + } } diff --git a/mofa-engine-core/src/backends/openai_compat.rs b/mofa-engine-core/src/backends/openai_compat.rs index 45713a4..cc6fcef 100644 --- a/mofa-engine-core/src/backends/openai_compat.rs +++ b/mofa-engine-core/src/backends/openai_compat.rs @@ -110,6 +110,10 @@ struct ChatChoice { #[derive(Debug, Deserialize)] struct ChatUsage { total_tokens: Option, + #[serde(default)] + prompt_tokens: Option, + #[serde(default)] + completion_tokens: Option, } #[derive(Debug, Serialize)] @@ -335,7 +339,10 @@ impl OpenAiCompatProvider { .and_then(|c| c.message.as_ref()) .map(|m| m.content.clone()); - let tokens = chat.usage.and_then(|u| u.total_tokens); + let (tokens, prompt_tokens, completion_tokens) = match chat.usage { + Some(u) => (u.total_tokens, u.prompt_tokens, u.completion_tokens), + None => (None, None, None), + }; let duration_ms = start.elapsed().as_millis() as u64; Ok(InferenceResponse { @@ -346,8 +353,11 @@ impl OpenAiCompatProvider { duration_ms, request_id: request.request_id.clone(), tokens_used: tokens, + prompt_tokens, + completion_tokens, fallback_used: false, routing_reason: None, + ..Default::default() }) } @@ -430,6 +440,7 @@ impl OpenAiCompatProvider { tokens_used: None, fallback_used: false, routing_reason: None, + ..Default::default() }) } @@ -505,6 +516,7 @@ impl OpenAiCompatProvider { tokens_used: None, fallback_used: false, routing_reason: None, + ..Default::default() }) } } diff --git a/mofa-engine-core/src/config.rs b/mofa-engine-core/src/config.rs index f39927c..4309760 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -299,6 +299,13 @@ pub struct ProviderConfig { /// Directory for `local_tts` audio artifacts (default: the system temp dir). #[serde(default)] pub output_dir: Option, + /// USD price per 1000 prompt (input) tokens. Used for cost tracking; `0` + /// (the default, e.g. local models) means no cost is attributed. + #[serde(default)] + pub price_input_per_1k: f64, + /// USD price per 1000 completion (output) tokens. + #[serde(default)] + pub price_output_per_1k: f64, } impl Default for ProviderConfig { @@ -316,6 +323,8 @@ impl Default for ProviderConfig { args: Vec::new(), output_format: None, output_dir: None, + price_input_per_1k: 0.0, + price_output_per_1k: 0.0, } } } diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index ded37b3..b4c6a09 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -4,7 +4,7 @@ //! circuit breaking, reservation-based memory admission, concurrency control, //! idle eviction, and observability. -use std::collections::{HashSet, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, OnceLock, Weak}; use std::time::{Duration, Instant}; @@ -116,6 +116,8 @@ pub struct Engine { event_tx: broadcast::Sender, /// Bounded-cardinality metrics counters. metrics: EngineMetrics, + /// Per-provider token pricing: name → (USD/1k prompt, USD/1k completion). + pricing: HashMap, /// Canonicalized allowlist of roots for request `input_file` paths; empty /// means any local path is accepted. input_roots: Vec, @@ -128,6 +130,22 @@ fn remaining(deadline: Instant) -> Duration { deadline.saturating_duration_since(Instant::now()) } +/// USD cost from per-1k-token prices. `None` when both prices are zero (free / +/// local models), so no spurious zero cost is attributed. +fn estimate_cost( + price_in_per_1k: f64, + price_out_per_1k: f64, + prompt_tokens: Option, + completion_tokens: Option, +) -> Option { + if price_in_per_1k == 0.0 && price_out_per_1k == 0.0 { + return None; + } + let prompt = prompt_tokens.unwrap_or(0) as f64; + let completion = completion_tokens.unwrap_or(0) as f64; + Some(prompt / 1000.0 * price_in_per_1k + completion / 1000.0 * price_out_per_1k) +} + impl Engine { /// Create and initialize a new engine from configuration. /// @@ -213,6 +231,19 @@ impl Engine { Duration::from_secs(config.artifacts.retention_secs), ); + // Per-provider token pricing for cost tracking. + let pricing = config + .providers + .iter() + .filter(|p| p.enabled) + .map(|p| { + ( + p.name.clone(), + (p.price_input_per_1k, p.price_output_per_1k), + ) + }) + .collect::>(); + let engine = Arc::new(Self { providers, models: DashMap::new(), @@ -236,6 +267,7 @@ impl Engine { weak_self: OnceLock::new(), event_tx, metrics: EngineMetrics::default(), + pricing, input_roots, started_at: Instant::now(), }); @@ -394,20 +426,42 @@ impl Engine { models } - /// Run inference, recording request-level metrics around the attempt. + /// Run inference, recording request-level metrics + token/cost usage. pub async fn invoke(&self, req: InferenceRequest) -> Result { let started = Instant::now(); - let result = self.invoke_inner(req).await; + let mut 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), + match &mut result { + Ok(resp) => { + resp.cost_usd = + self.compute_cost(&resp.provider, resp.prompt_tokens, resp.completion_tokens); + self.metrics + .record_request(true, duration_ms, resp.fallback_used); + self.metrics.record_usage( + &resp.provider, + resp.prompt_tokens, + resp.completion_tokens, + resp.cost_usd, + ); + } Err(_) => self.metrics.record_request(false, duration_ms, false), } result } + /// Estimated USD cost for a request served by `provider`, from the + /// configured per-1k-token prices. `None` when no price is configured + /// (e.g. local models) or no tokens were reported. + fn compute_cost( + &self, + provider: &str, + prompt_tokens: Option, + completion_tokens: Option, + ) -> Option { + let (price_in, price_out) = self.pricing.get(provider).copied()?; + estimate_cost(price_in, price_out, prompt_tokens, completion_tokens) + } + /// Run inference. /// /// Builds a single ranked candidate plan and walks it in order: the first @@ -767,9 +821,20 @@ impl Engine { self.record_transition(&req); self.emit_request_completed(&req.request_id, duration_ms, true); self.metrics.record_request(true, duration_ms, false); + let cost_usd = + self.compute_cost(&provider_name, resp.prompt_tokens, resp.completion_tokens); + self.metrics.record_usage( + &provider_name, + resp.prompt_tokens, + resp.completion_tokens, + cost_usd, + ); StreamChunk::Completed { duration_ms, tokens_used: resp.tokens_used, + prompt_tokens: resp.prompt_tokens, + completion_tokens: resp.completion_tokens, + cost_usd, file: resp.file, fallback_used: false, routing_reason: Some(routing_reason), @@ -1101,6 +1166,7 @@ impl Engine { params: serde_json::Value::Null, hint_next: None, request_id: String::new(), + ..Default::default() }; let budget = Some(self.memory.budget_bytes()); let Some(best) = Router::route_ranked(&all, &route_req, &providers, budget) @@ -1814,6 +1880,7 @@ mod tests { tokens_used: Some(1), fallback_used: false, routing_reason: None, + ..Default::default() }) } } @@ -1860,6 +1927,7 @@ mod tests { weak_self: OnceLock::new(), event_tx, metrics: EngineMetrics::default(), + pricing: HashMap::new(), input_roots: Vec::new(), started_at: Instant::now(), }); @@ -1891,6 +1959,7 @@ mod tests { params: serde_json::Value::Null, hint_next: None, request_id: "test".into(), + ..Default::default() } } @@ -1900,6 +1969,17 @@ mod tests { assert!(engine.capabilities().await.is_empty()); } + #[test] + fn cost_estimation_from_prices() { + // $0.01/1k in, $0.03/1k out; 1000 prompt + 500 completion tokens. + let cost = estimate_cost(0.01, 0.03, Some(1000), Some(500)).unwrap(); + assert!((cost - (0.01 + 0.015)).abs() < 1e-9); + // Zero prices (free / local model) → no cost attributed. + assert!(estimate_cost(0.0, 0.0, Some(1000), Some(500)).is_none()); + // Missing token counts default to zero. + assert_eq!(estimate_cost(0.01, 0.03, None, None), Some(0.0)); + } + #[tokio::test] async fn status_reports_uptime() { let engine = Engine::new(minimal_config()).await; @@ -2026,6 +2106,31 @@ mod tests { assert!(!resp.fallback_used); } + #[tokio::test] + async fn local_only_never_routes_to_cloud() { + // S5 privacy guardrail: a local_only request must be served locally or + // fail — never fall through to a cloud provider. + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible).with_model( + "gpt", + Capability::Chat, + ModelResidency::Remote, + 0, + 32, + ), + ); + let engine = build_engine(vec![cloud], 1000, 0); + engine.refresh_resources().await; + + let mut req = chat_request(None); + req.locality = mofa_kernel::Locality::LocalOnly; + // Only a cloud model exists → the request fails rather than leaving the device. + assert!(matches!( + engine.invoke(req).await, + Err(EngineError::NoCapableModel(_)) + )); + } + #[tokio::test] async fn named_request_routes_exactly() { let local = Arc::new( @@ -2614,6 +2719,7 @@ mod tests { tokens_used: Some(self.deltas.len() as u32), fallback_used: false, routing_reason: None, + ..Default::default() }) } async fn stream( diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs index 3eda177..f671355 100644 --- a/mofa-engine-core/src/metrics.rs +++ b/mofa-engine-core/src/metrics.rs @@ -7,10 +7,22 @@ use std::sync::atomic::{AtomicU64, Ordering}; +use dashmap::DashMap; + /// 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]; +/// Per-provider token and cost aggregation (dual-track observability). +#[derive(Default)] +struct ProviderUsage { + prompt_tokens: AtomicU64, + completion_tokens: AtomicU64, + /// Accumulated cost in micro-USD (millionths of a dollar) to avoid float + /// atomics; rendered back to USD. + cost_micro_usd: AtomicU64, +} + /// Process-wide engine counters. #[derive(Default)] pub struct EngineMetrics { @@ -25,6 +37,8 @@ pub struct EngineMetrics { evictions_total: AtomicU64, idle_unloads_total: AtomicU64, load_failures_total: AtomicU64, + /// Token/cost totals keyed by provider (bounded by the fixed provider set). + per_provider: DashMap, } impl EngineMetrics { @@ -69,6 +83,31 @@ impl EngineMetrics { self.requests_total.load(Ordering::Relaxed) } + /// Record per-provider token usage and cost for one successful request. + pub fn record_usage( + &self, + provider: &str, + prompt_tokens: Option, + completion_tokens: Option, + cost_usd: Option, + ) { + let entry = self.per_provider.entry(provider.to_string()).or_default(); + if let Some(p) = prompt_tokens { + entry.prompt_tokens.fetch_add(p as u64, Ordering::Relaxed); + } + if let Some(c) = completion_tokens { + entry + .completion_tokens + .fetch_add(c as u64, Ordering::Relaxed); + } + if let Some(cost) = cost_usd + && cost > 0.0 + { + let micro = (cost * 1_000_000.0).round() as u64; + entry.cost_micro_usd.fetch_add(micro, Ordering::Relaxed); + } + } + /// Render the Prometheus text-exposition body for these counters. /// /// `gauges` are point-in-time values pulled from the engine (memory, model @@ -209,6 +248,34 @@ impl EngineMetrics { )); } + // Dual-track observability: per-provider token and cost totals. + out.push_str( + "# HELP mofa_tokens_total Tokens processed, by provider and direction.\n\ + # TYPE mofa_tokens_total counter\n", + ); + for entry in self.per_provider.iter() { + let provider = escape_label_value(entry.key()); + out.push_str(&format!( + "mofa_tokens_total{{provider=\"{provider}\",direction=\"prompt\"}} {}\n", + entry.prompt_tokens.load(Ordering::Relaxed) + )); + out.push_str(&format!( + "mofa_tokens_total{{provider=\"{provider}\",direction=\"completion\"}} {}\n", + entry.completion_tokens.load(Ordering::Relaxed) + )); + } + out.push_str( + "# HELP mofa_cost_usd_total Estimated spend in USD, by provider.\n\ + # TYPE mofa_cost_usd_total counter\n", + ); + for entry in self.per_provider.iter() { + let usd = entry.cost_micro_usd.load(Ordering::Relaxed) as f64 / 1_000_000.0; + out.push_str(&format!( + "mofa_cost_usd_total{{provider=\"{}\"}} {usd}\n", + escape_label_value(entry.key()) + )); + } + out } } @@ -275,6 +342,24 @@ mod tests { assert!(text.contains("mofa_provider_up{provider=\"ollama\"} 1")); } + #[test] + fn per_provider_token_and_cost_render() { + let m = EngineMetrics::default(); + m.record_usage("openai", Some(100), Some(40), Some(0.0021)); + m.record_usage("openai", Some(50), Some(10), Some(0.0009)); + m.record_usage("ollama", Some(30), Some(20), None); // local → no cost + + let text = m.render_prometheus(&MetricsGauges::default()); + assert!(text.contains("mofa_tokens_total{provider=\"openai\",direction=\"prompt\"} 150")); + assert!( + text.contains("mofa_tokens_total{provider=\"openai\",direction=\"completion\"} 50") + ); + assert!(text.contains("mofa_tokens_total{provider=\"ollama\",direction=\"prompt\"} 30")); + // 0.0021 + 0.0009 = 0.003 USD accumulated for openai; ollama has none. + assert!(text.contains("mofa_cost_usd_total{provider=\"openai\"} 0.003")); + assert!(text.contains("mofa_cost_usd_total{provider=\"ollama\"} 0")); + } + #[test] fn provider_label_values_are_escaped() { let m = EngineMetrics::default(); diff --git a/mofa-engine-core/src/router.rs b/mofa-engine-core/src/router.rs index 0c036b1..0acb427 100644 --- a/mofa-engine-core/src/router.rs +++ b/mofa-engine-core/src/router.rs @@ -3,10 +3,14 @@ //! The router applies hard constraints first, then scores valid candidates. use mofa_kernel::{ - BackendHealth, Capability, CostTier, InferenceRequest, ModelCard, ModelResidency, ModelStatus, - ProviderKind, + BackendHealth, Capability, CostTier, InferenceRequest, Locality, ModelCard, ModelResidency, + ModelStatus, ProviderKind, }; +/// Score bonus that guarantees a `prefer_local` request ranks every local +/// candidate ahead of any cloud one. Larger than any achievable base score. +const LOCAL_FIRST_BONUS: i64 = 1_000_000_000; + /// Provider facts needed by routing. #[derive(Debug, Clone)] pub struct RoutingProvider { @@ -86,6 +90,12 @@ impl Router { continue; } + // Hard locality constraint (privacy / data-residency guardrail): a + // `local_only` request must never leave the device. + if request.locality == Locality::LocalOnly && !provider.kind.is_local() { + continue; + } + if model.availability == mofa_kernel::ModelAvailability::Unavailable { continue; } @@ -100,7 +110,13 @@ impl Router { continue; } - let score = Self::score(model, desired_cap, provider); + let mut score = Self::score(model, desired_cap, provider); + // `prefer_local` ranks every local candidate strictly ahead of any + // cloud one (cloud stays available only as a fallback) by adding a + // bonus larger than any achievable base score. + if request.locality == Locality::PreferLocal && provider.kind.is_local() { + score = score.saturating_add(LOCAL_FIRST_BONUS); + } if score <= 0 { continue; } @@ -305,6 +321,7 @@ mod tests { params: serde_json::Value::Null, hint_next: None, request_id: "test".into(), + ..Default::default() } } @@ -553,4 +570,83 @@ mod tests { providers[0].health = BackendHealth::Unavailable; assert!(Router::route(&models, &request(Capability::Chat, None), &providers).is_none()); } + + #[test] + fn local_only_excludes_cloud_candidates() { + 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 mut req = request(Capability::Chat, None); + req.locality = Locality::LocalOnly; + let ranked = Router::route_ranked(&models, &req, &providers, None); + // Only the local candidate survives the hard constraint. + assert_eq!(ranked.len(), 1); + assert_eq!(ranked[0].model.name, "local"); + } + + #[test] + fn local_only_with_no_local_returns_empty() { + let models = vec![make_model( + "cloud", + "openai", + Capability::Chat, + ModelResidency::Remote, + CostTier::High, + )]; + let providers = vec![provider("openai", ProviderKind::OpenAiCompatible, 10)]; + let mut req = request(Capability::Chat, None); + req.locality = Locality::LocalOnly; + // No local candidate → no route (engine turns this into NoCapableModel). + assert!(Router::route_ranked(&models, &req, &providers, None).is_empty()); + } + + #[test] + fn prefer_local_ranks_local_first_but_keeps_cloud_fallback() { + // A cloud model that is "hotter" than the local one would normally win on + // residency; prefer_local must still rank the local candidate first while + // keeping the cloud one as a fallback. + let models = vec![ + make_model( + "cloud", + "openai", + Capability::Chat, + ModelResidency::Remote, + CostTier::High, + ), + make_model( + "local", + "ollama", + Capability::Chat, + ModelResidency::Unloaded, + CostTier::Free, + ), + ]; + let providers = vec![ + provider("openai", ProviderKind::OpenAiCompatible, 10), + provider("ollama", ProviderKind::Ollama, 1), + ]; + let mut req = request(Capability::Chat, None); + req.locality = Locality::PreferLocal; + let ranked = Router::route_ranked(&models, &req, &providers, None); + assert_eq!(ranked.len(), 2); + assert_eq!(ranked[0].model.name, "local"); + assert_eq!(ranked[1].model.name, "cloud"); + } } diff --git a/mofa-engine-sdk/src/client.rs b/mofa-engine-sdk/src/client.rs index 1b51f0c..289d107 100644 --- a/mofa-engine-sdk/src/client.rs +++ b/mofa-engine-sdk/src/client.rs @@ -333,6 +333,7 @@ mod tests { params: serde_json::Value::Null, hint_next: None, request_id: "t".into(), + ..Default::default() }; assert!(matches!( engine.invoke(req), diff --git a/mofa-kernel/src/types.rs b/mofa-kernel/src/types.rs index 3eabccb..c22e549 100644 --- a/mofa-kernel/src/types.rs +++ b/mofa-kernel/src/types.rs @@ -356,6 +356,22 @@ pub enum FallbackPolicy { AllowNamed, } +/// Locality preference for routing (privacy / data-residency guardrail). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum Locality { + /// Normal ranking: local models are preferred but a cloud model can be + /// selected or used as a fallback. + #[default] + Auto, + /// Local models are ranked strictly ahead of any cloud model, which remains + /// available only as a fallback. + PreferLocal, + /// Hard constraint: only local models may serve the request. If none can, + /// the request fails rather than leaving the device. + LocalOnly, +} + /// A request to the engine for inference. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct InferenceRequest { @@ -370,6 +386,9 @@ pub struct InferenceRequest { /// Named fallback policy. #[serde(default)] pub fallback_policy: FallbackPolicy, + /// Locality preference (privacy guardrail). + #[serde(default)] + pub locality: Locality, /// Conversation messages. #[serde(default)] pub messages: Vec, @@ -385,12 +404,30 @@ pub struct InferenceRequest { pub request_id: String, } +impl Default for InferenceRequest { + fn default() -> Self { + Self { + capability: None, + model: None, + app_id: None, + session_id: None, + fallback_policy: FallbackPolicy::default(), + locality: Locality::default(), + messages: Vec::new(), + input_file: None, + params: serde_json::Value::Null, + hint_next: None, + request_id: generate_request_id(), + } + } +} + fn generate_request_id() -> String { Uuid::new_v4().to_string() } /// Response from an inference call. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct InferenceResponse { /// Text output (for chat, ASR, etc.). pub text: Option, @@ -404,8 +441,17 @@ pub struct InferenceResponse { pub duration_ms: u64, /// Request correlation ID. pub request_id: String, - /// Token usage, if reported by provider. + /// Total token usage, if reported by provider. pub tokens_used: Option, + /// Prompt/input tokens, if reported by provider. + #[serde(default)] + pub prompt_tokens: Option, + /// Completion/output tokens, if reported by provider. + #[serde(default)] + pub completion_tokens: Option, + /// Estimated cost in USD, if a price is configured for the serving provider. + #[serde(default)] + pub cost_usd: Option, /// Whether the response came from a fallback candidate. #[serde(default)] pub fallback_used: bool, @@ -440,8 +486,17 @@ pub enum StreamChunk { Completed { /// Wall-clock duration in milliseconds. duration_ms: u64, - /// Token usage, if reported. + /// Total token usage, if reported. tokens_used: Option, + /// Prompt/input tokens, if reported. + #[serde(default)] + prompt_tokens: Option, + /// Completion/output tokens, if reported. + #[serde(default)] + completion_tokens: Option, + /// Estimated cost in USD, if a price is configured for the provider. + #[serde(default)] + cost_usd: Option, /// File output path (for TTS, image gen, etc.). file: Option, /// Whether a fallback candidate served the request. From 39d52c28a93afa3815d1e662dd7414376faa2f42 Mon Sep 17 00:00:00 2001 From: aayank13 Date: Sun, 26 Jul 2026 02:27:53 +0530 Subject: [PATCH 07/22] feat: Enhance routing logic and introduce reasoning capabilities - Refactor router to implement new locality preferences (`Prefer::Local`, `Prefer::Cloud`, `Prefer::Auto`) for improved routing decisions. - Introduce reasoning tier matching bonus to prioritize models based on reasoning effort. - Update tests to cover new routing scenarios, including confidential data handling and reasoning tier selection. - Add new examples demonstrating the engine's capabilities in code review, document AI, and meeting brief scenarios. - Implement structured error handling for failed routing attempts, providing detailed failure chains for better debugging. - Extend message structure to support multimodal inputs with image references. --- Cargo.lock | 791 +++++++++++++++++- config.example.toml | 68 ++ docs/openapi.yaml | 60 +- mofa-engine-app/src/main.rs | 1 + mofa-engine-core/Cargo.toml | 3 + mofa-engine-core/src/backends.rs | 4 + mofa-engine-core/src/backends/liter_llm.rs | 719 ++++++++++++++++ mofa-engine-core/src/backends/local_asr.rs | 464 ++++++++++ mofa-engine-core/src/backends/local_tts.rs | 2 + mofa-engine-core/src/backends/ollama.rs | 6 +- .../src/backends/openai_compat.rs | 275 +++++- mofa-engine-core/src/config.rs | 21 +- mofa-engine-core/src/engine.rs | 187 ++++- mofa-engine-core/src/metrics.rs | 51 +- mofa-engine-core/src/router.rs | 186 +++- mofa-engine-sdk/examples/README.md | 30 + mofa-engine-sdk/examples/code_review.rs | 119 +++ mofa-engine-sdk/examples/doc_ai.rs | 62 ++ mofa-engine-sdk/examples/meeting_brief.rs | 127 +++ mofa-engine-sdk/src/server.rs | 4 +- mofa-kernel/src/error.rs | 115 ++- mofa-kernel/src/lib.rs | 2 +- mofa-kernel/src/traits.rs | 4 +- mofa-kernel/src/types.rs | 210 ++++- 24 files changed, 3376 insertions(+), 135 deletions(-) create mode 100644 mofa-engine-core/src/backends/liter_llm.rs create mode 100644 mofa-engine-core/src/backends/local_asr.rs create mode 100644 mofa-engine-sdk/examples/README.md create mode 100644 mofa-engine-sdk/examples/code_review.rs create mode 100644 mofa-engine-sdk/examples/doc_ai.rs create mode 100644 mofa-engine-sdk/examples/meeting_brief.rs diff --git a/Cargo.lock b/Cargo.lock index 1a44e43..1a5dd3d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,30 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "ahash" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "891477e0c6a8957309ee5c45a6368af3ae14bb510732d2684ffa19af310920f9" +dependencies = [ + "getrandom 0.2.17", + "once_cell", + "version_check", +] + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "version_check", + "zerocopy", +] + [[package]] name = "aho-corasick" version = "1.1.4" @@ -76,6 +100,21 @@ version = "1.0.102" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +[[package]] +name = "arc-swap" +version = "1.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c049c0be4daef0b145cb3555416b3b8ef5b7888a38aea1a3a155801fe7b0810b" +dependencies = [ + "rustversion", +] + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + [[package]] name = "async-trait" version = "0.1.89" @@ -84,7 +123,7 @@ checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -106,7 +145,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" dependencies = [ "axum-core", - "base64", + "base64 0.22.1", "bytes", "form_urlencoded", "futures-util", @@ -183,12 +222,30 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b25655df2c3cdd83c5e5b293b88acd880332b2ddadd7c30ac43144fdc0033da9" + [[package]] name = "bitflags" version = "2.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c4512299f36f043ab09a583e57bceb5a5aab7a73db1805848e8fef3c9e8c78b3" +[[package]] +name = "bitvec" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddcec3d12c579d40898fe0a9a358a803c23e9c52ca3c425707f81c9436211837" +dependencies = [ + "funty", + "radium", + "tap", + "wyz", +] + [[package]] name = "block-buffer" version = "0.10.4" @@ -198,12 +255,58 @@ dependencies = [ "generic-array", ] +[[package]] +name = "borsh" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a88b7ea17d208c4193f2c1e6de3c35fe71f98c96982d5ced308bdcc749ff6e1f" +dependencies = [ + "borsh-derive", + "bytes", + "cfg_aliases", +] + +[[package]] +name = "borsh-derive" +version = "1.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8f347189c62a579b8cd5f80714efa178f52e461dc2e6d701d264f5ff22e566c" +dependencies = [ + "once_cell", + "proc-macro-crate", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "bumpalo" version = "3.20.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" +[[package]] +name = "bytecheck" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23cdc57ce23ac53c931e88a43d06d070a6fd142f2617be5855eb75efc9beb1c2" +dependencies = [ + "bytecheck_derive", + "ptr_meta", + "simdutf8", +] + +[[package]] +name = "bytecheck_derive" +version = "0.6.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3db406d29fbcd95542e92559bed4d8ad92636d1ca8b3b72ede10b4bcc010e659" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "bytes" version = "1.11.1" @@ -226,6 +329,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" +[[package]] +name = "cfg_aliases" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f079e83a288787bcd14a6aea84cee5c87a67c5a3e660c30f557a3d24761b3527" + [[package]] name = "chrono" version = "0.4.44" @@ -271,7 +380,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -286,6 +395,16 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + [[package]] name = "core-foundation" version = "0.9.4" @@ -396,7 +515,7 @@ checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -472,6 +591,12 @@ dependencies = [ "percent-encoding", ] +[[package]] +name = "funty" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" + [[package]] name = "futures-channel" version = "0.3.32" @@ -487,6 +612,23 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-io" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4577ecaa3c4f96589d473f679a71b596316f6641bc350038b962a5daf0085d7a" + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "futures-sink" version = "0.3.32" @@ -506,12 +648,42 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-core", + "futures-io", + "futures-macro", "futures-sink", "futures-task", + "memchr", "pin-project-lite", "slab", ] +[[package]] +name = "gcp_auth" +version = "0.12.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d27dbcc645b60b8e7f6e2868a9d7102ece97d1bb49c1288b5321fcc67f7260" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "http", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "ring", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "thiserror", + "tokio", + "tracing", + "tracing-futures", + "url", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -577,6 +749,15 @@ dependencies = [ "tracing", ] +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" +dependencies = [ + "ahash 0.7.8", +] + [[package]] name = "hashbrown" version = "0.14.5" @@ -604,7 +785,7 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "headers-core", "http", @@ -705,6 +886,7 @@ dependencies = [ "hyper", "hyper-util", "rustls", + "rustls-native-certs", "tokio", "tokio-rustls", "tower-service", @@ -732,7 +914,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -914,6 +1096,55 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "jni" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5efd9a482cf3a427f00d6b35f14332adc7902ce91efb778580e180ff90fa3498" +dependencies = [ + "cfg-if", + "combine", + "jni-macros", + "jni-sys", + "log", + "simd_cesu8", + "thiserror", + "walkdir", + "windows-link", +] + +[[package]] +name = "jni-macros" +version = "0.22.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a00109accc170f0bdb141fed3e393c565b6f5e072365c3bd58f5b062591560a3" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "simd_cesu8", + "syn 2.0.117", +] + +[[package]] +name = "jni-sys" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6377a88cb3910bee9b0fa88d4f42e1d2da8e79915598f65fb0c7ee14c878af2" +dependencies = [ + "jni-sys-macros", +] + +[[package]] +name = "jni-sys-macros" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38c0b942f458fe50cdac086d2f946512305e5631e720728f2a61aabcd47a6264" +dependencies = [ + "quote", + "syn 2.0.117", +] + [[package]] name = "js-sys" version = "0.3.99" @@ -965,6 +1196,36 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "liter-llm" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "016fbd9e5fc66d4cb10eb5faa45242e54add93101816926fe9348bcfb0cb051a" +dependencies = [ + "ahash 0.8.12", + "arc-swap", + "async-trait", + "base64 0.22.1", + "bytes", + "dashmap", + "futures-core", + "gcp_auth", + "memchr", + "pin-project-lite", + "regex", + "reqwest 0.13.4", + "rust_decimal", + "rustls", + "secrecy", + "serde", + "serde_json", + "thiserror", + "tokio", + "tokio-util", + "toml 1.1.3+spec-1.1.0", + "url", +] + [[package]] name = "lock_api" version = "0.4.14" @@ -1050,18 +1311,21 @@ version = "0.1.0" dependencies = [ "anyhow", "async-trait", + "base64 0.23.0", "chrono", "dashmap", "dirs", + "liter-llm", "mofa-kernel", - "reqwest", + "reqwest 0.12.28", "serde", "serde_json", "sysinfo", "tempfile", "thiserror", "tokio", - "toml", + "tokio-stream", + "toml 0.8.23", "tracing", "uuid", ] @@ -1075,7 +1339,7 @@ dependencies = [ "chrono", "mofa-engine-core", "mofa-kernel", - "reqwest", + "reqwest 0.12.28", "serde", "serde_json", "thiserror", @@ -1186,7 +1450,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1242,6 +1506,26 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pin-project" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -1279,7 +1563,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn", + "syn 2.0.117", +] + +[[package]] +name = "proc-macro-crate" +version = "3.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f" +dependencies = [ + "toml_edit 0.25.13+spec-1.1.0", ] [[package]] @@ -1291,6 +1584,26 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "ptr_meta" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0738ccf7ea06b608c10564b31debd4f5bc5e197fc8bfe088f68ae5ce81e7a4f1" +dependencies = [ + "ptr_meta_derive", +] + +[[package]] +name = "ptr_meta_derive" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16b845dbfca988fa33db069c0e230574d15a3088f147a87b64c7589eb662c9ac" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + [[package]] name = "quote" version = "1.0.45" @@ -1312,14 +1625,41 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "radium" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc33ff2d4973d518d823d61aa239014831e521c75da58e3df4840d3f47749d09" + +[[package]] +name = "rand" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -1329,7 +1669,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -1361,6 +1710,18 @@ dependencies = [ "thiserror", ] +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + [[package]] name = "regex-automata" version = "0.4.14" @@ -1378,13 +1739,22 @@ version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +[[package]] +name = "rend" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71fe3824f5629716b1589be05dacd749f6aa084c87e00e016714a8cdfccc997c" +dependencies = [ + "bytecheck", +] + [[package]] name = "reqwest" version = "0.12.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-core", @@ -1420,6 +1790,51 @@ dependencies = [ "web-sys", ] +[[package]] +name = "reqwest" +version = "0.13.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "219c5811de6525e5416c7d5d53bb656d3afdbc6c5af816e0802bcfa42dbdc1c3" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-tls", + "hyper-util", + "js-sys", + "log", + "mime_guess", + "native-tls", + "percent-encoding", + "pin-project-lite", + "rustls", + "rustls-pki-types", + "rustls-platform-verifier", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-native-tls", + "tokio-rustls", + "tokio-util", + "tower", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams", + "web-sys", +] + [[package]] name = "ring" version = "0.17.14" @@ -1434,6 +1849,61 @@ dependencies = [ "windows-sys 0.52.0", ] +[[package]] +name = "rkyv" +version = "0.7.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2297bf9c81a3f0dc96bc9521370b88f054168c29826a75e89c55ff196e7ed6a1" +dependencies = [ + "bitvec", + "bytecheck", + "bytes", + "hashbrown 0.12.3", + "ptr_meta", + "rend", + "rkyv_derive", + "seahash", + "tinyvec", + "uuid", +] + +[[package]] +name = "rkyv_derive" +version = "0.7.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84d7b42d4b8d06048d3ac8db0eb31bcb942cbeb709f0b5f2b2ebde398d3038f5" +dependencies = [ + "proc-macro2", + "quote", + "syn 1.0.109", +] + +[[package]] +name = "rust_decimal" +version = "1.42.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2a24f50780bc85f09cc6ac299bdf1424302742d77221106859c9d8b102126a" +dependencies = [ + "arrayvec", + "borsh", + "bytes", + "num-traits", + "rand 0.8.7", + "rkyv", + "serde", + "serde_json", + "wasm-bindgen", +] + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + [[package]] name = "rustix" version = "1.1.4" @@ -1453,13 +1923,27 @@ version = "0.23.40" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" dependencies = [ + "log", "once_cell", + "ring", "rustls-pki-types", "rustls-webpki", "subtle", "zeroize", ] +[[package]] +name = "rustls-native-certs" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dab5152771c58876a2146916e53e35057e1a4dfa2b9df0f0305b07f611fdea4d" +dependencies = [ + "openssl-probe", + "rustls-pki-types", + "schannel", + "security-framework", +] + [[package]] name = "rustls-pki-types" version = "1.14.1" @@ -1469,6 +1953,33 @@ dependencies = [ "zeroize", ] +[[package]] +name = "rustls-platform-verifier" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d1e2536ce4f35f4846aa13bff16bd0ff40157cdb14cc056c7b14ba41233ba0" +dependencies = [ + "core-foundation 0.10.1", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls", + "rustls-native-certs", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + [[package]] name = "rustls-webpki" version = "0.103.13" @@ -1492,6 +2003,15 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + [[package]] name = "schannel" version = "0.1.29" @@ -1507,6 +2027,22 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "seahash" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c107b6f4780854c8b126e228ea8869f4d7b71260f962fefb57b996b8959ba6b" + +[[package]] +name = "secrecy" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e891af845473308773346dc847b2c23ee78fe442e0472ac50e22a18a93d3ae5a" +dependencies = [ + "serde", + "zeroize", +] + [[package]] name = "security-framework" version = "3.7.0" @@ -1563,7 +2099,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1599,6 +2135,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1647,6 +2192,22 @@ dependencies = [ "libc", ] +[[package]] +name = "simd_cesu8" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11031e251abf8611c80f460e19dbdeb54a66db918e49c65a7065b46ac7aec520" +dependencies = [ + "rustc_version", + "simdutf8", +] + +[[package]] +name = "simdutf8" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3a9fe34e3e7a50316060351f37187a3f546bce95496156754b601a5fa71b76e" + [[package]] name = "slab" version = "0.4.12" @@ -1687,6 +2248,17 @@ version = "2.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" +[[package]] +name = "syn" +version = "1.0.109" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + [[package]] name = "syn" version = "2.0.117" @@ -1715,7 +2287,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1752,6 +2324,12 @@ dependencies = [ "libc", ] +[[package]] +name = "tap" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" + [[package]] name = "tempfile" version = "3.27.0" @@ -1782,7 +2360,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1804,6 +2382,21 @@ dependencies = [ "zerovec", ] +[[package]] +name = "tinyvec" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + [[package]] name = "tokio" version = "1.52.3" @@ -1829,7 +2422,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1885,6 +2478,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -1896,9 +2490,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dc1beb996b9d83529a9e75c17a1686767d148d70663143c7854d8b4a09ced362" dependencies = [ "serde", - "serde_spanned", - "toml_datetime", - "toml_edit", + "serde_spanned 0.6.9", + "toml_datetime 0.6.11", + "toml_edit 0.22.27", +] + +[[package]] +name = "toml" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c96ecdfa941c8fc4fcaed14f99ada8ebed502eef533015095a07e3301d4c3c" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned 1.1.1", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "toml_writer", + "winnow 1.0.4", ] [[package]] @@ -1910,6 +2519,15 @@ dependencies = [ "serde", ] +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + [[package]] name = "toml_edit" version = "0.22.27" @@ -1918,10 +2536,31 @@ checksum = "41fe8c660ae4257887cf66394862d21dbca4a6ddd26f04a3560410406a2f819a" dependencies = [ "indexmap", "serde", - "serde_spanned", - "toml_datetime", + "serde_spanned 0.6.9", + "toml_datetime 0.6.11", "toml_write", - "winnow", + "winnow 0.7.15", +] + +[[package]] +name = "toml_edit" +version = "0.25.13+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6975367e4d2ef766d86af01ffad14b622fecc8d4357a998fbc4deb6e9bacaf9b" +dependencies = [ + "indexmap", + "toml_datetime 1.1.1+spec-1.1.0", + "toml_parser", + "winnow 1.0.4", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow 1.0.4", ] [[package]] @@ -1930,6 +2569,12 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" +[[package]] +name = "toml_writer" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" + [[package]] name = "tower" version = "0.5.3" @@ -1997,7 +2642,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2010,6 +2655,16 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "pin-project", + "tracing", +] + [[package]] name = "tracing-log" version = "0.2.0" @@ -2069,7 +2724,7 @@ dependencies = [ "http", "httparse", "log", - "rand", + "rand 0.9.4", "sha1", "thiserror", ] @@ -2157,6 +2812,16 @@ version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + [[package]] name = "want" version = "0.3.1" @@ -2199,6 +2864,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] @@ -2232,7 +2898,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wasm-bindgen-shared", ] @@ -2267,6 +2933,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" @@ -2289,6 +2968,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-root-certs" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b96554aa2acc8ccdb7e1c9a58a7a68dd5d13bccc69cd124cb09406db612a1c9b" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "winapi" version = "0.3.9" @@ -2305,6 +2993,15 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "winapi-x86_64-pc-windows-gnu" version = "0.4.0" @@ -2354,7 +3051,7 @@ checksum = "9107ddc059d5b6fbfbffdfa7a7fe3e22a226def0b2608f72e9d552763d3e1ad7" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2365,7 +3062,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2376,7 +3073,7 @@ checksum = "29bee4b38ea3cde66011baa44dba677c432a78593e202392d1e9070cf2a7fca7" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2387,7 +3084,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2525,6 +3222,15 @@ dependencies = [ "memchr", ] +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" +dependencies = [ + "memchr", +] + [[package]] name = "wit-bindgen" version = "0.51.0" @@ -2561,7 +3267,7 @@ dependencies = [ "heck", "indexmap", "prettyplease", - "syn", + "syn 2.0.117", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -2577,7 +3283,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -2625,6 +3331,15 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "wyz" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05f360fc0b24296329c78fda852a1e9ae82de9cf7b27dae4b7f62f118f77b9ed" +dependencies = [ + "tap", +] + [[package]] name = "yoke" version = "0.8.2" @@ -2644,7 +3359,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -2665,7 +3380,7 @@ checksum = "0b631b19d36a892ab55420c92dbc83ccd79274f25be714855d3074aa71cab639" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2685,7 +3400,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -2725,7 +3440,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] diff --git a/config.example.toml b/config.example.toml index f7de69b..0020bb6 100644 --- a/config.example.toml +++ b/config.example.toml @@ -73,6 +73,53 @@ price_output_per_1k = 0.01 name = "whisper-1" capability = "asr" +# Multi-vendor cloud gateway via the `liter-llm` crate (140+ vendors behind one +# OpenAI-style contract). Each model `name` is a liter-llm `provider/model` id, so +# one entry per vendor gives named + capability routing across them while MoFA keeps +# routing, failover, and dual-track cost/token observability. Leave `api_key` empty +# to let liter-llm read the vendor's env var (e.g. DEEPSEEK_API_KEY); set `base_url` +# only to override the vendor default (e.g. a private gateway). Supports chat +# (streaming) and image generation. +[[providers]] +name = "deepseek" +kind = "liter_llm" +base_url = "" # empty → liter-llm uses the vendor default endpoint +api_key = "${DEEPSEEK_API_KEY}" +priority = 10 +cost_tier = "low" +enabled = false # opt-in: set true once the vendor key is available +price_input_per_1k = 0.00027 +price_output_per_1k = 0.0011 + + # `reasoning_tier` lets `reasoning.effort` (low/medium/high) route to the + # matching model: a low-effort request prefers the cheap chat model, a + # high-effort one the stronger reasoner. + [[providers.models]] + name = "deepseek/deepseek-chat" + capability = "chat" + context_window = 65536 + reasoning_tier = "low" + + [[providers.models]] + name = "deepseek/deepseek-reasoner" + capability = "chat" + context_window = 65536 + reasoning_tier = "high" + +# A second liter-llm entry for image generation. One entry per vendor: this key +# authenticates its own models, and liter-llm formats each request per the model's +# `provider/` prefix (e.g. `openai/dall-e-3` → OpenAI images API). +# [[providers]] +# name = "openai-images" +# kind = "liter_llm" +# base_url = "" +# api_key = "${OPENAI_API_KEY}" +# priority = 10 +# cost_tier = "high" +# [[providers.models]] +# name = "openai/dall-e-3" +# capability = "image_gen" + # 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: @@ -93,3 +140,24 @@ enabled = false # opt-in: set true once the command is installed name = "kokoro" capability = "tts" memory_mb = 1024 + +# Local ASR process-adapter backend. Runs a local command (a FunASR or +# whisper.cpp-style CLI) once per transcription and returns the recognized text. +# Enables on-device transcription of long/confidential audio (`prefer=local`). +# Placeholders: +# {input} path to the source audio file +# {output} path the command must write the transcript to (omit to read stdout) +# {language} language hint from params.language (default `auto`) +[[providers]] +name = "local-asr" +kind = "local_asr" +base_url = "" # unused for local_asr, but the field is required +command = "funasr" +args = ["--input", "{input}", "--language", "{language}", "--output", "{output}"] +priority = 2 +enabled = false # opt-in: set true once the command is installed + + [[providers.models]] + name = "funasr" + capability = "asr" + memory_mb = 1024 diff --git a/docs/openapi.yaml b/docs/openapi.yaml index c6edc90..f503914 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -275,13 +275,42 @@ components: type: string enum: [capability_only, disabled, allow_named] default: capability_only - locality: + prefer: type: string - enum: [auto, prefer_local, local_only] + enum: [auto, local, cloud] default: auto description: > - Routing locality guardrail. `local_only` restricts the request to - local backends (privacy); `prefer_local` ranks local ahead of cloud. + Backend-locality preference / privacy guardrail. `local` is a hard + constraint restricting the request to local backends (confidential data + never leaves the device; fails rather than falling back to cloud); + `cloud` restricts to cloud backends; `auto` (default) ranks local ahead + of cloud while keeping cloud as a fallback. Accepts the legacy field + name `locality` and value `local_only` as aliases. + data_class: + type: string + enum: [public, internal, confidential] + default: public + description: > + Data-sensitivity class. `confidential` is a hard data-residency + constraint — the request is pinned to local backends regardless of + `prefer` (and overrides a conflicting `prefer=cloud`), failing rather + than sending sensitive data to the cloud. Every request's effective + locality is written to the `mofa::audit` log for traceability. + reasoning: + type: object + description: > + Deep-thinking controls for reasoning-capable models. Omit for standard + behavior. `effort` also drives inference-tier routing; when `include` + is true the thought chain is streamed as `reasoning` chunks (distinct + from answer `text`) rather than stripped. + properties: + effort: + type: string + enum: [low, medium, high] + default: medium + include: + type: boolean + default: false messages: type: array items: @@ -303,6 +332,13 @@ components: type: string content: type: string + images: + type: array + description: > + Image references for multimodal (VLM) requests — HTTP(S) URLs, `data:` + URLs, or local file paths. Omitted for text-only messages. + items: + type: string InferenceResponse: type: object properties: @@ -395,6 +431,22 @@ components: type: boolean source: type: string + failed_chain: + type: array + description: > + Complete per-candidate failure chain (present when failover was + attempted and every candidate failed). Lets an Agent decide whether to + retry or switch tiers rather than seeing only the last error. + items: + type: object + required: [provider, model, reason] + properties: + provider: { type: string } + model: { type: string } + reason: { type: string } + routing_reason: + type: string + description: Machine-readable routing reason for the primary candidate. AllocationSnapshot: type: object properties: diff --git a/mofa-engine-app/src/main.rs b/mofa-engine-app/src/main.rs index 63fcc15..6721fbc 100644 --- a/mofa-engine-app/src/main.rs +++ b/mofa-engine-app/src/main.rs @@ -132,6 +132,7 @@ async fn main() -> anyhow::Result<()> { messages: vec![Message { role: "user".into(), content: text, + ..Default::default() }], input_file: None, params: serde_json::Value::Null, diff --git a/mofa-engine-core/Cargo.toml b/mofa-engine-core/Cargo.toml index 0dcc507..ee1a1ab 100644 --- a/mofa-engine-core/Cargo.toml +++ b/mofa-engine-core/Cargo.toml @@ -22,3 +22,6 @@ sysinfo = { workspace = true } dashmap = { workspace = true } dirs = { workspace = true } tempfile = { workspace = true } +tokio-stream = { workspace = true } +liter-llm = "1.11.0" +base64 = "0.23.0" diff --git a/mofa-engine-core/src/backends.rs b/mofa-engine-core/src/backends.rs index 7783c49..4af3f83 100644 --- a/mofa-engine-core/src/backends.rs +++ b/mofa-engine-core/src/backends.rs @@ -1,9 +1,13 @@ //! Provider backend implementations. +pub mod liter_llm; +pub mod local_asr; pub mod local_tts; pub mod ollama; pub mod openai_compat; +pub use liter_llm::LiterLLMProvider; +pub use local_asr::LocalAsrProvider; pub use local_tts::LocalTtsProvider; pub use ollama::OllamaProvider; pub use openai_compat::OpenAiCompatProvider; diff --git a/mofa-engine-core/src/backends/liter_llm.rs b/mofa-engine-core/src/backends/liter_llm.rs new file mode 100644 index 0000000..0dd92ae --- /dev/null +++ b/mofa-engine-core/src/backends/liter_llm.rs @@ -0,0 +1,719 @@ +//! Multi-vendor cloud provider backed by the `liter-llm` crate. +//! +//! `liter-llm` exposes 140+ LLM vendors behind one OpenAI-style contract and +//! handles per-vendor request/response adaptation, auth, and the `provider/model` +//! routing convention. MoFA keeps ownership of routing, failover, memory, and +//! dual-track observability; this adapter just bridges the [`Provider`] trait to +//! liter-llm's [`LlmClient`]. +//! +//! A configured model `name` is the liter-llm model id (e.g. `openai/gpt-4o`, +//! `deepseek/deepseek-chat`). MoFA's canonical id becomes `/`, and +//! [`model_id_name`] recovers the `provider/model` string liter-llm expects. + +use async_trait::async_trait; +use base64::Engine as _; +use liter_llm::{ + AssistantContent, AssistantMessage, ChatCompletionRequest, ClientConfigBuilder, ContentPart, + CreateImageRequest, DefaultClient, ImageDetail, LlmClient, Message as LiterMessage, + SystemMessage, UserContent, UserMessage, +}; +use mofa_kernel::{ + BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelResidency, Provider, + ProviderKind, Reasoning, ReasoningEffort, StreamDelta, StreamSink, canonical_model_id, + model_id_name, +}; +use std::time::Instant; +use tokio_stream::StreamExt as _; + +use crate::config::ModelDef; + +/// A provider that reaches many cloud vendors through the `liter-llm` gateway. +pub struct LiterLLMProvider { + /// Display name. + name: String, + /// Configured models (each `name` is a liter-llm `provider/model` id). + models: Vec, + /// Cost tier for all models from this provider. + cost_tier: CostTier, + /// Directory for generated image artifacts. + output_dir: std::path::PathBuf, + /// The liter-llm client. `None` when no usable credentials were resolved at + /// construction, in which case the provider reports itself unavailable rather + /// than failing engine startup (mirrors the OpenAI-compatible backend). + client: Option, +} + +impl LiterLLMProvider { + /// Build a liter-llm provider. + /// + /// When `api_key` is empty, liter-llm resolves the vendor's environment + /// variable (e.g. `OPENAI_API_KEY`) from the first model's `provider/` prefix. + /// A non-empty `base_url` overrides the vendor default (e.g. a private gateway). + pub fn new( + name: impl Into, + api_key: Option, + base_url: impl Into, + models: Vec, + cost_tier: CostTier, + output_dir: Option, + ) -> Result { + let name = name.into(); + let key = api_key.unwrap_or_default(); + let has_key = !key.is_empty(); + + let mut builder = ClientConfigBuilder::new(key).load_env(!has_key); + let base_url = base_url.into(); + if !base_url.is_empty() { + builder = builder.base_url(base_url); + } + let config = builder.build(); + + // The model hint's `provider/` prefix selects the default vendor (and its + // env var when `load_env` is on). Per-request model prefixes still route. + let model_hint = models.first().map(|m| m.name.clone()); + let client = DefaultClient::new(config, model_hint.as_deref()).ok(); + if client.is_none() { + tracing::warn!( + provider = %name, + "liter-llm provider has no usable credentials; marked unavailable" + ); + } + + Ok(Self { + name, + 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, + }) + } + + fn client(&self) -> Result<&DefaultClient, EngineError> { + self.client + .as_ref() + .ok_or_else(|| EngineError::ProviderError { + provider: self.name.clone(), + detail: "no liter-llm credentials configured".into(), + }) + } + + fn provider_error(&self, e: impl std::fmt::Display) -> EngineError { + EngineError::ProviderError { + provider: self.name.clone(), + detail: e.to_string(), + } + } + + /// Map MoFA conversation messages to liter-llm's typed text `Message`s. + fn to_liter_messages(request: &InferenceRequest) -> Vec { + request + .messages + .iter() + .map(|m| text_message(&m.role, m.content.clone())) + .collect() + } + + /// Build a chat request from MoFA text messages, applying `max_tokens` and + /// the reasoning effort tier. Errors if there are no messages. + fn build_chat_request( + &self, + model_name: &str, + request: &InferenceRequest, + ) -> Result { + let messages = Self::to_liter_messages(request); + if messages.is_empty() { + return Err(EngineError::InvalidRequest("no messages provided".into())); + } + Ok(ChatCompletionRequest { + model: model_name.to_string(), + messages, + max_tokens: request.params.get("max_tokens").and_then(|v| v.as_u64()), + reasoning_effort: to_liter_effort(request.reasoning), + ..Default::default() + }) + } + + async fn invoke_chat( + &self, + model_name: &str, + request: &InferenceRequest, + start: Instant, + ) -> Result { + let req = self.build_chat_request(model_name, request)?; + + let resp = self + .client()? + .chat(req) + .await + .map_err(|e| self.provider_error(e))?; + + let text = resp.choices.first().and_then(|c| c.message.text()); + let (prompt_tokens, completion_tokens, tokens_used) = split_usage(resp.usage); + + Ok(InferenceResponse { + text, + file: None, + 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, + prompt_tokens, + completion_tokens, + fallback_used: false, + routing_reason: None, + ..Default::default() + }) + } + + async fn invoke_image_gen( + &self, + model_name: &str, + request: &InferenceRequest, + start: Instant, + ) -> Result { + let prompt = request + .messages + .first() + .map(|m| m.content.clone()) + .or_else(|| { + request + .params + .get("prompt") + .and_then(|v| v.as_str()) + .map(String::from) + }) + .filter(|p| !p.is_empty()) + .ok_or_else(|| { + EngineError::InvalidRequest("image generation requires a prompt".into()) + })?; + + let size = request + .params + .get("size") + .and_then(|v| v.as_str()) + .map(String::from); + + let req = CreateImageRequest { + prompt, + model: Some(model_name.to_string()), + n: Some(1), + size, + // Prefer base64 so we can persist a managed local artifact; vendors that + // ignore this still return a URL, which we surface as-is. + response_format: Some("b64_json".into()), + ..Default::default() + }; + + let resp = self + .client()? + .image_generate(req) + .await + .map_err(|e| self.provider_error(e))?; + + let image = resp + .data + .into_iter() + .next() + .ok_or_else(|| EngineError::ProviderError { + provider: self.name.clone(), + detail: "image generation returned no data".into(), + })?; + + let (text, file) = if let Some(b64) = image.b64_json { + let bytes = base64::engine::general_purpose::STANDARD + .decode(b64.as_bytes()) + .map_err(|e| self.provider_error(format!("image base64 decode error: {e}")))?; + let path = self + .output_dir + .join(format!("mofa_image_{}.png", uuid::Uuid::new_v4())); + tokio::fs::write(&path, &bytes) + .await + .map_err(|e| EngineError::Internal(format!("image write error: {e}")))?; + (None, Some(path.to_string_lossy().to_string())) + } else { + (image.url, None) + }; + + Ok(InferenceResponse { + text, + file, + 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, + ..Default::default() + }) + } + + /// Vision-language understanding (S3): send the conversation as multimodal + /// chat, attaching each message's `images` as image parts with a `detail` + /// tier (`low | high | auto`, from `params.detail`) that passes through to the + /// vendor's image billing. + async fn invoke_vlm( + &self, + model_name: &str, + request: &InferenceRequest, + start: Instant, + ) -> Result { + let detail = match request.params.get("detail").and_then(|v| v.as_str()) { + Some("low") => ImageDetail::Low, + Some("high") => ImageDetail::High, + _ => ImageDetail::Auto, + }; + + let mut messages = Vec::with_capacity(request.messages.len()); + for m in &request.messages { + if m.images.is_empty() { + messages.push(text_message(&m.role, m.content.clone())); + continue; + } + // A multimodal turn: text first, then each image as a detail-tagged part. + let mut parts = Vec::with_capacity(m.images.len() + 1); + if !m.content.is_empty() { + parts.push(ContentPart::text(m.content.clone())); + } + for image in &m.images { + parts.push(ContentPart::image_with_detail( + image_to_url(image).await?, + detail.clone(), + )); + } + messages.push(LiterMessage::user_with_parts(parts)); + } + if messages.is_empty() { + return Err(EngineError::InvalidRequest("no messages provided".into())); + } + + let req = ChatCompletionRequest { + model: model_name.to_string(), + messages, + max_tokens: request.params.get("max_tokens").and_then(|v| v.as_u64()), + reasoning_effort: to_liter_effort(request.reasoning), + ..Default::default() + }; + + let resp = self + .client()? + .chat(req) + .await + .map_err(|e| self.provider_error(e))?; + + let text = resp.choices.first().and_then(|c| c.message.text()); + let (prompt_tokens, completion_tokens, tokens_used) = split_usage(resp.usage); + + Ok(InferenceResponse { + text, + file: None, + 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, + prompt_tokens, + completion_tokens, + fallback_used: false, + routing_reason: None, + ..Default::default() + }) + } +} + +/// Map a MoFA reasoning request to liter-llm's effort enum. +fn to_liter_effort(reasoning: Option) -> Option { + reasoning.map(|r| match r.effort { + ReasoningEffort::Low => liter_llm::ReasoningEffort::Low, + ReasoningEffort::Medium => liter_llm::ReasoningEffort::Medium, + ReasoningEffort::High => liter_llm::ReasoningEffort::High, + }) +} + +/// Build a plain-text liter-llm message for the given role. +fn text_message(role: &str, content: String) -> LiterMessage { + match role { + "system" | "developer" => LiterMessage::System(SystemMessage { + content: UserContent::Text(content), + ..Default::default() + }), + "assistant" => LiterMessage::Assistant(AssistantMessage { + content: Some(AssistantContent::Text(content)), + ..Default::default() + }), + _ => LiterMessage::User(UserMessage { + content: UserContent::Text(content), + ..Default::default() + }), + } +} + +/// Resolve an image reference to a URL liter-llm accepts. HTTP(S) and `data:` +/// URLs pass through; a local file path is read and encoded as a `data:` URL. +async fn image_to_url(image: &str) -> Result { + if image.starts_with("http://") || image.starts_with("https://") || image.starts_with("data:") { + return Ok(image.to_string()); + } + let bytes = tokio::fs::read(image) + .await + .map_err(|e| EngineError::InvalidRequest(format!("cannot read image '{image}': {e}")))?; + let mime = match std::path::Path::new(image) + .extension() + .and_then(|e| e.to_str()) + .map(str::to_ascii_lowercase) + .as_deref() + { + Some("png") => "image/png", + Some("webp") => "image/webp", + Some("gif") => "image/gif", + _ => "image/jpeg", + }; + let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); + Ok(format!("data:{mime};base64,{b64}")) +} + +/// Split a liter-llm `Usage` into `(prompt, completion, total)` as `Option`, +/// falling back to `prompt + completion` when the vendor omits a total. +fn split_usage(usage: Option) -> (Option, Option, Option) { + match usage { + None => (None, None, None), + Some(u) => { + let prompt = u.prompt_tokens as u32; + let completion = u.completion_tokens as u32; + let total = if u.total_tokens > 0 { + u.total_tokens as u32 + } else { + prompt + completion + }; + (Some(prompt), Some(completion), Some(total)) + } + } +} + +#[async_trait] +impl Provider for LiterLLMProvider { + fn name(&self) -> &str { + &self.name + } + + fn kind(&self) -> ProviderKind { + ProviderKind::LiterLlm + } + + fn features(&self) -> Vec { + vec![BackendFeature::Discovery] + } + + 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, self.cost_tier); + card.id = canonical_model_id(&self.name, &m.name); + card.availability = ModelAvailability::Configured; + card.residency = ModelResidency::Remote; + card.context_window = m.context_window.unwrap_or(4096); + card.memory_estimate_bytes = m.memory_mb.unwrap_or(0) * 1024 * 1024; + card.execution.max_concurrency = 32; + card.reasoning_tier = m + .reasoning_tier + .as_deref() + .and_then(ReasoningEffort::from_str_loose); + card.refresh_status(); + Some(card) + }) + .collect(); + Ok(cards) + } + + async fn health(&self) -> Result { + let Some(client) = &self.client else { + return Ok(BackendHealth::Unavailable); + }; + // A flaky/unsupported model-list endpoint should not kill routing, so a + // reachable-but-erroring vendor stays Degraded (still routable) rather than + // Unavailable; invoke-time failures are what trip the circuit breaker. + match client.list_models().await { + Ok(_) => Ok(BackendHealth::Healthy), + Err(_) => Ok(BackendHealth::Degraded), + } + } + + async fn load(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: canonical_model_id(&self.name, model_id_name(model_id)), + residency: ModelResidency::Remote, + memory_bytes: Some(0), + changed: false, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: canonical_model_id(&self.name, model_id_name(model_id)), + residency: ModelResidency::Remote, + memory_bytes: Some(0), + changed: false, + }) + } + + 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::Chat); + let start = Instant::now(); + + let supports = self.models.iter().any(|m| { + m.name == model_name && Capability::from_str_loose(&m.capability) == Some(capability) + }); + if !supports { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' model '{}' does not support {capability}", + self.name, model_name + ))); + } + + match capability { + Capability::Chat => self.invoke_chat(model_name, request, start).await, + Capability::Vlm => self.invoke_vlm(model_name, request, start).await, + Capability::ImageGen => self.invoke_image_gen(model_name, request, start).await, + other => Err(EngineError::UnsupportedOperation(format!( + "provider '{}' does not support {other} via liter-llm", + self.name + ))), + } + } + + /// Real per-token streaming via liter-llm's `chat_stream`, forwarding each + /// `delta.content` through `sink`. Non-chat capabilities fall back to a + /// single-shot emit through [`invoke`](Self::invoke). + async fn stream( + &self, + model_id: &str, + request: &InferenceRequest, + sink: StreamSink, + ) -> Result { + let capability = request.capability.unwrap_or(Capability::Chat); + if capability != Capability::Chat { + let response = self.invoke(model_id, request).await?; + if let Some(text) = &response.text + && !text.is_empty() + { + let _ = sink.send(StreamDelta::Text(text.clone())).await; + } + return Ok(response); + } + + let model_name = model_id_name(model_id); + let req = self.build_chat_request(model_name, request)?; + + let start = Instant::now(); + let mut stream = self + .client()? + .chat_stream(req) + .await + .map_err(|e| self.provider_error(e))?; + + let mut full = String::new(); + let mut usage = None; + while let Some(item) = stream.next().await { + let chunk = item.map_err(|e| self.provider_error(e))?; + if chunk.usage.is_some() { + usage = chunk.usage; + } + let Some(choice) = chunk.choices.into_iter().next() else { + continue; + }; + // Thought-chain deltas are forwarded distinctly from answer text so + // callers can display/audit the reasoning separately (S2). + if let Some(rc) = choice.delta.reasoning_content + && !rc.is_empty() + { + let _ = sink.send(StreamDelta::Reasoning(rc)).await; + } + if let Some(content) = choice.delta.content + && !content.is_empty() + { + full.push_str(&content); + let _ = sink.send(StreamDelta::Text(content)).await; + } + } + + let (prompt_tokens, completion_tokens, tokens_used) = split_usage(usage); + Ok(InferenceResponse { + text: Some(full), + file: None, + 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, + prompt_tokens, + completion_tokens, + fallback_used: false, + routing_reason: None, + ..Default::default() + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn model(name: &str, cap: &str) -> ModelDef { + ModelDef { + name: name.into(), + capability: cap.into(), + context_window: None, + memory_mb: None, + ..Default::default() + } + } + + #[test] + fn kind_and_locality() { + let p = LiterLLMProvider::new( + "liter", + Some("sk-test".into()), + "", + vec![model("openai/gpt-4o", "chat")], + CostTier::Medium, + None, + ) + .unwrap(); + assert_eq!(p.kind(), ProviderKind::LiterLlm); + // Cloud gateway: never treated as on-device. + assert!(!p.kind().is_local()); + } + + #[tokio::test] + async fn discover_maps_provider_prefixed_models() { + let p = LiterLLMProvider::new( + "liter", + Some("sk-test".into()), + "", + vec![ + model("openai/gpt-4o", "chat"), + model("openai/dall-e-3", "image_gen"), + ], + CostTier::Medium, + None, + ) + .unwrap(); + let cards = p.discover().await.unwrap(); + assert_eq!(cards.len(), 2); + // Canonical id is `/`, and model_id_name recovers the + // `provider/model` string liter-llm routes on. + assert_eq!(cards[0].id, "liter/openai/gpt-4o"); + assert_eq!(model_id_name(&cards[0].id), "openai/gpt-4o"); + assert_eq!(cards[0].capability, Capability::Chat); + assert_eq!(cards[1].capability, Capability::ImageGen); + assert_eq!(cards[0].residency, ModelResidency::Remote); + } + + #[tokio::test] + async fn missing_credentials_report_unavailable() { + // Empty key + a vendor whose env var is unset → no client → Unavailable, + // and startup still succeeds (no hard failure). + let p = LiterLLMProvider::new( + "liter", + None, + "", + vec![model("openai/gpt-4o", "chat")], + CostTier::Medium, + None, + ) + .unwrap(); + // In CI OPENAI_API_KEY is typically unset; if it is set, the client builds + // and health is not Unavailable — accept either, but never panic. + let health = p.health().await.unwrap(); + assert!(matches!( + health, + BackendHealth::Unavailable | BackendHealth::Healthy | BackendHealth::Degraded + )); + } + + #[test] + fn usage_split_falls_back_to_sum() { + let (p, c, t) = split_usage(Some(liter_llm::Usage { + prompt_tokens: 5, + completion_tokens: 3, + total_tokens: 0, + ..Default::default() + })); + assert_eq!((p, c, t), (Some(5), Some(3), Some(8))); + assert_eq!(split_usage(None), (None, None, None)); + } + + #[test] + fn effort_maps_to_liter_enum() { + use liter_llm::ReasoningEffort as L; + let mk = |e| { + Some(Reasoning { + effort: e, + include: false, + }) + }; + assert!(matches!( + to_liter_effort(mk(ReasoningEffort::Low)), + Some(L::Low) + )); + assert!(matches!( + to_liter_effort(mk(ReasoningEffort::Medium)), + Some(L::Medium) + )); + assert!(matches!( + to_liter_effort(mk(ReasoningEffort::High)), + Some(L::High) + )); + assert!(to_liter_effort(None).is_none()); + } + + #[tokio::test] + async fn image_url_and_data_url_pass_through() { + assert_eq!( + image_to_url("https://example.com/a.png").await.unwrap(), + "https://example.com/a.png" + ); + assert_eq!( + image_to_url("data:image/png;base64,AAAA").await.unwrap(), + "data:image/png;base64,AAAA" + ); + } + + #[tokio::test] + async fn local_image_file_is_encoded_as_data_url() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("pic.png"); + tokio::fs::write(&path, b"\x89PNG\r\n").await.unwrap(); + let url = image_to_url(path.to_str().unwrap()).await.unwrap(); + assert!(url.starts_with("data:image/png;base64,")); + } + + #[tokio::test] + async fn discover_includes_vlm_capability() { + let p = LiterLLMProvider::new( + "liter", + Some("sk-test".into()), + "", + vec![model("openai/gpt-4o", "vlm")], + CostTier::Medium, + None, + ) + .unwrap(); + let cards = p.discover().await.unwrap(); + assert_eq!(cards[0].capability, Capability::Vlm); + } +} diff --git a/mofa-engine-core/src/backends/local_asr.rs b/mofa-engine-core/src/backends/local_asr.rs new file mode 100644 index 0000000..7da8568 --- /dev/null +++ b/mofa-engine-core/src/backends/local_asr.rs @@ -0,0 +1,464 @@ +//! Local ASR process-adapter backend. +//! +//! Runs a configured local command — a FunASR or whisper.cpp-style CLI — to +//! transcribe an audio file, returning the recognized text. This is the S1 +//! Meeting Brief enabler: long, confidential audio can be transcribed on-device +//! (`prefer=local`) without leaving the machine. +//! +//! Runtime- and device-specific concerns stay behind this `Provider` boundary, +//! so the engine treats a local ASR model like any other backend: it is +//! discovered, admitted through memory reservation, warmed, idle-unloaded, and +//! can fail over to a cloud ASR backend (non-confidential only) when the local +//! one is unavailable. +//! +//! ## Command contract +//! +//! The command is spawned once per transcription (a cold, stateless process). +//! Argument templates may contain: +//! - `{input}` — path to the source audio file (required), +//! - `{output}` — path the command must write the transcript text to; when +//! present the transcript is read back from this file, otherwise the +//! command's stdout is taken as the transcript, +//! - `{language}` — language hint from `params.language` (default `auto`). +//! +//! Child processes are spawned with `kill_on_drop(true)`, so an inference +//! timeout that drops the future terminates the process rather than leaking it. + +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 ASR command. +pub struct LocalAsrProvider { + /// Display name. + name: String, + /// Program to execute per transcription. + command: String, + /// Argument template with `{input}`, `{output}`, and `{language}` placeholders. + args: Vec, + /// Directory for transient transcript files. + output_dir: PathBuf, + /// Configured models this backend serves. + models: Vec, +} + +impl LocalAsrProvider { + /// Create a new local ASR process adapter. + pub fn new( + name: impl Into, + command: impl Into, + args: Vec, + output_dir: Option, + models: Vec, + ) -> Self { + Self { + name: name.into(), + command: command.into(), + args, + 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 transcribe( + &self, + model_name: &str, + request: &InferenceRequest, + start: std::time::Instant, + ) -> Result { + let input_file = request + .input_file + .as_deref() + .filter(|p| !p.is_empty()) + .ok_or_else(|| EngineError::InvalidRequest("ASR requires input_file".into()))?; + + let language = request + .params + .get("language") + .and_then(|v| v.as_str()) + .unwrap_or("auto"); + + // Materialize a transcript output path only if the template references it; + // otherwise the command's stdout is the transcript. + let uses_output = self.args.iter().any(|a| a.contains("{output}")); + let output_file = uses_output.then(|| { + self.output_dir + .join(format!("mofa_asr_{}.txt", uuid::Uuid::new_v4())) + }); + let output_str = output_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("{input}", input_file) + .replace("{output}", &output_str) + .replace("{language}", language) + }) + .collect(); + + let result = Command::new(&self.command) + .args(&resolved_args) + .kill_on_drop(true) + .output() + .await; + + let output = match result { + Ok(output) => output, + Err(e) => { + if let Some(path) = &output_file { + let _ = tokio::fs::remove_file(path).await; + } + return Err(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); + if let Some(path) = &output_file { + let _ = tokio::fs::remove_file(path).await; + } + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: format!( + "ASR command exited with {}: {}", + output.status, + stderr.trim() + ), + }); + } + + // Read the transcript from the output file (if templated) or from stdout. + let transcript = match &output_file { + Some(path) => { + let text = tokio::fs::read_to_string(path).await; + let _ = tokio::fs::remove_file(path).await; + text.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("cannot read ASR transcript: {e}"), + })? + } + None => String::from_utf8_lossy(&output.stdout).into_owned(), + }; + + let transcript = transcript.trim().to_string(); + if transcript.is_empty() { + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "ASR command produced no transcript".into(), + }); + } + + Ok(InferenceResponse { + text: Some(transcript), + file: None, + 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, + ..Default::default() + }) + } +} + +#[async_trait] +impl Provider for LocalAsrProvider { + fn name(&self) -> &str { + &self.name + } + + fn kind(&self) -> ProviderKind { + ProviderKind::LocalAsr + } + + 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!("ASR 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::Asr); + + if capability != Capability::Asr { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' only supports asr, not {capability}", + self.name + ))); + } + if !self.model_supports(model_name, Capability::Asr) { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' model '{model_name}' does not support asr", + self.name + ))); + } + + self.transcribe(model_name, request, std::time::Instant::now()) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn asr_models() -> Vec { + vec![ModelDef { + name: "fixture".into(), + capability: "asr".into(), + context_window: None, + memory_mb: Some(256), + ..Default::default() + }] + } + + fn asr_request(input_file: Option<&str>) -> InferenceRequest { + InferenceRequest { + capability: Some(Capability::Asr), + input_file: input_file.map(String::from), + request_id: "test".into(), + ..Default::default() + } + } + + #[test] + fn metadata_and_kind() { + let p = LocalAsrProvider::new("local-asr", "sh", vec![], None, asr_models()); + assert_eq!(p.kind(), ProviderKind::LocalAsr); + assert!(p.kind().is_local()); + assert_eq!(p.name(), "local-asr"); + } + + #[tokio::test] + async fn discover_reports_configured_asr_model() { + let p = LocalAsrProvider::new("local-asr", "sh", vec![], None, asr_models()); + let cards = p.discover().await.unwrap(); + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].id, "local-asr/fixture"); + assert_eq!(cards[0].capability, Capability::Asr); + assert_eq!(cards[0].memory_estimate_bytes, 256 * 1024 * 1024); + } + + #[tokio::test] + async fn transcribes_from_stdout() { + let dir = tempfile::tempdir().unwrap(); + let audio = dir.path().join("clip.wav"); + std::fs::write(&audio, b"fake-audio").unwrap(); + // A deterministic "ASR" fixture: echo a transcript to stdout. + let p = LocalAsrProvider::new( + "local-asr", + "sh", + vec!["-c".into(), "echo 'hello from audio'".into()], + Some(dir.path().to_string_lossy().to_string()), + asr_models(), + ); + + let resp = p + .invoke( + "local-asr/fixture", + &asr_request(Some(audio.to_str().unwrap())), + ) + .await + .unwrap(); + assert_eq!(resp.text.as_deref(), Some("hello from audio")); + } + + #[tokio::test] + async fn transcribes_from_output_file_with_input_and_language() { + let dir = tempfile::tempdir().unwrap(); + let audio = dir.path().join("clip.wav"); + std::fs::write(&audio, b"fake-audio").unwrap(); + // Fixture writes ":" to the output file so we can + // assert the placeholders were substituted and the file was read back. + let p = LocalAsrProvider::new( + "local-asr", + "sh", + vec![ + "-c".into(), + "printf '%s:%s' \"$1\" \"$(basename \"$2\")\" > \"$3\"".into(), + "sh".into(), + "{language}".into(), + "{input}".into(), + "{output}".into(), + ], + Some(dir.path().to_string_lossy().to_string()), + asr_models(), + ); + + let mut req = asr_request(Some(audio.to_str().unwrap())); + req.params = serde_json::json!({ "language": "en" }); + let resp = p.invoke("local-asr/fixture", &req).await.unwrap(); + assert_eq!(resp.text.as_deref(), Some("en:clip.wav")); + } + + #[tokio::test] + async fn missing_input_file_is_invalid_request() { + let p = LocalAsrProvider::new("local-asr", "sh", vec![], None, asr_models()); + let err = p + .invoke("local-asr/fixture", &asr_request(None)) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + } + + #[tokio::test] + async fn command_failure_is_a_retryable_provider_error() { + let dir = tempfile::tempdir().unwrap(); + let audio = dir.path().join("clip.wav"); + std::fs::write(&audio, b"x").unwrap(); + let p = LocalAsrProvider::new( + "local-asr", + "sh", + vec!["-c".into(), "echo boom >&2; exit 4".into()], + Some(dir.path().to_string_lossy().to_string()), + asr_models(), + ); + let err = p + .invoke( + "local-asr/fixture", + &asr_request(Some(audio.to_str().unwrap())), + ) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + assert!(err.retryable()); + } + + #[tokio::test] + async fn empty_transcript_is_a_provider_error() { + let dir = tempfile::tempdir().unwrap(); + let audio = dir.path().join("clip.wav"); + std::fs::write(&audio, b"x").unwrap(); + let p = LocalAsrProvider::new( + "local-asr", + "sh", + vec!["-c".into(), "true".into()], // no output + Some(dir.path().to_string_lossy().to_string()), + asr_models(), + ); + let err = p + .invoke( + "local-asr/fixture", + &asr_request(Some(audio.to_str().unwrap())), + ) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + } + + #[tokio::test] + async fn rejects_unsupported_capability() { + let p = LocalAsrProvider::new("local-asr", "sh", vec![], None, asr_models()); + let mut req = asr_request(Some("/tmp/x.wav")); + req.capability = Some(Capability::Chat); + let err = p.invoke("local-asr/fixture", &req).await.unwrap_err(); + assert!(matches!(err, EngineError::UnsupportedOperation(_))); + } +} diff --git a/mofa-engine-core/src/backends/local_tts.rs b/mofa-engine-core/src/backends/local_tts.rs index 53db6b4..d225c8c 100644 --- a/mofa-engine-core/src/backends/local_tts.rs +++ b/mofa-engine-core/src/backends/local_tts.rs @@ -325,6 +325,7 @@ mod tests { capability: "tts".into(), context_window: None, memory_mb: Some(64), + ..Default::default() }] } @@ -338,6 +339,7 @@ mod tests { messages: vec![Message { role: "user".into(), content: text.into(), + ..Default::default() }], input_file: None, params: serde_json::Value::Null, diff --git a/mofa-engine-core/src/backends/ollama.rs b/mofa-engine-core/src/backends/ollama.rs index fd19ef2..406e7d9 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -6,7 +6,7 @@ use async_trait::async_trait; use mofa_kernel::{ BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelResidency, Provider, - ProviderKind, StreamSink, canonical_model_id, model_id_name, + ProviderKind, StreamDelta, StreamSink, canonical_model_id, model_id_name, }; use reqwest::Client; use serde::{Deserialize, Serialize}; @@ -475,7 +475,7 @@ impl Provider for OllamaProvider { if let Some(chunk) = parse_stream_line(&line) && let Some(delta) = apply(chunk, &mut full) { - let _ = sink.send(delta).await; + let _ = sink.send(StreamDelta::Text(delta)).await; } } } @@ -483,7 +483,7 @@ impl Provider for OllamaProvider { if let Some(chunk) = parse_stream_line(&buf) && let Some(delta) = apply(chunk, &mut full) { - let _ = sink.send(delta).await; + let _ = sink.send(StreamDelta::Text(delta)).await; } let tokens_used = match (prompt_tokens, completion_tokens) { diff --git a/mofa-engine-core/src/backends/openai_compat.rs b/mofa-engine-core/src/backends/openai_compat.rs index cc6fcef..48cbe9d 100644 --- a/mofa-engine-core/src/backends/openai_compat.rs +++ b/mofa-engine-core/src/backends/openai_compat.rs @@ -6,7 +6,7 @@ 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, + ProviderKind, ReasoningEffort, StreamDelta, StreamSink, canonical_model_id, model_id_name, }; use reqwest::Client; use serde::{Deserialize, Serialize}; @@ -82,12 +82,23 @@ impl OpenAiCompatProvider { } } -#[derive(Debug, Serialize)] +#[derive(Debug, Serialize, Default)] struct ChatCompletionRequest { model: String, messages: Vec, #[serde(skip_serializing_if = "Option::is_none")] max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + stream: Option, + #[serde(skip_serializing_if = "Option::is_none")] + stream_options: Option, +} + +/// Ask the server to emit a final usage-only chunk during streaming (OpenAI +/// only reports token counts in a stream when this is set). +#[derive(Debug, Serialize, Default)] +struct StreamOptions { + include_usage: bool, } #[derive(Debug, Serialize, Deserialize)] @@ -116,6 +127,59 @@ struct ChatUsage { completion_tokens: Option, } +/// One `chat.completion.chunk` object from an OpenAI-style SSE stream. +#[derive(Debug, Deserialize)] +struct ChatCompletionChunk { + #[serde(default)] + choices: Vec, + /// Present only on the terminal usage chunk (requires `stream_options`). + #[serde(default)] + usage: Option, +} + +#[derive(Debug, Deserialize)] +struct ChatChunkChoice { + #[serde(default)] + delta: ChatDelta, +} + +#[derive(Debug, Deserialize, Default)] +struct ChatDelta { + #[serde(default)] + content: Option, +} + +/// A parsed SSE event from the chat stream. +enum SseEvent { + /// A `chat.completion.chunk` payload. + Chunk(ChatCompletionChunk), + /// The terminal `data: [DONE]` sentinel. + Done, +} + +/// Parse one SSE event block (the bytes preceding a blank-line delimiter). +/// +/// Concatenates any `data:` field lines, ignoring comments/`event:` lines and +/// tolerating `\r\n`. Blank blocks and unparseable payloads yield `None`. +fn parse_sse_event(block: &[u8]) -> Option { + let text = std::str::from_utf8(block).ok()?; + let mut data = String::new(); + for line in text.lines() { + if let Some(rest) = line.trim_start().strip_prefix("data:") { + data.push_str(rest.trim()); + } + } + if data.is_empty() { + return None; + } + if data == "[DONE]" { + return Some(SseEvent::Done); + } + serde_json::from_str::(&data) + .ok() + .map(SseEvent::Chunk) +} + #[derive(Debug, Serialize)] struct TtsRequest { model: String, @@ -152,6 +216,10 @@ impl Provider for OpenAiCompatProvider { card.context_window = m.context_window.unwrap_or(4096); card.memory_estimate_bytes = m.memory_mb.unwrap_or(0) * 1024 * 1024; card.execution.max_concurrency = 32; + card.reasoning_tier = m + .reasoning_tier + .as_deref() + .and_then(ReasoningEffort::from_str_loose); card.refresh_status(); Some(card) }) @@ -239,6 +307,147 @@ impl Provider for OpenAiCompatProvider { ))), } } + + /// Real per-token streaming over OpenAI-style Server-Sent Events. + /// + /// Overrides the pseudo-streaming default: sends `stream: true`, reads the + /// response body incrementally, and forwards each `choices[].delta.content` + /// token through `sink` as it arrives. Only chat streams token-by-token; + /// other capabilities fall back to a single-shot emit via [`invoke`]. + async fn stream( + &self, + model_id: &str, + request: &InferenceRequest, + sink: StreamSink, + ) -> Result { + let capability = request.capability.unwrap_or(Capability::Chat); + if capability != Capability::Chat { + let response = self.invoke(model_id, request).await?; + if let Some(text) = &response.text + && !text.is_empty() + { + let _ = sink.send(StreamDelta::Text(text.clone())).await; + } + return Ok(response); + } + + let model_name = model_id_name(model_id); + let messages: Vec = request + .messages + .iter() + .map(|m| ChatMessage { + role: m.role.clone(), + content: m.content.clone(), + }) + .collect(); + if messages.is_empty() { + return Err(EngineError::InvalidRequest("no messages provided".into())); + } + let max_tokens = request + .params + .get("max_tokens") + .and_then(|v| v.as_u64()) + .map(|v| v as u32); + + let body = ChatCompletionRequest { + model: model_name.to_string(), + messages, + max_tokens, + stream: Some(true), + stream_options: Some(StreamOptions { + include_usage: true, + }), + }; + + let url = format!("{}/chat/completions", self.base_url); + let start = std::time::Instant::now(); + let mut resp = self + .client + .post(&url) + .bearer_auth(&self.api_key) + .json(&body) + .send() + .await + .map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: e.to_string(), + })?; + + if !resp.status().is_success() { + let status = resp.status(); + let text = resp.text().await.unwrap_or_default(); + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("HTTP {status}: {text}"), + }); + } + + // OpenAI SSE: events are delimited by a blank line (`\n\n`). Buffer bytes, + // emit each token's `delta.content` as its event completes, and capture the + // token counts from the terminal usage-only chunk. + let mut buf: Vec = Vec::new(); + let mut full = String::new(); + let mut prompt_tokens = None; + let mut completion_tokens = None; + let mut total_tokens = None; + + let mut apply = |event: SseEvent, full: &mut String| -> Option { + let SseEvent::Chunk(chunk) = event else { + return None; + }; + if let Some(u) = chunk.usage { + prompt_tokens = u.prompt_tokens; + completion_tokens = u.completion_tokens; + total_tokens = u.total_tokens; + } + let content = chunk.choices.into_iter().next()?.delta.content?; + if content.is_empty() { + return None; + } + full.push_str(&content); + Some(content) + }; + + while let Some(bytes) = resp.chunk().await.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("stream read error: {e}"), + })? { + buf.extend_from_slice(&bytes); + while let Some(idx) = buf.windows(2).position(|w| w == b"\n\n") { + let block: Vec = buf.drain(..idx + 2).collect(); + if let Some(event) = parse_sse_event(&block) + && let Some(delta) = apply(event, &mut full) + { + let _ = sink.send(StreamDelta::Text(delta)).await; + } + } + } + // Any trailing event without a final blank line. + if let Some(event) = parse_sse_event(&buf) + && let Some(delta) = apply(event, &mut full) + { + let _ = sink.send(StreamDelta::Text(delta)).await; + } + + let tokens_used = total_tokens.or(match (prompt_tokens, completion_tokens) { + (None, None) => None, + (p, c) => Some(p.unwrap_or(0) + c.unwrap_or(0)), + }); + Ok(InferenceResponse { + text: Some(full), + file: None, + 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, + prompt_tokens, + completion_tokens, + fallback_used: false, + routing_reason: None, + ..Default::default() + }) + } } impl OpenAiCompatProvider { @@ -259,6 +468,7 @@ impl OpenAiCompatProvider { content: "hi".into(), }], max_tokens: Some(1), + ..Default::default() }; let url = format!("{}/chat/completions", self.base_url); @@ -302,6 +512,7 @@ impl OpenAiCompatProvider { model: model_name.to_string(), messages, max_tokens, + ..Default::default() }; let url = format!("{}/chat/completions", self.base_url); @@ -537,12 +748,14 @@ mod tests { capability: "chat".into(), context_window: Some(8192), memory_mb: None, + ..Default::default() }, ModelDef { name: "model-b".into(), capability: "tts".into(), context_window: None, memory_mb: Some(512), + ..Default::default() }, ], CostTier::Medium, @@ -565,4 +778,62 @@ mod tests { assert_eq!(p.kind(), ProviderKind::OpenAiCompatible); assert_eq!(p.name(), "x"); } + + #[test] + fn sse_events_accumulate_text_and_usage() { + // A realistic OpenAI `stream: true` sequence: two content deltas, an empty + // delta, a terminal usage-only chunk, and the `[DONE]` sentinel. + let blocks: Vec<&[u8]> = vec![ + br#"data: {"choices":[{"delta":{"role":"assistant","content":"He"}}]}"#, + br#"data: {"choices":[{"delta":{"content":"llo"}}]}"#, + br#"data: {"choices":[{"delta":{"content":""}}]}"#, + br#"data: {"choices":[{"delta":{}}],"usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8}}"#, + b"data: [DONE]", + ]; + + let mut full = String::new(); + let mut deltas = Vec::new(); + let mut prompt = None; + let mut completion = None; + let mut total = None; + for block in blocks { + match parse_sse_event(block) { + Some(SseEvent::Chunk(chunk)) => { + if let Some(u) = chunk.usage { + prompt = u.prompt_tokens; + completion = u.completion_tokens; + total = u.total_tokens; + } + if let Some(content) = chunk + .choices + .into_iter() + .next() + .and_then(|c| c.delta.content) + && !content.is_empty() + { + full.push_str(&content); + deltas.push(content); + } + } + Some(SseEvent::Done) | None => {} + } + } + + assert_eq!(deltas, vec!["He", "llo"]); + assert_eq!(full, "Hello"); + assert_eq!(prompt, Some(5)); + assert_eq!(completion, Some(3)); + assert_eq!(total, Some(8)); + } + + #[test] + fn parse_sse_event_handles_done_blank_and_garbage() { + assert!(matches!( + parse_sse_event(b"data: [DONE]"), + Some(SseEvent::Done) + )); + assert!(parse_sse_event(b"").is_none()); + assert!(parse_sse_event(b": keep-alive comment").is_none()); + assert!(parse_sse_event(b"data: not json").is_none()); + } } diff --git a/mofa-engine-core/src/config.rs b/mofa-engine-core/src/config.rs index 4309760..b32ad1f 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -342,7 +342,7 @@ fn default_enabled() -> bool { } /// A model definition within a provider config. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct ModelDef { /// Model name / identifier pub name: String, @@ -352,6 +352,10 @@ pub struct ModelDef { pub context_window: Option, /// Estimated memory in megabytes pub memory_mb: Option, + /// Reasoning tier (`low`/`medium`/`high`) this model serves, for + /// `reasoning.effort` tier routing (S2). Omit for non-reasoning models. + #[serde(default)] + pub reasoning_tier: Option, } impl EngineConfig { @@ -508,30 +512,35 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }, ModelDef { name: "gpt-4o-mini".into(), capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }, ModelDef { name: "tts-1".into(), capability: "tts".into(), context_window: None, memory_mb: None, + ..Default::default() }, ModelDef { name: "tts-1-hd".into(), capability: "tts".into(), context_window: None, memory_mb: None, + ..Default::default() }, ModelDef { name: "whisper-1".into(), capability: "asr".into(), context_window: None, memory_mb: None, + ..Default::default() }, ], enabled: true, @@ -554,12 +563,14 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(64000), memory_mb: None, + ..Default::default() }, ModelDef { name: "deepseek-reasoner".into(), capability: "chat".into(), context_window: Some(64000), memory_mb: None, + ..Default::default() }, ], enabled: true, @@ -582,18 +593,21 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(131072), memory_mb: None, + ..Default::default() }, ModelDef { name: "qwen-turbo".into(), capability: "chat".into(), context_window: Some(131072), memory_mb: None, + ..Default::default() }, ModelDef { name: "qwen-max".into(), capability: "chat".into(), context_window: Some(32768), memory_mb: None, + ..Default::default() }, ], enabled: true, @@ -615,6 +629,7 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }], enabled: true, ..Default::default() @@ -635,6 +650,7 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }], enabled: true, ..Default::default() @@ -655,6 +671,7 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }], enabled: true, ..Default::default() @@ -692,6 +709,8 @@ impl ProviderConfig { "ollama" => Ok(ProviderKind::Ollama), "openai_compatible" => Ok(ProviderKind::OpenAiCompatible), "local_tts" => Ok(ProviderKind::LocalTts), + "local_asr" => Ok(ProviderKind::LocalAsr), + "liter_llm" => Ok(ProviderKind::LiterLlm), other => Err(EngineError::Config(format!( "unknown provider kind '{}' for provider '{}'", other, self.name diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index b4c6a09..7858d3d 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -12,14 +12,16 @@ use std::time::{Duration, Instant}; use dashmap::DashMap; use mofa_kernel::{ BackendHealth, BackendStatus, Capability, CostTier, EngineError, EngineEvent, EngineStatus, - FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, ModelResidency, Provider, - ProviderHealth, ProviderKind, StreamChunk, model_id_name, + FailedAttempt, FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, ModelResidency, + Provider, ProviderHealth, ProviderKind, StreamChunk, StreamDelta, model_id_name, }; use serde::{Deserialize, Serialize}; use tokio::sync::{Mutex as AsyncMutex, Semaphore, broadcast, mpsc}; use tokio::task::{AbortHandle, JoinHandle}; -use crate::backends::{LocalTtsProvider, OllamaProvider, OpenAiCompatProvider}; +use crate::backends::{ + LiterLLMProvider, LocalAsrProvider, LocalTtsProvider, OllamaProvider, OpenAiCompatProvider, +}; use crate::circuit_breaker::{CircuitBreakerConfig, CircuitBreakerRegistry, CircuitState}; use crate::config::{EngineConfig, PreflightConfig, TimeoutConfig}; use crate::memory::{AllocationSnapshot, MemoryManager}; @@ -186,6 +188,14 @@ impl Engine { cost_tier, config.artifacts.dir.clone(), )?), + ProviderKind::LiterLlm => Arc::new(LiterLLMProvider::new( + &pc.name, + pc.api_key.clone(), + &pc.base_url, + pc.models.clone(), + cost_tier, + config.artifacts.dir.clone(), + )?), ProviderKind::LocalTts => { let command = pc.command.clone().ok_or_else(|| { EngineError::Config(format!( @@ -202,6 +212,21 @@ impl Engine { pc.models.clone(), )) } + ProviderKind::LocalAsr => { + let command = pc.command.clone().ok_or_else(|| { + EngineError::Config(format!( + "provider '{}' (local_asr) requires a command", + pc.name + )) + })?; + Arc::new(LocalAsrProvider::new( + &pc.name, + command, + pc.args.clone(), + pc.output_dir.clone(), + pc.models.clone(), + )) + } _ => { return Err(EngineError::Config(format!( "provider '{}' uses unsupported provider kind", @@ -429,6 +454,10 @@ impl Engine { /// Run inference, recording request-level metrics + token/cost usage. pub async fn invoke(&self, req: InferenceRequest) -> Result { let started = Instant::now(); + // Capture audit-relevant fields before `req` is consumed by inference. + let req_prefer = req.prefer; + let req_data_class = req.data_class; + let req_session = req.session_id.clone(); let mut result = self.invoke_inner(req).await; let duration_ms = started.elapsed().as_millis() as u64; match &mut result { @@ -437,12 +466,26 @@ impl Engine { self.compute_cost(&resp.provider, resp.prompt_tokens, resp.completion_tokens); self.metrics .record_request(true, duration_ms, resp.fallback_used); + let is_local = self.provider_is_local(&resp.provider); self.metrics.record_usage( &resp.provider, + is_local, resp.prompt_tokens, resp.completion_tokens, resp.cost_usd, ); + // Data-flow audit (S5): record where each request was actually + // served so residency is traceable by request/session over time. + tracing::info!( + target: "mofa::audit", + request_id = %resp.request_id, + session_id = req_session.as_deref().unwrap_or(""), + data_class = ?req_data_class, + prefer = ?req_prefer, + provider = %resp.provider, + locality = if is_local { "local" } else { "cloud" }, + "request served" + ); } Err(_) => self.metrics.record_request(false, duration_ms, false), } @@ -462,6 +505,15 @@ impl Engine { estimate_cost(price_in, price_out, prompt_tokens, completion_tokens) } + /// Whether the named provider serves models on-device (local track). Unknown + /// providers are treated as cloud so metrics never mislabel spend as local. + fn provider_is_local(&self, provider: &str) -> bool { + self.providers + .iter() + .find(|p| p.name == provider) + .is_some_and(|p| p.kind.is_local()) + } + /// Run inference. /// /// Builds a single ranked candidate plan and walks it in order: the first @@ -488,17 +540,27 @@ impl Engine { ))); } + // Machine-readable routing reason of the primary candidate, surfaced on + // failure so an Agent can see why this plan was chosen. + let routing_reason = candidates.first().map(|d| d.reason.clone()); let mut last_err: Option = None; + // Complete per-candidate failure chain (PRD §4.2), in attempt order. + let mut chain: Vec = Vec::new(); for (idx, decision) in candidates.iter().enumerate() { + let model_id = decision.model.id.clone(); + let provider_name = decision.model.provider.clone(); + if remaining(overall_deadline).is_zero() { - last_err = Some(EngineError::Timeout( - "overall request deadline exceeded".into(), - )); + let e = EngineError::Timeout("overall request deadline exceeded".into()); + chain.push(FailedAttempt { + provider: provider_name, + model: model_id_name(&model_id).to_string(), + reason: e.to_string(), + }); + last_err = Some(e); 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 @@ -522,13 +584,30 @@ impl Engine { return Err(e); } tracing::warn!("candidate '{model_id}' failed (retryable): {e}"); + chain.push(FailedAttempt { + provider: provider_name, + model: model_id_name(&model_id).to_string(), + reason: e.to_string(), + }); last_err = Some(e); } } } - Err(last_err - .unwrap_or_else(|| EngineError::NoCapableModel(Self::requested_capability(&req)))) + let last = last_err + .unwrap_or_else(|| EngineError::NoCapableModel(Self::requested_capability(&req))); + // Wrap the exhausted-candidates case so the caller sees the whole chain, + // not just the last error. A lone failure keeps its original shape. + if chain.len() <= 1 { + return Err(last); + } + Err(EngineError::Failover { + code: last.code(), + message: last.to_string(), + retryable: last.retryable(), + chain, + routing_reason, + }) } fn requested_capability(req: &InferenceRequest) -> String { @@ -791,14 +870,19 @@ impl Engine { }) .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); + // Forward the provider's deltas into the output stream as they arrive, + // mapping text vs thought-chain deltas to their respective chunk types. + // 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() { + let chunk = match delta { + StreamDelta::Text(delta) => StreamChunk::Text { delta }, + StreamDelta::Reasoning(delta) => StreamChunk::Reasoning { delta }, + }; + if forward_out.send(chunk).await.is_err() { break; } } @@ -825,6 +909,7 @@ impl Engine { self.compute_cost(&provider_name, resp.prompt_tokens, resp.completion_tokens); self.metrics.record_usage( &provider_name, + self.provider_is_local(&provider_name), resp.prompt_tokens, resp.completion_tokens, cost_usd, @@ -1954,6 +2039,7 @@ mod tests { messages: vec![Message { role: "user".into(), content: "hello".into(), + ..Default::default() }], input_file: None, params: serde_json::Value::Null, @@ -2123,7 +2209,7 @@ mod tests { engine.refresh_resources().await; let mut req = chat_request(None); - req.locality = mofa_kernel::Locality::LocalOnly; + req.prefer = mofa_kernel::Prefer::Local; // Only a cloud model exists → the request fails rather than leaving the device. assert!(matches!( engine.invoke(req).await, @@ -2179,6 +2265,35 @@ mod tests { assert!(resp.fallback_used); } + #[tokio::test] + async fn exhausted_failover_returns_full_failed_chain() { + // Both the local and cloud candidate fail retryably; the engine exhausts + // the plan and returns a Failover error carrying *every* attempt (PRD §4.2), + // not just the last one — so an Agent can see the whole chain. + 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::FailRetryable), + ); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + let err = engine.invoke(chat_request(None)).await.unwrap_err(); + let info = err.info(); + assert_eq!(info.failed_chain.len(), 2); + // Local is tried first (locality bias), then the cloud fallback. + assert_eq!(info.failed_chain[0].provider, "ollama"); + assert_eq!(info.failed_chain[0].model, "qwen"); + assert_eq!(info.failed_chain[1].provider, "openai"); + assert!(info.retryable); + assert!(info.routing_reason.is_some()); + } + #[tokio::test] async fn local_tts_failure_falls_back_to_cloud_tts() { // A local TTS backend (preferred by locality) crashes mid-synthesis; the @@ -2663,9 +2778,11 @@ mod tests { /// A provider that emits several real text deltas, to prove the engine /// forwards incremental output in order (not just the compatibility path). + /// `reasoning` deltas (if any) are emitted first as thought-chain increments. struct MultiDeltaProvider { name: String, deltas: Vec, + reasoning: Vec, } #[async_trait] @@ -2728,8 +2845,11 @@ mod tests { request: &InferenceRequest, sink: mofa_kernel::StreamSink, ) -> Result { + for r in &self.reasoning { + let _ = sink.send(StreamDelta::Reasoning(r.clone())).await; + } for d in &self.deltas { - let _ = sink.send(d.clone()).await; + let _ = sink.send(StreamDelta::Text(d.clone())).await; } self.invoke(model_id, request).await } @@ -2768,6 +2888,7 @@ mod tests { let provider = Arc::new(MultiDeltaProvider { name: "ollama".into(), deltas: vec!["Hello".into(), ", ".into(), "world".into()], + reasoning: vec![], }); let engine = build_engine(vec![provider], 1000, 0); engine.refresh_resources().await; @@ -2790,6 +2911,38 @@ mod tests { } } + #[tokio::test] + async fn stream_separates_reasoning_from_answer_deltas() { + // S2: the thought chain surfaces as `Reasoning` chunks, distinct from the + // final answer's `Text` chunks, and both precede the terminal `Completed`. + let provider = Arc::new(MultiDeltaProvider { + name: "ollama".into(), + deltas: vec!["4".into()], + reasoning: vec!["2+2".into(), " = ".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 reasoning: Vec = chunks + .iter() + .filter_map(|c| match c { + StreamChunk::Reasoning { delta } => Some(delta.clone()), + _ => None, + }) + .collect(); + let text: Vec = chunks + .iter() + .filter_map(|c| match c { + StreamChunk::Text { delta } => Some(delta.clone()), + _ => None, + }) + .collect(); + assert_eq!(reasoning, vec!["2+2", " = "]); + assert_eq!(text, vec!["4"]); + assert!(matches!(chunks.last(), Some(StreamChunk::Completed { .. }))); + } + #[tokio::test] async fn input_path_allowlist_blocks_paths_outside_roots() { let root = tempfile::tempdir().unwrap(); diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs index f671355..9f7408d 100644 --- a/mofa-engine-core/src/metrics.rs +++ b/mofa-engine-core/src/metrics.rs @@ -21,6 +21,9 @@ struct ProviderUsage { /// Accumulated cost in micro-USD (millionths of a dollar) to avoid float /// atomics; rendered back to USD. cost_micro_usd: AtomicU64, + /// Track this provider serves — `"local"` or `"cloud"` (a provider is always + /// one or the other). Set on first use; drives the dual-track `locality` label. + locality: std::sync::OnceLock<&'static str>, } /// Process-wide engine counters. @@ -84,14 +87,21 @@ impl EngineMetrics { } /// Record per-provider token usage and cost for one successful request. + /// + /// `is_local` tags the series with the dual-track `locality` label so the + /// dashboard can compare on-device vs cloud spend/throughput (PRD §5.3). pub fn record_usage( &self, provider: &str, + is_local: bool, prompt_tokens: Option, completion_tokens: Option, cost_usd: Option, ) { let entry = self.per_provider.entry(provider.to_string()).or_default(); + entry + .locality + .get_or_init(|| if is_local { "local" } else { "cloud" }); if let Some(p) = prompt_tokens { entry.prompt_tokens.fetch_add(p as u64, Ordering::Relaxed); } @@ -248,30 +258,33 @@ impl EngineMetrics { )); } - // Dual-track observability: per-provider token and cost totals. + // Dual-track observability: per-provider token and cost totals, tagged by + // `locality` (local/cloud) so on-device and cloud tracks are comparable. out.push_str( - "# HELP mofa_tokens_total Tokens processed, by provider and direction.\n\ + "# HELP mofa_tokens_total Tokens processed, by provider, locality, and direction.\n\ # TYPE mofa_tokens_total counter\n", ); for entry in self.per_provider.iter() { let provider = escape_label_value(entry.key()); + let locality = entry.locality.get().copied().unwrap_or("unknown"); out.push_str(&format!( - "mofa_tokens_total{{provider=\"{provider}\",direction=\"prompt\"}} {}\n", + "mofa_tokens_total{{provider=\"{provider}\",locality=\"{locality}\",direction=\"prompt\"}} {}\n", entry.prompt_tokens.load(Ordering::Relaxed) )); out.push_str(&format!( - "mofa_tokens_total{{provider=\"{provider}\",direction=\"completion\"}} {}\n", + "mofa_tokens_total{{provider=\"{provider}\",locality=\"{locality}\",direction=\"completion\"}} {}\n", entry.completion_tokens.load(Ordering::Relaxed) )); } out.push_str( - "# HELP mofa_cost_usd_total Estimated spend in USD, by provider.\n\ + "# HELP mofa_cost_usd_total Estimated spend in USD, by provider and locality.\n\ # TYPE mofa_cost_usd_total counter\n", ); for entry in self.per_provider.iter() { let usd = entry.cost_micro_usd.load(Ordering::Relaxed) as f64 / 1_000_000.0; + let locality = entry.locality.get().copied().unwrap_or("unknown"); out.push_str(&format!( - "mofa_cost_usd_total{{provider=\"{}\"}} {usd}\n", + "mofa_cost_usd_total{{provider=\"{}\",locality=\"{locality}\"}} {usd}\n", escape_label_value(entry.key()) )); } @@ -345,19 +358,23 @@ mod tests { #[test] fn per_provider_token_and_cost_render() { let m = EngineMetrics::default(); - m.record_usage("openai", Some(100), Some(40), Some(0.0021)); - m.record_usage("openai", Some(50), Some(10), Some(0.0009)); - m.record_usage("ollama", Some(30), Some(20), None); // local → no cost + m.record_usage("openai", false, Some(100), Some(40), Some(0.0021)); + m.record_usage("openai", false, Some(50), Some(10), Some(0.0009)); + m.record_usage("ollama", true, Some(30), Some(20), None); // local → no cost let text = m.render_prometheus(&MetricsGauges::default()); - assert!(text.contains("mofa_tokens_total{provider=\"openai\",direction=\"prompt\"} 150")); - assert!( - text.contains("mofa_tokens_total{provider=\"openai\",direction=\"completion\"} 50") - ); - assert!(text.contains("mofa_tokens_total{provider=\"ollama\",direction=\"prompt\"} 30")); - // 0.0021 + 0.0009 = 0.003 USD accumulated for openai; ollama has none. - assert!(text.contains("mofa_cost_usd_total{provider=\"openai\"} 0.003")); - assert!(text.contains("mofa_cost_usd_total{provider=\"ollama\"} 0")); + assert!(text.contains( + "mofa_tokens_total{provider=\"openai\",locality=\"cloud\",direction=\"prompt\"} 150" + )); + assert!(text.contains( + "mofa_tokens_total{provider=\"openai\",locality=\"cloud\",direction=\"completion\"} 50" + )); + assert!(text.contains( + "mofa_tokens_total{provider=\"ollama\",locality=\"local\",direction=\"prompt\"} 30" + )); + // 0.0021 + 0.0009 = 0.003 USD accumulated for openai; ollama (local) has none. + assert!(text.contains("mofa_cost_usd_total{provider=\"openai\",locality=\"cloud\"} 0.003")); + assert!(text.contains("mofa_cost_usd_total{provider=\"ollama\",locality=\"local\"} 0")); } #[test] diff --git a/mofa-engine-core/src/router.rs b/mofa-engine-core/src/router.rs index 0acb427..41d0794 100644 --- a/mofa-engine-core/src/router.rs +++ b/mofa-engine-core/src/router.rs @@ -3,13 +3,15 @@ //! The router applies hard constraints first, then scores valid candidates. use mofa_kernel::{ - BackendHealth, Capability, CostTier, InferenceRequest, Locality, ModelCard, ModelResidency, - ModelStatus, ProviderKind, + BackendHealth, Capability, CostTier, InferenceRequest, ModelCard, ModelResidency, ModelStatus, + Prefer, ProviderKind, }; -/// Score bonus that guarantees a `prefer_local` request ranks every local -/// candidate ahead of any cloud one. Larger than any achievable base score. -const LOCAL_FIRST_BONUS: i64 = 1_000_000_000; +/// Bonus applied to a candidate whose reasoning tier matches the request's +/// `reasoning.effort`. Sized to sit *below* the locality term (≈10 000, so +/// local-first still holds) yet *above* the cost/priority terms (≤2 500 each, so +/// effort selects the right tier among same-locality candidates). +const TIER_MATCH_BONUS: i64 = 5_000; /// Provider facts needed by routing. #[derive(Debug, Clone)] @@ -90,9 +92,18 @@ impl Router { continue; } - // Hard locality constraint (privacy / data-residency guardrail): a - // `local_only` request must never leave the device. - if request.locality == Locality::LocalOnly && !provider.kind.is_local() { + // Hard backend-locality constraint (privacy / data-residency + // guardrail). `prefer=local` — and any `confidential` request — must + // never leave the device; `prefer=cloud` restricts to cloud. When + // filtering leaves no candidate the engine turns the empty result into + // a typed `NoCapableModel` (fail-not-fallback). Confidential wins over + // a conflicting `prefer=cloud`: privacy is never traded for preference. + let force_local = + request.prefer == Prefer::Local || request.data_class.requires_local(); + if force_local && !provider.kind.is_local() { + continue; + } + if !force_local && request.prefer == Prefer::Cloud && provider.kind.is_local() { continue; } @@ -110,12 +121,18 @@ impl Router { continue; } + // Under `prefer=auto` (the default), the seven-dimensional score already + // biases toward local via the locality term while keeping cloud as a + // fallback; the hard cases above have removed the ineligible candidates. let mut score = Self::score(model, desired_cap, provider); - // `prefer_local` ranks every local candidate strictly ahead of any - // cloud one (cloud stays available only as a fallback) by adding a - // bonus larger than any achievable base score. - if request.locality == Locality::PreferLocal && provider.kind.is_local() { - score = score.saturating_add(LOCAL_FIRST_BONUS); + // Reasoning-tier routing (S2): when the request declares an effort, a + // model whose configured tier matches gets a bonus placed *below* + // locality (so local-first still holds) but *above* cost/priority, so + // effort selects the right tier among same-locality candidates. + if let Some(reasoning) = request.reasoning + && model.reasoning_tier == Some(reasoning.effort) + { + score = score.saturating_add(TIER_MATCH_BONUS); } if score <= 0 { continue; @@ -279,7 +296,10 @@ impl Router { #[cfg(test)] mod tests { use super::*; - use mofa_kernel::{ExecutionState, FallbackPolicy, Message, ModelAvailability, ModelResidency}; + use mofa_kernel::{ + DataClass, ExecutionState, FallbackPolicy, Message, ModelAvailability, ModelResidency, + Reasoning, ReasoningEffort, + }; fn make_model( name: &str, @@ -316,6 +336,7 @@ mod tests { messages: vec![Message { role: "user".into(), content: "hi".into(), + ..Default::default() }], input_file: None, params: serde_json::Value::Null, @@ -594,7 +615,7 @@ mod tests { provider("ollama", ProviderKind::Ollama, 1), ]; let mut req = request(Capability::Chat, None); - req.locality = Locality::LocalOnly; + req.prefer = Prefer::Local; let ranked = Router::route_ranked(&models, &req, &providers, None); // Only the local candidate survives the hard constraint. assert_eq!(ranked.len(), 1); @@ -612,16 +633,16 @@ mod tests { )]; let providers = vec![provider("openai", ProviderKind::OpenAiCompatible, 10)]; let mut req = request(Capability::Chat, None); - req.locality = Locality::LocalOnly; + req.prefer = Prefer::Local; // No local candidate → no route (engine turns this into NoCapableModel). assert!(Router::route_ranked(&models, &req, &providers, None).is_empty()); } #[test] - fn prefer_local_ranks_local_first_but_keeps_cloud_fallback() { - // A cloud model that is "hotter" than the local one would normally win on - // residency; prefer_local must still rank the local candidate first while - // keeping the cloud one as a fallback. + fn prefer_auto_keeps_both_local_and_cloud_candidates() { + // The default (`prefer=auto`) applies no hard filter: both a local and a + // cloud candidate remain routable, with the local one biased ahead via the + // locality score while cloud stays available as a fallback. let models = vec![ make_model( "cloud", @@ -634,7 +655,7 @@ mod tests { "local", "ollama", Capability::Chat, - ModelResidency::Unloaded, + ModelResidency::Loaded, CostTier::Free, ), ]; @@ -643,10 +664,129 @@ mod tests { provider("ollama", ProviderKind::Ollama, 1), ]; let mut req = request(Capability::Chat, None); - req.locality = Locality::PreferLocal; + req.prefer = Prefer::Auto; let ranked = Router::route_ranked(&models, &req, &providers, None); assert_eq!(ranked.len(), 2); assert_eq!(ranked[0].model.name, "local"); - assert_eq!(ranked[1].model.name, "cloud"); + } + + #[test] + fn cloud_only_filters_out_local_candidates() { + // `prefer=cloud` is the mirror hard constraint: local candidates are dropped + // entirely, leaving only the cloud model. + 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 mut req = request(Capability::Chat, None); + req.prefer = Prefer::Cloud; + let ranked = Router::route_ranked(&models, &req, &providers, None); + assert_eq!(ranked.len(), 1); + assert_eq!(ranked[0].model.name, "cloud"); + } + + #[test] + fn confidential_data_pins_local_even_when_prefer_is_cloud() { + // S5 moat: a confidential request must never leave the device, and privacy + // wins over a conflicting `prefer=cloud` — only the local model survives. + 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 mut req = request(Capability::Chat, None); + req.data_class = DataClass::Confidential; + req.prefer = Prefer::Cloud; // conflicting preference is overridden + let ranked = Router::route_ranked(&models, &req, &providers, None); + assert_eq!(ranked.len(), 1); + assert_eq!(ranked[0].model.name, "local"); + } + + #[test] + fn reasoning_effort_routes_to_matching_tier() { + // Two local chat models identical except for reasoning tier; a high-effort + // request must select the high-tier model, a low-effort one the low-tier. + let mut small = make_model( + "small", + "ollama", + Capability::Chat, + ModelResidency::Loaded, + CostTier::Free, + ); + small.reasoning_tier = Some(ReasoningEffort::Low); + let mut big = make_model( + "big", + "ollama", + Capability::Chat, + ModelResidency::Loaded, + CostTier::Free, + ); + big.reasoning_tier = Some(ReasoningEffort::High); + let models = vec![small, big]; + let providers = vec![provider("ollama", ProviderKind::Ollama, 1)]; + + let mut high = request(Capability::Chat, None); + high.reasoning = Some(Reasoning { + effort: ReasoningEffort::High, + include: false, + }); + let ranked = Router::route_ranked(&models, &high, &providers, None); + assert_eq!(ranked[0].model.name, "big"); + + let mut low = request(Capability::Chat, None); + low.reasoning = Some(Reasoning { + effort: ReasoningEffort::Low, + include: false, + }); + let ranked = Router::route_ranked(&models, &low, &providers, None); + assert_eq!(ranked[0].model.name, "small"); + } + + #[test] + fn confidential_with_no_local_fails_rather_than_leaking() { + let models = vec![make_model( + "cloud", + "openai", + Capability::Chat, + ModelResidency::Remote, + CostTier::High, + )]; + let providers = vec![provider("openai", ProviderKind::OpenAiCompatible, 10)]; + let mut req = request(Capability::Chat, None); + req.data_class = DataClass::Confidential; + // No local candidate → no route (engine turns this into NoCapableModel), + // never a silent cloud fallback. + assert!(Router::route_ranked(&models, &req, &providers, None).is_empty()); } } diff --git a/mofa-engine-sdk/examples/README.md b/mofa-engine-sdk/examples/README.md new file mode 100644 index 0000000..aa3a7f9 --- /dev/null +++ b/mofa-engine-sdk/examples/README.md @@ -0,0 +1,30 @@ +# MoFA Engine — scenario examples + +Runnable demos that exercise the engine end-to-end through the native SDK +(`EmbeddedEngine`) and the async `Engine`. Each maps to a PRD product scenario and +its acceptance anchor ("replicate the demo = scenario accepted"). All degrade +gracefully: with no capable backend they print the structured error (including any +`failed_chain`) instead of panicking. + +| Example | Scenario | Engine capabilities exercised | +|---|---|---| +| `code_review.rs` | **S2** Code/PR Review | `reasoning.effort` (tier routing) + **streamed thought chain** (`Reasoning` vs `Text` chunks) + token/cost | +| `doc_ai.rs` | **S3** Document/Screenshot AI | `Vlm` understanding, multimodal `Message.images`, `detail` billing tier | +| `meeting_brief.rs` | **S1** Meeting → Minutes + Brief | local-first `Asr → Chat → Tts` pipeline, `prefer=local` + `data_class=confidential`, `hint_next` warmup | + +## Running + +```bash +# S2 — needs a chat/reasoning backend (Ollama or a cloud key) +cargo run -p mofa-engine-sdk --example code_review + +# S3 — needs a VLM-capable backend (cloud via liter-llm); pass an image +cargo run -p mofa-engine-sdk --example doc_ai -- ./receipt.jpg + +# S1 — needs local ASR + a chat model + TTS; pass a recording +cargo run -p mofa-engine-sdk --example meeting_brief -- ./meeting.wav +``` + +Configuration is loaded from `./config.toml` (see `config.example.toml`) or falls +back to environment auto-detection (`OPENAI_API_KEY`, `DEEPSEEK_API_KEY`, …) and a +local Ollama at `127.0.0.1:11434`. diff --git a/mofa-engine-sdk/examples/code_review.rs b/mofa-engine-sdk/examples/code_review.rs new file mode 100644 index 0000000..29d7de1 --- /dev/null +++ b/mofa-engine-sdk/examples/code_review.rs @@ -0,0 +1,119 @@ +//! Scenario S2 — Code/PR Review Agent. +//! +//! Demonstrates the engine's deep-thinking path: a high-effort `reasoning` +//! request whose **thought chain streams as `Reasoning` chunks, distinct from the +//! final answer's `Text` chunks**, with token/cost reported on completion. When a +//! model declares a `reasoning_tier`, `effort=high` routes to the stronger tier. +//! +//! Run (needs a chat/reasoning backend, e.g. Ollama or a cloud key): +//! cargo run -p mofa-engine-sdk --example code_review +//! +//! The demo degrades gracefully: with no capable backend it prints the structured +//! error (including any `failed_chain`) instead of panicking. + +use std::io::Write; + +use mofa_engine_core::{Engine, EngineConfig}; +use mofa_kernel::{Capability, InferenceRequest, Message, Reasoning, ReasoningEffort, StreamChunk}; + +const DIFF: &str = r#"--- a/auth.rs ++++ b/auth.rs +@@ +- let token = req.header("authorization").unwrap(); ++ let token = req.header("authorization").unwrap_or(""); + if token == expected { grant() } else { deny() } +"#; + +#[tokio::main] +async fn main() { + let config = EngineConfig::load(None); + let engine = match Engine::try_new(config).await { + Ok(e) => e, + Err(e) => { + eprintln!("engine init failed: {e}"); + return; + } + }; + engine.refresh_resources().await; + + let prompt = format!( + "Review the following diff. Annotate issues (severity: blocker/major/minor) \ + and suggest fixes:\n{DIFF}" + ); + let request = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![ + Message { + role: "system".into(), + content: "You are a senior code reviewer. Think step by step, then \ + give a concise annotated report." + .into(), + ..Default::default() + }, + Message { + role: "user".into(), + content: prompt, + ..Default::default() + }, + ], + // High effort routes to the strongest reasoning tier (see `reasoning_tier` + // in config) and surfaces the thought chain. + reasoning: Some(Reasoning { + effort: ReasoningEffort::High, + include: true, + }), + ..Default::default() + }; + + println!("== S2 Code Review — streaming (reasoning vs answer) ==\n"); + let mut rx = engine.invoke_stream(request); + let mut in_thoughts = false; + let mut in_answer = false; + while let Some(chunk) = rx.recv().await { + match chunk { + StreamChunk::Started { + model_used, + provider, + .. + } => { + println!("[routed to {provider}/{model_used}]\n"); + } + StreamChunk::Reasoning { delta } => { + if !in_thoughts { + print!("🧠 thinking: "); + in_thoughts = true; + } + print!("{delta}"); + let _ = std::io::stdout().flush(); + } + StreamChunk::Text { delta } => { + if !in_answer { + println!("\n\n📝 report:"); + in_answer = true; + } + print!("{delta}"); + let _ = std::io::stdout().flush(); + } + StreamChunk::Completed { + tokens_used, + cost_usd, + .. + } => { + println!( + "\n\n[done — tokens: {:?}, est. cost: {:?}]", + tokens_used, cost_usd + ); + } + StreamChunk::Error(info) => { + eprintln!("\nerror [{:?}]: {}", info.code, info.message); + for attempt in &info.failed_chain { + eprintln!( + " tried {}/{}: {}", + attempt.provider, attempt.model, attempt.reason + ); + } + } + _ => {} + } + } +} diff --git a/mofa-engine-sdk/examples/doc_ai.rs b/mofa-engine-sdk/examples/doc_ai.rs new file mode 100644 index 0000000..e7dad9b --- /dev/null +++ b/mofa-engine-sdk/examples/doc_ai.rs @@ -0,0 +1,62 @@ +//! Scenario S3 — Document/Screenshot AI. +//! +//! Demonstrates vision-language understanding: a photo (receipt / screenshot) plus +//! a question, extracted into **structured data**. Images ride on +//! `Message.images` (URL, `data:` URL, or local path) and the `detail` tier +//! (`low | high | auto`) passes through to the vendor's image billing. +//! +//! Run (needs a VLM-capable backend, e.g. a cloud key via liter-llm): +//! cargo run -p mofa-engine-sdk --example doc_ai -- ./receipt.jpg +//! +//! With no capable backend it prints the structured error instead of panicking. + +use mofa_engine_core::EngineConfig; +use mofa_engine_sdk::EmbeddedEngine; +use mofa_kernel::{Capability, InferenceRequest, Message}; + +fn main() { + // The image reference: a CLI arg, or a placeholder URL for a dry run. + let image = std::env::args() + .nth(1) + .unwrap_or_else(|| "https://example.com/receipt.jpg".to_string()); + + let config = EngineConfig::load(None); + let engine = match EmbeddedEngine::new(config) { + Ok(e) => e, + Err(e) => { + eprintln!("engine init failed: {e}"); + return; + } + }; + engine.refresh(); + + let request = InferenceRequest { + capability: Some(Capability::Vlm), + messages: vec![Message { + role: "user".into(), + content: "Extract amount, date, and category from this receipt. \ + Respond as JSON." + .into(), + images: vec![image.clone()], + }], + // `detail=low` keeps batch cost down; edge cases can re-run at `high`. + params: serde_json::json!({ "detail": "low" }), + ..Default::default() + }; + + println!("== S3 Document AI — VLM extract =="); + println!("image: {image}\n"); + match engine.invoke(request) { + Ok(resp) => { + println!("provider: {}", resp.provider); + println!("extracted: {}", resp.text.as_deref().unwrap_or("(no text)")); + if let Some(cost) = resp.cost_usd { + println!("est. cost: ${cost:.6}"); + } + } + Err(e) => { + let info = e.info(); + eprintln!("error [{:?}]: {}", info.code, info.message); + } + } +} diff --git a/mofa-engine-sdk/examples/meeting_brief.rs b/mofa-engine-sdk/examples/meeting_brief.rs new file mode 100644 index 0000000..2ce2de8 --- /dev/null +++ b/mofa-engine-sdk/examples/meeting_brief.rs @@ -0,0 +1,127 @@ +//! Scenario S1 — Meeting/Lecture → Minutes + Audio Brief. +//! +//! Demonstrates a confidential, local-first pipeline: long-audio **ASR → +//! structured minutes (chat) → 30s audio brief (TTS)**. The whole chain is pinned +//! on-device with `prefer=local` + `data_class=confidential`, so a confidential +//! recording never leaves the machine (it fails rather than falling back to cloud). +//! `hint_next="tts"` warms the TTS model while the minutes are generated. +//! +//! Run (needs local ASR + a chat model + local/cloud TTS configured): +//! cargo run -p mofa-engine-sdk --example meeting_brief -- ./meeting.wav +//! +//! Each stage degrades gracefully: a missing backend prints the structured error +//! (a confidential stage that can't run locally fails on purpose) and stops. + +use mofa_engine_core::EngineConfig; +use mofa_engine_sdk::EmbeddedEngine; +use mofa_kernel::{Capability, DataClass, InferenceRequest, InferenceResponse, Message, Prefer}; + +fn main() { + let audio = std::env::args() + .nth(1) + .unwrap_or_else(|| "meeting.wav".into()); + + let config = EngineConfig::load(None); + let engine = match EmbeddedEngine::new(config) { + Ok(e) => e, + Err(e) => { + eprintln!("engine init failed: {e}"); + return; + } + }; + engine.refresh(); + + println!("== S1 Meeting Brief — confidential, local-first =="); + println!("audio: {audio}\n"); + + // 1. ASR — long audio → transcript, pinned local (confidential). + let asr = InferenceRequest { + capability: Some(Capability::Asr), + input_file: Some(audio), + prefer: Prefer::Local, + data_class: DataClass::Confidential, + hint_next: Some("chat".into()), + ..Default::default() + }; + let transcript = match run(&engine, asr, "ASR") { + Some(r) => r.text.unwrap_or_default(), + None => return, + }; + println!("transcript: {}\n", preview(&transcript)); + + // 2. Chat — extract structured minutes; warm TTS meanwhile. + let chat = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![ + Message { + role: "system".into(), + content: "Extract structured meeting minutes: resolutions, todos, \ + risks, and responsible people." + .into(), + ..Default::default() + }, + Message { + role: "user".into(), + content: transcript, + ..Default::default() + }, + ], + prefer: Prefer::Local, + data_class: DataClass::Confidential, + hint_next: Some("tts".into()), + ..Default::default() + }; + let minutes = match run(&engine, chat, "Chat (minutes)") { + Some(r) => r.text.unwrap_or_default(), + None => return, + }; + println!("minutes:\n{}\n", preview(&minutes)); + + // 3. TTS — minutes → a short audio brief. + let tts = InferenceRequest { + capability: Some(Capability::Tts), + messages: vec![Message { + role: "user".into(), + content: minutes, + ..Default::default() + }], + prefer: Prefer::Local, + data_class: DataClass::Confidential, + params: serde_json::json!({ "voice": "zh-female-1" }), + ..Default::default() + }; + if let Some(r) = run(&engine, tts, "TTS (brief)") { + println!("audio brief: {}", r.file.as_deref().unwrap_or("(no file)")); + } +} + +/// Invoke one stage, printing routing/cost on success or the structured error +/// (including any failover chain) on failure. Returns `None` to stop the pipeline. +fn run(engine: &EmbeddedEngine, req: InferenceRequest, stage: &str) -> Option { + match engine.invoke(req) { + Ok(resp) => { + println!( + "[{stage}] served by {} in {}ms", + resp.provider, resp.duration_ms + ); + Some(resp) + } + Err(e) => { + let info = e.info(); + eprintln!("[{stage}] error [{:?}]: {}", info.code, info.message); + for a in &info.failed_chain { + eprintln!(" tried {}/{}: {}", a.provider, a.model, a.reason); + } + None + } + } +} + +fn preview(s: &str) -> String { + let s = s.trim(); + if s.chars().count() > 200 { + format!("{}…", s.chars().take(200).collect::()) + } else { + s.to_string() + } +} diff --git a/mofa-engine-sdk/src/server.rs b/mofa-engine-sdk/src/server.rs index 2af2965..9c20294 100644 --- a/mofa-engine-sdk/src/server.rs +++ b/mofa-engine-sdk/src/server.rs @@ -271,7 +271,9 @@ fn error_status(e: &EngineError) -> StatusCode { EngineError::InvalidRequest(_) | EngineError::UnsupportedOperation(_) => { StatusCode::BAD_REQUEST } - EngineError::CircuitOpen(_) => StatusCode::SERVICE_UNAVAILABLE, + EngineError::CircuitOpen(_) | EngineError::Failover { .. } => { + StatusCode::SERVICE_UNAVAILABLE + } EngineError::Timeout(_) => StatusCode::GATEWAY_TIMEOUT, _ => StatusCode::INTERNAL_SERVER_ERROR, } diff --git a/mofa-kernel/src/error.rs b/mofa-kernel/src/error.rs index 85dfd9e..9a13343 100644 --- a/mofa-kernel/src/error.rs +++ b/mofa-kernel/src/error.rs @@ -27,6 +27,18 @@ pub enum ErrorCode { Internal, } +/// One failed candidate in a failover chain: which provider/model was tried and +/// why it failed. Lets an Agent see *all* attempts, not just the last error. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct FailedAttempt { + /// Provider that was tried. + pub provider: String, + /// Model that was tried (canonical short name). + pub model: String, + /// Human-readable failure reason. + pub reason: String, +} + /// Structured error body suitable for HTTP and SDKs. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ErrorInfo { @@ -38,6 +50,12 @@ pub struct ErrorInfo { pub retryable: bool, /// Optional source provider/backend. pub source: Option, + /// Complete per-candidate failure chain (empty unless failover was attempted). + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub failed_chain: Vec, + /// Machine-readable routing reason for the primary candidate, when known. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub routing_reason: Option, } /// Top-level engine error. @@ -89,6 +107,23 @@ pub enum EngineError { /// An internal/unexpected error. #[error("internal error: {0}")] Internal(String), + + /// Every routing candidate was exhausted during failover. Carries the full + /// per-candidate failure chain so an Agent can decide whether to retry or + /// switch tiers, rather than seeing only the last error. + #[error("all candidates failed: {message}")] + Failover { + /// Code of the final (last-tried) failure. + code: ErrorCode, + /// Aggregate human-readable message (the last failure's message). + message: String, + /// Whether retrying the whole request might succeed. + retryable: bool, + /// Per-candidate failures, in attempt order. + chain: Vec, + /// Machine-readable routing reason for the primary candidate. + routing_reason: Option, + }, } impl EngineError { @@ -104,19 +139,21 @@ impl EngineError { Self::UnsupportedOperation(_) => ErrorCode::UnsupportedOperation, Self::Config(_) => ErrorCode::Config, Self::Internal(_) => ErrorCode::Internal, + Self::Failover { code, .. } => *code, } } /// Whether retrying the same request may succeed. pub fn retryable(&self) -> bool { - matches!( - self, + match self { Self::ProviderError { .. } - | Self::CircuitOpen(_) - | Self::MemoryPressure { .. } - | Self::Timeout(_) - | Self::Internal(_) - ) + | Self::CircuitOpen(_) + | Self::MemoryPressure { .. } + | Self::Timeout(_) + | Self::Internal(_) => true, + Self::Failover { retryable, .. } => *retryable, + _ => false, + } } /// Optional source provider/backend for this error. @@ -129,11 +166,30 @@ impl EngineError { /// Convert to a structured error body. pub fn info(&self) -> ErrorInfo { + if let Self::Failover { + code, + message, + retryable, + chain, + routing_reason, + } = self + { + return ErrorInfo { + code: *code, + message: message.clone(), + retryable: *retryable, + source: chain.last().map(|a| a.provider.clone()), + failed_chain: chain.clone(), + routing_reason: routing_reason.clone(), + }; + } ErrorInfo { code: self.code(), message: self.to_string(), retryable: self.retryable(), source: self.source_name().map(ToOwned::to_owned), + failed_chain: Vec::new(), + routing_reason: None, } } } @@ -161,5 +217,50 @@ mod tests { assert_eq!(info.code, ErrorCode::CircuitOpen); assert!(info.retryable); assert_eq!(info.source.as_deref(), Some("ollama")); + // Non-failover errors carry an empty chain that is omitted from JSON. + assert!(info.failed_chain.is_empty()); + assert!( + !serde_json::to_string(&info) + .unwrap() + .contains("failed_chain") + ); + } + + #[test] + fn failover_info_carries_full_chain() { + let e = EngineError::Failover { + code: ErrorCode::ProviderError, + message: "provider 'openai' error: 503".into(), + retryable: true, + chain: vec![ + FailedAttempt { + provider: "ollama".into(), + model: "qwen".into(), + reason: "circuit open".into(), + }, + FailedAttempt { + provider: "openai".into(), + model: "gpt-4o".into(), + reason: "503".into(), + }, + ], + routing_reason: Some("resident local preferred".into()), + }; + let info = e.info(); + assert_eq!(info.code, ErrorCode::ProviderError); + assert!(info.retryable); + assert_eq!(info.failed_chain.len(), 2); + assert_eq!(info.failed_chain[0].provider, "ollama"); + assert_eq!(info.failed_chain[1].model, "gpt-4o"); + assert_eq!( + info.routing_reason.as_deref(), + Some("resident local preferred") + ); + // `source` points at the last-tried provider. + assert_eq!(info.source.as_deref(), Some("openai")); + + let json = serde_json::to_string(&info).unwrap(); + assert!(json.contains("failed_chain")); + assert!(json.contains("routing_reason")); } } diff --git a/mofa-kernel/src/lib.rs b/mofa-kernel/src/lib.rs index 79acdcc..31dbca7 100644 --- a/mofa-kernel/src/lib.rs +++ b/mofa-kernel/src/lib.rs @@ -8,6 +8,6 @@ pub mod error; pub mod traits; pub mod types; -pub use error::{EngineError, ErrorCode, ErrorInfo}; +pub use error::{EngineError, ErrorCode, ErrorInfo, FailedAttempt}; pub use traits::Provider; pub use types::*; diff --git a/mofa-kernel/src/traits.rs b/mofa-kernel/src/traits.rs index 6767793..581b366 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, StreamSink, + ProviderKind, StreamDelta, StreamSink, }; /// A model provider backend. @@ -62,7 +62,7 @@ pub trait Provider: Send + Sync { if let Some(text) = &response.text && !text.is_empty() { - let _ = sink.send(text.clone()).await; + let _ = sink.send(StreamDelta::Text(text.clone())).await; } Ok(response) } diff --git a/mofa-kernel/src/types.rs b/mofa-kernel/src/types.rs index c22e549..c3fdf69 100644 --- a/mofa-kernel/src/types.rs +++ b/mofa-kernel/src/types.rs @@ -68,6 +68,11 @@ pub enum ProviderKind { OpenAiCompatible, /// Local process-adapter backend (e.g. an MLX/Kokoro or Piper TTS CLI). LocalTts, + /// Local process-adapter ASR backend (e.g. a FunASR or whisper.cpp CLI). + LocalAsr, + /// Multi-vendor cloud gateway via the `liter-llm` crate (143+ providers, + /// unified OpenAI-style contract). + LiterLlm, } impl ProviderKind { @@ -75,7 +80,7 @@ impl ProviderKind { /// 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) + matches!(self, Self::Ollama | Self::LocalTts | Self::LocalAsr) } } @@ -268,6 +273,11 @@ pub struct ModelCard { pub context_window: u32, /// Estimated memory footprint in bytes. pub memory_estimate_bytes: u64, + /// Reasoning tier this model serves, when configured. Lets `reasoning.effort` + /// route `low | medium | high` to cheaper/stronger models (S2). `None` for + /// non-tiered models. + #[serde(default)] + pub reasoning_tier: Option, } impl ModelCard { @@ -297,6 +307,7 @@ impl ModelCard { cost_tier, context_window: 4096, memory_estimate_bytes: 0, + reasoning_tier: None, } } @@ -335,12 +346,16 @@ pub fn model_id_name(model_id: &str) -> &str { } /// A single message in a conversation. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct Message { /// Role: "system", "user", "assistant". pub role: String, - /// Message content. + /// Message text content. pub content: String, + /// Image references for multimodal (VLM) requests — HTTP(S) URLs, `data:` + /// URLs, or local file paths. Empty for text-only messages. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub images: Vec, } /// Named-model fallback behavior. @@ -356,20 +371,84 @@ pub enum FallbackPolicy { AllowNamed, } -/// Locality preference for routing (privacy / data-residency guardrail). +/// Backend-locality preference for routing (privacy / data-residency guardrail). +/// +/// Matches the PRD `prefer` request field: `auto | local | cloud`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] #[serde(rename_all = "snake_case")] -pub enum Locality { - /// Normal ranking: local models are preferred but a cloud model can be - /// selected or used as a fallback. +pub enum Prefer { + /// Default seven-dimensional scoring: local models are preferred but a cloud + /// model can be selected or used as a fallback. #[default] Auto, - /// Local models are ranked strictly ahead of any cloud model, which remains - /// available only as a fallback. - PreferLocal, - /// Hard constraint: only local models may serve the request. If none can, - /// the request fails rather than leaving the device. - LocalOnly, + /// Hard constraint: only local models may serve the request. If none can, the + /// request fails rather than leaving the device (fail-not-fallback). + #[serde(alias = "local_only")] + Local, + /// Hard constraint: only cloud models may serve the request. + Cloud, +} + +/// Extended-thinking effort for reasoning-capable models. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum ReasoningEffort { + /// Minimal deliberation — cheapest, routes to the lightest tier. + Low, + /// Balanced deliberation (default). + #[default] + Medium, + /// Maximum deliberation — routes to the strongest reasoning tier. + High, +} + +impl ReasoningEffort { + /// Parse a loose config/string value (`low`/`medium`/`high`, any case). + pub fn from_str_loose(s: &str) -> Option { + match s.trim().to_ascii_lowercase().as_str() { + "low" => Some(Self::Low), + "medium" | "med" => Some(Self::Medium), + "high" => Some(Self::High), + _ => None, + } + } +} + +/// Data-sensitivity classification for a request (S5 privacy moat). +/// +/// `Confidential` is a hard data-residency constraint: the request is pinned to +/// local backends regardless of `prefer`, and fails rather than sending sensitive +/// data to the cloud. Every request's effective locality is written to the audit +/// log so data flow is traceable. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum DataClass { + /// No residency constraint (default). + #[default] + Public, + /// Organization-internal; no hard constraint, but audited. + Internal, + /// Sensitive: must never leave the device (implies local-only routing). + Confidential, +} + +impl DataClass { + /// Whether this class forbids sending the request to a cloud backend. + pub fn requires_local(self) -> bool { + matches!(self, Self::Confidential) + } +} + +/// Deep-thinking controls for a request (S2 Code/PR Review). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] +pub struct Reasoning { + /// How much thinking effort to spend; also used for tier routing. + #[serde(default)] + pub effort: ReasoningEffort, + /// Whether the thought chain should be surfaced (streamed as `reasoning` + /// increments and returned) rather than stripped from the output. + #[serde(default)] + pub include: bool, } /// A request to the engine for inference. @@ -386,9 +465,17 @@ pub struct InferenceRequest { /// Named fallback policy. #[serde(default)] pub fallback_policy: FallbackPolicy, - /// Locality preference (privacy guardrail). + /// Backend-locality preference / privacy guardrail (`auto | local | cloud`). + #[serde(default, alias = "locality")] + pub prefer: Prefer, + /// Data-sensitivity class. `confidential` pins the request to local backends + /// regardless of `prefer` (privacy moat). #[serde(default)] - pub locality: Locality, + pub data_class: DataClass, + /// Deep-thinking controls (effort tier + thought-chain visibility). `None` + /// keeps standard behavior. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub reasoning: Option, /// Conversation messages. #[serde(default)] pub messages: Vec, @@ -412,7 +499,9 @@ impl Default for InferenceRequest { app_id: None, session_id: None, fallback_policy: FallbackPolicy::default(), - locality: Locality::default(), + prefer: Prefer::default(), + data_class: DataClass::default(), + reasoning: None, messages: Vec::new(), input_file: None, params: serde_json::Value::Null, @@ -482,6 +571,13 @@ pub enum StreamChunk { /// Text delta appended to the output so far. delta: String, }, + /// An incremental piece of the model's thought chain (reasoning-capable + /// models). Kept distinct from `Text` so callers can display or audit the + /// thought chain separately from the final answer (S2). + Reasoning { + /// Thought-chain delta. + delta: String, + }, /// Terminal success event with aggregate metadata and any file output. Completed { /// Wall-clock duration in milliseconds. @@ -508,11 +604,32 @@ pub enum StreamChunk { Error(ErrorInfo), } -/// Channel a provider uses to emit incremental text deltas while streaming. +/// One incremental delta a provider pushes while streaming: either final-answer +/// text or a piece of the model's thought chain. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum StreamDelta { + /// A piece of the final answer. + Text(String), + /// A piece of the thought chain (reasoning-capable models). + Reasoning(String), +} + +impl StreamDelta { + /// Borrow the delta text regardless of kind. + pub fn as_str(&self) -> &str { + match self { + Self::Text(s) | Self::Reasoning(s) => s, + } + } +} + +/// Channel a provider uses to emit incremental 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; +/// providers push [`StreamDelta`] items here and return the final aggregate +/// response. `Text` deltas become `StreamChunk::Text`, `Reasoning` deltas become +/// `StreamChunk::Reasoning`. +pub type StreamSink = tokio::sync::mpsc::Sender; /// Result of a lifecycle operation. #[derive(Debug, Clone, Serialize, Deserialize)] @@ -732,4 +849,59 @@ mod tests { let json = serde_json::to_string(&event).unwrap(); assert!(json.contains("model_status_changed")); } + + #[test] + fn message_images_are_backward_compatible() { + // A legacy text-only message (no `images`) still deserializes, and a + // text-only message serializes without an `images` field. + let m: Message = serde_json::from_str(r#"{"role":"user","content":"hi"}"#).unwrap(); + assert!(m.images.is_empty()); + assert!(!serde_json::to_string(&m).unwrap().contains("images")); + + // A multimodal message round-trips its image references. + let mm: Message = serde_json::from_str( + r#"{"role":"user","content":"what is this?","images":["https://x/a.png"]}"#, + ) + .unwrap(); + assert_eq!(mm.images, vec!["https://x/a.png"]); + } + + #[test] + fn prefer_and_reasoning_deserialize_from_contract() { + // PRD wire contract: `prefer: local` + `reasoning: { effort: high }`. + let req: InferenceRequest = serde_json::from_str( + r#"{"prefer":"local","reasoning":{"effort":"high","include":true}}"#, + ) + .unwrap(); + assert_eq!(req.prefer, Prefer::Local); + let r = req.reasoning.unwrap(); + assert_eq!(r.effort, ReasoningEffort::High); + assert!(r.include); + // Legacy alias still accepted. + let legacy: InferenceRequest = + serde_json::from_str(r#"{"locality":"local_only"}"#).unwrap(); + assert_eq!(legacy.prefer, Prefer::Local); + } + + #[test] + fn data_class_defaults_public_and_confidential_requires_local() { + let default: InferenceRequest = serde_json::from_str("{}").unwrap(); + assert_eq!(default.data_class, DataClass::Public); + assert!(!default.data_class.requires_local()); + + let sensitive: InferenceRequest = + serde_json::from_str(r#"{"data_class":"confidential"}"#).unwrap(); + assert_eq!(sensitive.data_class, DataClass::Confidential); + assert!(sensitive.data_class.requires_local()); + } + + #[test] + fn reasoning_stream_chunk_has_distinct_tag() { + let chunk = StreamChunk::Reasoning { + delta: "thinking".into(), + }; + let json = serde_json::to_string(&chunk).unwrap(); + assert!(json.contains(r#""type":"reasoning""#)); + assert!(json.contains(r#""delta":"thinking""#)); + } } From dc445fb0283ae3e4a5e73432ef4d89e57e26529c Mon Sep 17 00:00:00 2001 From: aayank13 Date: Thu, 30 Jul 2026 00:45:37 +0530 Subject: [PATCH 08/22] feat: add LocalImageGen provider and enhance ASR with speaker diarization - Introduced `LocalImageGen` provider to support local image generation. - Enhanced `ProviderConfig` to include `diarize_args` for ASR speaker diarization. - Updated `Engine` to handle new provider and incorporate cost estimation for requests. - Implemented a quality gate for video outputs, ensuring integrity and motion checks. - Added example for orchestrating an explainer video with a hard quality gate. - Updated existing examples to include cost control and speaker diarization features. --- config.example.toml | 40 ++ docs/openapi.yaml | 16 + mofa-engine-core/src/backends.rs | 2 + mofa-engine-core/src/backends/local_asr.rs | 126 ++++- .../src/backends/local_image_gen.rs | 518 ++++++++++++++++++ mofa-engine-core/src/backends/ollama.rs | 84 ++- mofa-engine-core/src/config.rs | 7 + mofa-engine-core/src/engine.rs | 235 +++++++- mofa-engine-core/src/lib.rs | 1 + mofa-engine-core/src/quality_gate.rs | 501 +++++++++++++++++ mofa-engine-sdk/examples/README.md | 7 +- mofa-engine-sdk/examples/code_review.rs | 4 + mofa-engine-sdk/examples/explainer_video.rs | 144 +++++ mofa-engine-sdk/examples/meeting_brief.rs | 4 +- mofa-kernel/src/types.rs | 15 +- 15 files changed, 1683 insertions(+), 21 deletions(-) create mode 100644 mofa-engine-core/src/backends/local_image_gen.rs create mode 100644 mofa-engine-core/src/quality_gate.rs create mode 100644 mofa-engine-sdk/examples/explainer_video.rs diff --git a/config.example.toml b/config.example.toml index 0020bb6..0b996f7 100644 --- a/config.example.toml +++ b/config.example.toml @@ -47,6 +47,17 @@ base_url = "http://127.0.0.1:11434" priority = 1 cost_tier = "free" + # Ollama models are auto-discovered, so listing them is optional. Declare a + # model here only to annotate it — currently to give a *local* reasoning model + # a `reasoning_tier` (low/medium/high) so `reasoning.effort` routes to it. + # `name` must match the Ollama model name exactly (as `ollama list` shows it). + # Local reasoning tiers win over cloud ones of equal effort (local-first), so a + # high-effort code review can stay on-device. + [[providers.models]] + name = "deepseek-r1:8b" + capability = "chat" + reasoning_tier = "high" + # Cloud OpenAI-compatible backend. Models are declared explicitly. [[providers]] name = "openai" @@ -154,6 +165,9 @@ kind = "local_asr" base_url = "" # unused for local_asr, but the field is required command = "funasr" args = ["--input", "{input}", "--language", "{language}", "--output", "{output}"] +# Appended only when a request sets params.diarize = true (S1 speaker separation); +# placeholder-substituted like `args`. Leave unset if the CLI can't diarize. +diarize_args = ["--diarize"] priority = 2 enabled = false # opt-in: set true once the command is installed @@ -161,3 +175,29 @@ enabled = false # opt-in: set true once the command is installed name = "funasr" capability = "asr" memory_mb = 1024 + +# Local image-generation process-adapter backend. Runs a local Stable Diffusion +# CLI once per image and returns the produced file. Enables fully-offline scene +# generation for the flagship S4 pipeline (`prefer=local`, cost ~= $0), failing +# over to a cloud image backend (liter_llm) when unavailable. Placeholders: +# {prompt} {negative_prompt} {width} {height} {output} +[[providers]] +name = "local-sd" +kind = "local_image_gen" +base_url = "" # unused for local_image_gen, but the field is required +command = "sd" # e.g. stable-diffusion.cpp's `sd`, or a wrapper script +args = [ + "--prompt", "{prompt}", + "--negative-prompt", "{negative_prompt}", + "--width", "{width}", + "--height", "{height}", + "--output", "{output}", +] +output_format = "png" +priority = 2 +enabled = false # opt-in: set true once the command is installed + + [[providers.models]] + name = "sd" + capability = "image_gen" + memory_mb = 4096 diff --git a/docs/openapi.yaml b/docs/openapi.yaml index f503914..94583a3 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -311,6 +311,17 @@ components: include: type: boolean default: false + max_cost_usd: + type: number + format: double + description: > + Per-request spend ceiling in USD. A candidate whose estimated cost + exceeds this is skipped during routing (and listed in the error's + `failed_chain`), so spend stays bounded and cheaper/local models win. + Free and local models estimate to $0 and are always affordable, so + `0.0` means "free/local only". Omit for no ceiling. The estimate is a + pre-flight approximation (token counts from input size); the + authoritative cost is reported per response as `cost_usd`. messages: type: array items: @@ -320,6 +331,11 @@ components: description: Path on the engine host. params: type: object + description: > + Capability-specific passthrough options. Recognized keys include + `language` and `diarize` (ASR speaker separation, S1), `detail` + (VLM billing tier), `voice` (TTS), `size` (ImageGen), and + `max_tokens` (also used for the `max_cost_usd` estimate). hint_next: type: string request_id: diff --git a/mofa-engine-core/src/backends.rs b/mofa-engine-core/src/backends.rs index 4af3f83..b68939b 100644 --- a/mofa-engine-core/src/backends.rs +++ b/mofa-engine-core/src/backends.rs @@ -2,12 +2,14 @@ pub mod liter_llm; pub mod local_asr; +pub mod local_image_gen; pub mod local_tts; pub mod ollama; pub mod openai_compat; pub use liter_llm::LiterLLMProvider; pub use local_asr::LocalAsrProvider; +pub use local_image_gen::LocalImageGenProvider; pub use local_tts::LocalTtsProvider; pub use ollama::OllamaProvider; pub use openai_compat::OpenAiCompatProvider; diff --git a/mofa-engine-core/src/backends/local_asr.rs b/mofa-engine-core/src/backends/local_asr.rs index 7da8568..8d0f34e 100644 --- a/mofa-engine-core/src/backends/local_asr.rs +++ b/mofa-engine-core/src/backends/local_asr.rs @@ -21,6 +21,11 @@ //! command's stdout is taken as the transcript, //! - `{language}` — language hint from `params.language` (default `auto`). //! +//! When a request sets `params.diarize = true` (S1 speaker separation) the +//! configured `diarize_args` are appended (with the same placeholder +//! substitution); if none are configured the request degrades to a plain +//! transcript rather than failing. +//! //! Child processes are spawned with `kill_on_drop(true)`, so an inference //! timeout that drops the future terminates the process rather than leaking it. @@ -44,6 +49,11 @@ pub struct LocalAsrProvider { command: String, /// Argument template with `{input}`, `{output}`, and `{language}` placeholders. args: Vec, + /// Extra arguments appended (with the same placeholder substitution) when a + /// request asks for speaker diarization via `params.diarize = true`. Empty + /// when the configured CLI does not do speaker separation, in which case a + /// diarize request degrades to a plain transcript. + diarize_args: Vec, /// Directory for transient transcript files. output_dir: PathBuf, /// Configured models this backend serves. @@ -63,6 +73,7 @@ impl LocalAsrProvider { name: name.into(), command: command.into(), args, + diarize_args: Vec::new(), output_dir: output_dir .filter(|s| !s.is_empty()) .map(PathBuf::from) @@ -71,6 +82,39 @@ impl LocalAsrProvider { } } + /// Set the arguments appended when a request enables speaker diarization + /// (`params.diarize = true`). These are placeholder-substituted like the base + /// `args`, so a CLI flag such as `--speaker-diarize` (or a value derived from + /// `{input}`/`{output}`) can be supplied. + pub fn with_diarize_args(mut self, diarize_args: Vec) -> Self { + self.diarize_args = diarize_args; + self + } + + /// Build the resolved argument vector for one transcription, substituting the + /// `{input}`/`{output}`/`{language}` placeholders and, when `diarize` is set, + /// appending the (also-substituted) diarization arguments. Pure over the + /// provider's configured templates, so placeholder + diarize wiring is + /// unit-testable without spawning a process. + fn resolve_args( + &self, + input: &str, + output: &str, + language: &str, + diarize: bool, + ) -> Vec { + let subst = |arg: &str| { + arg.replace("{input}", input) + .replace("{output}", output) + .replace("{language}", language) + }; + let mut resolved: Vec = self.args.iter().map(|a| subst(a)).collect(); + if diarize { + resolved.extend(self.diarize_args.iter().map(|a| subst(a))); + } + resolved + } + /// 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 { @@ -109,9 +153,27 @@ impl LocalAsrProvider { .and_then(|v| v.as_str()) .unwrap_or("auto"); + // Speaker diarization (S1): requested per-call, honored only when the + // backend has diarize arguments configured (otherwise transcribe plainly). + let diarize = request + .params + .get("diarize") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false); + if diarize && self.diarize_args.is_empty() { + tracing::debug!( + provider = %self.name, + "diarization requested but no diarize_args configured; transcribing without speaker separation" + ); + } + // Materialize a transcript output path only if the template references it; // otherwise the command's stdout is the transcript. - let uses_output = self.args.iter().any(|a| a.contains("{output}")); + let uses_output = self + .args + .iter() + .chain(self.diarize_args.iter()) + .any(|a| a.contains("{output}")); let output_file = uses_output.then(|| { self.output_dir .join(format!("mofa_asr_{}.txt", uuid::Uuid::new_v4())) @@ -121,15 +183,7 @@ impl LocalAsrProvider { .map(|p| p.to_string_lossy().to_string()) .unwrap_or_default(); - let resolved_args: Vec = self - .args - .iter() - .map(|arg| { - arg.replace("{input}", input_file) - .replace("{output}", &output_str) - .replace("{language}", language) - }) - .collect(); + let resolved_args = self.resolve_args(input_file, &output_str, language, diarize); let result = Command::new(&self.command) .args(&resolved_args) @@ -453,6 +507,58 @@ mod tests { assert!(matches!(err, EngineError::ProviderError { .. })); } + #[test] + fn resolve_args_appends_diarize_args_only_when_requested() { + let p = LocalAsrProvider::new( + "local-asr", + "asr", + vec!["{input}".into(), "--lang".into(), "{language}".into()], + None, + asr_models(), + ) + .with_diarize_args(vec!["--diarize".into(), "--out".into(), "{output}".into()]); + + // Without diarization: base args only. + let plain = p.resolve_args("clip.wav", "out.txt", "en", false); + assert_eq!(plain, vec!["clip.wav", "--lang", "en"]); + + // With diarization: the diarize args are appended, placeholders substituted. + let diarized = p.resolve_args("clip.wav", "out.txt", "en", true); + assert_eq!( + diarized, + vec!["clip.wav", "--lang", "en", "--diarize", "--out", "out.txt"] + ); + } + + #[tokio::test] + async fn diarize_flag_from_params_reaches_the_command() { + let dir = tempfile::tempdir().unwrap(); + let audio = dir.path().join("clip.wav"); + std::fs::write(&audio, b"fake-audio").unwrap(); + // The fixture echoes its first argument, so the transcript reveals whether + // the diarize marker was passed through from `params.diarize`. + let p = LocalAsrProvider::new( + "local-asr", + "sh", + vec!["-c".into(), "echo \"$1\"".into(), "sh".into()], + Some(dir.path().to_string_lossy().to_string()), + asr_models(), + ) + .with_diarize_args(vec!["SPEAKERS".into()]); + + // diarize = true → the marker is appended and echoed back. + let mut req = asr_request(Some(audio.to_str().unwrap())); + req.params = serde_json::json!({ "diarize": true }); + let resp = p.invoke("local-asr/fixture", &req).await.unwrap(); + assert_eq!(resp.text.as_deref(), Some("SPEAKERS")); + + // diarize omitted → no marker (empty first arg). + let plain = asr_request(Some(audio.to_str().unwrap())); + let resp = p.invoke("local-asr/fixture", &plain).await.unwrap_err(); + // No first arg → empty stdout → treated as an empty transcript. + assert!(matches!(resp, EngineError::ProviderError { .. })); + } + #[tokio::test] async fn rejects_unsupported_capability() { let p = LocalAsrProvider::new("local-asr", "sh", vec![], None, asr_models()); diff --git a/mofa-engine-core/src/backends/local_image_gen.rs b/mofa-engine-core/src/backends/local_image_gen.rs new file mode 100644 index 0000000..ce5853c --- /dev/null +++ b/mofa-engine-core/src/backends/local_image_gen.rs @@ -0,0 +1,518 @@ +//! Local image-generation process-adapter backend. +//! +//! Runs a configured local command — a Stable Diffusion CLI (e.g. `sd`/sd.cpp, +//! `mflux`, or a wrapper script) — to render an image from a text prompt, +//! returning the produced file as a managed artifact. This is the flagship S4 +//! enabler for **fully offline** operation: with a local SD backend, an explainer +//! video's scene images can be generated on-device (`prefer=local`, cost ≈ $0) +//! rather than falling back to a cloud image API. +//! +//! Device- and runtime-specific concerns stay behind this `Provider` boundary, +//! so the engine treats a local image-gen model like any other backend: it is +//! discovered, memory-managed, warmed, idle-unloaded, and can fail over to a +//! cloud image backend (via liter-llm) when the local one is unavailable. +//! +//! ## Command contract +//! +//! The command is spawned once per image (a cold, stateless process). Argument +//! templates may contain: +//! - `{prompt}` — the text prompt (required), +//! - `{output}` — path the command must write the image to (required), +//! - `{negative_prompt}` — from `params.negative_prompt` (default empty), +//! - `{width}` / `{height}` — parsed from `params.size` (e.g. `"1024x1024"`, +//! the default), or `params.width` / `params.height`. +//! +//! Child processes are spawned with `kill_on_drop(true)`, so an inference timeout +//! that drops the future terminates the render rather than leaking it. + +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; + +/// Default output image dimensions when a request does not specify a size. +const DEFAULT_EDGE: u32 = 1024; + +/// A process-adapter provider that shells out to a local image-generation command. +pub struct LocalImageGenProvider { + /// Display name. + name: String, + /// Program to execute per image. + command: String, + /// Argument template with `{prompt}`, `{output}`, `{negative_prompt}`, + /// `{width}`, and `{height}` placeholders. + args: Vec, + /// Output image extension/container (e.g. `png`, `jpg`). + output_format: String, + /// Directory for generated artifacts. + output_dir: PathBuf, + /// Configured models this backend serves. + models: Vec, +} + +impl LocalImageGenProvider { + /// Create a new local image-generation 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(|| "png".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) + }) + } + + /// Resolve the argument template for one render, substituting the prompt, + /// output path, negative prompt, and dimensions. Pure over the configured + /// template, so placeholder wiring is unit-testable without spawning. + fn resolve_args( + &self, + prompt: &str, + output: &str, + negative_prompt: &str, + width: u32, + height: u32, + ) -> Vec { + self.args + .iter() + .map(|arg| { + arg.replace("{prompt}", prompt) + .replace("{output}", output) + .replace("{negative_prompt}", negative_prompt) + .replace("{width}", &width.to_string()) + .replace("{height}", &height.to_string()) + }) + .collect() + } + + async fn generate( + &self, + model_name: &str, + request: &InferenceRequest, + start: std::time::Instant, + ) -> Result { + let prompt = request + .messages + .first() + .map(|m| m.content.clone()) + .or_else(|| { + request + .params + .get("prompt") + .and_then(|v| v.as_str()) + .map(String::from) + }) + .filter(|t| !t.trim().is_empty()) + .ok_or_else(|| EngineError::InvalidRequest("image_gen requires a prompt".into()))?; + + let negative_prompt = request + .params + .get("negative_prompt") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let (width, height) = image_dimensions(&request.params); + + let output_path = self.output_dir.join(format!( + "mofa_img_{}.{}", + uuid::Uuid::new_v4(), + self.output_format + )); + let output_str = output_path.to_string_lossy().to_string(); + + let resolved_args = self.resolve_args(&prompt, &output_str, negative_prompt, width, height); + + let result = Command::new(&self.command) + .args(&resolved_args) + .kill_on_drop(true) + .output() + .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!( + "image_gen command exited with {}: {}", + output.status, + stderr.trim() + ), + }); + } + + // A successful exit must have produced a non-empty image file. + match tokio::fs::metadata(&output_path).await { + Ok(meta) if meta.len() > 0 => {} + _ => { + let _ = tokio::fs::remove_file(&output_path).await; + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "image_gen command produced no image 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, + ..Default::default() + }) + } +} + +/// Parse the requested image dimensions from a request's `params`, accepting a +/// combined `size` (`"WxH"`) or explicit `width`/`height`, and falling back to a +/// square [`DEFAULT_EDGE`]. Pure so the parsing is unit-testable. +fn image_dimensions(params: &serde_json::Value) -> (u32, u32) { + if let Some((w, h)) = params + .get("size") + .and_then(|v| v.as_str()) + .and_then(parse_size) + { + return (w, h); + } + let as_edge = |key: &str| { + params + .get(key) + .and_then(serde_json::Value::as_u64) + .filter(|v| *v > 0) + .map(|v| v as u32) + }; + ( + as_edge("width").unwrap_or(DEFAULT_EDGE), + as_edge("height").unwrap_or(DEFAULT_EDGE), + ) +} + +/// Parse a `"WIDTHxHEIGHT"` size string into positive dimensions (accepts `x` or +/// `X` as the separator). +fn parse_size(size: &str) -> Option<(u32, u32)> { + let (w, h) = size.split_once(['x', 'X'])?; + let w: u32 = w.trim().parse().ok()?; + let h: u32 = h.trim().parse().ok()?; + (w > 0 && h > 0).then_some((w, h)) +} + +#[async_trait] +impl Provider for LocalImageGenProvider { + fn name(&self) -> &str { + &self.name + } + + fn kind(&self) -> ProviderKind { + ProviderKind::LocalImageGen + } + + 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!("image_gen 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::ImageGen); + + if capability != Capability::ImageGen { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' only supports image_gen, not {capability}", + self.name + ))); + } + if !self.model_supports(model_name, Capability::ImageGen) { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' model '{model_name}' does not support image_gen", + self.name + ))); + } + + self.generate(model_name, request, std::time::Instant::now()) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mofa_kernel::Message; + + fn img_models() -> Vec { + vec![ModelDef { + name: "fixture".into(), + capability: "image_gen".into(), + context_window: None, + memory_mb: Some(2048), + ..Default::default() + }] + } + + fn img_request(prompt: &str) -> InferenceRequest { + InferenceRequest { + capability: Some(Capability::ImageGen), + messages: vec![Message { + role: "user".into(), + content: prompt.into(), + ..Default::default() + }], + request_id: "test".into(), + ..Default::default() + } + } + + #[test] + fn metadata_and_defaults() { + let p = LocalImageGenProvider::new("local-sd", "sh", vec![], None, None, img_models()); + assert_eq!(p.kind(), ProviderKind::LocalImageGen); + assert!(p.kind().is_local()); + assert_eq!(p.output_format, "png"); + } + + #[test] + fn parse_size_and_dimensions() { + assert_eq!(parse_size("1024x768"), Some((1024, 768))); + assert_eq!(parse_size("512X512"), Some((512, 512))); + assert_eq!(parse_size("bad"), None); + assert_eq!(parse_size("0x100"), None); + + // size wins; else width/height; else the square default. + assert_eq!( + image_dimensions(&serde_json::json!({ "size": "640x480" })), + (640, 480) + ); + assert_eq!( + image_dimensions(&serde_json::json!({ "width": 256, "height": 128 })), + (256, 128) + ); + assert_eq!( + image_dimensions(&serde_json::Value::Null), + (DEFAULT_EDGE, DEFAULT_EDGE) + ); + } + + #[test] + fn resolve_args_substitutes_all_placeholders() { + let p = LocalImageGenProvider::new( + "local-sd", + "sd", + vec![ + "--prompt".into(), + "{prompt}".into(), + "--neg".into(), + "{negative_prompt}".into(), + "-W".into(), + "{width}".into(), + "-H".into(), + "{height}".into(), + "-o".into(), + "{output}".into(), + ], + None, + None, + img_models(), + ); + let args = p.resolve_args("a cat", "/tmp/o.png", "blurry", 1024, 768); + assert_eq!( + args, + vec![ + "--prompt", + "a cat", + "--neg", + "blurry", + "-W", + "1024", + "-H", + "768", + "-o", + "/tmp/o.png" + ] + ); + } + + #[tokio::test] + async fn generates_image_artifact_via_fixture_command() { + let dir = tempfile::tempdir().unwrap(); + // A deterministic "SD" fixture: write the prompt + size into the output + // path, proving prompt/size placeholders reached the command. + let p = LocalImageGenProvider::new( + "local-sd", + "sh", + vec![ + "-c".into(), + "printf 'PNG:%s:%sx%s' \"$1\" \"$2\" \"$3\" > \"$4\"".into(), + "sh".into(), + "{prompt}".into(), + "{width}".into(), + "{height}".into(), + "{output}".into(), + ], + Some("png".into()), + Some(dir.path().to_string_lossy().to_string()), + img_models(), + ); + + let mut req = img_request("a neuron"); + req.params = serde_json::json!({ "size": "512x512" }); + let resp = p.invoke("local-sd/fixture", &req).await.unwrap(); + let file = resp.file.expect("an image artifact path"); + assert!(file.ends_with(".png")); + let produced = std::fs::read_to_string(&file).unwrap(); + assert_eq!(produced, "PNG:a neuron:512x512"); + } + + #[tokio::test] + async fn empty_prompt_is_invalid_request() { + let p = LocalImageGenProvider::new("local-sd", "sh", vec![], None, None, img_models()); + let err = p + .invoke("local-sd/fixture", &img_request(" ")) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + } + + #[tokio::test] + async fn command_failure_is_a_retryable_provider_error() { + let dir = tempfile::tempdir().unwrap(); + let p = LocalImageGenProvider::new( + "local-sd", + "sh", + vec!["-c".into(), "echo boom >&2; exit 5".into()], + None, + Some(dir.path().to_string_lossy().to_string()), + img_models(), + ); + let err = p + .invoke("local-sd/fixture", &img_request("x")) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + // Retryable so the engine can fail over to a cloud image backend. + assert!(err.retryable()); + } + + #[tokio::test] + async fn rejects_unsupported_capability() { + let p = LocalImageGenProvider::new("local-sd", "sh", vec![], None, None, img_models()); + let mut req = img_request("x"); + req.capability = Some(Capability::Chat); + let err = p.invoke("local-sd/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 406e7d9..a8225a6 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -6,13 +6,15 @@ use async_trait::async_trait; use mofa_kernel::{ BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelResidency, Provider, - ProviderKind, StreamDelta, StreamSink, canonical_model_id, model_id_name, + ProviderKind, ReasoningEffort, StreamDelta, StreamSink, canonical_model_id, model_id_name, }; use reqwest::Client; use serde::{Deserialize, Serialize}; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::time::{Duration, Instant}; +use crate::config::ModelDef; + /// Provider for a local Ollama instance. pub struct OllamaProvider { /// Display name. @@ -21,6 +23,13 @@ pub struct OllamaProvider { base_url: String, /// HTTP client. client: Client, + /// Config-supplied reasoning-tier overrides, keyed by Ollama model name + /// (e.g. `"deepseek-r1:8b"`). Ollama auto-discovers its models, so these are + /// applied as annotations on the discovered card rather than a replacement + /// enumeration — they let a *local* reasoning model participate in + /// `reasoning.effort` → tier routing (S2), which cloud-configured backends + /// already get from their explicit model list. + reasoning_tiers: HashMap, } impl OllamaProvider { @@ -29,7 +38,15 @@ impl OllamaProvider { /// 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 { + /// + /// `models` are optional per-model annotations; only entries carrying a + /// `reasoning_tier` are retained (as tier-routing overrides). All other model + /// metadata is auto-discovered from Ollama, so the list may be empty. + pub fn new( + name: impl Into, + base_url: impl Into, + models: &[ModelDef], + ) -> Result { let client = Client::builder() .no_proxy() .connect_timeout(Duration::from_secs(5)) @@ -37,10 +54,20 @@ impl OllamaProvider { .build() .map_err(|e| EngineError::Config(format!("failed to build Ollama HTTP client: {e}")))?; + let reasoning_tiers = models + .iter() + .filter_map(|m| { + let tier = m.reasoning_tier.as_deref()?; + let tier = ReasoningEffort::from_str_loose(tier)?; + Some((m.name.clone(), tier)) + }) + .collect(); + Ok(Self { name: name.into(), base_url: base_url.into(), client, + reasoning_tiers, }) } @@ -214,6 +241,9 @@ impl Provider for OllamaProvider { }; card.context_window = 4096; card.memory_estimate_bytes = m.size.unwrap_or(0); + // Config-supplied tier override lets this local model take part in + // `reasoning.effort` → tier routing (S2). + card.reasoning_tier = self.reasoning_tiers.get(&model_name).copied(); card.refresh_status(); Some(card) }) @@ -513,12 +543,58 @@ mod tests { #[test] fn provider_metadata() { - let p = OllamaProvider::new("test-ollama", "http://localhost:11434").unwrap(); + 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)); } + #[test] + fn config_reasoning_tiers_are_parsed_and_filtered() { + // Only models with a recognizable reasoning_tier become overrides; a + // plain chat model and an unparseable tier are dropped, so a card without + // an override simply keeps `reasoning_tier = None`. + let models = vec![ + ModelDef { + name: "deepseek-r1:8b".into(), + capability: "chat".into(), + reasoning_tier: Some("high".into()), + ..Default::default() + }, + ModelDef { + name: "qwen3:4b".into(), + capability: "chat".into(), + reasoning_tier: Some("low".into()), + ..Default::default() + }, + ModelDef { + name: "llama3:8b".into(), + capability: "chat".into(), + reasoning_tier: None, + ..Default::default() + }, + ModelDef { + name: "mystery:1b".into(), + capability: "chat".into(), + reasoning_tier: Some("not-a-tier".into()), + ..Default::default() + }, + ]; + let p = OllamaProvider::new("ollama", "http://localhost:11434", &models).unwrap(); + + assert_eq!( + p.reasoning_tiers.get("deepseek-r1:8b"), + Some(&ReasoningEffort::High) + ); + assert_eq!( + p.reasoning_tiers.get("qwen3:4b"), + Some(&ReasoningEffort::Low) + ); + assert!(!p.reasoning_tiers.contains_key("llama3:8b")); + assert!(!p.reasoning_tiers.contains_key("mystery:1b")); + assert_eq!(p.reasoning_tiers.len(), 2); + } + #[test] fn model_id_parse_accepts_old_and_new_format() { assert_eq!(model_id_name("ollama/llama3:8b"), "llama3:8b"); diff --git a/mofa-engine-core/src/config.rs b/mofa-engine-core/src/config.rs index b32ad1f..86a39e7 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -299,6 +299,11 @@ pub struct ProviderConfig { /// Directory for `local_tts` audio artifacts (default: the system temp dir). #[serde(default)] pub output_dir: Option, + /// Extra arguments a `local_asr` command receives when a request enables + /// speaker diarization (`params.diarize = true`). Placeholder-substituted like + /// `args`. Leave empty when the CLI cannot separate speakers. + #[serde(default)] + pub diarize_args: Vec, /// USD price per 1000 prompt (input) tokens. Used for cost tracking; `0` /// (the default, e.g. local models) means no cost is attributed. #[serde(default)] @@ -323,6 +328,7 @@ impl Default for ProviderConfig { args: Vec::new(), output_format: None, output_dir: None, + diarize_args: Vec::new(), price_input_per_1k: 0.0, price_output_per_1k: 0.0, } @@ -710,6 +716,7 @@ impl ProviderConfig { "openai_compatible" => Ok(ProviderKind::OpenAiCompatible), "local_tts" => Ok(ProviderKind::LocalTts), "local_asr" => Ok(ProviderKind::LocalAsr), + "local_image_gen" => Ok(ProviderKind::LocalImageGen), "liter_llm" => Ok(ProviderKind::LiterLlm), other => Err(EngineError::Config(format!( "unknown provider kind '{}' for provider '{}'", diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index 7858d3d..b61a4ec 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -20,7 +20,8 @@ use tokio::sync::{Mutex as AsyncMutex, Semaphore, broadcast, mpsc}; use tokio::task::{AbortHandle, JoinHandle}; use crate::backends::{ - LiterLLMProvider, LocalAsrProvider, LocalTtsProvider, OllamaProvider, OpenAiCompatProvider, + LiterLLMProvider, LocalAsrProvider, LocalImageGenProvider, LocalTtsProvider, OllamaProvider, + OpenAiCompatProvider, }; use crate::circuit_breaker::{CircuitBreakerConfig, CircuitBreakerRegistry, CircuitState}; use crate::config::{EngineConfig, PreflightConfig, TimeoutConfig}; @@ -148,6 +149,30 @@ fn estimate_cost( Some(prompt / 1000.0 * price_in_per_1k + completion / 1000.0 * price_out_per_1k) } +/// Completion tokens assumed when a request does not cap them via +/// `params.max_tokens` — a deliberately generous default so the pre-flight +/// budget estimate errs toward *over*-counting rather than blowing the ceiling. +const DEFAULT_COMPLETION_TOKEN_ESTIMATE: u32 = 1024; + +/// Rough prompt-token estimate for pre-flight budgeting, using the ~4-chars-per +/// token rule of thumb over the message contents. Approximate by design (the +/// authoritative token count only exists post-call); a floor of 1 avoids a +/// zero estimate for an empty prompt. +fn estimate_prompt_tokens(req: &InferenceRequest) -> u32 { + let chars: usize = req.messages.iter().map(|m| m.content.len()).sum(); + ((chars / 4) as u32).max(1) +} + +/// Completion-token estimate for pre-flight budgeting: the request's +/// `params.max_tokens` cap when present, otherwise [`DEFAULT_COMPLETION_TOKEN_ESTIMATE`]. +fn estimate_completion_tokens(req: &InferenceRequest) -> u32 { + req.params + .get("max_tokens") + .and_then(serde_json::Value::as_u64) + .map(|v| v as u32) + .unwrap_or(DEFAULT_COMPLETION_TOKEN_ESTIMATE) +} + impl Engine { /// Create and initialize a new engine from configuration. /// @@ -179,7 +204,9 @@ 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::Ollama => { + Arc::new(OllamaProvider::new(&pc.name, &pc.base_url, &pc.models)?) + } ProviderKind::OpenAiCompatible => Arc::new(OpenAiCompatProvider::with_output_dir( &pc.name, &pc.base_url, @@ -219,10 +246,29 @@ impl Engine { pc.name )) })?; - Arc::new(LocalAsrProvider::new( + Arc::new( + LocalAsrProvider::new( + &pc.name, + command, + pc.args.clone(), + pc.output_dir.clone(), + pc.models.clone(), + ) + .with_diarize_args(pc.diarize_args.clone()), + ) + } + ProviderKind::LocalImageGen => { + let command = pc.command.clone().ok_or_else(|| { + EngineError::Config(format!( + "provider '{}' (local_image_gen) requires a command", + pc.name + )) + })?; + Arc::new(LocalImageGenProvider::new( &pc.name, command, pc.args.clone(), + pc.output_format.clone(), pc.output_dir.clone(), pc.models.clone(), )) @@ -505,6 +551,24 @@ impl Engine { estimate_cost(price_in, price_out, prompt_tokens, completion_tokens) } + /// Conservative pre-flight cost estimate (USD) for serving `req` on + /// `provider`, used to enforce a `max_cost_usd` ceiling *before* spending. + /// + /// Free/local providers (no configured price) and providers with unknown + /// pricing estimate to `0.0`, so a ceiling never blocks a local model and an + /// unpriced provider is never spuriously skipped. Cloud providers are priced + /// from a rough token estimate (see [`estimate_prompt_tokens`] / + /// [`estimate_completion_tokens`]); the authoritative cost is still computed + /// post-call from real token counts by [`Self::compute_cost`]. + fn estimate_request_cost(&self, provider: &str, req: &InferenceRequest) -> f64 { + self.compute_cost( + provider, + Some(estimate_prompt_tokens(req)), + Some(estimate_completion_tokens(req)), + ) + .unwrap_or(0.0) + } + /// Whether the named provider serves models on-device (local track). Unknown /// providers are treated as cloud so metrics never mislabel spend as local. fn provider_is_local(&self, provider: &str) -> bool { @@ -561,6 +625,28 @@ impl Engine { break; } + // Budget ceiling (PRD §S2/W3): skip a candidate whose estimated cost + // would exceed the request's ceiling, recording it in the chain so an + // Agent can see it was priced out (and could raise the budget or fall + // back to a local model). Free/local candidates estimate to $0. + if let Some(ceiling) = req.max_cost_usd { + let est = self.estimate_request_cost(&provider_name, &req); + if est > ceiling { + tracing::info!( + "candidate '{model_id}' skipped: est cost ${est:.4} exceeds ceiling ${ceiling:.4}" + ); + chain.push(FailedAttempt { + provider: provider_name.clone(), + model: model_id_name(&model_id).to_string(), + reason: format!("estimated cost ${est:.4} exceeds budget ${ceiling:.4}"), + }); + last_err = Some(EngineError::NoCapableModel(format!( + "all candidates exceed cost ceiling ${ceiling:.4}" + ))); + continue; + } + } + let max_concurrency = decision.model.execution.max_concurrency; match self @@ -1974,6 +2060,15 @@ mod tests { providers: Vec>, budget_mb: u64, idle_secs: u64, + ) -> Arc { + build_engine_with_pricing(providers, budget_mb, idle_secs, HashMap::new()) + } + + fn build_engine_with_pricing( + providers: Vec>, + budget_mb: u64, + idle_secs: u64, + pricing: HashMap, ) -> Arc { let (event_tx, _) = broadcast::channel(256); let registered = providers @@ -2012,7 +2107,7 @@ mod tests { weak_self: OnceLock::new(), event_tx, metrics: EngineMetrics::default(), - pricing: HashMap::new(), + pricing, input_roots: Vec::new(), started_at: Instant::now(), }); @@ -2294,6 +2389,138 @@ mod tests { assert!(info.routing_reason.is_some()); } + #[test] + fn token_estimates_from_request() { + // ~4 chars/token over message content; floor of 1 for an empty prompt. + let mut req = chat_request(None); // content = "hello" (5 chars) + assert_eq!(estimate_prompt_tokens(&req), 1); + req.messages[0].content = "x".repeat(400); + assert_eq!(estimate_prompt_tokens(&req), 100); + // Completion defaults to the generous estimate, or the params cap. + assert_eq!( + estimate_completion_tokens(&req), + DEFAULT_COMPLETION_TOKEN_ESTIMATE + ); + req.params = serde_json::json!({ "max_tokens": 64 }); + assert_eq!(estimate_completion_tokens(&req), 64); + } + + #[tokio::test] + async fn budget_ceiling_prices_out_cloud_and_prefers_local() { + // A free local model and a priced cloud model both serve chat. A request + // whose per-call budget is below the cloud estimate must route to the + // free local model instead of the (pricier, higher-priority) cloud one. + 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", + Capability::Chat, + ModelResidency::Remote, + 0, + 32, + ), + ); + // openai: $0.01/1k in, $0.03/1k out → default estimate (1 + 1024 tokens) + // ≈ $0.0307, well above the $0.001 ceiling. ollama is free. + let pricing = HashMap::from([("openai".to_string(), (0.01, 0.03))]); + let engine = build_engine_with_pricing(vec![local, cloud], 1000, 0, pricing); + engine.refresh_resources().await; + + let mut req = chat_request(None); + req.max_cost_usd = Some(0.001); + let resp = engine.invoke(req).await.unwrap(); + assert_eq!(resp.provider, "ollama"); + } + + #[tokio::test] + async fn budget_ceiling_exhausts_when_all_priced_out() { + // Two priced cloud models, no local option; a sub-cent ceiling prices both + // out, so the request fails over the whole plan and the budget reason is + // recorded per candidate — the Agent sees *why* nothing ran, rather than + // the engine silently overspending. + let cloud_a = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible).with_model( + "gpt", + Capability::Chat, + ModelResidency::Remote, + 0, + 32, + ), + ); + let cloud_b = Arc::new( + TestProvider::new("deepseek", ProviderKind::OpenAiCompatible).with_model( + "r1", + Capability::Chat, + ModelResidency::Remote, + 0, + 32, + ), + ); + let pricing = HashMap::from([ + ("openai".to_string(), (0.01, 0.03)), + ("deepseek".to_string(), (0.005, 0.02)), + ]); + let engine = build_engine_with_pricing(vec![cloud_a, cloud_b], 1000, 0, pricing); + engine.refresh_resources().await; + + let mut req = chat_request(None); + req.max_cost_usd = Some(0.0001); + let err = engine.invoke(req).await.unwrap_err(); + let info = err.info(); + assert_eq!(info.code, mofa_kernel::ErrorCode::NoCapableModel); + assert_eq!(info.failed_chain.len(), 2); + assert!( + info.failed_chain + .iter() + .all(|a| a.reason.contains("exceeds budget")) + ); + + // Without the ceiling the same request succeeds on a cloud model. + let resp = engine.invoke(chat_request(None)).await.unwrap(); + assert!(resp.provider == "openai" || resp.provider == "deepseek"); + } + + #[tokio::test] + async fn prefer_local_routes_image_gen_to_local_backend() { + // The offline flagship (S4): with a local image-gen backend present, a + // `prefer=local` ImageGen request must route on-device rather than to the + // higher-priority cloud image API — verifying the new `LocalImageGen` kind + // participates in the local-first hard filter. + let local = Arc::new( + TestProvider::new("local-sd", ProviderKind::LocalImageGen).with_model( + "sd", + Capability::ImageGen, + ModelResidency::Loaded, + 10 * MB, + 1, + ), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible).with_model( + "dall-e", + Capability::ImageGen, + ModelResidency::Remote, + 0, + 32, + ), + ); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + let mut req = request(Capability::ImageGen, None, FallbackPolicy::default()); + req.prefer = mofa_kernel::Prefer::Local; + let resp = engine.invoke(req).await.unwrap(); + assert_eq!(resp.provider, "local-sd"); + } + #[tokio::test] async fn local_tts_failure_falls_back_to_cloud_tts() { // A local TTS backend (preferred by locality) crashes mid-synthesis; the diff --git a/mofa-engine-core/src/lib.rs b/mofa-engine-core/src/lib.rs index 9943486..13cbdf2 100644 --- a/mofa-engine-core/src/lib.rs +++ b/mofa-engine-core/src/lib.rs @@ -12,6 +12,7 @@ pub mod engine; pub mod memory; pub mod metrics; pub mod preflight; +pub mod quality_gate; pub mod router; pub mod subscription; diff --git a/mofa-engine-core/src/quality_gate.rs b/mofa-engine-core/src/quality_gate.rs new file mode 100644 index 0000000..dcf28d4 --- /dev/null +++ b/mofa-engine-core/src/quality_gate.rs @@ -0,0 +1,501 @@ +//! Post-production video quality gate (S4 flagship — "no gate, no output"). +//! +//! The flagship Explainer Video scenario produces a finished `mp4` by composing +//! generated scenes, narration, and captions. Before that artifact is accepted, +//! it must clear at least one **hard gate** so a broken render (empty/truncated +//! file, a still slideshow instead of an animation, or a video that does not +//! match its script) is never shipped as a "finished" product. +//! +//! This gate runs three checks: +//! 1. **Duration / integrity** — `ffprobe` confirms a real video stream with a +//! non-trivial duration and non-zero resolution (catches a failed compose). +//! 2. **Slideshow-risk** — `ffmpeg` downsamples the video to a handful of tiny +//! grayscale frames; the mean pixel change between consecutive frames is the +//! *motion*. A static slideshow shows ~zero motion (high risk); an animation +//! shows real motion (low risk). See [`slideshow_risk_from_frames`]. +//! 3. **VLM match** *(optional)* — a semantic image/text verdict supplied by the +//! caller (who runs the engine's `Vlm` capability over sampled frames). This +//! gate stays free of a circular engine dependency by taking the verdict as +//! an input rather than calling back into the engine. +//! +//! The gate **fails closed**: if `ffprobe`/`ffmpeg` are unavailable or a signal +//! cannot be measured, the corresponding check fails rather than passing on +//! faith — an un-certifiable video is not a finished video. Both external tools +//! are widely available (`brew install ffmpeg`); MoFA shells out to them rather +//! than re-implementing media decoding (PRD: "reuse, not self-developed"). +//! +//! # Example +//! ```no_run +//! # use mofa_engine_core::quality_gate::QualityGate; +//! # async fn run() { +//! let gate = QualityGate::new(); +//! let report = gate.check("final.mp4".as_ref(), Some(true)).await; +//! assert!(report.passed, "checks: {:?}", report.checks); +//! # } +//! ``` + +use std::path::Path; + +use serde::Serialize; +use tokio::process::Command; + +/// Edge length of the downsampled grayscale frames used for motion analysis. +const FRAME_EDGE: usize = 8; +/// Pixels per sampled frame (`FRAME_EDGE²`), i.e. bytes per frame in the raw +/// grayscale stream ffmpeg emits. +const FRAME_PIXELS: usize = FRAME_EDGE * FRAME_EDGE; +/// Number of frames to sample across the whole video for motion analysis. +const DEFAULT_SAMPLE_FRAMES: usize = 8; +/// Mean per-pixel change (0..1) treated as "clearly moving". Motion at or above +/// this maps to zero slideshow-risk; below it scales linearly toward full risk. +/// Deliberately small: even a real animation changes only a fraction of an 8×8 +/// frame between samples, so a modest change already rules out a still slideshow. +const MOTION_FULL_SCALE: f64 = 0.06; + +/// Pass/fail thresholds for the gate. +#[derive(Debug, Clone)] +pub struct QualityThresholds { + /// Minimum acceptable video duration in seconds. + pub min_duration_secs: f64, + /// Maximum acceptable slideshow-risk (0 = full motion, 1 = static). + pub max_slideshow_risk: f64, +} + +impl Default for QualityThresholds { + fn default() -> Self { + Self { + min_duration_secs: 1.0, + max_slideshow_risk: 0.85, + } + } +} + +/// One named check within a [`QualityReport`]. +#[derive(Debug, Clone, Serialize)] +pub struct QualityCheck { + /// Stable identifier (`duration`, `slideshow_risk`, `vlm_match`). + pub name: &'static str, + /// Whether this check passed. + pub passed: bool, + /// Human-readable explanation (measured value vs threshold, or why it could + /// not be evaluated). + pub detail: String, +} + +/// The result of running the gate over one video. +#[derive(Debug, Clone, Serialize)] +pub struct QualityReport { + /// Overall verdict: true only if every check passed. + pub passed: bool, + /// Measured duration in seconds, when it could be probed. + pub duration_secs: Option, + /// Measured slideshow-risk (0..1), when frames could be sampled. + pub slideshow_risk: Option, + /// Per-check breakdown, in evaluation order. + pub checks: Vec, +} + +/// Probed technical facts about a video file. +#[derive(Debug, Clone, Default, PartialEq)] +struct ProbeInfo { + /// Container/stream duration in seconds. + duration_secs: Option, + /// Whether at least one video stream with a non-zero resolution is present. + has_video: bool, +} + +/// Validates a finished video against technical + semantic quality checks by +/// shelling out to `ffprobe` and `ffmpeg`. +pub struct QualityGate { + /// `ffprobe` program (resolved on `PATH`). + ffprobe: String, + /// `ffmpeg` program (resolved on `PATH`). + ffmpeg: String, + /// Frames sampled across the video for motion analysis. + sample_frames: usize, + /// Pass/fail thresholds. + thresholds: QualityThresholds, +} + +impl Default for QualityGate { + fn default() -> Self { + Self::new() + } +} + +impl QualityGate { + /// A gate with default tools (`ffprobe`/`ffmpeg`) and thresholds. + pub fn new() -> Self { + Self { + ffprobe: "ffprobe".into(), + ffmpeg: "ffmpeg".into(), + sample_frames: DEFAULT_SAMPLE_FRAMES, + thresholds: QualityThresholds::default(), + } + } + + /// Override the pass/fail thresholds. + pub fn with_thresholds(mut self, thresholds: QualityThresholds) -> Self { + self.thresholds = thresholds; + self + } + + /// Run the gate over `video`, optionally folding in a caller-supplied VLM + /// image/text match verdict (`vlm_ok`). Never panics or errors: a tool or + /// signal that cannot be evaluated becomes a *failed* check (fail-closed), so + /// the returned [`QualityReport`] always reflects a safe verdict. + pub async fn check(&self, video: &Path, vlm_ok: Option) -> QualityReport { + let probe = self.probe(video).await.unwrap_or_default(); + // Only bother sampling frames if the file looks like a real video. + let slideshow_risk = if probe.has_video { + self.sample_slideshow_risk(video, probe.duration_secs).await + } else { + None + }; + self.evaluate(&probe, slideshow_risk, vlm_ok) + } + + /// Combine measured signals into a [`QualityReport`]. Pure: no I/O, so the + /// verdict logic is unit-testable independent of `ffprobe`/`ffmpeg`. + fn evaluate( + &self, + probe: &ProbeInfo, + slideshow_risk: Option, + vlm_ok: Option, + ) -> QualityReport { + let mut checks = Vec::new(); + + // 1. Duration / integrity. + let dur_pass = probe.has_video + && probe + .duration_secs + .is_some_and(|d| d >= self.thresholds.min_duration_secs); + checks.push(QualityCheck { + name: "duration", + passed: dur_pass, + detail: match (probe.has_video, probe.duration_secs) { + (false, _) => "no decodable video stream (ffprobe failed or empty file)".into(), + (true, Some(d)) => { + format!("{d:.2}s (min {:.2}s)", self.thresholds.min_duration_secs) + } + (true, None) => "video stream present but duration unknown".into(), + }, + }); + + // 2. Slideshow-risk. + let slide_pass = slideshow_risk.is_some_and(|r| r <= self.thresholds.max_slideshow_risk); + checks.push(QualityCheck { + name: "slideshow_risk", + passed: slide_pass, + detail: match slideshow_risk { + Some(r) => format!( + "risk {r:.2} (max {:.2}); {}", + self.thresholds.max_slideshow_risk, + if r <= self.thresholds.max_slideshow_risk { + "sufficient motion" + } else { + "looks like a static slideshow" + } + ), + None => "could not sample frames (ffmpeg unavailable or too few frames)".into(), + }, + }); + + // 3. VLM semantic match (only when the caller supplied a verdict). + if let Some(ok) = vlm_ok { + checks.push(QualityCheck { + name: "vlm_match", + passed: ok, + detail: if ok { + "VLM judged frames consistent with the script".into() + } else { + "VLM judged frames inconsistent with the script".into() + }, + }); + } + + QualityReport { + passed: checks.iter().all(|c| c.passed), + duration_secs: probe.duration_secs, + slideshow_risk, + checks, + } + } + + /// Probe duration + stream presence via a single `ffprobe` JSON call. + async fn probe(&self, video: &Path) -> Option { + let output = Command::new(&self.ffprobe) + .args([ + "-v", + "error", + "-show_entries", + "format=duration:stream=codec_type,width,height", + "-of", + "json", + ]) + .arg(video) + .kill_on_drop(true) + .output() + .await + .ok()?; + if !output.status.success() { + return None; + } + parse_ffprobe_json(&String::from_utf8_lossy(&output.stdout)) + } + + /// Sample `sample_frames` tiny grayscale frames across the video and reduce + /// them to a slideshow-risk score. Returns `None` when frames cannot be + /// produced (e.g. `ffmpeg` missing) or fewer than two are available. + async fn sample_slideshow_risk(&self, video: &Path, duration: Option) -> Option { + // Sample evenly across the whole clip: fps = frames / duration. Fall back + // to 1 fps when duration is unknown. + let fps = match duration { + Some(d) if d > 0.0 => (self.sample_frames as f64 / d).max(0.01), + _ => 1.0, + }; + let vf = format!("fps={fps:.4},scale={FRAME_EDGE}:{FRAME_EDGE},format=gray"); + let output = Command::new(&self.ffmpeg) + .args(["-v", "error", "-i"]) + .arg(video) + .args(["-vf", &vf, "-f", "rawvideo", "-pix_fmt", "gray", "-"]) + .kill_on_drop(true) + .output() + .await + .ok()?; + if !output.status.success() { + return None; + } + slideshow_risk_from_frames(&output.stdout, FRAME_PIXELS) + } +} + +/// Parse the `ffprobe -of json` payload into a [`ProbeInfo`]. Extracted as a pure +/// function so the field-plucking is unit-testable without invoking `ffprobe`. +fn parse_ffprobe_json(json: &str) -> Option { + let root: serde_json::Value = serde_json::from_str(json).ok()?; + + let duration_secs = root + .get("format") + .and_then(|f| f.get("duration")) + .and_then(|d| d.as_str()) + .and_then(|s| s.trim().parse::().ok()) + .filter(|d| d.is_finite() && *d > 0.0); + + let has_video = root + .get("streams") + .and_then(|s| s.as_array()) + .is_some_and(|streams| { + streams.iter().any(|st| { + st.get("codec_type").and_then(|c| c.as_str()) == Some("video") + && st + .get("width") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) + > 0 + && st + .get("height") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) + > 0 + }) + }); + + Some(ProbeInfo { + duration_secs, + has_video, + }) +} + +/// Reduce a raw grayscale frame stream (frames of `frame_pixels` bytes, back to +/// back) to a slideshow-risk score in `0..=1`. +/// +/// The *motion* between two consecutive frames is the mean absolute per-pixel +/// difference, normalized to `0..1`. Averaged over all consecutive pairs and +/// compared against [`MOTION_FULL_SCALE`], it yields `risk = 1 - motion_ratio`: +/// identical frames (a still slideshow) score risk `1.0`; frames that change by +/// [`MOTION_FULL_SCALE`] or more score risk `0.0`. Returns `None` when fewer than +/// two whole frames are present (motion is undefined). +fn slideshow_risk_from_frames(raw: &[u8], frame_pixels: usize) -> Option { + if frame_pixels == 0 { + return None; + } + let frames: Vec<&[u8]> = raw.chunks_exact(frame_pixels).collect(); + if frames.len() < 2 { + return None; + } + + let mut total_motion = 0.0; + for pair in frames.windows(2) { + let diff: u64 = pair[0] + .iter() + .zip(pair[1].iter()) + .map(|(a, b)| a.abs_diff(*b) as u64) + .sum(); + // Mean absolute difference for this pair, normalized to 0..1. + total_motion += diff as f64 / (frame_pixels as f64 * 255.0); + } + let motion = total_motion / (frames.len() - 1) as f64; + let motion_ratio = (motion / MOTION_FULL_SCALE).clamp(0.0, 1.0); + Some(1.0 - motion_ratio) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parse_ffprobe_json_extracts_duration_and_video_stream() { + let json = r#"{ + "streams": [ + {"codec_type": "audio"}, + {"codec_type": "video", "width": 1024, "height": 576} + ], + "format": {"duration": "12.34"} + }"#; + let info = parse_ffprobe_json(json).unwrap(); + assert_eq!(info.duration_secs, Some(12.34)); + assert!(info.has_video); + } + + #[test] + fn parse_ffprobe_json_rejects_zero_resolution_and_missing_duration() { + // A "video" stream with no resolution is not a real render; N/A duration + // must not parse to a bogus number. + let json = r#"{ + "streams": [{"codec_type": "video", "width": 0, "height": 0}], + "format": {"duration": "N/A"} + }"#; + let info = parse_ffprobe_json(json).unwrap(); + assert_eq!(info.duration_secs, None); + assert!(!info.has_video); + } + + #[test] + fn identical_frames_are_full_slideshow_risk() { + // Two identical (non-uniform) frames: zero motion → maximum risk. + let mut frame = vec![10u8; FRAME_PIXELS]; + frame[0] = 200; // some internal structure, still identical across frames + let mut raw = frame.clone(); + raw.extend_from_slice(&frame); + assert_eq!(slideshow_risk_from_frames(&raw, FRAME_PIXELS), Some(1.0)); + } + + #[test] + fn strongly_changing_frames_are_low_risk() { + // Black frame then white frame: maximal motion → zero risk. + let mut raw = vec![0u8; FRAME_PIXELS]; + raw.extend(std::iter::repeat_n(255u8, FRAME_PIXELS)); + assert_eq!(slideshow_risk_from_frames(&raw, FRAME_PIXELS), Some(0.0)); + } + + #[test] + fn slideshow_risk_needs_two_frames() { + assert_eq!(slideshow_risk_from_frames(&[1, 2, 3], 64), None); + assert_eq!(slideshow_risk_from_frames(&[], 64), None); + assert_eq!(slideshow_risk_from_frames(&[0u8; 128], 0), None); + } + + #[test] + fn small_motion_scales_between_extremes() { + // A change well below MOTION_FULL_SCALE lands near (but below) full risk, + // confirming the score is monotone rather than binary. + let base = vec![100u8; FRAME_PIXELS]; + let mut moved = base.clone(); + moved[0] = 130; // one pixel changes by 30/255 over 64 pixels → tiny motion + let mut raw = base.clone(); + raw.extend_from_slice(&moved); + let risk = slideshow_risk_from_frames(&raw, FRAME_PIXELS).unwrap(); + assert!(risk > 0.9 && risk < 1.0, "risk was {risk}"); + } + + /// End-to-end against real `ffmpeg`/`ffprobe`. Ignored by default (the tools + /// are not present on every CI runner); run explicitly with: + /// `cargo test -p mofa-engine-core --lib quality_gate -- --ignored`. + /// It generates its own fixtures, so it needs no checked-in media. + #[tokio::test] + #[ignore = "requires ffmpeg/ffprobe on PATH"] + async fn real_gate_distinguishes_animation_slideshow_and_broken() { + let dir = std::env::temp_dir().join(format!("mofa_qg_{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let animation = dir.join("animation.mp4"); + let slideshow = dir.join("slideshow.mp4"); + let broken = dir.join("broken.mp4"); + + let make = |args: &[&str]| { + std::process::Command::new("ffmpeg") + .args(["-v", "error", "-y"]) + .args(args) + .status() + .unwrap() + .success() + }; + assert!(make(&[ + "-f", + "lavfi", + "-i", + "testsrc=size=320x240:rate=24:duration=4", + "-pix_fmt", + "yuv420p", + animation.to_str().unwrap(), + ])); + assert!(make(&[ + "-f", + "lavfi", + "-i", + "color=c=blue:size=320x240:rate=24:duration=4", + "-pix_fmt", + "yuv420p", + slideshow.to_str().unwrap(), + ])); + std::fs::write(&broken, b"not a video").unwrap(); + + let gate = QualityGate::new(); + + // Animation: real motion → low slideshow-risk → passes the gate. + let anim = gate.check(&animation, Some(true)).await; + assert!(anim.passed, "animation report: {anim:?}"); + assert!(anim.slideshow_risk.unwrap() <= gate.thresholds.max_slideshow_risk); + + // Slideshow: a held still → high slideshow-risk → fails the slideshow check. + let slide = gate.check(&slideshow, Some(true)).await; + assert!(!slide.passed, "slideshow report: {slide:?}"); + assert!(slide.slideshow_risk.unwrap() > gate.thresholds.max_slideshow_risk); + + // Broken file: no decodable video stream → fails closed. + let bad = gate.check(&broken, None).await; + assert!(!bad.passed, "broken report: {bad:?}"); + assert!(!bad.checks[0].passed); + + let _ = std::fs::remove_dir_all(&dir); + } + + #[test] + fn evaluate_passes_only_when_all_checks_pass() { + let gate = QualityGate::new(); + let good = ProbeInfo { + duration_secs: Some(10.0), + has_video: true, + }; + + // All three checks green. + let report = gate.evaluate(&good, Some(0.2), Some(true)); + assert!(report.passed); + assert_eq!(report.checks.len(), 3); + + // A failed VLM verdict fails the whole gate. + let report = gate.evaluate(&good, Some(0.2), Some(false)); + assert!(!report.passed); + + // Excessive slideshow-risk fails the gate. + let report = gate.evaluate(&good, Some(0.95), Some(true)); + assert!(!report.passed); + + // Fail-closed: an un-probeable video (no signals) does not pass, and the + // optional VLM check is simply absent. + let report = gate.evaluate(&ProbeInfo::default(), None, None); + assert!(!report.passed); + assert_eq!(report.checks.len(), 2); + assert!(report.checks.iter().all(|c| !c.passed)); + } +} diff --git a/mofa-engine-sdk/examples/README.md b/mofa-engine-sdk/examples/README.md index aa3a7f9..f90a697 100644 --- a/mofa-engine-sdk/examples/README.md +++ b/mofa-engine-sdk/examples/README.md @@ -8,13 +8,18 @@ gracefully: with no capable backend they print the structured error (including a | Example | Scenario | Engine capabilities exercised | |---|---|---| -| `code_review.rs` | **S2** Code/PR Review | `reasoning.effort` (tier routing) + **streamed thought chain** (`Reasoning` vs `Text` chunks) + token/cost | +| `explainer_video.rs` | **S4** Explainer Video (flagship) | `Chat`→`ImageGen`→`Tts`→`Asr` orchestration + **hard quality gate** (ffprobe duration + slideshow-risk + VLM seam) — "no gate, no output" | +| `code_review.rs` | **S2** Code/PR Review | `reasoning.effort` (tier routing) + **streamed thought chain** (`Reasoning` vs `Text` chunks) + `max_cost_usd` budget ceiling + token/cost | | `doc_ai.rs` | **S3** Document/Screenshot AI | `Vlm` understanding, multimodal `Message.images`, `detail` billing tier | | `meeting_brief.rs` | **S1** Meeting → Minutes + Brief | local-first `Asr → Chat → Tts` pipeline, `prefer=local` + `data_class=confidential`, `hint_next` warmup | ## Running ```bash +# S4 — flagship. Orchestration degrades gracefully; pass a composed mp4 to run +# the real quality gate (needs ffmpeg/ffprobe on PATH). +cargo run -p mofa-engine-sdk --example explainer_video -- ./final.mp4 + # S2 — needs a chat/reasoning backend (Ollama or a cloud key) cargo run -p mofa-engine-sdk --example code_review diff --git a/mofa-engine-sdk/examples/code_review.rs b/mofa-engine-sdk/examples/code_review.rs index 29d7de1..dca64a9 100644 --- a/mofa-engine-sdk/examples/code_review.rs +++ b/mofa-engine-sdk/examples/code_review.rs @@ -62,6 +62,10 @@ async fn main() { effort: ReasoningEffort::High, include: true, }), + // Cost control (S2): cap spend at $0.05/review. A pricier cloud candidate + // is priced out (listed in `failed_chain`) so a cheaper/local reasoning + // model wins; a local distilled R1 estimates to $0 and always fits. + max_cost_usd: Some(0.05), ..Default::default() }; diff --git a/mofa-engine-sdk/examples/explainer_video.rs b/mofa-engine-sdk/examples/explainer_video.rs new file mode 100644 index 0000000..0a7f0de --- /dev/null +++ b/mofa-engine-sdk/examples/explainer_video.rs @@ -0,0 +1,144 @@ +//! Scenario S4 — Explainer Video (flagship). +//! +//! One-sentence topic → a finished `mp4`. MoFA is the **orchestration layer + a +//! thin pipeline runner**: it drives `Chat (+reasoning) → ImageGen → TTS → ASR` +//! and then enforces a **hard quality gate** on the composed video, so a broken +//! render (empty file, a still slideshow, or footage that doesn't match the +//! script) is never shipped as "finished" — "no gate, no output". +//! +//! The composition step itself (scenes + narration + captions → `mp4`) reuses +//! Remotion/FFmpeg and is **not** re-implemented here (PRD scope boundary). Pass +//! an already-composed `mp4` to see the real gate run end-to-end: +//! +//! cargo run -p mofa-engine-sdk --example explainer_video -- ./final.mp4 +//! +//! With no file and no backends the pipeline still runs, degrading gracefully and +//! explaining what each stage would produce. + +use std::path::Path; + +use mofa_engine_core::quality_gate::QualityGate; +use mofa_engine_core::{Engine, EngineConfig}; +use mofa_kernel::{Capability, InferenceRequest, Message, Prefer, Reasoning, ReasoningEffort}; + +const TOPIC: &str = "Make a 45s animation explainer: How neural networks learn"; + +#[tokio::main] +async fn main() { + let final_mp4 = std::env::args().nth(1); + + let config = EngineConfig::load(None); + let engine = match Engine::try_new(config).await { + Ok(e) => e, + Err(e) => { + eprintln!("engine init failed: {e}"); + return; + } + }; + engine.refresh_resources().await; + + println!("== S4 Explainer Video — orchestration + hard quality gate =="); + println!("topic: {TOPIC}\n"); + + // 1. Script (Chat + optional reasoning), warming ImageGen for the next stage. + let script = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![Message { + role: "user".into(), + content: format!("{TOPIC}\nWrite a short scene-by-scene narration script."), + ..Default::default() + }], + reasoning: Some(Reasoning { + effort: ReasoningEffort::Medium, + include: false, + }), + hint_next: Some("image_gen".into()), + ..Default::default() + }; + let script = stage(&engine, script, "Script (chat)").await; + + // 2. Scene images (ImageGen, prefer local SD when available). + let scene_prompt = "flat vector animation frame: a neuron adjusting its weights"; + let image = InferenceRequest { + capability: Some(Capability::ImageGen), + messages: vec![Message { + role: "user".into(), + content: scene_prompt.into(), + ..Default::default() + }], + params: serde_json::json!({ "size": "1024x1024" }), + // Offline flagship: prefer local SD (cost ~= $0), fall back to cloud only + // under `auto`. Here we keep it local-first to showcase on-device gen. + prefer: Prefer::Local, + hint_next: Some("tts".into()), + ..Default::default() + }; + stage(&engine, image, "Scene image (image_gen)").await; + + // 3. Narration (TTS). + let narration = InferenceRequest { + capability: Some(Capability::Tts), + messages: vec![Message { + role: "user".into(), + content: script.clone().unwrap_or_else(|| TOPIC.into()), + ..Default::default() + }], + hint_next: Some("asr".into()), + ..Default::default() + }; + stage(&engine, narration, "Narration (tts)").await; + + // 4. (compose) — external: scenes + narration + captions → final.mp4 via + // Remotion/FFmpeg. Out of engine scope; supply the result as an argument. + println!("\n[compose] scenes + narration + captions → final.mp4 (external: Remotion/FFmpeg)"); + + // 5. Quality gate — the hard gate. Runs against a real mp4 when provided. + println!("\n== quality gate =="); + match final_mp4.as_deref() { + Some(path) => run_gate(Path::new(path)).await, + None => println!( + "no composed mp4 supplied — pass one to run the gate:\n \ + cargo run -p mofa-engine-sdk --example explainer_video -- ./final.mp4" + ), + } +} + +/// Run one orchestration stage, reporting where it was routed or the structured +/// error (the pipeline continues either way — this demo showcases orchestration, +/// not a live render). Returns the text output when present. +async fn stage(engine: &Engine, req: InferenceRequest, label: &str) -> Option { + match engine.invoke(req).await { + Ok(resp) => { + println!( + "[{label}] served by {} in {}ms", + resp.provider, resp.duration_ms + ); + resp.text + } + Err(e) => { + let info = e.info(); + println!("[{label}] skipped — {:?}: {}", info.code, info.message); + None + } + } +} + +/// Run the real quality gate over a composed video and print the verdict. The +/// optional VLM image/text check is omitted here (`None`); a full pipeline would +/// pass the engine's `Vlm` verdict over sampled frames. +async fn run_gate(path: &Path) { + let report = QualityGate::new().check(path, None).await; + for c in &report.checks { + println!( + " [{}] {}: {}", + if c.passed { "PASS" } else { "FAIL" }, + c.name, + c.detail + ); + } + if report.passed { + println!("verdict: ACCEPTED ✅"); + } else { + println!("verdict: REJECTED ❌ (no gate, no output)"); + } +} diff --git a/mofa-engine-sdk/examples/meeting_brief.rs b/mofa-engine-sdk/examples/meeting_brief.rs index 2ce2de8..68dd88b 100644 --- a/mofa-engine-sdk/examples/meeting_brief.rs +++ b/mofa-engine-sdk/examples/meeting_brief.rs @@ -34,12 +34,14 @@ fn main() { println!("== S1 Meeting Brief — confidential, local-first =="); println!("audio: {audio}\n"); - // 1. ASR — long audio → transcript, pinned local (confidential). + // 1. ASR — long audio → transcript, pinned local (confidential), with speaker + // diarization so the minutes can attribute resolutions/todos to people. let asr = InferenceRequest { capability: Some(Capability::Asr), input_file: Some(audio), prefer: Prefer::Local, data_class: DataClass::Confidential, + params: serde_json::json!({ "diarize": true }), hint_next: Some("chat".into()), ..Default::default() }; diff --git a/mofa-kernel/src/types.rs b/mofa-kernel/src/types.rs index c3fdf69..a95655a 100644 --- a/mofa-kernel/src/types.rs +++ b/mofa-kernel/src/types.rs @@ -70,6 +70,8 @@ pub enum ProviderKind { LocalTts, /// Local process-adapter ASR backend (e.g. a FunASR or whisper.cpp CLI). LocalAsr, + /// Local process-adapter image-generation backend (e.g. a Stable Diffusion CLI). + LocalImageGen, /// Multi-vendor cloud gateway via the `liter-llm` crate (143+ providers, /// unified OpenAI-style contract). LiterLlm, @@ -80,7 +82,10 @@ impl ProviderKind { /// 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 | Self::LocalAsr) + matches!( + self, + Self::Ollama | Self::LocalTts | Self::LocalAsr | Self::LocalImageGen + ) } } @@ -476,6 +481,13 @@ pub struct InferenceRequest { /// keeps standard behavior. #[serde(default, skip_serializing_if = "Option::is_none")] pub reasoning: Option, + /// Per-request spend ceiling in USD. A candidate whose *estimated* cost + /// exceeds this is skipped during routing (and recorded in `failed_chain`), + /// so spend stays bounded and cheaper/local models win. Free and local + /// models estimate to `$0` and are always affordable; `0.0` therefore means + /// "free/local only". `None` = no ceiling. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub max_cost_usd: Option, /// Conversation messages. #[serde(default)] pub messages: Vec, @@ -502,6 +514,7 @@ impl Default for InferenceRequest { prefer: Prefer::default(), data_class: DataClass::default(), reasoning: None, + max_cost_usd: None, messages: Vec::new(), input_file: None, params: serde_json::Value::Null, From 1585b4a2ed156bf8c69d7acf96bf83b4f2132e26 Mon Sep 17 00:00:00 2001 From: aayank13 Date: Thu, 30 Jul 2026 16:56:44 +0530 Subject: [PATCH 09/22] Refactor routing and model handling in the MoFA engine - Removed the `route` method from the Router struct, simplifying the routing logic. - Updated tests to use `route_ranked` directly for model selection. - Refactored the explainer video example to utilize the `ExplainerVideo` struct for orchestration stages. - Introduced the `MeetingBrief` struct for handling meeting-related orchestration, encapsulating the run and preview logic. - Moved error serialization logic into the `EmbeddedEngine` implementation for better encapsulation. - Changed the model ID handling to use a dedicated `ModelId` struct for canonical ID operations, improving clarity and maintainability. - Updated the server implementation to encapsulate HTTP server logic within a `Server` struct, enhancing organization and readability. - Refactored middleware functions to be methods of the `AppState` struct, promoting better structure and encapsulation. - Improved error handling and response formatting across various endpoints. --- mofa-engine-app/src/main.rs | 221 +++---- mofa-engine-core/src/artifacts.rs | 5 - mofa-engine-core/src/backends/liter_llm.rs | 188 +++--- mofa-engine-core/src/backends/local_asr.rs | 14 +- .../src/backends/local_image_gen.rs | 99 ++-- mofa-engine-core/src/backends/local_tts.rs | 14 +- mofa-engine-core/src/backends/ollama.rs | 48 +- .../src/backends/openai_compat.rs | 72 +-- mofa-engine-core/src/config.rs | 226 +++---- mofa-engine-core/src/engine.rs | 139 ++--- mofa-engine-core/src/lib.rs | 16 +- mofa-engine-core/src/memory.rs | 9 +- mofa-engine-core/src/metrics.rs | 29 +- mofa-engine-core/src/quality_gate.rs | 164 ++--- mofa-engine-core/src/router.rs | 69 ++- mofa-engine-sdk/examples/explainer_video.rs | 79 +-- mofa-engine-sdk/examples/meeting_brief.rs | 68 ++- mofa-engine-sdk/src/client.rs | 12 +- mofa-engine-sdk/src/lib.rs | 6 +- mofa-engine-sdk/src/server.rs | 561 +++++++++--------- mofa-kernel/src/types.rs | 69 ++- 21 files changed, 1123 insertions(+), 985 deletions(-) diff --git a/mofa-engine-app/src/main.rs b/mofa-engine-app/src/main.rs index 6721fbc..a75dc06 100644 --- a/mofa-engine-app/src/main.rs +++ b/mofa-engine-app/src/main.rs @@ -6,7 +6,7 @@ use clap::{Parser, Subcommand}; use mofa_engine_core::{Engine, EngineConfig}; -use mofa_engine_sdk::{DaemonClient, start_server}; +use mofa_engine_sdk::{DaemonClient, Server}; use mofa_kernel::{Capability, InferenceRequest, Message}; use std::path::PathBuf; use tracing_subscriber::{EnvFilter, prelude::*}; @@ -38,20 +38,6 @@ struct DaemonArgs { 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). @@ -80,112 +66,133 @@ enum Command { }, } -/// 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<()> { - init_tracing(); - let cli = Cli::parse(); + Cli::init_tracing(); + Cli::parse().run().await +} - 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), +impl Cli { + /// Dispatch the parsed command to its handler. Serve/validate run in-process; + /// the remaining subcommands act on a running daemon over the `/v1` API. + async fn run(self) -> anyhow::Result<()> { + let Cli { + config, + port, + command, + } = self; + match command { + None | Some(Command::Serve) => Self::serve(config, port).await, + Some(Command::ValidateConfig) => Self::validate_config(config), + Some(Command::Capabilities(d)) => { + Self::print_json(Self::daemon_client(d).capabilities().await) + } + Some(Command::Status(d)) => Self::print_json(Self::daemon_client(d).status().await), + Some(Command::Refresh(d)) => Self::print_json(Self::daemon_client(d).refresh().await), + Some(Command::Invoke { + daemon, + capability, model, - app_id: None, - session_id: None, - fallback_policy: Default::default(), - messages: vec![Message { - role: "user".into(), - content: text, + text, + }) => { + let request = InferenceRequest { + capability: capability.as_deref().and_then(Capability::from_str_loose), + model, + messages: vec![Message { + role: "user".into(), + content: text, + ..Default::default() + }], + request_id: String::new(), ..Default::default() - }], - input_file: None, - params: serde_json::Value::Null, - hint_next: None, - request_id: String::new(), - ..Default::default() - }; - print_json(daemon_client(daemon).invoke(&request).await) + }; + Self::print_json(Self::daemon_client(daemon).invoke(&request).await) + } } } -} -/// 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; + /// 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(); + } } - let host = config.listen.host.clone(); - let port = config.listen.port; - tracing::info!("MoFA Engine v{} starting", env!("CARGO_PKG_VERSION")); - let engine = Engine::try_new(config).await?; - start_server(engine, &host, port) - .await - .map_err(|e| anyhow::anyhow!("{e}")) -} + /// 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, + } + } -/// 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(()) + /// Run the daemon in-process. + 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")); + let engine = Engine::try_new(config).await?; + Server::start(engine, &host, port) + .await + .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}")), } - Err(e) => Err(anyhow::anyhow!("invalid configuration: {e}")), } -} -/// 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(()) + /// 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}")), } - Err(e) => Err(anyhow::anyhow!("{e}")), } } diff --git a/mofa-engine-core/src/artifacts.rs b/mofa-engine-core/src/artifacts.rs index 84a64ab..01ee5fb 100644 --- a/mofa-engine-core/src/artifacts.rs +++ b/mofa-engine-core/src/artifacts.rs @@ -30,11 +30,6 @@ impl ArtifactSweeper { } } - /// 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 diff --git a/mofa-engine-core/src/backends/liter_llm.rs b/mofa-engine-core/src/backends/liter_llm.rs index 0dd92ae..0b26cd6 100644 --- a/mofa-engine-core/src/backends/liter_llm.rs +++ b/mofa-engine-core/src/backends/liter_llm.rs @@ -8,7 +8,7 @@ //! //! A configured model `name` is the liter-llm model id (e.g. `openai/gpt-4o`, //! `deepseek/deepseek-chat`). MoFA's canonical id becomes `/`, and -//! [`model_id_name`] recovers the `provider/model` string liter-llm expects. +//! [`ModelId::name`] recovers the `provider/model` string liter-llm expects. use async_trait::async_trait; use base64::Engine as _; @@ -19,9 +19,8 @@ use liter_llm::{ }; use mofa_kernel::{ BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, - InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelResidency, Provider, - ProviderKind, Reasoning, ReasoningEffort, StreamDelta, StreamSink, canonical_model_id, - model_id_name, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, Reasoning, ReasoningEffort, StreamDelta, StreamSink, }; use std::time::Instant; use tokio_stream::StreamExt as _; @@ -113,7 +112,7 @@ impl LiterLLMProvider { request .messages .iter() - .map(|m| text_message(&m.role, m.content.clone())) + .map(|m| Self::text_message(&m.role, m.content.clone())) .collect() } @@ -132,7 +131,7 @@ impl LiterLLMProvider { model: model_name.to_string(), messages, max_tokens: request.params.get("max_tokens").and_then(|v| v.as_u64()), - reasoning_effort: to_liter_effort(request.reasoning), + reasoning_effort: Self::to_liter_effort(request.reasoning), ..Default::default() }) } @@ -152,7 +151,7 @@ impl LiterLLMProvider { .map_err(|e| self.provider_error(e))?; let text = resp.choices.first().and_then(|c| c.message.text()); - let (prompt_tokens, completion_tokens, tokens_used) = split_usage(resp.usage); + let (prompt_tokens, completion_tokens, tokens_used) = Self::split_usage(resp.usage); Ok(InferenceResponse { text, @@ -272,7 +271,7 @@ impl LiterLLMProvider { let mut messages = Vec::with_capacity(request.messages.len()); for m in &request.messages { if m.images.is_empty() { - messages.push(text_message(&m.role, m.content.clone())); + messages.push(Self::text_message(&m.role, m.content.clone())); continue; } // A multimodal turn: text first, then each image as a detail-tagged part. @@ -282,7 +281,7 @@ impl LiterLLMProvider { } for image in &m.images { parts.push(ContentPart::image_with_detail( - image_to_url(image).await?, + Self::image_to_url(image).await?, detail.clone(), )); } @@ -296,7 +295,7 @@ impl LiterLLMProvider { model: model_name.to_string(), messages, max_tokens: request.params.get("max_tokens").and_then(|v| v.as_u64()), - reasoning_effort: to_liter_effort(request.reasoning), + reasoning_effort: Self::to_liter_effort(request.reasoning), ..Default::default() }; @@ -307,7 +306,7 @@ impl LiterLLMProvider { .map_err(|e| self.provider_error(e))?; let text = resp.choices.first().and_then(|c| c.message.text()); - let (prompt_tokens, completion_tokens, tokens_used) = split_usage(resp.usage); + let (prompt_tokens, completion_tokens, tokens_used) = Self::split_usage(resp.usage); Ok(InferenceResponse { text, @@ -326,71 +325,78 @@ impl LiterLLMProvider { } } -/// Map a MoFA reasoning request to liter-llm's effort enum. -fn to_liter_effort(reasoning: Option) -> Option { - reasoning.map(|r| match r.effort { - ReasoningEffort::Low => liter_llm::ReasoningEffort::Low, - ReasoningEffort::Medium => liter_llm::ReasoningEffort::Medium, - ReasoningEffort::High => liter_llm::ReasoningEffort::High, - }) -} +/// Pure/IO conversions between MoFA and liter-llm types, grouped as private +/// associated functions rather than free functions. +impl LiterLLMProvider { + /// Map a MoFA reasoning request to liter-llm's effort enum. + fn to_liter_effort(reasoning: Option) -> Option { + reasoning.map(|r| match r.effort { + ReasoningEffort::Low => liter_llm::ReasoningEffort::Low, + ReasoningEffort::Medium => liter_llm::ReasoningEffort::Medium, + ReasoningEffort::High => liter_llm::ReasoningEffort::High, + }) + } -/// Build a plain-text liter-llm message for the given role. -fn text_message(role: &str, content: String) -> LiterMessage { - match role { - "system" | "developer" => LiterMessage::System(SystemMessage { - content: UserContent::Text(content), - ..Default::default() - }), - "assistant" => LiterMessage::Assistant(AssistantMessage { - content: Some(AssistantContent::Text(content)), - ..Default::default() - }), - _ => LiterMessage::User(UserMessage { - content: UserContent::Text(content), - ..Default::default() - }), + /// Build a plain-text liter-llm message for the given role. + fn text_message(role: &str, content: String) -> LiterMessage { + match role { + "system" | "developer" => LiterMessage::System(SystemMessage { + content: UserContent::Text(content), + ..Default::default() + }), + "assistant" => LiterMessage::Assistant(AssistantMessage { + content: Some(AssistantContent::Text(content)), + ..Default::default() + }), + _ => LiterMessage::User(UserMessage { + content: UserContent::Text(content), + ..Default::default() + }), + } } -} -/// Resolve an image reference to a URL liter-llm accepts. HTTP(S) and `data:` -/// URLs pass through; a local file path is read and encoded as a `data:` URL. -async fn image_to_url(image: &str) -> Result { - if image.starts_with("http://") || image.starts_with("https://") || image.starts_with("data:") { - return Ok(image.to_string()); + /// Resolve an image reference to a URL liter-llm accepts. HTTP(S) and `data:` + /// URLs pass through; a local file path is read and encoded as a `data:` URL. + async fn image_to_url(image: &str) -> Result { + if image.starts_with("http://") + || image.starts_with("https://") + || image.starts_with("data:") + { + return Ok(image.to_string()); + } + let bytes = tokio::fs::read(image).await.map_err(|e| { + EngineError::InvalidRequest(format!("cannot read image '{image}': {e}")) + })?; + let mime = match std::path::Path::new(image) + .extension() + .and_then(|e| e.to_str()) + .map(str::to_ascii_lowercase) + .as_deref() + { + Some("png") => "image/png", + Some("webp") => "image/webp", + Some("gif") => "image/gif", + _ => "image/jpeg", + }; + let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); + Ok(format!("data:{mime};base64,{b64}")) } - let bytes = tokio::fs::read(image) - .await - .map_err(|e| EngineError::InvalidRequest(format!("cannot read image '{image}': {e}")))?; - let mime = match std::path::Path::new(image) - .extension() - .and_then(|e| e.to_str()) - .map(str::to_ascii_lowercase) - .as_deref() - { - Some("png") => "image/png", - Some("webp") => "image/webp", - Some("gif") => "image/gif", - _ => "image/jpeg", - }; - let b64 = base64::engine::general_purpose::STANDARD.encode(&bytes); - Ok(format!("data:{mime};base64,{b64}")) -} -/// Split a liter-llm `Usage` into `(prompt, completion, total)` as `Option`, -/// falling back to `prompt + completion` when the vendor omits a total. -fn split_usage(usage: Option) -> (Option, Option, Option) { - match usage { - None => (None, None, None), - Some(u) => { - let prompt = u.prompt_tokens as u32; - let completion = u.completion_tokens as u32; - let total = if u.total_tokens > 0 { - u.total_tokens as u32 - } else { - prompt + completion - }; - (Some(prompt), Some(completion), Some(total)) + /// Split a liter-llm `Usage` into `(prompt, completion, total)` as `Option`, + /// falling back to `prompt + completion` when the vendor omits a total. + fn split_usage(usage: Option) -> (Option, Option, Option) { + match usage { + None => (None, None, None), + Some(u) => { + let prompt = u.prompt_tokens as u32; + let completion = u.completion_tokens as u32; + let total = if u.total_tokens > 0 { + u.total_tokens as u32 + } else { + prompt + completion + }; + (Some(prompt), Some(completion), Some(total)) + } } } } @@ -417,7 +423,7 @@ impl Provider for LiterLLMProvider { let cap = Capability::from_str_loose(&m.capability)?; let mut card = ModelCard::new(self.name.clone(), m.name.clone(), cap, self.cost_tier); - card.id = canonical_model_id(&self.name, &m.name); + card.id = ModelId::canonical(&self.name, &m.name); card.availability = ModelAvailability::Configured; card.residency = ModelResidency::Remote; card.context_window = m.context_window.unwrap_or(4096); @@ -449,7 +455,7 @@ impl Provider for LiterLLMProvider { async fn load(&self, model_id: &str) -> Result { Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_id_name(model_id)), + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), residency: ModelResidency::Remote, memory_bytes: Some(0), changed: false, @@ -458,7 +464,7 @@ impl Provider for LiterLLMProvider { async fn unload(&self, model_id: &str) -> Result { Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_id_name(model_id)), + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), residency: ModelResidency::Remote, memory_bytes: Some(0), changed: false, @@ -470,7 +476,7 @@ impl Provider for LiterLLMProvider { model_id: &str, request: &InferenceRequest, ) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let capability = request.capability.unwrap_or(Capability::Chat); let start = Instant::now(); @@ -515,7 +521,7 @@ impl Provider for LiterLLMProvider { return Ok(response); } - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let req = self.build_chat_request(model_name, request)?; let start = Instant::now(); @@ -550,7 +556,7 @@ impl Provider for LiterLLMProvider { } } - let (prompt_tokens, completion_tokens, tokens_used) = split_usage(usage); + let (prompt_tokens, completion_tokens, tokens_used) = Self::split_usage(usage); Ok(InferenceResponse { text: Some(full), file: None, @@ -614,10 +620,10 @@ mod tests { .unwrap(); let cards = p.discover().await.unwrap(); assert_eq!(cards.len(), 2); - // Canonical id is `/`, and model_id_name recovers the + // Canonical id is `/`, and ModelId::name recovers the // `provider/model` string liter-llm routes on. assert_eq!(cards[0].id, "liter/openai/gpt-4o"); - assert_eq!(model_id_name(&cards[0].id), "openai/gpt-4o"); + assert_eq!(ModelId::name(&cards[0].id), "openai/gpt-4o"); assert_eq!(cards[0].capability, Capability::Chat); assert_eq!(cards[1].capability, Capability::ImageGen); assert_eq!(cards[0].residency, ModelResidency::Remote); @@ -647,14 +653,14 @@ mod tests { #[test] fn usage_split_falls_back_to_sum() { - let (p, c, t) = split_usage(Some(liter_llm::Usage { + let (p, c, t) = LiterLLMProvider::split_usage(Some(liter_llm::Usage { prompt_tokens: 5, completion_tokens: 3, total_tokens: 0, ..Default::default() })); assert_eq!((p, c, t), (Some(5), Some(3), Some(8))); - assert_eq!(split_usage(None), (None, None, None)); + assert_eq!(LiterLLMProvider::split_usage(None), (None, None, None)); } #[test] @@ -667,28 +673,32 @@ mod tests { }) }; assert!(matches!( - to_liter_effort(mk(ReasoningEffort::Low)), + LiterLLMProvider::to_liter_effort(mk(ReasoningEffort::Low)), Some(L::Low) )); assert!(matches!( - to_liter_effort(mk(ReasoningEffort::Medium)), + LiterLLMProvider::to_liter_effort(mk(ReasoningEffort::Medium)), Some(L::Medium) )); assert!(matches!( - to_liter_effort(mk(ReasoningEffort::High)), + LiterLLMProvider::to_liter_effort(mk(ReasoningEffort::High)), Some(L::High) )); - assert!(to_liter_effort(None).is_none()); + assert!(LiterLLMProvider::to_liter_effort(None).is_none()); } #[tokio::test] async fn image_url_and_data_url_pass_through() { assert_eq!( - image_to_url("https://example.com/a.png").await.unwrap(), + LiterLLMProvider::image_to_url("https://example.com/a.png") + .await + .unwrap(), "https://example.com/a.png" ); assert_eq!( - image_to_url("data:image/png;base64,AAAA").await.unwrap(), + LiterLLMProvider::image_to_url("data:image/png;base64,AAAA") + .await + .unwrap(), "data:image/png;base64,AAAA" ); } @@ -698,7 +708,9 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("pic.png"); tokio::fs::write(&path, b"\x89PNG\r\n").await.unwrap(); - let url = image_to_url(path.to_str().unwrap()).await.unwrap(); + let url = LiterLLMProvider::image_to_url(path.to_str().unwrap()) + .await + .unwrap(); assert!(url.starts_with("data:image/png;base64,")); } diff --git a/mofa-engine-core/src/backends/local_asr.rs b/mofa-engine-core/src/backends/local_asr.rs index 8d0f34e..d4fd68b 100644 --- a/mofa-engine-core/src/backends/local_asr.rs +++ b/mofa-engine-core/src/backends/local_asr.rs @@ -34,8 +34,8 @@ 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, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, }; use tokio::process::Command; @@ -282,7 +282,7 @@ impl Provider for LocalAsrProvider { 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.id = ModelId::canonical(&self.name, &m.name); card.availability = ModelAvailability::Configured; card.residency = ModelResidency::Unloaded; card.memory_estimate_bytes = m.memory_mb.unwrap_or(0) * 1024 * 1024; @@ -309,7 +309,7 @@ impl Provider for LocalAsrProvider { detail: format!("ASR command '{}' not found", self.command), }); } - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let estimate = self .models .iter() @@ -317,7 +317,7 @@ impl Provider for LocalAsrProvider { .and_then(|m| m.memory_mb) .map(|mb| mb * 1024 * 1024); Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_name), + model_id: ModelId::canonical(&self.name, model_name), residency: ModelResidency::Loaded, memory_bytes: estimate, changed: true, @@ -326,7 +326,7 @@ impl Provider for LocalAsrProvider { async fn unload(&self, model_id: &str) -> Result { Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_id_name(model_id)), + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), residency: ModelResidency::Unloaded, memory_bytes: Some(0), changed: true, @@ -338,7 +338,7 @@ impl Provider for LocalAsrProvider { model_id: &str, request: &InferenceRequest, ) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let capability = request.capability.unwrap_or(Capability::Asr); if capability != Capability::Asr { diff --git a/mofa-engine-core/src/backends/local_image_gen.rs b/mofa-engine-core/src/backends/local_image_gen.rs index ce5853c..015d01d 100644 --- a/mofa-engine-core/src/backends/local_image_gen.rs +++ b/mofa-engine-core/src/backends/local_image_gen.rs @@ -30,8 +30,8 @@ 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, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, }; use tokio::process::Command; @@ -150,7 +150,7 @@ impl LocalImageGenProvider { .get("negative_prompt") .and_then(|v| v.as_str()) .unwrap_or(""); - let (width, height) = image_dimensions(&request.params); + let (width, height) = Self::image_dimensions(&request.params); let output_path = self.output_dir.join(format!( "mofa_img_{}.{}", @@ -212,37 +212,40 @@ impl LocalImageGenProvider { } } -/// Parse the requested image dimensions from a request's `params`, accepting a -/// combined `size` (`"WxH"`) or explicit `width`/`height`, and falling back to a -/// square [`DEFAULT_EDGE`]. Pure so the parsing is unit-testable. -fn image_dimensions(params: &serde_json::Value) -> (u32, u32) { - if let Some((w, h)) = params - .get("size") - .and_then(|v| v.as_str()) - .and_then(parse_size) - { - return (w, h); +/// Pure image-size parsing, grouped as private associated functions. +impl LocalImageGenProvider { + /// Requested image dimensions from a request's `params`, accepting a combined + /// `size` (`"WxH"`) or explicit `width`/`height`, and falling back to a square + /// [`DEFAULT_EDGE`]. + fn image_dimensions(params: &serde_json::Value) -> (u32, u32) { + if let Some((w, h)) = params + .get("size") + .and_then(|v| v.as_str()) + .and_then(Self::parse_size) + { + return (w, h); + } + let as_edge = |key: &str| { + params + .get(key) + .and_then(serde_json::Value::as_u64) + .filter(|v| *v > 0) + .map(|v| v as u32) + }; + ( + as_edge("width").unwrap_or(DEFAULT_EDGE), + as_edge("height").unwrap_or(DEFAULT_EDGE), + ) } - let as_edge = |key: &str| { - params - .get(key) - .and_then(serde_json::Value::as_u64) - .filter(|v| *v > 0) - .map(|v| v as u32) - }; - ( - as_edge("width").unwrap_or(DEFAULT_EDGE), - as_edge("height").unwrap_or(DEFAULT_EDGE), - ) -} -/// Parse a `"WIDTHxHEIGHT"` size string into positive dimensions (accepts `x` or -/// `X` as the separator). -fn parse_size(size: &str) -> Option<(u32, u32)> { - let (w, h) = size.split_once(['x', 'X'])?; - let w: u32 = w.trim().parse().ok()?; - let h: u32 = h.trim().parse().ok()?; - (w > 0 && h > 0).then_some((w, h)) + /// Parse a `"WIDTHxHEIGHT"` size string into positive dimensions (accepts `x` + /// or `X` as the separator). + fn parse_size(size: &str) -> Option<(u32, u32)> { + let (w, h) = size.split_once(['x', 'X'])?; + let w: u32 = w.trim().parse().ok()?; + let h: u32 = h.trim().parse().ok()?; + (w > 0 && h > 0).then_some((w, h)) + } } #[async_trait] @@ -272,7 +275,7 @@ impl Provider for LocalImageGenProvider { 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.id = ModelId::canonical(&self.name, &m.name); card.availability = ModelAvailability::Configured; card.residency = ModelResidency::Unloaded; card.memory_estimate_bytes = m.memory_mb.unwrap_or(0) * 1024 * 1024; @@ -299,7 +302,7 @@ impl Provider for LocalImageGenProvider { detail: format!("image_gen command '{}' not found", self.command), }); } - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let estimate = self .models .iter() @@ -307,7 +310,7 @@ impl Provider for LocalImageGenProvider { .and_then(|m| m.memory_mb) .map(|mb| mb * 1024 * 1024); Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_name), + model_id: ModelId::canonical(&self.name, model_name), residency: ModelResidency::Loaded, memory_bytes: estimate, changed: true, @@ -316,7 +319,7 @@ impl Provider for LocalImageGenProvider { async fn unload(&self, model_id: &str) -> Result { Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_id_name(model_id)), + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), residency: ModelResidency::Unloaded, memory_bytes: Some(0), changed: true, @@ -328,7 +331,7 @@ impl Provider for LocalImageGenProvider { model_id: &str, request: &InferenceRequest, ) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let capability = request.capability.unwrap_or(Capability::ImageGen); if capability != Capability::ImageGen { @@ -387,22 +390,30 @@ mod tests { #[test] fn parse_size_and_dimensions() { - assert_eq!(parse_size("1024x768"), Some((1024, 768))); - assert_eq!(parse_size("512X512"), Some((512, 512))); - assert_eq!(parse_size("bad"), None); - assert_eq!(parse_size("0x100"), None); + assert_eq!( + LocalImageGenProvider::parse_size("1024x768"), + Some((1024, 768)) + ); + assert_eq!( + LocalImageGenProvider::parse_size("512X512"), + Some((512, 512)) + ); + assert_eq!(LocalImageGenProvider::parse_size("bad"), None); + assert_eq!(LocalImageGenProvider::parse_size("0x100"), None); // size wins; else width/height; else the square default. assert_eq!( - image_dimensions(&serde_json::json!({ "size": "640x480" })), + LocalImageGenProvider::image_dimensions(&serde_json::json!({ "size": "640x480" })), (640, 480) ); assert_eq!( - image_dimensions(&serde_json::json!({ "width": 256, "height": 128 })), + LocalImageGenProvider::image_dimensions( + &serde_json::json!({ "width": 256, "height": 128 }) + ), (256, 128) ); assert_eq!( - image_dimensions(&serde_json::Value::Null), + LocalImageGenProvider::image_dimensions(&serde_json::Value::Null), (DEFAULT_EDGE, DEFAULT_EDGE) ); } diff --git a/mofa-engine-core/src/backends/local_tts.rs b/mofa-engine-core/src/backends/local_tts.rs index d225c8c..510e99c 100644 --- a/mofa-engine-core/src/backends/local_tts.rs +++ b/mofa-engine-core/src/backends/local_tts.rs @@ -29,8 +29,8 @@ 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, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, }; use tokio::process::Command; @@ -237,7 +237,7 @@ impl Provider for LocalTtsProvider { 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.id = ModelId::canonical(&self.name, &m.name); card.availability = ModelAvailability::Configured; card.residency = ModelResidency::Unloaded; card.memory_estimate_bytes = m.memory_mb.unwrap_or(0) * 1024 * 1024; @@ -264,7 +264,7 @@ impl Provider for LocalTtsProvider { detail: format!("TTS command '{}' not found", self.command), }); } - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let estimate = self .models .iter() @@ -272,7 +272,7 @@ impl Provider for LocalTtsProvider { .and_then(|m| m.memory_mb) .map(|mb| mb * 1024 * 1024); Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_name), + model_id: ModelId::canonical(&self.name, model_name), residency: ModelResidency::Loaded, memory_bytes: estimate, changed: true, @@ -281,7 +281,7 @@ impl Provider for LocalTtsProvider { async fn unload(&self, model_id: &str) -> Result { Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_id_name(model_id)), + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), residency: ModelResidency::Unloaded, memory_bytes: Some(0), changed: true, @@ -293,7 +293,7 @@ impl Provider for LocalTtsProvider { model_id: &str, request: &InferenceRequest, ) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let capability = request.capability.unwrap_or(Capability::Tts); if capability != Capability::Tts { diff --git a/mofa-engine-core/src/backends/ollama.rs b/mofa-engine-core/src/backends/ollama.rs index a8225a6..032a845 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -5,8 +5,8 @@ use async_trait::async_trait; use mofa_kernel::{ BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, - InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelResidency, Provider, - ProviderKind, ReasoningEffort, StreamDelta, StreamSink, canonical_model_id, model_id_name, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, ReasoningEffort, StreamDelta, StreamSink, }; use reqwest::Client; use serde::{Deserialize, Serialize}; @@ -159,13 +159,15 @@ struct OllamaStreamChunk { prompt_eval_count: Option, } -/// Parse one NDJSON line from an Ollama stream. Blank/whitespace lines and -/// unparseable fragments yield `None`. -fn parse_stream_line(line: &[u8]) -> Option { - if line.iter().all(u8::is_ascii_whitespace) { - return None; +impl OllamaProvider { + /// Parse one NDJSON line from an Ollama stream. Blank/whitespace lines and + /// unparseable fragments yield `None`. + fn parse_stream_line(line: &[u8]) -> Option { + if line.iter().all(u8::is_ascii_whitespace) { + return None; + } + serde_json::from_slice(line).ok() } - serde_json::from_slice(line).ok() } #[async_trait] @@ -232,7 +234,7 @@ impl Provider for OllamaProvider { Capability::Chat, CostTier::Free, ); - card.id = canonical_model_id(&self.name, &model_name); + card.id = ModelId::canonical(&self.name, &model_name); card.availability = ModelAvailability::Discovered; card.residency = if loaded.contains(&model_name) { ModelResidency::Loaded @@ -266,7 +268,7 @@ impl Provider for OllamaProvider { } async fn load(&self, model_id: &str) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let body = OllamaChatRequest { model: model_name.to_string(), messages: vec![OllamaMessage { @@ -298,7 +300,7 @@ impl Provider for OllamaProvider { } Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_name), + model_id: ModelId::canonical(&self.name, model_name), residency: ModelResidency::Loaded, memory_bytes: None, changed: true, @@ -306,7 +308,7 @@ impl Provider for OllamaProvider { } async fn unload(&self, model_id: &str) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let body = OllamaChatRequest { model: model_name.to_string(), messages: vec![OllamaMessage { @@ -339,7 +341,7 @@ impl Provider for OllamaProvider { } Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_name), + model_id: ModelId::canonical(&self.name, model_name), residency: ModelResidency::Unloaded, memory_bytes: Some(0), changed: true, @@ -351,7 +353,7 @@ impl Provider for OllamaProvider { model_id: &str, request: &InferenceRequest, ) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let messages: Vec = request .messages @@ -433,7 +435,7 @@ impl Provider for OllamaProvider { request: &InferenceRequest, sink: StreamSink, ) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let messages: Vec = request .messages .iter() @@ -502,7 +504,7 @@ impl Provider for OllamaProvider { buf.extend_from_slice(&bytes); while let Some(pos) = buf.iter().position(|&b| b == b'\n') { let line: Vec = buf.drain(..=pos).collect(); - if let Some(chunk) = parse_stream_line(&line) + if let Some(chunk) = Self::parse_stream_line(&line) && let Some(delta) = apply(chunk, &mut full) { let _ = sink.send(StreamDelta::Text(delta)).await; @@ -510,7 +512,7 @@ impl Provider for OllamaProvider { } } // Any trailing bytes without a final newline. - if let Some(chunk) = parse_stream_line(&buf) + if let Some(chunk) = Self::parse_stream_line(&buf) && let Some(delta) = apply(chunk, &mut full) { let _ = sink.send(StreamDelta::Text(delta)).await; @@ -597,8 +599,8 @@ mod tests { #[test] fn model_id_parse_accepts_old_and_new_format() { - assert_eq!(model_id_name("ollama/llama3:8b"), "llama3:8b"); - assert_eq!(model_id_name("ollama::llama3:8b"), "llama3:8b"); + assert_eq!(ModelId::name("ollama/llama3:8b"), "llama3:8b"); + assert_eq!(ModelId::name("ollama::llama3:8b"), "llama3:8b"); } #[test] @@ -619,7 +621,7 @@ mod tests { let mut prompt = None; let mut completion = None; for line in lines { - if let Some(chunk) = parse_stream_line(line) { + if let Some(chunk) = OllamaProvider::parse_stream_line(line) { if chunk.done { prompt = chunk.prompt_eval_count; completion = chunk.eval_count; @@ -642,8 +644,8 @@ mod tests { #[test] fn parse_stream_line_ignores_blank_and_garbage() { - assert!(parse_stream_line(b"").is_none()); - assert!(parse_stream_line(b" \n").is_none()); - assert!(parse_stream_line(b"not json").is_none()); + assert!(OllamaProvider::parse_stream_line(b"").is_none()); + assert!(OllamaProvider::parse_stream_line(b" \n").is_none()); + assert!(OllamaProvider::parse_stream_line(b"not json").is_none()); } } diff --git a/mofa-engine-core/src/backends/openai_compat.rs b/mofa-engine-core/src/backends/openai_compat.rs index 48cbe9d..b0b07a1 100644 --- a/mofa-engine-core/src/backends/openai_compat.rs +++ b/mofa-engine-core/src/backends/openai_compat.rs @@ -5,8 +5,8 @@ use async_trait::async_trait; use mofa_kernel::{ BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, - InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelResidency, Provider, - ProviderKind, ReasoningEffort, StreamDelta, StreamSink, canonical_model_id, model_id_name, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, ReasoningEffort, StreamDelta, StreamSink, }; use reqwest::Client; use serde::{Deserialize, Serialize}; @@ -33,7 +33,9 @@ pub struct OpenAiCompatProvider { } impl OpenAiCompatProvider { - /// Create a new OpenAI-compatible provider. + /// Create a new OpenAI-compatible provider with default (temp-dir) artifact + /// output. Test-support: the engine factory uses [`Self::with_output_dir`]. + #[cfg(test)] pub fn new( name: impl Into, base_url: impl Into, @@ -157,27 +159,29 @@ enum SseEvent { Done, } -/// Parse one SSE event block (the bytes preceding a blank-line delimiter). -/// -/// Concatenates any `data:` field lines, ignoring comments/`event:` lines and -/// tolerating `\r\n`. Blank blocks and unparseable payloads yield `None`. -fn parse_sse_event(block: &[u8]) -> Option { - let text = std::str::from_utf8(block).ok()?; - let mut data = String::new(); - for line in text.lines() { - if let Some(rest) = line.trim_start().strip_prefix("data:") { - data.push_str(rest.trim()); +impl OpenAiCompatProvider { + /// Parse one SSE event block (the bytes preceding a blank-line delimiter). + /// + /// Concatenates any `data:` field lines, ignoring comments/`event:` lines and + /// tolerating `\r\n`. Blank blocks and unparseable payloads yield `None`. + fn parse_sse_event(block: &[u8]) -> Option { + let text = std::str::from_utf8(block).ok()?; + let mut data = String::new(); + for line in text.lines() { + if let Some(rest) = line.trim_start().strip_prefix("data:") { + data.push_str(rest.trim()); + } } + if data.is_empty() { + return None; + } + if data == "[DONE]" { + return Some(SseEvent::Done); + } + serde_json::from_str::(&data) + .ok() + .map(SseEvent::Chunk) } - if data.is_empty() { - return None; - } - if data == "[DONE]" { - return Some(SseEvent::Done); - } - serde_json::from_str::(&data) - .ok() - .map(SseEvent::Chunk) } #[derive(Debug, Serialize)] @@ -210,7 +214,7 @@ impl Provider for OpenAiCompatProvider { let cap = Capability::from_str_loose(&m.capability)?; let mut card = ModelCard::new(self.name.clone(), m.name.clone(), cap, self.cost_tier); - card.id = canonical_model_id(&self.name, &m.name); + card.id = ModelId::canonical(&self.name, &m.name); card.availability = ModelAvailability::Configured; card.residency = ModelResidency::Remote; card.context_window = m.context_window.unwrap_or(4096); @@ -262,7 +266,7 @@ impl Provider for OpenAiCompatProvider { async fn load(&self, model_id: &str) -> Result { Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_id_name(model_id)), + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), residency: ModelResidency::Remote, memory_bytes: Some(0), changed: false, @@ -271,7 +275,7 @@ impl Provider for OpenAiCompatProvider { async fn unload(&self, model_id: &str) -> Result { Ok(LifecycleResult { - model_id: canonical_model_id(&self.name, model_id_name(model_id)), + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), residency: ModelResidency::Remote, memory_bytes: Some(0), changed: false, @@ -283,7 +287,7 @@ impl Provider for OpenAiCompatProvider { model_id: &str, request: &InferenceRequest, ) -> Result { - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let capability = request.capability.unwrap_or(Capability::Chat); let start = std::time::Instant::now(); @@ -331,7 +335,7 @@ impl Provider for OpenAiCompatProvider { return Ok(response); } - let model_name = model_id_name(model_id); + let model_name = ModelId::name(model_id); let messages: Vec = request .messages .iter() @@ -415,7 +419,7 @@ impl Provider for OpenAiCompatProvider { buf.extend_from_slice(&bytes); while let Some(idx) = buf.windows(2).position(|w| w == b"\n\n") { let block: Vec = buf.drain(..idx + 2).collect(); - if let Some(event) = parse_sse_event(&block) + if let Some(event) = Self::parse_sse_event(&block) && let Some(delta) = apply(event, &mut full) { let _ = sink.send(StreamDelta::Text(delta)).await; @@ -423,7 +427,7 @@ impl Provider for OpenAiCompatProvider { } } // Any trailing event without a final blank line. - if let Some(event) = parse_sse_event(&buf) + if let Some(event) = Self::parse_sse_event(&buf) && let Some(delta) = apply(event, &mut full) { let _ = sink.send(StreamDelta::Text(delta)).await; @@ -797,7 +801,7 @@ mod tests { let mut completion = None; let mut total = None; for block in blocks { - match parse_sse_event(block) { + match OpenAiCompatProvider::parse_sse_event(block) { Some(SseEvent::Chunk(chunk)) => { if let Some(u) = chunk.usage { prompt = u.prompt_tokens; @@ -829,11 +833,11 @@ mod tests { #[test] fn parse_sse_event_handles_done_blank_and_garbage() { assert!(matches!( - parse_sse_event(b"data: [DONE]"), + OpenAiCompatProvider::parse_sse_event(b"data: [DONE]"), Some(SseEvent::Done) )); - assert!(parse_sse_event(b"").is_none()); - assert!(parse_sse_event(b": keep-alive comment").is_none()); - assert!(parse_sse_event(b"data: not json").is_none()); + assert!(OpenAiCompatProvider::parse_sse_event(b"").is_none()); + assert!(OpenAiCompatProvider::parse_sse_event(b": keep-alive comment").is_none()); + assert!(OpenAiCompatProvider::parse_sse_event(b"data: not json").is_none()); } } diff --git a/mofa-engine-core/src/config.rs b/mofa-engine-core/src/config.rs index 86a39e7..9fde243 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -44,7 +44,7 @@ pub struct ArtifactConfig { 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")] + #[serde(default = "ArtifactConfig::default_artifact_retention")] pub retention_secs: u64, } @@ -52,13 +52,16 @@ impl Default for ArtifactConfig { fn default() -> Self { Self { dir: None, - retention_secs: default_artifact_retention(), + retention_secs: Self::default_artifact_retention(), } } } -fn default_artifact_retention() -> u64 { - 3600 +impl ArtifactConfig { + /// Retain generated artifacts for one hour by default. + fn default_artifact_retention() -> u64 { + 3600 + } } /// Security and file-access constraints. @@ -75,28 +78,32 @@ pub struct SecurityConfig { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ListenConfig { /// Bind address - #[serde(default = "default_host")] + #[serde(default = "ListenConfig::default_host")] pub host: String, /// Bind port - #[serde(default = "default_port")] + #[serde(default = "ListenConfig::default_port")] pub port: u16, } impl Default for ListenConfig { fn default() -> Self { Self { - host: default_host(), - port: default_port(), + host: Self::default_host(), + port: Self::default_port(), } } } -fn default_host() -> String { - "127.0.0.1".into() -} +impl ListenConfig { + /// Bind loopback by default — the API is exposed to the local machine only + /// unless the operator opts into a wider host. + fn default_host() -> String { + "127.0.0.1".into() + } -fn default_port() -> u16 { - 8420 + fn default_port() -> u16 { + 8420 + } } /// Memory management configuration. @@ -105,7 +112,7 @@ pub struct MemoryConfig { /// Memory budget in megabytes (None = auto-detect from system) pub budget_mb: Option, /// Seconds of idle time before evicting a model - #[serde(default = "default_idle_timeout")] + #[serde(default = "MemoryConfig::default_idle_timeout")] pub idle_timeout_secs: u64, } @@ -113,13 +120,16 @@ impl Default for MemoryConfig { fn default() -> Self { Self { budget_mb: None, - idle_timeout_secs: default_idle_timeout(), + idle_timeout_secs: Self::default_idle_timeout(), } } } -fn default_idle_timeout() -> u64 { - 120 +impl MemoryConfig { + /// Evict a model after two minutes idle by default. + fn default_idle_timeout() -> u64 { + 120 + } } /// Timeouts for the distinct phases of a request, all in seconds. @@ -131,34 +141,34 @@ fn default_idle_timeout() -> u64 { pub struct TimeoutConfig { /// Overall budget for a single `invoke`, spanning queueing, loading, and /// inference across every candidate attempted. - #[serde(default = "default_request_timeout")] + #[serde(default = "TimeoutConfig::default_request_timeout")] pub request_secs: u64, /// Maximum time a request may wait for a concurrency permit before failing. - #[serde(default = "default_queue_timeout")] + #[serde(default = "TimeoutConfig::default_queue_timeout")] pub queue_secs: u64, /// Maximum time to load/warm a single model. - #[serde(default = "default_load_timeout")] + #[serde(default = "TimeoutConfig::default_load_timeout")] pub load_secs: u64, /// Maximum time for a single inference call to a provider. - #[serde(default = "default_inference_timeout")] + #[serde(default = "TimeoutConfig::default_inference_timeout")] pub inference_secs: u64, /// Maximum time for a single provider discovery probe. - #[serde(default = "default_discovery_timeout")] + #[serde(default = "TimeoutConfig::default_discovery_timeout")] pub discovery_secs: u64, /// Maximum time for a single provider health probe. - #[serde(default = "default_health_timeout")] + #[serde(default = "TimeoutConfig::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(), + request_secs: Self::default_request_timeout(), + queue_secs: Self::default_queue_timeout(), + load_secs: Self::default_load_timeout(), + inference_secs: Self::default_inference_timeout(), + discovery_secs: Self::default_discovery_timeout(), + health_secs: Self::default_health_timeout(), } } } @@ -190,72 +200,80 @@ impl TimeoutConfig { } } -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 +impl TimeoutConfig { + 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")] + #[serde(default = "PreflightConfig::default_true")] pub enabled: bool, /// Whether the engine learns capability transition history. - #[serde(default = "default_true")] + #[serde(default = "PreflightConfig::default_true")] pub history_learning: bool, /// Whether the engine speculatively warms models from hints/history/subscriptions. - #[serde(default = "default_true")] + #[serde(default = "PreflightConfig::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")] + #[serde(default = "PreflightConfig::default_confidence_threshold")] pub confidence_threshold: f64, /// Minimum number of observed transitions before history is trusted. - #[serde(default = "default_min_samples")] + #[serde(default = "PreflightConfig::default_min_samples")] pub min_samples: u64, /// Default lifetime of a capability subscription, in seconds. - #[serde(default = "default_subscription_ttl")] + #[serde(default = "PreflightConfig::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(), + enabled: Self::default_true(), + history_learning: Self::default_true(), + speculative_warming: Self::default_true(), + confidence_threshold: Self::default_confidence_threshold(), + min_samples: Self::default_min_samples(), + subscription_ttl_secs: Self::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 +impl PreflightConfig { + /// Preflight and its sub-switches are on by default; the operator opts out. + fn default_true() -> bool { + true + } + /// Require moderately strong evidence before a learned prediction warms a model. + fn default_confidence_threshold() -> f64 { + 0.6 + } + /// Ignore a transition until it has been seen a few times, to avoid warming on noise. + fn default_min_samples() -> u64 { + 3 + } + /// Keep a capability subscription warm for an hour by default. + fn default_subscription_ttl() -> u64 { + 3600 + } } /// Configuration for a single provider. @@ -270,16 +288,16 @@ pub struct ProviderConfig { /// API key (supports `${ENV_VAR}` syntax) pub api_key: Option, /// Routing priority (lower = preferred; 1 = local, 10 = cloud) - #[serde(default = "default_priority")] + #[serde(default = "ProviderConfig::default_priority")] pub priority: u8, /// Cost tier string - #[serde(default = "default_cost_tier")] + #[serde(default = "ProviderConfig::default_cost_tier")] pub cost_tier: String, /// Explicit model definitions #[serde(default)] pub models: Vec, /// Whether this provider is enabled - #[serde(default = "default_enabled")] + #[serde(default = "ProviderConfig::default_enabled")] pub enabled: bool, /// Program to execute for a `local_tts` process-adapter backend. /// @@ -320,10 +338,10 @@ impl Default for ProviderConfig { kind: String::new(), base_url: String::new(), api_key: None, - priority: default_priority(), - cost_tier: default_cost_tier(), + priority: Self::default_priority(), + cost_tier: Self::default_cost_tier(), models: Vec::new(), - enabled: default_enabled(), + enabled: Self::default_enabled(), command: None, args: Vec::new(), output_format: None, @@ -335,16 +353,19 @@ impl Default for ProviderConfig { } } -fn default_priority() -> u8 { - 5 -} - -fn default_cost_tier() -> String { - "medium".into() -} - -fn default_enabled() -> bool { - true +impl ProviderConfig { + /// A middling routing priority, so an unprioritized provider sits between the + /// conventional local (1) and cloud (10) bands. + fn default_priority() -> u8 { + 5 + } + fn default_cost_tier() -> String { + "medium".into() + } + /// Providers are enabled unless explicitly disabled. + fn default_enabled() -> bool { + true + } } /// A model definition within a provider config. @@ -422,7 +443,7 @@ impl EngineConfig { // Resolve env vars in api_key fields for provider in &mut config.providers { if let Some(ref key) = provider.api_key { - provider.api_key = Some(resolve_env_var(key)?); + provider.api_key = Some(ProviderConfig::resolve_env_var(key)?); } } @@ -696,19 +717,21 @@ impl EngineConfig { } } -/// Resolve `${ENV_VAR}` patterns in a string. -fn resolve_env_var(s: &str) -> Result { - if let Some(rest) = s.strip_prefix("${") - && let Some(var_name) = rest.strip_suffix('}') - { - return std::env::var(var_name).map_err(|_| { - EngineError::Config(format!("environment variable '{var_name}' is not set")) - }); +impl ProviderConfig { + /// Resolve a `${ENV_VAR}` reference to the environment variable's value, + /// passing any other string through unchanged. Lets a config keep secrets + /// (e.g. `api_key = "${OPENAI_API_KEY}"`) out of the file itself. + fn resolve_env_var(s: &str) -> Result { + if let Some(rest) = s.strip_prefix("${") + && let Some(var_name) = rest.strip_suffix('}') + { + return std::env::var(var_name).map_err(|_| { + EngineError::Config(format!("environment variable '{var_name}' is not set")) + }); + } + Ok(s.to_string()) } - Ok(s.to_string()) -} -impl ProviderConfig { /// Parse the configured provider kind. pub fn provider_kind(&self) -> Result { match self.kind.as_str() { @@ -732,7 +755,10 @@ mod tests { #[test] fn resolve_env_var_plain() { - assert_eq!(resolve_env_var("plain-key").unwrap(), "plain-key"); + assert_eq!( + ProviderConfig::resolve_env_var("plain-key").unwrap(), + "plain-key" + ); } #[test] @@ -740,7 +766,7 @@ mod tests { // SAFETY: this test runs single-threaded and restores the env var. unsafe { std::env::set_var("TEST_MOFA_KEY", "resolved-value") }; assert_eq!( - resolve_env_var("${TEST_MOFA_KEY}").unwrap(), + ProviderConfig::resolve_env_var("${TEST_MOFA_KEY}").unwrap(), "resolved-value" ); unsafe { std::env::remove_var("TEST_MOFA_KEY") }; @@ -749,7 +775,7 @@ mod tests { #[test] fn unresolved_env_var_errors() { unsafe { std::env::remove_var("MISSING_MOFA_KEY") }; - assert!(resolve_env_var("${MISSING_MOFA_KEY}").is_err()); + assert!(ProviderConfig::resolve_env_var("${MISSING_MOFA_KEY}").is_err()); } #[test] diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index b61a4ec..de92013 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -12,8 +12,8 @@ use std::time::{Duration, Instant}; use dashmap::DashMap; use mofa_kernel::{ BackendHealth, BackendStatus, Capability, CostTier, EngineError, EngineEvent, EngineStatus, - FailedAttempt, FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, ModelResidency, - Provider, ProviderHealth, ProviderKind, StreamChunk, StreamDelta, model_id_name, + FailedAttempt, FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, ModelId, + ModelResidency, Provider, ProviderHealth, ProviderKind, StreamChunk, StreamDelta, }; use serde::{Deserialize, Serialize}; use tokio::sync::{Mutex as AsyncMutex, Semaphore, broadcast, mpsc}; @@ -128,49 +128,54 @@ pub struct Engine { started_at: Instant, } -/// Remaining time until `deadline`, saturating at zero. -fn remaining(deadline: Instant) -> Duration { - deadline.saturating_duration_since(Instant::now()) -} +/// Pure routing/cost helpers, grouped as private associated functions rather than +/// free functions so they share the `Engine` namespace with the logic that uses them. +impl Engine { + /// Completion tokens assumed when a request does not cap them via + /// `params.max_tokens` — a deliberately generous default so the pre-flight + /// budget estimate errs toward *over*-counting rather than blowing the ceiling. + const DEFAULT_COMPLETION_TOKEN_ESTIMATE: u32 = 1024; -/// USD cost from per-1k-token prices. `None` when both prices are zero (free / -/// local models), so no spurious zero cost is attributed. -fn estimate_cost( - price_in_per_1k: f64, - price_out_per_1k: f64, - prompt_tokens: Option, - completion_tokens: Option, -) -> Option { - if price_in_per_1k == 0.0 && price_out_per_1k == 0.0 { - return None; - } - let prompt = prompt_tokens.unwrap_or(0) as f64; - let completion = completion_tokens.unwrap_or(0) as f64; - Some(prompt / 1000.0 * price_in_per_1k + completion / 1000.0 * price_out_per_1k) -} + /// Remaining time until `deadline`, saturating at zero. + fn remaining(deadline: Instant) -> Duration { + deadline.saturating_duration_since(Instant::now()) + } -/// Completion tokens assumed when a request does not cap them via -/// `params.max_tokens` — a deliberately generous default so the pre-flight -/// budget estimate errs toward *over*-counting rather than blowing the ceiling. -const DEFAULT_COMPLETION_TOKEN_ESTIMATE: u32 = 1024; - -/// Rough prompt-token estimate for pre-flight budgeting, using the ~4-chars-per -/// token rule of thumb over the message contents. Approximate by design (the -/// authoritative token count only exists post-call); a floor of 1 avoids a -/// zero estimate for an empty prompt. -fn estimate_prompt_tokens(req: &InferenceRequest) -> u32 { - let chars: usize = req.messages.iter().map(|m| m.content.len()).sum(); - ((chars / 4) as u32).max(1) -} + /// USD cost from per-1k-token prices. `None` when both prices are zero (free / + /// local models), so no spurious zero cost is attributed. + fn estimate_cost( + price_in_per_1k: f64, + price_out_per_1k: f64, + prompt_tokens: Option, + completion_tokens: Option, + ) -> Option { + if price_in_per_1k == 0.0 && price_out_per_1k == 0.0 { + return None; + } + let prompt = prompt_tokens.unwrap_or(0) as f64; + let completion = completion_tokens.unwrap_or(0) as f64; + Some(prompt / 1000.0 * price_in_per_1k + completion / 1000.0 * price_out_per_1k) + } -/// Completion-token estimate for pre-flight budgeting: the request's -/// `params.max_tokens` cap when present, otherwise [`DEFAULT_COMPLETION_TOKEN_ESTIMATE`]. -fn estimate_completion_tokens(req: &InferenceRequest) -> u32 { - req.params - .get("max_tokens") - .and_then(serde_json::Value::as_u64) - .map(|v| v as u32) - .unwrap_or(DEFAULT_COMPLETION_TOKEN_ESTIMATE) + /// Rough prompt-token estimate for pre-flight budgeting, using the ~4-chars-per + /// token rule of thumb over the message contents. Approximate by design (the + /// authoritative token count only exists post-call); a floor of 1 avoids a + /// zero estimate for an empty prompt. + fn estimate_prompt_tokens(req: &InferenceRequest) -> u32 { + let chars: usize = req.messages.iter().map(|m| m.content.len()).sum(); + ((chars / 4) as u32).max(1) + } + + /// Completion-token estimate for pre-flight budgeting: the request's + /// `params.max_tokens` cap when present, otherwise + /// [`Self::DEFAULT_COMPLETION_TOKEN_ESTIMATE`]. + fn estimate_completion_tokens(req: &InferenceRequest) -> u32 { + req.params + .get("max_tokens") + .and_then(serde_json::Value::as_u64) + .map(|v| v as u32) + .unwrap_or(Self::DEFAULT_COMPLETION_TOKEN_ESTIMATE) + } } impl Engine { @@ -548,7 +553,7 @@ impl Engine { completion_tokens: Option, ) -> Option { let (price_in, price_out) = self.pricing.get(provider).copied()?; - estimate_cost(price_in, price_out, prompt_tokens, completion_tokens) + Engine::estimate_cost(price_in, price_out, prompt_tokens, completion_tokens) } /// Conservative pre-flight cost estimate (USD) for serving `req` on @@ -563,8 +568,8 @@ impl Engine { fn estimate_request_cost(&self, provider: &str, req: &InferenceRequest) -> f64 { self.compute_cost( provider, - Some(estimate_prompt_tokens(req)), - Some(estimate_completion_tokens(req)), + Some(Engine::estimate_prompt_tokens(req)), + Some(Engine::estimate_completion_tokens(req)), ) .unwrap_or(0.0) } @@ -614,11 +619,11 @@ impl Engine { let model_id = decision.model.id.clone(); let provider_name = decision.model.provider.clone(); - if remaining(overall_deadline).is_zero() { + if Engine::remaining(overall_deadline).is_zero() { let e = EngineError::Timeout("overall request deadline exceeded".into()); chain.push(FailedAttempt { provider: provider_name, - model: model_id_name(&model_id).to_string(), + model: ModelId::name(&model_id).to_string(), reason: e.to_string(), }); last_err = Some(e); @@ -637,7 +642,7 @@ impl Engine { ); chain.push(FailedAttempt { provider: provider_name.clone(), - model: model_id_name(&model_id).to_string(), + model: ModelId::name(&model_id).to_string(), reason: format!("estimated cost ${est:.4} exceeds budget ${ceiling:.4}"), }); last_err = Some(EngineError::NoCapableModel(format!( @@ -672,7 +677,7 @@ impl Engine { tracing::warn!("candidate '{model_id}' failed (retryable): {e}"); chain.push(FailedAttempt { provider: provider_name, - model: model_id_name(&model_id).to_string(), + model: ModelId::name(&model_id).to_string(), reason: e.to_string(), }); last_err = Some(e); @@ -796,7 +801,7 @@ impl Engine { // 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 queue_budget = Engine::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(_)) => { @@ -817,7 +822,7 @@ impl Engine { // 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 inference_budget = Engine::remaining(overall_deadline).min(self.timeouts.inference()); let start = Instant::now(); let result = tokio::time::timeout(inference_budget, provider.invoke(model_id, req)).await; let duration_ms = start.elapsed().as_millis() as u64; @@ -928,7 +933,7 @@ impl Engine { } let sem = self.semaphore_for(&model_id, max_concurrency); - let queue_budget = remaining(overall_deadline).min(self.timeouts.queue()); + let queue_budget = Engine::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(_)) => { @@ -951,7 +956,7 @@ impl Engine { let _ = out .send(StreamChunk::Started { request_id: req.request_id.clone(), - model_used: model_id_name(&model_id).to_string(), + model_used: ModelId::name(&model_id).to_string(), provider: provider_name.clone(), }) .await; @@ -974,7 +979,7 @@ impl Engine { } }); - let inference_budget = remaining(overall_deadline).min(self.timeouts.inference()); + let inference_budget = Engine::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)) @@ -1860,7 +1865,7 @@ mod tests { use async_trait::async_trait; use mofa_kernel::{ BackendFeature, Capability, ExecutionState, LifecycleResult, Message, ModelAvailability, - canonical_model_id, + ModelId, }; use std::sync::atomic::AtomicUsize; @@ -1963,7 +1968,7 @@ mod tests { .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.id = ModelId::canonical(&self.name, name); card.availability = ModelAvailability::Discovered; card.residency = *residency; card.memory_estimate_bytes = *mem; @@ -1997,7 +2002,7 @@ mod tests { let mem = self .template .iter() - .find(|(n, ..)| canonical_model_id(&self.name, n) == model_id) + .find(|(n, ..)| ModelId::canonical(&self.name, n) == model_id) .map(|(.., mem, _)| *mem); Ok(LifecycleResult { model_id: model_id.into(), @@ -2044,7 +2049,7 @@ mod tests { Ok(InferenceResponse { text: Some("ok".into()), file: None, - model_used: mofa_kernel::model_id_name(model_id).into(), + model_used: mofa_kernel::ModelId::name(model_id).into(), provider: self.name.clone(), duration_ms: 1, request_id: request.request_id.clone(), @@ -2153,12 +2158,12 @@ mod tests { #[test] fn cost_estimation_from_prices() { // $0.01/1k in, $0.03/1k out; 1000 prompt + 500 completion tokens. - let cost = estimate_cost(0.01, 0.03, Some(1000), Some(500)).unwrap(); + let cost = Engine::estimate_cost(0.01, 0.03, Some(1000), Some(500)).unwrap(); assert!((cost - (0.01 + 0.015)).abs() < 1e-9); // Zero prices (free / local model) → no cost attributed. - assert!(estimate_cost(0.0, 0.0, Some(1000), Some(500)).is_none()); + assert!(Engine::estimate_cost(0.0, 0.0, Some(1000), Some(500)).is_none()); // Missing token counts default to zero. - assert_eq!(estimate_cost(0.01, 0.03, None, None), Some(0.0)); + assert_eq!(Engine::estimate_cost(0.01, 0.03, None, None), Some(0.0)); } #[tokio::test] @@ -2393,16 +2398,16 @@ mod tests { fn token_estimates_from_request() { // ~4 chars/token over message content; floor of 1 for an empty prompt. let mut req = chat_request(None); // content = "hello" (5 chars) - assert_eq!(estimate_prompt_tokens(&req), 1); + assert_eq!(Engine::estimate_prompt_tokens(&req), 1); req.messages[0].content = "x".repeat(400); - assert_eq!(estimate_prompt_tokens(&req), 100); + assert_eq!(Engine::estimate_prompt_tokens(&req), 100); // Completion defaults to the generous estimate, or the params cap. assert_eq!( - estimate_completion_tokens(&req), - DEFAULT_COMPLETION_TOKEN_ESTIMATE + Engine::estimate_completion_tokens(&req), + Engine::DEFAULT_COMPLETION_TOKEN_ESTIMATE ); req.params = serde_json::json!({ "max_tokens": 64 }); - assert_eq!(estimate_completion_tokens(&req), 64); + assert_eq!(Engine::estimate_completion_tokens(&req), 64); } #[tokio::test] @@ -3056,7 +3061,7 @@ mod tests { Ok(InferenceResponse { text: Some(self.deltas.concat()), file: None, - model_used: mofa_kernel::model_id_name(model_id).into(), + model_used: mofa_kernel::ModelId::name(model_id).into(), provider: self.name.clone(), duration_ms: 1, request_id: request.request_id.clone(), diff --git a/mofa-engine-core/src/lib.rs b/mofa-engine-core/src/lib.rs index 13cbdf2..2336fd4 100644 --- a/mofa-engine-core/src/lib.rs +++ b/mofa-engine-core/src/lib.rs @@ -4,16 +4,20 @@ //! memory management, circuit breaker, preflight prediction, //! and the main `Engine` orchestrator. -pub mod artifacts; -pub mod backends; -pub mod circuit_breaker; +// Crate-internal machinery: reachable across the crate but not part of the +// published surface, so it can evolve freely. +pub(crate) mod artifacts; +pub(crate) mod backends; +pub(crate) mod circuit_breaker; +pub(crate) mod memory; +pub(crate) mod metrics; +pub(crate) mod router; + +// Public surface consumed by the SDK, the app, and examples. pub mod config; pub mod engine; -pub mod memory; -pub mod metrics; pub mod preflight; pub mod quality_gate; -pub mod router; pub mod subscription; pub use config::EngineConfig; diff --git a/mofa-engine-core/src/memory.rs b/mofa-engine-core/src/memory.rs index c874fe3..a071488 100644 --- a/mofa-engine-core/src/memory.rs +++ b/mofa-engine-core/src/memory.rs @@ -105,6 +105,8 @@ impl MemoryManager { } /// Whether `needed_bytes` currently fits within the budget without eviction. + /// Test-support: production paths reserve atomically via [`Self::try_reserve`]. + #[cfg(test)] pub fn can_fit(&self, needed_bytes: u64) -> bool { self.available_bytes() >= needed_bytes } @@ -137,10 +139,9 @@ impl MemoryManager { 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. + /// Unconditionally record a reservation, ignoring the budget. Test-support: + /// production paths reserve atomically via [`Self::try_reserve`]. + #[cfg(test)] pub fn allocate(&self, model_id: &str, bytes: u64) { let mut allocs = self.lock(); match allocs.get_mut(model_id) { diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs index 9f7408d..afb52c7 100644 --- a/mofa-engine-core/src/metrics.rs +++ b/mofa-engine-core/src/metrics.rs @@ -81,11 +81,6 @@ impl EngineMetrics { .fetch_add(1, Ordering::Relaxed); } - /// Total requests observed. - pub fn requests_total(&self) -> u64 { - self.requests_total.load(Ordering::Relaxed) - } - /// Record per-provider token usage and cost for one successful request. /// /// `is_local` tags the series with the dual-track `locality` label so the @@ -253,7 +248,7 @@ impl EngineMetrics { for (provider, up) in &gauges.provider_up { out.push_str(&format!( "mofa_provider_up{{provider=\"{}\"}} {}\n", - escape_label_value(provider), + Self::escape_label_value(provider), u8::from(*up) )); } @@ -265,7 +260,7 @@ impl EngineMetrics { # TYPE mofa_tokens_total counter\n", ); for entry in self.per_provider.iter() { - let provider = escape_label_value(entry.key()); + let provider = Self::escape_label_value(entry.key()); let locality = entry.locality.get().copied().unwrap_or("unknown"); out.push_str(&format!( "mofa_tokens_total{{provider=\"{provider}\",locality=\"{locality}\",direction=\"prompt\"}} {}\n", @@ -285,7 +280,7 @@ impl EngineMetrics { let locality = entry.locality.get().copied().unwrap_or("unknown"); out.push_str(&format!( "mofa_cost_usd_total{{provider=\"{}\",locality=\"{locality}\"}} {usd}\n", - escape_label_value(entry.key()) + Self::escape_label_value(entry.key()) )); } @@ -293,14 +288,16 @@ impl EngineMetrics { } } -/// 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") +impl EngineMetrics { + /// 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. diff --git a/mofa-engine-core/src/quality_gate.rs b/mofa-engine-core/src/quality_gate.rs index dcf28d4..3350868 100644 --- a/mofa-engine-core/src/quality_gate.rs +++ b/mofa-engine-core/src/quality_gate.rs @@ -241,7 +241,7 @@ impl QualityGate { if !output.status.success() { return None; } - parse_ffprobe_json(&String::from_utf8_lossy(&output.stdout)) + QualityGate::parse_ffprobe_json(&String::from_utf8_lossy(&output.stdout)) } /// Sample `sample_frames` tiny grayscale frames across the video and reduce @@ -266,78 +266,82 @@ impl QualityGate { if !output.status.success() { return None; } - slideshow_risk_from_frames(&output.stdout, FRAME_PIXELS) + QualityGate::slideshow_risk_from_frames(&output.stdout, FRAME_PIXELS) } } -/// Parse the `ffprobe -of json` payload into a [`ProbeInfo`]. Extracted as a pure -/// function so the field-plucking is unit-testable without invoking `ffprobe`. -fn parse_ffprobe_json(json: &str) -> Option { - let root: serde_json::Value = serde_json::from_str(json).ok()?; - - let duration_secs = root - .get("format") - .and_then(|f| f.get("duration")) - .and_then(|d| d.as_str()) - .and_then(|s| s.trim().parse::().ok()) - .filter(|d| d.is_finite() && *d > 0.0); - - let has_video = root - .get("streams") - .and_then(|s| s.as_array()) - .is_some_and(|streams| { - streams.iter().any(|st| { - st.get("codec_type").and_then(|c| c.as_str()) == Some("video") - && st - .get("width") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) - > 0 - && st - .get("height") - .and_then(serde_json::Value::as_u64) - .unwrap_or(0) - > 0 - }) - }); - - Some(ProbeInfo { - duration_secs, - has_video, - }) -} +/// Pure parsing/scoring helpers, grouped as private associated functions so the +/// signal extraction shares the `QualityGate` namespace and stays unit-testable +/// without invoking `ffprobe`/`ffmpeg`. +impl QualityGate { + /// Parse the `ffprobe -of json` payload into a [`ProbeInfo`]. + fn parse_ffprobe_json(json: &str) -> Option { + let root: serde_json::Value = serde_json::from_str(json).ok()?; + + let duration_secs = root + .get("format") + .and_then(|f| f.get("duration")) + .and_then(|d| d.as_str()) + .and_then(|s| s.trim().parse::().ok()) + .filter(|d| d.is_finite() && *d > 0.0); + + let has_video = root + .get("streams") + .and_then(|s| s.as_array()) + .is_some_and(|streams| { + streams.iter().any(|st| { + st.get("codec_type").and_then(|c| c.as_str()) == Some("video") + && st + .get("width") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) + > 0 + && st + .get("height") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0) + > 0 + }) + }); -/// Reduce a raw grayscale frame stream (frames of `frame_pixels` bytes, back to -/// back) to a slideshow-risk score in `0..=1`. -/// -/// The *motion* between two consecutive frames is the mean absolute per-pixel -/// difference, normalized to `0..1`. Averaged over all consecutive pairs and -/// compared against [`MOTION_FULL_SCALE`], it yields `risk = 1 - motion_ratio`: -/// identical frames (a still slideshow) score risk `1.0`; frames that change by -/// [`MOTION_FULL_SCALE`] or more score risk `0.0`. Returns `None` when fewer than -/// two whole frames are present (motion is undefined). -fn slideshow_risk_from_frames(raw: &[u8], frame_pixels: usize) -> Option { - if frame_pixels == 0 { - return None; - } - let frames: Vec<&[u8]> = raw.chunks_exact(frame_pixels).collect(); - if frames.len() < 2 { - return None; + Some(ProbeInfo { + duration_secs, + has_video, + }) } - let mut total_motion = 0.0; - for pair in frames.windows(2) { - let diff: u64 = pair[0] - .iter() - .zip(pair[1].iter()) - .map(|(a, b)| a.abs_diff(*b) as u64) - .sum(); - // Mean absolute difference for this pair, normalized to 0..1. - total_motion += diff as f64 / (frame_pixels as f64 * 255.0); + /// Reduce a raw grayscale frame stream (frames of `frame_pixels` bytes, back to + /// back) to a slideshow-risk score in `0..=1`. + /// + /// The *motion* between two consecutive frames is the mean absolute per-pixel + /// difference, normalized to `0..1`. Averaged over all consecutive pairs and + /// compared against [`MOTION_FULL_SCALE`], it yields `risk = 1 - motion_ratio`: + /// identical frames (a still slideshow) score risk `1.0`; frames that change by + /// [`MOTION_FULL_SCALE`] or more score risk `0.0`. Returns `None` when fewer than + /// two whole frames are present (motion is undefined). + fn slideshow_risk_from_frames(raw: &[u8], frame_pixels: usize) -> Option { + if frame_pixels == 0 { + return None; + } + let frames: Vec<&[u8]> = raw.chunks_exact(frame_pixels).collect(); + if frames.len() < 2 { + return None; + } + + let mut total_motion = 0.0; + for pair in frames.windows(2) { + let diff: u64 = pair[0] + .iter() + .zip(pair[1].iter()) + .map(|(a, b)| a.abs_diff(*b) as u64) + .sum(); + // Mean absolute difference for this pair, normalized to 0..1. + total_motion += diff as f64 / (frame_pixels as f64 * 255.0); + } + let motion = total_motion / (frames.len() - 1) as f64; + let motion_ratio = (motion / MOTION_FULL_SCALE).clamp(0.0, 1.0); + Some(1.0 - motion_ratio) } - let motion = total_motion / (frames.len() - 1) as f64; - let motion_ratio = (motion / MOTION_FULL_SCALE).clamp(0.0, 1.0); - Some(1.0 - motion_ratio) } #[cfg(test)] @@ -353,7 +357,7 @@ mod tests { ], "format": {"duration": "12.34"} }"#; - let info = parse_ffprobe_json(json).unwrap(); + let info = QualityGate::parse_ffprobe_json(json).unwrap(); assert_eq!(info.duration_secs, Some(12.34)); assert!(info.has_video); } @@ -366,7 +370,7 @@ mod tests { "streams": [{"codec_type": "video", "width": 0, "height": 0}], "format": {"duration": "N/A"} }"#; - let info = parse_ffprobe_json(json).unwrap(); + let info = QualityGate::parse_ffprobe_json(json).unwrap(); assert_eq!(info.duration_secs, None); assert!(!info.has_video); } @@ -378,7 +382,10 @@ mod tests { frame[0] = 200; // some internal structure, still identical across frames let mut raw = frame.clone(); raw.extend_from_slice(&frame); - assert_eq!(slideshow_risk_from_frames(&raw, FRAME_PIXELS), Some(1.0)); + assert_eq!( + QualityGate::slideshow_risk_from_frames(&raw, FRAME_PIXELS), + Some(1.0) + ); } #[test] @@ -386,14 +393,23 @@ mod tests { // Black frame then white frame: maximal motion → zero risk. let mut raw = vec![0u8; FRAME_PIXELS]; raw.extend(std::iter::repeat_n(255u8, FRAME_PIXELS)); - assert_eq!(slideshow_risk_from_frames(&raw, FRAME_PIXELS), Some(0.0)); + assert_eq!( + QualityGate::slideshow_risk_from_frames(&raw, FRAME_PIXELS), + Some(0.0) + ); } #[test] fn slideshow_risk_needs_two_frames() { - assert_eq!(slideshow_risk_from_frames(&[1, 2, 3], 64), None); - assert_eq!(slideshow_risk_from_frames(&[], 64), None); - assert_eq!(slideshow_risk_from_frames(&[0u8; 128], 0), None); + assert_eq!( + QualityGate::slideshow_risk_from_frames(&[1, 2, 3], 64), + None + ); + assert_eq!(QualityGate::slideshow_risk_from_frames(&[], 64), None); + assert_eq!( + QualityGate::slideshow_risk_from_frames(&[0u8; 128], 0), + None + ); } #[test] @@ -405,7 +421,7 @@ mod tests { moved[0] = 130; // one pixel changes by 30/255 over 64 pixels → tiny motion let mut raw = base.clone(); raw.extend_from_slice(&moved); - let risk = slideshow_risk_from_frames(&raw, FRAME_PIXELS).unwrap(); + let risk = QualityGate::slideshow_risk_from_frames(&raw, FRAME_PIXELS).unwrap(); assert!(risk > 0.9 && risk < 1.0, "risk was {risk}"); } diff --git a/mofa-engine-core/src/router.rs b/mofa-engine-core/src/router.rs index 41d0794..83ce95e 100644 --- a/mofa-engine-core/src/router.rs +++ b/mofa-engine-core/src/router.rs @@ -43,19 +43,6 @@ pub struct RouteDecision<'a> { 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. @@ -157,15 +144,6 @@ impl Router { decisions } - /// Compatibility helper returning just the selected model. - pub fn select_model<'a>( - models: &'a [ModelCard], - request: &InferenceRequest, - providers: &[RoutingProvider], - ) -> Option<&'a ModelCard> { - 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 { @@ -410,9 +388,12 @@ mod tests { ), ]; let providers = vec![provider("ollama", ProviderKind::Ollama, 1)]; - let selected = Router::route(&models, &request(Capability::Chat, None), &providers) - .unwrap() - .model; + let selected = + Router::route_ranked(&models, &request(Capability::Chat, None), &providers, None) + .into_iter() + .next() + .unwrap() + .model; assert_eq!(selected.name, "b"); } @@ -438,9 +419,12 @@ mod tests { provider("openai", ProviderKind::OpenAiCompatible, 10), provider("ollama", ProviderKind::Ollama, 1), ]; - let selected = Router::route(&models, &request(Capability::Chat, None), &providers) - .unwrap() - .model; + let selected = + Router::route_ranked(&models, &request(Capability::Chat, None), &providers, None) + .into_iter() + .next() + .unwrap() + .model; assert_eq!(selected.name, "local"); } @@ -463,9 +447,12 @@ mod tests { ), ]; let providers = vec![provider("p", ProviderKind::OpenAiCompatible, 5)]; - let selected = Router::route(&models, &request(Capability::Tts, None), &providers) - .unwrap() - .model; + let selected = + Router::route_ranked(&models, &request(Capability::Tts, None), &providers, None) + .into_iter() + .next() + .unwrap() + .model; assert_eq!(selected.name, "tts"); } @@ -488,9 +475,16 @@ mod tests { ), ]; let providers = vec![provider("p", ProviderKind::OpenAiCompatible, 5)]; - let selected = Router::route(&models, &request(Capability::Chat, Some("b")), &providers) - .unwrap() - .model; + let selected = Router::route_ranked( + &models, + &request(Capability::Chat, Some("b")), + &providers, + None, + ) + .into_iter() + .next() + .unwrap() + .model; assert_eq!(selected.name, "b"); } @@ -589,7 +583,12 @@ mod tests { )]; let mut providers = vec![provider("p", ProviderKind::OpenAiCompatible, 5)]; providers[0].health = BackendHealth::Unavailable; - assert!(Router::route(&models, &request(Capability::Chat, None), &providers).is_none()); + assert!( + Router::route_ranked(&models, &request(Capability::Chat, None), &providers, None) + .into_iter() + .next() + .is_none() + ); } #[test] diff --git a/mofa-engine-sdk/examples/explainer_video.rs b/mofa-engine-sdk/examples/explainer_video.rs index 0a7f0de..eb9f918 100644 --- a/mofa-engine-sdk/examples/explainer_video.rs +++ b/mofa-engine-sdk/examples/explainer_video.rs @@ -55,7 +55,7 @@ async fn main() { hint_next: Some("image_gen".into()), ..Default::default() }; - let script = stage(&engine, script, "Script (chat)").await; + let script = ExplainerVideo::stage(&engine, script, "Script (chat)").await; // 2. Scene images (ImageGen, prefer local SD when available). let scene_prompt = "flat vector animation frame: a neuron adjusting its weights"; @@ -73,7 +73,7 @@ async fn main() { hint_next: Some("tts".into()), ..Default::default() }; - stage(&engine, image, "Scene image (image_gen)").await; + ExplainerVideo::stage(&engine, image, "Scene image (image_gen)").await; // 3. Narration (TTS). let narration = InferenceRequest { @@ -86,7 +86,7 @@ async fn main() { hint_next: Some("asr".into()), ..Default::default() }; - stage(&engine, narration, "Narration (tts)").await; + ExplainerVideo::stage(&engine, narration, "Narration (tts)").await; // 4. (compose) — external: scenes + narration + captions → final.mp4 via // Remotion/FFmpeg. Out of engine scope; supply the result as an argument. @@ -95,7 +95,7 @@ async fn main() { // 5. Quality gate — the hard gate. Runs against a real mp4 when provided. println!("\n== quality gate =="); match final_mp4.as_deref() { - Some(path) => run_gate(Path::new(path)).await, + Some(path) => ExplainerVideo::run_gate(Path::new(path)).await, None => println!( "no composed mp4 supplied — pass one to run the gate:\n \ cargo run -p mofa-engine-sdk --example explainer_video -- ./final.mp4" @@ -103,42 +103,47 @@ async fn main() { } } -/// Run one orchestration stage, reporting where it was routed or the structured -/// error (the pipeline continues either way — this demo showcases orchestration, -/// not a live render). Returns the text output when present. -async fn stage(engine: &Engine, req: InferenceRequest, label: &str) -> Option { - match engine.invoke(req).await { - Ok(resp) => { +/// Namespace for this demo's helpers. +struct ExplainerVideo; + +impl ExplainerVideo { + /// Run one orchestration stage, reporting where it was routed or the structured + /// error (the pipeline continues either way — this demo showcases orchestration, + /// not a live render). Returns the text output when present. + async fn stage(engine: &Engine, req: InferenceRequest, label: &str) -> Option { + match engine.invoke(req).await { + Ok(resp) => { + println!( + "[{label}] served by {} in {}ms", + resp.provider, resp.duration_ms + ); + resp.text + } + Err(e) => { + let info = e.info(); + println!("[{label}] skipped — {:?}: {}", info.code, info.message); + None + } + } + } + + /// Run the real quality gate over a composed video and print the verdict. The + /// optional VLM image/text check is omitted here (`None`); a full pipeline would + /// pass the engine's `Vlm` verdict over sampled frames. + async fn run_gate(path: &Path) { + let report = QualityGate::new().check(path, None).await; + for c in &report.checks { println!( - "[{label}] served by {} in {}ms", - resp.provider, resp.duration_ms + " [{}] {}: {}", + if c.passed { "PASS" } else { "FAIL" }, + c.name, + c.detail ); - resp.text } - Err(e) => { - let info = e.info(); - println!("[{label}] skipped — {:?}: {}", info.code, info.message); - None + if report.passed { + println!("verdict: ACCEPTED ✅"); + } else { + println!("verdict: REJECTED ❌ (no gate, no output)"); } } } - -/// Run the real quality gate over a composed video and print the verdict. The -/// optional VLM image/text check is omitted here (`None`); a full pipeline would -/// pass the engine's `Vlm` verdict over sampled frames. -async fn run_gate(path: &Path) { - let report = QualityGate::new().check(path, None).await; - for c in &report.checks { - println!( - " [{}] {}: {}", - if c.passed { "PASS" } else { "FAIL" }, - c.name, - c.detail - ); - } - if report.passed { - println!("verdict: ACCEPTED ✅"); - } else { - println!("verdict: REJECTED ❌ (no gate, no output)"); - } -} diff --git a/mofa-engine-sdk/examples/meeting_brief.rs b/mofa-engine-sdk/examples/meeting_brief.rs index 68dd88b..eebc54b 100644 --- a/mofa-engine-sdk/examples/meeting_brief.rs +++ b/mofa-engine-sdk/examples/meeting_brief.rs @@ -45,11 +45,11 @@ fn main() { hint_next: Some("chat".into()), ..Default::default() }; - let transcript = match run(&engine, asr, "ASR") { + let transcript = match MeetingBrief::run(&engine, asr, "ASR") { Some(r) => r.text.unwrap_or_default(), None => return, }; - println!("transcript: {}\n", preview(&transcript)); + println!("transcript: {}\n", MeetingBrief::preview(&transcript)); // 2. Chat — extract structured minutes; warm TTS meanwhile. let chat = InferenceRequest { @@ -73,11 +73,11 @@ fn main() { hint_next: Some("tts".into()), ..Default::default() }; - let minutes = match run(&engine, chat, "Chat (minutes)") { + let minutes = match MeetingBrief::run(&engine, chat, "Chat (minutes)") { Some(r) => r.text.unwrap_or_default(), None => return, }; - println!("minutes:\n{}\n", preview(&minutes)); + println!("minutes:\n{}\n", MeetingBrief::preview(&minutes)); // 3. TTS — minutes → a short audio brief. let tts = InferenceRequest { @@ -92,38 +92,48 @@ fn main() { params: serde_json::json!({ "voice": "zh-female-1" }), ..Default::default() }; - if let Some(r) = run(&engine, tts, "TTS (brief)") { + if let Some(r) = MeetingBrief::run(&engine, tts, "TTS (brief)") { println!("audio brief: {}", r.file.as_deref().unwrap_or("(no file)")); } } -/// Invoke one stage, printing routing/cost on success or the structured error -/// (including any failover chain) on failure. Returns `None` to stop the pipeline. -fn run(engine: &EmbeddedEngine, req: InferenceRequest, stage: &str) -> Option { - match engine.invoke(req) { - Ok(resp) => { - println!( - "[{stage}] served by {} in {}ms", - resp.provider, resp.duration_ms - ); - Some(resp) - } - Err(e) => { - let info = e.info(); - eprintln!("[{stage}] error [{:?}]: {}", info.code, info.message); - for a in &info.failed_chain { - eprintln!(" tried {}/{}: {}", a.provider, a.model, a.reason); +/// Namespace for this demo's helpers. +struct MeetingBrief; + +impl MeetingBrief { + /// Invoke one stage, printing routing/cost on success or the structured error + /// (including any failover chain) on failure. Returns `None` to stop the pipeline. + fn run( + engine: &EmbeddedEngine, + req: InferenceRequest, + stage: &str, + ) -> Option { + match engine.invoke(req) { + Ok(resp) => { + println!( + "[{stage}] served by {} in {}ms", + resp.provider, resp.duration_ms + ); + Some(resp) + } + Err(e) => { + let info = e.info(); + eprintln!("[{stage}] error [{:?}]: {}", info.code, info.message); + for a in &info.failed_chain { + eprintln!(" tried {}/{}: {}", a.provider, a.model, a.reason); + } + None } - None } } -} -fn preview(s: &str) -> String { - let s = s.trim(); - if s.chars().count() > 200 { - format!("{}…", s.chars().take(200).collect::()) - } else { - s.to_string() + /// Truncate long text for a readable preview. + fn preview(s: &str) -> String { + let s = s.trim(); + if s.chars().count() > 200 { + format!("{}…", s.chars().take(200).collect::()) + } else { + s.to_string() + } } } diff --git a/mofa-engine-sdk/src/client.rs b/mofa-engine-sdk/src/client.rs index 289d107..b3dedc1 100644 --- a/mofa-engine-sdk/src/client.rs +++ b/mofa-engine-sdk/src/client.rs @@ -154,16 +154,18 @@ impl EmbeddedEngine { /// 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())))?; + .map_err(|e| Self::err_json(EngineError::InvalidRequest(e.to_string())))?; self.invoke(request) .map(|resp| serde_json::to_string(&resp).unwrap_or_default()) - .map_err(err_json) + .map_err(Self::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()) +impl EmbeddedEngine { + /// 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). diff --git a/mofa-engine-sdk/src/lib.rs b/mofa-engine-sdk/src/lib.rs index 8ed550e..1025cd0 100644 --- a/mofa-engine-sdk/src/lib.rs +++ b/mofa-engine-sdk/src/lib.rs @@ -11,8 +11,10 @@ //! exposes. pub mod client; -pub mod dashboard; pub mod server; +// Internal: the embedded dashboard HTML, served by `server`. +pub(crate) mod dashboard; + pub use client::{ClientError, DaemonClient, EmbeddedEngine}; -pub use server::start_server; +pub use server::Server; diff --git a/mofa-engine-sdk/src/server.rs b/mofa-engine-sdk/src/server.rs index 9c20294..99aad1e 100644 --- a/mofa-engine-sdk/src/server.rs +++ b/mofa-engine-sdk/src/server.rs @@ -43,81 +43,92 @@ struct AppState { 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)"); - } +/// The MoFA Engine HTTP server — the public entry point to the versioned `/v1` +/// API, SSE streaming, and the embedded dashboard. +pub struct Server; + +impl Server { + /// 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( + 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 state = AppState { + engine, + started_at: std::time::Instant::now(), + api_token, + }; - let app = build_app(state); + let app = AppState::build_app(state); - let addr = format!("{host}:{port}"); - tracing::info!("MoFA Engine listening on http://{addr}"); + 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?; + let listener = tokio::net::TcpListener::bind(&addr).await?; + axum::serve(listener, app).await?; - Ok(()) + 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()) - // 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) +impl AppState { + /// 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(AppState::capabilities_handler)) + .route("/v1/invoke", post(AppState::invoke_handler)) + .route("/v1/invoke/stream", post(AppState::invoke_stream_handler)) + .route("/v1/status", get(AppState::status_handler)) + .route("/v1/memory", get(AppState::memory_handler)) + .route("/v1/lifecycle", get(AppState::lifecycle_handler)) + .route("/v1/preflight", get(AppState::preflight_handler)) + .route( + "/v1/subscriptions", + get(AppState::list_subscriptions_handler).post(AppState::subscribe_handler), + ) + .route( + "/v1/subscriptions/{id}", + delete(AppState::unsubscribe_handler), + ) + .route("/v1/events", get(AppState::events_handler)) + .route("/v1/discovery/refresh", post(AppState::refresh_handler)) + .route("/v1/models/load", post(AppState::load_model_handler)) + .route("/v1/models/unload", post(AppState::unload_model_handler)) + .route_layer(middleware::from_fn_with_state( + state.clone(), + AppState::auth_middleware, + )); + + Router::new() + // Public (unauthenticated) routes. + .route("/", get(AppState::dashboard_handler)) + .route("/health", get(AppState::health_handler)) + .route("/metrics", get(AppState::metrics_handler)) + .merge(api) + .layer(DefaultBodyLimit::max(MAX_BODY_BYTES)) + .layer(CorsLayer::permissive()) + .layer(TraceLayer::new_for_http()) + // Outermost, so the request-id span encloses the trace layer and every + // response (including framework rejections) carries `x-request-id`. + .layer(middleware::from_fn(AppState::correlation_middleware)) + .with_state(state) + } } /// Maximum accepted length of a client-supplied `x-request-id`. A longer (or @@ -125,58 +136,60 @@ fn build_app(state: AppState) -> Router { /// or allocation cost through the correlation header. const MAX_REQUEST_ID_LEN: usize = 128; -/// 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()); - - 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); +impl AppState { + /// 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()); + + 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 } - 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() + /// 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| AppState::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; + /// 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 } - diff == 0 } /// Health check response. @@ -187,52 +200,55 @@ struct HealthResponse { uptime_secs: u64, } -async fn health_handler(State(state): State) -> Json { - Json(HealthResponse { - status: "ok", - version: env!("CARGO_PKG_VERSION"), - uptime_secs: state.started_at.elapsed().as_secs(), - }) -} +impl AppState { + async fn health_handler(State(state): State) -> Json { + Json(HealthResponse { + status: "ok", + version: env!("CARGO_PKG_VERSION"), + uptime_secs: state.started_at.elapsed().as_secs(), + }) + } -async fn capabilities_handler(State(state): State) -> Json> { - let caps = state.engine.capabilities().await; - Json(caps) -} + async fn capabilities_handler( + State(state): State, + ) -> Json> { + let caps = state.engine.capabilities().await; + Json(caps) + } -async fn invoke_handler( - State(state): State, - Json(req): Json, -) -> Result, (StatusCode, Json)> { - match state.engine.invoke(req).await { - Ok(resp) => Ok(Json(resp)), - Err(e) => { - let status = match &e { - mofa_kernel::EngineError::NoCapableModel(_) => StatusCode::NOT_FOUND, - mofa_kernel::EngineError::InvalidRequest(_) => StatusCode::BAD_REQUEST, - mofa_kernel::EngineError::CircuitOpen(_) => StatusCode::SERVICE_UNAVAILABLE, - mofa_kernel::EngineError::Timeout(_) => StatusCode::GATEWAY_TIMEOUT, - mofa_kernel::EngineError::UnsupportedOperation(_) => StatusCode::BAD_REQUEST, - mofa_kernel::EngineError::Config(_) => StatusCode::INTERNAL_SERVER_ERROR, - _ => StatusCode::INTERNAL_SERVER_ERROR, - }; - Err((status, Json(e.info()))) + async fn invoke_handler( + State(state): State, + Json(req): Json, + ) -> Result, (StatusCode, Json)> { + match state.engine.invoke(req).await { + Ok(resp) => Ok(Json(resp)), + Err(e) => { + let status = match &e { + mofa_kernel::EngineError::NoCapableModel(_) => StatusCode::NOT_FOUND, + mofa_kernel::EngineError::InvalidRequest(_) => StatusCode::BAD_REQUEST, + mofa_kernel::EngineError::CircuitOpen(_) => StatusCode::SERVICE_UNAVAILABLE, + mofa_kernel::EngineError::Timeout(_) => StatusCode::GATEWAY_TIMEOUT, + mofa_kernel::EngineError::UnsupportedOperation(_) => StatusCode::BAD_REQUEST, + mofa_kernel::EngineError::Config(_) => StatusCode::INTERNAL_SERVER_ERROR, + _ => StatusCode::INTERNAL_SERVER_ERROR, + }; + Err((status, Json(e.info()))) + } } } -} -/// 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| { + /// 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. @@ -241,41 +257,42 @@ async fn invoke_stream_handler( }); Ok(Event::default().data(data)) }); - Sse::new(stream).keep_alive( - KeepAlive::new() - .interval(Duration::from_secs(15)) - .text("ping"), - ) -} + 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) -} + async fn status_handler(State(state): State) -> Json { + let status = state.engine.status().await; + Json(status) + } -async fn refresh_handler(State(state): State) -> Json { - state.engine.refresh_resources().await; - Json(state.engine.status().await) -} + async fn refresh_handler(State(state): State) -> Json { + state.engine.refresh_resources().await; + Json(state.engine.status().await) + } -/// Prometheus text-exposition metrics. Public (no auth) so scrapers can reach it. -async fn metrics_handler(State(state): State) -> impl IntoResponse { - let body = state.engine.metrics_prometheus(); - ([(header::CONTENT_TYPE, "text/plain; version=0.0.4")], body) -} + /// Prometheus text-exposition metrics. Public (no auth) so scrapers can reach it. + async fn metrics_handler(State(state): State) -> 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(_) | EngineError::Failover { .. } => { - StatusCode::SERVICE_UNAVAILABLE + /// 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(_) | EngineError::Failover { .. } => { + StatusCode::SERVICE_UNAVAILABLE + } + EngineError::Timeout(_) => StatusCode::GATEWAY_TIMEOUT, + _ => StatusCode::INTERNAL_SERVER_ERROR, } - EngineError::Timeout(_) => StatusCode::GATEWAY_TIMEOUT, - _ => StatusCode::INTERNAL_SERVER_ERROR, } } @@ -292,48 +309,52 @@ struct ModelActionResponse { 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()))), +impl AppState { + /// `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((AppState::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()))), + /// `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((AppState::error_status(&e), Json(e.info()))), + } } -} -async fn memory_handler(State(state): State) -> Json { - Json(state.engine.memory_report()) -} + 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 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 preflight_handler(State(state): State) -> Json { + Json(state.engine.preflight_stats()) + } -async fn list_subscriptions_handler(State(state): State) -> Json> { - Json(state.engine.subscriptions()) + async fn list_subscriptions_handler( + State(state): State, + ) -> Json> { + Json(state.engine.subscriptions()) + } } /// Body for `POST /v1/subscriptions`. @@ -354,19 +375,21 @@ 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()))); +impl AppState { + 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 })) } - 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)] @@ -374,40 +397,42 @@ 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 })) -} +impl AppState { + 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>> { - let rx = state.engine.subscribe_events(); - let stream = BroadcastStream::new(rx).filter_map(|result| match result { - Ok(event) => { - let data = serde_json::to_string(&event).unwrap_or_default(); - Some(Ok(Event::default().data(data))) - } - Err(_) => None, - }); + async fn events_handler( + State(state): State, + ) -> Sse>> { + let rx = state.engine.subscribe_events(); + let stream = BroadcastStream::new(rx).filter_map(|result| match result { + Ok(event) => { + let data = serde_json::to_string(&event).unwrap_or_default(); + Some(Ok(Event::default().data(data))) + } + Err(_) => None, + }); - Sse::new(stream).keep_alive( - KeepAlive::new() - .interval(Duration::from_secs(15)) - .text("ping"), - ) -} + Sse::new(stream).keep_alive( + KeepAlive::new() + .interval(Duration::from_secs(15)) + .text("ping"), + ) + } -async fn dashboard_handler() -> Html<&'static str> { - Html(dashboard::DASHBOARD_HTML) + async fn dashboard_handler() -> Html<&'static str> { + Html(dashboard::DASHBOARD_HTML) + } } #[cfg(test)] @@ -449,7 +474,7 @@ mod tests { #[tokio::test] async fn metrics_endpoint_serves_prometheus_text() { - let app = build_app(test_state(None).await); + let app = AppState::build_app(test_state(None).await); let resp = app .oneshot( HttpRequest::builder() @@ -475,7 +500,7 @@ mod tests { #[tokio::test] async fn responses_carry_a_request_id_header() { - let app = build_app(test_state(None).await); + let app = AppState::build_app(test_state(None).await); let resp = app .oneshot( HttpRequest::builder() @@ -490,7 +515,7 @@ mod tests { #[tokio::test] async fn provided_request_id_is_echoed_back() { - let app = build_app(test_state(None).await); + let app = AppState::build_app(test_state(None).await); let resp = app .oneshot( HttpRequest::builder() @@ -506,7 +531,7 @@ mod tests { #[tokio::test] async fn auth_gate_rejects_missing_token_but_allows_public_routes() { - let app = build_app(test_state(Some("secret")).await); + let app = AppState::build_app(test_state(Some("secret")).await); // A /v1 route without the token is unauthorized. let resp = app diff --git a/mofa-kernel/src/types.rs b/mofa-kernel/src/types.rs index a95655a..52f44a8 100644 --- a/mofa-kernel/src/types.rs +++ b/mofa-kernel/src/types.rs @@ -300,7 +300,7 @@ impl ModelCard { let execution = ExecutionState::default(); let status = ModelStatus::from_state(availability, residency, &execution); Self { - id: canonical_model_id(&provider, &name), + id: ModelId::canonical(&provider, &name), name, provider, capability, @@ -327,27 +327,38 @@ impl ModelCard { } } -/// Build a canonical model identifier. -pub fn canonical_model_id(provider: &str, model: &str) -> String { - format!("{provider}/{model}") -} +/// Namespace for canonical model-identifier operations. +/// +/// Model identifiers have the canonical form `provider/model` (a legacy +/// `provider::model` form is still accepted when parsing). Construction and +/// parsing live here as associated functions on a zero-sized type so that all id +/// handling has a single, discoverable home rather than scattered free functions. +pub struct ModelId; + +impl ModelId { + /// Build a canonical `provider/model` identifier. + pub fn canonical(provider: &str, model: &str) -> String { + format!("{provider}/{model}") + } -/// Extract the provider segment from a canonical model identifier. -/// Accepts both the canonical `provider/model` and legacy `provider::model` forms. -pub fn model_id_provider(model_id: &str) -> Option<&str> { - model_id - .split_once('/') - .map(|(provider, _)| provider) - .or_else(|| model_id.split_once("::").map(|(provider, _)| provider)) -} + /// Extract the provider segment, accepting the canonical `provider/model` and + /// the legacy `provider::model` forms. + pub fn provider(model_id: &str) -> Option<&str> { + model_id + .split_once('/') + .map(|(provider, _)| provider) + .or_else(|| model_id.split_once("::").map(|(provider, _)| provider)) + } -/// Extract the model-name segment from a canonical model identifier. -pub fn model_id_name(model_id: &str) -> &str { - model_id - .split_once('/') - .map(|(_, model)| model) - .or_else(|| model_id.split_once("::").map(|(_, model)| model)) - .unwrap_or(model_id) + /// Extract the model-name segment, falling back to the whole id when it + /// carries no provider prefix. + pub fn name(model_id: &str) -> &str { + model_id + .split_once('/') + .map(|(_, model)| model) + .or_else(|| model_id.split_once("::").map(|(_, model)| model)) + .unwrap_or(model_id) + } } /// A single message in a conversation. @@ -499,7 +510,7 @@ pub struct InferenceRequest { /// Hint about what capability will be needed next. pub hint_next: Option, /// Unique request identifier. - #[serde(default = "generate_request_id")] + #[serde(default = "InferenceRequest::generate_request_id")] pub request_id: String, } @@ -519,13 +530,17 @@ impl Default for InferenceRequest { input_file: None, params: serde_json::Value::Null, hint_next: None, - request_id: generate_request_id(), + request_id: Self::generate_request_id(), } } } -fn generate_request_id() -> String { - Uuid::new_v4().to_string() +impl InferenceRequest { + /// Generate a fresh unique request identifier. Also the serde default for + /// `request_id`, so a deserialized request without one still gets a unique id. + fn generate_request_id() -> String { + Uuid::new_v4().to_string() + } } /// Response from an inference call. @@ -831,9 +846,9 @@ mod tests { #[test] fn canonical_id_uses_slash() { - assert_eq!(canonical_model_id("ollama", "qwen"), "ollama/qwen"); - assert_eq!(model_id_name("ollama/qwen"), "qwen"); - assert_eq!(model_id_name("ollama::qwen"), "qwen"); + assert_eq!(ModelId::canonical("ollama", "qwen"), "ollama/qwen"); + assert_eq!(ModelId::name("ollama/qwen"), "qwen"); + assert_eq!(ModelId::name("ollama::qwen"), "qwen"); } #[test] From e2733a904fc0e77de4c8ea2c41f75476093fd5b8 Mon Sep 17 00:00:00 2001 From: aayank13 Date: Thu, 30 Jul 2026 21:09:44 +0530 Subject: [PATCH 10/22] refactor: change visibility of several structs and methods to crate-level --- mofa-engine-core/src/config.rs | 4 ++-- mofa-engine-core/src/engine.rs | 9 ++------- mofa-engine-core/src/preflight.rs | 8 ++++---- mofa-engine-core/src/subscription.rs | 5 +++-- 4 files changed, 11 insertions(+), 15 deletions(-) diff --git a/mofa-engine-core/src/config.rs b/mofa-engine-core/src/config.rs index 9fde243..799ffc3 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -451,7 +451,7 @@ impl EngineConfig { } /// Validate configuration before constructing the engine. - pub fn validate(&self) -> Result<(), EngineError> { + pub(crate) 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 {}", @@ -733,7 +733,7 @@ impl ProviderConfig { } /// Parse the configured provider kind. - pub fn provider_kind(&self) -> Result { + pub(crate) fn provider_kind(&self) -> Result { match self.kind.as_str() { "ollama" => Ok(ProviderKind::Ollama), "openai_compatible" => Ok(ProviderKind::OpenAiCompatible), diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index de92013..6afb071 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -1742,8 +1742,8 @@ impl Engine { } } - /// Get backend status snapshots. - pub fn backend_statuses(&self) -> Vec { + /// Snapshot every provider's backend status, for `status()`. + fn backend_statuses(&self) -> Vec { self.providers .iter() .map(|registered| BackendStatus { @@ -1838,11 +1838,6 @@ impl Engine { 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 { diff --git a/mofa-engine-core/src/preflight.rs b/mofa-engine-core/src/preflight.rs index 70012b4..9bbae73 100644 --- a/mofa-engine-core/src/preflight.rs +++ b/mofa-engine-core/src/preflight.rs @@ -22,10 +22,10 @@ const DECAY_INTERVAL: u64 = 100; /// 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__"; +pub(crate) const GLOBAL_SCOPE: &str = "__global__"; /// Per-scope Markov chain transition tracker with decay. -pub struct PreflightPredictor { +pub(crate) struct PreflightPredictor { inner: Mutex>, } @@ -104,7 +104,7 @@ impl ScopeChain { /// Prediction result. #[derive(Debug, Clone, PartialEq)] -pub struct Prediction { +pub(crate) struct Prediction { /// Predicted next capability. pub capability: Capability, /// Confidence score (0.0 – 1.0). @@ -220,7 +220,7 @@ impl Default for PreflightPredictor { /// Live counters describing Preflight effectiveness. #[derive(Debug, Default)] -pub struct PreflightMetrics { +pub(crate) struct PreflightMetrics { warms_started: AtomicU64, warms_completed: AtomicU64, warms_failed: AtomicU64, diff --git a/mofa-engine-core/src/subscription.rs b/mofa-engine-core/src/subscription.rs index 5e30fa1..e6e63da 100644 --- a/mofa-engine-core/src/subscription.rs +++ b/mofa-engine-core/src/subscription.rs @@ -14,7 +14,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; /// Registry of active capability subscriptions. -pub struct SubscriptionRegistry { +pub(crate) struct SubscriptionRegistry { inner: Mutex>, next_id: AtomicU64, } @@ -106,7 +106,8 @@ impl SubscriptionRegistry { .collect() } - /// Whether `capability` is currently subscribed. + /// Whether `capability` is currently subscribed. Test-support. + #[cfg(test)] pub fn is_subscribed(&self, capability: Capability) -> bool { self.active_capabilities().contains(&capability) } From fca76778b37dbad703fcea45ea055d96b6e27e31 Mon Sep 17 00:00:00 2001 From: aayank13 Date: Fri, 31 Jul 2026 04:02:03 +0530 Subject: [PATCH 11/22] feat: enhance error handling and add API token management in the dashboard --- mofa-engine-core/src/backends/liter_llm.rs | 11 +- mofa-engine-core/src/backends/ollama.rs | 99 ++++++++++--- .../src/backends/openai_compat.rs | 86 +++++++++-- mofa-engine-core/src/engine.rs | 133 ++++++++++++++---- mofa-engine-core/src/metrics.rs | 7 +- mofa-engine-sdk/src/client.rs | 88 ++++++++++++ mofa-engine-sdk/src/dashboard.rs | 67 ++++++--- mofa-engine-sdk/src/server.rs | 47 ++++++- mofa-kernel/src/types.rs | 4 + 9 files changed, 453 insertions(+), 89 deletions(-) diff --git a/mofa-engine-core/src/backends/liter_llm.rs b/mofa-engine-core/src/backends/liter_llm.rs index 0b26cd6..a9ca1e1 100644 --- a/mofa-engine-core/src/backends/liter_llm.rs +++ b/mofa-engine-core/src/backends/liter_llm.rs @@ -542,17 +542,22 @@ impl Provider for LiterLLMProvider { continue; }; // Thought-chain deltas are forwarded distinctly from answer text so - // callers can display/audit the reasoning separately (S2). + // callers can display/audit the reasoning separately (S2). A send + // error means the receiver was dropped (client disconnected), so we + // stop consuming the upstream stream rather than paying for unread tokens. if let Some(rc) = choice.delta.reasoning_content && !rc.is_empty() + && sink.send(StreamDelta::Reasoning(rc)).await.is_err() { - let _ = sink.send(StreamDelta::Reasoning(rc)).await; + break; } if let Some(content) = choice.delta.content && !content.is_empty() { full.push_str(&content); - let _ = sink.send(StreamDelta::Text(content)).await; + if sink.send(StreamDelta::Text(content)).await.is_err() { + break; + } } } diff --git a/mofa-engine-core/src/backends/ollama.rs b/mofa-engine-core/src/backends/ollama.rs index 032a845..f0bdd73 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -3,6 +3,7 @@ //! Communicates with a local Ollama instance via its HTTP API. use async_trait::async_trait; +use base64::Engine as _; use mofa_kernel::{ BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, @@ -127,6 +128,11 @@ struct OllamaChatRequest { struct OllamaMessage { role: String, content: String, + /// Base64-encoded image data for multimodal (vision) models such as `llava`. + /// Ollama expects raw base64 (no `data:` prefix); empty for text-only turns, + /// in which case the field is omitted entirely. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + images: Vec, } #[derive(Debug, Deserialize)] @@ -168,6 +174,56 @@ impl OllamaProvider { } serde_json::from_slice(line).ok() } + + /// Convert MoFA conversation messages to Ollama's wire form, resolving any + /// attached image references to the raw base64 a vision model (e.g. `llava`) + /// expects. A message with no images serializes as a plain text turn. + /// + /// Note: Ollama models are discovered as [`Capability::Chat`], so the router + /// only sends chat requests here; attaching images lets a locally-pulled + /// vision model still answer an image-bearing chat turn on-device. + async fn to_ollama_messages(&self, request: &InferenceRequest) -> Vec { + let mut out = Vec::with_capacity(request.messages.len()); + for m in &request.messages { + let mut images = Vec::with_capacity(m.images.len()); + for image in &m.images { + if let Some(b64) = Self::resolve_image_b64(image).await { + images.push(b64); + } + } + out.push(OllamaMessage { + role: m.role.clone(), + content: m.content.clone(), + images, + }); + } + out + } + + /// Resolve one image reference to the raw base64 payload Ollama accepts. + /// A `data:` URL is stripped to its payload; a local file is read and + /// encoded; an `http(s)` URL is skipped (Ollama cannot fetch remote images) + /// with a warning. Unreadable references are skipped rather than failing the + /// whole request. + async fn resolve_image_b64(image: &str) -> Option { + if image.starts_with("data:") { + // `data:;base64,` → keep only the base64 payload. + return image + .split_once(',') + .map(|(_, payload)| payload.to_string()); + } + if image.starts_with("http://") || image.starts_with("https://") { + tracing::warn!("Ollama cannot fetch remote image URL '{image}'; skipping"); + return None; + } + match tokio::fs::read(image).await { + Ok(bytes) => Some(base64::engine::general_purpose::STANDARD.encode(bytes)), + Err(e) => { + tracing::warn!("cannot read local image '{image}': {e}; skipping"); + None + } + } + } } #[async_trait] @@ -228,6 +284,12 @@ impl Provider for OllamaProvider { return None; } + // TODO: Ollama's `/api/tags` does not report modality, so every + // non-embedding model is typed as `Chat`. A locally-pulled vision + // model therefore won't be routed for a `Vlm` capability request, + // though it can still answer an image-bearing *chat* turn (see + // `to_ollama_messages`). Modality inference from the model family is + // deferred. let mut card = ModelCard::new( self.name.clone(), model_name.clone(), @@ -274,6 +336,7 @@ impl Provider for OllamaProvider { messages: vec![OllamaMessage { role: "user".into(), content: " ".into(), + images: Vec::new(), }], stream: false, keep_alive: Some("5m".into()), @@ -314,6 +377,7 @@ impl Provider for OllamaProvider { messages: vec![OllamaMessage { role: "user".into(), content: " ".into(), + images: Vec::new(), }], stream: false, keep_alive: Some("0".into()), @@ -355,15 +419,7 @@ impl Provider for OllamaProvider { ) -> Result { let model_name = ModelId::name(model_id); - let messages: Vec = request - .messages - .iter() - .map(|m| OllamaMessage { - role: m.role.clone(), - content: m.content.clone(), - }) - .collect(); - + let messages = self.to_ollama_messages(request).await; if messages.is_empty() { return Err(EngineError::InvalidRequest("no messages provided".into())); } @@ -436,14 +492,7 @@ impl Provider for OllamaProvider { sink: StreamSink, ) -> Result { let model_name = ModelId::name(model_id); - let messages: Vec = request - .messages - .iter() - .map(|m| OllamaMessage { - role: m.role.clone(), - content: m.content.clone(), - }) - .collect(); + let messages = self.to_ollama_messages(request).await; if messages.is_empty() { return Err(EngineError::InvalidRequest("no messages provided".into())); } @@ -497,17 +546,23 @@ impl Provider for OllamaProvider { Some(content) }; - while let Some(bytes) = resp.chunk().await.map_err(|e| EngineError::ProviderError { - provider: self.name.clone(), - detail: format!("stream read error: {e}"), - })? { + 'read: while let Some(bytes) = + resp.chunk().await.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("stream read error: {e}"), + })? + { buf.extend_from_slice(&bytes); while let Some(pos) = buf.iter().position(|&b| b == b'\n') { let line: Vec = buf.drain(..=pos).collect(); if let Some(chunk) = Self::parse_stream_line(&line) && let Some(delta) = apply(chunk, &mut full) + && sink.send(StreamDelta::Text(delta)).await.is_err() { - let _ = sink.send(StreamDelta::Text(delta)).await; + // The receiver was dropped (client disconnected): stop draining + // the upstream generation rather than paying for tokens no one + // will read. + break 'read; } } } diff --git a/mofa-engine-core/src/backends/openai_compat.rs b/mofa-engine-core/src/backends/openai_compat.rs index b0b07a1..a806d2f 100644 --- a/mofa-engine-core/src/backends/openai_compat.rs +++ b/mofa-engine-core/src/backends/openai_compat.rs @@ -162,14 +162,21 @@ enum SseEvent { impl OpenAiCompatProvider { /// Parse one SSE event block (the bytes preceding a blank-line delimiter). /// - /// Concatenates any `data:` field lines, ignoring comments/`event:` lines and - /// tolerating `\r\n`. Blank blocks and unparseable payloads yield `None`. + /// Joins the event's `data:` field lines with newlines per the SSE spec + /// (a single leading space after `data:` is stripped), ignoring + /// comments/`event:` lines and tolerating `\r\n` (`str::lines` drops the + /// terminator). Blank blocks and unparseable payloads yield `None`. fn parse_sse_event(block: &[u8]) -> Option { let text = std::str::from_utf8(block).ok()?; let mut data = String::new(); for line in text.lines() { - if let Some(rest) = line.trim_start().strip_prefix("data:") { - data.push_str(rest.trim()); + if let Some(rest) = line.strip_prefix("data:") { + if !data.is_empty() { + data.push('\n'); + } + // Per the spec a single space immediately after the colon is part + // of the delimiter, not the value; any further whitespace is content. + data.push_str(rest.strip_prefix(' ').unwrap_or(rest)); } } if data.is_empty() { @@ -182,6 +189,25 @@ impl OpenAiCompatProvider { .ok() .map(SseEvent::Chunk) } + + /// Length (in bytes) of the first complete SSE event frame in `buf`, + /// including its blank-line delimiter, or `None` if no complete frame is + /// buffered yet. Recognizes both `\n\n` and the spec-permitted `\r\n\r\n` + /// (as `\n` followed by `\r\n`), so a CRLF-emitting server or proxy streams + /// correctly rather than buffering the whole body. + fn sse_frame_len(buf: &[u8]) -> Option { + for i in 0..buf.len() { + if buf[i] != b'\n' { + continue; + } + match buf.get(i + 1) { + Some(b'\n') => return Some(i + 2), + Some(b'\r') if buf.get(i + 2) == Some(&b'\n') => return Some(i + 3), + _ => {} + } + } + None + } } #[derive(Debug, Serialize)] @@ -412,17 +438,21 @@ impl Provider for OpenAiCompatProvider { Some(content) }; - while let Some(bytes) = resp.chunk().await.map_err(|e| EngineError::ProviderError { - provider: self.name.clone(), - detail: format!("stream read error: {e}"), - })? { + 'read: while let Some(bytes) = + resp.chunk().await.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("stream read error: {e}"), + })? + { buf.extend_from_slice(&bytes); - while let Some(idx) = buf.windows(2).position(|w| w == b"\n\n") { - let block: Vec = buf.drain(..idx + 2).collect(); + while let Some(frame_len) = Self::sse_frame_len(&buf) { + let block: Vec = buf.drain(..frame_len).collect(); if let Some(event) = Self::parse_sse_event(&block) && let Some(delta) = apply(event, &mut full) + && sink.send(StreamDelta::Text(delta)).await.is_err() { - let _ = sink.send(StreamDelta::Text(delta)).await; + // Receiver dropped (client disconnected): stop draining upstream. + break 'read; } } } @@ -840,4 +870,38 @@ mod tests { assert!(OpenAiCompatProvider::parse_sse_event(b": keep-alive comment").is_none()); assert!(OpenAiCompatProvider::parse_sse_event(b"data: not json").is_none()); } + + #[test] + fn sse_frame_len_detects_lf_and_crlf_delimiters() { + // `\n\n` frame: `"data: a\n\n"` is 9 bytes (the drained frame incl. both + // newlines), leaving `"rest"`. + assert_eq!( + OpenAiCompatProvider::sse_frame_len(b"data: a\n\nrest"), + Some(9) + ); + // `\r\n\r\n` frame (CRLF server/proxy): `"data: a\r\n\r\n"` is 11 bytes. + assert_eq!( + OpenAiCompatProvider::sse_frame_len(b"data: a\r\n\r\nrest"), + Some(11) + ); + // A lone trailing `\n` is not a complete frame yet. + assert_eq!(OpenAiCompatProvider::sse_frame_len(b"data: a\n"), None); + } + + #[test] + fn parse_sse_event_joins_multiple_data_lines() { + // Two `data:` lines are joined with a newline, yielding valid JSON. + let block = b"data: {\"choices\":[{\"delta\":\ndata: {\"content\":\"hi\"}}]}"; + let Some(SseEvent::Chunk(chunk)) = OpenAiCompatProvider::parse_sse_event(block) else { + panic!("expected a parsed chunk from joined data lines"); + }; + assert_eq!( + chunk + .choices + .into_iter() + .next() + .and_then(|c| c.delta.content), + Some("hi".to_string()) + ); + } } diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index 6afb071..a014755 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -11,9 +11,10 @@ use std::time::{Duration, Instant}; use dashmap::DashMap; use mofa_kernel::{ - BackendHealth, BackendStatus, Capability, CostTier, EngineError, EngineEvent, EngineStatus, - FailedAttempt, FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, ModelId, - ModelResidency, Provider, ProviderHealth, ProviderKind, StreamChunk, StreamDelta, + BackendHealth, BackendStatus, Capability, CostTier, DataClass, EngineError, EngineEvent, + EngineStatus, FailedAttempt, FallbackPolicy, InferenceRequest, InferenceResponse, ModelCard, + ModelId, ModelResidency, Prefer, Provider, ProviderHealth, ProviderKind, StreamChunk, + StreamDelta, }; use serde::{Deserialize, Serialize}; use tokio::sync::{Mutex as AsyncMutex, Semaphore, broadcast, mpsc}; @@ -162,7 +163,10 @@ impl Engine { /// authoritative token count only exists post-call); a floor of 1 avoids a /// zero estimate for an empty prompt. fn estimate_prompt_tokens(req: &InferenceRequest) -> u32 { - let chars: usize = req.messages.iter().map(|m| m.content.len()).sum(); + // Count Unicode scalar values, not bytes: a byte count would over-count + // multi-byte (e.g. CJK) prompts several-fold and could spuriously price a + // request out of its `max_cost_usd` ceiling. + let chars: usize = req.messages.iter().map(|m| m.content.chars().count()).sum(); ((chars / 4) as u32).max(1) } @@ -410,23 +414,37 @@ impl Engine { 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; + // Reconcile with the engine's current view. Capture the + // residency *and* the live execution counters so a + // rediscovery that lands mid-request does not reset an + // in-flight model's `active_requests` (which would corrupt + // its Busy status and capacity score, and let a later + // `end_execution` saturate to zero). `.map` drops the shard + // guard immediately so the later insert cannot deadlock. + let previous = self + .models + .get(&card.id) + .map(|c| (c.residency, c.execution.active_requests)); + if let Some((prev_residency, active_requests)) = previous { + // Freshly discovered cards start with zero in-flight + // requests; carry the live count forward. The new card's + // `max_concurrency` is kept so config changes still apply. + card.execution.active_requests = active_requests; + match prev_residency { + // A load is in flight; don't let discovery regress it. + ModelResidency::Loading => { + card.residency = ModelResidency::Loading; + } + // The backend reports the model is no longer resident, + // so release the reservation we still held for it. + ModelResidency::Loaded + if !matches!(card.residency, ModelResidency::Loaded) => + { + self.memory.deallocate(&card.id); + freed = true; + } + _ => {} } - // 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); @@ -527,15 +545,12 @@ impl Engine { ); // Data-flow audit (S5): record where each request was actually // served so residency is traceable by request/session over time. - tracing::info!( - target: "mofa::audit", - request_id = %resp.request_id, - session_id = req_session.as_deref().unwrap_or(""), - data_class = ?req_data_class, - prefer = ?req_prefer, - provider = %resp.provider, - locality = if is_local { "local" } else { "cloud" }, - "request served" + self.audit_served( + &resp.request_id, + &resp.provider, + req_prefer, + req_data_class, + req_session.as_deref(), ); } Err(_) => self.metrics.record_request(false, duration_ms, false), @@ -583,6 +598,30 @@ impl Engine { .is_some_and(|p| p.kind.is_local()) } + /// Emit the S5 data-flow audit record for one served request: where it ran + /// and under what residency constraint. Both the unary [`invoke`](Self::invoke) + /// and the streaming path call this so *every* served request — not just + /// non-streamed ones — is traceable by request/session over time. + fn audit_served( + &self, + request_id: &str, + provider: &str, + prefer: Prefer, + data_class: DataClass, + session_id: Option<&str>, + ) { + tracing::info!( + target: "mofa::audit", + request_id = %request_id, + session_id = session_id.unwrap_or(""), + data_class = ?data_class, + prefer = ?prefer, + provider = %provider, + locality = if self.provider_is_local(provider) { "local" } else { "cloud" }, + "request served" + ); + } + /// Run inference. /// /// Builds a single ranked candidate plan and walks it in order: the first @@ -906,6 +945,21 @@ impl Engine { let routing_reason = decision.reason.clone(); let overall_deadline = Instant::now() + self.timeouts.request(); + // Budget ceiling parity with `invoke`: streaming targets a single + // candidate, so if that candidate's estimated cost exceeds the request's + // ceiling the stream fails up front rather than spending. Free/local + // candidates estimate to $0 and are always affordable. + if let Some(ceiling) = req.max_cost_usd { + let est = self.estimate_request_cost(&provider_name, &req); + if est > ceiling { + let err = EngineError::NoCapableModel(format!( + "estimated cost ${est:.4} exceeds budget ${ceiling:.4}" + )); + let _ = out.send(StreamChunk::Error(err.info())).await; + return; + } + } + if !self.circuit_breakers.allow_request(&provider_name) { let err = EngineError::CircuitOpen(provider_name.clone()); let _ = out.send(StreamChunk::Error(err.info())).await; @@ -1005,6 +1059,15 @@ impl Engine { resp.completion_tokens, cost_usd, ); + // Data-flow audit (S5): streamed requests are traced exactly like + // unary ones, so the residency guarantee has no streaming blind spot. + self.audit_served( + &req.request_id, + &provider_name, + req.prefer, + req.data_class, + req.session_id.as_deref(), + ); StreamChunk::Completed { duration_ms, tokens_used: resp.tokens_used, @@ -1145,7 +1208,17 @@ impl Engine { .get(model_id) .and_then(|card| self.find_provider(&card.provider)); if let Some(provider) = provider { - let _ = provider.unload(model_id).await; + // Bound the backend unload with the load timeout: during memory-pressure + // admission this runs under `load_gate`, so a hung backend unload must + // not wedge every concurrent load. A timed-out (or failed) unload still + // falls through to release the engine's own accounting below. + match tokio::time::timeout(self.timeouts.load(), provider.unload(model_id)).await { + Ok(Ok(_)) => {} + Ok(Err(e)) => tracing::warn!("backend unload of '{model_id}' failed: {e}"), + Err(_) => { + tracing::warn!("backend unload of '{model_id}' timed out; releasing accounting") + } + } } // Frees the reservation and emits residency/memory events. self.set_model_residency(model_id, ModelResidency::Unloaded); diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs index afb52c7..a243256 100644 --- a/mofa-engine-core/src/metrics.rs +++ b/mofa-engine-core/src/metrics.rs @@ -290,13 +290,14 @@ impl EngineMetrics { impl EngineMetrics { /// 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. + /// double-quote, and line breaks must be escaped so a provider name cannot + /// produce malformed output. A carriage return is normalized to `\n` (the + /// format defines no `\r` escape) so it cannot terminate the line early. fn escape_label_value(value: &str) -> String { value .replace('\\', "\\\\") .replace('"', "\\\"") - .replace('\n', "\\n") + .replace(['\r', '\n'], "\\n") } } diff --git a/mofa-engine-sdk/src/client.rs b/mofa-engine-sdk/src/client.rs index b3dedc1..80ce261 100644 --- a/mofa-engine-sdk/src/client.rs +++ b/mofa-engine-sdk/src/client.rs @@ -219,6 +219,19 @@ impl DaemonClient { Self::decode(resp).await } + async fn post_json( + &self, + path: &str, + body: serde_json::Value, + ) -> Result { + let resp = self + .authed(self.http.post(self.url(path)).json(&body)) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } + async fn decode( resp: reqwest::Response, ) -> Result { @@ -294,6 +307,81 @@ impl DaemonClient { .map_err(|e| ClientError::Transport(e.to_string()))?; Self::decode(resp).await } + + /// `POST /v1/subscriptions` — register a capability subscription and warm its + /// models; returns the new subscription id. + pub async fn subscribe( + &self, + app_id: Option, + session_id: Option, + capabilities: Vec, + ttl_secs: Option, + ) -> Result { + #[derive(serde::Deserialize)] + struct SubscribeResponse { + id: u64, + } + let body = serde_json::json!({ + "app_id": app_id, + "session_id": session_id, + "capabilities": capabilities, + "ttl_secs": ttl_secs, + }); + let resp: SubscribeResponse = self.post_json("/v1/subscriptions", body).await?; + Ok(resp.id) + } + + /// `DELETE /v1/subscriptions/{id}` — remove a subscription; returns whether + /// one existed. A not-found id is reported as `Ok(false)`, not an error. + pub async fn unsubscribe(&self, id: u64) -> Result { + let resp = self + .authed( + self.http + .delete(self.url(&format!("/v1/subscriptions/{id}"))), + ) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + if resp.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(false); + } + #[derive(serde::Deserialize)] + struct UnsubscribeResponse { + removed: bool, + } + let resp: UnsubscribeResponse = Self::decode(resp).await?; + Ok(resp.removed) + } + + /// `POST /v1/models/load` — manually warm a model; returns whether backend + /// state changed. + pub async fn load_model(&self, model_id: &str) -> Result { + #[derive(serde::Deserialize)] + struct ModelActionResponse { + changed: bool, + } + let body = serde_json::json!({ "model_id": model_id }); + let resp: ModelActionResponse = self.post_json("/v1/models/load", body).await?; + Ok(resp.changed) + } + + /// `POST /v1/models/unload` — manually unload a model; returns whether the + /// model was known. + pub async fn unload_model(&self, model_id: &str) -> Result { + #[derive(serde::Deserialize)] + struct ModelActionResponse { + changed: bool, + } + let body = serde_json::json!({ "model_id": model_id }); + let resp: ModelActionResponse = self.post_json("/v1/models/unload", body).await?; + Ok(resp.changed) + } + + // Streaming endpoints (`POST /v1/invoke/stream` and `GET /v1/events`) are + // intentionally not part of this typed client: SSE consumption is + // caller-specific (framing, backpressure, reconnection), so a caller drives + // them with `reqwest` directly, or uses the in-process `EmbeddedEngine` / + // `Engine::invoke_stream` for a typed `StreamChunk` receiver. } #[cfg(test)] diff --git a/mofa-engine-sdk/src/dashboard.rs b/mofa-engine-sdk/src/dashboard.rs index a3aaee6..f1872bb 100644 --- a/mofa-engine-sdk/src/dashboard.rs +++ b/mofa-engine-sdk/src/dashboard.rs @@ -500,6 +500,7 @@ body {
-- -- +
@@ -605,6 +606,33 @@ body { + + diff --git a/mofa-studio/src/main.rs b/mofa-studio/src/main.rs new file mode 100644 index 0000000..254460f --- /dev/null +++ b/mofa-studio/src/main.rs @@ -0,0 +1,421 @@ +//! mofa-studio — the out-of-the-box creative studio. +//! +//! cargo run -p mofa-studio # then open http://127.0.0.1:8787 +//! +//! This is the "usable, validated, users-directly-experience-it" deliverable: a +//! single binary that boots the MoFA engine and serves a browser UI where anyone +//! types a prompt and gets a **real image or video** back — no SDK, no CLI flags, +//! no code. It is deliberately *not* another example script; it is the framework +//! wearing a face a non-developer can use. +//! +//! ## What it proves about the framework +//! +//! Every generation goes through the ordinary engine contract +//! (`Capability::ImageGen` / `Capability::VideoGen`), so the UI can surface the +//! things that make MoFA more than a thin API wrapper: **which provider served +//! the request, whether it ran local or cloud, and what it cost**. That is the +//! "practical value" made visible. +//! +//! ## Bring-your-own model (the validation path) +//! +//! Offline, chat is local Ollama and there is no zero-config image/video backend, +//! so those buttons report honestly that no backend is configured. Set a cloud +//! key and the same UI renders real media: +//! +//! AGNES_API_KEY=... AGNES_BASE_URL=https://.../v1 cargo run -p mofa-studio +//! +//! `AGNES_*` wires the mentor-provided free omni-modal gateway (Agnes AI) as an +//! OpenAI-compatible provider for chat + image; `ARK_API_KEY` (see the +//! `cloud_video_gen` backend) wires Seedance for video. The model ids are +//! overridable via `AGNES_CHAT_MODEL` / `AGNES_IMAGE_MODEL`. +//! +//! TODO(video): once Agnes's video endpoint shape is confirmed, wire it here the +//! same way (either as a `cloud_video_gen` base_url or a dedicated adapter). + +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; + +use axum::{ + Json, Router, + extract::State, + http::StatusCode, + response::{Html, IntoResponse}, + routing::{get, post}, +}; +use clap::Parser; +use mofa_engine_core::{Engine, EngineConfig}; +use mofa_kernel::{Capability, InferenceRequest, Message, Prefer}; +use serde::{Deserialize, Serialize}; +use tower_http::services::ServeDir; + +/// Prompt → image / video, in the browser, via the MoFA engine. +#[derive(Parser, Debug)] +#[command(name = "mofa-studio", version, about)] +struct Cli { + /// Port for the local web UI. + #[arg(long, default_value_t = 8787)] + port: u16, + + /// Use a specific engine config instead of the auto-provisioned one. + #[arg(short, long)] + config: Option, +} + +#[tokio::main] +async fn main() { + tracing_subscriber::fmt() + .with_env_filter( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "mofa_studio=info,mofa_engine_core=warn".into()), + ) + .init(); + + if let Err(e) = run().await { + eprintln!("error: {e}"); + std::process::exit(1); + } +} + +async fn run() -> Result<(), String> { + let cli = Cli::parse(); + + // The engine writes every artifact into this directory; we then serve it + // statically so the browser can load the produced image/video by URL. + let artifacts_dir = studio_artifacts_dir()?; + + let config_path = match &cli.config { + Some(p) => Some(p.clone()), + None => Some(provision_config(&artifacts_dir)?), + }; + let engine = Engine::try_new(EngineConfig::load(config_path.as_deref())) + .await + .map_err(|e| format!("engine init failed: {e}"))?; + engine.refresh_resources().await; + + let state = AppState { + engine, + artifacts_dir: artifacts_dir.clone(), + }; + + // The whole app: the embedded UI, a tiny JSON API, and static artifact serving. + let app = Router::new() + .route("/", get(index)) + .route("/api/capabilities", get(capabilities)) + .route("/api/generate", post(generate)) + .nest_service("/artifacts", ServeDir::new(&artifacts_dir)) + .with_state(state); + + let addr = SocketAddr::from(([127, 0, 0, 1], cli.port)); + let listener = tokio::net::TcpListener::bind(addr) + .await + .map_err(|e| format!("binding {addr}: {e}"))?; + + println!("\n 🎨 mofa-studio — open http://{addr}\n"); + println!(" artifacts: {}\n", artifacts_dir.display()); + axum::serve(listener, app) + .await + .map_err(|e| format!("server error: {e}")) +} + +// ============================================================================== +// Shared state +// ============================================================================== + +#[derive(Clone)] +struct AppState { + engine: Arc, + artifacts_dir: PathBuf, +} + +// ============================================================================== +// API — capabilities +// ============================================================================== + +/// One row per (capability, provider) the engine can currently serve. The UI uses +/// this to tell the user *before* they click whether image/video is wired up. +#[derive(Serialize)] +struct CapabilityRow { + capability: String, + provider: String, + available: bool, + local: bool, +} + +async fn capabilities(State(state): State) -> Json> { + let rows = state + .engine + .capabilities() + .await + .into_iter() + .map(|c| CapabilityRow { + capability: c.capability.to_string(), + provider: c.provider, + available: matches!( + c.availability, + mofa_kernel::ModelAvailability::Discovered + | mofa_kernel::ModelAvailability::Configured + ), + local: matches!( + c.residency, + mofa_kernel::ModelResidency::Loaded + | mofa_kernel::ModelResidency::Unloaded + | mofa_kernel::ModelResidency::Loading + ), + }) + .collect(); + Json(rows) +} + +// ============================================================================== +// API — generate +// ============================================================================== + +#[derive(Deserialize)] +struct GenerateRequest { + prompt: String, + /// `"image"` or `"video"`. + mode: String, +} + +#[derive(Serialize)] +struct GenerateResponse { + ok: bool, + /// URL of the produced artifact under `/artifacts`, when successful. + url: Option, + provider: Option, + /// `true` when served on-device (cost ≈ $0), `false` for a metered cloud call. + local: Option, + cost_usd: Option, + duration_ms: Option, + /// Human-readable error when `ok` is false (e.g. "no image backend configured"). + error: Option, +} + +async fn generate( + State(state): State, + Json(req): Json, +) -> impl IntoResponse { + let prompt = req.prompt.trim(); + if prompt.is_empty() { + return ( + StatusCode::BAD_REQUEST, + Json(GenerateResponse::err("a prompt is required")), + ); + } + + // Map the UI mode to an engine capability. Both are cloud-heavy, so we let the + // router use cloud (`prefer=auto`) rather than pinning local. + let (capability, params) = match req.mode.as_str() { + "image" => ( + Capability::ImageGen, + serde_json::json!({ "size": "1024x1024" }), + ), + "video" => ( + Capability::VideoGen, + serde_json::json!({ "ratio": "16:9", "resolution": "720p", "duration": 5 }), + ), + other => { + return ( + StatusCode::BAD_REQUEST, + Json(GenerateResponse::err(&format!("unknown mode '{other}'"))), + ); + } + }; + + let request = InferenceRequest { + capability: Some(capability), + messages: vec![Message { + role: "user".into(), + content: prompt.to_string(), + ..Default::default() + }], + params, + prefer: Prefer::Auto, + ..Default::default() + }; + + match state.engine.invoke(request).await { + Ok(resp) => { + let url = resp + .file + .as_deref() + .and_then(|f| artifact_url(&state.artifacts_dir, f)); + // Locality is a property of where the model ran, not of cost (a cloud + // model can bill $0). Read it from the serving provider's residency: + // a remote card means cloud, anything else is on-device. + let local = state + .engine + .capabilities() + .await + .iter() + .find(|c| c.provider == resp.provider) + .map(|c| c.residency != mofa_kernel::ModelResidency::Remote); + ( + StatusCode::OK, + Json(GenerateResponse { + ok: url.is_some(), + error: if url.is_some() { + None + } else { + Some("backend returned no downloadable artifact".into()) + }, + url, + provider: Some(resp.provider), + local, + cost_usd: resp.cost_usd, + duration_ms: Some(resp.duration_ms), + }), + ) + } + Err(e) => { + // A missing backend is the common, expected case offline — report it + // as a clean message the UI can show, not a 500. + let info = e.info(); + ( + StatusCode::OK, + Json(GenerateResponse::err(&format!( + "{:?}: {}", + info.code, info.message + ))), + ) + } + } +} + +impl GenerateResponse { + fn err(msg: &str) -> Self { + Self { + ok: false, + url: None, + provider: None, + local: None, + cost_usd: None, + duration_ms: None, + error: Some(msg.to_string()), + } + } +} + +/// Turn an engine artifact path into a `/artifacts/` URL, but only if the +/// file really lives under our served directory (defence against serving a path +/// a provider wrote elsewhere). +fn artifact_url(dir: &std::path::Path, file: &str) -> Option { + let path = std::path::Path::new(file); + let name = path.file_name()?.to_str()?; + let canonical_dir = std::fs::canonicalize(dir).ok()?; + let canonical_file = std::fs::canonicalize(path).ok()?; + canonical_file + .starts_with(&canonical_dir) + .then(|| format!("/artifacts/{name}")) +} + +// ============================================================================== +// UI +// ============================================================================== + +async fn index() -> Html<&'static str> { + Html(INDEX_HTML) +} + +/// The whole front-end: one self-contained page, no build step, no external +/// assets. Kept intentionally small — the point is to *drive* the engine, not to +/// be a design showcase. +const INDEX_HTML: &str = include_str!("index.html"); + +// ============================================================================== +// Config provisioning +// ============================================================================== + +/// The directory every generated artifact is written to (and served from). Uses +/// the OS cache dir so repeated runs share it and it survives a `cargo clean`. +fn studio_artifacts_dir() -> Result { + let dir = dirs::cache_dir() + .unwrap_or_else(std::env::temp_dir) + .join("mofa-studio") + .join("artifacts"); + std::fs::create_dir_all(&dir).map_err(|e| format!("creating {}: {e}", dir.display()))?; + Ok(dir) +} + +/// Write the engine config: always local Ollama for chat + the studio artifact +/// dir; add the Agnes gateway (chat + image) when `AGNES_API_KEY` is present, so +/// the exact same binary is offline-friendly *and* cloud-validated depending only +/// on the environment. +fn provision_config(artifacts_dir: &std::path::Path) -> Result { + let dir = dirs::cache_dir() + .unwrap_or_else(std::env::temp_dir) + .join("mofa-studio"); + std::fs::create_dir_all(&dir).map_err(|e| format!("creating {}: {e}", dir.display()))?; + + let mut toml = format!( + "# Auto-generated by mofa-studio.\n\ + # Cloud video generation is a slow, poll-until-done task (a clip can take + # several minutes), so the request/inference budgets are widened well past + # the engine defaults — otherwise a video that IS rendering server-side is + # killed mid-poll. `[timeouts]` is in seconds.\n\ + [timeouts]\n\ + request_secs = 900\n\ + inference_secs = 840\n\n\ + [artifacts]\n\ + dir = {artifacts:?}\n\n\ + [[providers]]\n\ + name = \"ollama\"\n\ + kind = \"ollama\"\n\ + base_url = \"http://127.0.0.1:11434\"\n\ + priority = 2\n\ + cost_tier = \"free\"\n", + artifacts = artifacts_dir.to_string_lossy(), + ); + + // Agnes AI (mentor-provided free omni-modal gateway): OpenAI-compatible, so it + // slots in as an `openai_compatible` provider for chat + image. Only added when + // a key is present — otherwise the studio stays a clean offline build. + if let Ok(key) = std::env::var("AGNES_API_KEY") + && !key.trim().is_empty() + { + let base = std::env::var("AGNES_BASE_URL") + .unwrap_or_else(|_| "https://apihub.agnes-ai.com/v1".into()); + let chat = std::env::var("AGNES_CHAT_MODEL").unwrap_or_else(|_| "agnes-2.5-flash".into()); + let image = + std::env::var("AGNES_IMAGE_MODEL").unwrap_or_else(|_| "agnes-image-2.1-flash".into()); + let video = + std::env::var("AGNES_VIDEO_MODEL").unwrap_or_else(|_| "agnes-video-v2.0".into()); + // Chat + image share one OpenAI-compatible provider… + toml.push_str(&format!( + "\n[[providers]]\n\ + name = \"agnes\"\n\ + kind = \"openai_compatible\"\n\ + base_url = {base:?}\n\ + api_key = \"env:AGNES_API_KEY\"\n\ + priority = 1\n\ + cost_tier = \"low\"\n\ + [[providers.models]]\n\ + name = {chat:?}\n\ + capability = \"chat\"\n\ + [[providers.models]]\n\ + name = {image:?}\n\ + capability = \"image_gen\"\n", + )); + // …video is a separate task-based endpoint, so it needs the cloud_video_gen + // backend with the `agnes` dialect (submit /videos → poll → download). + toml.push_str(&format!( + "\n[[providers]]\n\ + name = \"agnes-video\"\n\ + kind = \"cloud_video_gen\"\n\ + dialect = \"agnes\"\n\ + base_url = {base:?}\n\ + api_key = \"env:AGNES_API_KEY\"\n\ + priority = 1\n\ + cost_tier = \"low\"\n\ + [[providers.models]]\n\ + name = {video:?}\n\ + capability = \"video_gen\"\n", + )); + tracing::info!("Agnes gateway configured (chat + image + video)"); + } + + let config = dir.join("config.toml"); + std::fs::write(&config, toml).map_err(|e| format!("writing {}: {e}", config.display()))?; + Ok(config) +}