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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub enum SdkLoadMode {
23 Lightweight,
25 WithPackageIndex,
27}
28
29#[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#[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}