Skip to main content

weft_core/
sdk.rs

1use crate::api::openai_compat::AppState;
2use crate::app::{AppProfile, CorePolicy, GenerationStoreMap, PackageIndex, ResolvedAppMap};
3use crate::config::store::load_config;
4use crate::defaults::{DefaultErrorHandler, DefaultRouter, FailoverSelector};
5use crate::package::{
6    build_service_config, discover_runtime_packages, DiscoveredPackage, PackageInfo,
7    PackageManager, PackageRuntime,
8};
9use crate::pipeline::Pipeline;
10use crate::process::ProcessManager;
11use crate::types::{ChatRequest, ChatResponse};
12use crate::vkeys::VirtualKeyStore;
13use anyhow::Result;
14use axum::Router;
15use std::path::PathBuf;
16use std::sync::{Arc, Mutex as StdMutex};
17use std::time::Duration;
18use tokio::sync::RwLock;
19
20/// Controls how much package metadata the SDK bootstrap loads.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum SdkLoadMode {
23    /// Embedded mode: avoid package index resolution and external package-index network calls.
24    Lightweight,
25    /// Full mode: resolve the package index using the same package-index loader used by the main runtime.
26    WithPackageIndex,
27}
28
29/// Minimal options for embedding or constructing WEFT core services.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub struct WeftCoreOptions {
32    pub config_path: PathBuf,
33    pub repo_root: PathBuf,
34    pub load_mode: SdkLoadMode,
35    pub start_services: bool,
36}
37
38impl WeftCoreOptions {
39    pub fn new(config_path: impl Into<PathBuf>, repo_root: impl Into<PathBuf>) -> Self {
40        Self {
41            config_path: config_path.into(),
42            repo_root: repo_root.into(),
43            load_mode: SdkLoadMode::Lightweight,
44            start_services: false,
45        }
46    }
47
48    pub fn with_load_mode(mut self, load_mode: SdkLoadMode) -> Self {
49        self.load_mode = load_mode;
50        self
51    }
52
53    pub fn with_package_index(mut self) -> Self {
54        self.load_mode = SdkLoadMode::WithPackageIndex;
55        self
56    }
57
58    pub fn with_start_services(mut self, start_services: bool) -> Self {
59        self.start_services = start_services;
60        self
61    }
62}
63
64impl Default for WeftCoreOptions {
65    fn default() -> Self {
66        Self {
67            config_path: PathBuf::from("config/config.toml"),
68            repo_root: PathBuf::from("."),
69            load_mode: SdkLoadMode::Lightweight,
70            start_services: false,
71        }
72    }
73}
74
75fn default_provider_name(config: &crate::config::AppConfig) -> String {
76    config
77        .routing
78        .default_provider
79        .clone()
80        .or_else(|| {
81            config
82                .providers
83                .first()
84                .map(|provider| provider.name.clone())
85        })
86        .unwrap_or_default()
87}
88
89fn build_pipeline(config: &crate::config::AppConfig) -> Result<Pipeline> {
90    Ok(Pipeline {
91        router: Arc::new(DefaultRouter {
92            default_provider: default_provider_name(config),
93        }),
94        key_selector: Arc::new(FailoverSelector),
95        transforms: Arc::new(crate::defaults::transforms::TransformRegistry::with_defaults()),
96        error_handler: Arc::new(DefaultErrorHandler {
97            max_retries: config.fallback.retry_count,
98        }),
99        http_client: reqwest::Client::builder()
100            .connect_timeout(Duration::from_secs(10))
101            .timeout(Duration::from_secs(90))
102            .http1_only()
103            .build()?,
104    })
105}
106
107async fn load_package_index_for_mode(
108    options: &WeftCoreOptions,
109    data_dir: &str,
110    configured_source_url: Option<&str>,
111) -> PackageIndex {
112    match options.load_mode {
113        SdkLoadMode::Lightweight => PackageIndex::default(),
114        SdkLoadMode::WithPackageIndex => {
115            crate::app::resolve_package_index(&options.repo_root, data_dir, configured_source_url)
116                .await
117        }
118    }
119}
120
121async fn register_discovered_runtime_plugins(
122    options: &WeftCoreOptions,
123    package_index: &PackageIndex,
124    process_manager: &Arc<ProcessManager>,
125) -> (
126    PackageManager,
127    Vec<crate::api::openai_compat::ChatProviderInfo>,
128    Vec<DiscoveredPackage>,
129) {
130    let mut package_manager = PackageManager::new();
131    let discovered_packages = match options.load_mode {
132        SdkLoadMode::Lightweight => Vec::new(),
133        SdkLoadMode::WithPackageIndex => {
134            discover_runtime_packages(&options.repo_root, package_index)
135        }
136    };
137
138    for package in &discovered_packages {
139        let manifest = &package.manifest;
140        package_manager.register(PackageInfo {
141            name: manifest.package_info.name.clone(),
142            version: Some(manifest.package_info.version.clone()),
143            overrides: vec![],
144            enabled: true,
145            has_ui: false,
146            description: Some(manifest.package_info.description.clone()),
147        });
148
149        if package.runtime == PackageRuntime::Service {
150            if let Ok(service_config) = build_service_config(package) {
151                process_manager.register(service_config).await;
152            }
153        }
154    }
155
156    let mut chat_providers: Vec<crate::api::openai_compat::ChatProviderInfo> = discovered_packages
157        .iter()
158        .filter(|package| {
159            package
160                .manifest
161                .resolved_provides()
162                .contains(&"chat_channel".to_string())
163        })
164        .map(|package| crate::api::openai_compat::ChatProviderInfo {
165            name: package.manifest.package_info.name.clone(),
166            endpoint: package
167                .manifest
168                .resolved_chat_endpoint()
169                .unwrap_or_else(|| "/chat".to_string()),
170            description: package.manifest.package_info.description.clone(),
171        })
172        .collect();
173    chat_providers.sort_by(|left, right| left.name.cmp(&right.name));
174    chat_providers
175        .dedup_by(|left, right| left.name == right.name && left.endpoint == right.endpoint);
176
177    (package_manager, chat_providers, discovered_packages)
178}
179
180/// Lightweight reusable SDK handle around the existing core application state.
181#[derive(Clone)]
182pub struct WeftCore {
183    state: AppState,
184}
185
186impl WeftCore {
187    pub async fn load(options: WeftCoreOptions) -> Result<Self> {
188        let config = load_config(&options.config_path)?;
189
190        let shared_config = Arc::new(RwLock::new(config.clone()));
191        let shared_pipeline = Arc::new(build_pipeline(&config)?);
192
193        let process_manager = Arc::new(ProcessManager::new());
194        for service in &config.services {
195            process_manager.register(service.clone()).await;
196        }
197        let vkey_store = Arc::new(VirtualKeyStore::new());
198        vkey_store.load_from_config(&config.virtual_keys);
199
200        let package_index = load_package_index_for_mode(
201            &options,
202            &config.core.data_dir,
203            config.registry.package_source_url.as_deref(),
204        )
205        .await;
206        let (package_manager, chat_providers, discovered_packages) =
207            register_discovered_runtime_plugins(&options, &package_index, &process_manager).await;
208        if options.start_services {
209            process_manager.start_auto().await;
210        }
211        let mut capability_registry =
212            crate::app::build_capability_registry(&discovered_packages, &ResolvedAppMap::new());
213        crate::app::merge_core_capabilities(&mut capability_registry);
214        let (_shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
215        drop(shutdown_rx);
216
217        Ok(Self {
218            state: AppState {
219                config: shared_config,
220                config_path: options.config_path,
221                pipeline: shared_pipeline,
222                process_manager,
223                vkey_store,
224                package_manager: Arc::new(RwLock::new(package_manager)),
225                wasm_handle: Arc::new(RwLock::new(None)),
226                native_handle: Arc::new(RwLock::new(None)),
227                resolved_apps: Arc::new(RwLock::new(ResolvedAppMap::new())),
228                capability_registry: Arc::new(RwLock::new(capability_registry)),
229                active_profile: Arc::new(RwLock::new(AppProfile::Safe)),
230                core_policy: Arc::new(CorePolicy::default_policy()),
231                generation_store: Arc::new(RwLock::new(GenerationStoreMap::new())),
232                package_index: Arc::new(package_index),
233                repo_root: options.repo_root.clone(),
234                data_dir: options.repo_root.join("data"),
235                runtime_token: None,
236                runtime_token_path: None,
237                chat_providers: Arc::new(RwLock::new(chat_providers)),
238                shutdown_tx: Arc::new(StdMutex::new(Some(_shutdown_tx))),
239                stream_buffer: Arc::new(StdMutex::new(std::collections::HashMap::new())),
240            },
241        })
242    }
243
244    pub fn from_state(state: AppState) -> Self {
245        Self { state }
246    }
247
248    pub fn state(&self) -> &AppState {
249        &self.state
250    }
251
252    pub fn into_state(self) -> AppState {
253        self.state
254    }
255
256    pub fn router(&self) -> Router {
257        crate::api::build_router(self.state.clone())
258    }
259
260    pub async fn start_services(&self) {
261        self.state.process_manager.start_auto().await;
262    }
263
264    pub async fn shutdown(&self) -> Result<()> {
265        if let Some(shutdown_tx) = self.state.shutdown_tx.lock().unwrap().take() {
266            let _ = shutdown_tx.send(());
267        }
268        self.state.process_manager.stop_all().await;
269        Ok(())
270    }
271
272    pub async fn chat(&self, request: ChatRequest) -> Result<ChatResponse> {
273        let config = self.state.config.read().await;
274        self.state.pipeline.execute(&request, &config).await
275    }
276
277    pub async fn models(&self) -> serde_json::Value {
278        let config = self.state.config.read().await;
279        let models: Vec<serde_json::Value> = config
280            .providers
281            .iter()
282            .flat_map(|provider| {
283                provider.models.iter().map(move |model| {
284                    serde_json::json!({
285                        "id": model,
286                        "object": "model",
287                        "owned_by": provider.name,
288                    })
289                })
290            })
291            .collect();
292
293        serde_json::json!({
294            "object": "list",
295            "data": models,
296        })
297    }
298}
299
300#[cfg(test)]
301mod tests {
302    use super::{WeftCore, WeftCoreOptions, SdkLoadMode};
303    use axum::body::{to_bytes, Body};
304    use axum::http::{Request, StatusCode};
305    use serde_json::json;
306    use tempfile::TempDir;
307    use tower::util::ServiceExt;
308
309    fn write_package_index(dir: &TempDir) {
310        let packages_dir = dir.path().join("packages");
311        std::fs::create_dir_all(&packages_dir).expect("packages dir created");
312        std::fs::write(
313            packages_dir.join("index.toml"),
314            r#"
315version = 1
316revision = "test-sdk"
317source_url = "local://sdk-test"
318
319[[package_sources]]
320name = "sdk-test-runtime"
321kind = "service"
322package_kind = "provider"
323runtime_provider = "sdk-test-runtime"
324current_source = "packages/official/sdk-test-runtime"
325trusted = true
326provides = ["chat_channel"]
327"#,
328        )
329        .expect("package index written");
330    }
331
332    fn write_runtime_plugin(dir: &TempDir) {
333        let package_dir = dir
334            .path()
335            .join("packages")
336            .join("official")
337            .join("sdk-test-runtime");
338        std::fs::create_dir_all(&package_dir).expect("package dir created");
339        std::fs::write(package_dir.join("run.ps1"), "Write-Output sdk-test-runtime")
340            .expect("package entry written");
341        std::fs::write(
342            package_dir.join("package.toml"),
343            r#"
344[package_info]
345name = "sdk-test-runtime"
346version = "0.1.0"
347description = "SDK test runtime"
348entry = "run.ps1"
349provides = ["chat_channel"]
350chat_endpoint = "/sdk-chat"
351
352[package]
353runtime = "service"
354entry = "run.ps1"
355provides = ["chat_channel"]
356chat_endpoint = "/sdk-chat"
357
358[runtime_contract]
359startup_mode = "persistent"
360restart_policy = "manual"
361"#,
362        )
363        .expect("package manifest written");
364    }
365
366    fn write_minimal_config(dir: &TempDir) -> std::path::PathBuf {
367        let config_dir = dir.path().join("config");
368        std::fs::create_dir_all(&config_dir).expect("config dir created");
369        let config_path = config_dir.join("config.toml");
370        std::fs::write(
371            &config_path,
372            r#"
373[core]
374host = "127.0.0.1"
375port = 0
376data_dir = "data"
377
378[routing]
379default_provider = "mock"
380
381[fallback]
382retry_count = 0
383
384[[providers]]
385name = "mock"
386base_url = "http://127.0.0.1:9/v1"
387format = "openai"
388models = ["mock-chat"]
389
390[[providers.keys]]
391value = "sk-test"
392
393[[virtual_keys]]
394key = "vk-test"
395provider = "mock"
396model = "mock-chat"
397"#,
398        )
399        .expect("config written");
400        config_path
401    }
402
403    #[tokio::test]
404    async fn load_constructs_minimal_embedded_core_without_network() {
405        let dir = TempDir::new().expect("temp dir");
406        let config_path = write_minimal_config(&dir);
407
408        let options = WeftCoreOptions::new(&config_path, dir.path());
409        assert_eq!(options.load_mode, SdkLoadMode::Lightweight);
410        assert!(!options.start_services);
411
412        let core = WeftCore::load(options)
413            .await
414            .expect("core loads from minimal config");
415
416        assert_eq!(core.state().config_path, config_path);
417        assert!(core.state().resolved_apps.read().await.is_empty());
418        assert!(core.state().chat_providers.read().await.is_empty());
419        assert!(core.state().package_manager.read().await.list().is_empty());
420        assert!(core.state().process_manager.all_statuses().await.is_empty());
421        assert!(core.state().wasm_handle.read().await.is_none());
422        assert!(core.state().native_handle.read().await.is_none());
423        assert!(core.state().package_index.package_sources.is_empty());
424
425        let models = core.models().await;
426        assert_eq!(models["object"], "list");
427        assert_eq!(models["data"][0]["id"], "mock-chat");
428        assert_eq!(models["data"][0]["owned_by"], "mock");
429    }
430
431    #[tokio::test]
432    async fn router_serves_models_from_loaded_core() {
433        let dir = TempDir::new().expect("temp dir");
434        let config_path = write_minimal_config(&dir);
435        let core = WeftCore::load(WeftCoreOptions::new(&config_path, dir.path()))
436            .await
437            .expect("core loads from minimal config");
438
439        let response = core
440            .router()
441            .oneshot(
442                Request::builder()
443                    .uri("/v1/models")
444                    .body(Body::empty())
445                    .expect("request"),
446            )
447            .await
448            .expect("response");
449
450        assert_eq!(response.status(), StatusCode::OK);
451        let body = to_bytes(response.into_body(), usize::MAX)
452            .await
453            .expect("body bytes");
454        let payload: serde_json::Value = serde_json::from_slice(&body).expect("json payload");
455        assert_eq!(payload["object"], "list");
456        assert_eq!(payload["data"][0]["id"], "mock-chat");
457        assert_eq!(payload["data"][0]["owned_by"], "mock");
458    }
459
460    #[tokio::test]
461    async fn chat_returns_error_without_real_network() {
462        let dir = TempDir::new().expect("temp dir");
463        let config_path = write_minimal_config(&dir);
464        let core = WeftCore::load(WeftCoreOptions::new(&config_path, dir.path()))
465            .await
466            .expect("core loads from minimal config");
467
468        let response = core
469            .router()
470            .oneshot(
471                Request::builder()
472                    .method("POST")
473                    .uri("/v1/chat/completions")
474                    .header("content-type", "application/json")
475                    .body(Body::from(
476                        json!({
477                            "model": "mock-chat",
478                            "messages": [{"role": "user", "content": "hello"}]
479                        })
480                        .to_string(),
481                    ))
482                    .expect("request"),
483            )
484            .await
485            .expect("response");
486
487        assert_eq!(response.status(), StatusCode::BAD_GATEWAY);
488        let body = to_bytes(response.into_body(), usize::MAX)
489            .await
490            .expect("body bytes");
491        let payload: serde_json::Value = serde_json::from_slice(&body).expect("json payload");
492        assert_eq!(payload["error"]["type"], "proxy_error");
493        assert_eq!(payload["error"]["code"], "bad_gateway");
494        assert!(!payload["error"]["message"]
495            .as_str()
496            .expect("error message")
497            .is_empty());
498    }
499
500    #[tokio::test]
501    async fn with_package_index_loads_runtime_plugin_metadata_without_starting_services() {
502        let dir = TempDir::new().expect("temp dir");
503        let config_path = write_minimal_config(&dir);
504        write_package_index(&dir);
505        write_runtime_plugin(&dir);
506
507        let core =
508            WeftCore::load(WeftCoreOptions::new(&config_path, dir.path()).with_package_index())
509                .await
510                .expect("core loads with package index");
511
512        assert_eq!(core.state().package_index.package_sources.len(), 1);
513        assert!(core
514            .state()
515            .package_manager
516            .read()
517            .await
518            .get("sdk-test-runtime")
519            .is_some());
520        assert_eq!(core.state().chat_providers.read().await.len(), 1);
521        assert_eq!(
522            core.state()
523                .process_manager
524                .status("sdk-test-runtime")
525                .await
526                .expect("service registered"),
527            crate::process::ServiceStatus::Stopped
528        );
529        assert!(core.state().wasm_handle.read().await.is_none());
530        assert!(core.state().native_handle.read().await.is_none());
531    }
532
533    #[tokio::test]
534    async fn shutdown_is_noop_for_unstarted_sdk_core() {
535        let dir = TempDir::new().expect("temp dir");
536        let config_path = write_minimal_config(&dir);
537        let core = WeftCore::load(WeftCoreOptions::new(&config_path, dir.path()))
538            .await
539            .expect("core loads from minimal config");
540
541        core.shutdown().await.expect("first shutdown succeeds");
542        core.shutdown().await.expect("second shutdown succeeds");
543        assert!(core.state().process_manager.all_statuses().await.is_empty());
544    }
545}