pub mod commands;
pub mod host;
mod ledger;
pub mod messages;
pub mod run;
pub mod tasks;
#[cfg(test)]
mod tests;
use std::collections::{BTreeMap, HashMap, HashSet};
use std::os::fd::{AsFd, BorrowedFd};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard, OnceLock, PoisonError, RwLock};
use std::time::Duration;
use talon_core::audit::SharedAuditSink;
use talon_core::channel::Channel;
use talon_core::config::TalonConfig;
use talon_core::paths::group_folder_path;
use talon_core::protocol::SessionUsage;
use talon_core::types::{ChatMetadata, NewMessage, RegisteredGroup};
use talon_core::{Error, Result};
use talon_db::Db;
use talon_scheduler::{GroupQueue, TaskScheduler};
use talon_security::redact::Redactor;
use tokio::sync::Notify;
use tokio::time::Instant;
use crate::agent::{ContainerLauncher, LlmAccess};
use crate::allowlist::AllowlistSource;
use crate::output_limits::OutputLimits;
use crate::safe_fs;
use talon_security::instructions::{AGENT_INSTRUCTIONS_FILE, InstructionFiles};
pub const LAST_TIMESTAMP_KEY: &str = "last_timestamp";
pub const LAST_AGENT_TIMESTAMP_KEY: &str = "last_agent_timestamp";
const AGENT_MD: &str = AGENT_INSTRUCTIONS_FILE;
const MAX_TEMPLATE_BYTES: u64 = 256 * 1024;
#[derive(Debug, Clone)]
pub struct Settings {
pub task_close_delay: Duration,
pub command_close_delay: Duration,
pub token_ttl_margin: Duration,
pub shutdown_grace: Duration,
pub output_limits: OutputLimits,
}
impl Default for Settings {
fn default() -> Self {
Self {
task_close_delay: Duration::from_secs(10),
command_close_delay: Duration::from_secs(10),
token_ttl_margin: Duration::from_secs(5 * 60),
shutdown_grace: Duration::from_secs(10),
output_limits: OutputLimits::default(),
}
}
}
pub struct OrchestratorDeps {
pub config: Arc<TalonConfig>,
pub db: Db,
pub queue: GroupQueue,
pub launcher: Arc<dyn ContainerLauncher>,
pub llm: Option<Arc<dyn LlmAccess>>,
pub allowlist: AllowlistSource,
pub redactor: Arc<Redactor>,
pub audit: SharedAuditSink,
pub settings: Settings,
}
#[derive(Debug, Default)]
pub(crate) struct State {
pub(crate) groups: HashMap<String, RegisteredGroup>,
pub(crate) persisted: HashSet<String>,
pub(crate) sessions: HashMap<String, String>,
pub(crate) last_timestamp: String,
pub(crate) last_agent_timestamp: BTreeMap<String, String>,
pub(crate) session_epochs: HashMap<String, u64>,
pub(crate) token_deadlines: HashMap<String, Instant>,
pub(crate) run_tokens: HashMap<String, (String, Duration)>,
pub(crate) ledgers: HashMap<String, Arc<Mutex<ledger::LaneLedger>>>,
pub(crate) committed: HashMap<String, String>,
}
pub struct Orchestrator {
pub(crate) config: Arc<TalonConfig>,
pub(crate) settings: Settings,
pub(crate) db: Db,
pub(crate) queue: GroupQueue,
pub(crate) launcher: Arc<dyn ContainerLauncher>,
pub(crate) llm: Option<Arc<dyn LlmAccess>>,
pub(crate) allowlist: AllowlistSource,
pub(crate) redactor: Arc<Redactor>,
pub(crate) audit: SharedAuditSink,
pub(crate) state: Mutex<State>,
channels: RwLock<Vec<Arc<dyn Channel>>>,
scheduler: OnceLock<Arc<TaskScheduler>>,
pub(crate) wake: Notify,
accepting: AtomicBool,
pub(crate) loop_running: AtomicBool,
}
impl std::fmt::Debug for Orchestrator {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Orchestrator")
.field("assistant", &self.config.assistant_name)
.field("queue", &self.queue)
.finish_non_exhaustive()
}
}
impl Orchestrator {
pub fn new(deps: OrchestratorDeps) -> Result<Arc<Self>> {
let state = load_state(&deps.db)?;
tracing::info!(group_count = state.groups.len(), "State loaded");
Ok(Arc::new(Self {
config: deps.config,
settings: deps.settings,
db: deps.db,
queue: deps.queue,
launcher: deps.launcher,
llm: deps.llm,
allowlist: deps.allowlist,
redactor: deps.redactor,
audit: deps.audit,
state: Mutex::new(state),
channels: RwLock::new(Vec::new()),
scheduler: OnceLock::new(),
wake: Notify::new(),
accepting: AtomicBool::new(true),
loop_running: AtomicBool::new(false),
}))
}
pub(crate) fn lock(&self) -> MutexGuard<'_, State> {
self.state.lock().unwrap_or_else(PoisonError::into_inner)
}
pub fn config(&self) -> &TalonConfig {
&self.config
}
pub fn queue(&self) -> &GroupQueue {
&self.queue
}
pub fn db(&self) -> &Db {
&self.db
}
pub fn set_channels(&self, channels: Vec<Arc<dyn Channel>>) {
*self
.channels
.write()
.unwrap_or_else(PoisonError::into_inner) = channels;
}
pub fn channels(&self) -> Vec<Arc<dyn Channel>> {
self.channels
.read()
.unwrap_or_else(PoisonError::into_inner)
.clone()
}
pub fn find_channel(&self, jid: &str) -> Option<Arc<dyn Channel>> {
let channels = self.channels.read().unwrap_or_else(PoisonError::into_inner);
talon_router::find_channel(&channels, jid).cloned()
}
pub fn set_scheduler(&self, scheduler: Arc<TaskScheduler>) {
if self.scheduler.set(scheduler).is_err() {
tracing::warn!("task scheduler already installed; ignoring");
}
}
pub fn scheduler(&self) -> Option<&Arc<TaskScheduler>> {
self.scheduler.get()
}
pub fn wake_scheduler(&self) {
if let Some(s) = self.scheduler.get() {
s.wake();
}
}
pub fn wake_message_loop(&self) {
self.wake.notify_one();
}
pub fn stop_accepting(&self) {
self.accepting.store(false, Ordering::SeqCst);
self.wake.notify_one();
}
pub fn is_accepting(&self) -> bool {
self.accepting.load(Ordering::SeqCst)
}
pub fn groups_snapshot(&self) -> HashMap<String, RegisteredGroup> {
self.lock().groups.clone()
}
pub fn group(&self, jid: &str) -> Option<RegisteredGroup> {
self.lock().groups.get(jid).cloned()
}
pub fn session_for_key(&self, key: &str) -> Option<String> {
self.lock().sessions.get(key).cloned()
}
pub(crate) fn save_state(&self) {
let (last, agent) = {
let st = self.lock();
let mut cursors = st.last_agent_timestamp.clone();
for (jid, committed) in &st.committed {
cursors.insert(jid.clone(), committed.clone());
}
(st.last_timestamp.clone(), serde_json::to_string(&cursors))
};
if let Err(e) = self.db.set_router_state(LAST_TIMESTAMP_KEY, &last) {
tracing::error!(error = %e, "failed to save last_timestamp");
}
match agent {
Ok(json) => {
if let Err(e) = self.db.set_router_state(LAST_AGENT_TIMESTAMP_KEY, &json) {
tracing::error!(error = %e, "failed to save last_agent_timestamp");
}
}
Err(e) => tracing::error!(error = %e, "failed to encode last_agent_timestamp"),
}
}
pub fn cursor(&self, jid: &str) -> String {
self.lock()
.last_agent_timestamp
.get(jid)
.cloned()
.unwrap_or_default()
}
pub(crate) fn set_cursor(&self, jid: &str, ts: &str) {
self.lock()
.last_agent_timestamp
.insert(jid.to_string(), ts.to_string());
self.save_state();
}
pub(crate) fn get_or_recover_cursor(&self, jid: &str) -> String {
let existing = self.cursor(jid);
if !existing.is_empty() {
return existing;
}
match self
.db
.get_last_bot_message_timestamp(jid, &self.config.assistant_name)
{
Ok(Some(bot_ts)) if !bot_ts.is_empty() => {
tracing::info!(chat_jid = jid, recovered_from = %bot_ts, "Recovered message cursor from last bot reply");
self.set_cursor(jid, &bot_ts);
bot_ts
}
Ok(_) => String::new(),
Err(e) => {
tracing::warn!(chat_jid = jid, error = %e, "cursor recovery failed");
String::new()
}
}
}
#[cfg(test)]
pub(crate) fn persist_session(
&self,
group: &RegisteredGroup,
session_id: &str,
fingerprint: &str,
usage: Option<&SessionUsage>,
) {
self.store_session(group, session_id, fingerprint, usage, None);
}
pub(crate) fn session_epoch(&self, key: &str) -> u64 {
self.lock().session_epochs.get(key).copied().unwrap_or(0)
}
pub(crate) fn persist_run_session(
&self,
group: &RegisteredGroup,
session_id: &str,
fingerprint: &str,
usage: Option<&SessionUsage>,
epoch: u64,
) {
self.store_session(group, session_id, fingerprint, usage, Some(epoch));
}
fn store_session(
&self,
group: &RegisteredGroup,
session_id: &str,
fingerprint: &str,
usage: Option<&SessionUsage>,
epoch: Option<u64>,
) {
if !group.should_persist_session() || session_id.is_empty() {
return;
}
let key = group.session_key().to_string();
{
let mut st = self.lock();
let current = st.session_epochs.get(&key).copied().unwrap_or(0);
if epoch.is_some_and(|e| e != current) {
tracing::info!(group = %group.name, "Session was cleared during the run; not keeping the run's session");
return;
}
st.sessions.insert(key.clone(), session_id.to_string());
}
if let Err(e) = self.db.set_session(&key, session_id, Some(fingerprint)) {
tracing::error!(group = %group.name, error = %e, "failed to persist session");
return;
}
let context = usage.map(SessionUsage::context_size).unwrap_or(0);
if context > 0
&& let Err(e) = self.db.record_session_usage(&key, context)
{
tracing::warn!(group = %group.name, error = %e, "Failed to record session usage");
}
}
pub(crate) fn clear_session_for(&self, group: &RegisteredGroup) {
let key = group.session_key().to_string();
let had = {
let mut st = self.lock();
*st.session_epochs.entry(key.clone()).or_default() += 1;
st.sessions.remove(&key).is_some()
};
if let Err(e) = self.db.delete_session(&key) {
tracing::error!(group = %group.name, error = %e, "failed to delete session");
}
if had {
self.audit.record(talon_core::audit::event(
"session.cleared",
&group.folder,
Some(&key),
"info",
serde_json::json!({ "group": group.name }),
));
}
}
pub(crate) fn store_message(&self, msg: &NewMessage) {
if self.db.store_message(msg).is_ok() {
return;
}
let meta = ChatMetadata {
chat_jid: msg.chat_jid.clone(),
timestamp: msg.timestamp.clone(),
name: None,
channel: None,
is_group: None,
};
let result = self
.db
.store_chat_metadata(&meta)
.and_then(|()| self.db.store_message(msg));
if let Err(e) = result {
tracing::error!(chat_jid = %msg.chat_jid, error = %e, "failed to store message");
}
}
pub(crate) fn register_group_inner(
&self,
jid: &str,
group: RegisteredGroup,
persist: bool,
create_template: bool,
) -> Result<()> {
let group_dir = match group_folder_path(&self.config.groups_dir, &group.folder) {
Ok(dir) => dir,
Err(e) => {
tracing::warn!(jid, folder = %group.folder, error = %e, "Rejecting group registration with invalid folder");
return Err(e);
}
};
for key in [group.session_key.as_deref(), group.ipc_key.as_deref()]
.into_iter()
.flatten()
{
talon_core::paths::assert_valid_group_folder(key)?;
}
let mut persisted = false;
if persist {
match self.db.set_registered_group(jid, &group) {
Ok(()) => persisted = true,
Err(Error::InvalidInput(reason)) => {
tracing::debug!(jid, %reason, "lane shares its folder; keeping it in memory");
}
Err(e) => {
tracing::error!(jid, error = %e, "failed to persist group registration");
return Err(e);
}
}
}
{
let mut st = self.lock();
if persisted {
st.persisted.insert(jid.to_string());
}
st.groups.insert(jid.to_string(), group.clone());
}
self.prepare_group_dir(&group_dir, &group, create_template);
tracing::info!(jid, name = %group.name, folder = %group.folder, persisted, "Group registered");
Ok(())
}
fn prepare_group_dir(&self, group_dir: &Path, group: &RegisteredGroup, create_template: bool) {
let groups_dir = &self.config.groups_dir;
let prepared = std::fs::create_dir_all(groups_dir)
.and_then(|()| safe_fs::open_root(groups_dir))
.and_then(|groups| {
let dir = safe_fs::ensure_dir_at(groups.as_fd(), &group.folder)?;
Ok((groups, dir))
});
let (groups, dir) = match prepared {
Ok(fds) => fds,
Err(e) => {
tracing::warn!(folder = %group.folder, dir = %group_dir.display(), error = %e, "cannot create group folder");
return;
}
};
if let Err(e) = safe_fs::ensure_dir_at(dir.as_fd(), "logs") {
tracing::warn!(folder = %group.folder, error = %e, "cannot create group logs folder");
}
if create_template {
self.install_instructions(groups.as_fd(), dir.as_fd(), group);
}
}
fn install_instructions(
&self,
groups: BorrowedFd<'_>,
dir: BorrowedFd<'_>,
group: &RegisteredGroup,
) {
match InstructionFiles::inspect_at(dir) {
Ok(files) if files.is_vacant() => {}
Ok(files) => {
if files.is_conflict() {
tracing::warn!(
folder = %group.folder,
fix = talon_security::instructions::CONFLICT_FIX,
"group folder has both AGENT.md and CLAUDE.md with content; both left untouched, runs there are refused"
);
}
return;
}
Err(e) => {
tracing::warn!(folder = %group.folder, error = %e, "cannot inspect the group instructions");
return;
}
}
let template_folder = if group.is_main { "main" } else { "global" };
let Some(mut content) = read_template(groups, template_folder) else {
return;
};
let name = &self.config.assistant_name;
for shipped in [
talon_core::trigger::DEFAULT_ASSISTANT_NAME,
talon_core::trigger::LEGACY_ASSISTANT_NAME,
] {
if name == shipped {
continue;
}
let heading = format!("# {shipped}");
content = content
.lines()
.map(|l| {
if l == heading {
format!("# {name}")
} else {
l.to_string()
}
})
.collect::<Vec<_>>()
.join("\n")
+ if content.ends_with('\n') { "\n" } else { "" };
content = content.replace(&format!("You are {shipped}"), &format!("You are {name}"));
}
match safe_fs::create_new_at(dir, AGENT_MD, content.as_bytes()) {
Ok(()) => tracing::info!(folder = %group.folder, "Created AGENT.md from template"),
Err(e) if e.kind() == std::io::ErrorKind::AlreadyExists => {
tracing::debug!(folder = %group.folder, "AGENT.md appeared concurrently; kept");
}
Err(e) => tracing::warn!(folder = %group.folder, error = %e, "cannot create AGENT.md"),
}
}
pub(crate) fn unregister_lane(&self, jid: &str) -> bool {
let mut st = self.lock();
!st.persisted.contains(jid) && st.groups.remove(jid).is_some()
}
pub async fn send_outbound(&self, jid: &str, raw: &str) -> Result<bool> {
let redacted = self.redactor.redact(raw);
let channels = self.channels();
talon_router::send_formatted(&channels, jid, &redacted).await
}
}
fn read_template(groups: BorrowedFd<'_>, template_folder: &str) -> Option<String> {
let folder = match safe_fs::open_dir_at(groups, template_folder) {
Ok(fd) => fd,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return None,
Err(e) => {
tracing::warn!(template = template_folder, error = %e, "unsafe instructions template folder; not copied");
return None;
}
};
let name = match InstructionFiles::inspect_at(folder.as_fd()).map(|f| f.source()) {
Ok(Ok(Some(name))) => name,
Ok(Ok(None)) => return None,
Ok(Err(found)) => {
tracing::warn!(template = template_folder, reason = %found, "unusable instructions template; not copied");
return None;
}
Err(e) => {
tracing::warn!(template = template_folder, error = %e, "cannot inspect the instructions template");
return None;
}
};
let file = match safe_fs::open_regular_at(folder.as_fd(), name, MAX_TEMPLATE_BYTES) {
Ok((file, _)) => file,
Err(safe_fs::OpenRefusal::Missing) => return None,
Err(safe_fs::OpenRefusal::Rejected(reason)) => {
tracing::warn!(template = template_folder, file = name, %reason, "unsafe instructions template; not copied");
return None;
}
};
match safe_fs::read_capped(file, MAX_TEMPLATE_BYTES).map(String::from_utf8) {
Ok(Ok(text)) => Some(text),
Ok(Err(_)) => {
tracing::warn!(
template = template_folder,
"instructions template is not UTF-8; not copied"
);
None
}
Err(e) => {
tracing::warn!(template = template_folder, error = %e, "cannot read the instructions template");
None
}
}
}
fn load_state(db: &Db) -> Result<State> {
let last_timestamp = db.get_router_state(LAST_TIMESTAMP_KEY)?.unwrap_or_default();
let last_agent_timestamp = match db.get_router_state(LAST_AGENT_TIMESTAMP_KEY)? {
None => BTreeMap::new(),
Some(raw) if raw.trim().is_empty() => BTreeMap::new(),
Some(raw) => match serde_json::from_str::<BTreeMap<String, String>>(&raw) {
Ok(map) => map,
Err(_) => {
tracing::warn!("Corrupted last_agent_timestamp in DB, resetting");
BTreeMap::new()
}
},
};
let sessions = db.get_all_sessions()?;
let groups = db.get_all_registered_groups()?;
let persisted = groups.keys().cloned().collect();
Ok(State {
groups,
persisted,
sessions,
last_timestamp,
last_agent_timestamp,
..State::default()
})
}