talon/orchestrator/
run.rs

1//! One agent run (message runs and the container part of scheduled-task
2//! runs).
3//!
4//! 1. Effective provider: `.env` defaults merged with the group's config.
5//!    Without vLLM (or without the LLM gateway) the run fails before any
6//!    container starts, with an error that says what to configure.
7//! 2. Session: resume the lane's session unless the provider fingerprint
8//!    stored with it differs (a transcript produced by another model is
9//!    not resumed — the session is dropped instead).
10//! 3. Snapshots: `current_tasks.json` / `available_groups.json` in the
11//!    lane's IPC namespace. Tasks are filtered per lane (MYTH-03): a
12//!    non-main namespace only sees tasks whose chat JID belongs to one of
13//!    its own lanes, never those of other lanes sharing the group folder,
14//!    and only those are loaded from the database (N-3).
15//! 4. A per-run gateway token (TTL = the run's maximum lifetime + margin)
16//!    goes into `ContainerInput.llm` — stdin only — and is revoked when the
17//!    run ends. Its label is the lane's IPC key (per-lane prefix-cache salt
18//!    and usage accounting in the gateway; N-6).
19//! 5. Output frames stream to the caller's handler in order over a
20//!    bounded channel (backpressure); every frame passes the run's
21//!    [`OutputGuard`] first (result text capped per frame, run aborted —
22//!    container stopped, rest dropped — when it exceeds its per-run caps;
23//!    MYTH-07). Session ids are persisted as they arrive, and only when
24//!    they are valid (at most 128 bytes of `[A-Za-z0-9_-]`; N-5).
25//! 6. Errors that say the stored session is gone clear it so the queue's
26//!    retry starts fresh.
27
28use std::collections::HashMap;
29use std::sync::{Arc, LazyLock, Mutex, PoisonError};
30use std::time::Duration;
31
32use regex::Regex;
33use talon_container::{RunOutcome, RunSpec};
34use talon_core::protocol::{ContainerInput, ContainerOutput, OutputStatus};
35use talon_core::provider::{effective_vllm, merge_container_config, provider_fingerprint};
36use talon_core::types::{AvailableGroup, RegisteredGroup, ScheduledTask};
37use talon_router::AgentOutputHandler;
38use talon_scheduler::RunningContainer;
39use tokio::sync::mpsc;
40use tokio::time::Instant;
41
42use super::Orchestrator;
43use crate::agent::{AgentOutcome, StartHook};
44use crate::output_limits::{
45    Admission, OUTPUT_CHANNEL_CAPACITY, OutputGuard, cap_text, sanitize_session_id,
46};
47
48/// How long an aborted run's container gets to stop before the run is left
49/// to the runner's own timeout handling.
50const ABORT_STOP_TIMEOUT: Duration = Duration::from_secs(30);
51
52/// The container of the current run, once started (stopped on abort).
53type ContainerSlot = Arc<Mutex<Option<Arc<dyn RunningContainer>>>>;
54
55/// Error reported when no vLLM provider is configured.
56pub const VLLM_NOT_CONFIGURED: &str = "vLLM provider is not configured: set TALON_VLLM_BASE_URL and \
57     TALON_VLLM_MODEL in .env (run `talon doctor`)";
58/// Error reported when the LLM gateway is not running.
59pub const GATEWAY_NOT_RUNNING: &str =
60    "LLM gateway is not running, so containers cannot reach vLLM (run `talon doctor`)";
61
62/// Errors meaning the stored session no longer exists
63/// (`/no conversation found|ENOENT.*\.jsonl|session.*not found/i`).
64static STALE_SESSION: LazyLock<Regex> = LazyLock::new(|| {
65    // Static pattern: compiling it cannot fail.
66    #[allow(clippy::expect_used)]
67    Regex::new(r"(?i)no conversation found|ENOENT.*\.jsonl|session.*not found")
68        .expect("static stale-session regex")
69});
70
71/// Whether `error` says the resumed session is stale/corrupt.
72pub fn is_stale_session_error(error: &str) -> bool {
73    STALE_SESSION.is_match(error)
74}
75
76/// How a run treats sessions and what it tells the container.
77#[derive(Debug, Clone, Default, PartialEq, Eq)]
78pub struct RunOptions {
79    /// `ContainerInput.is_scheduled_task`.
80    pub scheduled_task: bool,
81    /// Pre-flight script of a scheduled task.
82    pub script: Option<String>,
83    /// Resume the lane's stored session (if any).
84    pub resume_session: bool,
85    /// Persist new session ids and clear stale / foreign-provider sessions
86    /// (message runs; scheduled tasks only borrow the session).
87    pub manage_session: bool,
88}
89
90impl RunOptions {
91    /// A chat (message-lane) run.
92    pub fn message() -> Self {
93        Self {
94            scheduled_task: false,
95            script: None,
96            resume_session: true,
97            manage_session: true,
98        }
99    }
100
101    /// A scheduled-task run.
102    pub fn task(script: Option<String>, resume_group_session: bool) -> Self {
103        Self {
104            scheduled_task: true,
105            script: script.filter(|s| !s.trim().is_empty()),
106            resume_session: resume_group_session,
107            manage_session: false,
108        }
109    }
110}
111
112impl Orchestrator {
113    /// Groups the agent may see (`available_groups.json`).
114    pub fn available_groups(&self) -> Vec<AvailableGroup> {
115        let groups = self.groups_snapshot();
116        match self.db.get_all_chats() {
117            Ok(chats) => talon_router::available_groups(&chats, &groups),
118            Err(e) => {
119                tracing::warn!(error = %e, "cannot list chats for available_groups");
120                Vec::new()
121            }
122        }
123    }
124
125    /// The tasks a namespace's snapshot shows, loaded with the narrowest
126    /// query (see [`snapshot_task_query`]). A database error yields an
127    /// empty list (logged).
128    pub(crate) fn snapshot_tasks(&self, query: &TaskQuery) -> Vec<ScheduledTask> {
129        let loaded = match query {
130            TaskQuery::All => self.db.get_all_tasks(),
131            TaskQuery::ForJids(jids) => self.db.get_tasks_for_jids(jids),
132        };
133        match loaded {
134            Ok(tasks) => tasks,
135            Err(e) => {
136                tracing::warn!(error = %e, "cannot list tasks for snapshot");
137                Vec::new()
138            }
139        }
140    }
141
142    /// Write `current_tasks.json` into IPC namespace `ipc_key`. A main
143    /// namespace sees every task; any other only the tasks of its own lanes
144    /// (`extra_jid`: the running lane, in case it is not registered), and
145    /// only those are loaded from the database (N-3: a container start of
146    /// a non-main lane never loads the whole task table).
147    pub(crate) fn write_lane_tasks_snapshot_for(
148        &self,
149        ipc_key: &str,
150        groups: &HashMap<String, RegisteredGroup>,
151        extra_jid: Option<&str>,
152    ) {
153        // Main-ness comes from the registry only, never from `extra_jid`.
154        let (mut jids, is_main) = lane_namespace(groups, ipc_key);
155        if let Some(jid) = extra_jid {
156            if !jids.iter().any(|j| j == jid) {
157                jids.push(jid.to_string());
158            }
159        }
160        let tasks = self.snapshot_tasks(&snapshot_task_query(is_main, &jids));
161        if let Err(e) = talon_container::write_lane_tasks_snapshot(
162            &self.config.data_dir,
163            ipc_key,
164            is_main,
165            &jids,
166            &tasks,
167        ) {
168            tracing::warn!(ipc_key, error = %e, "failed to write tasks snapshot");
169        }
170    }
171
172    /// Write both snapshots for a run of `group` in lane `chat_jid`.
173    fn write_snapshots(&self, group: &RegisteredGroup, chat_jid: &str) {
174        let groups = self.groups_snapshot();
175        let extra = (groups.get(chat_jid).map(RegisteredGroup::ipc_key) != Some(group.ipc_key()))
176            .then_some(chat_jid);
177        self.write_lane_tasks_snapshot_for(group.ipc_key(), &groups, extra);
178        let available = self.available_groups();
179        if let Err(e) = talon_container::write_groups_snapshot(
180            &self.config.data_dir,
181            &group.folder,
182            group.is_main,
183            &available,
184            Some(group.ipc_key()),
185        ) {
186            tracing::warn!(group = %group.name, error = %e, "failed to write groups snapshot");
187        }
188    }
189
190    /// Stop a run whose output exceeded its caps (MYTH-07): log, audit and
191    /// stop the container (bounded by [`ABORT_STOP_TIMEOUT`]).
192    async fn abort_run(
193        &self,
194        group: &RegisteredGroup,
195        chat_jid: &str,
196        slot: &ContainerSlot,
197        reason: &str,
198        guard: &OutputGuard,
199    ) {
200        tracing::error!(
201            group = %group.name,
202            chat_jid,
203            frames = guard.frames(),
204            result_bytes = guard.result_bytes(),
205            %reason,
206            "agent output limit exceeded; stopping the container"
207        );
208        self.audit.record(talon_core::audit::event(
209            "container.output_limit",
210            &group.folder,
211            Some(chat_jid),
212            "warn",
213            serde_json::json!({
214                "reason": reason,
215                "frames": guard.frames(),
216                "resultBytes": guard.result_bytes(),
217            }),
218        ));
219        let container = slot.lock().unwrap_or_else(PoisonError::into_inner).take();
220        match container {
221            Some(container) => {
222                match tokio::time::timeout(ABORT_STOP_TIMEOUT, container.stop()).await {
223                    Ok(Ok(())) => {}
224                    Ok(Err(e)) => {
225                        tracing::warn!(container = container.name(), error = %e, "cannot stop container")
226                    }
227                    Err(_) => {
228                        tracing::warn!(container = container.name(), "container stop timed out")
229                    }
230                }
231            }
232            None => tracing::warn!(chat_jid, "aborted run has no registered container"),
233        }
234    }
235
236    /// The session a run should resume (including the provider check).
237    fn resolve_session(
238        &self,
239        group: &RegisteredGroup,
240        fingerprint: &str,
241        opts: &RunOptions,
242    ) -> Option<String> {
243        if !opts.resume_session || !group.should_persist_session() {
244            return None;
245        }
246        let key = group.session_key();
247        let session_id = self.session_for_key(key)?;
248        match self.db.get_session_provider(key) {
249            Ok(Some(stored)) if stored != fingerprint => {
250                tracing::warn!(
251                    group = %group.name,
252                    session_provider = %stored,
253                    provider = %fingerprint,
254                    "Provider changed since session was created — starting a fresh session"
255                );
256                if opts.manage_session {
257                    self.clear_session_for(group);
258                }
259                None
260            }
261            Ok(_) => Some(session_id),
262            Err(e) => {
263                tracing::warn!(group = %group.name, error = %e, "cannot read session provider; starting fresh");
264                None
265            }
266        }
267    }
268
269    /// Run the agent for `group` in lane `chat_jid`, streaming every output
270    /// frame to `sink`. Never fails; see [`AgentOutcome`].
271    pub async fn run_agent(
272        &self,
273        group: &RegisteredGroup,
274        prompt: &str,
275        chat_jid: &str,
276        opts: RunOptions,
277        sink: &mut dyn AgentOutputHandler,
278    ) -> AgentOutcome {
279        let merged = merge_container_config(
280            &self.config.provider_defaults,
281            group.container_config.as_ref(),
282        );
283        let Some(vllm) = effective_vllm(&merged).cloned() else {
284            tracing::error!(group = %group.name, "{VLLM_NOT_CONFIGURED}");
285            return AgentOutcome::not_started(VLLM_NOT_CONFIGURED);
286        };
287        let Some(llm) = self.llm.clone() else {
288            tracing::error!(group = %group.name, "{GATEWAY_NOT_RUNNING}");
289            return AgentOutcome::not_started(GATEWAY_NOT_RUNNING);
290        };
291        let fingerprint = provider_fingerprint(&merged);
292        let session_id = self.resolve_session(group, &fingerprint, &opts);
293
294        self.write_snapshots(group, chat_jid);
295
296        let ttl = self
297            .launcher
298            .hard_timeout(&merged)
299            .saturating_add(self.config.idle_timeout)
300            .saturating_add(self.settings.token_ttl_margin);
301        // Label = the lane's IPC key, not the folder: the gateway derives
302        // the per-label prefix-cache salt and usage stats from it, so lanes
303        // sharing a folder (CoPilot analysts) get separate caches (N-6).
304        let (token, endpoint) = llm.issue(group.ipc_key(), ttl, &vllm);
305        let input = ContainerInput {
306            prompt: prompt.to_string(),
307            session_id: session_id.clone(),
308            group_folder: group.folder.clone(),
309            chat_jid: chat_jid.to_string(),
310            is_main: group.is_main,
311            is_scheduled_task: opts.scheduled_task,
312            assistant_name: Some(self.config.assistant_name.clone()),
313            script: opts.script.clone(),
314            llm: Some(endpoint),
315            ephemeral: group.ephemeral_session,
316            timezone: Some(self.config.timezone_name.clone()),
317        };
318        let spec = RunSpec::new(group.clone(), input).with_config(merged);
319
320        if !opts.scheduled_task {
321            let mut st = self.lock();
322            st.token_deadlines
323                .insert(chat_jid.to_string(), Instant::now() + ttl);
324            st.run_tokens
325                .insert(chat_jid.to_string(), (token.clone(), ttl));
326        }
327        let slot: ContainerSlot = Arc::new(Mutex::new(None));
328        let hook: StartHook = {
329            let queue = self.queue.clone();
330            let jid = chat_jid.to_string();
331            let ipc_key = group.ipc_key().to_string();
332            let slot = slot.clone();
333            Box::new(move |container| {
334                *slot.lock().unwrap_or_else(PoisonError::into_inner) = Some(container.clone());
335                queue.register_process(&jid, container, &ipc_key);
336            })
337        };
338
339        let limits = self.settings.output_limits;
340        let (tx, rx) = mpsc::channel(OUTPUT_CHANNEL_CAPACITY);
341        let launch = self.launcher.launch(spec, tx, hook);
342        let consume = async {
343            let mut rx = rx;
344            let mut guard = OutputGuard::new(limits);
345            let mut last_result: Option<String> = None;
346            let mut aborted: Option<String> = None;
347            while let Some(output) = rx.recv().await {
348                let output = match guard.admit(output) {
349                    Admission::Deliver(output) => output,
350                    Admission::Abort(reason) => {
351                        aborted = Some(reason);
352                        break;
353                    }
354                };
355                if opts.manage_session {
356                    if let Some(id) = output.new_session_id.as_deref().filter(|s| !s.is_empty()) {
357                        self.persist_session(group, id, &fingerprint, output.usage.as_ref());
358                    }
359                }
360                if let Some(result) = output.result.as_ref().filter(|r| !r.is_empty()) {
361                    last_result = Some(result.clone());
362                }
363                sink.on_output(&output).await;
364            }
365            // Dropping the receiver makes every further send fail: the
366            // launcher discards the rest of an aborted run's frames.
367            drop(rx);
368            if let Some(reason) = &aborted {
369                self.abort_run(group, chat_jid, &slot, reason, &guard).await;
370                sink.on_output(&ContainerOutput::error(reason.clone()))
371                    .await;
372            }
373            (last_result, aborted)
374        };
375        let (mut result, (streamed_result, aborted)) = tokio::join!(launch, consume);
376        if let Some(reason) = aborted {
377            result.status = OutputStatus::Error;
378            result.error = Some(reason);
379        }
380        // The runner's own copies are as untrusted as the frames: the
381        // session id is validated again before it is persisted (N-5).
382        result.new_session_id = sanitize_session_id(result.new_session_id.take());
383        result.error = result
384            .error
385            .as_deref()
386            .map(|e| cap_text(e, limits.max_error_bytes));
387        result.last_result = result
388            .last_result
389            .as_deref()
390            .map(|r| cap_text(r, limits.max_result_bytes));
391
392        llm.revoke(&token);
393        if !opts.scheduled_task {
394            let mut st = self.lock();
395            st.token_deadlines.remove(chat_jid);
396            st.run_tokens.remove(chat_jid);
397        }
398
399        if opts.manage_session {
400            // Usually already persisted from the streamed frame; rewriting
401            // it would reset the usage recorded with it.
402            if let Some(id) = result
403                .new_session_id
404                .as_deref()
405                .filter(|id| self.session_for_key(group.session_key()).as_deref() != Some(*id))
406            {
407                self.persist_session(group, id, &fingerprint, result.usage.as_ref());
408            }
409        }
410
411        if result.status == OutputStatus::Error {
412            let error = result
413                .error
414                .clone()
415                .unwrap_or_else(|| "Unknown error".to_string());
416            if opts.manage_session && session_id.is_some() && is_stale_session_error(&error) {
417                tracing::warn!(
418                    group = %group.name,
419                    error = %error,
420                    "Stale session detected — clearing for next retry"
421                );
422                self.clear_session_for(group);
423            }
424            tracing::error!(group = %group.name, error = %error, "Container agent error");
425        }
426
427        AgentOutcome {
428            status: result.status,
429            error: result.error.clone(),
430            last_result: streamed_result.or_else(|| result.last_result.clone()),
431            new_session_id: result.new_session_id.clone(),
432            usage: result.usage,
433            launched: !matches!(
434                result.outcome,
435                RunOutcome::SetupFailed | RunOutcome::SpawnFailed
436            ),
437        }
438    }
439}
440
441/// Which tasks a `current_tasks.json` snapshot loads from the database.
442#[derive(Debug, Clone, PartialEq, Eq)]
443pub enum TaskQuery {
444    /// Every task (main namespaces only).
445    All,
446    /// Only the tasks whose chat JID is one of these (a non-main
447    /// namespace's own lanes; empty = none, no database access).
448    ForJids(Vec<String>),
449}
450
451/// The snapshot query of a namespace with lanes `jids`: [`TaskQuery::All`]
452/// for main, otherwise a query filtered to `jids` (N-3) — never the whole
453/// table, so a lane with many tasks of its own cannot make every other
454/// lane's container start load them.
455pub fn snapshot_task_query(is_main: bool, jids: &[String]) -> TaskQuery {
456    if is_main {
457        TaskQuery::All
458    } else {
459        TaskQuery::ForJids(jids.to_vec())
460    }
461}
462
463/// The lanes of IPC namespace `ipc_key`: every registered JID whose
464/// `ipc_key()` equals it (sorted), and whether the namespace counts as main.
465/// It does only when every lane in it is main: lanes that disagree get the
466/// restricted view (the IPC watcher refuses such a namespace altogether).
467pub fn lane_namespace(
468    groups: &HashMap<String, RegisteredGroup>,
469    ipc_key: &str,
470) -> (Vec<String>, bool) {
471    let mut jids: Vec<String> = Vec::new();
472    let mut all_main = true;
473    for (jid, group) in groups {
474        if group.ipc_key() == ipc_key {
475            all_main &= group.is_main;
476            jids.push(jid.clone());
477        }
478    }
479    jids.sort();
480    let is_main = all_main && !jids.is_empty();
481    (jids, is_main)
482}
483
484#[cfg(test)]
485mod tests {
486    use super::*;
487
488    #[test]
489    fn stale_session_patterns_match_ts() {
490        assert!(is_stale_session_error(
491            "No conversation found with session ID abc"
492        ));
493        assert!(is_stale_session_error(
494            "ENOENT: no such file or directory, open '/home/agent/.talon/sessions/x.jsonl'"
495        ));
496        assert!(is_stale_session_error("session 1234 not found"));
497        assert!(is_stale_session_error("Session Not Found"));
498        assert!(!is_stale_session_error("Container exited with code 1"));
499        assert!(!is_stale_session_error("not found: model"));
500    }
501
502    fn lane(folder: &str, ipc_key: Option<&str>, is_main: bool) -> RegisteredGroup {
503        RegisteredGroup {
504            name: folder.into(),
505            folder: folder.into(),
506            trigger: "@Talon".into(),
507            added_at: "2026-01-01T00:00:00.000Z".into(),
508            container_config: None,
509            requires_trigger: None,
510            is_main,
511            session_key: None,
512            ipc_key: ipc_key.map(str::to_string),
513            ephemeral_session: false,
514            trusted_session_commands: false,
515        }
516    }
517
518    #[test]
519    fn lane_jids_are_the_registered_jids_of_one_namespace() {
520        let mut groups = HashMap::new();
521        groups.insert("tg:main".to_string(), lane("main", None, true));
522        groups.insert("http:copilot".to_string(), lane("copilot", None, false));
523        groups.insert(
524            "http:copilot:u:alice".to_string(),
525            lane("copilot", Some("copilot-u-alice"), false),
526        );
527        groups.insert(
528            "http:copilot:u:bob".to_string(),
529            lane("copilot", Some("copilot-u-bob"), false),
530        );
531        groups.insert(
532            "http:copilot:investigate".to_string(),
533            lane("copilot", Some("copilot-investigate"), false),
534        );
535        groups.insert(
536            "http:copilot:investigate:2".to_string(),
537            lane("copilot", Some("copilot-investigate"), false),
538        );
539
540        assert_eq!(
541            lane_namespace(&groups, "copilot-u-alice"),
542            (vec!["http:copilot:u:alice".to_string()], false)
543        );
544        // The shared folder does not merge lanes: only the base lane uses
545        // the folder as its key.
546        assert_eq!(
547            lane_namespace(&groups, "copilot"),
548            (vec!["http:copilot".to_string()], false)
549        );
550        assert_eq!(
551            lane_namespace(&groups, "copilot-investigate"),
552            (
553                vec![
554                    "http:copilot:investigate".to_string(),
555                    "http:copilot:investigate:2".to_string()
556                ],
557                false
558            )
559        );
560        assert_eq!(
561            lane_namespace(&groups, "main"),
562            (vec!["tg:main".to_string()], true)
563        );
564        assert_eq!(lane_namespace(&groups, "nobody"), (Vec::new(), false));
565
566        // A non-main lane sharing main's namespace demotes it.
567        groups.insert("tg:intruder".to_string(), lane("main", None, false));
568        let (jids, is_main) = lane_namespace(&groups, "main");
569        assert_eq!(jids.len(), 2);
570        assert!(!is_main);
571    }
572
573    #[test]
574    fn non_main_snapshots_query_only_their_own_lanes() {
575        let jids = vec!["http:copilot:u:alice".to_string()];
576        assert_eq!(snapshot_task_query(true, &jids), TaskQuery::All);
577        assert_eq!(
578            snapshot_task_query(false, &jids),
579            TaskQuery::ForJids(jids.clone())
580        );
581        assert_eq!(
582            snapshot_task_query(false, &[]),
583            TaskQuery::ForJids(Vec::new())
584        );
585    }
586
587    #[test]
588    fn run_options() {
589        let m = RunOptions::message();
590        assert!(m.resume_session && m.manage_session && !m.scheduled_task);
591        let t = RunOptions::task(Some("  ".into()), true);
592        assert!(t.scheduled_task && t.resume_session && !t.manage_session);
593        assert_eq!(t.script, None);
594    }
595}
talon/orchestrator/run.rs

talon/orchestrator/
run.rs
