diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..9d488e8 --- /dev/null +++ b/.env.example @@ -0,0 +1,21 @@ +# Example environment for mofa-studio (and any app using the Agnes / Ark backends). +# +# Copy this to `.env.local` (gitignored), fill in your key, then: +# source .env.local && cargo run -p mofa-studio +# +# `.env.local` is NOT auto-loaded by the app — you `source` it so the values land +# in the process environment, where the engine resolves `api_key = "env:VAR"`. + +# --- Agnes AI (free omni-modal gateway) — chat + image + video --- +# Get a free key at https://agnes-ai.com. This one key enables all three. +export AGNES_API_KEY="" + +# Endpoint + model ids (defaults shown; override only if the docs change): +export AGNES_BASE_URL="https://apihub.agnes-ai.com/v1" +export AGNES_CHAT_MODEL="agnes-2.5-flash" +export AGNES_IMAGE_MODEL="agnes-image-2.1-flash" +export AGNES_VIDEO_MODEL="agnes-video-v2.0" # or "agnes-video-2.5" + +# --- Seedance via Volcengine Ark (alternative video backend) --- +# Only needed if you use the Ark dialect instead of Agnes for video. +# export ARK_API_KEY="" diff --git a/.gitignore b/.gitignore index 7ee12e8..65e5840 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,7 @@ target/ .env +.env.local +.env.*.local *.swp __pycache__/ node_modules/ diff --git a/Cargo.lock b/Cargo.lock index 022e0d5..d494900 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", @@ -661,6 +842,12 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "http-range-header" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9171a2ea8a68358193d15dd5d70c1c10a2afc3e7e4c5bc92bc9f025cebd7359c" + [[package]] name = "httparse" version = "1.10.1" @@ -705,6 +892,7 @@ dependencies = [ "hyper", "hyper-util", "rustls", + "rustls-native-certs", "tokio", "tokio-rustls", "tower-service", @@ -732,7 +920,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 +1102,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 +1202,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" @@ -1036,6 +1303,9 @@ dependencies = [ "clap", "mofa-engine-core", "mofa-engine-sdk", + "mofa-kernel", + "serde", + "serde_json", "tokio", "tracing", "tracing-subscriber", @@ -1047,18 +1317,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", ] @@ -1072,14 +1345,29 @@ dependencies = [ "chrono", "mofa-engine-core", "mofa-kernel", + "reqwest 0.12.28", "serde", "serde_json", + "subtle", "thiserror", "tokio", "tokio-stream", "tower", "tower-http", "tracing", + "uuid", +] + +[[package]] +name = "mofa-explainer" +version = "0.1.0" +dependencies = [ + "clap", + "dirs", + "mofa-engine-core", + "mofa-kernel", + "serde_json", + "tokio", ] [[package]] @@ -1090,9 +1378,52 @@ dependencies = [ "serde", "serde_json", "thiserror", + "tokio", "uuid", ] +[[package]] +name = "mofa-podcast" +version = "0.1.0" +dependencies = [ + "clap", + "dirs", + "mofa-engine-core", + "mofa-kernel", + "reqwest 0.12.28", + "serde_json", + "tokio", +] + +[[package]] +name = "mofa-review" +version = "0.1.0" +dependencies = [ + "clap", + "dirs", + "mofa-engine-core", + "mofa-kernel", + "serde_json", + "tokio", +] + +[[package]] +name = "mofa-studio" +version = "0.1.0" +dependencies = [ + "axum", + "clap", + "dirs", + "mofa-engine-core", + "mofa-kernel", + "serde", + "serde_json", + "tokio", + "tower-http", + "tracing", + "tracing-subscriber", +] + [[package]] name = "native-tls" version = "0.2.18" @@ -1180,7 +1511,7 @@ checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1236,6 +1567,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" @@ -1273,7 +1624,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]] @@ -1285,6 +1645,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" @@ -1307,13 +1687,40 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" [[package]] -name = "rand" -version = "0.9.4" +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]] @@ -1323,7 +1730,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]] @@ -1355,6 +1771,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" @@ -1372,13 +1800,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", @@ -1414,6 +1851,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" @@ -1428,6 +1910,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" @@ -1447,13 +1984,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" @@ -1463,6 +2014,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" @@ -1486,6 +2064,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" @@ -1501,6 +2088,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" @@ -1557,7 +2160,7 @@ checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1593,6 +2196,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" @@ -1641,6 +2253,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" @@ -1681,6 +2309,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" @@ -1709,7 +2348,7 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1746,6 +2385,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" @@ -1776,7 +2421,7 @@ checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1798,6 +2443,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" @@ -1823,7 +2483,7 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -1879,6 +2539,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] @@ -1890,9 +2551,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]] @@ -1904,6 +2580,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" @@ -1912,10 +2597,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]] @@ -1924,6 +2630,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" @@ -1948,10 +2660,19 @@ checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "bitflags", "bytes", + "futures-core", "futures-util", "http", "http-body", + "http-body-util", + "http-range-header", + "httpdate", + "mime", + "mime_guess", + "percent-encoding", "pin-project-lite", + "tokio", + "tokio-util", "tower", "tower-layer", "tower-service", @@ -1991,7 +2712,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2004,6 +2725,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" @@ -2063,7 +2794,7 @@ dependencies = [ "http", "httparse", "log", - "rand", + "rand 0.9.4", "sha1", "thiserror", ] @@ -2151,6 +2882,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" @@ -2193,6 +2934,7 @@ dependencies = [ "cfg-if", "once_cell", "rustversion", + "serde", "wasm-bindgen-macro", "wasm-bindgen-shared", ] @@ -2226,7 +2968,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wasm-bindgen-shared", ] @@ -2261,6 +3003,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" @@ -2283,6 +3038,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" @@ -2299,6 +3063,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" @@ -2348,7 +3121,7 @@ checksum = "9107ddc059d5b6fbfbffdfa7a7fe3e22a226def0b2608f72e9d552763d3e1ad7" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2359,7 +3132,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2370,7 +3143,7 @@ checksum = "29bee4b38ea3cde66011baa44dba677c432a78593e202392d1e9070cf2a7fca7" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2381,7 +3154,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2519,6 +3292,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" @@ -2555,7 +3337,7 @@ dependencies = [ "heck", "indexmap", "prettyplease", - "syn", + "syn 2.0.117", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -2571,7 +3353,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn", + "syn 2.0.117", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -2619,6 +3401,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" @@ -2638,7 +3429,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -2659,7 +3450,7 @@ checksum = "0b631b19d36a892ab55420c92dbc83ccd79274f25be714855d3074aa71cab639" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] @@ -2679,7 +3470,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", "synstructure", ] @@ -2719,7 +3510,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn", + "syn 2.0.117", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 6e7adfc..370277f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,10 @@ members = [ "mofa-engine-core", "mofa-engine-sdk", "mofa-engine-app", + "mofa-podcast", + "mofa-explainer", + "mofa-review", + "mofa-studio", ] [workspace.package] @@ -27,7 +31,7 @@ reqwest = { version = "0.12", features = ["json", "multipart"] } axum = { version = "0.8", features = ["ws"] } axum-extra = { version = "0.10", features = ["typed-header"] } tower = "0.5" -tower-http = { version = "0.6", features = ["cors", "trace"] } +tower-http = { version = "0.6", features = ["cors", "trace", "fs"] } # Error handling thiserror = "2" diff --git a/config.example.toml b/config.example.toml new file mode 100644 index 0000000..52a9882 --- /dev/null +++ b/config.example.toml @@ -0,0 +1,316 @@ +# 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. +# +# Security-relevant environment variables (not TOML keys): +# MOFA_API_TOKEN When set, every `/v1` route requires +# `Authorization: Bearer `. Strongly recommended +# whenever `host` is not loopback. +# MOFA_CORS_ALLOW_ORIGINS Comma-separated origin allowlist for cross-origin +# browser access to `/v1` (e.g. +# "https://app.example.com"). Unset/empty = same-origin +# only, which is the safe default. +# MOFA_LOG_FORMAT=json Emit structured JSON logs (levels honor RUST_LOG). + +[listen] +host = "127.0.0.1" # loopback by default; remote access is opt-in +port = 8420 + +[memory] +# Omit budget_mb to auto-detect (70% of system RAM). +budget_mb = 12288 +idle_timeout_secs = 120 # unload a resident model after this many idle seconds + +[timeouts] +request_secs = 240 # overall budget for one /v1/invoke, across all candidates +queue_secs = 30 # max wait for a concurrency permit +load_secs = 60 # max time to load/warm one model +inference_secs = 180 # max time for one provider inference call +discovery_secs = 8 # per-provider discovery probe +health_secs = 5 # per-provider health probe + +[preflight] +enabled = true # master switch for predictive warming +history_learning = true # learn capability transition history +speculative_warming = true # warm models from hints/history/subscriptions +confidence_threshold = 0.6 # min confidence (0.0-1.0) before a history prediction warms +min_samples = 3 # min observed transitions before history is trusted +subscription_ttl_secs = 3600 + +[artifacts] +# dir = "/tmp" # where generated audio/artifacts are scanned (default: system temp) +retention_secs = 3600 # delete engine-generated artifacts older than this; 0 disables + +[security] +# Allowlist of directory roots that request `input_file` paths must fall under. +# Empty (default) allows any local path — appropriate only on a trusted machine. +# input_roots = ["/Users/me/mofa-workspace"] +input_roots = [] + +# Local Ollama (auto-discovered models). Priority is lower = preferred. +[[providers]] +name = "ollama" +kind = "ollama" +base_url = "http://127.0.0.1:11434" +priority = 1 +cost_tier = "free" + + # 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" +kind = "openai_compatible" +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" + capability = "chat" + context_window = 128000 + + [[providers.models]] + name = "tts-1" + capability = "tts" + + [[providers.models]] + 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" + +# Built-in zero-config voice: when this config declares NO TTS backend, the engine +# automatically registers the operating system's native voice (macOS `say`, Linux +# `espeak-ng`/`espeak`) as a low-priority local fallback — so local audio scenarios +# work with nothing installed. Any TTS backend you declare below outranks it. No +# configuration is needed to use it; it appears as provider `system-tts`. + +# Local TTS process-adapter backend. Runs a local command once per synthesis and +# returns the produced audio file. Recommended stack (S1/S6): Qwen3-TTS + Crane +# as the primary voice, with Kokoro as a lighter fallback — configure both and +# let priority pick Qwen3 first, failing over to Kokoro if it is unavailable. +# 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 +# {voice} voice name from params.voice (empty if unset) +# {speed} speaking rate from params.speed (empty if unset) +# {format} container from params.format; also sets the {output} extension +# +# Primary: Qwen3-TTS via a Crane wrapper CLI. +# [[providers]] +# name = "qwen3-tts" +# kind = "local_tts" +# base_url = "" +# command = "crane-tts" # a wrapper that runs Qwen3-TTS through Crane +# args = ["--text-file", "{text_file}", "--out", "{output}"] +# output_format = "wav" +# priority = 1 # preferred over the Kokoro fallback below +# enabled = false +# [[providers.models]] +# name = "qwen3-tts" +# capability = "tts" +# memory_mb = 2048 +# +# Fallback: Kokoro (lighter, faster to load). +[[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 # fallback tier: used when the Qwen3 primary is down +enabled = false # opt-in: set true once the command is installed + + [[providers.models]] + 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}"] +# 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 + + [[providers.models]] + 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 + +# Local video-generation process-adapter backend. Runs a local text-to-video CLI +# (e.g. an AnimateDiff / Stable-Video-Diffusion or Wan-style wrapper) once per +# clip and returns the produced file (S4). This is the on-device video path; the +# cloud_video_gen backend below is the API path. Verify output with the quality +# gate before presenting it as finished. Placeholders: +# {prompt} {negative_prompt} {width} {height} {seconds} {fps} {output} +# [[providers]] +# name = "local-video" +# kind = "local_video_gen" +# base_url = "" # unused for local_video_gen, but the field is required +# command = "svd" # a text-to-video wrapper script +# args = [ +# "--prompt", "{prompt}", +# "--width", "{width}", "--height", "{height}", +# "--seconds", "{seconds}", "--fps", "{fps}", +# "--output", "{output}", +# ] +# output_format = "mp4" +# priority = 2 +# enabled = false # opt-in: set true once the command is installed +# [[providers.models]] +# name = "svd" +# capability = "video_gen" +# memory_mb = 8192 + +# Cloud video-generation backend — the *API-level* video path (S4). Speaks the +# Volcengine Ark / BytePlus "content generation tasks" contract that ByteDance's +# Seedance models use: the engine submits a task, polls it to completion, then +# downloads the finished clip as a managed artifact — all behind the same +# `video_gen` capability as the local adapter, so routing and `prefer` work the +# same way (this is a *cloud* backend, so `prefer = "local"` skips it). +# +# A `dialect` selects the vendor contract: "ark" (default) or "agnes" (the Agnes +# AI gateway's `/videos` task API). The API key is read from this field or the +# dialect's env var (ARK_API_KEY / AGNES_API_KEY); with no key the provider +# registers but reports itself unavailable, so a keyless config still boots. +# Per-request knobs go in `params`: `ratio` ("16:9"), `resolution` ("1080p"), +# `duration`/`seconds`, `fps`, `seed`, `image_url` (for image-to-video), plus +# `width`/`height`/`num_frames` for the Agnes dialect. +# [[providers]] +# name = "seedance" +# kind = "cloud_video_gen" +# dialect = "ark" # or "agnes" +# base_url = "" # default per dialect; for Ark BytePlus (intl) use +# # "https://ark.ap-southeast.bytepluses.com/api/v3" +# api_key = "env:ARK_API_KEY" +# cost_tier = "high" +# priority = 3 +# enabled = false # opt-in: set true once a key is configured +# [[providers.models]] +# name = "doubao-seedance-1-0-pro-250528" # Ark: a Seedance model id +# capability = "video_gen" +# +# Agnes video variant (free gateway; same AGNES_API_KEY as chat/image): +# [[providers]] +# name = "agnes-video" +# kind = "cloud_video_gen" +# dialect = "agnes" +# base_url = "https://apihub.agnes-ai.com/v1" +# api_key = "env:AGNES_API_KEY" +# cost_tier = "low" +# priority = 3 +# enabled = false +# [[providers.models]] +# name = "agnes-video-v2.0" # or "agnes-video-2.5" +# capability = "video_gen" + +# Local embedding model via Ollama (RAG enabler, S3). Ollama auto-discovers pulled +# models; a model whose name contains `embed` (e.g. `nomic-embed-text`, +# `mxbai-embed-large`) is typed as an `embedding` capability automatically — just +# `ollama pull nomic-embed-text`. Cloud embeddings work too: add an +# `openai_compatible` or `liter_llm` model with `capability = "embedding"` +# (e.g. `text-embedding-3-small`). Send input via `params.input` (a string or an +# array) or the message contents; the response carries `embedding` vectors. diff --git a/docs/acceptance.md b/docs/acceptance.md deleted file mode 100644 index 5766815..0000000 --- a/docs/acceptance.md +++ /dev/null @@ -1,28 +0,0 @@ -# MoFA Engine Acceptance Tracking - -This file converts the RFC acceptance checklist into implementation evidence. - -## Stage 0-2 Baseline - -- [x] Rust CI gates exist for formatting, Clippy, and tests. -- [x] Deterministic mock-backed core tests do not require Ollama or cloud keys. -- [x] Provider configuration is validated before production startup. -- [x] Provider health is tracked separately from circuit breaker state. -- [x] Manual discovery refresh is available through the engine and HTTP API. -- [x] Stale discovered models are removed during refresh. -- [x] Canonical model IDs use `provider/model`. -- [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. - diff --git a/docs/decisions/0001-stage-0-2-baseline.md b/docs/decisions/0001-stage-0-2-baseline.md deleted file mode 100644 index bcb92e4..0000000 --- a/docs/decisions/0001-stage-0-2-baseline.md +++ /dev/null @@ -1,27 +0,0 @@ -# ADR 0001: Stage 0-2 Engine Baseline - -## Status - -Accepted for the Stage 0-2 implementation baseline. - -## Context - -The prototype proves the shape of MoFA Engine but overloads one model status field, treats HTTP server code as an SDK, silently skips invalid provider configuration, and does not expose a manual discovery refresh path. - -## Decisions - -- Keep `/v1/invoke`, `/v1/capabilities`, `/v1/status`, and `/v1/events` compatible while strengthening the internal contracts. -- Add `/v1/discovery/refresh` for manual resource refresh. -- Use canonical model identifiers in the form `provider/model`. -- Keep short model names only when they resolve uniquely. -- Treat named-model fallback as strict by default; callers must opt into named fallback. -- Split backend health, model availability, residency, and execution state. -- Default the daemon bind address to `127.0.0.1`. -- Make production startup use checked configuration and fail on invalid provider kinds or missing configured secrets. -- Keep cloud lifecycle operations as explicit no-ops returning `ModelResidency::Remote`. - -## Consequences - -- Existing demo clients continue to work with the current request shape. -- The engine can now explain discovery, health, routing, and lifecycle state more precisely. -- Future stages can add memory reservations and richer SDK bindings without changing the basic registry model again. \ No newline at end of file diff --git a/docs/openapi.yaml b/docs/openapi.yaml index d42476d..dcda577 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -2,38 +2,126 @@ 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: deterministic capability routing with failover, + reservation-based memory and lifecycle management, Preflight + (hint/subscription/history) warming, dual-track (local/cloud) observability, + real token streaming, and the stateful Responses API. + + + ## Versioning + + Business endpoints are served under the `/v1` prefix, which is the stability + contract for this API. Operational endpoints (`/`, `/health`, `/metrics`) are + intentionally unversioned, since they are process/infrastructure contracts + rather than part of the evolving API surface. + + + ## Authentication + + When the daemon is started with the `MOFA_API_TOKEN` environment variable, + every `/v1` route requires an `Authorization: Bearer ` header and a + request without a valid token receives `401`. When `MOFA_API_TOKEN` is unset + (appropriate only for a trusted local machine) the `/v1` API is open. The + public operational routes (`/`, `/health`, `/metrics`) are never gated. + + + Every response carries an `x-request-id` header for log correlation; a + client-supplied `x-request-id` (up to 128 chars) is echoed back, otherwise the + engine generates one. +# Bearer auth applies to every operation by default; the public operational +# routes below opt out with an empty `security: []`. +security: + - bearerAuth: [] paths: + /: + get: + summary: Embedded dashboard (HTML) + security: [] + responses: + "200": + description: The dashboard single-page app. + content: + text/html: + schema: + type: string /health: get: summary: Health check + security: [] responses: "200": description: Engine process health. + content: + application/json: + schema: + $ref: "#/components/schemas/HealthResponse" + /metrics: + get: + summary: Prometheus text-exposition metrics + description: > + Bounded-cardinality engine metrics (requests, failures, fallback, + load/unload/eviction counters, a request-latency histogram, per-provider + token/cost totals tagged by locality, and memory / model / provider + gauges). Public — not behind the `/v1` auth gate, so scrapers can reach it. + security: [] + responses: + "200": + description: Prometheus exposition text. + content: + text/plain: + schema: + type: string /v1/capabilities: get: summary: List registered model capabilities responses: "200": description: Registered model cards. + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/ModelCard" + "401": + $ref: "#/components/responses/Unauthorized" /v1/status: get: summary: Get engine status responses: "200": description: Engine, backend, and memory status. + content: + application/json: + schema: + $ref: "#/components/schemas/EngineStatus" + "401": + $ref: "#/components/responses/Unauthorized" /v1/events: get: summary: Subscribe to server-sent engine events responses: "200": description: SSE event stream. + content: + text/event-stream: + schema: + type: string + "401": + $ref: "#/components/responses/Unauthorized" /v1/discovery/refresh: post: summary: Refresh backend health and model discovery responses: "200": description: Updated engine status. + content: + application/json: + schema: + $ref: "#/components/schemas/EngineStatus" + "401": + $ref: "#/components/responses/Unauthorized" /v1/invoke: post: summary: Invoke a model by capability or name @@ -56,14 +144,25 @@ paths: application/json: schema: $ref: "#/components/schemas/ErrorInfo" + "401": + $ref: "#/components/responses/Unauthorized" "404": description: No model satisfies the request. content: application/json: schema: $ref: "#/components/schemas/ErrorInfo" + "500": + description: > + Internal error (provider error, memory pressure, or misconfiguration). + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" "503": - description: Provider circuit is open. + description: > + Provider circuit is open, or every failover candidate was exhausted + (the error's `failed_chain` lists each attempt). content: application/json: schema: @@ -74,8 +173,322 @@ 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` (and, for reasoning-capable + models, `reasoning`) 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" + "401": + $ref: "#/components/responses/Unauthorized" + /v1/responses: + post: + summary: One turn of the stateful Responses API (multi-turn deep reasoning) + description: > + Runs one conversation turn. The engine stores the running message history + keyed by the returned `id`; continue the conversation by sending that id + as `previous_response_id` on the next turn instead of resending history. + The body is a superset of InferenceRequest (all routing knobs apply) plus + `previous_response_id`, `instructions` (system prompt for a new + conversation), and `input` (a single-user-message shorthand). + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ResponsesRequest" + responses: + "200": + description: The reply plus the conversation id to continue from. + content: + application/json: + schema: + $ref: "#/components/schemas/ResponsesResponse" + "400": + description: No input, or an unknown/expired previous_response_id. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + description: No model satisfies the request. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + "500": + description: Internal error while serving the turn. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + "503": + description: Provider circuit open, or every failover candidate exhausted. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + "504": + description: Request timed out. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + /v1/responses/{id}: + get: + summary: Stored message history for a conversation + parameters: + - name: id + in: path + required: true + schema: + type: string + responses: + "200": + description: The stored messages, in order. + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/Message" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + description: Unknown conversation id. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + delete: + summary: Forget a stored conversation + parameters: + - name: id + in: path + required: true + schema: + type: string + responses: + "200": + description: The conversation was removed. + content: + application/json: + schema: + $ref: "#/components/schemas/RemovedResponse" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + description: Unknown conversation id. + content: + application/json: + schema: + $ref: "#/components/schemas/RemovedResponse" + /v1/memory: + get: + summary: Current memory accounting (budget, reservations, leases) + responses: + "200": + description: Memory report. + content: + application/json: + schema: + $ref: "#/components/schemas/MemoryReport" + "401": + $ref: "#/components/responses/Unauthorized" + /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" + "401": + $ref: "#/components/responses/Unauthorized" + /v1/preflight: + get: + summary: Preflight effectiveness counters + responses: + "200": + description: Preflight statistics. + content: + application/json: + schema: + $ref: "#/components/schemas/PreflightStats" + "401": + $ref: "#/components/responses/Unauthorized" + /v1/subscriptions: + get: + summary: List active capability subscriptions + responses: + "200": + description: Active subscriptions. + content: + application/json: + schema: + type: array + items: + $ref: "#/components/schemas/SubscriptionInfo" + "401": + $ref: "#/components/responses/Unauthorized" + 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 + required: [id] + properties: + id: + type: integer + "400": + description: Empty or invalid capability list. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + "401": + $ref: "#/components/responses/Unauthorized" + /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: + $ref: "#/components/schemas/RemovedResponse" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + description: No such subscription. + content: + application/json: + schema: + $ref: "#/components/schemas/RemovedResponse" + /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 model. + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" + "401": + $ref: "#/components/responses/Unauthorized" + "500": + description: Load failed (provider error, memory pressure, or timeout). + 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; `changed` is false when the model 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" + "401": + $ref: "#/components/responses/Unauthorized" components: + securitySchemes: + bearerAuth: + type: http + scheme: bearer + description: > + Required on `/v1` routes when the daemon is started with `MOFA_API_TOKEN`. + responses: + Unauthorized: + description: > + Missing or invalid bearer token (only when the daemon is started with + `MOFA_API_TOKEN`). + content: + application/json: + schema: + $ref: "#/components/schemas/ErrorInfo" schemas: + HealthResponse: + type: object + required: [status, version, uptime_secs] + properties: + status: + type: string + version: + type: string + uptime_secs: + type: integer InferenceRequest: type: object properties: @@ -93,6 +506,53 @@ components: type: string enum: [capability_only, disabled, allow_named] default: capability_only + prefer: + type: string + enum: [auto, local, cloud] + default: auto + description: > + 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 + 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: @@ -102,6 +562,14 @@ components: description: Path on the engine host. params: type: object + description: > + Capability-specific passthrough options. Recognized keys include + `input` (TTS/embedding text; string or array for embedding), + `voice` (TTS), `language` and `diarize` (ASR speaker separation, S1), + `detail` (VLM billing tier), `prompt` / `negative_prompt` / `size` + (`WxH`) / `width` / `height` (ImageGen and VideoGen), `seconds` and + `fps` (VideoGen), and `max_tokens` (also used for the `max_cost_usd` + estimate). hint_next: type: string request_id: @@ -114,6 +582,53 @@ 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 + ResponsesRequest: + type: object + description: > + A superset of InferenceRequest for the stateful Responses API. All + InferenceRequest fields apply (flattened at the top level); the fields + below are additional. + allOf: + - $ref: "#/components/schemas/InferenceRequest" + - type: object + properties: + previous_response_id: + type: string + nullable: true + description: > + Continue the conversation stored under this prior response id. + Omit to start a new conversation; an unknown id is a 400. + instructions: + type: string + nullable: true + description: System prompt, applied only when starting a new conversation. + input: + type: string + nullable: true + description: Shorthand for a single new user message this turn. + ResponsesResponse: + type: object + description: > + A superset of InferenceResponse: the same fields plus the conversation + `id` and `message_count`. + allOf: + - $ref: "#/components/schemas/InferenceResponse" + - type: object + required: [id, message_count] + properties: + id: + type: string + description: Id of this response; pass as previous_response_id to continue. + message_count: + type: integer + description: Total messages now stored in the conversation. InferenceResponse: type: object properties: @@ -121,6 +636,16 @@ components: type: [string, "null"] file: type: [string, "null"] + description: Output file path (TTS, image gen, video gen). + embedding: + type: [array, "null"] + description: > + Embedding vectors for the `embedding` capability — one row per input, + in input order. Omitted for non-embedding responses. + items: + type: array + items: + type: number model_used: type: string provider: @@ -131,19 +656,352 @@ 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: 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 + description: > + A piece of the model's thought chain (reasoning-capable models), + emitted distinctly from answer `text` so callers can display or audit + it separately. Only produced when the request sets `reasoning.include`. + required: [type, delta] + properties: + type: { const: reasoning } + delta: { type: string } + - type: object + required: [type, duration_ms, fallback_used] + properties: + 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"] } + - type: object + description: > + Terminal error event, carrying an ErrorInfo body flattened alongside + the `type` tag. + required: [type, code, message, retryable] + properties: + type: { const: error } + code: + type: string + enum: + [ + no_capable_model, + provider_error, + circuit_open, + memory_pressure, + invalid_request, + timeout, + unsupported_operation, + config, + internal, + ] + message: { type: string } + retryable: { type: boolean } + source: { type: [string, "null"] } + failed_chain: + type: array + items: + $ref: "#/components/schemas/FailedAttempt" + routing_reason: { 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 + RemovedResponse: + type: object + required: [removed] + properties: + removed: + type: boolean + FailedAttempt: + type: object + required: [provider, model, reason] + properties: + provider: { type: string } + model: { type: string } + reason: { type: string } ErrorInfo: type: object required: [code, message, retryable] properties: code: type: string + enum: + [ + no_capable_model, + provider_error, + circuit_open, + memory_pressure, + invalid_request, + timeout, + unsupported_operation, + config, + internal, + ] message: type: string retryable: 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: + $ref: "#/components/schemas/FailedAttempt" + routing_reason: + type: string + description: Machine-readable routing reason for the primary candidate. + ExecutionState: + type: object + required: [active_requests, max_concurrency] + properties: + active_requests: + type: integer + max_concurrency: + type: integer + ModelCard: + type: object + properties: + id: + type: string + description: Canonical `provider/model` identifier. + name: + type: string + provider: + type: string + capability: + type: string + enum: [chat, tts, asr, image_gen, video_gen, vlm, embedding] + capabilities: + type: array + items: + type: string + enum: [chat, tts, asr, image_gen, video_gen, vlm, embedding] + status: + type: string + enum: [cold, warming, hot, busy, failed] + availability: + type: string + enum: [discovered, configured, unavailable] + residency: + type: string + enum: [unknown, unloaded, loading, loaded, unloading, remote] + execution: + $ref: "#/components/schemas/ExecutionState" + cost_tier: + type: string + enum: [free, low, medium, high] + context_window: + type: integer + memory_estimate_bytes: + type: integer + reasoning_tier: + type: [string, "null"] + enum: [low, medium, high, null] + description: > + Reasoning tier this model serves, for `reasoning.effort` tier routing. + Null for non-tiered models. + ProviderHealth: + type: object + required: [name, healthy, circuit_state] + properties: + name: + type: string + healthy: + type: boolean + circuit_state: + type: string + BackendStatus: + type: object + required: [name, kind, health, circuit_state, features] + properties: + name: + type: string + kind: + type: string + enum: + [ + ollama, + openai_compatible, + local_tts, + local_asr, + local_image_gen, + local_video_gen, + liter_llm, + ] + health: + type: string + enum: [unknown, healthy, degraded, unavailable] + circuit_state: + type: string + description: closed, open, or half_open. + features: + type: array + items: + type: string + enum: + [ + discovery, + load, + unload, + residency_inspection, + streaming, + memory_reporting, + ] + EngineStatus: + type: object + properties: + total_models: + type: integer + loaded_models: + type: integer + providers: + type: integer + memory_used_bytes: + type: integer + memory_budget_bytes: + type: integer + uptime_secs: + type: integer + provider_health: + type: array + items: + $ref: "#/components/schemas/ProviderHealth" + backends: + type: array + items: + $ref: "#/components/schemas/BackendStatus" + AllocationSnapshot: + type: object + properties: + model_id: + type: string + reserved_bytes: + type: integer + observed_bytes: + type: [integer, "null"] + leases: + type: integer + MemoryReport: + type: object + properties: + used_bytes: + type: integer + budget_bytes: + type: integer + available_bytes: + type: integer + allocations: + type: array + items: + $ref: "#/components/schemas/AllocationSnapshot" + LifecycleRecord: + type: object + properties: + seq: + type: integer + at_ms: + type: integer + description: Milliseconds since engine start (monotonic). + model_id: + type: string + event: + type: string + description: load, unload, evict, idle_unload, load_failed, etc. + detail: + type: [string, "null"] + PreflightStats: + type: object + properties: + warms_started: + type: integer + warms_completed: + type: integer + warms_failed: + type: integer + warms_skipped: + type: integer + predictions: + type: integer + hits: + type: integer + misses: + type: integer + SubscriptionInfo: + type: object + properties: + id: + type: integer + app_id: + type: [string, "null"] + session_id: + type: [string, "null"] + capabilities: + type: array + items: + type: string + enum: [chat, tts, asr, image_gen, video_gen, vlm, embedding] + expires_in_secs: + type: [integer, "null"] + SubscribeRequest: + type: object + required: [capabilities] + properties: + app_id: + type: string + session_id: + type: string + capabilities: + type: array + items: + type: string + enum: [chat, tts, asr, image_gen, video_gen, vlm, embedding] + ttl_secs: + type: integer + description: Optional lifetime in seconds; omitted means until removed. diff --git a/mofa-engine-app/Cargo.toml b/mofa-engine-app/Cargo.toml index 6e7bc16..66ad1a9 100644 --- a/mofa-engine-app/Cargo.toml +++ b/mofa-engine-app/Cargo.toml @@ -6,9 +6,12 @@ license.workspace = true description = "MoFA Engine binary entry point" [dependencies] +mofa-kernel = { workspace = true } mofa-engine-core = { workspace = true } mofa-engine-sdk = { workspace = true } tokio = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } clap = { workspace = true } tracing = { workspace = true } tracing-subscriber = { workspace = true } diff --git a/mofa-engine-app/src/main.rs b/mofa-engine-app/src/main.rs index 2f163c8..a75dc06 100644 --- a/mofa-engine-app/src/main.rs +++ b/mofa-engine-app/src/main.rs @@ -1,59 +1,198 @@ //! 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, Server}; +use mofa_kernel::{Capability, InferenceRequest, Message}; use std::path::PathBuf; -use tracing_subscriber::{EnvFilter, fmt}; +use tracing_subscriber::{EnvFilter, prelude::*}; /// MoFA Engine — multimodal AI model orchestration #[derive(Parser, Debug)] #[command(name = "mofa-engine", version, about)] struct Cli { /// Path to config.toml (default: auto-detect) - #[arg(short, long)] + #[arg(short, long, global = true)] config: Option, - /// Override listen port + /// Override listen port (serve mode) #[arg(short, long)] port: Option, + + #[command(subcommand)] + command: Option, +} + +/// Base URL argument shared by the daemon-facing subcommands. +#[derive(clap::Args, Debug)] +struct DaemonArgs { + /// Base URL of the running engine daemon. + #[arg(long, default_value = "http://127.0.0.1:8420")] + url: String, + /// Bearer token for a secured daemon (defaults to $MOFA_API_TOKEN). + #[arg(long)] + token: Option, +} + +#[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, + }, } #[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(); - - let cli = Cli::parse(); - - // Load and validate configuration. - let mut config = EngineConfig::load_checked(cli.config.as_deref())?; - - // CLI port override - if let Some(port) = cli.port { - config.listen.port = port; + Cli::init_tracing(); + Cli::parse().run().await +} + +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, + 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() + }; + Self::print_json(Self::daemon_client(daemon).invoke(&request).await) + } + } } - let host = config.listen.host.clone(); - let port = config.listen.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); - tracing::info!("MoFA Engine v{} starting", env!("CARGO_PKG_VERSION")); + 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(); + } + } + + /// 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, + } + } - // Create engine - let engine = Engine::try_new(config).await?; + /// 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; - // Start server - start_server(engine, &host, port) - .await - .map_err(|e| anyhow::anyhow!("{e}"))?; + 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}")) + } - 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/Cargo.toml b/mofa-engine-core/Cargo.toml index 0dcc507..6cec2e7 100644 --- a/mofa-engine-core/Cargo.toml +++ b/mofa-engine-core/Cargo.toml @@ -22,3 +22,10 @@ sysinfo = { workspace = true } dashmap = { workspace = true } dirs = { workspace = true } tempfile = { workspace = true } +tokio-stream = { workspace = true } +# Pinned exactly: liter-llm is a young crate whose provider-routing surface we +# depend on directly, so upgrades are a deliberate reviewed step (re-vet the +# changelog and our routing assumptions) rather than a floating semver bump. +# Vetted 1.11.0 for the initial engine landing. +liter-llm = "=1.11.0" +base64 = "0.23.0" diff --git a/mofa-engine-core/src/artifacts.rs b/mofa-engine-core/src/artifacts.rs new file mode 100644 index 0000000..1679cfd --- /dev/null +++ b/mofa-engine-core/src/artifacts.rs @@ -0,0 +1,170 @@ +//! 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 one artifact directory — by default a +//! mofa-owned subdirectory of the system temp dir ([`default_artifact_dir`]), +//! never the shared temp dir itself. 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, and the mofa-owned default keeps even those deletions away from +//! other tenants of a world-writable shared temp dir. + +use std::path::PathBuf; +use std::time::{Duration, SystemTime}; + +/// Name prefix stamped on every engine-generated artifact. +pub(crate) const ARTIFACT_PREFIX: &str = "mofa_"; + +/// Directory the engine writes artifacts into when none is configured: a +/// mofa-OWNED subdirectory of the system temp dir — never the shared temp +/// dir itself. The [`ArtifactSweeper`] deletes by name prefix, and sweeping +/// a world-writable shared directory would happily remove another tenant's +/// (or another app's) `mofa_*` files. Writers and the sweeper resolve +/// through this one function so they always agree on the location. +pub(crate) fn default_artifact_dir() -> PathBuf { + std::env::temp_dir().join("mofa_artifacts") +} + +/// Resolve a configured artifact directory, falling back to +/// [`default_artifact_dir`] when unset (or blank). +pub(crate) fn resolve_artifact_dir(configured: Option) -> PathBuf { + configured + .filter(|s| !s.trim().is_empty()) + .map(PathBuf::from) + .unwrap_or_else(default_artifact_dir) +} + +/// Resolve the artifact directory and create it on a best-effort basis, so a +/// writer into a not-yet-existing default subdirectory does not fail. A +/// creation failure is logged, not fatal — the write itself will surface the +/// error with more context if the directory truly cannot be used. +pub(crate) fn ensure_artifact_dir(configured: Option) -> PathBuf { + let dir = resolve_artifact_dir(configured); + if let Err(e) = std::fs::create_dir_all(&dir) { + tracing::warn!( + dir = %dir.display(), + error = %e, + "failed to create artifact directory" + ); + } + dir +} + +/// Deletes stale engine artifacts from a directory. +#[derive(Debug, Clone)] +pub(crate) struct ArtifactSweeper { + dir: PathBuf, + retention: Duration, +} + +impl ArtifactSweeper { + /// Create a sweeper for `dir` (defaulting to [`default_artifact_dir`], a + /// mofa-owned subdirectory of the system temp) that removes engine + /// artifacts older than `retention`. + pub(crate) fn new(dir: Option, retention: Duration) -> Self { + Self { + dir: dir.unwrap_or_else(default_artifact_dir), + retention, + } + } + + /// The retention window; artifacts at least this old are removed. + pub(crate) 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(crate) 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); + } + + #[test] + fn default_artifact_dir_is_a_mofa_owned_subdirectory() { + // #4 review: the default must never be the shared system temp dir + // itself — the sweeper deletes by name prefix, and sweeping a + // world-writable shared directory would delete any other tenant's + // `mofa_*` files. + let dir = resolve_artifact_dir(None); + assert_eq!(dir, std::env::temp_dir().join("mofa_artifacts")); + assert_ne!(dir, std::env::temp_dir()); + // A configured dir is honored verbatim. + assert_eq!( + resolve_artifact_dir(Some("/tmp/custom-mofa".into())), + PathBuf::from("/tmp/custom-mofa") + ); + } + + #[test] + fn sweeper_defaults_to_the_mofa_owned_subdirectory() { + let sweeper = ArtifactSweeper::new(None, Duration::ZERO); + assert_eq!(sweeper.dir, std::env::temp_dir().join("mofa_artifacts")); + } +} diff --git a/mofa-engine-core/src/backends.rs b/mofa-engine-core/src/backends.rs index b4332fa..94d7a2a 100644 --- a/mofa-engine-core/src/backends.rs +++ b/mofa-engine-core/src/backends.rs @@ -1,7 +1,21 @@ //! Provider backend implementations. -pub mod ollama; -pub mod openai_compat; +pub(crate) mod cloud_video_gen; +pub(crate) mod liter_llm; +pub(crate) mod local_asr; +pub(crate) mod local_image_gen; +pub(crate) mod local_tts; +pub(crate) mod local_video_gen; +pub(crate) mod ollama; +pub(crate) mod openai_compat; +pub(crate) mod system_tts; -pub use ollama::OllamaProvider; -pub use openai_compat::OpenAiCompatProvider; +pub(crate) use cloud_video_gen::CloudVideoGenProvider; +pub(crate) use liter_llm::LiterLLMProvider; +pub(crate) use local_asr::LocalAsrProvider; +pub(crate) use local_image_gen::LocalImageGenProvider; +pub(crate) use local_tts::LocalTtsProvider; +pub(crate) use local_video_gen::LocalVideoGenProvider; +pub(crate) use ollama::OllamaProvider; +pub(crate) use openai_compat::OpenAiCompatProvider; +pub(crate) use system_tts::SystemTtsProvider; diff --git a/mofa-engine-core/src/backends/cloud_video_gen.rs b/mofa-engine-core/src/backends/cloud_video_gen.rs new file mode 100644 index 0000000..326202f --- /dev/null +++ b/mofa-engine-core/src/backends/cloud_video_gen.rs @@ -0,0 +1,931 @@ +//! Cloud video-generation backend (task-based APIs; see [`VideoDialect`]). +//! +//! This is the *API-level* video path the local process-adapter +//! ([`LocalVideoGenProvider`](super::local_video_gen)) could not offer: a real +//! text-to-video (and image-to-video) model reached over HTTP. Two vendor +//! dialects are supported behind one backend: **Ark** (the Volcengine Ark / +//! BytePlus "content generation tasks" contract that ByteDance's **Seedance** +//! models speak) and **Agnes** (the Agnes AI gateway's `/videos` task API). So a +//! scenario can render a genuine clip from a prompt without any local generator +//! installed. +//! +//! ## Why this is a distinct backend +//! +//! Our multi-vendor gateway (`liter-llm`) exposes chat and image endpoints but +//! **no video endpoint**, and Seedance is not synchronous the way chat/image are: +//! generation is a *task* you submit, poll for completion, and then download. +//! That submit → poll → download shape does not fit the chat/image providers, so +//! it lives here behind the same [`Provider`] boundary as every other backend — +//! the engine discovers it, routes to it, and manages it identically. +//! +//! ## API contract +//! +//! Submit a task (`POST {base_url}/contents/generations/tasks`): +//! +//! ```json +//! { "model": "", +//! "content": [ { "type": "text", +//! "text": " --ratio 16:9 --resolution 1080p --duration 5" } ] } +//! ``` +//! +//! For image-to-video an extra `{"type":"image_url","image_url":{"url":…}}` item is +//! appended. The response carries a task `id`. We then poll +//! `GET {base_url}/contents/generations/tasks/{id}` until `status` is `succeeded` +//! (yielding `content.video_url`) or `failed`, and finally download the video into +//! a managed artifact. The produced clip should still clear the +//! [`quality_gate`](crate::quality_gate) before being presented as an S4 artifact. +//! +//! ## Parameters (`request.params`) +//! +//! Seedance takes generation knobs as text-command suffixes appended to the +//! prompt, so we translate the familiar keys into that form: +//! - `ratio` (e.g. `"16:9"`), `resolution` (e.g. `"1080p"`), +//! - `duration` / `seconds` (integer seconds), `fps`, `seed`, +//! - `image_url` — a reference image for image-to-video, +//! - `poll_interval_secs` — override the default poll cadence. + +use std::path::PathBuf; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use mofa_kernel::{ + BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, +}; +use reqwest::Client; + +use crate::config::ModelDef; + +/// Default Ark endpoint (Volcengine, China). Operators pointing at BytePlus +/// (international) or a private gateway override this via `base_url`. +const DEFAULT_ARK_BASE_URL: &str = "https://ark.cn-beijing.volces.com/api/v3"; +/// Default Agnes AI gateway endpoint. +const DEFAULT_AGNES_BASE_URL: &str = "https://apihub.agnes-ai.com/v1"; +/// How often to poll a running task when the request does not say otherwise. +const DEFAULT_POLL_INTERVAL: Duration = Duration::from_secs(3); +/// Upper bound on poll attempts. Generation is slow but not unbounded; the engine +/// also enforces its own inference timeout, whose dropped future cancels the poll +/// loop. This cap is a backstop so a stuck task cannot poll forever. +const MAX_POLL_ATTEMPTS: u32 = 200; + +/// Which vendor's task-based video API this provider speaks. Both share the +/// submit → poll → download shape but differ in endpoint paths, request body, and +/// where the finished URL sits in the poll response. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum VideoDialect { + /// Volcengine Ark / BytePlus (ByteDance Seedance): `/contents/generations/tasks`, + /// prompt carried as a `content` array, URL at `content.video_url`. + Ark, + /// Agnes AI gateway: `/videos`, flat `{prompt,width,height,num_frames,frame_rate}` + /// body, URL at `metadata.url`. + Agnes, +} + +impl VideoDialect { + /// Parse the config `dialect` string; unknown/empty defaults to [`Ark`]. + fn from_str_loose(s: &str) -> Self { + match s.trim().to_ascii_lowercase().as_str() { + "agnes" => Self::Agnes, + _ => Self::Ark, + } + } + + /// The environment variable this dialect resolves a missing key from. + fn api_key_env(self) -> &'static str { + match self { + Self::Ark => "ARK_API_KEY", + Self::Agnes => "AGNES_API_KEY", + } + } + + fn default_base_url(self) -> &'static str { + match self { + Self::Ark => DEFAULT_ARK_BASE_URL, + Self::Agnes => DEFAULT_AGNES_BASE_URL, + } + } +} + +/// A provider that renders video through a task-based cloud API (see [`VideoDialect`]). +pub(crate) struct CloudVideoGenProvider { + /// Display name. + name: String, + /// Which vendor dialect to speak. + dialect: VideoDialect, + /// API root, e.g. `https://ark.cn-beijing.volces.com/api/v3`. + base_url: String, + /// Bearer token. Empty means "no credentials" → the provider reports itself + /// unavailable rather than failing engine startup (mirrors the other cloud + /// backends), so an offline/keyless config still boots. + api_key: String, + /// Configured Seedance model ids this backend serves. + models: Vec, + /// Cost tier applied to this provider's models. + cost_tier: CostTier, + /// Directory for downloaded video artifacts. + output_dir: PathBuf, + /// Shared HTTP client (connection reuse across submit/poll/download). + client: Client, +} + +/// The state of a generation task as read from a poll response. +#[derive(Debug, PartialEq, Eq)] +enum TaskState { + /// Queued or running — keep polling. + Pending, + /// Finished; carries the downloadable video URL. + Succeeded(String), + /// Terminal failure; carries the vendor's reason. + Failed(String), +} + +impl CloudVideoGenProvider { + /// Build a cloud video provider for the given `dialect` (`"ark"` or `"agnes"`). + /// + /// An empty `api_key` falls back to the dialect's environment variable + /// ([`VideoDialect::api_key_env`]); if that too is empty the provider is + /// constructed but stays unavailable, so a config that lists it without a key + /// does not break engine startup. + pub(crate) fn new( + name: impl Into, + dialect: &str, + base_url: impl Into, + api_key: Option, + models: Vec, + cost_tier: CostTier, + output_dir: Option, + ) -> Result { + let name = name.into(); + let dialect = VideoDialect::from_str_loose(dialect); + let base_url = { + let b = base_url.into(); + if b.is_empty() { + dialect.default_base_url().to_string() + } else { + // Tolerate a trailing slash so `{base}/…` never doubles it. + b.trim_end_matches('/').to_string() + } + }; + + let key_env = dialect.api_key_env(); + let api_key = api_key + .filter(|k| !k.is_empty()) + .or_else(|| std::env::var(key_env).ok()) + .unwrap_or_default(); + if api_key.is_empty() { + tracing::warn!( + provider = %name, + "cloud video provider has no {key_env}; marked unavailable" + ); + } + + // Video renders are slow; give the per-request calls a generous read + // timeout. The overall wall-clock is still bounded by MAX_POLL_ATTEMPTS + // and the engine's inference timeout. + let client = Client::builder() + .connect_timeout(Duration::from_secs(10)) + .timeout(Duration::from_secs(120)) + .build() + .map_err(|e| { + EngineError::Config(format!("failed to build cloud video HTTP client: {e}")) + })?; + + Ok(Self { + name, + dialect, + base_url, + api_key, + models, + cost_tier, + output_dir: crate::artifacts::ensure_artifact_dir(output_dir), + client, + }) + } + + /// Whether usable credentials were resolved. + fn is_available(&self) -> bool { + !self.api_key.is_empty() + } + + fn provider_error(&self, detail: impl std::fmt::Display) -> EngineError { + EngineError::ProviderError { + provider: self.name.clone(), + detail: detail.to_string(), + } + } + + /// 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) + }) + } + + // ========================================================================== + // Pure request/response mapping (unit-tested without the network) + // ========================================================================== + + /// Translate the familiar size/duration knobs into the Seedance text-command + /// suffix appended to the prompt (e.g. ` --ratio 16:9 --resolution 1080p`). + /// + /// Only keys actually present are emitted, so we never override a model's + /// defaults with guesses. + fn command_suffix(params: &serde_json::Value) -> String { + let mut suffix = String::new(); + let mut push_str = |flag: &str, key: &str| { + if let Some(v) = params.get(key).and_then(|v| v.as_str()) + && !v.trim().is_empty() + { + suffix.push_str(&format!(" --{flag} {}", v.trim())); + } + }; + push_str("ratio", "ratio"); + push_str("resolution", "resolution"); + + // `duration` is preferred; `seconds` is accepted as an alias so callers can + // reuse the local video-gen vocabulary. + let duration = params + .get("duration") + .or_else(|| params.get("seconds")) + .and_then(serde_json::Value::as_u64); + if let Some(d) = duration { + suffix.push_str(&format!(" --duration {d}")); + } + if let Some(fps) = params.get("fps").and_then(serde_json::Value::as_u64) { + suffix.push_str(&format!(" --fps {fps}")); + } + if let Some(seed) = params.get("seed").and_then(serde_json::Value::as_i64) { + suffix.push_str(&format!(" --seed {seed}")); + } + suffix + } + + /// Pixel dimensions for the Agnes body: an explicit `width`/`height`, a + /// `size` string (`"WxH"`), or a 3:2 720p-ish default the gateway will snap + /// to its nearest preset. + fn dimensions(params: &serde_json::Value) -> (u32, u32) { + if let (Some(w), Some(h)) = ( + Self::positive_u32(params, "width"), + Self::positive_u32(params, "height"), + ) { + return (w, h); + } + if let Some((w, h)) = params + .get("size") + .and_then(|v| v.as_str()) + .and_then(|s| s.split_once(['x', 'X'])) + .and_then(|(w, h)| Some((w.trim().parse().ok()?, h.trim().parse().ok()?))) + { + return (w, h); + } + (1152, 768) + } + + /// Read a strictly-positive `u32` from a params key, or `None`. + fn positive_u32(params: &serde_json::Value, key: &str) -> Option { + params + .get(key) + .and_then(serde_json::Value::as_u64) + .filter(|v| *v > 0) + .map(|v| v as u32) + } + + /// Build the JSON task body for the given [`VideoDialect`]. + fn build_task_body( + dialect: VideoDialect, + model_name: &str, + prompt: &str, + params: &serde_json::Value, + ) -> serde_json::Value { + match dialect { + VideoDialect::Ark => Self::build_ark_body(model_name, prompt, params), + VideoDialect::Agnes => Self::build_agnes_body(model_name, prompt, params), + } + } + + /// Ark/Seedance body: prompt (plus generation knobs as text-command suffixes) + /// carried in a `content` array. Adds an `image_url` item for image-to-video. + fn build_ark_body( + model_name: &str, + prompt: &str, + params: &serde_json::Value, + ) -> serde_json::Value { + let text = format!("{prompt}{}", Self::command_suffix(params)); + let mut content = vec![serde_json::json!({ "type": "text", "text": text })]; + if let Some(url) = params.get("image_url").and_then(|v| v.as_str()) + && !url.is_empty() + { + content.push(serde_json::json!({ + "type": "image_url", + "image_url": { "url": url } + })); + } + serde_json::json!({ "model": model_name, "content": content }) + } + + /// Agnes body: a flat object with explicit pixel dimensions and frame count. + /// The familiar `ratio`/`resolution`/`duration` knobs are translated into + /// `width`/`height`/`num_frames` (Agnes snaps to its nearest size preset). + fn build_agnes_body( + model_name: &str, + prompt: &str, + params: &serde_json::Value, + ) -> serde_json::Value { + let (width, height) = Self::dimensions(params); + let frame_rate = Self::positive_u32(params, "fps") + .or_else(|| Self::positive_u32(params, "frame_rate")) + .unwrap_or(24); + // Prefer an explicit num_frames; otherwise derive from duration × fps (+1, + // since generators count an inclusive final frame). Default ~5s. + let num_frames = Self::positive_u32(params, "num_frames").unwrap_or_else(|| { + let seconds = params + .get("duration") + .or_else(|| params.get("seconds")) + .and_then(serde_json::Value::as_u64) + .filter(|s| *s > 0) + .unwrap_or(5) as u32; + seconds * frame_rate + 1 + }); + serde_json::json!({ + "model": model_name, + "prompt": prompt, + "width": width, + "height": height, + "num_frames": num_frames, + "frame_rate": frame_rate, + }) + } + + /// Extract the task id from a submit response. + fn parse_task_id(body: &serde_json::Value) -> Result { + body.get("id") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(String::from) + .ok_or_else(|| format!("submit response missing task id: {body}")) + } + + /// Interpret a poll response into a [`TaskState`]. + /// + /// Unknown/absent statuses are treated as still-pending so a transient shape + /// we do not recognise makes us wait rather than crash — the poll cap and + /// engine timeout remain the safety net. + fn parse_task_state(body: &serde_json::Value) -> TaskState { + let status = body + .get("status") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_ascii_lowercase(); + match status.as_str() { + "succeeded" | "success" | "done" | "completed" => { + // The finished URL sits at `content.video_url` (Ark), a top-level + // `video_url`, or `metadata.url` (Agnes) — accept any of them. + let url = body + .get("content") + .and_then(|c| c.get("video_url")) + .or_else(|| body.get("video_url")) + .or_else(|| body.get("metadata").and_then(|m| m.get("url"))) + .and_then(|v| v.as_str()); + match url { + Some(u) if !u.is_empty() => TaskState::Succeeded(u.to_string()), + _ => TaskState::Failed("task succeeded but returned no video URL".into()), + } + } + "failed" | "error" | "cancelled" | "canceled" => { + let reason = body + .get("error") + .and_then(|e| e.get("message")) + .or_else(|| body.get("metadata").and_then(|m| m.get("message"))) + .and_then(|v| v.as_str()) + .unwrap_or("task failed without a reason"); + TaskState::Failed(reason.to_string()) + } + _ => TaskState::Pending, + } + } + + // ========================================================================== + // HTTP orchestration (submit → poll → download) + // ========================================================================== + + /// Endpoint paths differ per dialect (Ark nests under `contents/generations`). + fn tasks_path(&self) -> &'static str { + match self.dialect { + VideoDialect::Ark => "contents/generations/tasks", + VideoDialect::Agnes => "videos", + } + } + + async fn submit_task(&self, body: &serde_json::Value) -> Result { + let url = format!("{}/{}", self.base_url, self.tasks_path()); + let resp = self + .client + .post(&url) + .bearer_auth(&self.api_key) + .json(body) + .send() + .await + .map_err(|e| self.provider_error(format!("submit request failed: {e}")))?; + let json = self.read_json(resp, "submit").await?; + Self::parse_task_id(&json).map_err(|e| self.provider_error(e)) + } + + async fn poll_task(&self, task_id: &str) -> Result { + let url = format!("{}/{}/{task_id}", self.base_url, self.tasks_path()); + let resp = self + .client + .get(&url) + .bearer_auth(&self.api_key) + .send() + .await + .map_err(|e| self.provider_error(format!("poll request failed: {e}")))?; + self.read_json(resp, "poll").await + } + + /// Read a JSON body, turning a non-2xx status into a descriptive error that + /// includes the (truncated) response body — the vendor puts the actionable + /// reason there. + async fn read_json( + &self, + resp: reqwest::Response, + stage: &str, + ) -> Result { + let status = resp.status(); + let text = resp + .text() + .await + .map_err(|e| self.provider_error(format!("reading {stage} response body: {e}")))?; + if !status.is_success() { + let mut body = text; + body.truncate(500); + return Err(self.provider_error(format!("{stage} returned HTTP {status}: {body}"))); + } + serde_json::from_str(&text) + .map_err(|e| self.provider_error(format!("parsing {stage} response JSON: {e}"))) + } + + /// Download the finished clip into a managed artifact and return its path. + async fn download(&self, video_url: &str) -> Result { + let resp = self + .client + .get(video_url) + .send() + .await + .map_err(|e| self.provider_error(format!("downloading video: {e}")))?; + if !resp.status().is_success() { + return Err( + self.provider_error(format!("video download returned HTTP {}", resp.status())) + ); + } + let bytes = resp + .bytes() + .await + .map_err(|e| self.provider_error(format!("reading video bytes: {e}")))?; + if bytes.is_empty() { + return Err(self.provider_error("video download was empty")); + } + let path = self + .output_dir + .join(format!("mofa_video_{}.mp4", uuid::Uuid::new_v4())); + tokio::fs::write(&path, &bytes) + .await + .map_err(|e| EngineError::Internal(format!("video write error: {e}")))?; + Ok(path.to_string_lossy().to_string()) + } + + async fn generate( + &self, + model_name: &str, + request: &InferenceRequest, + start: Instant, + ) -> Result { + if !self.is_available() { + return Err(self.provider_error(format!( + "no {} credentials configured", + self.dialect.api_key_env() + ))); + } + + 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("video_gen requires a prompt".into()))?; + + let body = Self::build_task_body(self.dialect, model_name, &prompt, &request.params); + let task_id = self.submit_task(&body).await?; + tracing::info!(provider = %self.name, %task_id, "submitted video generation task"); + + let interval = request + .params + .get("poll_interval_secs") + .and_then(serde_json::Value::as_u64) + .filter(|s| *s > 0) + .map(Duration::from_secs) + .unwrap_or(DEFAULT_POLL_INTERVAL); + + // Poll until the task reaches a terminal state or we exhaust the backstop. + let mut video_url = None; + for _ in 0..MAX_POLL_ATTEMPTS { + tokio::time::sleep(interval).await; + let poll = self.poll_task(&task_id).await?; + match Self::parse_task_state(&poll) { + TaskState::Pending => continue, + TaskState::Succeeded(url) => { + video_url = Some(url); + break; + } + TaskState::Failed(reason) => { + return Err( + self.provider_error(format!("video task {task_id} failed: {reason}")) + ); + } + } + } + let video_url = video_url.ok_or_else(|| { + self.provider_error(format!( + "video task {task_id} did not finish within {MAX_POLL_ATTEMPTS} polls" + )) + })?; + + let file = self.download(&video_url).await?; + + Ok(InferenceResponse { + text: None, + file: Some(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() + }) + } +} + +#[async_trait] +impl Provider for CloudVideoGenProvider { + fn name(&self) -> &str { + &self.name + } + + fn kind(&self) -> ProviderKind { + ProviderKind::CloudVideoGen + } + + fn features(&self) -> Vec { + // No memory reporting: the model runs remotely and consumes no local RAM. + vec![ + BackendFeature::Discovery, + BackendFeature::Load, + BackendFeature::Unload, + ] + } + + 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 = ModelId::canonical(&self.name, &m.name); + card.availability = if self.is_available() { + ModelAvailability::Configured + } else { + ModelAvailability::Unavailable + }; + // Cloud models are remote: they hold no local model memory. + card.residency = ModelResidency::Remote; + card.refresh_status(); + Some(card) + }) + .collect(); + Ok(cards) + } + + async fn health(&self) -> Result { + if self.is_available() { + Ok(BackendHealth::Healthy) + } else { + Ok(BackendHealth::Unavailable) + } + } + + async fn load(&self, model_id: &str) -> Result { + // Nothing to warm: remote models are always "loaded" from our side. + Ok(LifecycleResult { + model_id: ModelId::canonical(&self.name, ModelId::name(model_id)), + residency: ModelResidency::Remote, + memory_bytes: Some(0), + changed: false, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(&self.name, ModelId::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 = ModelId::name(model_id); + let capability = request.capability.unwrap_or(Capability::VideoGen); + + if capability != Capability::VideoGen { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' only supports video_gen, not {capability}", + self.name + ))); + } + if !self.model_supports(model_name, Capability::VideoGen) { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' model '{model_name}' does not support video_gen", + self.name + ))); + } + + self.generate(model_name, request, Instant::now()).await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mofa_kernel::Message; + + fn video_models() -> Vec { + vec![ModelDef { + name: "seedance-1-0-pro".into(), + capability: "video_gen".into(), + context_window: None, + memory_mb: None, + ..Default::default() + }] + } + + fn provider(api_key: Option) -> CloudVideoGenProvider { + CloudVideoGenProvider::new( + "ark", + "ark", + "", + api_key, + video_models(), + CostTier::High, + Some(std::env::temp_dir().to_string_lossy().to_string()), + ) + .expect("HTTP client builds") + } + + fn video_request(prompt: &str) -> InferenceRequest { + InferenceRequest { + capability: Some(Capability::VideoGen), + messages: vec![Message { + role: "user".into(), + content: prompt.into(), + ..Default::default() + }], + request_id: "test".into(), + ..Default::default() + } + } + + #[test] + fn empty_base_url_defaults_per_dialect() { + let ark = provider(Some("k".into())); + assert_eq!(ark.base_url, DEFAULT_ARK_BASE_URL); + assert_eq!(ark.dialect, VideoDialect::Ark); + assert_eq!(ark.kind(), ProviderKind::CloudVideoGen); + // Cloud: not counted as a local backend for routing/memory. + assert!(!ark.kind().is_local()); + + let agnes = CloudVideoGenProvider::new( + "agnes-video", + "agnes", + "", + Some("k".into()), + video_models(), + CostTier::Low, + None, + ) + .unwrap(); + assert_eq!(agnes.base_url, DEFAULT_AGNES_BASE_URL); + assert_eq!(agnes.dialect, VideoDialect::Agnes); + } + + #[test] + fn trailing_slash_in_base_url_is_trimmed() { + let p = CloudVideoGenProvider::new( + "ark", + "ark", + "https://example.com/api/v3/", + Some("k".into()), + video_models(), + CostTier::High, + None, + ) + .unwrap(); + assert_eq!(p.base_url, "https://example.com/api/v3"); + } + + #[test] + fn command_suffix_only_emits_present_keys() { + let params = serde_json::json!({ + "ratio": "16:9", + "resolution": "1080p", + "duration": 5, + "fps": 24, + "seed": 42 + }); + assert_eq!( + CloudVideoGenProvider::command_suffix(¶ms), + " --ratio 16:9 --resolution 1080p --duration 5 --fps 24 --seed 42" + ); + // `seconds` is accepted as an alias for `duration`. + assert_eq!( + CloudVideoGenProvider::command_suffix(&serde_json::json!({ "seconds": 8 })), + " --duration 8" + ); + // Nothing present → no suffix, so we never override model defaults. + assert_eq!( + CloudVideoGenProvider::command_suffix(&serde_json::Value::Null), + "" + ); + } + + #[test] + fn ark_body_text_only() { + let body = CloudVideoGenProvider::build_task_body( + VideoDialect::Ark, + "seedance-1-0-pro", + "a rocket launch", + &serde_json::json!({ "ratio": "16:9" }), + ); + assert_eq!(body["model"], "seedance-1-0-pro"); + let content = body["content"].as_array().unwrap(); + assert_eq!(content.len(), 1); + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[0]["text"], "a rocket launch --ratio 16:9"); + } + + #[test] + fn ark_body_image_to_video_appends_image_item() { + let body = CloudVideoGenProvider::build_task_body( + VideoDialect::Ark, + "seedance-1-0-pro", + "make it move", + &serde_json::json!({ "image_url": "https://img/x.png" }), + ); + let content = body["content"].as_array().unwrap(); + assert_eq!(content.len(), 2); + assert_eq!(content[1]["type"], "image_url"); + assert_eq!(content[1]["image_url"]["url"], "https://img/x.png"); + } + + #[test] + fn agnes_body_is_flat_with_derived_frames() { + // duration 5 × fps 24 (+1) = 121 frames; explicit width/height pass through. + let body = CloudVideoGenProvider::build_task_body( + VideoDialect::Agnes, + "agnes-video-v2.0", + "a paper boat", + &serde_json::json!({ "width": 1152, "height": 768, "duration": 5, "fps": 24 }), + ); + assert_eq!(body["model"], "agnes-video-v2.0"); + assert_eq!(body["prompt"], "a paper boat"); + assert_eq!(body["width"], 1152); + assert_eq!(body["height"], 768); + assert_eq!(body["frame_rate"], 24); + assert_eq!(body["num_frames"], 121); + // Defaults when nothing is specified. + let d = CloudVideoGenProvider::build_task_body( + VideoDialect::Agnes, + "agnes-video-v2.0", + "x", + &serde_json::Value::Null, + ); + assert_eq!(d["width"], 1152); + assert_eq!(d["height"], 768); + assert_eq!(d["frame_rate"], 24); + assert_eq!(d["num_frames"], 121); + } + + #[test] + fn parse_task_id_reads_id_or_errors() { + assert_eq!( + CloudVideoGenProvider::parse_task_id(&serde_json::json!({ "id": "cgt-1" })).unwrap(), + "cgt-1" + ); + assert!(CloudVideoGenProvider::parse_task_id(&serde_json::json!({})).is_err()); + } + + #[test] + fn parse_task_state_covers_terminal_and_pending() { + // Success with a nested video URL. + assert_eq!( + CloudVideoGenProvider::parse_task_state(&serde_json::json!({ + "status": "succeeded", + "content": { "video_url": "https://v/clip.mp4" } + })), + TaskState::Succeeded("https://v/clip.mp4".into()) + ); + // Success without a URL is a failure, not a hang. + assert!(matches!( + CloudVideoGenProvider::parse_task_state(&serde_json::json!({ "status": "succeeded" })), + TaskState::Failed(_) + )); + // Explicit failure carries the vendor reason. + assert_eq!( + CloudVideoGenProvider::parse_task_state(&serde_json::json!({ + "status": "failed", + "error": { "message": "nsfw" } + })), + TaskState::Failed("nsfw".into()) + ); + // Running / unknown → keep polling. + assert_eq!( + CloudVideoGenProvider::parse_task_state(&serde_json::json!({ "status": "running" })), + TaskState::Pending + ); + assert_eq!( + CloudVideoGenProvider::parse_task_state(&serde_json::json!({})), + TaskState::Pending + ); + } + + #[test] + fn parse_task_state_agnes_shape() { + // Agnes: `completed` status, URL nested under `metadata.url`. + assert_eq!( + CloudVideoGenProvider::parse_task_state(&serde_json::json!({ + "status": "completed", + "metadata": { "url": "https://out/clip.mp4" } + })), + TaskState::Succeeded("https://out/clip.mp4".into()) + ); + // Agnes progress states keep polling. + assert_eq!( + CloudVideoGenProvider::parse_task_state( + &serde_json::json!({ "status": "in_progress", "progress": 30 }) + ), + TaskState::Pending + ); + } + + #[tokio::test] + async fn keyless_provider_is_unavailable_and_refuses_to_generate() { + // Ensure the env fallback can't accidentally make this available. + // SAFETY: single-threaded test; no other thread reads the env concurrently. + unsafe { + std::env::remove_var("ARK_API_KEY"); + } + let p = provider(None); + assert!(!p.is_available()); + assert_eq!(p.health().await.unwrap(), BackendHealth::Unavailable); + let cards = p.discover().await.unwrap(); + assert_eq!(cards[0].availability, ModelAvailability::Unavailable); + let err = p + .invoke("ark/seedance-1-0-pro", &video_request("x")) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + } + + #[tokio::test] + async fn rejects_non_video_capability() { + let p = provider(Some("k".into())); + let mut req = video_request("x"); + req.capability = Some(Capability::Chat); + let err = p.invoke("ark/seedance-1-0-pro", &req).await.unwrap_err(); + assert!(matches!(err, EngineError::UnsupportedOperation(_))); + } + + #[tokio::test] + async fn empty_prompt_is_invalid_request() { + let p = provider(Some("k".into())); + let err = p + .invoke("ark/seedance-1-0-pro", &video_request(" ")) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + } +} 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..2518a66 --- /dev/null +++ b/mofa-engine-core/src/backends/liter_llm.rs @@ -0,0 +1,805 @@ +//! 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 +//! [`ModelId::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, EmbeddingInput, EmbeddingRequest as LiterEmbeddingRequest, + ImageDetail, LlmClient, Message as LiterMessage, SystemMessage, UserContent, UserMessage, +}; +use mofa_kernel::{ + BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, Reasoning, ReasoningEffort, StreamDelta, StreamSink, +}; +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(crate) 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(crate) 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: crate::artifacts::ensure_artifact_dir(output_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| Self::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: Self::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) = Self::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() + }) + } + + /// Text embedding (S3/RAG enabler): embed one or more inputs in a single + /// call. Inputs come from `params.input` (string or array) or, failing that, + /// the message contents — shared with the OpenAI-compatible backend so both + /// resolve inputs identically. + async fn invoke_embedding( + &self, + model_name: &str, + request: &InferenceRequest, + start: Instant, + ) -> Result { + let inputs = + crate::backends::openai_compat::OpenAiCompatProvider::embedding_inputs(request); + if inputs.is_empty() { + return Err(EngineError::InvalidRequest( + "embedding requires text input (params.input or messages)".into(), + )); + } + + let req = LiterEmbeddingRequest { + model: model_name.to_string(), + input: EmbeddingInput::Multiple(inputs), + ..Default::default() + }; + + let mut resp = self + .client()? + .embed(req) + .await + .map_err(|e| self.provider_error(e))?; + + // Providers may return rows out of batch order; restore input order. + resp.data.sort_by_key(|d| d.index); + let vectors: Vec> = resp.data.into_iter().map(|d| d.embedding).collect(); + if vectors.is_empty() { + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "embedding response contained no vectors".into(), + }); + } + + // Embeddings bill input tokens only (completion is always zero). + let (prompt_tokens, tokens_used) = match resp.usage { + Some(u) => { + let prompt = u.prompt_tokens as u32; + let total = if u.total_tokens > 0 { + u.total_tokens as u32 + } else { + prompt + }; + (Some(prompt), Some(total)) + } + None => (None, None), + }; + + Ok(InferenceResponse { + text: None, + file: None, + embedding: Some(vectors), + 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: Some(0), + 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(Self::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( + Self::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: Self::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) = Self::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() + }) + } +} + +/// 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() + }), + } + } + + /// 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 = ModelId::canonical(&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: ModelId::canonical(&self.name, ModelId::name(model_id)), + residency: ModelResidency::Remote, + memory_bytes: Some(0), + changed: false, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(&self.name, ModelId::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 = ModelId::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, + Capability::Embedding => self.invoke_embedding(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 = ModelId::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). 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() + { + break; + } + if let Some(content) = choice.delta.content + && !content.is_empty() + { + full.push_str(&content); + if sink.send(StreamDelta::Text(content)).await.is_err() { + break; + } + } + } + + let (prompt_tokens, completion_tokens, tokens_used) = Self::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 ModelId::name recovers the + // `provider/model` string liter-llm routes on. + assert_eq!(cards[0].id, "liter/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); + } + + #[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) = 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!(LiterLLMProvider::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!( + LiterLLMProvider::to_liter_effort(mk(ReasoningEffort::Low)), + Some(L::Low) + )); + assert!(matches!( + LiterLLMProvider::to_liter_effort(mk(ReasoningEffort::Medium)), + Some(L::Medium) + )); + assert!(matches!( + LiterLLMProvider::to_liter_effort(mk(ReasoningEffort::High)), + Some(L::High) + )); + assert!(LiterLLMProvider::to_liter_effort(None).is_none()); + } + + #[tokio::test] + async fn image_url_and_data_url_pass_through() { + assert_eq!( + LiterLLMProvider::image_to_url("https://example.com/a.png") + .await + .unwrap(), + "https://example.com/a.png" + ); + assert_eq!( + LiterLLMProvider::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 = LiterLLMProvider::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..12779b4 --- /dev/null +++ b/mofa-engine-core/src/backends/local_asr.rs @@ -0,0 +1,567 @@ +//! 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`). +//! +//! 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. + +use std::path::{Path, PathBuf}; + +use async_trait::async_trait; +use mofa_kernel::{ + BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, +}; +use tokio::process::Command; + +use crate::config::ModelDef; + +/// A process-adapter provider that shells out to a local ASR command. +pub(crate) struct LocalAsrProvider { + /// Display name. + name: String, + /// Program to execute per transcription. + 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. + models: Vec, +} + +impl LocalAsrProvider { + /// Create a new local ASR process adapter. + pub(crate) fn new( + name: impl Into, + command: impl Into, + args: Vec, + output_dir: Option, + models: Vec, + ) -> Self { + Self { + name: name.into(), + command: command.into(), + args, + diarize_args: Vec::new(), + output_dir: crate::artifacts::ensure_artifact_dir(output_dir), + models, + } + } + + /// 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(crate) 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 { + 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"); + + // 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() + .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())) + }); + let output_str = output_file + .as_ref() + .map(|p| p.to_string_lossy().to_string()) + .unwrap_or_default(); + + let resolved_args = self.resolve_args(input_file, &output_str, language, diarize); + + 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 = 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; + 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 = ModelId::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: ModelId::canonical(&self.name, model_name), + residency: ModelResidency::Loaded, + memory_bytes: estimate, + changed: true, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(&self.name, ModelId::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 = ModelId::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 { .. })); + } + + #[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()); + 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_image_gen.rs b/mofa-engine-core/src/backends/local_image_gen.rs new file mode 100644 index 0000000..c32cfa2 --- /dev/null +++ b/mofa-engine-core/src/backends/local_image_gen.rs @@ -0,0 +1,526 @@ +//! 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, ModelId, ModelResidency, + Provider, ProviderKind, +}; +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(crate) 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(crate) 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: crate::artifacts::ensure_artifact_dir(output_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) = Self::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() + }) + } +} + +/// 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), + ) + } + + /// 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 = 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; + 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 = ModelId::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: ModelId::canonical(&self.name, model_name), + residency: ModelResidency::Loaded, + memory_bytes: estimate, + changed: true, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(&self.name, ModelId::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 = ModelId::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!( + 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!( + LocalImageGenProvider::image_dimensions(&serde_json::json!({ "size": "640x480" })), + (640, 480) + ); + assert_eq!( + LocalImageGenProvider::image_dimensions( + &serde_json::json!({ "width": 256, "height": 128 }) + ), + (256, 128) + ); + assert_eq!( + LocalImageGenProvider::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/local_tts.rs b/mofa-engine-core/src/backends/local_tts.rs new file mode 100644 index 0000000..b8eef31 --- /dev/null +++ b/mofa-engine-core/src/backends/local_tts.rs @@ -0,0 +1,544 @@ +//! 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, ModelId, ModelResidency, + Provider, ProviderKind, +}; +use tokio::process::Command; + +use crate::config::ModelDef; + +/// A process-adapter provider that shells out to a local TTS command. +pub(crate) struct LocalTtsProvider { + /// Display name. + name: String, + /// Program to execute per synthesis. + command: String, + /// Argument template with `{text}`, `{text_file}`, `{output}`, `{voice}`, + /// `{speed}`, and `{format}` 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(crate) 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: crate::artifacts::ensure_artifact_dir(output_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()))?; + + // Per-request voice/speed/format (S1/S6). The request's `params` can select + // a voice and speaking rate, and choose the output container; each flows to + // the command via its matching placeholder. `format` also picks the artifact + // extension when set, so a caller asking for mp3 gets `…​.mp3`, not the + // backend default. Absent params substitute to empty (a no-op in the args). + let voice = request + .params + .get("voice") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_string(); + // `speed` may arrive as a number (1.5) or a string ("1.5"); normalize both + // to a plain string for substitution. + let speed = request + .params + .get("speed") + .map(|v| { + v.as_str() + .map(String::from) + .unwrap_or_else(|| v.to_string()) + }) + .unwrap_or_default(); + let format = request + .params + .get("format") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .unwrap_or(&self.output_format) + .to_string(); + + let output_path = + self.output_dir + .join(format!("mofa_tts_{}.{}", uuid::Uuid::new_v4(), 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) + .replace("{voice}", &voice) + .replace("{speed}", &speed) + .replace("{format}", &format) + }) + .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, + ..Default::default() + }) + } +} + +#[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 = 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; + 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 = ModelId::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: ModelId::canonical(&self.name, model_name), + residency: ModelResidency::Loaded, + memory_bytes: estimate, + changed: true, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(&self.name, ModelId::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 = ModelId::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), + ..Default::default() + }] + } + + 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(), + ..Default::default() + }], + input_file: None, + params: serde_json::Value::Null, + hint_next: None, + request_id: "test".into(), + ..Default::default() + } + } + + #[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 passes_voice_and_format_through_params() { + let dir = tempfile::tempdir().unwrap(); + // Fixture: write the substituted `{voice}` into `{output}`, so we can assert + // the placeholder reached the command. The request's `format` should also + // override the backend's default extension. + let p = LocalTtsProvider::new( + "local-tts", + "sh", + vec![ + "-c".into(), + "printf '%s' \"$1\" > \"$2\"".into(), + "sh".into(), + "{voice}".into(), + "{output}".into(), + ], + Some("wav".into()), // backend default… + Some(dir.path().to_string_lossy().to_string()), + tts_models(), + ); + + let mut req = tts_request("hello"); + req.params = serde_json::json!({ "voice": "narrator", "format": "mp3" }); + let resp = p.invoke("local-tts/fixture", &req).await.unwrap(); + let file = resp.file.expect("a TTS artifact path"); + + // …request `format=mp3` overrode the `wav` default extension: + assert!(file.ends_with(".mp3"), "got {file}"); + // …and `{voice}` was substituted from params: + assert_eq!(std::fs::read_to_string(&file).unwrap(), "narrator"); + } + + #[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/local_video_gen.rs b/mofa-engine-core/src/backends/local_video_gen.rs new file mode 100644 index 0000000..ce2ec97 --- /dev/null +++ b/mofa-engine-core/src/backends/local_video_gen.rs @@ -0,0 +1,549 @@ +//! Local video-generation process-adapter backend. +//! +//! Runs a configured local command — a text-to-video CLI (e.g. an AnimateDiff / +//! Stable-Video-Diffusion wrapper, or a `Wan`/`mochi`-style generator) — to +//! render a short clip from a text prompt, returning the produced file as a +//! managed artifact. This completes the S4 Creative Generation track: with a +//! local backend an explainer clip can be generated on-device (`prefer=local`, +//! cost ≈ $0); liter-llm exposes no video endpoint, so unlike image generation +//! there is no cloud fallback and this local adapter is the whole path. +//! +//! Device- and runtime-specific concerns stay behind this `Provider` boundary, +//! so the engine treats a local video-gen model like any other backend: it is +//! discovered, memory-managed, warmed, and idle-unloaded. Because generation is +//! typically slow, an operator should size the `inference_secs` timeout +//! accordingly; the child is spawned with `kill_on_drop(true)` so a timeout that +//! drops the future terminates the render rather than leaking it. +//! +//! ## Command contract +//! +//! The command is spawned once per clip (a cold, stateless process). Argument +//! templates may contain: +//! - `{prompt}` — the text prompt (required), +//! - `{output}` — path the command must write the video to (required), +//! - `{negative_prompt}` — from `params.negative_prompt` (default empty), +//! - `{width}` / `{height}` — from `params.size` (`"WxH"`) or +//! `params.width` / `params.height` (default 512×512), +//! - `{seconds}` — target duration from `params.seconds` (default 4), +//! - `{fps}` — target frame rate from `params.fps` (default 16). +//! +//! The produced clip should be verified with the [`quality_gate`](crate::quality_gate) +//! before it is presented as a finished S4 artifact ("no gate, no output"). + +use std::path::{Path, PathBuf}; + +use async_trait::async_trait; +use mofa_kernel::{ + BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, +}; +use tokio::process::Command; + +use crate::config::ModelDef; + +/// Default output dimensions when a request does not specify a size. +const DEFAULT_EDGE: u32 = 512; +/// Default clip length in seconds. +const DEFAULT_SECONDS: u32 = 4; +/// Default frame rate. +const DEFAULT_FPS: u32 = 16; + +/// A process-adapter provider that shells out to a local video-generation command. +pub(crate) struct LocalVideoGenProvider { + /// Display name. + name: String, + /// Program to execute per clip. + command: String, + /// Argument template with `{prompt}`, `{output}`, `{negative_prompt}`, + /// `{width}`, `{height}`, `{seconds}`, and `{fps}` placeholders. + args: Vec, + /// Output video extension/container (e.g. `mp4`, `webm`). + output_format: String, + /// Directory for generated artifacts. + output_dir: PathBuf, + /// Configured models this backend serves. + models: Vec, +} + +impl LocalVideoGenProvider { + /// Create a new local video-generation process adapter. + pub(crate) 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(|| "mp4".into()), + output_dir: crate::artifacts::ensure_artifact_dir(output_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. Pure over the configured + /// template, so placeholder wiring is unit-testable without spawning. + #[allow(clippy::too_many_arguments)] + fn resolve_args( + &self, + prompt: &str, + output: &str, + negative_prompt: &str, + width: u32, + height: u32, + seconds: u32, + fps: 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()) + .replace("{seconds}", &seconds.to_string()) + .replace("{fps}", &fps.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("video_gen requires a prompt".into()))?; + + let negative_prompt = request + .params + .get("negative_prompt") + .and_then(|v| v.as_str()) + .unwrap_or(""); + let (width, height) = Self::dimensions(&request.params); + let seconds = Self::positive_u32(&request.params, "seconds").unwrap_or(DEFAULT_SECONDS); + let fps = Self::positive_u32(&request.params, "fps").unwrap_or(DEFAULT_FPS); + + let output_path = self.output_dir.join(format!( + "mofa_video_{}.{}", + 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, + seconds, + fps, + ); + + 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!( + "video_gen command exited with {}: {}", + output.status, + stderr.trim() + ), + }); + } + + // A successful exit must have produced a non-empty video 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: "video_gen command produced no video 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() + }) + } +} + +/// Pure parameter parsing, grouped as private associated functions. +impl LocalVideoGenProvider { + /// Requested dimensions from `params`, accepting a combined `size` (`"WxH"`) + /// or explicit `width`/`height`, falling back to a square [`DEFAULT_EDGE`]. + fn 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); + } + ( + Self::positive_u32(params, "width").unwrap_or(DEFAULT_EDGE), + Self::positive_u32(params, "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)) + } + + /// Read a strictly-positive `u32` from a params key, or `None`. + fn positive_u32(params: &serde_json::Value, key: &str) -> Option { + params + .get(key) + .and_then(serde_json::Value::as_u64) + .filter(|v| *v > 0) + .map(|v| v as u32) + } +} + +#[async_trait] +impl Provider for LocalVideoGenProvider { + fn name(&self) -> &str { + &self.name + } + + fn kind(&self) -> ProviderKind { + ProviderKind::LocalVideoGen + } + + 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 = 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; + 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!("video_gen command '{}' not found", self.command), + }); + } + let model_name = ModelId::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: ModelId::canonical(&self.name, model_name), + residency: ModelResidency::Loaded, + memory_bytes: estimate, + changed: true, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(&self.name, ModelId::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 = ModelId::name(model_id); + let capability = request.capability.unwrap_or(Capability::VideoGen); + + if capability != Capability::VideoGen { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' only supports video_gen, not {capability}", + self.name + ))); + } + if !self.model_supports(model_name, Capability::VideoGen) { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{}' model '{model_name}' does not support video_gen", + self.name + ))); + } + + self.generate(model_name, request, std::time::Instant::now()) + .await + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mofa_kernel::Message; + + fn video_models() -> Vec { + vec![ModelDef { + name: "fixture".into(), + capability: "video_gen".into(), + context_window: None, + memory_mb: Some(4096), + ..Default::default() + }] + } + + fn video_request(prompt: &str) -> InferenceRequest { + InferenceRequest { + capability: Some(Capability::VideoGen), + messages: vec![Message { + role: "user".into(), + content: prompt.into(), + ..Default::default() + }], + request_id: "test".into(), + ..Default::default() + } + } + + #[test] + fn metadata_and_defaults() { + let p = LocalVideoGenProvider::new("local-video", "sh", vec![], None, None, video_models()); + assert_eq!(p.kind(), ProviderKind::LocalVideoGen); + assert!(p.kind().is_local()); + assert_eq!(p.output_format, "mp4"); + } + + #[test] + fn dimensions_and_positive_params() { + assert_eq!( + LocalVideoGenProvider::dimensions(&serde_json::json!({ "size": "640x480" })), + (640, 480) + ); + assert_eq!( + LocalVideoGenProvider::dimensions(&serde_json::json!({ "width": 256, "height": 128 })), + (256, 128) + ); + assert_eq!( + LocalVideoGenProvider::dimensions(&serde_json::Value::Null), + (DEFAULT_EDGE, DEFAULT_EDGE) + ); + // Zero is rejected as non-positive and falls back to the default. + assert_eq!( + LocalVideoGenProvider::positive_u32(&serde_json::json!({ "fps": 0 }), "fps"), + None + ); + assert_eq!( + LocalVideoGenProvider::positive_u32(&serde_json::json!({ "fps": 24 }), "fps"), + Some(24) + ); + } + + #[test] + fn resolve_args_substitutes_all_placeholders() { + let p = LocalVideoGenProvider::new( + "local-video", + "gen", + vec![ + "--prompt".into(), + "{prompt}".into(), + "-W".into(), + "{width}".into(), + "-H".into(), + "{height}".into(), + "--secs".into(), + "{seconds}".into(), + "--fps".into(), + "{fps}".into(), + "-o".into(), + "{output}".into(), + ], + None, + None, + video_models(), + ); + let args = p.resolve_args("a rocket", "/tmp/o.mp4", "", 512, 512, 4, 16); + assert_eq!( + args, + vec![ + "--prompt", + "a rocket", + "-W", + "512", + "-H", + "512", + "--secs", + "4", + "--fps", + "16", + "-o", + "/tmp/o.mp4" + ] + ); + } + + #[tokio::test] + async fn generates_video_artifact_via_fixture_command() { + let dir = tempfile::tempdir().unwrap(); + // A deterministic "generator" fixture: write the prompt + geometry to the + // output path, proving the placeholders reached the command. + let p = LocalVideoGenProvider::new( + "local-video", + "sh", + vec![ + "-c".into(), + "printf 'MP4:%s:%sx%s@%s' \"$1\" \"$2\" \"$3\" \"$4\" > \"$5\"".into(), + "sh".into(), + "{prompt}".into(), + "{width}".into(), + "{height}".into(), + "{fps}".into(), + "{output}".into(), + ], + Some("mp4".into()), + Some(dir.path().to_string_lossy().to_string()), + video_models(), + ); + + let mut req = video_request("a neuron firing"); + req.params = serde_json::json!({ "size": "320x240", "fps": 24 }); + let resp = p.invoke("local-video/fixture", &req).await.unwrap(); + let file = resp.file.expect("a video artifact path"); + assert!(file.ends_with(".mp4")); + let produced = std::fs::read_to_string(&file).unwrap(); + assert_eq!(produced, "MP4:a neuron firing:320x240@24"); + } + + #[tokio::test] + async fn empty_prompt_is_invalid_request() { + let p = LocalVideoGenProvider::new("local-video", "sh", vec![], None, None, video_models()); + let err = p + .invoke("local-video/fixture", &video_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 = LocalVideoGenProvider::new( + "local-video", + "sh", + vec!["-c".into(), "echo boom >&2; exit 7".into()], + None, + Some(dir.path().to_string_lossy().to_string()), + video_models(), + ); + let err = p + .invoke("local-video/fixture", &video_request("x")) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + assert!(err.retryable()); + } + + #[tokio::test] + async fn rejects_unsupported_capability() { + let p = LocalVideoGenProvider::new("local-video", "sh", vec![], None, None, video_models()); + let mut req = video_request("x"); + req.capability = Some(Capability::Chat); + let err = p.invoke("local-video/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..054a7bf 100644 --- a/mofa-engine-core/src/backends/ollama.rs +++ b/mofa-engine-core/src/backends/ollama.rs @@ -3,41 +3,73 @@ //! 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, ModelResidency, Provider, - ProviderKind, canonical_model_id, model_id_name, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, ReasoningEffort, StreamDelta, StreamSink, }; use reqwest::Client; use serde::{Deserialize, Serialize}; -use std::collections::HashSet; -use std::time::Duration; +use std::collections::{HashMap, HashSet}; +use std::time::{Duration, Instant}; + +use crate::config::ModelDef; /// Provider for a local Ollama instance. -pub struct OllamaProvider { +pub(crate) struct OllamaProvider { /// Display name. name: String, /// Base URL. 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 { /// Create a new Ollama provider. - pub fn new(name: impl Into, base_url: impl Into) -> Self { + /// + /// Fails (rather than panicking or silently dropping the configured + /// timeouts/`no_proxy`) if the system TLS/HTTP stack cannot build a client, + /// so the engine surfaces a clean startup error instead of crashing. + /// + /// `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(crate) fn new( + name: impl Into, + base_url: impl Into, + models: &[ModelDef], + ) -> Result { let client = Client::builder() .no_proxy() .connect_timeout(Duration::from_secs(5)) .timeout(Duration::from_secs(180)) .build() - .unwrap_or_default(); + .map_err(|e| EngineError::Config(format!("failed to build Ollama HTTP client: {e}")))?; + + 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(); - Self { + Ok(Self { name: name.into(), base_url: base_url.into(), client, - } + reasoning_tiers, + }) } async fn loaded_models(&self) -> HashSet { @@ -96,18 +128,193 @@ 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)] #[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, } +/// Request to Ollama's `/api/embed` endpoint (batch-capable). +#[derive(Debug, Serialize)] +struct OllamaEmbedRequest { + model: String, + input: Vec, +} + +/// Response from Ollama's `/api/embed` endpoint. +#[derive(Debug, Deserialize)] +struct OllamaEmbedResponse { + /// One vector per input, in input order. + #[serde(default)] + embeddings: Vec>, + /// Prompt/input tokens, when reported. + #[serde(default)] + prompt_eval_count: 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, +} + +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() + } + + /// 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 + } + + /// Embed one or more inputs via Ollama's `/api/embed`. Inputs are resolved + /// the same way as the cloud backends (`params.input` string/array, else the + /// message contents), so embedding requests are portable across providers. + async fn embed( + &self, + model_name: &str, + request: &InferenceRequest, + ) -> Result { + let inputs = + crate::backends::openai_compat::OpenAiCompatProvider::embedding_inputs(request); + if inputs.is_empty() { + return Err(EngineError::InvalidRequest( + "embedding requires text input (params.input or messages)".into(), + )); + } + + let body = OllamaEmbedRequest { + model: model_name.to_string(), + input: inputs, + }; + let url = format!("{}/api/embed", self.base_url); + let start = std::time::Instant::now(); + let 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!("embed HTTP {status}: {text}"), + }); + } + + let embed: OllamaEmbedResponse = + resp.json().await.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("embed parse error: {e}"), + })?; + if embed.embeddings.is_empty() { + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "embed response contained no vectors".into(), + }); + } + + Ok(InferenceResponse { + text: None, + file: None, + embedding: Some(embed.embeddings), + 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: embed.prompt_eval_count, + prompt_tokens: embed.prompt_eval_count, + completion_tokens: Some(0), + fallback_used: false, + routing_reason: None, + ..Default::default() + }) + } + + /// 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] impl Provider for OllamaProvider { fn name(&self) -> &str { @@ -162,17 +369,33 @@ impl Provider for OllamaProvider { .filter_map(|m| { let model_name = m.name.or(m.model)?; let lower = model_name.to_lowercase(); - if lower.contains("embed") || lower.contains(":cloud") || lower.contains("-cloud") { + // Ollama's cloud-hosted tags are not local models; skip them. + if lower.contains(":cloud") || lower.contains("-cloud") { return None; } + // The tags endpoint reports no modality, so we infer capability + // from the name: models named `*embed*` serve `Embedding` + // (e.g. `nomic-embed-text`, `mxbai-embed-large`) and everything + // else is typed `Chat`. + // + // TODO: this same lack of modality metadata means a locally-pulled + // *vision* model is still typed `Chat`, so it won't be routed for an + // explicit `Vlm` request (though it can answer an image-bearing chat + // turn — see `to_ollama_messages`). Family-based modality inference + // is deferred. + let capability = if lower.contains("embed") { + Capability::Embedding + } else { + Capability::Chat + }; let mut card = ModelCard::new( self.name.clone(), model_name.clone(), - Capability::Chat, + capability, 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 @@ -181,6 +404,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) }) @@ -203,12 +429,13 @@ 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 { role: "user".into(), content: " ".into(), + images: Vec::new(), }], stream: false, keep_alive: Some("5m".into()), @@ -235,7 +462,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, @@ -243,12 +470,13 @@ 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 { role: "user".into(), content: " ".into(), + images: Vec::new(), }], stream: false, keep_alive: Some("0".into()), @@ -276,7 +504,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, @@ -288,17 +516,15 @@ 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 - .iter() - .map(|m| OllamaMessage { - role: m.role.clone(), - content: m.content.clone(), - }) - .collect(); + // Embedding is a distinct endpoint (`/api/embed`), so dispatch it before + // building a chat body. + if request.capability == Some(Capability::Embedding) { + return self.embed(model_name, request).await; + } + let messages = self.to_ollama_messages(request).await; if messages.is_empty() { return Err(EngineError::InvalidRequest("no messages provided".into())); } @@ -341,6 +567,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), @@ -349,9 +581,129 @@ 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 = ModelId::name(model_id); + + // Embedding is a distinct, non-incremental endpoint (`/api/embed`): there + // are no token deltas to stream, so fall back to a single-shot invoke and + // return its vectors — mirroring the cloud backends, whose `stream` also + // delegates non-chat capabilities to `invoke`. + if request.capability == Some(Capability::Embedding) { + return self.embed(model_name, request).await; + } + + let messages = self.to_ollama_messages(request).await; + 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) + }; + + '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() + { + // The receiver was dropped (client disconnected): stop draining + // the upstream generation rather than paying for tokens no one + // will read. + break 'read; + } + } + } + // Any trailing bytes without a final newline. + if let Some(chunk) = Self::parse_stream_line(&buf) + && let Some(delta) = apply(chunk, &mut full) + { + let _ = sink.send(StreamDelta::Text(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() }) } } @@ -362,15 +714,107 @@ 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)); } + #[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"); - 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] + 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) = OllamaProvider::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!(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 60a1ba9..1349e39 100644 --- a/mofa-engine-core/src/backends/openai_compat.rs +++ b/mofa-engine-core/src/backends/openai_compat.rs @@ -1,12 +1,14 @@ //! Generic OpenAI-compatible provider backend. //! -//! Works with APIs that follow OpenAI-style chat, TTS, and ASR contracts. +//! Works with APIs that follow OpenAI-style chat, TTS, ASR, embedding, and +//! image-generation (`/images/generations`) contracts — e.g. the Agnes AI +//! gateway used to validate the media scenarios end-to-end. 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, ReasoningEffort, StreamDelta, StreamSink, }; use reqwest::Client; use serde::{Deserialize, Serialize}; @@ -15,7 +17,7 @@ use std::time::Duration; use crate::config::ModelDef; /// A provider for any OpenAI-compatible API. -pub struct OpenAiCompatProvider { +pub(crate) struct OpenAiCompatProvider { /// Display name. name: String, /// Base URL. @@ -26,42 +28,79 @@ 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, } 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, 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 + /// mofa-owned default artifact directory 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(crate) fn with_output_dir( + name: impl Into, + base_url: impl Into, + api_key: impl Into, + models: Vec, + cost_tier: CostTier, + output_dir: Option, + ) -> Result { let client = Client::builder() .connect_timeout(Duration::from_secs(10)) .timeout(Duration::from_secs(120)) .build() - .unwrap_or_default(); + .map_err(|e| { + EngineError::Config(format!( + "failed to build OpenAI-compatible HTTP client: {e}" + )) + })?; - Self { + Ok(Self { name: name.into(), base_url: base_url.into(), api_key: api_key.into(), models, cost_tier, + output_dir: crate::artifacts::ensure_artifact_dir(output_dir), client, - } + }) } } -#[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)] @@ -84,6 +123,122 @@ struct ChatChoice { #[derive(Debug, Deserialize)] struct ChatUsage { total_tokens: Option, + #[serde(default)] + prompt_tokens: Option, + #[serde(default)] + 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, +} + +impl OpenAiCompatProvider { + /// Parse one SSE event block (the bytes preceding a blank-line delimiter). + /// + /// 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.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() { + return None; + } + if data == "[DONE]" { + return Some(SseEvent::Done); + } + serde_json::from_str::(&data) + .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 + } + + /// Collect the input strings to embed from a request: an explicit + /// `params.input` (a string or an array of strings) takes precedence, + /// otherwise every non-empty message content is embedded in order. Empty + /// input is rejected by the caller. + /// + /// Shared by the Ollama and liter-llm backends so embedding requests resolve + /// their inputs identically across providers. + pub(crate) fn embedding_inputs(request: &InferenceRequest) -> Vec { + match request.params.get("input") { + Some(serde_json::Value::String(s)) if !s.is_empty() => return vec![s.clone()], + Some(serde_json::Value::Array(arr)) => { + let inputs: Vec = arr + .iter() + .filter_map(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(String::from) + .collect(); + if !inputs.is_empty() { + return inputs; + } + } + _ => {} + } + request + .messages + .iter() + .map(|m| m.content.clone()) + .filter(|c| !c.trim().is_empty()) + .collect() + } } #[derive(Debug, Serialize)] @@ -94,6 +249,62 @@ struct TtsRequest { response_format: String, } +/// OpenAI-style `/images/generations` request. +#[derive(Debug, Serialize)] +struct ImageGenRequest { + model: String, + prompt: String, + n: u32, + #[serde(skip_serializing_if = "Option::is_none")] + size: Option, + /// Optional and omitted by default: some gateways (e.g. Agnes) reject + /// `response_format` with a 400, and when it is absent they return an image + /// `url` — which we download into a managed artifact anyway. Set it only when + /// the caller explicitly asks (`params.response_format`, e.g. `"b64_json"`) + /// for a server known to honor it. + #[serde(skip_serializing_if = "Option::is_none")] + response_format: Option, +} + +#[derive(Debug, Deserialize)] +struct ImageGenResponse { + data: Vec, +} + +#[derive(Debug, Deserialize)] +struct ImageGenData { + #[serde(default)] + b64_json: Option, + #[serde(default)] + url: Option, +} + +/// OpenAI-style `/embeddings` request: one or more input strings embedded in a +/// single call (the API accepts a string or an array; we always send an array). +#[derive(Debug, Serialize)] +struct EmbeddingRequest { + model: String, + input: Vec, +} + +#[derive(Debug, Deserialize)] +struct EmbeddingResponse { + #[serde(default)] + data: Vec, + usage: Option, +} + +#[derive(Debug, Deserialize)] +struct EmbeddingData { + #[serde(default)] + embedding: Vec, + /// Position in the batch; the API may return rows out of order, so we sort by + /// it. `Option` (not a `#[serde(default)]` `0`) so we can tell a real index + /// from an absent one — collapsing every absent index to `0` would turn the + /// sort into a silent no-op and could mismatch rows to inputs. + index: Option, +} + #[async_trait] impl Provider for OpenAiCompatProvider { fn name(&self) -> &str { @@ -116,12 +327,16 @@ 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); 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) }) @@ -164,7 +379,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, @@ -173,7 +388,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, @@ -185,7 +400,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(); @@ -203,15 +418,170 @@ impl Provider for OpenAiCompatProvider { Capability::Chat => self.invoke_chat(model_name, request, start).await, Capability::Tts => self.invoke_tts(model_name, request, start).await, Capability::Asr => self.invoke_asr(model_name, request, start).await, + Capability::ImageGen => self.invoke_image_gen(model_name, request, start).await, + Capability::Embedding => self.invoke_embedding(model_name, request, start).await, other => Err(EngineError::UnsupportedOperation(format!( "provider '{}' does not support {other}", self.name ))), } } + + /// 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 = ModelId::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) + }; + + '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(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() + { + // Receiver dropped (client disconnected): stop draining upstream. + break 'read; + } + } + } + // Any trailing event without a final blank line. + if let Some(event) = Self::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 { + /// Fallback health probe for gateways whose `GET /models` returns 404: send a + /// minimal (`max_tokens: 1`) chat completion and treat a success as healthy. + /// + /// Note: this issues a *real, billable* request. It runs only when the cheap + /// `/models` probe is unavailable (404), and only at startup and on explicit + /// `discovery/refresh`, so the cost is a single 1-token completion per probe + /// rather than per inference — an intentional trade to detect reachability for + /// providers that do not expose a free model-list endpoint. async fn health_via_chat(&self) -> bool { let Some(first_chat_model) = self .models @@ -229,6 +599,7 @@ impl OpenAiCompatProvider { content: "hi".into(), }], max_tokens: Some(1), + ..Default::default() }; let url = format!("{}/chat/completions", self.base_url); @@ -272,6 +643,7 @@ impl OpenAiCompatProvider { model: model_name.to_string(), messages, max_tokens, + ..Default::default() }; let url = format!("{}/chat/completions", self.base_url); @@ -309,7 +681,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 { @@ -320,8 +695,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() }) } @@ -385,8 +763,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(); @@ -401,6 +782,144 @@ impl OpenAiCompatProvider { tokens_used: None, fallback_used: false, routing_reason: None, + ..Default::default() + }) + } + + /// Generate an image via the OpenAI-compatible `/images/generations` endpoint + /// and store it as a managed artifact. + /// + /// We request `b64_json` so we can always write a local file; when a gateway + /// ignores that and returns a `url` instead, we download the bytes ourselves + /// so the caller uniformly gets a `file` path (never a bare remote URL it + /// would have to fetch separately). + async fn invoke_image_gen( + &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(|p| !p.trim().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 response_format = request + .params + .get("response_format") + .and_then(|v| v.as_str()) + .map(String::from); + + let body = ImageGenRequest { + model: model_name.to_string(), + prompt, + n: 1, + size, + response_format, + }; + + let url = format!("{}/images/generations", self.base_url); + let 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!("image_gen HTTP {status}: {text}"), + }); + } + + let parsed: ImageGenResponse = + resp.json().await.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("image_gen decode error: {e}"), + })?; + let image = parsed + .data + .into_iter() + .next() + .ok_or_else(|| EngineError::ProviderError { + provider: self.name.clone(), + detail: "image_gen returned no data".into(), + })?; + + // Resolve to bytes: decode inline base64, or download the returned URL. + let bytes = + if let Some(b64) = image.b64_json { + use base64::Engine as _; + base64::engine::general_purpose::STANDARD + .decode(b64.as_bytes()) + .map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("image base64 decode error: {e}"), + })? + } else if let Some(remote) = image.url { + let dl = self.client.get(&remote).send().await.map_err(|e| { + EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("image download error: {e}"), + } + })?; + dl.bytes() + .await + .map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("image download read error: {e}"), + })? + .to_vec() + } else { + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "image_gen returned neither b64_json nor url".into(), + }); + }; + + 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}")))?; + + Ok(InferenceResponse { + text: None, + file: Some(path.to_string_lossy().to_string()), + 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() }) } @@ -415,7 +934,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}")) })?; @@ -476,6 +995,91 @@ impl OpenAiCompatProvider { tokens_used: None, fallback_used: false, routing_reason: None, + ..Default::default() + }) + } + + async fn invoke_embedding( + &self, + model_name: &str, + request: &InferenceRequest, + start: std::time::Instant, + ) -> Result { + let inputs = Self::embedding_inputs(request); + if inputs.is_empty() { + return Err(EngineError::InvalidRequest( + "embedding requires text input (params.input or messages)".into(), + )); + } + + let body = EmbeddingRequest { + model: model_name.to_string(), + input: inputs, + }; + + let url = format!("{}/embeddings", self.base_url); + let 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!("embedding HTTP {status}: {text}"), + }); + } + + let mut parsed: EmbeddingResponse = + resp.json().await.map_err(|e| EngineError::ProviderError { + provider: self.name.clone(), + detail: format!("embedding parse error: {e}"), + })?; + + // The API may return rows out of batch order; restore input order — but + // only when every row actually carries an index. If any row omits it we + // can't reorder reliably, so we trust the provider's returned order + // (OpenAI-compatible servers return rows in input order) rather than + // reordering on missing data. + if parsed.data.iter().all(|d| d.index.is_some()) { + parsed.data.sort_by_key(|d| d.index); + } + let vectors: Vec> = parsed.data.into_iter().map(|d| d.embedding).collect(); + if vectors.is_empty() { + return Err(EngineError::ProviderError { + provider: self.name.clone(), + detail: "embedding response contained no vectors".into(), + }); + } + + let (prompt_tokens, tokens_used) = match parsed.usage { + // Embeddings bill input tokens only, so total == prompt. + Some(u) => (u.prompt_tokens, u.total_tokens.or(u.prompt_tokens)), + None => (None, None), + }; + Ok(InferenceResponse { + text: None, + file: None, + embedding: Some(vectors), + 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: Some(0), + fallback_used: false, + routing_reason: None, + ..Default::default() }) } } @@ -496,16 +1100,19 @@ 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, - ); + ) + .unwrap(); let cards = provider.discover().await.unwrap(); assert_eq!(cards.len(), 2); @@ -518,8 +1125,137 @@ 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"); } + + #[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 OpenAiCompatProvider::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!( + OpenAiCompatProvider::parse_sse_event(b"data: [DONE]"), + Some(SseEvent::Done) + )); + 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()); + } + + #[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 embedding_inputs_prefers_params_then_messages() { + use mofa_kernel::Message; + + // An explicit `params.input` array wins over messages. + let mut req = InferenceRequest { + capability: Some(Capability::Embedding), + messages: vec![Message { + role: "user".into(), + content: "from message".into(), + ..Default::default() + }], + ..Default::default() + }; + req.params = serde_json::json!({ "input": ["a", "b", ""] }); + assert_eq!(OpenAiCompatProvider::embedding_inputs(&req), vec!["a", "b"]); // empties dropped + + // A single-string `params.input` embeds one row. + req.params = serde_json::json!({ "input": "solo" }); + assert_eq!(OpenAiCompatProvider::embedding_inputs(&req), vec!["solo"]); + + // With no `params.input`, non-empty message contents are embedded in order. + req.params = serde_json::Value::Null; + assert_eq!( + OpenAiCompatProvider::embedding_inputs(&req), + vec!["from message"] + ); + + // Nothing to embed → empty (the caller turns this into InvalidRequest). + let empty = InferenceRequest { + capability: Some(Capability::Embedding), + ..Default::default() + }; + assert!(OpenAiCompatProvider::embedding_inputs(&empty).is_empty()); + } + + #[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/backends/system_tts.rs b/mofa-engine-core/src/backends/system_tts.rs new file mode 100644 index 0000000..b5902e4 --- /dev/null +++ b/mofa-engine-core/src/backends/system_tts.rs @@ -0,0 +1,471 @@ +//! Built-in, zero-config local TTS using the operating system's native voice. +//! +//! Every other TTS path ([`super::local_tts`]) shells out to a command the user +//! must install (Crane/Kokoro). This provider needs nothing installed: it uses the +//! voice already present on the machine — `say` on macOS, `espeak-ng`/`espeak` on +//! Linux — so a fresh checkout has a working local voice out of the box. The +//! engine auto-registers it (see `Engine::try_new`) as a low-priority fallback +//! whenever a real config declares no TTS backend, so configured voices still win. +//! +//! ## Output format +//! +//! The native tool produces an intermediate (AIFF on macOS, WAV on Linux) which is +//! transcoded to the requested `params.format`. `wav` is the zero-dependency +//! default (macOS `afconvert` and Linux `espeak` both produce it natively); `mp3` +//! is honored when `lame` or `ffmpeg` is available, otherwise the response falls +//! back to `wav` and its `file` carries the real extension. `voice` and `speed` +//! (`params.voice`/`params.speed`) pass through to the native tool. + +use std::path::{Path, PathBuf}; + +use async_trait::async_trait; +use mofa_kernel::{ + BackendFeature, BackendHealth, Capability, CostTier, EngineError, InferenceRequest, + InferenceResponse, LifecycleResult, ModelAvailability, ModelCard, ModelId, ModelResidency, + Provider, ProviderKind, +}; +use tokio::process::Command; + +/// The canonical name and model id for the built-in system voice. +pub(crate) const SYSTEM_TTS_NAME: &str = "system-tts"; +const SYSTEM_TTS_MODEL: &str = "system"; + +/// Which OS-native TTS tool this build will drive, resolved once at construction. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Voice { + /// macOS `say` (+ `afconvert` for WAV/M4A). + MacSay, + /// Linux `espeak-ng` or `espeak` (writes WAV directly). + Espeak, + /// No system voice available on this machine. + None, +} + +/// A zero-config TTS provider backed by the operating system's voice. +pub(crate) struct SystemTtsProvider { + output_dir: PathBuf, + voice: Voice, +} + +impl SystemTtsProvider { + /// Build the provider, detecting the OS-native voice tool once. + pub(crate) fn new(output_dir: Option) -> Self { + Self { + output_dir: crate::artifacts::ensure_artifact_dir(output_dir), + voice: detect_voice(), + } + } + + /// Whether this machine has a usable system voice. + pub(crate) fn is_available(&self) -> bool { + self.voice != Voice::None + } + + async fn synthesize( + &self, + request: &InferenceRequest, + start: std::time::Instant, + ) -> Result { + // Text comes from the first message, or `params.input` (parity with local_tts). + 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 voice_name = request + .params + .get("voice") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()); + // Optional speaking rate in words-per-minute (`params.speed` as a number or + // numeric string). Left to the tool's default when absent or unparseable. + let rate_wpm = request.params.get("speed").and_then(|v| { + v.as_f64() + .or_else(|| v.as_str().and_then(|s| s.parse::().ok())) + }); + let want_format = request + .params + .get("format") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .unwrap_or("wav") + .to_ascii_lowercase(); + + // A per-request scratch space; dropped (and cleaned) at the end. + let stem = uuid::Uuid::new_v4().to_string(); + let text_file = self.output_dir.join(format!("mofa_tts_{stem}.txt")); + tokio::fs::write(&text_file, &text) + .await + .map_err(|e| EngineError::Internal(format!("cannot write TTS input file: {e}")))?; + + // Synthesize to the tool's native container, then transcode to the request's + // format (falling back to the native format when no encoder is available). + let result = self + .render(&text_file, voice_name, rate_wpm, &stem, &want_format) + .await; + let _ = tokio::fs::remove_file(&text_file).await; + let output = result?; + + Ok(InferenceResponse { + text: None, + file: Some(output.to_string_lossy().to_string()), + model_used: SYSTEM_TTS_MODEL.to_string(), + provider: SYSTEM_TTS_NAME.to_string(), + 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() + }) + } + + /// Run the native voice and transcode to `want_format`, returning the final + /// audio path. The returned file's extension reflects what was actually + /// produced (which may differ from `want_format` when no encoder exists). + async fn render( + &self, + text_file: &Path, + voice: Option<&str>, + rate_wpm: Option, + stem: &str, + want_format: &str, + ) -> Result { + // 1. Native synthesis → an intermediate file. + let native = match self.voice { + Voice::MacSay => { + let aiff = self.output_dir.join(format!("mofa_tts_{stem}.aiff")); + let mut cmd = Command::new("say"); + if let Some(v) = voice { + cmd.arg("-v").arg(v); + } + if let Some(wpm) = rate_wpm.filter(|w| *w > 0.0) { + cmd.arg("-r").arg(format!("{}", wpm.round() as i64)); + } + cmd.arg("-f").arg(text_file).arg("-o").arg(&aiff); + run(cmd).await?; + aiff + } + Voice::Espeak => { + let wav = self.output_dir.join(format!("mofa_tts_{stem}.wav")); + let program = if which("espeak-ng") { + "espeak-ng" + } else { + "espeak" + }; + let mut cmd = Command::new(program); + if let Some(v) = voice { + cmd.arg("-v").arg(v); + } + if let Some(wpm) = rate_wpm.filter(|w| *w > 0.0) { + cmd.arg("-s").arg(format!("{}", wpm.round() as i64)); + } + cmd.arg("-w").arg(&wav).arg("-f").arg(text_file); + run(cmd).await?; + wav + } + Voice::None => { + return Err(EngineError::ProviderError { + provider: SYSTEM_TTS_NAME.into(), + detail: "no system voice available on this platform".into(), + }); + } + }; + + // A successful command must have left a non-empty file. + match tokio::fs::metadata(&native).await { + Ok(m) if m.len() > 0 => {} + _ => { + let _ = tokio::fs::remove_file(&native).await; + return Err(EngineError::ProviderError { + provider: SYSTEM_TTS_NAME.into(), + detail: "system voice produced no audio".into(), + }); + } + } + + // 2. Transcode to the requested format when it differs from the native one. + let native_ext = native + .extension() + .and_then(|e| e.to_str()) + .unwrap_or("") + .to_ascii_lowercase(); + if want_format == native_ext { + return Ok(native); + } + + let target = self + .output_dir + .join(format!("mofa_tts_{stem}.{want_format}")); + let transcoded = self.transcode(&native, &target, want_format).await; + match transcoded { + Ok(()) => { + let _ = tokio::fs::remove_file(&native).await; + Ok(target) + } + // No encoder for the requested format: keep the native file so the + // caller still gets real, playable audio (with its true extension). + Err(_) => { + let _ = tokio::fs::remove_file(&target).await; + Ok(native) + } + } + } + + /// Transcode `src` → `dst` in `format`, preferring built-in tools. Returns an + /// error (handled by the caller as "keep the native file") when no suitable + /// encoder is available. + async fn transcode(&self, src: &Path, dst: &Path, format: &str) -> Result<(), EngineError> { + // macOS `afconvert` covers wav/aiff/m4a with zero extra installs. + if which("afconvert") { + let spec = match format { + "wav" => Some(("WAVE", "LEI16")), + "aiff" | "aif" => Some(("AIFF", "BEI16")), + "m4a" | "aac" => Some(("m4af", "aac")), + _ => None, + }; + if let Some((file_fmt, data_fmt)) = spec { + let mut cmd = Command::new("afconvert"); + cmd.arg("-f") + .arg(file_fmt) + .arg("-d") + .arg(data_fmt) + .arg(src) + .arg(dst); + return run(cmd).await; + } + } + + // mp3 (and anything afconvert can't do) via lame or ffmpeg when present. + if format == "mp3" && which("lame") { + let mut cmd = Command::new("lame"); + cmd.arg("--quiet").arg(src).arg(dst); + return run(cmd).await; + } + if which("ffmpeg") { + let mut cmd = Command::new("ffmpeg"); + cmd.arg("-y") + .arg("-loglevel") + .arg("error") + .arg("-i") + .arg(src) + .arg(dst); + return run(cmd).await; + } + + Err(EngineError::ProviderError { + provider: SYSTEM_TTS_NAME.into(), + detail: format!("no encoder available for '{format}'"), + }) + } +} + +#[async_trait] +impl Provider for SystemTtsProvider { + fn name(&self) -> &str { + SYSTEM_TTS_NAME + } + + fn kind(&self) -> ProviderKind { + // Reported as local TTS so routing/observability treat it like any other + // on-device voice (local, free). + ProviderKind::LocalTts + } + + fn features(&self) -> Vec { + vec![ + BackendFeature::Discovery, + BackendFeature::Load, + BackendFeature::Unload, + ] + } + + async fn discover(&self) -> Result, EngineError> { + if !self.is_available() { + return Ok(vec![]); + } + let mut card = ModelCard::new( + SYSTEM_TTS_NAME.to_string(), + SYSTEM_TTS_MODEL.to_string(), + Capability::Tts, + CostTier::Free, + ); + card.id = ModelId::canonical(SYSTEM_TTS_NAME, SYSTEM_TTS_MODEL); + card.availability = ModelAvailability::Configured; + card.residency = ModelResidency::Unloaded; + card.refresh_status(); + Ok(vec![card]) + } + + async fn health(&self) -> Result { + if self.is_available() { + Ok(BackendHealth::Healthy) + } else { + Ok(BackendHealth::Unavailable) + } + } + + async fn load(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(SYSTEM_TTS_NAME, ModelId::name(model_id)), + residency: ModelResidency::Loaded, + memory_bytes: Some(0), + changed: true, + }) + } + + async fn unload(&self, model_id: &str) -> Result { + Ok(LifecycleResult { + model_id: ModelId::canonical(SYSTEM_TTS_NAME, ModelId::name(model_id)), + residency: ModelResidency::Unloaded, + memory_bytes: Some(0), + changed: true, + }) + } + + async fn invoke( + &self, + _model_id: &str, + request: &InferenceRequest, + ) -> Result { + let capability = request.capability.unwrap_or(Capability::Tts); + if capability != Capability::Tts { + return Err(EngineError::UnsupportedOperation(format!( + "provider '{SYSTEM_TTS_NAME}' only supports tts, not {capability}" + ))); + } + self.synthesize(request, std::time::Instant::now()).await + } +} + +// ============================================================================== +// Platform detection helpers +// ============================================================================== + +/// Detect the OS-native voice tool available on this machine. +fn detect_voice() -> Voice { + if cfg!(target_os = "macos") && which("say") { + Voice::MacSay + } else if which("espeak-ng") || which("espeak") { + Voice::Espeak + } else { + Voice::None + } +} + +/// Whether a bare command name resolves on `PATH`. +fn which(program: &str) -> bool { + std::env::var_os("PATH") + .map(|paths| { + std::env::split_paths(&paths).any(|dir| { + let p = dir.join(program); + p.is_file() + }) + }) + .unwrap_or(false) +} + +/// Run a prepared command, mapping a non-zero exit or spawn failure to a +/// retryable provider error (so the engine can fail over to a cloud voice). +async fn run(mut cmd: Command) -> Result<(), EngineError> { + let output = cmd + .kill_on_drop(true) + .output() + .await + .map_err(|e| EngineError::ProviderError { + provider: SYSTEM_TTS_NAME.into(), + detail: format!("failed to spawn command: {e}"), + })?; + if output.status.success() { + Ok(()) + } else { + Err(EngineError::ProviderError { + provider: SYSTEM_TTS_NAME.into(), + detail: format!( + "command exited with {}: {}", + output.status, + String::from_utf8_lossy(&output.stderr).trim() + ), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use mofa_kernel::Message; + + fn tts_request(text: &str, params: serde_json::Value) -> InferenceRequest { + InferenceRequest { + capability: Some(Capability::Tts), + messages: vec![Message { + role: "user".into(), + content: text.into(), + ..Default::default() + }], + params, + request_id: "test".into(), + ..Default::default() + } + } + + #[test] + fn detects_a_voice_on_this_platform() { + // CI/dev macOS always has `say`; the point is detection doesn't panic and + // availability agrees with `discover`/`health`. + let p = SystemTtsProvider::new(None); + assert_eq!(p.kind(), ProviderKind::LocalTts); + assert_eq!(p.name(), SYSTEM_TTS_NAME); + } + + #[tokio::test] + async fn discover_and_health_agree_on_availability() { + let p = SystemTtsProvider::new(None); + let cards = p.discover().await.unwrap(); + let health = p.health().await.unwrap(); + if p.is_available() { + assert_eq!(cards.len(), 1); + assert_eq!(cards[0].capability, Capability::Tts); + assert_eq!(cards[0].id, "system-tts/system"); + assert_eq!(health, BackendHealth::Healthy); + } else { + assert!(cards.is_empty()); + assert_eq!(health, BackendHealth::Unavailable); + } + } + + #[tokio::test] + async fn empty_text_is_an_invalid_request() { + let p = SystemTtsProvider::new(None); + let err = p + .invoke( + "system-tts/system", + &tts_request(" ", serde_json::Value::Null), + ) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + } + + #[tokio::test] + async fn rejects_non_tts_capability() { + let p = SystemTtsProvider::new(None); + let mut req = tts_request("hi", serde_json::Value::Null); + req.capability = Some(Capability::Chat); + let err = p.invoke("system-tts/system", &req).await.unwrap_err(); + assert!(matches!(err, EngineError::UnsupportedOperation(_))); + } + + // A real synthesis test would depend on `say`/`espeak` being installed and is + // covered end-to-end by the `mofa-podcast`/`mofa-explainer` app runs; omitted + // here to keep the unit tests hermetic. +} diff --git a/mofa-engine-core/src/circuit_breaker.rs b/mofa-engine-core/src/circuit_breaker.rs index d62f650..5134ca5 100644 --- a/mofa-engine-core/src/circuit_breaker.rs +++ b/mofa-engine-core/src/circuit_breaker.rs @@ -9,7 +9,7 @@ use std::time::{Duration, Instant}; /// Circuit breaker state. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum CircuitState { +pub(crate) enum CircuitState { /// Normal operation — requests flow through. Closed, /// Provider is deemed unhealthy — all requests fail fast. @@ -30,7 +30,7 @@ impl std::fmt::Display for CircuitState { /// Configuration for the circuit breaker. #[derive(Debug, Clone)] -pub struct CircuitBreakerConfig { +pub(crate) struct CircuitBreakerConfig { /// Number of consecutive failures before opening the circuit. pub failure_threshold: u32, /// Seconds to wait in Open before transitioning to HalfOpen. @@ -50,6 +50,12 @@ struct ProviderBreaker { state: CircuitState, failure_count: u32, last_failure: Option, + /// Whether a half-open probe has been admitted and is awaiting its outcome. + /// Guarantees that exactly one probe is in flight while half-open. + probe_in_flight: bool, + /// When the in-flight half-open probe was admitted, used to self-heal if the + /// caller never records the probe's outcome. + probe_started: Option, config: CircuitBreakerConfig, } @@ -59,27 +65,60 @@ impl ProviderBreaker { state: CircuitState::Closed, failure_count: 0, last_failure: None, + probe_in_flight: false, + probe_started: None, config, } } + /// How long to wait for a half-open probe to report an outcome before + /// assuming it was dropped and admitting a fresh one. Floored at 1s so a + /// zero cool-down (used in tests) still admits exactly one probe at a time. + fn probe_timeout(&self) -> Duration { + Duration::from_secs(self.config.cool_down_secs.max(1)) + } + /// Check if a request is allowed through. + /// + /// In `HalfOpen`, admits exactly one probe at a time: the first caller after + /// the cool-down wins, and concurrent callers fail fast until that probe's + /// outcome is recorded (which moves the circuit to `Closed` or `Open`). If an + /// admitted probe never records an outcome (e.g. its task was dropped), the + /// slot self-heals after `probe_timeout` so the provider is not black-holed + /// forever. fn allow_request(&mut self) -> bool { match self.state { CircuitState::Closed => true, CircuitState::Open => { - // Check if cool-down period has elapsed + // Check if cool-down period has elapsed. if let Some(last) = self.last_failure && last.elapsed() >= Duration::from_secs(self.config.cool_down_secs) { self.state = CircuitState::HalfOpen; - tracing::info!("circuit breaker → half_open (cool-down elapsed)"); + self.probe_in_flight = true; + self.probe_started = Some(Instant::now()); + tracing::info!("circuit breaker → half_open (admitting one probe)"); return true; } false } CircuitState::HalfOpen => { - // Allow exactly one probe request + if self.probe_in_flight { + // A probe is being tested; reject everyone else — unless the + // probe has gone silent past its timeout, in which case admit + // a fresh one so a dropped probe cannot wedge the circuit. + let stale = self + .probe_started + .map(|t| t.elapsed() >= self.probe_timeout()) + .unwrap_or(true); + if stale { + self.probe_started = Some(Instant::now()); + return true; + } + return false; + } + self.probe_in_flight = true; + self.probe_started = Some(Instant::now()); true } } @@ -87,8 +126,21 @@ impl ProviderBreaker { /// Record a successful request. fn record_success(&mut self) { - self.failure_count = 0; - self.state = CircuitState::Closed; + match self.state { + // The authorized probe succeeded → recover. + CircuitState::HalfOpen => { + self.state = CircuitState::Closed; + self.failure_count = 0; + self.probe_in_flight = false; + self.probe_started = None; + } + CircuitState::Closed => { + self.failure_count = 0; + } + // A late success from a request dispatched before the circuit opened + // must not reopen the gates; the cool-down governs recovery. + CircuitState::Open => {} + } } /// Record a failed request. @@ -104,26 +156,28 @@ impl ProviderBreaker { } } CircuitState::HalfOpen => { - // Probe failed — back to Open + // Probe failed — back to Open and release the probe slot. self.state = CircuitState::Open; + self.probe_in_flight = false; + self.probe_started = None; tracing::warn!("circuit breaker → open (probe failed)"); } CircuitState::Open => { - // Already open + // Already open. } } } } /// Manages circuit breakers for all providers. -pub struct CircuitBreakerRegistry { +pub(crate) struct CircuitBreakerRegistry { breakers: Mutex>, config: CircuitBreakerConfig, } impl CircuitBreakerRegistry { /// Create a new registry with the given configuration. - pub fn new(config: CircuitBreakerConfig) -> Self { + pub(crate) fn new(config: CircuitBreakerConfig) -> Self { Self { breakers: Mutex::new(HashMap::new()), config, @@ -131,7 +185,7 @@ impl CircuitBreakerRegistry { } /// Check if a request to the given provider is allowed. - pub fn allow_request(&self, provider: &str) -> bool { + pub(crate) fn allow_request(&self, provider: &str) -> bool { let mut breakers = self.breakers.lock().unwrap_or_else(|e| e.into_inner()); let breaker = breakers .entry(provider.to_string()) @@ -140,7 +194,7 @@ impl CircuitBreakerRegistry { } /// Record a successful request for a provider. - pub fn record_success(&self, provider: &str) { + pub(crate) fn record_success(&self, provider: &str) { let mut breakers = self.breakers.lock().unwrap_or_else(|e| e.into_inner()); if let Some(b) = breakers.get_mut(provider) { b.record_success(); @@ -148,7 +202,7 @@ impl CircuitBreakerRegistry { } /// Record a failed request for a provider. - pub fn record_failure(&self, provider: &str) { + pub(crate) fn record_failure(&self, provider: &str) { let mut breakers = self.breakers.lock().unwrap_or_else(|e| e.into_inner()); let breaker = breakers .entry(provider.to_string()) @@ -157,7 +211,7 @@ impl CircuitBreakerRegistry { } /// Get the current state for a provider. - pub fn state(&self, provider: &str) -> CircuitState { + pub(crate) fn state(&self, provider: &str) -> CircuitState { let breakers = self.breakers.lock().unwrap_or_else(|e| e.into_inner()); breakers .get(provider) @@ -228,6 +282,90 @@ mod tests { assert_eq!(reg.state("p"), CircuitState::Closed); } + #[test] + fn half_open_admits_exactly_one_probe() { + let cfg = CircuitBreakerConfig { + failure_threshold: 1, + cool_down_secs: 0, + }; + let reg = CircuitBreakerRegistry::new(cfg); + + reg.record_failure("p"); + assert_eq!(reg.state("p"), CircuitState::Open); + std::thread::sleep(Duration::from_millis(10)); + + // First caller after cool-down wins the single probe slot. + assert!(reg.allow_request("p")); + assert_eq!(reg.state("p"), CircuitState::HalfOpen); + + // Concurrent callers while the probe is in flight are rejected. + assert!(!reg.allow_request("p")); + assert!(!reg.allow_request("p")); + + // Once the probe succeeds, the circuit closes and admits traffic again. + reg.record_success("p"); + assert_eq!(reg.state("p"), CircuitState::Closed); + assert!(reg.allow_request("p")); + } + + #[test] + fn half_open_probe_slot_is_released_on_reopen() { + let cfg = CircuitBreakerConfig { + failure_threshold: 1, + cool_down_secs: 0, + }; + let reg = CircuitBreakerRegistry::new(cfg); + + reg.record_failure("p"); + std::thread::sleep(Duration::from_millis(10)); + assert!(reg.allow_request("p")); // → HalfOpen, probe in flight + assert!(!reg.allow_request("p")); // rejected while in flight + reg.record_failure("p"); // probe fails → Open, slot released + assert_eq!(reg.state("p"), CircuitState::Open); + + // After another cool-down, a fresh single probe is admitted. + std::thread::sleep(Duration::from_millis(10)); + assert!(reg.allow_request("p")); + assert!(!reg.allow_request("p")); + } + + #[test] + fn late_success_does_not_reopen_an_open_circuit() { + // A success recorded for a request that was in flight before the circuit + // opened must not flip it back to Closed. + let reg = CircuitBreakerRegistry::new(test_config()); + reg.record_failure("p"); + reg.record_failure("p"); + reg.record_failure("p"); + assert_eq!(reg.state("p"), CircuitState::Open); + + reg.record_success("p"); + assert_eq!(reg.state("p"), CircuitState::Open); + assert!(!reg.allow_request("p")); + } + + #[test] + fn stuck_half_open_probe_self_heals() { + // A probe that never records an outcome must not black-hole the provider. + let cfg = CircuitBreakerConfig { + failure_threshold: 1, + cool_down_secs: 1, // probe_timeout floors at 1s + }; + let reg = CircuitBreakerRegistry::new(cfg); + reg.record_failure("p"); + std::thread::sleep(Duration::from_millis(1100)); + + // First probe admitted, then it goes silent (no success/failure recorded). + assert!(reg.allow_request("p")); + assert_eq!(reg.state("p"), CircuitState::HalfOpen); + assert!(!reg.allow_request("p")); // still within probe timeout + + // After the probe timeout, a fresh probe is admitted rather than the + // circuit staying wedged forever. + std::thread::sleep(Duration::from_millis(1100)); + assert!(reg.allow_request("p")); + } + #[test] fn failure_in_half_open_reopens_circuit() { let cfg = CircuitBreakerConfig { diff --git a/mofa-engine-core/src/config.rs b/mofa-engine-core/src/config.rs index a31c82c..f596717 100644 --- a/mofa-engine-core/src/config.rs +++ b/mofa-engine-core/src/config.rs @@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize}; use std::path::{Path, PathBuf}; /// Top-level engine configuration. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct EngineConfig { /// Network listen settings #[serde(default)] @@ -19,37 +19,92 @@ pub struct EngineConfig { /// Memory management settings #[serde(default)] pub memory: MemoryConfig, + /// Operation timeout settings + #[serde(default)] + pub timeouts: TimeoutConfig, + /// Preflight (predictive warming) settings + #[serde(default)] + pub preflight: PreflightConfig, + /// Generated-artifact retention settings + #[serde(default)] + pub artifacts: ArtifactConfig, + /// Security and file-access settings + #[serde(default)] + pub security: SecurityConfig, /// Provider definitions #[serde(default)] pub providers: Vec, } +/// Retention policy for engine-generated artifacts (e.g. TTS audio files). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ArtifactConfig { + /// Directory scanned for cleanup (default: the mofa-owned `mofa_artifacts` + /// subdirectory of 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 = "ArtifactConfig::default_artifact_retention")] + pub retention_secs: u64, +} + +impl Default for ArtifactConfig { + fn default() -> Self { + Self { + dir: None, + retention_secs: Self::default_artifact_retention(), + } + } +} + +impl ArtifactConfig { + /// Retain generated artifacts for one hour by default. + 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 { /// 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. @@ -58,7 +113,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, } @@ -66,13 +121,163 @@ 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. +/// +/// 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 = "TimeoutConfig::default_request_timeout")] + pub request_secs: u64, + /// Maximum time a request may wait for a concurrency permit before failing. + #[serde(default = "TimeoutConfig::default_queue_timeout")] + pub queue_secs: u64, + /// Maximum time to load/warm a single model. + #[serde(default = "TimeoutConfig::default_load_timeout")] + pub load_secs: u64, + /// Maximum time for a single inference call to a provider. + #[serde(default = "TimeoutConfig::default_inference_timeout")] + pub inference_secs: u64, + /// Maximum time for a single provider discovery probe. + #[serde(default = "TimeoutConfig::default_discovery_timeout")] + pub discovery_secs: u64, + /// Maximum time for a single provider health probe. + #[serde(default = "TimeoutConfig::default_health_timeout")] + pub health_secs: u64, +} + +impl Default for TimeoutConfig { + fn default() -> Self { + Self { + 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(), + } + } +} + +// These `Duration` accessors are consumed only by the engine (in-crate), so they +// are `pub(crate)`; the struct itself and its `_secs` fields stay `pub` for +// external configuration/serialization. +impl TimeoutConfig { + /// Overall request budget as a `Duration`. + pub(crate) fn request(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.request_secs) + } + /// Queue-wait budget as a `Duration`. + pub(crate) fn queue(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.queue_secs) + } + /// Load budget as a `Duration`. + pub(crate) fn load(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.load_secs) + } + /// Inference budget as a `Duration`. + pub(crate) fn inference(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.inference_secs) + } + /// Discovery probe budget as a `Duration`. + pub(crate) fn discovery(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.discovery_secs) + } + /// Health probe budget as a `Duration`. + pub(crate) fn health(&self) -> std::time::Duration { + std::time::Duration::from_secs(self.health_secs) + } +} + +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 = "PreflightConfig::default_true")] + pub enabled: bool, + /// Whether the engine learns capability transition history. + #[serde(default = "PreflightConfig::default_true")] + pub history_learning: bool, + /// Whether the engine speculatively warms models from hints/history/subscriptions. + #[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 = "PreflightConfig::default_confidence_threshold")] + pub confidence_threshold: f64, + /// Minimum number of observed transitions before history is trusted. + #[serde(default = "PreflightConfig::default_min_samples")] + pub min_samples: u64, + /// Default lifetime of a capability subscription, in seconds. + #[serde(default = "PreflightConfig::default_subscription_ttl")] + pub subscription_ttl_secs: u64, +} + +impl Default for PreflightConfig { + fn default() -> Self { + Self { + 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(), + } + } +} + +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. @@ -87,33 +292,95 @@ 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. + /// + /// 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, + /// Vendor dialect for a `cloud_video_gen` backend: `"ark"` (default, the + /// Volcengine Ark / Seedance task API) or `"agnes"` (the Agnes AI gateway's + /// `/videos` task API). Ignored by other provider kinds. + #[serde(default)] + pub dialect: String, + /// Directory for `local_tts` audio artifacts (default: the mofa-owned + /// `mofa_artifacts` subdirectory of 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)] + pub price_input_per_1k: f64, + /// USD price per 1000 completion (output) tokens. + #[serde(default)] + pub price_output_per_1k: f64, } -fn default_priority() -> u8 { - 5 -} - -fn default_cost_tier() -> String { - "medium".into() +impl Default for ProviderConfig { + fn default() -> Self { + Self { + name: String::new(), + kind: String::new(), + base_url: String::new(), + api_key: None, + priority: Self::default_priority(), + cost_tier: Self::default_cost_tier(), + models: Vec::new(), + enabled: Self::default_enabled(), + command: None, + args: Vec::new(), + output_format: None, + dialect: String::new(), + output_dir: None, + diarize_args: Vec::new(), + price_input_per_1k: 0.0, + price_output_per_1k: 0.0, + } + } } -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. -#[derive(Debug, Clone, Serialize, Deserialize)] +#[derive(Debug, Clone, Default, Serialize, Deserialize)] pub struct ModelDef { /// Model name / identifier pub name: String, @@ -123,6 +390,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 { @@ -183,7 +454,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)?); } } @@ -191,7 +462,23 @@ 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 {}", + self.preflight.confidence_threshold + ))); + } + // Each configured allowlist root must resolve. Otherwise a typo would be + // silently dropped, leaving an *empty* allowlist that permits every path + // — a security downgrade rather than the intended restriction. + for root in &self.security.input_roots { + if std::fs::canonicalize(root).is_err() { + return Err(EngineError::Config(format!( + "security.input_roots entry '{root}' does not exist or cannot be resolved" + ))); + } + } for provider in &self.providers { if !provider.enabled { continue; @@ -205,6 +492,28 @@ impl EngineConfig { provider.name ))); } + // Every local process-adapter kind is spawned from a configured + // `command` and serves an explicit model list, so validate all three + // uniformly: catch a missing command or empty model list here (at + // `validate-config` time) rather than later at engine construction or + // silently discovering nothing. + if matches!( + provider.kind.as_str(), + "local_tts" | "local_asr" | "local_image_gen" | "local_video_gen" + ) { + if provider.command.as_deref().unwrap_or_default().is_empty() { + return Err(EngineError::Config(format!( + "provider '{}' ({}) requires a non-empty command", + provider.name, provider.kind + ))); + } + if provider.models.is_empty() { + return Err(EngineError::Config(format!( + "provider '{}' ({}) must declare at least one model", + provider.name, provider.kind + ))); + } + } for model in &provider.models { if mofa_kernel::Capability::from_str_loose(&model.capability).is_none() { return Err(EngineError::Config(format!( @@ -231,6 +540,7 @@ impl EngineConfig { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }); // OpenAI @@ -248,33 +558,39 @@ 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, + ..Default::default() }); } @@ -293,15 +609,18 @@ 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, + ..Default::default() }); } @@ -320,21 +639,25 @@ 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, + ..Default::default() }); } @@ -352,8 +675,10 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }], enabled: true, + ..Default::default() }); } @@ -371,8 +696,10 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }], enabled: true, + ..Default::default() }); } @@ -390,37 +717,58 @@ impl EngineConfig { capability: "chat".into(), context_window: Some(128000), memory_mb: None, + ..Default::default() }], enabled: true, + ..Default::default() }); } EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers, } } } -/// 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 { + // Two accepted indirection forms resolve to an environment variable: + // - `${VAR}` (shell-style), and + // - `env:VAR` (the terser form used throughout our example configs). + // Anything else is treated as a literal secret. A real API key never + // matches these prefixes, so this cannot accidentally shadow one. + let var_name = s + .strip_prefix("${") + .and_then(|rest| rest.strip_suffix('}')) + .or_else(|| s.strip_prefix("env:")); + if let Some(var_name) = var_name { + 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 { + pub(crate) fn provider_kind(&self) -> Result { match self.kind.as_str() { "ollama" => Ok(ProviderKind::Ollama), "openai_compatible" => Ok(ProviderKind::OpenAiCompatible), + "local_tts" => Ok(ProviderKind::LocalTts), + "local_asr" => Ok(ProviderKind::LocalAsr), + "local_image_gen" => Ok(ProviderKind::LocalImageGen), + "local_video_gen" => Ok(ProviderKind::LocalVideoGen), + "cloud_video_gen" => Ok(ProviderKind::CloudVideoGen), + "liter_llm" => Ok(ProviderKind::LiterLlm), other => Err(EngineError::Config(format!( "unknown provider kind '{}' for provider '{}'", other, self.name @@ -435,7 +783,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] @@ -443,16 +794,30 @@ 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") }; } + #[test] + fn resolve_env_var_colon_syntax() { + // The terser `env:VAR` form (used across our example configs) resolves the + // same way as `${VAR}`. + // SAFETY: this test runs single-threaded and restores the env var. + unsafe { std::env::set_var("TEST_MOFA_KEY2", "resolved-colon") }; + assert_eq!( + ProviderConfig::resolve_env_var("env:TEST_MOFA_KEY2").unwrap(), + "resolved-colon" + ); + unsafe { std::env::remove_var("TEST_MOFA_KEY2") }; + } + #[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()); + assert!(ProviderConfig::resolve_env_var("env:MISSING_MOFA_KEY").is_err()); } #[test] @@ -472,6 +837,10 @@ mod tests { let cfg = EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers: vec![ProviderConfig { name: "test".into(), kind: "ollama".into(), @@ -481,6 +850,7 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }], }; let toml_str = toml::to_string_pretty(&cfg).unwrap(); @@ -489,11 +859,98 @@ 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_requires_command_and_models_for_all_local_kinds() { + // Every local process-adapter kind must be validated uniformly: a missing + // command, or an empty model list, is rejected at validate-config time. + for kind in [ + "local_tts", + "local_asr", + "local_image_gen", + "local_video_gen", + ] { + let base = |command: Option<&str>, models: Vec| EngineConfig { + providers: vec![ProviderConfig { + name: format!("{kind}-provider"), + kind: kind.into(), + base_url: String::new(), + command: command.map(String::from), + models, + ..Default::default() + }], + ..Default::default() + }; + let a_model = || { + vec![ModelDef { + name: "fixture".into(), + capability: match kind { + "local_tts" => "tts", + "local_asr" => "asr", + "local_video_gen" => "video_gen", + _ => "image_gen", + } + .into(), + ..Default::default() + }] + }; + + // Missing command → rejected. + assert!( + base(None, a_model()).validate().is_err(), + "{kind} without a command should fail validation" + ); + // Empty model list → rejected. + assert!( + base(Some("some-cli"), vec![]).validate().is_err(), + "{kind} with no models should fail validation" + ); + // Command + at least one model → valid. + base(Some("some-cli"), a_model()) + .validate() + .unwrap_or_else(|e| panic!("{kind} with command and model should validate: {e}")); + } + } + #[test] fn validate_rejects_unknown_provider_kind() { let cfg = EngineConfig { listen: ListenConfig::default(), memory: MemoryConfig::default(), + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: ArtifactConfig::default(), + security: SecurityConfig::default(), providers: vec![ProviderConfig { name: "bad".into(), kind: "bad_kind".into(), @@ -503,6 +960,7 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: true, + ..Default::default() }], }; assert!(cfg.validate().is_err()); diff --git a/mofa-engine-core/src/engine.rs b/mofa-engine-core/src/engine.rs index e656d6c..912dbfe 100644 --- a/mofa-engine-core/src/engine.rs +++ b/mofa-engine-core/src/engine.rs @@ -1,26 +1,47 @@ //! 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::{HashMap, 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, DataClass, EngineError, EngineEvent, + EngineStatus, FailedAttempt, FallbackPolicy, InferenceRequest, InferenceResponse, Message, + ModelCard, ModelId, ModelResidency, Prefer, Provider, ProviderHealth, ProviderKind, + ResponsesRequest, ResponsesResponse, StreamChunk, StreamDelta, +}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Mutex as AsyncMutex, Semaphore, broadcast, mpsc}; +use tokio::task::{AbortHandle, JoinHandle}; + +use crate::backends::{ + CloudVideoGenProvider, LiterLLMProvider, LocalAsrProvider, LocalImageGenProvider, + LocalTtsProvider, LocalVideoGenProvider, OllamaProvider, OpenAiCompatProvider, + SystemTtsProvider, }; -use tokio::sync::broadcast; - -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::metrics::{EngineMetrics, MetricsGauges}; +use crate::preflight::{GLOBAL_SCOPE, PreflightMetrics, PreflightPredictor, PreflightStats}; +use crate::router::{RouteDecision, Router, RoutingProvider}; +use crate::subscription::{SubscriptionInfo, SubscriptionRegistry}; + +/// Maximum number of lifecycle records retained in the rolling history. +const LIFECYCLE_CAPACITY: usize = 256; +/// Maximum number of in-flight predictions tracked for hit/miss accounting. +/// Bounds memory against an unbounded stream of unique scope identifiers. +const MAX_PENDING_PREDICTIONS: usize = 4096; +/// Maximum number of stored Responses-API conversations. Each continued turn +/// mints a new id (so a caller can branch from any prior response), so this +/// bounds total memory by evicting the oldest conversation when full. +const MAX_CONVERSATIONS: usize = 4096; #[derive(Clone)] struct RegisteredProvider { @@ -30,6 +51,41 @@ struct RegisteredProvider { provider: Arc, } +/// A single entry in the model lifecycle history. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct LifecycleRecord { + /// Monotonic sequence number. + pub seq: u64, + /// Milliseconds since engine start (monotonic). + pub at_ms: u64, + /// Affected model. + pub model_id: String, + /// Event kind: `load`, `unload`, `evict`, `idle_unload`, `load_failed`, etc. + pub event: String, + /// Optional human-readable detail. + pub detail: Option, +} + +/// A stored Responses-API conversation: the full running message history and +/// when it was created (for oldest-first eviction). +struct StoredConversation { + messages: Vec, + created: Instant, +} + +/// A snapshot of the engine's memory accounting. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MemoryReport { + /// Bytes reserved across all models. + pub used_bytes: u64, + /// Total budget. + pub budget_bytes: u64, + /// Bytes still free within the budget. + pub available_bytes: u64, + /// Per-model allocations. + pub allocations: Vec, +} + /// The main MoFA Engine orchestrator. pub struct Engine { /// Named providers. @@ -38,18 +94,109 @@ 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, + /// Stored Responses-API conversations, keyed by response id. Bounded by + /// [`MAX_CONVERSATIONS`] via oldest-first eviction. + conversations: DashMap, + /// Operation timeouts. + timeouts: TimeoutConfig, + /// Preflight configuration. + preflight_config: PreflightConfig, + /// Idle timeout before a resident model is auto-unloaded. + idle_timeout: Duration, + /// Per-model concurrency admission. + semaphores: DashMap>, + /// Serializes the evict-then-reserve admission critical section. + load_gate: AsyncMutex<()>, + /// Rolling lifecycle history. + lifecycle: Mutex>, + /// Lifecycle sequence counter. + lifecycle_seq: AtomicU64, + /// Background idle-eviction task; aborted on drop. + idle_task: Mutex>>, + /// Background artifact-cleanup task; aborted on drop. + artifact_task: Mutex>>, + /// Weak self-reference for background tasks. + weak_self: OnceLock>, /// Event broadcast channel. event_tx: broadcast::Sender, + /// Bounded-cardinality metrics counters. + metrics: EngineMetrics, + /// 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, /// Engine start time. started_at: Instant, } +/// 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; + + /// Remaining time until `deadline`, saturating at zero. + 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) + } + + /// 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 { + // 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) + } + + /// 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 { /// Create and initialize a new engine from configuration. /// @@ -69,6 +216,7 @@ impl Engine { let memory = MemoryManager::new(config.memory.budget_mb); let circuit_breakers = CircuitBreakerRegistry::new(CircuitBreakerConfig::default()); let preflight = PreflightPredictor::new(); + let idle_timeout = Duration::from_secs(config.memory.idle_timeout_secs); let mut providers = Vec::new(); for pc in &config.providers { @@ -80,14 +228,100 @@ impl Engine { let kind = pc.provider_kind()?; let cost_tier = CostTier::from_str_loose(&pc.cost_tier); let provider: Arc = match kind { - ProviderKind::Ollama => Arc::new(OllamaProvider::new(&pc.name, &pc.base_url)), - ProviderKind::OpenAiCompatible => Arc::new(OpenAiCompatProvider::new( + ProviderKind::Ollama => { + Arc::new(OllamaProvider::new(&pc.name, &pc.base_url, &pc.models)?) + } + 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::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!( + "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(), + )) + } + 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(), + ) + .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(), + )) + } + ProviderKind::LocalVideoGen => { + let command = pc.command.clone().ok_or_else(|| { + EngineError::Config(format!( + "provider '{}' (local_video_gen) requires a command", + pc.name + )) + })?; + Arc::new(LocalVideoGenProvider::new( + &pc.name, + command, + pc.args.clone(), + pc.output_format.clone(), + pc.output_dir.clone(), + pc.models.clone(), + )) + } + ProviderKind::CloudVideoGen => Arc::new(CloudVideoGenProvider::new( + &pc.name, + &pc.dialect, + &pc.base_url, + pc.api_key.clone(), + pc.models.clone(), + cost_tier, + config.artifacts.dir.clone(), + )?), _ => { return Err(EngineError::Config(format!( "provider '{}' uses unsupported provider kind", @@ -104,6 +338,63 @@ impl Engine { }); } + // Zero-config local voice: when a real config is present but declares no TTS + // backend, auto-register the OS-native voice (macOS `say` / Linux `espeak`) + // as a low-priority fallback. This gives local-first audio scenarios (S1/S4/ + // S6) a working voice with nothing to install, while any configured TTS — + // having a lower priority number — still wins the route. An empty config is + // left untouched (no providers in, no providers out). + let has_configured_provider = config.providers.iter().any(|pc| pc.enabled); + let declares_tts = config.providers.iter().any(|pc| { + pc.enabled + && pc.models.iter().any(|m| { + mofa_kernel::Capability::from_str_loose(&m.capability) + == Some(mofa_kernel::Capability::Tts) + }) + }); + if has_configured_provider && !declares_tts { + let system_tts = SystemTtsProvider::new(config.artifacts.dir.clone()); + if system_tts.is_available() { + tracing::info!( + "no TTS backend configured; registering built-in system voice as fallback" + ); + providers.push(RegisteredProvider { + name: crate::backends::system_tts::SYSTEM_TTS_NAME.to_string(), + kind: ProviderKind::LocalTts, + // A large priority number is the least-preferred tier, so a + // configured voice always outranks this fallback. + priority: 99, + provider: Arc::new(system_tts), + }); + } + } + + // 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), + ); + + // 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(), @@ -111,11 +402,32 @@ impl Engine { memory, circuit_breakers, preflight, + preflight_metrics: PreflightMetrics::default(), + subscriptions: SubscriptionRegistry::new(), + warming: DashMap::new(), + pending_predictions: DashMap::new(), + conversations: DashMap::new(), + timeouts: config.timeouts.clone(), + preflight_config: config.preflight.clone(), + idle_timeout, + semaphores: DashMap::new(), + load_gate: AsyncMutex::new(()), + lifecycle: Mutex::new(VecDeque::with_capacity(LIFECYCLE_CAPACITY)), + lifecycle_seq: AtomicU64::new(0), + idle_task: Mutex::new(None), + artifact_task: Mutex::new(None), + weak_self: OnceLock::new(), event_tx, + metrics: EngineMetrics::default(), + pricing, + input_roots, started_at: Instant::now(), }); + let _ = engine.weak_self.set(Arc::downgrade(&engine)); engine.refresh_resources().await; + engine.spawn_idle_eviction(); + engine.spawn_artifact_sweep(artifact_sweeper); Ok(engine) } @@ -127,6 +439,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 +447,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 +476,54 @@ 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. 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; + } + _ => {} + } + } 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 +544,7 @@ impl Engine { } async fn refresh_health(&self) { + let health_timeout = self.timeouts.health(); let handles = self .providers .iter() @@ -203,7 +552,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 +578,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 +592,394 @@ impl Engine { models } - /// Run inference: route → circuit-check → load → invoke → optional fallback. + /// 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 { + 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); + 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. + 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), + } + result + } + + // ----- Responses API (stateful multi-turn) -------------------------------- + + /// Run one turn of the stateful Responses API (S2 deep multi-turn reasoning). + /// + /// Seeds the turn with the conversation stored under `previous_response_id` + /// (or a fresh history with the optional `instructions` as a system message), + /// appends this turn's explicit `messages` and the `input` shorthand, routes + /// through the normal [`invoke`](Self::invoke) path (so routing, failover, + /// metrics, cost, and the S5 audit all apply), then appends the assistant + /// reply and stores the updated history under a **new** id. The caller passes + /// that id back as `previous_response_id` to continue the conversation. + /// + /// Continuing mints a fresh id per turn (so a caller can branch from any prior + /// response); the store is bounded by [`MAX_CONVERSATIONS`] via oldest-first + /// eviction. TODO: a single very long chain stores one full-history snapshot + /// per turn (quadratic in turns); a future revision can share a suffix or cap + /// retained turns without changing this contract. + pub async fn respond(&self, req: ResponsesRequest) -> Result { + let ResponsesRequest { + previous_response_id, + instructions, + input, + mut request, + } = req; + + // 1. Seed history: continue a stored conversation, or start a fresh one + // (with optional system instructions). An unknown prior id is a hard + // error rather than a silent fresh start, so an expired/evicted chain + // surfaces to the caller. + let mut messages: Vec = match previous_response_id.as_deref() { + Some(prev) => self + .conversations + .get(prev) + .map(|c| c.messages.clone()) + .ok_or_else(|| { + EngineError::InvalidRequest(format!( + "unknown previous_response_id '{prev}' (expired or never existed)" + )) + })?, + None => match instructions.filter(|s| !s.trim().is_empty()) { + Some(system) => vec![Message { + role: "system".into(), + content: system, + ..Default::default() + }], + None => Vec::new(), + }, + }; + + // 2. Append this turn's new messages, then the `input` shorthand as a + // user turn. + messages.append(&mut request.messages); + if let Some(text) = input.filter(|t| !t.trim().is_empty()) { + messages.push(Message { + role: "user".into(), + content: text, + ..Default::default() + }); + } + if messages.iter().all(|m| m.content.trim().is_empty()) { + return Err(EngineError::InvalidRequest( + "responses request has no input (set `input`, `messages`, or continue a conversation)".into(), + )); + } + + // 3. Route + invoke with the accumulated history. Default to chat when the + // caller does not pin a capability, since Responses is a chat surface. + request.messages = messages.clone(); + if request.capability.is_none() { + request.capability = Some(Capability::Chat); + } + let response = self.invoke(request).await?; + + // 4. Append the assistant reply and persist the updated history under a + // new response id. + if let Some(text) = response.text.clone().filter(|t| !t.is_empty()) { + messages.push(Message { + role: "assistant".into(), + content: text, + ..Default::default() + }); + } + let id = uuid::Uuid::new_v4().to_string(); + let message_count = messages.len(); + self.store_conversation(id.clone(), messages); + + Ok(ResponsesResponse { + id, + message_count, + response, + }) + } + + /// Persist a conversation under `id`, evicting the oldest stored conversation + /// first when at [`MAX_CONVERSATIONS`] capacity so memory stays bounded. + fn store_conversation(&self, id: String, messages: Vec) { + if self.conversations.len() >= MAX_CONVERSATIONS + && !self.conversations.contains_key(&id) + && let Some(victim) = self + .conversations + .iter() + .min_by_key(|e| e.created) + .map(|e| e.key().clone()) + { + self.conversations.remove(&victim); + } + self.conversations.insert( + id, + StoredConversation { + messages, + created: Instant::now(), + }, + ); + } + + /// The stored message history for a conversation id, if still retained. + pub fn conversation_messages(&self, id: &str) -> Option> { + self.conversations.get(id).map(|c| c.messages.clone()) + } + + /// Forget a stored conversation. Returns whether one existed. + pub fn delete_conversation(&self, id: &str) -> bool { + self.conversations.remove(id).is_some() + } + + /// Number of stored conversations. Test-only diagnostic, so it is neither + /// public API nor compiled into release builds. + #[cfg(test)] + fn conversation_count(&self) -> usize { + self.conversations.len() + } + + /// 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()?; + 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(Engine::estimate_prompt_tokens(req)), + Some(Engine::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 { + self.providers + .iter() + .find(|p| p.name == provider) + .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 + /// candidate is the primary selection, and each subsequent one is a failover. + /// Only *retryable* failures advance to the next candidate, so malformed or + /// unsupported requests fail immediately rather than masquerading as + /// transient errors. The candidate list itself encodes the fallback policy. + async fn invoke_inner(&self, req: InferenceRequest) -> Result { self.reject_ambiguous_short_name(&req)?; + self.check_input_path(&req)?; + let overall_deadline = Instant::now() + self.timeouts.request(); + + // Confirm the previous request's prediction for this scope before a new + // one is formed during this request's admission. + self.confirm_prediction(&req); let all_models = self.capabilities().await; let providers = self.routing_providers(); - let selected = Router::route(&all_models, &req, &providers).ok_or_else(|| { - let cap_str = req - .capability - .map(|c| c.to_string()) - .unwrap_or_else(|| "any".into()); - EngineError::NoCapableModel(cap_str) - })?; + let candidates = self.build_candidates(&all_models, &req, &providers); - let model_id = selected.model.id.clone(); - let provider_name = selected.model.provider.clone(); - let routing_reason = selected.reason.clone(); + if candidates.is_empty() { + return Err(EngineError::NoCapableModel(Self::requested_capability( + &req, + ))); + } - match self - .try_invoke(&model_id, &provider_name, &req, Some(routing_reason)) - .await - { - Ok(resp) => Ok(resp), - Err(primary_err) if self.can_fallback(&req) => { - tracing::warn!("primary model '{model_id}' failed: {primary_err}, trying fallback"); - let mut fallback_req = req.clone(); - fallback_req.model = None; - let fallback_models = self - .capabilities() - .await - .into_iter() - .filter(|m| m.id != model_id && m.provider != provider_name) - .collect::>(); - let Some(fallback) = Router::route(&fallback_models, &fallback_req, &providers) - else { - return Err(primary_err); - }; - let fb_id = fallback.model.id.clone(); - let fb_provider = fallback.model.provider.clone(); - let mut resp = self - .try_invoke(&fb_id, &fb_provider, &fallback_req, Some(fallback.reason)) - .await - .map_err(|_| primary_err)?; - resp.fallback_used = true; - Ok(resp) + // 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 Engine::remaining(overall_deadline).is_zero() { + let e = EngineError::Timeout("overall request deadline exceeded".into()); + chain.push(FailedAttempt { + provider: provider_name, + model: ModelId::name(&model_id).to_string(), + reason: e.to_string(), + }); + last_err = Some(e); + 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: ModelId::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 + .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}"); + chain.push(FailedAttempt { + provider: provider_name, + model: ModelId::name(&model_id).to_string(), + reason: e.to_string(), + }); + last_err = Some(e); + } + } + } + + 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 { + req.capability + .map(|c| c.to_string()) + .unwrap_or_else(|| "any".into()) + } + + /// Build the ordered candidate plan, applying the request's fallback policy. + fn build_candidates<'a>( + &self, + all_models: &'a [ModelCard], + req: &InferenceRequest, + providers: &[RoutingProvider], + ) -> Vec> { + let budget = Some(self.memory.budget_bytes()); + let primary = Router::route_ranked(all_models, req, providers, budget); + + match (req.model.is_some(), req.fallback_policy) { + // Named request, strict (default) or fallback disabled: only the named model. + (true, FallbackPolicy::CapabilityOnly | FallbackPolicy::Disabled) => primary, + // Named request with explicit opt-in: the named model first, then any + // capability-compatible candidate. + (true, FallbackPolicy::AllowNamed) => { + let mut cands = primary; + let mut cap_req = req.clone(); + cap_req.model = None; + for d in Router::route_ranked(all_models, &cap_req, providers, budget) { + if !cands.iter().any(|c| c.model.id == d.model.id) { + cands.push(d); + } + } + cands } - Err(primary_err) => Err(primary_err), + // Capability request, fallback disabled: only the single best. + (false, FallbackPolicy::Disabled) => primary.into_iter().take(1).collect(), + // Capability request, default/allow: the full ranked list. + (false, _) => primary, } } @@ -317,22 +1012,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,169 +1038,862 @@ 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 = 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(_)) => { + 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 = Engine::remaining(overall_deadline).min(self.timeouts.inference()); let start = Instant::now(); - let invoke_timeout = Duration::from_secs(180); - let result = tokio::time::timeout(invoke_timeout, provider.invoke(model_id, req)).await; + let result = tokio::time::timeout(inference_budget, provider.invoke(model_id, req)).await; let duration_ms = start.elapsed().as_millis() as u64; self.end_execution(model_id); + drop(permit); match result { Ok(Ok(mut resp)) => { - resp.routing_reason = routing_reason; + resp.routing_reason = Some(routing_reason); self.circuit_breakers.record_success(provider_name); - if let Some(cap) = req.capability { - self.preflight.record(cap); - } - let _ = self.event_tx.send(EngineEvent::RequestCompleted { - request_id: req.request_id.clone(), - duration_ms, - success: true, - }); + self.record_transition(req); + self.emit_request_completed(&req.request_id, duration_ms, true); Ok(resp) } Ok(Err(e)) => { self.circuit_breakers.record_failure(provider_name); - let _ = self.event_tx.send(EngineEvent::RequestCompleted { - request_id: req.request_id.clone(), - duration_ms, - success: false, - }); + self.emit_request_completed(&req.request_id, duration_ms, false); Err(e) } Err(_) => { self.circuit_breakers.record_failure(provider_name); - let _ = self.event_tx.send(EngineEvent::RequestCompleted { - request_id: req.request_id.clone(), - duration_ms: invoke_timeout.as_millis() as u64, - success: false, - }); + self.emit_request_completed( + &req.request_id, + inference_budget.as_millis() as u64, + false, + ); Err(EngineError::Timeout(format!( - "provider '{}' did not respond within {}s", - provider_name, - invoke_timeout.as_secs() + "provider '{provider_name}' did not respond within {}s", + inference_budget.as_secs() ))) } } } + /// Run inference and stream typed output chunks. + /// + /// Returns a receiver that yields a `Started` chunk, then `Text` deltas as + /// the backend produces them, then a terminal `Completed` or `Error`. Errors + /// are delivered in-band as [`StreamChunk::Error`] rather than as a `Result`, + /// so a consumer only has to read one channel. + /// + /// Streaming targets the single best candidate. Unlike [`invoke`](Self::invoke) + /// it does not fail over once output has begun, since partial output cannot + /// be un-sent; a failure before the first token surfaces as `Error`. + pub fn invoke_stream(&self, req: InferenceRequest) -> mpsc::Receiver { + let (tx, rx) = mpsc::channel(64); + if let Some(weak) = self.weak() { + tokio::spawn(async move { + if let Some(engine) = weak.upgrade() { + engine.run_stream(req, tx).await; + } + }); + } + rx + } + + /// Emit a terminal `Error` chunk for a streaming request that failed *before* + /// generation began, recording the failed-request metric so a pre-dispatch + /// failure (circuit open, no capable model, admission timeout, …) is counted + /// exactly like the equivalent unary `invoke` failure rather than being + /// invisible to `mofa_requests_total`. The terminal match after generation + /// starts records metrics itself (with the real duration), so this helper is + /// only used on the pre-dispatch error paths. + async fn fail_stream(&self, out: &mpsc::Sender, err: EngineError) { + self.metrics.record_request(false, 0, false); + let _ = out.send(StreamChunk::Error(err.info())).await; + } + + /// 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)) + { + self.fail_stream(&out, e).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)); + self.fail_stream(&out, err).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(); + + // 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}" + )); + self.fail_stream(&out, err).await; + return; + } + } + + if !self.circuit_breakers.allow_request(&provider_name) { + let err = EngineError::CircuitOpen(provider_name.clone()); + self.fail_stream(&out, err).await; + return; + } + let Some(provider) = self.find_provider(&provider_name) else { + let err = EngineError::Internal(format!("provider '{provider_name}' not found")); + self.fail_stream(&out, err).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); + self.fail_stream(&out, e).await; + return; + } + + let sem = self.semaphore_for(&model_id, max_concurrency); + 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(_)) => { + self.emit_request_completed(&req.request_id, 0, false); + let err = EngineError::Internal("concurrency semaphore closed".into()); + self.fail_stream(&out, err).await; + return; + } + Err(_) => { + self.emit_request_completed(&req.request_id, 0, false); + let err = EngineError::Timeout(format!("queue wait exceeded for '{model_id}'")); + self.fail_stream(&out, err).await; + return; + } + }; + + self.begin_execution(&model_id); + self.trigger_preflight(&req); + + let _ = out + .send(StreamChunk::Started { + request_id: req.request_id.clone(), + model_used: ModelId::name(&model_id).to_string(), + provider: provider_name.clone(), + }) + .await; + + // 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 { + 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; + } + } + }); + + 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)) + .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); + let cost_usd = + 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, + ); + // 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, + prompt_tokens: resp.prompt_tokens, + completion_tokens: resp.completion_tokens, + cost_usd, + file: resp.file, + fallback_used: false, + routing_reason: Some(routing_reason), + } + } + Ok(Err(e)) => { + self.circuit_breakers.record_failure(&provider_name); + self.emit_request_completed(&req.request_id, duration_ms, false); + self.metrics.record_request(false, duration_ms, false); + StreamChunk::Error(e.info()) + } + Err(_) => { + self.circuit_breakers.record_failure(&provider_name); + let budget_ms = inference_budget.as_millis() as u64; + self.emit_request_completed(&req.request_id, budget_ms, false); + self.metrics.record_request(false, budget_ms, false); + let err = EngineError::Timeout(format!( + "provider '{provider_name}' did not respond within {}s", + inference_budget.as_secs() + )); + StreamChunk::Error(err.info()) + } + }; + let _ = out.send(terminal).await; + } + + /// Ensure a model is resident, reserving memory before loading. + /// + /// The evict-then-reserve step runs under `load_gate` so concurrent loads + /// cannot jointly overcommit memory; the slow backend load then runs outside + /// the gate so independent models can load in parallel. async fn ensure_loaded( &self, model_id: &str, provider_name: &str, provider: &Arc, ) -> Result<(), EngineError> { - let residency = self - .models - .get(model_id) - .map(|m| m.residency) - .unwrap_or(ModelResidency::Unknown); - if matches!(residency, ModelResidency::Loaded | ModelResidency::Remote) { + if self.is_resident(model_id) { + self.memory.touch(model_id); return Ok(()); } - self.set_model_residency(model_id, ModelResidency::Loading); - let load_result = tokio::time::timeout(Duration::from_secs(30), provider.load(model_id)) - .await - .map_err(|_| { - self.circuit_breakers.record_failure(provider_name); - self.set_model_residency(model_id, ModelResidency::Unloaded); - EngineError::Timeout(format!("load timeout for {provider_name}/{model_id}")) - })? - .inspect_err(|_| { - self.circuit_breakers.record_failure(provider_name); - self.set_model_residency(model_id, ModelResidency::Unloaded); - })?; + let estimate = self + .models + .get(model_id) + .map(|m| m.memory_estimate_bytes) + .unwrap_or(0); - self.set_model_residency(model_id, load_result.residency); - if matches!(load_result.residency, ModelResidency::Loaded) - && let Some(card) = self.models.get(model_id) { - let bytes = load_result - .memory_bytes - .unwrap_or(card.memory_estimate_bytes); - if bytes > 0 { - if !self.memory.can_fit(bytes) { - self.evict_to_fit(model_id, bytes).await; + let _gate = self.load_gate.lock().await; + // Another task may have loaded this model while we waited for the gate. + if self.is_resident(model_id) { + self.memory.touch(model_id); + return Ok(()); + } + if estimate > 0 { + self.admit_memory(model_id, estimate).await?; + } + self.set_model_residency(model_id, ModelResidency::Loading); + } + + let load_result = + match tokio::time::timeout(self.timeouts.load(), provider.load(model_id)).await { + Ok(Ok(r)) => r, + Ok(Err(e)) => { + self.circuit_breakers.record_failure(provider_name); + self.rollback_reservation(model_id, estimate); + self.record_lifecycle(model_id, "load_failed", Some(&e.to_string())); + return Err(e); } - if !self.memory.can_fit(bytes) { - let _ = provider.unload(model_id).await; - self.set_model_residency(model_id, ModelResidency::Unloaded); - return Err(EngineError::MemoryPressure { - need: bytes, - available: self.memory.available_bytes(), - }); + Err(_) => { + self.circuit_breakers.record_failure(provider_name); + self.rollback_reservation(model_id, estimate); + self.record_lifecycle(model_id, "load_timeout", None); + return Err(EngineError::Timeout(format!( + "load timeout for {provider_name}/{model_id}" + ))); } - self.memory.allocate(model_id, bytes); - let _ = self.event_tx.send(EngineEvent::MemoryChanged { - used_bytes: self.memory.used_bytes(), - total_bytes: self.memory.budget_bytes(), - }); + }; + + self.set_model_residency(model_id, load_result.residency); + if matches!(load_result.residency, ModelResidency::Loaded) { + if let Some(observed) = load_result.memory_bytes + && observed > 0 + { + self.memory.reconcile(model_id, observed); } + // Mark freshly loaded so it is not the immediate LRU victim before its + // first lease is taken. + self.memory.touch(model_id); + self.emit_memory_changed(); + self.record_lifecycle(model_id, "load", None); } Ok(()) } - /// Evict least-recently-used idle models until `bytes` can fit, or until - /// nothing more can be freed. Models with in-flight requests and the - /// incoming model itself are never evicted. - async fn evict_to_fit(&self, incoming_id: &str, bytes: u64) { - let mut last_victim: Option = None; - while !self.memory.can_fit(bytes) { - let protected = self.protected_model_ids(incoming_id); + /// Reserve `bytes` for `model_id`, evicting idle models until it fits. + /// Caller must hold `load_gate`. + async fn admit_memory(&self, model_id: &str, bytes: u64) -> Result<(), EngineError> { + if self.memory.try_reserve(model_id, bytes) { + return Ok(()); + } + loop { + let mut protected = self.eviction_protected_ids(); + protected.push(model_id.to_string()); let Some(victim) = self.memory.lru_candidate(&protected) else { break; }; - // Defensive: if the previous eviction freed nothing, stop instead of spinning. - if last_victim.as_deref() == Some(victim.as_str()) { - break; - } - - let provider = self - .models - .get(&victim) - .and_then(|card| self.find_provider(&card.provider)); - if let Some(provider) = provider { - let _ = provider.unload(&victim).await; + self.unload_model(&victim, "memory_pressure").await; + if self.memory.try_reserve(model_id, bytes) { + return Ok(()); } - // Frees the allocation and emits ModelResidencyChanged/MemoryChanged. - self.set_model_residency(&victim, ModelResidency::Unloaded); - tracing::info!("evicted '{victim}' to admit {bytes} bytes"); - last_victim = Some(victim); } + Err(EngineError::MemoryPressure { + need: bytes, + available: self.memory.available_bytes(), + }) } - /// Model IDs that must never be evicted: the incoming model plus any model - /// with in-flight requests. - fn protected_model_ids(&self, incoming_id: &str) -> Vec { - let mut protected = vec![incoming_id.to_string()]; - protected.extend( - self.models - .iter() - .filter(|card| card.execution.active_requests > 0) - .map(|card| card.key().clone()), - ); - protected + fn rollback_reservation(&self, model_id: &str, estimate: u64) { + if estimate > 0 { + self.memory.deallocate(model_id); + } + self.set_model_residency(model_id, ModelResidency::Unloaded); + self.emit_memory_changed(); } - fn find_provider(&self, name: &str) -> Option> { - self.providers - .iter() - .find(|registered| registered.name == name) - .map(|registered| Arc::clone(®istered.provider)) + /// Unload a model from a backend and release its accounting. + async fn unload_model(&self, model_id: &str, reason: &str) { + let provider = self + .models + .get(model_id) + .and_then(|card| self.find_provider(&card.provider)); + if let Some(provider) = provider { + // 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); + // Distinguish the cause so metrics count each correctly: an idle sweep is + // an idle-unload, an operator action is a plain unload, and everything + // else (memory pressure) is an eviction. + let event = match reason { + "idle_timeout" => "idle_unload", + "manual" => "unload", + _ => "evict", + }; + self.record_lifecycle(model_id, event, Some(reason)); + let _ = self.event_tx.send(EngineEvent::ModelEvicted { + model_id: model_id.to_string(), + reason: reason.to_string(), + }); + tracing::info!("unloaded '{model_id}' ({reason})"); } - fn routing_providers(&self) -> Vec { - self.providers + /// Model IDs that must never be evicted by memory pressure or the idle sweep. + /// + /// Leased (in-flight) models are already excluded inside the memory manager; + /// this additionally protects: + /// - models whose load is currently in flight (`Loading`), so a concurrent + /// admission cannot unload a model out from under an active load and leave + /// its reservation orphaned; and + /// - resident models whose capability is covered by an active subscription + /// (the RFC "resident keep-alive" policy). + /// + /// If every candidate is protected, admission returns structured memory + /// pressure rather than breaking a subscription or a load. + fn eviction_protected_ids(&self) -> Vec { + let subscribed = self.subscriptions.active_capabilities(); + self.models + .iter() + .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(); + } + } + } + + /// Remember a prediction for later hit/miss confirmation, bounding the map + /// so caller-supplied scope keys cannot grow it without limit. + fn remember_prediction(&self, scope: String, capability: Capability) { + if self.pending_predictions.len() >= MAX_PENDING_PREDICTIONS + && !self.pending_predictions.contains_key(&scope) + && let Some(victim) = self + .pending_predictions + .iter() + .next() + .map(|e| e.key().clone()) + { + self.pending_predictions.remove(&victim); + } + self.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(), + ..Default::default() + }; + 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. + /// + /// An explicit per-call `ttl` wins; when `None` is passed the configured + /// `preflight.subscription_ttl_secs` default bounds the subscription (0 + /// disables expiry) so it cannot pin models warm forever. + /// + /// Returns the new subscription id. + pub fn subscribe( + &self, + app_id: Option, + session_id: Option, + capabilities: Vec, + ttl: Option, + ) -> u64 { + let ttl = ttl.or_else(|| { + let secs = self.preflight_config.subscription_ttl_secs; + (secs > 0).then(|| Duration::from_secs(secs)) + }); + let id = self.subscriptions.subscribe( + app_id.clone(), + session_id.clone(), + capabilities.clone(), + ttl, + ); + for cap in capabilities { + self.warm_capability(cap, "subscription", app_id.clone(), session_id.clone()); + } + id + } + + /// Remove a subscription by id. Returns whether one existed. + pub fn unsubscribe(&self, id: u64) -> bool { + self.subscriptions.unsubscribe(id) + } + + /// List all active subscriptions. + pub fn subscriptions(&self) -> Vec { + self.subscriptions.list() + } + + /// Snapshot of Preflight effectiveness counters. + pub fn preflight_stats(&self) -> PreflightStats { + self.preflight_metrics.snapshot() + } + + /// Background task: unload models idle beyond the configured timeout. + fn spawn_idle_eviction(&self) { + if self.idle_timeout.is_zero() { + return; + } + let Some(weak) = self.weak_self.get().cloned() else { + return; + }; + // Tick often enough to act soon after the timeout without busy-looping. + let tick = Duration::from_secs((self.idle_timeout.as_secs().max(2) / 2).clamp(1, 15)); + let handle = tokio::spawn(async move { + let mut ticker = tokio::time::interval(tick); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + ticker.tick().await; + let Some(engine) = weak.upgrade() else { + break; + }; + engine.run_idle_sweep().await; + } + }); + *self.idle_task.lock().unwrap_or_else(|e| e.into_inner()) = Some(handle); + } + + /// Background task: periodically delete stale engine artifacts. + fn spawn_artifact_sweep(&self, sweeper: crate::artifacts::ArtifactSweeper) { + let retention = sweeper.retention(); + if retention.is_zero() { + return; // retention 0 disables cleanup + } + let Some(weak) = self.weak_self.get().cloned() else { + return; + }; + // Sweep on roughly half the retention interval, bounded so tests and + // long-lived daemons both behave. + let tick = Duration::from_secs((retention.as_secs().max(2) / 2).clamp(1, 300)); + let handle = tokio::spawn(async move { + let mut ticker = tokio::time::interval(tick); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + ticker.tick().await; + if weak.upgrade().is_none() { + break; + } + let removed = sweeper.sweep(); + if removed > 0 { + tracing::debug!("artifact sweep removed {removed} stale file(s)"); + } + } + }); + *self.artifact_task.lock().unwrap_or_else(|e| e.into_inner()) = Some(handle); + } + + /// Reject a request whose local file inputs fall outside the configured + /// allowlist — `input_file` AND any `messages[].images` entry that names a + /// local path. http(s)/data: image URLs are not local paths and skip the + /// check. 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(()); + } + if let Some(path) = req.input_file.as_deref() { + self.check_path_within_roots(path, "input_file")?; + } + // Image references are resolved by the backends with a raw filesystem + // read (ollama.rs base64-encodes them; liter_llm.rs then UPLOADS the + // contents to a cloud provider) — an unchecked path is a local-file + // read primitive that becomes data exfiltration, so they go through + // the exact same allowlist as `input_file` before any routing. + for message in &req.messages { + for image in &message.images { + if image.starts_with("http://") + || image.starts_with("https://") + || image.starts_with("data:") + { + continue; + } + self.check_path_within_roots(image, "images[]")?; + } + } + Ok(()) + } + + /// Validate one local path against the configured allowlist. + fn check_path_within_roots(&self, path: &str, field: &str) -> Result<(), EngineError> { + let canonical = std::fs::canonicalize(path).map_err(|_| { + EngineError::InvalidRequest(format!("{field} '{path}' cannot be resolved")) + })?; + if self + .input_roots + .iter() + .any(|root| canonical.starts_with(root)) + { + Ok(()) + } else { + Err(EngineError::InvalidRequest(format!( + "{field} '{path}' is outside the allowed roots" + ))) + } + } + + /// Unload every model that has been idle past the timeout. + async fn run_idle_sweep(&self) { + let protected = self.eviction_protected_ids(); + let victims = self.memory.idle_candidates(self.idle_timeout, &protected); + for victim in victims { + // Re-check under no lock: only unload models that are genuinely + // resident and not leased right now. + let resident = self + .models + .get(&victim) + .map(|c| matches!(c.residency, ModelResidency::Loaded)) + .unwrap_or(false); + if !resident || self.memory.lease_count(&victim) > 0 { + continue; + } + self.unload_model(&victim, "idle_timeout").await; + } + } + + fn is_resident(&self, model_id: &str) -> bool { + self.models + .get(model_id) + .map(|m| matches!(m.residency, ModelResidency::Loaded | ModelResidency::Remote)) + .unwrap_or(false) + } + + fn semaphore_for(&self, model_id: &str, max_concurrency: u32) -> Arc { + let permits = max_concurrency.max(1) as usize; + self.semaphores + .entry(model_id.to_string()) + .or_insert_with(|| Arc::new(Semaphore::new(permits))) + .clone() + } + + fn find_provider(&self, name: &str) -> Option> { + self.providers + .iter() + .find(|registered| registered.name == name) + .map(|registered| Arc::clone(®istered.provider)) + } + + fn routing_providers(&self) -> Vec { + self.providers .iter() .map(|registered| { let health = self @@ -544,7 +1926,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 +1943,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 +1973,63 @@ 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>) { + self.metrics.record_lifecycle(event); + let record = LifecycleRecord { + seq: self.lifecycle_seq.fetch_add(1, Ordering::Relaxed), + at_ms: self.started_at.elapsed().as_millis() as u64, + model_id: model_id.to_string(), + event: event.to_string(), + detail: detail.map(ToString::to_string), + }; + let mut hist = self.lifecycle.lock().unwrap_or_else(|e| e.into_inner()); + if hist.len() >= LIFECYCLE_CAPACITY { + hist.pop_front(); + } + hist.push_back(record); + } + + /// Return the rolling lifecycle history, oldest first. + pub fn lifecycle_history(&self) -> Vec { + self.lifecycle + .lock() + .unwrap_or_else(|e| e.into_inner()) + .iter() + .cloned() + .collect() + } + + /// Return a snapshot of the engine's memory accounting. + pub fn memory_report(&self) -> MemoryReport { + MemoryReport { + used_bytes: self.memory.used_bytes(), + budget_bytes: self.memory.budget_bytes(), + available_bytes: self.memory.available_bytes(), + allocations: self.memory.snapshot(), + } + } + /// Get a snapshot of the engine status. pub async fn status(&self) -> EngineStatus { let all_models = self.capabilities().await; @@ -631,8 +2062,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 { @@ -654,9 +2085,92 @@ impl Engine { self.event_tx.subscribe() } - /// Engine uptime in seconds. - pub fn uptime_secs(&self) -> u64 { - self.started_at.elapsed().as_secs() + /// 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) + } +} + +impl Drop for Engine { + fn drop(&mut self) { + for task in [&self.idle_task, &self.artifact_task] { + if let Some(handle) = task.lock().unwrap_or_else(|e| e.into_inner()).take() { + handle.abort(); + } + } + // Cancel any in-flight speculative warm tasks. + for entry in self.warming.iter() { + entry.value().abort(); + } } } @@ -666,11 +2180,16 @@ 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; + + use crate::config::{ + EngineConfig, ListenConfig, MemoryConfig, ModelDef, PreflightConfig, ProviderConfig, + SecurityConfig, TimeoutConfig, }; - use std::sync::atomic::{AtomicUsize, Ordering}; - use crate::config::{EngineConfig, ListenConfig, MemoryConfig, ProviderConfig}; + const MB: u64 = 1024 * 1024; fn minimal_config() -> EngineConfig { EngineConfig { @@ -679,51 +2198,104 @@ mod tests { budget_mb: Some(100), idle_timeout_secs: 60, }, + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), providers: vec![], } } - #[derive(Default)] - struct MockProvider { + /// Behavior knobs for the configurable test provider. + #[derive(Clone, Copy)] + enum InvokeBehavior { + Ok, + SlowOkMs(u64), + FailRetryable, + FailInvalid, + } + + struct TestProvider { name: String, - discover_calls: AtomicUsize, + kind: ProviderKind, + /// (model_name, capability, residency, mem_bytes, max_concurrency) + template: Vec<(String, Capability, ModelResidency, u64, u32)>, + behavior: InvokeBehavior, + load_should_fail: bool, + in_flight: Arc, + max_in_flight: Arc, + invoke_calls: Arc, } - impl MockProvider { - fn new(name: &str) -> Self { + impl TestProvider { + fn new(name: &str, kind: ProviderKind) -> Self { Self { name: name.into(), - discover_calls: AtomicUsize::new(0), + kind, + template: Vec::new(), + behavior: InvokeBehavior::Ok, + load_should_fail: false, + in_flight: Arc::new(AtomicUsize::new(0)), + max_in_flight: Arc::new(AtomicUsize::new(0)), + invoke_calls: Arc::new(AtomicUsize::new(0)), } } + + fn with_model( + mut self, + name: &str, + cap: Capability, + residency: ModelResidency, + mem_bytes: u64, + max_concurrency: u32, + ) -> Self { + self.template + .push((name.into(), cap, residency, mem_bytes, max_concurrency)); + self + } + + fn behavior(mut self, behavior: InvokeBehavior) -> Self { + self.behavior = behavior; + self + } + + fn failing_load(mut self) -> Self { + self.load_should_fail = true; + self + } } #[async_trait] - impl Provider for MockProvider { + impl Provider for TestProvider { fn name(&self) -> &str { &self.name } - fn kind(&self) -> ProviderKind { - ProviderKind::OpenAiCompatible + self.kind } - fn features(&self) -> Vec { - vec![BackendFeature::Discovery] + vec![BackendFeature::Discovery, BackendFeature::Load] } async fn discover(&self) -> Result, EngineError> { - self.discover_calls.fetch_add(1, Ordering::SeqCst); - let mut card = ModelCard::new(&self.name, "mock-chat", Capability::Chat, CostTier::Low); - card.id = canonical_model_id(&self.name, "mock-chat"); - card.availability = ModelAvailability::Configured; - card.residency = ModelResidency::Remote; - card.execution = ExecutionState { - active_requests: 0, - max_concurrency: 4, - }; - card.refresh_status(); - Ok(vec![card]) + let cards = self + .template + .iter() + .map(|(name, cap, residency, mem, conc)| { + let mut card = ModelCard::new(&self.name, name, *cap, CostTier::Low); + card.id = ModelId::canonical(&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 +2303,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, ..)| ModelId::canonical(&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,62 +2341,145 @@ 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, - 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(), tokens_used: Some(1), fallback_used: false, routing_reason: None, + ..Default::default() }) } } - fn engine_with_provider(provider: Arc) -> Arc { + fn build_engine( + 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); - 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(), + conversations: DashMap::new(), + timeouts: TimeoutConfig::default(), + preflight_config: PreflightConfig::default(), + idle_timeout: Duration::from_secs(idle_secs), + semaphores: DashMap::new(), + load_gate: AsyncMutex::new(()), + lifecycle: Mutex::new(VecDeque::new()), + lifecycle_seq: AtomicU64::new(0), + idle_task: Mutex::new(None), + artifact_task: Mutex::new(None), + weak_self: OnceLock::new(), event_tx, + metrics: EngineMetrics::default(), + pricing, + input_roots: Vec::new(), started_at: Instant::now(), - }) + }); + let _ = engine.weak_self.set(Arc::downgrade(&engine)); + engine.spawn_idle_eviction(); + engine } fn chat_request(model: Option<&str>) -> InferenceRequest { + request(Capability::Chat, model, FallbackPolicy::default()) + } + + fn request( + capability: Capability, + model: Option<&str>, + fallback_policy: FallbackPolicy, + ) -> InferenceRequest { InferenceRequest { - capability: Some(Capability::Chat), + capability: Some(capability), model: model.map(str::to_owned), app_id: None, session_id: None, - fallback_policy: FallbackPolicy::default(), + fallback_policy, messages: vec![Message { role: "user".into(), content: "hello".into(), + ..Default::default() }], input_file: None, params: serde_json::Value::Null, hint_next: None, request_id: "test".into(), + ..Default::default() } } #[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()); + } + + #[test] + fn cost_estimation_from_prices() { + // $0.01/1k in, $0.03/1k out; 1000 prompt + 500 completion tokens. + 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!(Engine::estimate_cost(0.0, 0.0, Some(1000), Some(500)).is_none()); + // Missing token counts default to zero. + assert_eq!(Engine::estimate_cost(0.01, 0.03, None, None), Some(0.0)); } #[tokio::test] @@ -826,18 +2497,92 @@ mod tests { assert!(matches!(result, Err(EngineError::NoCapableModel(_)))); } + #[tokio::test] + async fn system_voice_backs_a_tts_less_config() { + // A config with one enabled, non-TTS provider gains the built-in system + // voice as a fallback wherever the OS provides one (macOS `say` / Linux + // `espeak`). `true` stands in as a harmless, network-free provider command. + let mut config = minimal_config(); + config.providers = vec![ProviderConfig { + name: "img".into(), + kind: "local_image_gen".into(), + command: Some("true".into()), + enabled: true, + models: vec![ModelDef { + name: "sd".into(), + capability: "image_gen".into(), + ..Default::default() + }], + ..Default::default() + }]; + let engine = Engine::new(config).await; + engine.refresh_resources().await; + + let voice_available = + crate::backends::system_tts::SystemTtsProvider::new(None).is_available(); + let injected = engine + .capabilities() + .await + .iter() + .any(|c| c.id == "system-tts/system"); + assert_eq!(injected, voice_available); + + // An empty config is left untouched: no providers in, no voice out. + let empty = Engine::new(minimal_config()).await; + empty.refresh_resources().await; + assert!( + empty + .capabilities() + .await + .iter() + .all(|c| c.id != "system-tts/system") + ); + } + + #[tokio::test] + async fn cloud_video_gen_config_exposes_a_video_capability() { + // A `cloud_video_gen` provider (Seedance / Ark) wires through config → the + // backend factory → discovery, exposing a first-class `video_gen` + // capability. A dummy key keeps it network-free: `health()` and `discover()` + // for this backend never touch the wire, so the card is `Configured` + // without a live API. + let mut config = minimal_config(); + config.providers = vec![ProviderConfig { + name: "seedance".into(), + kind: "cloud_video_gen".into(), + api_key: Some("test-key".into()), + cost_tier: "high".into(), + enabled: true, + models: vec![ModelDef { + name: "doubao-seedance-1-0-pro".into(), + capability: "video_gen".into(), + ..Default::default() + }], + ..Default::default() + }]; + let engine = Engine::new(config).await; + engine.refresh_resources().await; + + let card = engine + .capabilities() + .await + .into_iter() + .find(|c| c.provider == "seedance") + .expect("the configured cloud video provider is discovered"); + assert_eq!(card.capability, mofa_kernel::Capability::VideoGen); + // Cloud model: remote residency, not counted against local memory. + assert_eq!(card.residency, ModelResidency::Remote); + } + #[tokio::test] 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 +2599,10 @@ mod tests { budget_mb: Some(100), idle_timeout_secs: 60, }, + timeouts: TimeoutConfig::default(), + preflight: PreflightConfig::default(), + artifacts: Default::default(), + security: Default::default(), providers: vec![ProviderConfig { name: "disabled-ollama".into(), kind: "ollama".into(), @@ -863,99 +2612,1198 @@ mod tests { cost_tier: "free".into(), models: vec![], enabled: false, + ..Default::default() }], }; let engine = Engine::new(config).await; - let status = engine.status().await; - assert_eq!(status.providers, 0); + assert_eq!(engine.status().await.providers, 0); } #[tokio::test] - async fn mock_provider_refresh_and_invoke_are_deterministic() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); + async fn discovers_and_invokes_deterministically() { + let provider = Arc::new( + TestProvider::new("cloud", ProviderKind::OpenAiCompatible).with_model( + "mock-chat", + Capability::Chat, + ModelResidency::Remote, + 0, + 4, + ), + ); + let engine = build_engine(vec![provider], 100, 0); engine.refresh_resources().await; let caps = engine.capabilities().await; assert_eq!(caps.len(), 1); - assert_eq!(caps[0].id, "mock/mock-chat"); + assert_eq!(caps[0].id, "cloud/mock-chat"); let resp = engine.invoke(chat_request(None)).await.unwrap(); assert_eq!(resp.text.as_deref(), Some("ok")); - assert_eq!(resp.provider, "mock"); + assert_eq!(resp.provider, "cloud"); assert!(resp.routing_reason.is_some()); + assert!(!resp.fallback_used); } #[tokio::test] - async fn ambiguous_short_model_name_is_rejected() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); - let mut a = ModelCard::new("a", "same", Capability::Chat, CostTier::Free); - a.residency = ModelResidency::Remote; - a.refresh_status(); - let mut b = ModelCard::new("b", "same", Capability::Chat, CostTier::Free); - b.residency = ModelResidency::Remote; - b.refresh_status(); - engine.models.insert(a.id.clone(), a); - engine.models.insert(b.id.clone(), b); + async fn responses_api_accumulates_multi_turn_history() { + use mofa_kernel::{ResponsesRequest, ResponsesResponse}; + + 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; + + // Turn 1: a fresh conversation with system instructions + user input. + let first: ResponsesResponse = engine + .respond(ResponsesRequest { + instructions: Some("be terse".into()), + input: Some("hello".into()), + ..Default::default() + }) + .await + .unwrap(); + // system + user + assistant("ok") = 3 messages stored under a fresh id. + assert_eq!(first.message_count, 3); + assert_eq!(first.response.text.as_deref(), Some("ok")); + assert!(!first.id.is_empty()); + assert_eq!(engine.conversation_count(), 1); + + // Turn 2: continue from the prior response id with a new user input. + let second = engine + .respond(ResponsesRequest { + previous_response_id: Some(first.id.clone()), + input: Some("again".into()), + ..Default::default() + }) + .await + .unwrap(); + // Prior 3 + new user + new assistant = 5, and the new turn gets its own id. + assert_eq!(second.message_count, 5); + assert_ne!(second.id, first.id); + + // The stored history carries the full accumulated dialogue in order. + let history = engine.conversation_messages(&second.id).unwrap(); + assert_eq!(history.len(), 5); + assert_eq!(history[0].role, "system"); + assert_eq!(history[1].content, "hello"); + assert_eq!(history[3].content, "again"); + + // An unknown prior id is a hard (non-retryable) error, not a fresh start. + let err = engine + .respond(ResponsesRequest { + previous_response_id: Some("does-not-exist".into()), + input: Some("x".into()), + ..Default::default() + }) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + assert!(!err.retryable()); + // A turn with no input at all is rejected. + assert!(matches!( + engine.respond(ResponsesRequest::default()).await, + Err(EngineError::InvalidRequest(_)) + )); + + // Deleting a conversation forgets it. + assert!(engine.delete_conversation(&second.id)); + assert!(engine.conversation_messages(&second.id).is_none()); + } + + #[tokio::test] + async fn ambiguous_short_model_name_is_rejected() { + 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 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.prefer = mofa_kernel::Prefer::Local; + // 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( + 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 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()); + } + + #[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!(Engine::estimate_prompt_tokens(&req), 1); + req.messages[0].content = "x".repeat(400); + assert_eq!(Engine::estimate_prompt_tokens(&req), 100); + // Completion defaults to the generous estimate, or the params cap. + assert_eq!( + Engine::estimate_completion_tokens(&req), + Engine::DEFAULT_COMPLETION_TOKEN_ESTIMATE + ); + req.params = serde_json::json!({ "max_tokens": 64 }); + assert_eq!(Engine::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 + // engine must fail over to the configured cloud TTS model. This is RFC + // acceptance step 8 exercised deterministically without a real backend. + let local_tts = Arc::new( + TestProvider::new("local-tts", ProviderKind::LocalTts) + .with_model( + "kokoro", + Capability::Tts, + ModelResidency::Loaded, + 10 * MB, + 1, + ) + .behavior(InvokeBehavior::FailRetryable), + ); + let cloud_tts = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible) + .with_model("tts-1", Capability::Tts, ModelResidency::Remote, 0, 32) + .behavior(InvokeBehavior::Ok), + ); + let local_calls = local_tts.invoke_calls.clone(); + let engine = build_engine(vec![local_tts, cloud_tts], 1000, 0); + engine.refresh_resources().await; + + let resp = engine + .invoke(request(Capability::Tts, None, FallbackPolicy::default())) + .await + .unwrap(); + // The local backend was tried first, then the engine fell over to cloud. + assert_eq!(local_calls.load(Ordering::SeqCst), 1); + assert_eq!(resp.provider, "openai"); + assert_eq!(resp.model_used, "tts-1"); + assert!(resp.fallback_used); + } + + #[tokio::test] + async fn invalid_request_does_not_fall_back() { + let local = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model("qwen", Capability::Chat, ModelResidency::Loaded, 10 * MB, 1) + .behavior(InvokeBehavior::FailInvalid), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible) + .with_model("gpt", Capability::Chat, ModelResidency::Remote, 0, 32) + .behavior(InvokeBehavior::Ok), + ); + let cloud_calls = cloud.invoke_calls.clone(); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + let err = engine.invoke(chat_request(None)).await.unwrap_err(); + assert!(matches!(err, EngineError::InvalidRequest(_))); + // The cloud candidate must never have been tried. + assert_eq!(cloud_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn named_strict_does_not_fall_back() { + let local = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model("qwen", Capability::Chat, ModelResidency::Loaded, 10 * MB, 1) + .behavior(InvokeBehavior::FailRetryable), + ); + let cloud = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible) + .with_model("gpt", Capability::Chat, ModelResidency::Remote, 0, 32) + .behavior(InvokeBehavior::Ok), + ); + let cloud_calls = cloud.invoke_calls.clone(); + let engine = build_engine(vec![local, cloud], 1000, 0); + engine.refresh_resources().await; + + // Default policy is strict for named requests. + let err = engine.invoke(chat_request(Some("qwen"))).await.unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + assert_eq!(cloud_calls.load(Ordering::SeqCst), 0); + } + + #[tokio::test] + async fn ten_concurrent_requests_respect_concurrency_limit() { + let provider = Arc::new( + TestProvider::new("cloud", ProviderKind::OpenAiCompatible) + .with_model("m", Capability::Chat, ModelResidency::Remote, 0, 4) + .behavior(InvokeBehavior::SlowOkMs(40)), + ); + let max_seen = provider.max_in_flight.clone(); + let engine = build_engine(vec![provider], 1000, 0); + engine.refresh_resources().await; + + let mut handles = Vec::new(); + for _ in 0..10 { + let engine = Arc::clone(&engine); + handles.push(tokio::spawn(async move { + engine.invoke(chat_request(None)).await + })); + } + for h in handles { + assert!(h.await.unwrap().is_ok()); + } + // The semaphore must have capped concurrency at the model's limit. + assert!( + max_seen.load(Ordering::SeqCst) <= 4, + "over-admitted requests" + ); + } + #[tokio::test] async fn memory_pressure_evicts_lru_idle_model() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); // 100 MB budget + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 0); // 100 MB budget - // Two resident, idle models at 40 MB each (80 MB used). for name in ["old", "new"] { - let mut card = ModelCard::new("mock", name, Capability::Chat, CostTier::Low); + let mut card = ModelCard::new("ollama", name, Capability::Chat, CostTier::Low); card.residency = ModelResidency::Loaded; card.refresh_status(); let id = card.id.clone(); engine.models.insert(id.clone(), card); - engine.memory.allocate(&id, 40 * 1024 * 1024); + engine.memory.allocate(&id, 40 * MB); } - // Touch "new" so "old" is the least-recently-used victim. - engine.memory.touch("mock/new"); - assert_eq!(engine.memory.used_bytes(), 80 * 1024 * 1024); + engine.memory.touch("ollama/new"); // make "old" the LRU victim + assert_eq!(engine.memory.used_bytes(), 80 * MB); - // Admitting another 40 MB needs 20 MB freed → evict exactly the LRU model. - engine.evict_to_fit("mock/incoming", 40 * 1024 * 1024).await; + // Admitting 40 MB needs 20 MB freed → evict exactly the LRU model. + engine + .admit_memory("ollama/incoming", 40 * MB) + .await + .unwrap(); - assert!(engine.memory.can_fit(40 * 1024 * 1024)); assert_eq!( - engine.models.get("mock/old").unwrap().residency, + engine.models.get("ollama/old").unwrap().residency, ModelResidency::Unloaded ); assert_eq!( - engine.models.get("mock/new").unwrap().residency, + engine.models.get("ollama/new").unwrap().residency, ModelResidency::Loaded ); + assert!(engine.memory.lease_count("ollama/incoming") == 0); } #[tokio::test] - async fn busy_models_are_never_evicted() { - let provider = Arc::new(MockProvider::new("mock")); - let engine = engine_with_provider(provider); // 100 MB budget + async fn leased_models_are_never_evicted() { + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 0); - // A single resident model with an in-flight request, filling the budget. - let mut card = ModelCard::new("mock", "busy", Capability::Chat, CostTier::Low); + let mut card = ModelCard::new("ollama", "busy", Capability::Chat, CostTier::Low); card.residency = ModelResidency::Loaded; - card.execution.active_requests = 1; card.refresh_status(); let id = card.id.clone(); engine.models.insert(id.clone(), card); - engine.memory.allocate(&id, 90 * 1024 * 1024); + engine.memory.allocate(&id, 90 * MB); + engine.memory.lease(&id); // in-flight inference + + // Cannot evict the only (leased) model → memory pressure. + let err = engine + .admit_memory("ollama/incoming", 40 * MB) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::MemoryPressure { .. })); + assert_eq!( + engine.models.get("ollama/busy").unwrap().residency, + ModelResidency::Loaded + ); + } + + #[tokio::test] + async fn load_failure_rolls_back_reservation() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 10 * MB, + 1, + ) + .failing_load(), + ); + let engine = build_engine(vec![provider], 100, 0); + engine.refresh_resources().await; + + let err = engine.invoke(chat_request(None)).await.unwrap_err(); + assert!(matches!(err, EngineError::ProviderError { .. })); + // The reservation must have been rolled back. + assert_eq!(engine.memory.used_bytes(), 0); + assert_eq!( + engine.models.get("ollama/qwen").unwrap().residency, + ModelResidency::Unloaded + ); + } - // It cannot be evicted, so the budget stays full. - engine.evict_to_fit("mock/incoming", 40 * 1024 * 1024).await; + #[tokio::test] + async fn successful_load_reserves_and_records_memory() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 20 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 100, 0); + engine.refresh_resources().await; - assert!(!engine.memory.can_fit(40 * 1024 * 1024)); + engine.invoke(chat_request(None)).await.unwrap(); + assert_eq!(engine.memory.used_bytes(), 20 * MB); assert_eq!( - engine.models.get("mock/busy").unwrap().residency, + engine.models.get("ollama/qwen").unwrap().residency, ModelResidency::Loaded ); + // A load event must be in the lifecycle history. + assert!( + engine + .lifecycle_history() + .iter() + .any(|r| r.model_id == "ollama/qwen" && r.event == "load") + ); + } + + #[tokio::test] + async fn loading_model_is_not_evicted() { + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 0); + + // A model whose load is in flight, holding most of the budget. + let mut card = ModelCard::new("ollama", "loading", Capability::Chat, CostTier::Low); + card.residency = ModelResidency::Loading; + card.refresh_status(); + let id = card.id.clone(); + engine.models.insert(id.clone(), card); + engine.memory.allocate(&id, 90 * MB); + + // It must be protected, so admission cannot satisfy the request. + assert!(engine.eviction_protected_ids().contains(&id)); + let err = engine + .admit_memory("ollama/incoming", 40 * MB) + .await + .unwrap_err(); + assert!(matches!(err, EngineError::MemoryPressure { .. })); + assert_eq!( + engine.models.get(&id).unwrap().residency, + ModelResidency::Loading + ); + } + + #[tokio::test] + async fn rediscovery_releases_reservation_for_unloaded_model() { + // The provider reports the model as Unloaded on discovery. + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 100, 0); + + // The engine currently believes the model is loaded and reserved. + let mut card = ModelCard::new("ollama", "qwen", Capability::Chat, CostTier::Low); + card.residency = ModelResidency::Loaded; + card.refresh_status(); + engine.models.insert(card.id.clone(), card); + engine.memory.allocate("ollama/qwen", 10 * MB); + assert_eq!(engine.memory.used_bytes(), 10 * MB); + + // Rediscovery reconciles: the backend says Unloaded, so the reservation + // is released to match reality. + engine.refresh_resources().await; + assert_eq!( + engine.models.get("ollama/qwen").unwrap().residency, + ModelResidency::Unloaded + ); + assert_eq!(engine.memory.used_bytes(), 0); + } + + async fn wait_until(mut cond: impl FnMut() -> bool) -> bool { + // Up to ~4s, returning as soon as the condition holds. + for _ in 0..400 { + if cond() { + return true; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + cond() + } + + fn scoped_request(cap: Capability, app: &str, session: &str) -> InferenceRequest { + let mut r = request(cap, None, FallbackPolicy::default()); + r.app_id = Some(app.into()); + r.session_id = Some(session.into()); + r + } + + #[tokio::test] + async fn hint_warms_next_model_concurrently() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model("qwen", Capability::Chat, ModelResidency::Loaded, 10 * MB, 1) + .with_model( + "kokoro", + Capability::Tts, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 1000, 0); + engine.refresh_resources().await; + + let mut req = chat_request(None); + req.hint_next = Some("tts".into()); + engine.invoke(req).await.unwrap(); + + let engine_ref = Arc::clone(&engine); + let warmed = wait_until(move || engine_ref.is_resident("ollama/kokoro")).await; + assert!(warmed, "the hinted TTS model should have been warmed"); + assert!(engine.preflight_stats().warms_started >= 1); + } + + #[tokio::test] + async fn subscription_warms_and_protects_model() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "kokoro", + Capability::Tts, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 1000, 1); + engine.refresh_resources().await; + + engine.subscribe(Some("mofa-fm".into()), None, vec![Capability::Tts], None); + + let engine_ref = Arc::clone(&engine); + let warmed = wait_until(move || engine_ref.is_resident("ollama/kokoro")).await; + assert!(warmed, "subscription should warm its capability's model"); + // A subscribed resident model is protected from eviction. + assert!( + engine + .eviction_protected_ids() + .contains(&"ollama/kokoro".to_string()) + ); + } + + #[tokio::test] + async fn history_predicts_and_confirms_hits() { + let chat = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama).with_model( + "qwen", + Capability::Chat, + ModelResidency::Loaded, + 10 * MB, + 4, + ), + ); + let tts = Arc::new( + TestProvider::new("openai", ProviderKind::OpenAiCompatible).with_model( + "tts-1", + Capability::Tts, + ModelResidency::Remote, + 0, + 32, + ), + ); + let engine = build_engine(vec![chat, tts], 1000, 0); + engine.refresh_resources().await; + + // Build a Chat → Tts habit within one session. + for _ in 0..5 { + engine + .invoke(scoped_request(Capability::Chat, "fm", "s1")) + .await + .unwrap(); + engine + .invoke(scoped_request(Capability::Tts, "fm", "s1")) + .await + .unwrap(); + } + + let stats = engine.preflight_stats(); + assert!(stats.predictions > 0, "history should produce predictions"); + assert!(stats.hits > 0, "predictions should be confirmed as hits"); + } + + /// The headline article-to-podcast flow across Stages 3-5: a hinted chat + /// call warms TTS concurrently, the follow-up TTS call is served by the + /// already-warm model, and both then unload on idle. + #[tokio::test] + async fn article_to_podcast_flow() { + let provider = Arc::new( + TestProvider::new("ollama", ProviderKind::Ollama) + .with_model( + "qwen", + Capability::Chat, + ModelResidency::Unloaded, + 20 * MB, + 1, + ) + .with_model( + "kokoro", + Capability::Tts, + ModelResidency::Unloaded, + 10 * MB, + 1, + ), + ); + let engine = build_engine(vec![provider], 1000, 1); + engine.refresh_resources().await; + + // Step 1: LLM translation, hinting that TTS is next. + let mut chat = scoped_request(Capability::Chat, "mofa-fm", "s1"); + chat.hint_next = Some("tts".into()); + let chat_resp = engine.invoke(chat).await.unwrap(); + assert_eq!(chat_resp.model_used, "qwen"); + + // The hint should have warmed kokoro concurrently. + let engine_ref = Arc::clone(&engine); + assert!( + wait_until(move || engine_ref.is_resident("ollama/kokoro")).await, + "TTS should be warmed by the hint before it is requested" + ); + + // Step 2: TTS synthesis — served by the already-resident model. + let tts_resp = engine + .invoke(scoped_request(Capability::Tts, "mofa-fm", "s1")) + .await + .unwrap(); + assert_eq!(tts_resp.model_used, "kokoro"); + assert!(!tts_resp.fallback_used); + assert!(engine.preflight_stats().warms_started >= 1); + + // Step 3: both models unload once idle, releasing memory. + let engine_ref = Arc::clone(&engine); + assert!( + wait_until(move || engine_ref.memory.used_bytes() == 0).await, + "idle models should unload and release memory" + ); + } + + #[tokio::test] + async fn idle_sweep_unloads_stale_models() { + let provider = Arc::new(TestProvider::new("ollama", ProviderKind::Ollama)); + let engine = build_engine(vec![provider], 100, 1); // 1s idle timeout + + let mut card = ModelCard::new("ollama", "qwen", Capability::Chat, CostTier::Low); + card.residency = ModelResidency::Loaded; + card.refresh_status(); + let id = card.id.clone(); + engine.models.insert(id.clone(), card); + engine.memory.allocate(&id, 30 * MB); + + // The idle timeout uses monotonic last-access; wait past it then sweep. + tokio::time::sleep(Duration::from_millis(1100)).await; + engine.run_idle_sweep().await; + + assert_eq!( + engine.models.get(&id).unwrap().residency, + ModelResidency::Unloaded + ); + assert_eq!(engine.memory.used_bytes(), 0); + assert!( + engine + .lifecycle_history() + .iter() + .any(|r| r.event == "idle_unload") + ); + } + + async fn collect_stream(mut rx: mpsc::Receiver) -> Vec { + let mut chunks = Vec::new(); + while let Some(c) = rx.recv().await { + chunks.push(c); + } + chunks + } + + /// A provider that emits several real text deltas, to prove the engine + /// forwards incremental output in order (not just the compatibility path). + /// `reasoning` deltas (if any) are emitted first as thought-chain increments. + struct MultiDeltaProvider { + name: String, + deltas: Vec, + reasoning: 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::ModelId::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, + ..Default::default() + }) + } + async fn stream( + &self, + model_id: &str, + 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(StreamDelta::Text(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()], + reasoning: vec![], + }); + 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 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(); + 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 image_paths_outside_input_roots_are_rejected() { + // #4 review (security): `messages[].images` entries that name LOCAL + // FILES must pass the same `security.input_roots` validation as + // `input_file` — the ollama path base64-reads them and the liter-llm + // path uploads their contents to a cloud provider, so an unchecked + // path is a local-file-read primitive that becomes exfiltration. + let root = tempfile::tempdir().unwrap(); + let allowed = root.path().join("frame.png"); + std::fs::write(&allowed, b"x").unwrap(); + let outside = std::env::temp_dir().join("mofa_outside_allowlist.png"); + 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; + + let image_req = |image: String| { + let mut req = request(Capability::Vlm, None, FallbackPolicy::default()); + req.messages = vec![Message { + role: "user".into(), + content: "describe this".into(), + images: vec![image], + }]; + req + }; + + // A local image path outside the allowlist is rejected before routing. + assert!(matches!( + engine + .invoke(image_req(outside.to_string_lossy().to_string())) + .await, + Err(EngineError::InvalidRequest(_)) + )); + // A path inside the allowlist passes the check (then fails later only + // because no VLM model is registered). + assert!(matches!( + engine + .invoke(image_req(allowed.to_string_lossy().to_string())) + .await, + Err(EngineError::NoCapableModel(_)) + )); + // http(s) and data: URLs are not local paths — the allowlist does not + // apply to them. + for remote in ["https://example.com/a.png", "data:image/png;base64,AAAA"] { + assert!( + matches!( + engine.invoke(image_req(remote.into())).await, + Err(EngineError::NoCapableModel(_)) + ), + "{remote} must not be treated as a local path" + ); + } + } + + #[tokio::test] + async fn subscription_without_explicit_ttl_uses_config_default() { + // #4 review: `PreflightConfig::subscription_ttl_secs` must actually + // bound subscriptions — previously nothing read it, so a subscription + // that omitted the TTL pinned its models warm forever. + let config = EngineConfig { + preflight: PreflightConfig { + subscription_ttl_secs: 120, + ..PreflightConfig::default() + }, + ..minimal_config() + }; + let engine = Engine::new(config).await; + + engine.subscribe(None, None, vec![Capability::Chat], None); + let info = engine.subscriptions(); + assert_eq!(info.len(), 1); + let ttl = info[0] + .expires_in_secs + .expect("the configured default TTL must apply when none is given"); + assert!(ttl > 110 && ttl <= 120, "ttl {ttl} should be ~120"); + + // An explicit per-call TTL still wins over the config default. + engine.subscribe( + None, + None, + vec![Capability::Tts], + Some(Duration::from_secs(5)), + ); + let explicit = engine + .subscriptions() + .into_iter() + .find(|s| s.capabilities == vec![Capability::Tts]) + .and_then(|s| s.expires_in_secs) + .expect("explicit TTL reported"); + assert!(explicit <= 5, "explicit TTL must win, got {explicit}"); + } + + #[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:?}"), + } + } + + #[tokio::test] + async fn stream_predispatch_failure_is_counted_in_metrics() { + // A stream rejected *before* generation begins (here: no capable model) + // must still be counted as a handled+failed request, exactly like the + // equivalent unary `invoke` failure — otherwise the streaming path + // undercounts failures in `mofa_requests_total` / `_failed_total`. + let engine = build_engine(vec![], 100, 0); + let chunks = collect_stream(engine.invoke_stream(chat_request(None))).await; + assert!(matches!(chunks.as_slice(), [StreamChunk::Error(_)])); + + let metrics = engine.metrics_prometheus(); + assert!( + metrics.contains("mofa_requests_total 1"), + "pre-dispatch stream failure was not counted as a handled request:\n{metrics}" + ); + assert!( + metrics.contains("mofa_requests_failed_total 1"), + "pre-dispatch stream failure was not counted as a failed request:\n{metrics}" + ); } } diff --git a/mofa-engine-core/src/lib.rs b/mofa-engine-core/src/lib.rs index f747665..2336fd4 100644 --- a/mofa-engine-core/src/lib.rs +++ b/mofa-engine-core/src/lib.rs @@ -4,13 +4,21 @@ //! memory management, circuit breaker, preflight prediction, //! and the main `Engine` orchestrator. -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 preflight; -pub mod router; +pub mod quality_gate; +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..aa3fa7d 100644 --- a/mofa-engine-core/src/memory.rs +++ b/mofa-engine-core/src/memory.rs @@ -1,32 +1,72 @@ -//! Memory manager with LRU eviction. +//! Reservation-based memory manager with LRU eviction. //! -//! Tracks per-model memory allocations, detects real system memory, -//! and evicts least-recently-used models when under pressure. +//! Memory is admitted *before* a model is loaded: callers atomically reserve +//! capacity, evict idle models if needed, and only then ask the backend to +//! load. This prevents the time-of-check/time-of-use race where two concurrent +//! loads each observe enough free memory and together overcommit the device. +//! +//! Three quantities are tracked per model: +//! - **reserved** — the bytes admitted against the budget (a conservative +//! estimate, later reconciled toward observed usage); +//! - **observed** — actual usage reported by the backend, when available; +//! - **leases** — in-flight inferences; a model with active leases is never +//! evicted, even under pressure. +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Mutex; -use std::time::Instant; +use std::time::{Duration, Instant}; -/// Tracks memory allocations and enforces a budget. -pub struct MemoryManager { - /// Total budget in bytes +/// Tracks per-model memory reservations and enforces a global budget. +pub(crate) struct MemoryManager { + /// Total budget in bytes. budget_bytes: u64, - /// Per-model allocations: model_id → (bytes, last_access) + /// Per-model allocations. allocations: Mutex>, } struct Allocation { - bytes: u64, + /// Bytes admitted against the budget. Accounting always uses this figure. + reserved_bytes: u64, + /// Actual usage reported by the backend, when known. + observed_bytes: Option, + /// Monotonic time of the last access, for LRU and idle-timeout decisions. last_access: Instant, + /// Active inference leases. A model with `leases > 0` is never evicted. + leases: u32, +} + +impl Allocation { + fn new(reserved_bytes: u64) -> Self { + Self { + reserved_bytes, + observed_bytes: None, + last_access: Instant::now(), + leases: 0, + } + } +} + +/// A point-in-time view of one model's memory accounting. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct AllocationSnapshot { + /// Model identifier. + pub model_id: String, + /// Bytes reserved against the budget. + pub reserved_bytes: u64, + /// Bytes observed by the backend, if reported. + pub observed_bytes: Option, + /// Active inference leases. + pub leases: u32, } impl MemoryManager { /// Create a new memory manager. /// /// If `budget_mb` is `None`, auto-detect from system RAM (use 70% of total). - pub fn new(budget_mb: Option) -> Self { + pub(crate) fn new(budget_mb: Option) -> Self { let budget_bytes = match budget_mb { - Some(mb) => mb * 1024 * 1024, + Some(mb) => mb.saturating_mul(1024 * 1024), None => Self::detect_system_memory(), }; @@ -45,69 +85,167 @@ 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 { + pub(crate) fn budget_bytes(&self) -> u64 { self.budget_bytes } - /// Currently used bytes. - 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() + /// Currently reserved bytes across all models. + pub(crate) fn used_bytes(&self) -> u64 { + self.lock().values().map(|a| a.reserved_bytes).sum() } - /// Available bytes. - pub fn available_bytes(&self) -> u64 { + /// Bytes still available within the budget. + pub(crate) fn available_bytes(&self) -> u64 { self.budget_bytes.saturating_sub(self.used_bytes()) } - /// Record that a model is using memory. Updates last-access time. - 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(), - }, - ); + /// Whether `needed_bytes` currently fits within the budget without eviction. + /// Test-support (this module only): production paths reserve atomically via + /// [`Self::try_reserve`]. + #[cfg(test)] + fn can_fit(&self, needed_bytes: u64) -> bool { + self.available_bytes() >= needed_bytes } - /// Remove a model's allocation. - pub fn deallocate(&self, model_id: &str) { - let mut allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); - allocs.remove(model_id); + /// 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(crate) fn try_reserve(&self, model_id: &str, bytes: u64) -> bool { + let mut allocs = self.lock(); + let used_by_others: u64 = allocs + .iter() + .filter(|(id, _)| id.as_str() != model_id) + .map(|(_, a)| a.reserved_bytes) + .sum(); + if used_by_others.saturating_add(bytes) > self.budget_bytes { + return false; + } + match allocs.get_mut(model_id) { + Some(existing) => { + existing.reserved_bytes = bytes; + existing.last_access = Instant::now(); + } + None => { + allocs.insert(model_id.to_string(), Allocation::new(bytes)); + } + } + true } - /// 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()); + /// Unconditionally record a reservation, ignoring the budget. Test-support + /// (used by other modules' tests, hence `pub(crate)`): production paths + /// reserve atomically via [`Self::try_reserve`]. + #[cfg(test)] + pub(crate) fn allocate(&self, model_id: &str, bytes: u64) { + 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)); + } + } + } + + /// Release a model's reservation entirely. + pub(crate) fn deallocate(&self, model_id: &str) { + self.lock().remove(model_id); + } + + /// 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(crate) 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(crate) 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(crate) 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(crate) 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(crate) 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). - pub fn lru_candidate(&self, protected: &[String]) -> Option { - let allocs = self.allocations.lock().unwrap_or_else(|e| e.into_inner()); - allocs + /// Models in `protected`, and any model holding an active lease, are never + /// returned. + pub(crate) fn lru_candidate(&self, protected: &[String]) -> Option { + 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(crate) 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(crate) 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 +253,130 @@ impl MemoryManager { mod tests { use super::*; + const MB: u64 = 1024 * 1024; + #[test] fn allocate_and_track() { - let mm = MemoryManager::new(Some(100)); // 100 MB - assert_eq!(mm.budget_bytes(), 100 * 1024 * 1024); + let mm = MemoryManager::new(Some(100)); + assert_eq!(mm.budget_bytes(), 100 * MB); assert_eq!(mm.used_bytes(), 0); - mm.allocate("model-a", 10 * 1024 * 1024); - assert_eq!(mm.used_bytes(), 10 * 1024 * 1024); + mm.allocate("model-a", 10 * MB); + assert_eq!(mm.used_bytes(), 10 * MB); - mm.allocate("model-b", 20 * 1024 * 1024); - assert_eq!(mm.used_bytes(), 30 * 1024 * 1024); + mm.allocate("model-b", 20 * MB); + assert_eq!(mm.used_bytes(), 30 * MB); mm.deallocate("model-a"); - assert_eq!(mm.used_bytes(), 20 * 1024 * 1024); + assert_eq!(mm.used_bytes(), 20 * MB); + } + + #[test] + fn try_reserve_respects_budget() { + let mm = MemoryManager::new(Some(100)); + assert!(mm.try_reserve("a", 60 * MB)); + // 60 used, only 40 left → a 50 MB reservation must fail. + assert!(!mm.try_reserve("b", 50 * MB)); + assert!(mm.try_reserve("b", 40 * MB)); + assert_eq!(mm.used_bytes(), 100 * MB); + } + + #[test] + fn try_reserve_same_model_accounts_delta_not_sum() { + let mm = MemoryManager::new(Some(100)); + assert!(mm.try_reserve("a", 60 * MB)); + // Re-reserving the same model at 80 MB replaces, not adds → fits. + assert!(mm.try_reserve("a", 80 * MB)); + assert_eq!(mm.used_bytes(), 80 * MB); + } + + #[test] + fn reconcile_raises_reservation_to_observed() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("a", 10 * MB); + mm.reconcile("a", 25 * MB); + assert_eq!(mm.used_bytes(), 25 * MB); + // A lower observed value does not shrink the conservative reservation. + mm.reconcile("a", 5 * MB); + assert_eq!(mm.used_bytes(), 25 * MB); } #[test] fn can_fit_check() { let mm = MemoryManager::new(Some(100)); - mm.allocate("m1", 90 * 1024 * 1024); - assert!(!mm.can_fit(20 * 1024 * 1024)); - assert!(mm.can_fit(10 * 1024 * 1024)); + mm.allocate("m1", 90 * MB); + assert!(!mm.can_fit(20 * MB)); + assert!(mm.can_fit(10 * MB)); } #[test] fn lru_eviction_order() { let mm = MemoryManager::new(Some(100)); - mm.allocate("old", 10 * 1024 * 1024); - // tiny delay so timestamps differ - std::thread::sleep(std::time::Duration::from_millis(10)); - mm.allocate("new", 10 * 1024 * 1024); - - let candidate = mm.lru_candidate(&[]).unwrap(); - assert_eq!(candidate, "old"); + mm.allocate("old", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("new", 10 * MB); + assert_eq!(mm.lru_candidate(&[]).unwrap(), "old"); } #[test] fn lru_respects_protected() { let mm = MemoryManager::new(Some(100)); - mm.allocate("protected", 10 * 1024 * 1024); - std::thread::sleep(std::time::Duration::from_millis(10)); - mm.allocate("evictable", 10 * 1024 * 1024); + mm.allocate("protected", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("evictable", 10 * MB); + assert_eq!( + mm.lru_candidate(&["protected".into()]).unwrap(), + "evictable" + ); + } - let candidate = mm.lru_candidate(&["protected".into()]).unwrap(); - assert_eq!(candidate, "evictable"); + #[test] + fn leased_models_are_never_lru_candidates() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("leased", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("free", 10 * MB); + mm.lease("leased"); + // "leased" is older but is protected by its lease. + assert_eq!(mm.lru_candidate(&[]).unwrap(), "free"); + mm.release_lease("leased"); + assert_eq!(mm.lease_count("leased"), 0); } #[test] fn touch_updates_access_time() { let mm = MemoryManager::new(Some(100)); - mm.allocate("a", 10 * 1024 * 1024); - std::thread::sleep(std::time::Duration::from_millis(10)); - mm.allocate("b", 10 * 1024 * 1024); - - // 'a' is older, but touch it to make it newer - std::thread::sleep(std::time::Duration::from_millis(10)); + mm.allocate("a", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); + mm.allocate("b", 10 * MB); + std::thread::sleep(Duration::from_millis(10)); mm.touch("a"); + assert_eq!(mm.lru_candidate(&[]).unwrap(), "b"); + } + + #[test] + fn idle_candidates_reports_only_stale_unleased() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("stale", 10 * MB); + mm.allocate("leased_stale", 10 * MB); + mm.lease("leased_stale"); + std::thread::sleep(Duration::from_millis(20)); + mm.allocate("fresh", 10 * MB); - let candidate = mm.lru_candidate(&[]).unwrap(); - assert_eq!(candidate, "b"); + let idle = mm.idle_candidates(Duration::from_millis(10), &[]); + assert_eq!(idle, vec!["stale".to_string()]); + } + + #[test] + fn snapshot_is_sorted_and_complete() { + let mm = MemoryManager::new(Some(100)); + mm.allocate("b", 20 * MB); + mm.allocate("a", 10 * MB); + mm.reconcile("a", 15 * MB); + let snap = mm.snapshot(); + assert_eq!(snap.len(), 2); + assert_eq!(snap[0].model_id, "a"); + assert_eq!(snap[0].observed_bytes, Some(15 * MB)); + assert_eq!(snap[1].model_id, "b"); } } diff --git a/mofa-engine-core/src/metrics.rs b/mofa-engine-core/src/metrics.rs new file mode 100644 index 0000000..5a28a6d --- /dev/null +++ b/mofa-engine-core/src/metrics.rs @@ -0,0 +1,400 @@ +//! 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}; + +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, + /// 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. +#[derive(Default)] +pub(crate) 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, + /// Token/cost totals keyed by provider (bounded by the fixed provider set). + per_provider: DashMap, +} + +impl EngineMetrics { + /// Record the outcome of one logical inference request. + pub(crate) 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(crate) 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); + } + + /// 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(crate) 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); + } + 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 + /// counts, preflight, and per-provider routability) sampled at scrape time. + pub(crate) 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", + Self::escape_label_value(provider), + u8::from(*up) + )); + } + + // 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, locality, and direction.\n\ + # TYPE mofa_tokens_total counter\n", + ); + for entry in self.per_provider.iter() { + 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", + entry.prompt_tokens.load(Ordering::Relaxed) + )); + out.push_str(&format!( + "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 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=\"{}\",locality=\"{locality}\"}} {usd}\n", + Self::escape_label_value(entry.key()) + )); + } + + out + } +} + +impl EngineMetrics { + /// Escape a Prometheus label value per the exposition format: backslash, + /// 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(['\r', '\n'], "\\n") + } +} + +/// Point-in-time gauge values sampled from the engine at scrape time. +#[derive(Debug, Default, Clone)] +pub(crate) 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 per_provider_token_and_cost_render() { + let m = EngineMetrics::default(); + 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\",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] + fn provider_label_values_are_escaped() { + let m = EngineMetrics::default(); + let gauges = MetricsGauges { + provider_up: vec![("we\"ird\\name".into(), true)], + ..Default::default() + }; + let text = m.render_prometheus(&gauges); + // The quote and backslash must be escaped so the exposition stays valid. + assert!(text.contains(r#"mofa_provider_up{provider="we\"ird\\name"} 1"#)); + } + + #[test] + fn latency_buckets_are_cumulative() { + let m = EngineMetrics::default(); + m.record_request(true, 10, false); // le=50 + m.record_request(true, 80, false); // le=100 + let text = m.render_prometheus(&MetricsGauges::default()); + // le=50 covers 1, le=100 covers both. + assert!(text.contains("mofa_request_duration_ms_bucket{le=\"50\"} 1")); + assert!(text.contains("mofa_request_duration_ms_bucket{le=\"100\"} 2")); + } +} diff --git a/mofa-engine-core/src/preflight.rs b/mofa-engine-core/src/preflight.rs index 4a4a4c3..62d6f67 100644 --- a/mofa-engine-core/src/preflight.rs +++ b/mofa-engine-core/src/preflight.rs @@ -1,103 +1,214 @@ -//! Preflight prediction via Markov chain. +//! Preflight prediction via per-scope Markov chains. //! -//! Tracks transitions between capabilities (e.g. Chat → Tts) and -//! predicts the next capability so the engine can pre-warm models. +//! Tracks transitions between capabilities (e.g. `Chat → Tts`) so the engine can +//! pre-warm the next model before it is requested. History is keyed by an +//! application/session **scope** so one application's behavior never pollutes +//! another's predictions, with a global chain as a fallback when a scope has not +//! yet accumulated enough samples. use mofa_kernel::Capability; +use serde::{Deserialize, Serialize}; use std::collections::HashMap; use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Instant; /// Decay factor applied every `DECAY_INTERVAL` transitions. const DECAY_FACTOR: f64 = 0.95; /// Number of transitions between decay applications. const DECAY_INTERVAL: u64 = 100; -/// Minimum observations before making a prediction. -const MIN_OBSERVATIONS: u64 = 3; +/// Maximum number of distinct (non-global) scopes retained. Scope keys come from +/// caller-supplied app/session IDs, so this bounds memory against an unbounded +/// stream of unique identifiers by evicting the least-recently-used scope. +const MAX_SCOPES: usize = 2048; +/// Scope key under which cross-application history is aggregated. +pub(crate) const GLOBAL_SCOPE: &str = "__global__"; -/// Markov chain transition tracker with decay. -pub struct PreflightPredictor { - inner: Mutex, +/// Per-scope Markov chain transition tracker with decay. +pub(crate) struct PreflightPredictor { + inner: Mutex>, } -struct PreflightInner { - /// Transition counts: from_capability → (to_capability → count) +struct ScopeChain { + /// Transition counts: from_capability → (to_capability → weight). transitions: HashMap>, - /// Total transitions observed + /// Total transitions observed in this scope. total_transitions: u64, - /// Last capability observed + /// Last capability observed in this scope. last_capability: Option, + /// Last time this scope was recorded into or predicted from (for LRU). + last_access: Instant, +} + +impl ScopeChain { + fn new() -> Self { + Self { + transitions: HashMap::new(), + total_transitions: 0, + last_capability: None, + last_access: Instant::now(), + } + } + + /// Observe a capability used within this scope, recording the `prev → this` + /// transition against the scope's own running order. + fn observe(&mut self, capability: Capability) { + self.last_access = Instant::now(); + if let Some(prev) = self.last_capability { + self.record_transition(prev, capability); + } + self.last_capability = Some(capability); + } + + /// Add a single `from → to` transition to this chain's counts, applying + /// decay on the interval. Unlike [`observe`](Self::observe) it does not + /// touch `last_capability`, so it can mirror a transition that already + /// happened elsewhere (e.g. into the global aggregate) without inventing a + /// spurious edge from whatever this chain saw last. + fn record_transition(&mut self, from: Capability, to: Capability) { + self.last_access = Instant::now(); + *self + .transitions + .entry(from) + .or_default() + .entry(to) + .or_insert(0.0) += 1.0; + self.total_transitions += 1; + + if self.total_transitions.is_multiple_of(DECAY_INTERVAL) { + for counts in self.transitions.values_mut() { + for count in counts.values_mut() { + *count *= DECAY_FACTOR; + } + } + } + } + + fn predict(&self, current: Capability, min_samples: u64) -> Option { + let counts = self.transitions.get(¤t)?; + let total: f64 = counts.values().sum(); + // Compare in f64: `total` may be fractional after decay, and a lossy + // cast to u64 would be an unnecessary truncation. + if total < min_samples as f64 { + return None; + } + counts + .iter() + .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(cap, count)| Prediction { + capability: *cap, + confidence: count / total, + }) + } } /// Prediction result. -#[derive(Debug, Clone)] -pub struct Prediction { - /// Predicted next capability +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct Prediction { + /// Predicted next capability. pub capability: Capability, - /// Confidence score (0.0 - 1.0) + /// Confidence score (0.0 – 1.0). pub confidence: f64, } impl PreflightPredictor { /// Create a new predictor. - pub fn new() -> Self { + pub(crate) fn new() -> Self { Self { - inner: Mutex::new(PreflightInner { - transitions: HashMap::new(), - total_transitions: 0, - last_capability: None, - }), - } - } - - /// Record that a capability was just used. - /// This updates the transition counts from the previous capability. - pub fn record(&self, capability: Capability) { - let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); - - if let Some(prev) = inner.last_capability { - let entry = inner - .transitions - .entry(prev) - .or_default() - .entry(capability) - .or_insert(0.0); - *entry += 1.0; - - inner.total_transitions += 1; - - // Apply decay periodically - if inner.total_transitions.is_multiple_of(DECAY_INTERVAL) { - for counts in inner.transitions.values_mut() { - for count in counts.values_mut() { - *count *= DECAY_FACTOR; - } - } - } + inner: Mutex::new(HashMap::new()), } + } - inner.last_capability = Some(capability); + fn lock(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.inner.lock().unwrap_or_else(|e| e.into_inner()) } - /// Predict the most likely next capability given the current one. + /// Record that a capability was used within `scope`. /// - /// Returns `None` if there are not enough observations. - pub fn predict(&self, current: Capability) -> Option { - let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + /// The observation updates the scope's own chain, and the *same* per-scope + /// transition (if any) is mirrored into the global aggregate. The global + /// chain therefore only ever accumulates transitions that genuinely occurred + /// within some scope — an edge that spanned two different scopes (because + /// their requests interleaved) is never invented. + pub(crate) fn record(&self, scope: &str, capability: Capability) { + let mut scopes = self.lock(); + Self::evict_if_full(&mut scopes, scope); - let counts = inner.transitions.get(¤t)?; - let total: f64 = counts.values().sum(); + // Observe into the scope's own chain, capturing the transition edge it + // just formed (its previous capability, if any). + let edge_from = { + let chain = scopes + .entry(scope.to_string()) + .or_insert_with(ScopeChain::new); + let prev = chain.last_capability; + chain.observe(capability); + prev + }; - if (total as u64) < MIN_OBSERVATIONS { - return None; + // Mirror only that real edge into the global aggregate. + if scope != GLOBAL_SCOPE + && let Some(from) = edge_from + { + scopes + .entry(GLOBAL_SCOPE.to_string()) + .or_insert_with(ScopeChain::new) + .record_transition(from, capability); } + } - counts + /// Evict the least-recently-used non-global scope when at capacity and the + /// incoming scope is new. The global scope is never evicted. + fn evict_if_full(scopes: &mut HashMap, incoming: &str) { + if scopes.len() < MAX_SCOPES || scopes.contains_key(incoming) { + return; + } + if let Some(victim) = scopes .iter() - .max_by(|a, b| a.1.partial_cmp(b.1).unwrap_or(std::cmp::Ordering::Equal)) - .map(|(cap, count)| Prediction { - capability: *cap, - confidence: count / total, - }) + .filter(|(k, _)| k.as_str() != GLOBAL_SCOPE) + .min_by_key(|(_, c)| c.last_access) + .map(|(k, _)| k.clone()) + { + scopes.remove(&victim); + } + } + + /// Predict the next capability for `scope` given `current`. + /// + /// Tries the scope's own chain first; if it lacks enough samples or + /// confidence, falls back to the global chain. Returns `None` when neither + /// clears `min_samples` and `confidence_threshold`. + pub(crate) fn predict( + &self, + scope: &str, + current: Capability, + min_samples: u64, + confidence_threshold: f64, + ) -> Option { + let mut scopes = self.lock(); + // Touch the scope so an actively-predicted-from scope is not evicted. + if let Some(chain) = scopes.get_mut(scope) { + chain.last_access = Instant::now(); + } + let scoped = scopes + .get(scope) + .and_then(|chain| chain.predict(current, min_samples)) + .filter(|p| p.confidence >= confidence_threshold); + if scoped.is_some() { + return scoped; + } + if scope == GLOBAL_SCOPE { + return None; + } + scopes + .get(GLOBAL_SCOPE) + .and_then(|chain| chain.predict(current, min_samples)) + .filter(|p| p.confidence >= confidence_threshold) + } + + /// Number of scopes currently retained (test/diagnostic use). + #[cfg(test)] + fn scope_count(&self) -> usize { + self.lock().len() } } @@ -107,54 +218,220 @@ impl Default for PreflightPredictor { } } +/// Live counters describing Preflight effectiveness. +#[derive(Debug, Default)] +pub(crate) struct PreflightMetrics { + warms_started: AtomicU64, + warms_completed: AtomicU64, + warms_failed: AtomicU64, + warms_skipped: AtomicU64, + predictions: AtomicU64, + hits: AtomicU64, + misses: AtomicU64, +} + +/// Serializable snapshot of [`PreflightMetrics`]. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct PreflightStats { + /// Warm tasks spawned. + pub warms_started: u64, + /// Warm tasks that completed successfully. + pub warms_completed: u64, + /// Warm tasks that failed. + pub warms_failed: u64, + /// Warm requests skipped (deduplicated, disabled, or memory-unsafe). + pub warms_skipped: u64, + /// History/hint predictions produced. + pub predictions: u64, + /// Predictions later confirmed by the next request. + pub hits: u64, + /// Predictions the next request did not match. + pub misses: u64, +} + +impl PreflightMetrics { + /// A warm task was spawned. + pub(crate) fn warm_started(&self) { + self.warms_started.fetch_add(1, Ordering::Relaxed); + } + /// A warm task completed successfully. + pub(crate) fn warm_completed(&self) { + self.warms_completed.fetch_add(1, Ordering::Relaxed); + } + /// A warm task failed. + pub(crate) fn warm_failed(&self) { + self.warms_failed.fetch_add(1, Ordering::Relaxed); + } + /// A warm was skipped (deduplicated, disabled, or unsafe). + pub(crate) fn warm_skipped(&self) { + self.warms_skipped.fetch_add(1, Ordering::Relaxed); + } + /// A prediction was produced. + pub(crate) fn prediction(&self) { + self.predictions.fetch_add(1, Ordering::Relaxed); + } + /// A prior prediction matched the next request. + pub(crate) fn hit(&self) { + self.hits.fetch_add(1, Ordering::Relaxed); + } + /// A prior prediction did not match the next request. + pub(crate) fn miss(&self) { + self.misses.fetch_add(1, Ordering::Relaxed); + } + + /// Read all counters into a `PreflightStats`. + /// + /// Each counter is read independently with relaxed ordering, so under heavy + /// concurrent updates the fields may reflect slightly different instants; + /// they are monotonic and eventually consistent, which is what the metrics + /// consumers need. + pub(crate) fn snapshot(&self) -> PreflightStats { + PreflightStats { + warms_started: self.warms_started.load(Ordering::Relaxed), + warms_completed: self.warms_completed.load(Ordering::Relaxed), + warms_failed: self.warms_failed.load(Ordering::Relaxed), + warms_skipped: self.warms_skipped.load(Ordering::Relaxed), + predictions: self.predictions.load(Ordering::Relaxed), + hits: self.hits.load(Ordering::Relaxed), + misses: self.misses.load(Ordering::Relaxed), + } + } +} + #[cfg(test)] mod tests { use super::*; + const SAMPLES: u64 = 3; + const THRESHOLD: f64 = 0.6; + #[test] fn no_prediction_without_observations() { let p = PreflightPredictor::new(); - assert!(p.predict(Capability::Chat).is_none()); + assert!( + p.predict("app", Capability::Chat, SAMPLES, THRESHOLD) + .is_none() + ); } #[test] fn no_prediction_below_threshold() { let p = PreflightPredictor::new(); - p.record(Capability::Chat); - p.record(Capability::Tts); - // Only 1 transition Chat→Tts, need 3 - assert!(p.predict(Capability::Chat).is_none()); + p.record("app", Capability::Chat); + p.record("app", Capability::Tts); + // Only one Chat→Tts transition; need three samples. + assert!( + p.predict("app", Capability::Chat, SAMPLES, THRESHOLD) + .is_none() + ); } #[test] fn predicts_after_enough_observations() { let p = PreflightPredictor::new(); - - // Record Chat → Tts four times for _ in 0..4 { - p.record(Capability::Chat); - p.record(Capability::Tts); + p.record("app", Capability::Chat); + p.record("app", Capability::Tts); } - - let pred = p.predict(Capability::Chat).unwrap(); + let pred = p + .predict("app", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap(); assert_eq!(pred.capability, Capability::Tts); assert!(pred.confidence > 0.9); } #[test] - fn handles_multiple_successors() { + fn scopes_do_not_pollute_each_other() { let p = PreflightPredictor::new(); + // App A always does Chat → Tts. + for _ in 0..5 { + p.record("app-a", Capability::Chat); + p.record("app-a", Capability::Tts); + } + // App B always does Chat → Asr. + for _ in 0..5 { + p.record("app-b", Capability::Chat); + p.record("app-b", Capability::Asr); + } + assert_eq!( + p.predict("app-a", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap() + .capability, + Capability::Tts + ); + assert_eq!( + p.predict("app-b", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap() + .capability, + Capability::Asr + ); + } - // Chat → Tts (3x), Chat → Asr (1x) - for _ in 0..3 { - p.record(Capability::Chat); - p.record(Capability::Tts); + #[test] + fn unknown_scope_falls_back_to_global() { + let p = PreflightPredictor::new(); + for _ in 0..5 { + p.record("app-a", Capability::Chat); + p.record("app-a", Capability::Tts); } - p.record(Capability::Chat); - p.record(Capability::Asr); + // A brand-new scope has no history, but global aggregates app-a's. + let pred = p + .predict("fresh-scope", Capability::Chat, SAMPLES, THRESHOLD) + .unwrap(); + assert_eq!(pred.capability, Capability::Tts); + } - let pred = p.predict(Capability::Chat).unwrap(); + #[test] + fn global_chain_ignores_cross_scope_transitions() { + let p = PreflightPredictor::new(); + // Two scopes run concurrently with interleaved requests. App A always + // does Chat → Tts; app B always does Asr → Vlm. If the global chain + // observed the interleaved stream it would learn phantom edges like + // Chat → Asr (A's Chat followed by B's Asr), which never happened in any + // real sequence. + for _ in 0..6 { + p.record("app-a", Capability::Chat); + p.record("app-b", Capability::Asr); + p.record("app-a", Capability::Tts); + p.record("app-b", Capability::Vlm); + } + // A fresh scope falls back to the global aggregate. It must predict the + // real edge Chat → Tts, not the phantom Chat → Asr the old code learned. + let pred = p + .predict("fresh", Capability::Chat, SAMPLES, THRESHOLD) + .expect("global should predict the real transition"); assert_eq!(pred.capability, Capability::Tts); - assert!(pred.confidence > 0.7); + assert!(pred.confidence > 0.9); + } + + #[test] + fn scope_count_is_bounded() { + let p = PreflightPredictor::new(); + // Record far more distinct scopes than the cap allows. + for i in 0..(MAX_SCOPES + 200) { + let scope = format!("session-{i}"); + p.record(&scope, Capability::Chat); + p.record(&scope, Capability::Tts); + } + // Bounded by the cap plus the always-retained global scope. + assert!( + p.scope_count() <= MAX_SCOPES + 1, + "scope map grew to {} entries", + p.scope_count() + ); + } + + #[test] + fn metrics_snapshot_counts() { + let m = PreflightMetrics::default(); + m.warm_started(); + m.warm_started(); + m.warm_completed(); + m.hit(); + let snap = m.snapshot(); + assert_eq!(snap.warms_started, 2); + assert_eq!(snap.warms_completed, 1); + assert_eq!(snap.hits, 1); + assert_eq!(snap.misses, 0); } } diff --git a/mofa-engine-core/src/quality_gate.rs b/mofa-engine-core/src/quality_gate.rs new file mode 100644 index 0000000..3350868 --- /dev/null +++ b/mofa-engine-core/src/quality_gate.rs @@ -0,0 +1,517 @@ +//! 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; + } + QualityGate::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; + } + QualityGate::slideshow_risk_from_frames(&output.stdout, FRAME_PIXELS) + } +} + +/// 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 + }) + }); + + 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 = QualityGate::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 = QualityGate::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!( + QualityGate::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!( + QualityGate::slideshow_risk_from_frames(&raw, FRAME_PIXELS), + Some(0.0) + ); + } + + #[test] + fn slideshow_risk_needs_two_frames() { + 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] + 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 = QualityGate::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-core/src/router.rs b/mofa-engine-core/src/router.rs index 7a8e0c5..1c35178 100644 --- a/mofa-engine-core/src/router.rs +++ b/mofa-engine-core/src/router.rs @@ -4,12 +4,18 @@ use mofa_kernel::{ BackendHealth, Capability, CostTier, InferenceRequest, ModelCard, ModelResidency, ModelStatus, - ProviderKind, + Prefer, ProviderKind, }; +/// 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)] -pub struct RoutingProvider { +pub(crate) struct RoutingProvider { /// Provider name. pub name: String, /// Provider kind. @@ -24,7 +30,7 @@ pub struct RoutingProvider { /// Explainable route decision. #[derive(Debug, Clone)] -pub struct RouteDecision<'a> { +pub(crate) struct RouteDecision<'a> { /// Selected model. pub model: &'a ModelCard, /// Composite score. @@ -34,17 +40,25 @@ pub struct RouteDecision<'a> { } /// Selects the best model for a given request from a pool of candidates. -pub struct Router; +pub(crate) struct Router; impl Router { - /// Score and rank models, returning the best match and reason. - pub fn route<'a>( + /// 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(crate) fn route_ranked<'a>( models: &'a [ModelCard], request: &InferenceRequest, providers: &[RoutingProvider], - ) -> Option> { + 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()) { @@ -65,49 +79,78 @@ impl Router { continue; } + // 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; + } + if model.availability == mofa_kernel::ModelAvailability::Unavailable { 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; } - let score = Self::score(model, desired_cap, provider); + // 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); + // 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; } 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. - 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 { + u8::from(matches!( + model.residency, + ModelResidency::Loaded | ModelResidency::Remote + )) } fn matches_model_name(model: &ModelCard, target: Option<&str>) -> bool { @@ -172,11 +215,7 @@ impl Router { } fn locality_score(kind: ProviderKind) -> i64 { - match kind { - ProviderKind::Ollama => 100, - ProviderKind::OpenAiCompatible => 0, - _ => 0, - } + if kind.is_local() { 100 } else { 0 } } fn cost_score(tier: CostTier) -> i64 { @@ -198,7 +237,9 @@ impl Router { } fn priority_score(priority: u8) -> i64 { - 100_i64.saturating_sub(priority as i64) + // Lower configured priority is preferred; floor at 0 so a priority above + // 100 cannot contribute a negative term to the score. + (100_i64 - priority as i64).max(0) } fn health_score(health: BackendHealth) -> i64 { @@ -216,18 +257,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={:?}", @@ -245,7 +274,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, @@ -282,14 +314,61 @@ mod tests { messages: vec![Message { role: "user".into(), content: "hi".into(), + ..Default::default() }], input_file: None, params: serde_json::Value::Null, hint_next: None, request_id: "test".into(), + ..Default::default() } } + #[test] + #[ignore = "timing-sensitive perf guard; run explicitly with `--ignored` or as a benchmark"] + fn routing_decision_meets_latency_target() { + // RFC budget: a scheduling decision should take < 1ms (excluding load). + // Wall-clock assertions are sensitive to CI host load, so this is a + // perf guard run on demand rather than a gating unit test. + // Build a realistic pool: three providers, several models each. + let providers = vec![ + provider("ollama", ProviderKind::Ollama, 1), + provider("openai", ProviderKind::OpenAiCompatible, 10), + provider("deepseek", ProviderKind::OpenAiCompatible, 8), + ]; + let mut models = Vec::new(); + for (prov, count) in [("ollama", 8), ("openai", 6), ("deepseek", 6)] { + for i in 0..count { + models.push(make_model( + &format!("m{i}"), + prov, + Capability::Chat, + if i % 2 == 0 { + ModelResidency::Unloaded + } else { + ModelResidency::Loaded + }, + CostTier::Low, + )); + } + } + let req = request(Capability::Chat, None); + let budget = Some(16 * 1024 * 1024 * 1024); + + let iterations = 2000; + let start = std::time::Instant::now(); + for _ in 0..iterations { + let plan = Router::route_ranked(&models, &req, &providers, budget); + std::hint::black_box(&plan); + } + let avg = start.elapsed() / iterations; + assert!( + avg < std::time::Duration::from_millis(1), + "routing averaged {avg:?}/decision over {} models, exceeding the 1ms target", + models.len() + ); + } + #[test] fn prefers_loaded_over_unloaded() { let models = vec![ @@ -309,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"); } @@ -337,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"); } @@ -362,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"); } @@ -387,12 +475,103 @@ 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"); } + #[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( @@ -404,6 +583,209 @@ 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] + 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.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); + 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.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_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", + "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::Auto; + let ranked = Router::route_ranked(&models, &req, &providers, None); + assert_eq!(ranked.len(), 2); + assert_eq!(ranked[0].model.name, "local"); + } + + #[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-core/src/subscription.rs b/mofa-engine-core/src/subscription.rs new file mode 100644 index 0000000..60d32a8 --- /dev/null +++ b/mofa-engine-core/src/subscription.rs @@ -0,0 +1,198 @@ +//! Capability subscriptions. +//! +//! An application declares at startup which capabilities it will use. The engine +//! keeps subscribed capabilities' models warm and protects them from eviction +//! for the lifetime of the subscription (until it is explicitly removed or its +//! TTL expires). Subscriptions are owned by an `app_id`/`session_id` so they can +//! be reasoned about and cleaned up per application. + +use mofa_kernel::Capability; +use serde::{Deserialize, Serialize}; +use std::collections::HashSet; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::{Duration, Instant}; + +/// Registry of active capability subscriptions. +pub(crate) struct SubscriptionRegistry { + inner: Mutex>, + next_id: AtomicU64, +} + +#[derive(Clone)] +struct Subscription { + id: u64, + app_id: Option, + session_id: Option, + capabilities: Vec, + expires_at: Option, +} + +impl Subscription { + fn is_expired(&self, now: Instant) -> bool { + self.expires_at.is_some_and(|e| now >= e) + } +} + +/// Public view of a subscription. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct SubscriptionInfo { + /// Unique subscription id. + pub id: u64, + /// Owning application, if provided. + pub app_id: Option, + /// Owning session, if provided. + pub session_id: Option, + /// Subscribed capabilities. + pub capabilities: Vec, + /// Seconds until expiry, or `None` if the subscription does not expire. + pub expires_in_secs: Option, +} + +impl SubscriptionRegistry { + /// Create an empty registry. + pub(crate) 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(crate) 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(crate) fn unsubscribe(&self, id: u64) -> bool { + let mut subs = self.lock(); + let before = subs.len(); + subs.retain(|s| s.id != id); + subs.len() != before + } + + /// The set of capabilities currently kept warm by any live subscription. + /// + /// Prunes and reads under a single lock, so the result is a consistent + /// snapshot and expired entries are dropped exactly once. + pub(crate) fn active_capabilities(&self) -> HashSet { + let now = Instant::now(); + let mut subs = self.lock(); + subs.retain(|s| !s.is_expired(now)); + subs.iter() + .flat_map(|s| s.capabilities.iter().copied()) + .collect() + } + + /// Whether `capability` is currently subscribed. Test-support (this module only). + #[cfg(test)] + fn is_subscribed(&self, capability: Capability) -> bool { + self.active_capabilities().contains(&capability) + } + + /// List all live subscriptions. + pub(crate) fn list(&self) -> Vec { + let now = Instant::now(); + let mut subs = self.lock(); + subs.retain(|s| !s.is_expired(now)); + subs.iter() + .map(|s| SubscriptionInfo { + id: s.id, + app_id: s.app_id.clone(), + session_id: s.session_id.clone(), + capabilities: s.capabilities.clone(), + expires_in_secs: s + .expires_at + .map(|e| e.saturating_duration_since(now).as_secs()), + }) + .collect() + } +} + +impl Default for SubscriptionRegistry { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn subscribe_and_query() { + let reg = SubscriptionRegistry::new(); + let id = reg.subscribe( + Some("mofa-fm".into()), + None, + vec![Capability::Chat, Capability::Tts], + None, + ); + assert!(reg.is_subscribed(Capability::Chat)); + assert!(reg.is_subscribed(Capability::Tts)); + assert!(!reg.is_subscribed(Capability::Asr)); + assert_eq!(reg.list().len(), 1); + + assert!(reg.unsubscribe(id)); + assert!(!reg.is_subscribed(Capability::Chat)); + assert!(reg.list().is_empty()); + } + + #[test] + fn unsubscribe_unknown_is_false() { + let reg = SubscriptionRegistry::new(); + assert!(!reg.unsubscribe(999)); + } + + #[test] + fn expired_subscriptions_are_pruned() { + let reg = SubscriptionRegistry::new(); + reg.subscribe( + None, + None, + vec![Capability::Tts], + Some(Duration::from_millis(10)), + ); + assert!(reg.is_subscribed(Capability::Tts)); + std::thread::sleep(Duration::from_millis(20)); + assert!(!reg.is_subscribed(Capability::Tts)); + assert!(reg.list().is_empty()); + } + + #[test] + fn list_reports_remaining_ttl() { + let reg = SubscriptionRegistry::new(); + reg.subscribe( + None, + None, + vec![Capability::Chat], + Some(Duration::from_secs(3600)), + ); + let info = reg.list(); + assert_eq!(info.len(), 1); + let ttl = info[0].expires_in_secs.unwrap(); + assert!(ttl > 3590 && ttl <= 3600); + } +} diff --git a/mofa-engine-sdk/Cargo.toml b/mofa-engine-sdk/Cargo.toml index eb785c4..b4b6124 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,5 @@ thiserror = { workspace = true } tracing = { workspace = true } chrono = { workspace = true } tokio-stream = { workspace = true } +uuid = { workspace = true } +subtle = "2.6.1" diff --git a/mofa-engine-sdk/examples/README.md b/mofa-engine-sdk/examples/README.md new file mode 100644 index 0000000..feaa2d9 --- /dev/null +++ b/mofa-engine-sdk/examples/README.md @@ -0,0 +1,59 @@ +# 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. + +Start with `quickstart` — one command, one answer, and the routing/cost line that +is the whole pitch. The rest map to PRD product scenarios. + +| Example | Scenario | Engine capabilities exercised | +|---|---|---| +| `quickstart.rs` | — (5-minute entry) | boot + discover backends + one `Chat` call, printing **routed provider · local/cloud · latency · 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" | +| `video_gen.rs` | **S4** Text-to-Video | `VideoGen` at the API level via **Seedance** (Volcengine Ark task API: submit→poll→download) — one contract call returns a managed `mp4`; degrades to a typed error offline | +| `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 | +| `podcast.rs` | **S6** Podcast / Long-Audio | `Chat` rewrite → `Tts` narration, `hint_next` cross-capability warmup — a playable episode, offline | +| `provider_race.rs` | **S7** Multi-Vendor Access | same question across **every** discovered chat model, side-by-side **dual-track cost** (local `$0` vs metered cloud) + fault tolerance | +| `streaming_chat.rs` | — (true streaming) | real per-token `invoke_stream`, **time-to-first-token** the number pseudo-streaming can't move | +| `private_local.rs` | **S5** Privacy moat | `prefer=local` + `data_class=confidential` — **fail, don't leak**: pins on-device or errors, never silently cloud | + +## Running + +```bash +# Start here — offline, just needs a local Ollama with any chat model pulled. +cargo run -p mofa-engine-sdk --example quickstart + +# 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 + +# 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 + +# S6 — article → playable episode; needs a chat model + a TTS backend (both local). +cargo run -p mofa-engine-sdk --example podcast -- ./article.txt + +# S7 — race the same question across every discovered chat model (local + cloud). +cargo run -p mofa-engine-sdk --example provider_race + +# True per-token streaming, with time-to-first-token. Needs a chat backend. +cargo run -p mofa-engine-sdk --example streaming_chat + +# Privacy moat — pins confidential text on-device (fail, don't leak). +cargo run -p mofa-engine-sdk --example private_local +``` + +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..dca64a9 --- /dev/null +++ b/mofa-engine-sdk/examples/code_review.rs @@ -0,0 +1,123 @@ +//! 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, + }), + // 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() + }; + + 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/explainer_video.rs b/mofa-engine-sdk/examples/explainer_video.rs new file mode 100644 index 0000000..eb9f918 --- /dev/null +++ b/mofa-engine-sdk/examples/explainer_video.rs @@ -0,0 +1,149 @@ +//! 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 = 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"; + 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() + }; + ExplainerVideo::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() + }; + 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. + 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) => 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" + ), + } +} + +/// 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!( + " [{}] {}: {}", + 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 new file mode 100644 index 0000000..eebc54b --- /dev/null +++ b/mofa-engine-sdk/examples/meeting_brief.rs @@ -0,0 +1,139 @@ +//! 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), 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() + }; + let transcript = match MeetingBrief::run(&engine, asr, "ASR") { + Some(r) => r.text.unwrap_or_default(), + None => return, + }; + println!("transcript: {}\n", MeetingBrief::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 MeetingBrief::run(&engine, chat, "Chat (minutes)") { + Some(r) => r.text.unwrap_or_default(), + None => return, + }; + println!("minutes:\n{}\n", MeetingBrief::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) = MeetingBrief::run(&engine, tts, "TTS (brief)") { + println!("audio brief: {}", r.file.as_deref().unwrap_or("(no file)")); + } +} + +/// 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 + } + } + } + + /// 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/examples/podcast.rs b/mofa-engine-sdk/examples/podcast.rs new file mode 100644 index 0000000..6f07f67 --- /dev/null +++ b/mofa-engine-sdk/examples/podcast.rs @@ -0,0 +1,140 @@ +//! Scenario S6 — Podcast / Long-Audio Matrix. +//! +//! An article → a natural, colloquial **podcast episode you can play**. Two +//! stages, both local-first: rewrite the article into spoken-style script +//! (`Chat`), then synthesize narration (`Tts`). The rewrite carries +//! `hint_next="tts"`, so the engine warms the TTS model *while* the script is +//! being written — the second stage starts from a hot model instead of a cold +//! one (the cross-capability warmup that gateways proxying tokens don't have). +//! +//! This is the Rust port of `mofa-fm/article_to_podcast.py`, the repo's original +//! runnable demo. The whole chain runs offline on Ollama + a local TTS backend. +//! +//! Run (needs a chat model + a TTS backend; both can be local): +//! cargo run -p mofa-engine-sdk --example podcast -- ./article.txt +//! cargo run -p mofa-engine-sdk --example podcast # uses a built-in sample +//! +//! Each stage degrades gracefully: a missing backend prints the structured error +//! and stops rather than shipping a half-made episode. + +use mofa_engine_core::EngineConfig; +use mofa_engine_sdk::EmbeddedEngine; +use mofa_kernel::{Capability, InferenceRequest, InferenceResponse, Message}; + +/// A short built-in article so the demo runs with zero arguments. +const SAMPLE_ARTICLE: &str = "\ +Local-first AI is having a moment. As open models get small enough to run on a \ +laptop, developers are discovering they can keep sensitive data on-device, cut \ +cloud bills to zero for most requests, and still reach for a hosted model when a \ +task genuinely needs one. The catch has always been glue code: every provider \ +ships its own SDK, its own auth, its own error shapes. An orchestration layer \ +that unifies them — routing locally when it can and to the cloud when it must — \ +is what turns a pile of models into a product."; + +fn main() { + // Load the article from a file argument, or fall back to the built-in sample. + let article = match std::env::args().nth(1) { + Some(path) => match std::fs::read_to_string(&path) { + Ok(text) => text, + Err(e) => { + eprintln!("could not read {path}: {e}"); + return; + } + }, + None => SAMPLE_ARTICLE.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(); + + println!("== S6 Podcast — article → narrated episode ==\n"); + + // 1. Chat: rewrite the article into a spoken-style script. `hint_next="tts"` + // warms the synthesis model during this step so stage 2 has no cold start. + let rewrite = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![ + Message { + role: "system".into(), + content: "Rewrite the article as a natural, colloquial podcast \ + monologue under 200 words. No headings, just what the \ + host would say aloud." + .into(), + ..Default::default() + }, + Message { + role: "user".into(), + content: article, + ..Default::default() + }, + ], + hint_next: Some("tts".into()), + ..Default::default() + }; + let script = match Podcast::run(&engine, rewrite, "Chat (rewrite)") { + Some(r) => r.text.unwrap_or_default(), + None => return, + }; + println!("\nscript:\n{}\n", script.trim()); + + // 2. TTS: script → a playable audio file. Local Crane/Kokoro when available, + // cloud tts-1 as fallback under `auto`. + let narrate = InferenceRequest { + capability: Some(Capability::Tts), + messages: vec![Message { + role: "user".into(), + content: script, + ..Default::default() + }], + params: serde_json::json!({ "voice": "narrator", "format": "mp3" }), + ..Default::default() + }; + if let Some(r) = Podcast::run(&engine, narrate, "TTS (narration)") { + match r.file.as_deref() { + Some(path) => println!("\n🎧 episode: {path}"), + None => println!("\n(TTS returned no file path)"), + } + } +} + +/// Namespace for this demo's helpers. +struct Podcast; + +impl Podcast { + /// Invoke one stage, reporting where it routed (and cost) on success or the + /// structured error on failure. Returns `None` to stop the pipeline. + fn run( + engine: &EmbeddedEngine, + req: InferenceRequest, + stage: &str, + ) -> Option { + match engine.invoke(req) { + Ok(resp) => { + let cost = resp + .cost_usd + .map(|c| format!("${c:.6}")) + .unwrap_or_else(|| "$0.00 (local)".into()); + println!( + "[{stage}] served by {} in {}ms · {cost}", + 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 + } + } + } +} diff --git a/mofa-engine-sdk/examples/private_local.rs b/mofa-engine-sdk/examples/private_local.rs new file mode 100644 index 0000000..9b16ca3 --- /dev/null +++ b/mofa-engine-sdk/examples/private_local.rs @@ -0,0 +1,106 @@ +//! The privacy moat — `prefer=local` + `data_class=confidential`. +//! +//! MoFA's differentiator over any cloud gateway is a *hard* data-residency +//! boundary: a request marked `confidential` is pinned to local backends, and if +//! none can serve it the request **fails rather than falling back to the cloud**. +//! "Fail, not leak" is the guarantee a compliance team can actually sign off on. +//! +//! This demo sends the same prompt twice: +//! 1. `prefer=auto` — normal routing (may pick cloud if it scores best). +//! 2. `prefer=local` + — pinned on-device; the sensitive text never leaves +//! `data_class=confidential` the machine, even if that means failing. +//! +//! Run: +//! cargo run -p mofa-engine-sdk --example private_local +//! +//! To *see* the guarantee bite, run it with no local backend available (stop +//! Ollama): the confidential call returns a structured error while the `auto` +//! call may still succeed via cloud — exactly the point. + +use mofa_engine_core::EngineConfig; +use mofa_engine_sdk::EmbeddedEngine; +use mofa_kernel::{Capability, DataClass, InferenceRequest, Message, Prefer}; + +/// Stand-in for text a business would never want sent to a third-party API. +const SENSITIVE: &str = + "Summarize this internal note: Q3 layoffs affect the Shanghai team; do not disclose."; + +fn main() { + 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!("== privacy moat — fail, don't leak ==\n"); + + // 1. Default routing: the engine is free to choose cloud if it scores best. + let auto = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![Message { + role: "user".into(), + content: SENSITIVE.into(), + ..Default::default() + }], + prefer: Prefer::Auto, + ..Default::default() + }; + describe(&engine, auto, "prefer=auto (unconstrained)"); + + // 2. Confidential: pinned local. `data_class=confidential` alone pins it, and + // `prefer=local` states the intent explicitly. If no local model can serve, + // this returns an error with a `failed_chain` — it does not reach for cloud. + let confidential = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![Message { + role: "user".into(), + content: SENSITIVE.into(), + ..Default::default() + }], + prefer: Prefer::Local, + data_class: DataClass::Confidential, + ..Default::default() + }; + describe( + &engine, + confidential, + "prefer=local + confidential (pinned)", + ); +} + +/// Invoke and report where the request actually ran — local vs cloud is the whole +/// story here — or, on failure, that it declined to leave the device. +fn describe(engine: &EmbeddedEngine, req: InferenceRequest, label: &str) { + println!("── {label}"); + match engine.invoke(req) { + Ok(resp) => { + let cost = resp + .cost_usd + .map(|c| format!("${c:.6}")) + .unwrap_or_else(|| "$0.00 (local/free)".into()); + println!( + " served by {} · {}ms · {cost}", + resp.provider, resp.duration_ms + ); + if let Some(reason) = &resp.routing_reason { + println!(" routing: {reason}"); + } + println!( + " answer: {}\n", + resp.text.as_deref().unwrap_or("(no text)").trim() + ); + } + Err(e) => { + let info = e.info(); + println!(" declined — {:?}: {}", info.code, info.message); + for a in &info.failed_chain { + println!(" considered {}/{}: {}", a.provider, a.model, a.reason); + } + println!(" → the confidential text was NOT sent anywhere.\n"); + } + } +} diff --git a/mofa-engine-sdk/examples/provider_race.rs b/mofa-engine-sdk/examples/provider_race.rs new file mode 100644 index 0000000..440a68e --- /dev/null +++ b/mofa-engine-sdk/examples/provider_race.rs @@ -0,0 +1,125 @@ +//! Scenario S7 — Multi-Vendor Unified Access. +//! +//! The same question, sent to **every** chat model the engine discovered — local +//! and cloud — through one interface, one request type, one error shape. The +//! output is a side-by-side table: provider · latency · tokens · cost. This is +//! the infrastructure scenario made visible: adding a vendor is config, not code, +//! and the **dual-track cost** column shows local models landing at `$0` next to +//! metered cloud ones — the comparison a plain OpenAI-compatible proxy can't draw +//! because it never sees the local side. +//! +//! Run (more interesting with both a local Ollama and a cloud key configured): +//! cargo run -p mofa-engine-sdk --example provider_race +//! +//! Models that can't serve (unavailable, over budget) show up as a row with their +//! failure reason rather than aborting the race — fault tolerance on display. + +use std::time::Instant; + +use mofa_engine_core::EngineConfig; +use mofa_engine_sdk::EmbeddedEngine; +use mofa_kernel::{Capability, InferenceRequest, Message}; + +const QUESTION: &str = "Explain what a circuit breaker does, in two sentences."; + +fn main() { + let config = EngineConfig::load(None); + let engine = match EmbeddedEngine::new(config) { + Ok(e) => e, + Err(e) => { + eprintln!("engine init failed: {e}"); + return; + } + }; + engine.refresh(); + + // Every chat-capable model becomes a lane in the race. We pin each request to + // one model by id so the router targets exactly that backend instead of + // picking a winner — the whole point is to compare them. + let lanes: Vec<_> = engine + .capabilities() + .into_iter() + .filter(|c| c.capabilities.contains(&Capability::Chat)) + .collect(); + + println!("== S7 Multi-Vendor Race =="); + println!("question: {QUESTION}\n"); + if lanes.is_empty() { + println!("no chat backends discovered — start Ollama or set a cloud API key."); + return; + } + println!( + "{:<30} {:>8} {:>8} {:>12}", + "provider/model", "ms", "tokens", "cost" + ); + println!("{}", "-".repeat(64)); + + // Track the cheapest and fastest successful lanes for a closing summary. + let mut best_cost: Option<(String, f64)> = None; + let mut best_latency: Option<(String, u128)> = None; + + for card in &lanes { + let request = InferenceRequest { + capability: Some(Capability::Chat), + model: Some(card.id.clone()), + messages: vec![Message { + role: "user".into(), + content: QUESTION.into(), + ..Default::default() + }], + ..Default::default() + }; + + // Wall-clock the call ourselves too, so an error row still reports how long + // the attempt took before failing. + let started = Instant::now(); + match engine.invoke(request) { + Ok(resp) => { + let elapsed = started.elapsed().as_millis(); + let cost = resp.cost_usd.unwrap_or(0.0); + let cost_str = if cost > 0.0 { + format!("${cost:.6}") + } else { + "$0.00 local".into() + }; + println!( + "{:<30} {:>8} {:>8} {:>12}", + card.id, + resp.duration_ms, + resp.tokens_used + .map(|t| t.to_string()) + .unwrap_or_else(|| "-".into()), + cost_str, + ); + + // Fold this lane into the running winners. + match &best_cost { + Some((_, c)) if *c <= cost => {} + _ => best_cost = Some((card.id.clone(), cost)), + } + match &best_latency { + Some((_, ms)) if *ms <= elapsed => {} + _ => best_latency = Some((card.id.clone(), elapsed)), + } + } + Err(e) => { + let info = e.info(); + println!("{:<30} {:>8} {:>8} {:>12}", card.id, "-", "-", "failed"); + println!(" └─ {:?}: {}", info.code, info.message); + } + } + } + + println!(); + if let Some((id, cost)) = best_cost { + let label = if cost > 0.0 { + format!("${cost:.6}") + } else { + "$0.00 (local, free)".into() + }; + println!("cheapest: {id} — {label}"); + } + if let Some((id, ms)) = best_latency { + println!("fastest: {id} — {ms}ms"); + } +} diff --git a/mofa-engine-sdk/examples/quickstart.rs b/mofa-engine-sdk/examples/quickstart.rs new file mode 100644 index 0000000..88d62a2 --- /dev/null +++ b/mofa-engine-sdk/examples/quickstart.rs @@ -0,0 +1,94 @@ +//! Quickstart — the 5-minute "does it actually run?" demo. +//! +//! This is the shortest path from a fresh checkout to a real answer: boot the +//! embedded engine, show what backends it discovered, and run one chat request. +//! The point is not the answer — it's the line printed after it: **which +//! provider served the request, local or cloud, how long it took, and what it +//! cost**. That single line is the whole MoFA pitch (one interface, intelligent +//! routing, local-first, cost visible) made concrete. +//! +//! Run (offline: just needs a local Ollama with any chat model pulled): +//! ollama serve & # if not already running +//! ollama pull llama3.2 # or any chat model +//! cargo run -p mofa-engine-sdk --example quickstart +//! +//! With no backend at all 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() { + // Configuration comes from ./config.toml when present, otherwise environment + // auto-detection (cloud keys) + a local Ollama at 127.0.0.1:11434. + let config = EngineConfig::load(None); + let engine = match EmbeddedEngine::new(config) { + Ok(e) => e, + Err(e) => { + eprintln!("engine init failed: {e}"); + return; + } + }; + engine.refresh(); + + // Show the routing table the engine has to work with. On a laptop with only + // Ollama running this lists local models at cost tier "free"; add a cloud key + // and cloud models appear alongside them — same interface, both routable. + println!("== MoFA quickstart ==\n"); + let cards = engine.capabilities(); + if cards.is_empty() { + println!("no backends discovered — start Ollama or set a cloud API key, then retry.\n"); + } else { + println!("discovered {} model(s):", cards.len()); + for card in &cards { + println!( + " · {:<28} [{}] {:?}", + card.id, card.provider, card.cost_tier + ); + } + println!(); + } + + // One ordinary chat request. `prefer` is left at its default (`auto`), so the + // seven-dimensional router picks local-first and only reaches for cloud when a + // local model can't serve — nothing here is pinned by hand. + let request = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![Message { + role: "user".into(), + content: "In one sentence, what is MoFA Engine?".into(), + ..Default::default() + }], + ..Default::default() + }; + + match engine.invoke(request) { + Ok(resp) => { + println!( + "answer:\n {}\n", + resp.text.as_deref().unwrap_or("(no text)").trim() + ); + // The MoFA pitch, made concrete: where it ran and what it cost. + println!( + "routed to {}/{} · {}ms · tokens {:?} · cost {}", + resp.provider, + resp.model_used, + resp.duration_ms, + resp.tokens_used, + resp.cost_usd + .map(|c| format!("${c:.6}")) + .unwrap_or_else(|| "$0.00 (local/free)".into()), + ); + if let Some(reason) = &resp.routing_reason { + println!("why: {reason}"); + } + } + Err(e) => { + let info = e.info(); + eprintln!("error [{:?}]: {}", info.code, info.message); + for a in &info.failed_chain { + eprintln!(" tried {}/{}: {}", a.provider, a.model, a.reason); + } + } + } +} diff --git a/mofa-engine-sdk/examples/streaming_chat.rs b/mofa-engine-sdk/examples/streaming_chat.rs new file mode 100644 index 0000000..9953909 --- /dev/null +++ b/mofa-engine-sdk/examples/streaming_chat.rs @@ -0,0 +1,91 @@ +//! Streaming — real per-token output, and time-to-first-token. +//! +//! The engine streams genuine incremental tokens from the backend (Ollama NDJSON +//! and cloud SSE are both parsed into the internal `StreamChunk` stream), not a +//! whole answer sliced up after the fact. This demo makes that visible by +//! printing the **time to first token** — the number that only moves if streaming +//! is real — then the tokens as they arrive, and the token/cost total at the end. +//! +//! Run (needs a chat backend; local Ollama is fine): +//! cargo run -p mofa-engine-sdk --example streaming_chat +//! +//! With no backend it prints the structured error (including any `failed_chain`). + +use std::io::Write; +use std::time::Instant; + +use mofa_engine_core::{Engine, EngineConfig}; +use mofa_kernel::{Capability, InferenceRequest, Message, StreamChunk}; + +#[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 request = InferenceRequest { + capability: Some(Capability::Chat), + messages: vec![Message { + role: "user".into(), + content: "Write a short haiku about local-first inference.".into(), + ..Default::default() + }], + ..Default::default() + }; + + println!("== streaming chat — watch time-to-first-token ==\n"); + let started = Instant::now(); + let mut first_token_at: Option = None; + + let mut rx = engine.invoke_stream(request); + while let Some(chunk) = rx.recv().await { + match chunk { + StreamChunk::Started { + provider, + model_used, + .. + } => { + println!("[routed to {provider}/{model_used}]\n"); + } + StreamChunk::Text { delta } => { + // Record when the very first content token lands — the honest + // latency-to-first-token that pseudo-streaming can't improve. + if first_token_at.is_none() { + first_token_at = Some(started.elapsed().as_millis()); + } + print!("{delta}"); + let _ = std::io::stdout().flush(); + } + StreamChunk::Completed { + tokens_used, + cost_usd, + .. + } => { + let ttft = first_token_at + .map(|ms| format!("{ms}ms")) + .unwrap_or_else(|| "n/a".into()); + let cost = cost_usd + .map(|c| format!("${c:.6}")) + .unwrap_or_else(|| "$0.00 (local)".into()); + println!( + "\n\n[first token: {ttft} · total: {}ms · tokens {:?} · {cost}]", + started.elapsed().as_millis(), + tokens_used, + ); + } + StreamChunk::Error(info) => { + eprintln!("\nerror [{:?}]: {}", info.code, info.message); + for a in &info.failed_chain { + eprintln!(" tried {}/{}: {}", a.provider, a.model, a.reason); + } + } + _ => {} + } + } +} diff --git a/mofa-engine-sdk/examples/video_gen.rs b/mofa-engine-sdk/examples/video_gen.rs new file mode 100644 index 0000000..db2fff7 --- /dev/null +++ b/mofa-engine-sdk/examples/video_gen.rs @@ -0,0 +1,109 @@ +//! Scenario S4 (video) — text-to-video at the **API level**. +//! +//! The engine exposes `Capability::VideoGen` through the same contract as every +//! other capability, so a prompt can be turned into a real clip by a cloud video +//! model — ByteDance's **Seedance** via the Volcengine Ark / BytePlus task API — +//! without the caller knowing anything about task submission, polling, or the +//! download. Configure it (see `config.example.toml`, `kind = "cloud_video_gen"`) +//! and the request below returns a managed `mp4` artifact: +//! +//! [[providers]] +//! name = "seedance" +//! kind = "cloud_video_gen" +//! api_key = "env:ARK_API_KEY" +//! [[providers.models]] +//! name = "doubao-seedance-1-0-pro-250528" +//! capability = "video_gen" +//! +//! With no video backend configured (the default offline setup) the demo still +//! runs: it reports that no `video_gen` model is available and prints the typed +//! routing error, rather than pretending to render — the honest degradation the +//! rest of the scenarios follow. + +use mofa_engine_core::EngineConfig; +use mofa_engine_sdk::EmbeddedEngine; +use mofa_kernel::{Capability, InferenceRequest, Message, Prefer}; + +const PROMPT: &str = "A paper airplane glides over a sunlit desk, then banks toward the window. Cinematic, soft light."; + +fn main() { + 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!("== S4 Video — text-to-video via the video_gen capability =="); + println!("prompt: {PROMPT}\n"); + + // Surface what the engine can actually do: is any video_gen model reachable? + // This is the "support at the API level" check — the capability is first-class + // whether or not a backend happens to be configured right now. + let video_models: Vec<_> = engine + .capabilities() + .into_iter() + .filter(|c| c.capabilities.contains(&Capability::VideoGen)) + .collect(); + + if video_models.is_empty() { + println!("no video_gen backend configured."); + println!( + "add a cloud_video_gen (Seedance) or local_video_gen provider — see \ + config.example.toml — then re-run to render a real clip." + ); + } else { + println!("video_gen backends available:"); + for c in &video_models { + println!(" - {}/{} [{:?}]", c.provider, c.name, c.availability); + } + println!(); + } + + // Issue the request regardless: with a backend it renders and returns an mp4; + // without one it returns a typed `no_capable_model` error. Either way the + // *call site* is identical — that is the point of API-level support. + let request = InferenceRequest { + capability: Some(Capability::VideoGen), + messages: vec![Message { + role: "user".into(), + content: PROMPT.into(), + ..Default::default() + }], + // Seedance generation knobs travel in `params` (ratio / resolution / + // duration / fps / seed, and `image_url` for image-to-video). + params: serde_json::json!({ + "ratio": "16:9", + "resolution": "720p", + "duration": 5 + }), + // Video generation is inherently a cloud call here, so allow it. + prefer: Prefer::Auto, + ..Default::default() + }; + + match engine.invoke(request) { + Ok(resp) => { + println!( + "[video_gen] served by {} in {}ms", + resp.provider, resp.duration_ms + ); + match resp.file { + Some(path) => println!( + "clip saved → {path}\n(verify it with the quality gate before shipping)" + ), + None => println!("backend returned no file artifact"), + } + } + Err(e) => { + let info = e.info(); + println!( + "[video_gen] not rendered — {:?}: {}", + info.code, info.message + ); + } + } +} diff --git a/mofa-engine-sdk/src/client.rs b/mofa-engine-sdk/src/client.rs new file mode 100644 index 0000000..5e639e3 --- /dev/null +++ b/mofa-engine-sdk/src/client.rs @@ -0,0 +1,547 @@ +//! 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, Message, + ModelCard, ResponsesRequest, ResponsesResponse, +}; + +/// Errors returned by [`DaemonClient`]. +#[derive(Debug, thiserror::Error)] +pub enum ClientError { + /// Network/transport failure reaching the daemon. + #[error("transport error: {0}")] + Transport(String), + /// The daemon returned a structured engine error. + #[error("engine error: {}", .0.message)] + Engine(ErrorInfo), + /// The daemon returned a non-success status without a structured body + /// (e.g. a 401/404/413 from the framework rather than the handler). + #[error("http {status}: {body}")] + Status { + /// HTTP status code. + status: u16, + /// Raw response body. + body: String, + }, + /// A success response body could not be decoded into the expected type. + #[error("decode error: {0}")] + Decode(String), +} + +impl ClientError { + /// The structured engine error, if this was an engine-level failure. + pub fn engine_error(&self) -> Option<&ErrorInfo> { + match self { + Self::Engine(info) => Some(info), + _ => None, + } + } +} + +/// A synchronous, in-process facade over the engine for embedded/UniFFI use. +/// +/// Owns a multi-threaded Tokio runtime and blocks on each call, so it can be +/// driven from ordinary synchronous code. Background engine tasks (idle +/// eviction, speculative warming) run on the owned runtime for the lifetime of +/// this value. +pub struct EmbeddedEngine { + runtime: tokio::runtime::Runtime, + engine: Arc, +} + +impl EmbeddedEngine { + /// Build and initialize an embedded engine from configuration. + pub fn new(config: EngineConfig) -> Result { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .map_err(|e| EngineError::Internal(format!("failed to build runtime: {e}")))?; + let engine = runtime.block_on(Engine::try_new(config))?; + Ok(Self { runtime, engine }) + } + + /// List all known model cards. + pub fn capabilities(&self) -> Vec { + self.runtime.block_on(self.engine.capabilities()) + } + + /// Run one inference request to completion. + pub fn invoke(&self, request: InferenceRequest) -> Result { + self.runtime.block_on(self.engine.invoke(request)) + } + + /// Run one turn of the stateful Responses API (multi-turn deep reasoning). + /// The returned [`ResponsesResponse::id`] is passed as `previous_response_id` + /// to continue the conversation on the next turn. + pub fn respond(&self, request: ResponsesRequest) -> Result { + self.runtime.block_on(self.engine.respond(request)) + } + + /// Snapshot the engine status. + pub fn status(&self) -> EngineStatus { + self.runtime.block_on(self.engine.status()) + } + + /// Re-run discovery and health probes. + pub fn refresh(&self) { + self.runtime.block_on(self.engine.refresh_resources()); + } + + /// Register a capability subscription; returns its id. + pub fn subscribe( + &self, + app_id: Option, + session_id: Option, + capabilities: Vec, + ttl_secs: Option, + ) -> u64 { + // `Engine::subscribe` spawns background warm tasks, so it must run with + // the owned runtime entered on this thread — otherwise `tokio::spawn` + // panics when called from a plain (non-async) caller. + let _guard = self.runtime.enter(); + self.engine.subscribe( + app_id, + session_id, + capabilities, + ttl_secs.map(Duration::from_secs), + ) + } + + /// Remove a subscription by id. + pub fn unsubscribe(&self, id: u64) -> bool { + let _guard = self.runtime.enter(); + self.engine.unsubscribe(id) + } + + /// Access the underlying async engine for advanced, async-native use. + pub fn engine(&self) -> &Arc { + &self.engine + } +} + +/// JSON string facade — the surface a UniFFI binding exports to Python. +/// +/// UniFFI cannot carry `serde_json::Value`, borrowed types, or the full domain +/// structs across the FFI boundary without hand-written type maps. These +/// string-in / string-out methods sidestep that: they are stable, language +/// agnostic, and map onto UniFFI's `string` and `Result` +/// natively. A generated binding is a thin wrapper that forwards JSON. +impl EmbeddedEngine { + /// Capability list as a JSON array. + pub fn capabilities_json(&self) -> String { + serde_json::to_string(&self.capabilities()).unwrap_or_else(|_| "[]".into()) + } + + /// Engine status as a JSON object. + pub fn status_json(&self) -> String { + serde_json::to_string(&self.status()).unwrap_or_else(|_| "{}".into()) + } + + /// Parse an [`InferenceRequest`] from JSON, invoke, and return the response + /// as JSON. On failure the `Err` payload is a JSON [`ErrorInfo`] with a + /// stable code, so a binding can raise a typed exception. + pub fn invoke_json(&self, request_json: &str) -> Result { + let request: InferenceRequest = serde_json::from_str(request_json) + .map_err(|e| Self::err_json(EngineError::InvalidRequest(e.to_string())))?; + self.invoke(request) + .map(|resp| serde_json::to_string(&resp).unwrap_or_default()) + .map_err(Self::err_json) + } + + /// Parse a [`ResponsesRequest`] from JSON, run one turn, and return the + /// [`ResponsesResponse`] as JSON — the UniFFI-friendly stateful surface. On + /// failure the `Err` payload is a JSON [`ErrorInfo`] with a stable code. + pub fn respond_json(&self, request_json: &str) -> Result { + let request: ResponsesRequest = serde_json::from_str(request_json) + .map_err(|e| Self::err_json(EngineError::InvalidRequest(e.to_string())))?; + self.respond(request) + .map(|resp| serde_json::to_string(&resp).unwrap_or_default()) + .map_err(Self::err_json) + } +} + +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). +pub struct DaemonClient { + base_url: String, + http: reqwest::Client, + token: Option, +} + +impl DaemonClient { + /// Create a client for a daemon at `base_url` (e.g. `http://localhost:8420`). + pub fn new(base_url: impl Into) -> Self { + Self::with_client(base_url, reqwest::Client::new()) + } + + /// Create a client with a caller-provided HTTP client. + pub fn with_client(base_url: impl Into, http: reqwest::Client) -> Self { + let base_url = base_url.into().trim_end_matches('/').to_string(); + Self { + base_url, + http, + token: None, + } + } + + /// Attach a bearer token, sent on every request. Required when the daemon + /// runs with `MOFA_API_TOKEN` set. + pub fn with_token(mut self, token: impl Into) -> Self { + self.token = Some(token.into()); + self + } + + fn url(&self, path: &str) -> String { + format!("{}/{}", self.base_url, path.trim_start_matches('/')) + } + + /// Attach the bearer token to a request builder when one is configured. + fn authed(&self, rb: reqwest::RequestBuilder) -> reqwest::RequestBuilder { + match &self.token { + Some(token) => rb.bearer_auth(token), + None => rb, + } + } + + async fn get_json(&self, path: &str) -> Result { + let resp = self + .authed(self.http.get(self.url(path))) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } + + async fn 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 { + let status = resp.status(); + if status.is_success() { + return resp + .json::() + .await + .map_err(|e| ClientError::Decode(e.to_string())); + } + // Read the body once, then prefer a structured ErrorInfo; fall back to a + // status-bearing error for framework responses (401/404/413/…) so the + // status and reason are not lost. + let body = resp.text().await.unwrap_or_default(); + match serde_json::from_str::(&body) { + Ok(info) => Err(ClientError::Engine(info)), + Err(_) => Err(ClientError::Status { + status: status.as_u16(), + body, + }), + } + } + + /// `GET /v1/capabilities` + pub async fn capabilities(&self) -> Result, ClientError> { + self.get_json("/v1/capabilities").await + } + + /// `GET /v1/status` + pub async fn status(&self) -> Result { + self.get_json("/v1/status").await + } + + /// `GET /v1/memory` + pub async fn memory(&self) -> Result { + self.get_json("/v1/memory").await + } + + /// `GET /v1/lifecycle` + pub async fn lifecycle(&self) -> Result, ClientError> { + self.get_json("/v1/lifecycle").await + } + + /// `GET /v1/preflight` + pub async fn preflight(&self) -> Result { + self.get_json("/v1/preflight").await + } + + /// `GET /v1/subscriptions` + pub async fn subscriptions(&self) -> Result, ClientError> { + self.get_json("/v1/subscriptions").await + } + + /// `POST /v1/invoke` + pub async fn invoke( + &self, + request: &InferenceRequest, + ) -> Result { + let resp = self + .authed(self.http.post(self.url("/v1/invoke")).json(request)) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } + + /// `POST /v1/responses` — one turn of the stateful Responses API. The returned + /// [`ResponsesResponse::id`] is passed as `previous_response_id` to continue. + pub async fn respond( + &self, + request: &ResponsesRequest, + ) -> Result { + let resp = self + .authed(self.http.post(self.url("/v1/responses")).json(request)) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } + + /// `GET /v1/responses/{id}` — the stored message history for a conversation. + pub async fn conversation(&self, id: &str) -> Result, ClientError> { + self.get_json(&format!("/v1/responses/{id}")).await + } + + /// `DELETE /v1/responses/{id}` — forget a stored conversation; returns whether + /// one existed. A not-found id is reported as `Ok(false)`, not an error. + pub async fn delete_conversation(&self, id: &str) -> Result { + let resp = self + .authed(self.http.delete(self.url(&format!("/v1/responses/{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 DeleteResponse { + removed: bool, + } + let resp: DeleteResponse = Self::decode(resp).await?; + Ok(resp.removed) + } + + /// `POST /v1/discovery/refresh` + pub async fn refresh(&self) -> Result { + let resp = self + .authed(self.http.post(self.url("/v1/discovery/refresh"))) + .send() + .await + .map_err(|e| ClientError::Transport(e.to_string()))?; + Self::decode(resp).await + } + + /// `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)] +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(), + ..Default::default() + }; + 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 embedded_responds_and_reports_invalid_input() { + let engine = EmbeddedEngine::new(empty_config()).unwrap(); + + // A turn with input but no models still routes and fails with a typed + // no_capable_model error (the stateful path reaches routing). + let err = engine + .respond_json(r#"{"input":"hello","capability":"chat"}"#) + .unwrap_err(); + assert!(err.contains("no_capable_model"), "got: {err}"); + + // An empty turn is rejected before routing. + let err = engine.respond_json("{}").unwrap_err(); + assert!(err.contains("invalid_request"), "got: {err}"); + } + + #[test] + fn daemon_client_normalizes_urls() { + let client = DaemonClient::new("http://localhost:8420/"); + assert_eq!(client.url("/v1/status"), "http://localhost:8420/v1/status"); + assert_eq!(client.url("v1/status"), "http://localhost:8420/v1/status"); + } + + #[test] + fn with_token_is_retained_and_url_unaffected() { + let client = DaemonClient::new("http://localhost:8420").with_token("secret"); + assert_eq!(client.token.as_deref(), Some("secret")); + assert_eq!(client.url("/v1/status"), "http://localhost:8420/v1/status"); + } + + #[test] + fn client_error_status_preserves_code_and_body() { + let err = ClientError::Status { + status: 401, + body: "unauthorized".into(), + }; + assert!(err.to_string().contains("401")); + assert!(err.to_string().contains("unauthorized")); + assert!(err.engine_error().is_none()); + } +} diff --git a/mofa-engine-sdk/src/dashboard.rs b/mofa-engine-sdk/src/dashboard.rs index a3aaee6..8e0cf26 100644 --- a/mofa-engine-sdk/src/dashboard.rs +++ b/mofa-engine-sdk/src/dashboard.rs @@ -4,7 +4,7 @@ //! No external dependencies — vanilla JS + CSS with glassmorphism design. /// The complete dashboard HTML page. -pub const DASHBOARD_HTML: &str = r##" +pub(crate) const DASHBOARD_HTML: &str = r##" @@ -265,7 +265,9 @@ body { .cap-badge.chat { background: rgba(59, 130, 246, 0.15); color: var(--accent-blue); } .cap-badge.tts { background: rgba(139, 92, 246, 0.15); color: var(--accent-purple); } .cap-badge.asr { background: rgba(6, 182, 212, 0.15); color: var(--accent-cyan); } +.cap-badge.vlm { background: rgba(59, 130, 246, 0.15); color: var(--accent-cyan); } .cap-badge.image_gen { background: rgba(245, 158, 11, 0.15); color: var(--accent-yellow); } +.cap-badge.video_gen { background: rgba(239, 68, 68, 0.15); color: var(--accent-red); } .cap-badge.embedding { background: rgba(16, 185, 129, 0.15); color: var(--accent-green); } .model-provider { @@ -500,6 +502,7 @@ body {
-- -- +
@@ -563,7 +566,9 @@ body { + + @@ -605,6 +610,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) +}