1use std::collections::HashMap;
28use std::net::{IpAddr, SocketAddr};
29use std::sync::Arc;
30use std::time::Duration;
31
32use anyhow::{Context as _, anyhow, bail};
33use talon_container::{ContainerRunner, ContainerRuntime, LogRedactor, MountPolicy, RunnerConfig};
34use talon_core::audit::{SharedAuditSink, event};
35use talon_core::channel::{Channel, ChannelHost};
36use talon_core::config::TalonConfig;
37use talon_core::provider::{effective_vllm, merge_container_config};
38use talon_core::types::RegisteredGroup;
39use talon_db::{Db, DbAuditSink, DbOptions, default_db_path};
40use talon_http::{HttpChannel, HttpConfig};
41use talon_ipc::{IpcHost, IpcWatcher, IpcWatcherConfig};
42use talon_llm::gateway::{GatewayConfig, GatewayHandle, LlmGateway};
43use talon_scheduler::{GroupQueue, QueueConfig, SchedulerConfig, TaskScheduler};
44use talon_security::mounts::MountValidator;
45use talon_security::redact::Redactor;
46use talon_security::senders::REQUIRED_ENV;
47use tokio_util::sync::CancellationToken;
48
49use crate::agent::{GatewayLlm, LlmAccess, RunnerLauncher};
50use crate::allowlist::AllowlistSource;
51use crate::orchestrator::{Orchestrator, OrchestratorDeps, Settings};
52use crate::palace::host_palace_runner;
53
54pub const DEFAULT_CONTAINER_HOST: &str = "host.docker.internal";
57pub const ALLOW_NO_API_KEY_ENV: &str = "TALON_VLLM_ALLOW_NO_API_KEY";
59pub const VLLM_API_KEY_REQUIRED: &str = "TALON_VLLM_API_KEY is not set: refusing to start the LLM \
61 gateway. Without an API key vLLM serves anyone who can reach it, so agent containers could \
62 call it directly and bypass the gateway (per-run tokens, model allowlist, request filtering). \
63 Set TALON_VLLM_API_KEY in .env to vLLM's --api-key; only if containers provably cannot reach \
64 vLLM, set TALON_VLLM_ALLOW_NO_API_KEY=true. Every agent run fails until then — run `talon doctor`";
65const CHANNEL_DISCONNECT_TIMEOUT: Duration = Duration::from_secs(10);
67const LOOP_STOP_TIMEOUT: Duration = Duration::from_secs(5);
69
70pub fn log_audit_head(db: &Db, phase: &str) {
72 match (db.audit_chain_head(), db.audit_count()) {
73 (Ok(Some((id, hash))), Ok(count)) => {
74 tracing::info!(phase, head_id = id, head_hash = %hash, events = count, "audit chain head");
75 }
76 (Ok(None), _) => tracing::info!(phase, "audit chain empty"),
77 (Err(e), _) | (_, Err(e)) => {
78 tracing::warn!(phase, error = %e, "cannot read audit chain head")
79 }
80 }
81}
82
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum ApiKeyPolicy {
86 Present,
88 MissingAllowed,
90 Missing,
92}
93
94pub fn api_key_policy(config: &TalonConfig) -> ApiKeyPolicy {
96 let present = config
97 .env
98 .file_value("TALON_VLLM_API_KEY")
99 .is_some_and(|k| !k.trim().is_empty());
100 if present {
101 ApiKeyPolicy::Present
102 } else if config.env.get_bool(ALLOW_NO_API_KEY_ENV) {
103 ApiKeyPolicy::MissingAllowed
104 } else {
105 ApiKeyPolicy::Missing
106 }
107}
108
109fn parse_bool(raw: &str) -> Option<bool> {
111 match raw.trim().to_ascii_lowercase().as_str() {
112 "1" | "true" | "yes" | "on" => Some(true),
113 "0" | "false" | "no" | "off" => Some(false),
114 _ => None,
115 }
116}
117
118fn normalize_base(url: &str) -> String {
119 let u = url.trim().trim_end_matches('/');
120 u.strip_suffix("/v1").unwrap_or(u).to_ascii_lowercase()
121}
122
123pub fn gateway_config(
126 config: &TalonConfig,
127 groups: &HashMap<String, RegisteredGroup>,
128 bind: SocketAddr,
129) -> Option<GatewayConfig> {
130 let vllm = effective_vllm(&config.provider_defaults)?;
131 let api_key = config
132 .env
133 .file_value("TALON_VLLM_API_KEY")
134 .map(str::trim)
135 .filter(|k| !k.is_empty())
136 .map(str::to_string);
137 let mut gw = GatewayConfig::from_vllm(vllm, api_key);
138 gw.bind = bind;
139 if let Some(rpm) = config
140 .env
141 .get("TALON_LLM_GATEWAY_RPM")
142 .and_then(|v| v.trim().parse::<u32>().ok())
143 {
144 gw.requests_per_minute = rpm;
145 }
146 gw.audit_requests = config.env.get_bool("TALON_LLM_GATEWAY_AUDIT_REQUESTS");
147 if let Some(raw) = config.env.get("TALON_LLM_GATEWAY_MAX_CONCURRENT") {
148 match raw.trim().parse::<u32>() {
149 Ok(n) => gw.max_concurrent_requests = n,
150 Err(_) => tracing::warn!(
151 value = %raw,
152 default = gw.max_concurrent_requests,
153 "TALON_LLM_GATEWAY_MAX_CONCURRENT is not a number; using the default"
154 ),
155 }
156 }
157 if let Some(raw) = config.env.get("TALON_LLM_GATEWAY_CACHE_SALT") {
158 match parse_bool(&raw) {
159 Some(on) => gw.cache_salt = on,
160 None => tracing::warn!(
161 value = %raw,
162 default = gw.cache_salt,
163 "TALON_LLM_GATEWAY_CACHE_SALT is not true/false; using the default"
164 ),
165 }
166 }
167
168 let upstream = normalize_base(&vllm.base_url);
169 let mut jids: Vec<&String> = groups.keys().collect();
170 jids.sort();
171 for jid in jids {
172 let group = &groups[jid];
173 let merged =
174 merge_container_config(&config.provider_defaults, group.container_config.as_ref());
175 let Some(v) = effective_vllm(&merged) else {
176 continue;
177 };
178 if normalize_base(&v.base_url) != upstream {
179 tracing::warn!(
180 jid = %jid,
181 group = %group.name,
182 "group overrides the vLLM base URL; the LLM gateway only proxies to the host \
183 upstream (TALON_VLLM_BASE_URL), so its requests go there"
184 );
185 }
186 gw = gw.allow_model(v.model.clone());
187 if let Some(fast) = &v.fast_model {
188 gw = gw.allow_model(fast.clone());
189 }
190 }
191 Some(gw)
192}
193
194pub async fn gateway_bind_ip(
200 config: &TalonConfig,
201 runtime: &ContainerRuntime,
202) -> anyhow::Result<IpAddr> {
203 match config.env.get("TALON_LLM_GATEWAY_BIND") {
204 Some(raw) => raw
205 .trim()
206 .parse::<IpAddr>()
207 .with_context(|| format!("TALON_LLM_GATEWAY_BIND {raw:?} is not an IP address")),
208 None => Ok(talon_container::detect_gateway_bind_addr(runtime).await),
209 }
210}
211
212async fn start_gateway(
213 config: &TalonConfig,
214 runtime: &ContainerRuntime,
215 groups: &HashMap<String, RegisteredGroup>,
216 audit: SharedAuditSink,
217) -> anyhow::Result<Option<GatewayHandle>> {
218 if effective_vllm(&config.provider_defaults).is_none() {
219 tracing::error!(
220 "vLLM is not configured (TALON_VLLM_BASE_URL / TALON_VLLM_MODEL): the LLM gateway is \
221 not started and every agent run will fail until it is — run `talon doctor`"
222 );
223 return Ok(None);
224 }
225 match api_key_policy(config) {
226 ApiKeyPolicy::Present => {}
227 ApiKeyPolicy::MissingAllowed => tracing::warn!(
228 "TALON_VLLM_API_KEY is not set and {ALLOW_NO_API_KEY_ENV}=true: the gateway calls vLLM \
229 without an API key; make sure agent containers cannot reach vLLM directly"
230 ),
231 ApiKeyPolicy::Missing => {
232 tracing::error!("{VLLM_API_KEY_REQUIRED}");
233 return Ok(None);
234 }
235 }
236 let ip = gateway_bind_ip(config, runtime).await?;
237 let bind = SocketAddr::new(ip, config.llm_gateway_port);
238 let Some(gw) = gateway_config(config, groups, bind) else {
239 return Ok(None);
240 };
241 let handle = LlmGateway::start(gw, audit)
242 .await
243 .map_err(|e| anyhow!("cannot start the LLM gateway on {bind}: {e}"))?;
244 Ok(Some(handle))
245}
246
247pub async fn ensure_agent_network(
256 runtime: &ContainerRuntime,
257 runner_config: &RunnerConfig,
258) -> anyhow::Result<()> {
259 let Some(network) = runner_config.default_network.as_deref() else {
260 tracing::warn!(
261 "agent containers use the runtime's default bridge network (no dedicated \
262 talon-agents network): containers can reach each other and egress is not \
263 restricted by `talon setup egress`"
264 );
265 return Ok(());
266 };
267 runtime.ensure_agent_network(network).await.map_err(|e| {
268 anyhow!(
269 "cannot prepare the agent container network {network:?}: {e}. Run `talon setup \
270 egress` (creates the network and its egress firewall rules), or check that the \
271 container runtime can create networks"
272 )
273 })?;
274 tracing::info!(network, "agent network ready");
275 Ok(())
276}
277
278async fn wait_for_signal() -> &'static str {
279 use tokio::signal::unix::{SignalKind, signal};
280 let mut term = match signal(SignalKind::terminate()) {
281 Ok(s) => Some(s),
282 Err(e) => {
283 tracing::warn!(error = %e, "cannot install SIGTERM handler");
284 None
285 }
286 };
287 tokio::select! {
288 r = tokio::signal::ctrl_c() => {
289 if let Err(e) = r {
290 tracing::warn!(error = %e, "cannot listen for SIGINT");
291 std::future::pending::<()>().await;
292 }
293 "SIGINT"
294 }
295 () = async {
296 match term.as_mut() {
297 Some(t) => { t.recv().await; }
298 None => std::future::pending::<()>().await,
299 }
300 } => "SIGTERM",
301 }
302}
303
304async fn connect_channels(candidates: Vec<Arc<dyn Channel>>) -> Vec<Arc<dyn Channel>> {
305 let mut connected = Vec::new();
306 for channel in candidates {
307 match channel.connect().await {
308 Ok(()) => {
309 tracing::info!(channel = channel.name(), "channel connected");
310 connected.push(channel);
311 }
312 Err(e) => {
313 tracing::error!(channel = channel.name(), error = %e, "channel failed to connect")
314 }
315 }
316 }
317 connected
318}
319
320async fn join_with_timeout(name: &str, handle: tokio::task::JoinHandle<()>) {
321 let abort = handle.abort_handle();
322 match tokio::time::timeout(LOOP_STOP_TIMEOUT, handle).await {
323 Ok(Ok(())) => {}
324 Ok(Err(e)) => tracing::warn!(task = name, error = %e, "background task ended abnormally"),
325 Err(_) => {
326 tracing::warn!(
327 task = name,
328 "background task did not stop in time; aborting"
329 );
330 abort.abort();
331 }
332 }
333}
334
335pub async fn run(config: Arc<TalonConfig>, redactor: Arc<Redactor>) -> anyhow::Result<()> {
341 tracing::info!(version = crate::VERSION, root = %config.project_root.display(), "Talon starting");
342
343 let runtime = Arc::new(
345 ContainerRuntime::detect(&config.container_runtime)
346 .map_err(|e| anyhow!("container runtime: {e}"))?,
347 );
348 runtime.ensure_running().await.map_err(|e| anyhow!("{e}"))?;
349 let orphans = runtime.cleanup_orphans().await;
350 if !orphans.is_empty() {
351 tracing::info!(count = orphans.len(), names = ?orphans, "Stopped orphaned containers");
352 }
353 let runner_config = RunnerConfig::from_talon_config(&config);
354 ensure_agent_network(&runtime, &runner_config).await?;
355
356 let db = Db::open_with_options(
358 &default_db_path(&config.store_dir),
359 &DbOptions {
360 assistant_name: config.assistant_name.clone(),
361 },
362 )
363 .map_err(|e| anyhow!("cannot open database: {e}"))?;
364 tracing::info!("Database initialized");
365 let audit: SharedAuditSink = Arc::new(DbAuditSink::new(db.clone()));
366 log_audit_head(&db, "startup");
367 audit.record(event(
368 "talon.start",
369 "talon",
370 None,
371 "info",
372 serde_json::json!({ "version": crate::VERSION, "pid": std::process::id() }),
373 ));
374
375 let persisted_groups = db
377 .get_all_registered_groups()
378 .map_err(|e| anyhow!("cannot load registered groups: {e}"))?;
379 let gateway = start_gateway(&config, &runtime, &persisted_groups, audit.clone()).await?;
380 let container_host = config
381 .env
382 .get("TALON_LLM_GATEWAY_HOST")
383 .unwrap_or_else(|| DEFAULT_CONTAINER_HOST.to_string());
384 let llm: Option<Arc<dyn LlmAccess>> = gateway
385 .clone()
386 .map(|g| Arc::new(GatewayLlm::new(g, container_host)) as Arc<dyn LlmAccess>);
387
388 let mounts: Arc<dyn MountPolicy> = Arc::new(MountValidator::load(&config.mount_allowlist_path));
390 let log_redactor: Arc<dyn LogRedactor> = redactor.clone();
391 let agent_network = runner_config.default_network.clone();
392 let runner = ContainerRunner::new(
393 runner_config,
394 runtime.clone(),
395 mounts,
396 log_redactor,
397 audit.clone(),
398 );
399 let mut queue_config = QueueConfig::from_talon_config(&config);
400 queue_config.stop_on_shutdown = true;
403 let queue = GroupQueue::new(queue_config);
404 let settings = Settings::default();
405 let shutdown_grace = settings.shutdown_grace;
406 let orch = Orchestrator::new(OrchestratorDeps {
407 config: config.clone(),
408 db: db.clone(),
409 queue: queue.clone(),
410 launcher: Arc::new(RunnerLauncher::new(runner)),
411 llm,
412 allowlist: AllowlistSource::from_file(
413 config.sender_allowlist_path.clone(),
414 config.env.get_bool(REQUIRED_ENV),
415 ),
416 redactor: redactor.clone(),
417 audit: audit.clone(),
418 settings,
419 })
420 .map_err(|e| anyhow!("cannot load state: {e}"))?;
421 let scheduler = Arc::new(TaskScheduler::new(
422 SchedulerConfig::from_talon_config(&config),
423 Arc::new(db.clone()),
424 queue.clone(),
425 orch.clone(),
426 orch.clone(),
427 ));
428 orch.set_scheduler(scheduler.clone());
429 queue.set_processor(orch.clone());
430
431 let host: Arc<dyn ChannelHost> = orch.clone();
433 let http_config =
434 HttpConfig::from_talon_config(&config).map_err(|e| anyhow!("HTTP channel: {e}"))?;
435 let palace = host_palace_runner(&config, &http_config, &runtime, agent_network);
436 let mut candidates: Vec<Arc<dyn Channel>> = vec![Arc::new(HttpChannel::with_palace(
437 http_config,
438 host.clone(),
439 db.clone(),
440 audit.clone(),
441 palace,
442 ))];
443 candidates.extend(talon_channels::build_channels(
444 host.clone(),
445 &config.env,
446 audit.clone(),
447 ));
448 let channels = connect_channels(candidates).await;
449 if channels.is_empty() {
450 bail!("No channels connected");
451 }
452 orch.set_channels(channels.clone());
453
454 let cancel = CancellationToken::new();
456 let scheduler_task = tokio::spawn({
457 let scheduler = scheduler.clone();
458 let cancel = cancel.clone();
459 async move { scheduler.run(cancel).await }
460 });
461 let ipc_host: Arc<dyn IpcHost> = orch.clone();
462 let watcher = IpcWatcher::new(
463 IpcWatcherConfig::from_talon_config(&config),
464 db.clone(),
465 ipc_host,
466 audit.clone(),
467 );
468 let ipc_task = tokio::spawn({
469 let cancel = cancel.clone();
470 async move { watcher.run(cancel).await }
471 });
472 let recovered = orch.recover_pending_messages();
473 if !recovered.is_empty() {
474 tracing::info!(lanes = recovered.len(), "Recovery: lanes enqueued");
475 }
476 let loop_task = tokio::spawn(orch.clone().message_loop(cancel.clone()));
477
478 let signal = wait_for_signal().await;
480 tracing::info!(signal, "Shutdown signal received");
481 tokio::spawn(async {
482 let second = wait_for_signal().await;
483 tracing::warn!(signal = second, "second signal: exiting immediately");
484 std::process::exit(130);
485 });
486
487 orch.stop_accepting();
489 cancel.cancel();
490 let report = queue.shutdown(shutdown_grace).await;
491 tracing::info!(
492 drained = report.drained,
493 still_active = ?report.still_active,
494 stopped = ?report.stopped_containers,
495 dropped_pending = report.dropped_pending,
496 "Group queue shut down"
497 );
498 for channel in &channels {
499 match tokio::time::timeout(CHANNEL_DISCONNECT_TIMEOUT, channel.disconnect()).await {
500 Ok(Ok(())) => {}
501 Ok(Err(e)) => {
502 tracing::warn!(channel = channel.name(), error = %e, "channel disconnect failed")
503 }
504 Err(_) => tracing::warn!(channel = channel.name(), "channel disconnect timed out"),
505 }
506 }
507 join_with_timeout("message-loop", loop_task).await;
508 join_with_timeout("scheduler", scheduler_task).await;
509 join_with_timeout("ipc-watcher", ipc_task).await;
510 if let Some(gateway) = &gateway {
511 gateway.shutdown().await;
512 }
513 audit.record(event(
514 "talon.stop",
515 "talon",
516 None,
517 "info",
518 serde_json::json!({ "signal": signal, "drained": report.drained }),
519 ));
520 log_audit_head(&db, "shutdown");
521 tracing::info!("Talon stopped");
522 Ok(())
523}
524
525#[cfg(test)]
526mod tests {
527 use std::path::Path;
528
529 use talon_core::env::EnvFile;
530 use talon_core::types::{ContainerConfig, ProviderKind, VllmConfig};
531
532 use super::*;
533
534 fn cfg(content: &str) -> TalonConfig {
535 TalonConfig::from_env(
536 Path::new("/tmp/talon-test"),
537 EnvFile::from_str_content(content),
538 )
539 }
540
541 fn group_with_model(model: &str, base: &str) -> RegisteredGroup {
542 RegisteredGroup {
543 name: model.into(),
544 folder: model.into(),
545 trigger: "@Talon".into(),
546 added_at: "2026-01-01T00:00:00.000Z".into(),
547 container_config: Some(ContainerConfig {
548 provider: Some(ProviderKind::Vllm),
549 vllm: Some(VllmConfig {
550 base_url: base.into(),
551 model: model.into(),
552 fast_model: Some(format!("{model}-fast")),
553 max_context_tokens: None,
554 max_output_tokens: None,
555 temperature: None,
556 top_p: None,
557 enable_thinking: None,
558 request_timeout_secs: None,
559 }),
560 ..ContainerConfig::default()
561 }),
562 requires_trigger: None,
563 is_main: false,
564 session_key: None,
565 ipc_key: None,
566 ephemeral_session: false,
567 trusted_session_commands: false,
568 }
569 }
570
571 #[test]
572 fn no_gateway_without_vllm() {
573 let c = cfg("");
574 if std::env::var("TALON_VLLM_BASE_URL").is_ok() {
575 return; }
577 assert!(gateway_config(&c, &HashMap::new(), "127.0.0.1:0".parse().unwrap()).is_none());
578 }
579
580 #[test]
581 fn gateway_allows_group_models_and_keeps_key_secret() {
582 let c = cfg(
583 "TALON_VLLM_BASE_URL=http://vllm:8000/v1\nTALON_VLLM_MODEL=qwen\n\
584 TALON_VLLM_FAST_MODEL=qwen-small\nTALON_VLLM_API_KEY=sekret-key-123\n",
585 );
586 if effective_vllm(&c.provider_defaults).map(|v| v.model.as_str()) != Some("qwen") {
587 return; }
589 let mut groups = HashMap::new();
590 groups.insert(
591 "g1".to_string(),
592 group_with_model("llama", "http://vllm:8000"),
593 );
594 groups.insert(
595 "g2".to_string(),
596 group_with_model("other", "http://elsewhere:8000"),
597 );
598 let bind: SocketAddr = "127.0.0.1:3999".parse().unwrap();
599 let gw = gateway_config(&c, &groups, bind).unwrap();
600 assert_eq!(gw.bind, bind);
601 assert_eq!(gw.default_model, "qwen");
602 for m in ["qwen", "qwen-small", "llama", "llama-fast", "other"] {
603 assert!(gw.allowed_models.iter().any(|a| a == m), "{m} not allowed");
604 }
605 assert_eq!(gw.upstream_api_key.as_deref(), Some("sekret-key-123"));
606 assert!(!format!("{gw:?}").contains("sekret-key-123"));
607 }
608
609 #[test]
610 fn api_key_is_required_unless_explicitly_waived() {
611 if std::env::var(ALLOW_NO_API_KEY_ENV).is_ok() {
612 return; }
614 assert_eq!(
615 api_key_policy(&cfg("TALON_VLLM_API_KEY=abc\n")),
616 ApiKeyPolicy::Present
617 );
618 assert_eq!(api_key_policy(&cfg("")), ApiKeyPolicy::Missing);
619 assert_eq!(
620 api_key_policy(&cfg("TALON_VLLM_API_KEY=\" \"\n")),
621 ApiKeyPolicy::Missing
622 );
623 assert_eq!(
624 api_key_policy(&cfg("TALON_VLLM_ALLOW_NO_API_KEY=false\n")),
625 ApiKeyPolicy::Missing
626 );
627 assert_eq!(
628 api_key_policy(&cfg("TALON_VLLM_ALLOW_NO_API_KEY=true\n")),
629 ApiKeyPolicy::MissingAllowed
630 );
631 assert_eq!(
633 api_key_policy(&cfg(
634 "TALON_VLLM_API_KEY=k\nTALON_VLLM_ALLOW_NO_API_KEY=true\n"
635 )),
636 ApiKeyPolicy::Present
637 );
638 }
639
640 #[test]
641 fn gateway_concurrency_and_cache_salt_settings() {
642 let base = "TALON_VLLM_BASE_URL=http://vllm:8000/v1\nTALON_VLLM_MODEL=qwen\n";
643 if [
644 "TALON_VLLM_BASE_URL",
645 "TALON_VLLM_MODEL",
646 "TALON_LLM_GATEWAY_MAX_CONCURRENT",
647 "TALON_LLM_GATEWAY_CACHE_SALT",
648 ]
649 .iter()
650 .any(|k| std::env::var(k).is_ok())
651 {
652 return; }
654 let bind: SocketAddr = "127.0.0.1:3999".parse().unwrap();
655 let defaults = gateway_config(&cfg(base), &HashMap::new(), bind).unwrap();
656 assert!(defaults.cache_salt);
657 let gw = gateway_config(
658 &cfg(&format!(
659 "{base}TALON_LLM_GATEWAY_MAX_CONCURRENT=7\nTALON_LLM_GATEWAY_CACHE_SALT=false\n"
660 )),
661 &HashMap::new(),
662 bind,
663 )
664 .unwrap();
665 assert_eq!(gw.max_concurrent_requests, 7);
666 assert!(!gw.cache_salt);
667 let gw = gateway_config(
668 &cfg(&format!(
669 "{base}TALON_LLM_GATEWAY_MAX_CONCURRENT=lots\nTALON_LLM_GATEWAY_CACHE_SALT=maybe\n"
670 )),
671 &HashMap::new(),
672 bind,
673 )
674 .unwrap();
675 assert_eq!(gw.max_concurrent_requests, defaults.max_concurrent_requests);
676 assert!(gw.cache_salt);
677 assert_eq!(parse_bool(" Off "), Some(false));
678 assert_eq!(parse_bool("YES"), Some(true));
679 }
680
681 #[test]
682 fn base_url_normalization() {
683 assert_eq!(normalize_base("http://V:8000/v1/"), "http://v:8000");
684 assert_eq!(normalize_base("http://v:8000"), "http://v:8000");
685 }
686}