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        let session_epoch = self.session_epoch(group.session_key());
294
295        self.write_snapshots(group, chat_jid);
296
297        let ttl = self
298            .launcher
299            .hard_timeout(&merged)
300            .saturating_add(self.config.idle_timeout)
301            .saturating_add(self.settings.token_ttl_margin);
302        // Label = the lane's IPC key, not the folder: the gateway derives
303        // the per-label prefix-cache salt and usage stats from it, so lanes
304        // sharing a folder (CoPilot analysts) get separate caches (N-6).
305        let (token, endpoint) = llm.issue(group.ipc_key(), ttl, &vllm);
306        let input = ContainerInput {
307            prompt: prompt.to_string(),
308            session_id: session_id.clone(),
309            group_folder: group.folder.clone(),
310            chat_jid: chat_jid.to_string(),
311            is_main: group.is_main,
312            is_scheduled_task: opts.scheduled_task,
313            assistant_name: Some(self.config.assistant_name.clone()),
314            script: opts.script.clone(),
315            llm: Some(endpoint),
316            ephemeral: group.ephemeral_session,
317            timezone: Some(self.config.timezone_name.clone()),
318        };
319        let spec = RunSpec::new(group.clone(), input).with_config(merged);
320
321        if !opts.scheduled_task {
322            let mut st = self.lock();
323            st.token_deadlines
324                .insert(chat_jid.to_string(), Instant::now() + ttl);
325            st.run_tokens
326                .insert(chat_jid.to_string(), (token.clone(), ttl));
327        }
328        let slot: ContainerSlot = Arc::new(Mutex::new(None));
329        let hook: StartHook = {
330            let queue = self.queue.clone();
331            let jid = chat_jid.to_string();
332            let ipc_key = group.ipc_key().to_string();
333            let slot = slot.clone();
334            Box::new(move |container| {
335                *slot.lock().unwrap_or_else(PoisonError::into_inner) = Some(container.clone());
336                queue.register_process(&jid, container, &ipc_key);
337            })
338        };
339
340        let limits = self.settings.output_limits;
341        let (tx, rx) = mpsc::channel(OUTPUT_CHANNEL_CAPACITY);
342        let launch = self.launcher.launch(spec, tx, hook);
343        let consume = async {
344            let mut rx = rx;
345            let mut guard = OutputGuard::new(limits);
346            let mut last_result: Option<String> = None;
347            let mut aborted: Option<String> = None;
348            while let Some(output) = rx.recv().await {
349                let output = match guard.admit(output) {
350                    Admission::Deliver(output) => output,
351                    Admission::Abort(reason) => {
352                        aborted = Some(reason);
353                        break;
354                    }
355                };
356                if opts.manage_session {
357                    if let Some(id) = output.new_session_id.as_deref().filter(|s| !s.is_empty()) {
358                        self.persist_run_session(
359                            group,
360                            id,
361                            &fingerprint,
362                            output.usage.as_ref(),
363                            session_epoch,
364                        );
365                    }
366                }
367                if let Some(result) = output.result.as_ref().filter(|r| !r.is_empty()) {
368                    last_result = Some(result.clone());
369                }
370                sink.on_output(&output).await;
371            }
372            // Dropping the receiver makes every further send fail: the
373            // launcher discards the rest of an aborted run's frames.
374            drop(rx);
375            if let Some(reason) = &aborted {
376                self.abort_run(group, chat_jid, &slot, reason, &guard).await;
377                sink.on_output(&ContainerOutput::error(reason.clone()))
378                    .await;
379            }
380            (last_result, aborted)
381        };
382        let (mut result, (streamed_result, aborted)) = tokio::join!(launch, consume);
383        if let Some(reason) = aborted {
384            result.status = OutputStatus::Error;
385            result.error = Some(reason);
386        }
387        // The runner's own copies are as untrusted as the frames: the
388        // session id is validated again before it is persisted (N-5).
389        result.new_session_id = sanitize_session_id(result.new_session_id.take());
390        result.error = result
391            .error
392            .as_deref()
393            .map(|e| cap_text(e, limits.max_error_bytes));
394        result.last_result = result
395            .last_result
396            .as_deref()
397            .map(|r| cap_text(r, limits.max_result_bytes));
398
399        llm.revoke(&token);
400        if !opts.scheduled_task {
401            let mut st = self.lock();
402            st.token_deadlines.remove(chat_jid);
403            st.run_tokens.remove(chat_jid);
404        }
405
406        if opts.manage_session {
407            // Usually already persisted from the streamed frame; rewriting
408            // it would reset the usage recorded with it.
409            if let Some(id) = result
410                .new_session_id
411                .as_deref()
412                .filter(|id| self.session_for_key(group.session_key()).as_deref() != Some(*id))
413            {
414                self.persist_run_session(
415                    group,
416                    id,
417                    &fingerprint,
418                    result.usage.as_ref(),
419                    session_epoch,
420                );
421            }
422        }
423
424        if result.status == OutputStatus::Error {
425            let error = result
426                .error
427                .clone()
428                .unwrap_or_else(|| "Unknown error".to_string());
429            if opts.manage_session && session_id.is_some() && is_stale_session_error(&error) {
430                tracing::warn!(
431                    group = %group.name,
432                    error = %error,
433                    "Stale session detected — clearing for next retry"
434                );
435                self.clear_session_for(group);
436            }
437            tracing::error!(group = %group.name, error = %error, "Container agent error");
438        }
439
440        AgentOutcome {
441            status: result.status,
442            error: result.error.clone(),
443            last_result: streamed_result.or_else(|| result.last_result.clone()),
444            new_session_id: result.new_session_id.clone(),
445            usage: result.usage,
446            launched: !matches!(
447                result.outcome,
448                RunOutcome::SetupFailed | RunOutcome::SpawnFailed
449            ),
450        }
451    }
452}
453
454/// Which tasks a `current_tasks.json` snapshot loads from the database.
455#[derive(Debug, Clone, PartialEq, Eq)]
456pub enum TaskQuery {
457    /// Every task (main namespaces only).
458    All,
459    /// Only the tasks whose chat JID is one of these (a non-main
460    /// namespace's own lanes; empty = none, no database access).
461    ForJids(Vec<String>),
462}
463
464/// The snapshot query of a namespace with lanes `jids`: [`TaskQuery::All`]
465/// for main, otherwise a query filtered to `jids` (N-3) — never the whole
466/// table, so a lane with many tasks of its own cannot make every other
467/// lane's container start load them.
468pub fn snapshot_task_query(is_main: bool, jids: &[String]) -> TaskQuery {
469    if is_main {
470        TaskQuery::All
471    } else {
472        TaskQuery::ForJids(jids.to_vec())
473    }
474}
475
476/// The lanes of IPC namespace `ipc_key`: every registered JID whose
477/// `ipc_key()` equals it (sorted), and whether the namespace counts as main.
478/// It does only when every lane in it is main: lanes that disagree get the
479/// restricted view (the IPC watcher refuses such a namespace altogether).
480pub fn lane_namespace(
481    groups: &HashMap<String, RegisteredGroup>,
482    ipc_key: &str,
483) -> (Vec<String>, bool) {
484    let mut jids: Vec<String> = Vec::new();
485    let mut all_main = true;
486    for (jid, group) in groups {
487        if group.ipc_key() == ipc_key {
488            all_main &= group.is_main;
489            jids.push(jid.clone());
490        }
491    }
492    jids.sort();
493    let is_main = all_main && !jids.is_empty();
494    (jids, is_main)
495}
496
497#[cfg(test)]
498mod tests {
499    use super::*;
500
501    #[test]
502    fn stale_session_patterns_match_ts() {
503        assert!(is_stale_session_error(
504            "No conversation found with session ID abc"
505        ));
506        assert!(is_stale_session_error(
507            "ENOENT: no such file or directory, open '/home/agent/.talon/sessions/x.jsonl'"
508        ));
509        assert!(is_stale_session_error("session 1234 not found"));
510        assert!(is_stale_session_error("Session Not Found"));
511        assert!(!is_stale_session_error("Container exited with code 1"));
512        assert!(!is_stale_session_error("not found: model"));
513    }
514
515    fn lane(folder: &str, ipc_key: Option<&str>, is_main: bool) -> RegisteredGroup {
516        RegisteredGroup {
517            name: folder.into(),
518            folder: folder.into(),
519            trigger: "@Talon".into(),
520            added_at: "2026-01-01T00:00:00.000Z".into(),
521            container_config: None,
522            requires_trigger: None,
523            is_main,
524            session_key: None,
525            ipc_key: ipc_key.map(str::to_string),
526            ephemeral_session: false,
527            trusted_session_commands: false,
528        }
529    }
530
531    #[test]
532    fn lane_jids_are_the_registered_jids_of_one_namespace() {
533        let mut groups = HashMap::new();
534        groups.insert("tg:main".to_string(), lane("main", None, true));
535        groups.insert("http:copilot".to_string(), lane("copilot", None, false));
536        groups.insert(
537            "http:copilot:u:alice".to_string(),
538            lane("copilot", Some("copilot-u-alice"), false),
539        );
540        groups.insert(
541            "http:copilot:u:bob".to_string(),
542            lane("copilot", Some("copilot-u-bob"), false),
543        );
544        groups.insert(
545            "http:copilot:investigate".to_string(),
546            lane("copilot", Some("copilot-investigate"), false),
547        );
548        groups.insert(
549            "http:copilot:investigate:2".to_string(),
550            lane("copilot", Some("copilot-investigate"), false),
551        );
552
553        assert_eq!(
554            lane_namespace(&groups, "copilot-u-alice"),
555            (vec!["http:copilot:u:alice".to_string()], false)
556        );
557        // The shared folder does not merge lanes: only the base lane uses
558        // the folder as its key.
559        assert_eq!(
560            lane_namespace(&groups, "copilot"),
561            (vec!["http:copilot".to_string()], false)
562        );
563        assert_eq!(
564            lane_namespace(&groups, "copilot-investigate"),
565            (
566                vec![
567                    "http:copilot:investigate".to_string(),
568                    "http:copilot:investigate:2".to_string()
569                ],
570                false
571            )
572        );
573        assert_eq!(
574            lane_namespace(&groups, "main"),
575            (vec!["tg:main".to_string()], true)
576        );
577        assert_eq!(lane_namespace(&groups, "nobody"), (Vec::new(), false));
578
579        // A non-main lane sharing main's namespace demotes it.
580        groups.insert("tg:intruder".to_string(), lane("main", None, false));
581        let (jids, is_main) = lane_namespace(&groups, "main");
582        assert_eq!(jids.len(), 2);
583        assert!(!is_main);
584    }
585
586    #[test]
587    fn non_main_snapshots_query_only_their_own_lanes() {
588        let jids = vec!["http:copilot:u:alice".to_string()];
589        assert_eq!(snapshot_task_query(true, &jids), TaskQuery::All);
590        assert_eq!(
591            snapshot_task_query(false, &jids),
592            TaskQuery::ForJids(jids.clone())
593        );
594        assert_eq!(
595            snapshot_task_query(false, &[]),
596            TaskQuery::ForJids(Vec::new())
597        );
598    }
599
600    #[test]
601    fn run_options() {
602        let m = RunOptions::message();
603        assert!(m.resume_session && m.manage_session && !m.scheduled_task);
604        let t = RunOptions::task(Some("  ".into()), true);
605        assert!(t.scheduled_task && t.resume_session && !t.manage_session);
606        assert_eq!(t.script, None);
607    }
608}
talon/orchestrator/run.rs

talon/orchestrator/
run.rs
