basis_tasks/

state.rs

//! Durable per-agent state: metadata, terminal records, and their bounds.
//!
//! Three files carry an agent's task state, each atomic-replace JSON:
//! `meta.json` (bookkeeping and the recorded spawn request — never
//! conversation content, which is mentra's store's alone), `inbox.json`
//! (messages, see [`super::inbox`](mod@super::inbox)), and `terminal.json` — written as the
//! executor's **last** act, whose existence is the completion signal. An agent
//! is resumable iff its terminal record does not exist; every crash before
//! that write resolves toward resumable.

use std::{
    io,
    time::{SystemTime, UNIX_EPOCH},
};

use basis::{Bound, Effort, RunUsage, SystemPrompt};

use serde::{Deserialize, Serialize};
use serde_json::Value;

use crate::approve::Approve;
use crate::data_dir::{AgentPaths, write_private_atomic};

pub(crate) const MAX_MESSAGES: usize = 16;
pub(crate) const MAX_MESSAGE: usize = 256 * 1024;
pub(crate) const MAX_PROMPT: usize = 256 * 1024;
pub(crate) const MAX_EVENT_BYTES: usize = 32 * 1024;
pub(crate) const MAX_RESULT_BYTES: usize = 1024 * 1024;
/// How many tasks one workspace may hold at once. `spawn` refuses past it
/// (archiving an old agent directory is the way out); the ownership-policy
/// walks in `policy` and `attach::cancel_tree` are bounded by it too, so
/// corrupt or cyclic metadata cannot loop forever.
pub const MAX_TASKS: usize = 1024;
/// Byte cap on one agent's `events.jsonl`, standing in for the daemon-era
/// journal cap: on overflow a final notice line is appended and recording
/// stops — a run is never failed for observability.
pub(crate) const MAX_EVENTS_BYTES: u64 = 32 * 1024 * 1024;

/// The recorded spawn request, so a later attach can execute what spawn
/// accepted. Credentials are intentionally absent: the attaching process owns
/// its own environment, exactly as any other host does.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub(crate) struct RunOptions {
    pub provider: Option<String>,
    pub base_url: Option<String>,
    pub model: Option<String>,
    pub no_shell: bool,
    /// The host's say over the system prompt, held as the type it is:
    /// `Replace` displaces the workspace's context, `Append` follows it.
    ///
    /// Written typed (`{"replace": …}` / `{"append": …}`). The pre-0.6
    /// record spelled it as two flat strings, and a task directory outlives
    /// the binary that minted it, so the compat reader below still accepts
    /// the flat `Replace` half and [`folded`](Self::folded) picks up the
    /// legacy append field on load.
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        deserialize_with = "system_prompt_compat"
    )]
    pub system_prompt: Option<SystemPrompt>,
    /// The pre-0.6 append half, read-only: folded into
    /// [`system_prompt`](Self::system_prompt) on load and never written back.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub append_system_prompt: Option<String>,
    /// Typed, and byte-for-byte the strings the old writer spelled
    /// (`"low"` … `"max"`), so an old record needs no shim.
    pub effort: Option<Effort>,
    /// Typed, same story: `"always"` / `"never"` / `"prompt"` on the wire,
    /// unchanged — and an invalid mode is now unrepresentable, so corruption
    /// fails at decode, named, instead of somewhere downstream.
    pub approve: Approve,
    pub deadline_ms: Option<u64>,
    pub tool_budget: Option<usize>,
    pub token_budget: Option<u64>,
}

impl RunOptions {
    /// Folds the pre-0.6 two-string system-prompt spelling into the typed
    /// field. Called by [`load_meta`], so everything downstream sees one
    /// spelling however old the record is.
    fn folded(self) -> Self {
        let (system_prompt, append) = (
            self.system_prompt.clone(),
            self.append_system_prompt.clone(),
        );
        Self {
            system_prompt: system_prompt.or_else(|| append.map(SystemPrompt::Append)),
            append_system_prompt: None,
            ..self
        }
    }
}

/// Accepts the typed spelling and the pre-0.6 flat string, which was always
/// the `Replace` half — clap refused both flags at once before either was
/// recorded.
fn system_prompt_compat<'de, D>(deserializer: D) -> Result<Option<SystemPrompt>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    #[derive(Deserialize)]
    #[serde(untagged)]
    enum Compat {
        Typed(SystemPrompt),
        Legacy(String),
    }

    Ok(
        Option::<Compat>::deserialize(deserializer)?.map(|value| match value {
            Compat::Typed(system_prompt) => system_prompt,
            Compat::Legacy(replace) => SystemPrompt::Replace(replace),
        }),
    )
}

/// Accepts what the typed writer spells — the same `deadline` /
/// `tool_budget` / `token_budget` strings the old one used — and reads
/// anything else, like the `unknown` the old writer could emit for a bound it
/// had no name for, as no bound rather than an unloadable record.
fn bound_compat<'de, D>(deserializer: D) -> Result<Option<Bound>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let name = Option::<String>::deserialize(deserializer)?;
    Ok(name.and_then(|name| Bound::from_type_tag(&name)))
}

/// A task whose own work has finished, but whose attached children may still
/// lack terminal records. Recorded in `meta.json` before the settle pass, so a
/// kill between a child's terminal and the parent's leaves the parent
/// resumable with the model work already done.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "state", rename_all = "snake_case")]
#[non_exhaustive]
pub enum Terminal {
    /// The task produced its final answer.
    Succeeded { result: String },
    /// The task stopped with a failure message.
    Failed { error: String },
    /// The task stopped because cancellation was requested.
    Cancelled,
}

/// A terminal payload exactly as written, with the basis-owned fields this
/// build recognizes alongside it.
#[derive(Debug, Clone, PartialEq)]
pub struct TerminalRecord {
    /// The exact terminal or message-terminal JSON read from disk.
    pub raw: Value,
    /// The typed terminal state, or `None` when a newer writer used a state
    /// this build does not know.
    pub terminal: Option<Terminal>,
    /// The bound that stopped the run, when one did.
    pub stopped_by: Option<Bound>,
    /// Usage reported by the completed task, when present.
    pub usage: Option<RunUsage>,
    /// Whether the result text was truncated before it was recorded.
    pub result_truncated: bool,
}

impl TerminalRecord {
    pub(crate) fn from_raw(raw: Value) -> Self {
        let terminal = serde_json::from_value(raw.clone()).ok();
        let stopped_by = raw
            .get("stopped_by")
            .and_then(Value::as_str)
            .and_then(Bound::from_type_tag);
        let usage = raw
            .get("usage")
            .cloned()
            .and_then(|usage| serde_json::from_value(usage).ok());
        let result_truncated = raw
            .get("result_truncated")
            .and_then(Value::as_bool)
            .unwrap_or(false);
        Self {
            raw,
            terminal,
            stopped_by,
            usage,
            result_truncated,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MessageState {
    /// Accepted but not yet presented to the task.
    Pending,
    /// Presented to the task in the current turn.
    InFlight,
    /// Presented to a completed turn.
    Delivered,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MessageReply {
    /// The reply text.
    pub result: String,
    /// Whether the reply text was truncated.
    #[serde(default, skip_serializing_if = "is_false")]
    pub result_truncated: bool,
    /// The bound that stopped the reply's turn, when one did.
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        deserialize_with = "bound_compat"
    )]
    pub stopped_by: Option<Bound>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct MessageRecord {
    /// The durable message id.
    pub id: String,
    /// The message text.
    pub body: String,
    /// Whether the bounded public view shortened [`Self::body`].
    #[serde(default, skip_serializing_if = "is_false")]
    pub body_truncated: bool,
    /// Delivery state.
    pub state: MessageState,
    /// Milliseconds since the Unix epoch when the message was accepted.
    pub created_ms: u64,
    /// The correlated reply, when the task produced one.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reply: Option<MessageReply>,
}

/// The exact `basis inbox` payload with its bounded typed messages alongside.
#[derive(Debug, Clone, PartialEq)]
pub struct InboxRecord {
    /// The exact JSON shape printed by `basis inbox --json`.
    pub raw: Value,
    /// Bounded message views, oldest first.
    pub messages: Vec<MessageRecord>,
}

#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub(crate) struct TaskMeta {
    pub id: String,
    pub parent: Option<String>,
    pub detached: bool,
    pub workspace: String,
    /// The mentra resume key; empty until the first attach prepares the run.
    pub agent_id: String,
    /// The conversation this task was told to pick up, when `spawn` was given
    /// `--continue` or `--session`. The first attach resumes it instead of
    /// minting one, which is what lets a new handle carry an old dialogue —
    /// and why continuing a settled task is a new task rather than a message
    /// to a closed inbox (ADR-0019).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub continues: Option<String>,
    /// How many assistant turns the conversation already held when this task
    /// first attached to it.
    ///
    /// Zero for a task that opened its own conversation. Nonzero only for a
    /// continued one, where the previous answers are on the transcript from
    /// the very first turn — and where the resume recovery in `attach` would
    /// otherwise read one of them as this task's own answer and settle
    /// `succeeded` without ever asking its prompt.
    #[serde(default)]
    pub answered_before: usize,
    pub prompt: String,
    pub options: RunOptions,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub pending_terminal: Option<Terminal>,
    #[serde(default, skip_serializing_if = "is_false")]
    pub result_truncated: bool,
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        deserialize_with = "bound_compat"
    )]
    pub stopped_by: Option<Bound>,
    /// What every turn this task has driven reported spending, summed.
    ///
    /// Kept in `meta.json` rather than recomputed from the event journal
    /// because a task outlives the process that drove it: two attaches, each
    /// running turns, both add to one tally, and the journal may have been
    /// capped ([`MAX_EVENTS_BYTES`]) long before anyone asks.
    #[serde(default)]
    pub usage: RunUsage,
    pub deadline_at_ms: Option<u64>,
    pub created_ms: u64,
    /// When this task last advanced, as its executor records it: the attach
    /// that opened its conversation, every turn it banked, and the completion
    /// it settled on. Each write to this record is a step the task took, which
    /// is what lets one field carry both meanings — a bookkeeping rewrite that
    /// is *not* a step would have to say so rather than borrow this one.
    ///
    /// Read by [`super::tasks`](mod@super::tasks) to answer "the conversation
    /// I was just having", the question `--continue` asks. Defaulted for that
    /// reason: a record that omits it must still load, because a task that
    /// vanished from `basis list` after an upgrade would be a worse answer
    /// than a task ordered by the wrong clock. The reader falls back to
    /// `created_ms`.
    #[serde(default)]
    pub updated_ms: u64,
}

impl TaskMeta {
    pub(crate) fn new(
        id: String,
        parent: Option<String>,
        detached: bool,
        workspace: String,
        prompt: String,
        options: RunOptions,
        deadline_at_ms: Option<u64>,
    ) -> Self {
        let now = now_ms();
        Self {
            id,
            parent,
            detached,
            workspace,
            agent_id: String::new(),
            continues: None,
            answered_before: 0,
            prompt,
            options,
            pending_terminal: None,
            result_truncated: false,
            stopped_by: None,
            usage: RunUsage::default(),
            deadline_at_ms,
            created_ms: now,
            updated_ms: now,
        }
    }

    /// Records the conversation this task continues, when it continues one.
    ///
    /// A field rather than a pre-set `agent_id`, because the two mean
    /// different things to the executor: `agent_id` says *this task has
    /// attached before*, and taking a continued conversation for a resumed
    /// one is how a task settles on someone else's answer.
    #[must_use]
    pub(crate) fn continuing(self, agent_id: Option<String>) -> Self {
        Self {
            continues: agent_id,
            ..self
        }
    }

    pub(crate) fn deadline_passed(&self) -> bool {
        self.deadline_at_ms
            .is_some_and(|deadline| deadline <= now_ms())
    }

    /// The terminal payload the recorded completion earns: exactly the shape
    /// the daemon's journal produced, minted once into `terminal.json`.
    pub(crate) fn terminal_payload(&self) -> Option<Value> {
        let payload = match self.pending_terminal.as_ref()? {
            Terminal::Succeeded { result } => {
                let mut terminal = serde_json::json!({"state": "succeeded", "result": result});
                if self.result_truncated {
                    terminal["result_truncated"] = Value::Bool(true);
                }
                with_stopped_by(terminal, self.stopped_by)
            }
            Terminal::Failed { error } => with_stopped_by(
                serde_json::json!({"state": "failed", "error": error}),
                self.stopped_by,
            ),
            Terminal::Cancelled => serde_json::json!({"state": "cancelled"}),
        };
        Some(with_usage(payload, self.usage))
    }
}

fn with_stopped_by(mut payload: Value, stopped_by: Option<Bound>) -> Value {
    if let Some(stopped_by) = stopped_by {
        payload["stopped_by"] = serde_json::json!(stopped_by);
    }
    payload
}

/// Adds what the task spent, when it spent anything.
///
/// A record that names no usage is a task whose turns reported none — a
/// cancellation that never reached the model is the ordinary case — and a
/// `usage` object full of zeros would claim a measurement nobody made. The
/// same rule the finish line follows (`Event::RunFinished::usage`), so the
/// stream and the record cannot disagree.
fn with_usage(mut payload: Value, usage: RunUsage) -> Value {
    if usage != RunUsage::default() {
        payload["usage"] = serde_json::json!(usage);
    }
    payload
}

fn is_false(value: &bool) -> bool {
    !*value
}

pub(crate) fn load_meta(paths: &AgentPaths) -> Result<TaskMeta, String> {
    let bytes = std::fs::read(paths.meta())
        .map_err(|error| format!("read task metadata for {}: {error}", paths.dir().display()))?;
    let meta: TaskMeta = serde_json::from_slice(&bytes).map_err(|error| {
        format!(
            "decode task metadata for {}: {error}",
            paths.dir().display()
        )
    })?;

    // One spelling downstream, however old the record: the pre-0.6 system
    // prompt strings fold into the typed field on the way in.
    Ok(TaskMeta {
        options: meta.options.folded(),
        ..meta
    })
}

pub(crate) fn save_meta(paths: &AgentPaths, meta: &TaskMeta) -> Result<(), String> {
    let bytes =
        serde_json::to_vec(meta).map_err(|error| format!("encode task metadata: {error}"))?;
    write_private_atomic(&paths.meta(), &bytes)
        .map_err(|error| format!("persist task metadata: {error}"))
}

/// Reads the terminal record. `None` is the resumable state.
pub(crate) fn read_terminal(paths: &AgentPaths) -> Result<Option<Value>, String> {
    match std::fs::read(paths.terminal()) {
        Ok(bytes) => serde_json::from_slice(&bytes)
            .map(Some)
            .map_err(|error| format!("decode terminal record: {error}")),
        Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(None),
        Err(error) => Err(format!("read terminal record: {error}")),
    }
}

pub(crate) fn write_terminal(paths: &AgentPaths, payload: &Value) -> Result<(), String> {
    let bytes =
        serde_json::to_vec(payload).map_err(|error| format!("encode terminal record: {error}"))?;
    write_private_atomic(&paths.terminal(), &bytes)
        .map_err(|error| format!("persist terminal record: {error}"))
}

pub(crate) fn cancel_requested(paths: &AgentPaths) -> bool {
    paths.cancel_marker().exists()
}

/// Writes the cancel marker. Existence is the signal; the content is
/// diagnostic only.
pub(crate) fn request_cancel(paths: &AgentPaths, by: Option<&str>) -> Result<(), String> {
    let content = serde_json::json!({"requested_ms": now_ms(), "by": by});
    write_private_atomic(&paths.cancel_marker(), content.to_string().as_bytes())
        .map_err(|error| format!("record cancel request: {error}"))
}

pub(crate) fn bounded_text(mut value: String, limit: usize) -> (String, bool) {
    if value.len() <= limit {
        return (value, false);
    }
    let mut end = limit;
    while !value.is_char_boundary(end) {
        end -= 1;
    }
    value.truncate(end);
    (value, true)
}

/// Milliseconds since the epoch, clamped rather than panicking on a clock
/// before it. The one clock every durable record here is stamped with —
/// `created_ms`, `updated_ms`, a cancel marker's `requested_ms`, a message's
/// `created_ms` — and what a host rendering one of those alongside its own
/// notion of "now" (`basis list`'s own age column, for one) should read
/// `now` from, so the two clocks cannot disagree about the epoch.
pub fn now_ms() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_millis()
        .min(u128::from(u64::MAX)) as u64
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::data_dir::DataDir;

    fn agent(dir: &tempfile::TempDir) -> AgentPaths {
        let data = DataDir::from_path(dir.path()).unwrap();
        let paths = data
            .agent_dir("0123456789abcdef/0123456789abcdef0123456789abcdef")
            .unwrap();
        std::fs::create_dir_all(paths.dir()).unwrap();
        paths
    }

    fn meta(paths: &AgentPaths) -> TaskMeta {
        TaskMeta::new(
            "0123456789abcdef/0123456789abcdef0123456789abcdef".to_string(),
            None,
            true,
            "/repo".to_string(),
            "prompt".to_string(),
            RunOptions::default(),
            None,
        )
        .tap(paths)
    }

    trait Tap {
        fn tap(self, paths: &AgentPaths) -> Self;
    }
    impl Tap for TaskMeta {
        fn tap(self, paths: &AgentPaths) -> Self {
            save_meta(paths, &self).unwrap();
            self
        }
    }

    #[test]
    fn metadata_round_trips_through_its_file() {
        let dir = tempfile::tempdir().unwrap();
        let paths = agent(&dir);
        let saved = meta(&paths);
        assert_eq!(load_meta(&paths).unwrap(), saved);
    }

    #[test]
    fn terminal_records_are_absent_until_written_then_repeatably_observable() {
        let dir = tempfile::tempdir().unwrap();
        let paths = agent(&dir);
        assert!(read_terminal(&paths).unwrap().is_none());

        let mut record = meta(&paths);
        record.pending_terminal = Some(Terminal::Succeeded {
            result: "done".to_string(),
        });
        let payload = record.terminal_payload().unwrap();
        write_terminal(&paths, &payload).unwrap();

        assert_eq!(read_terminal(&paths).unwrap().unwrap(), payload);
        assert_eq!(read_terminal(&paths).unwrap().unwrap(), payload);
        assert_eq!(
            payload,
            serde_json::json!({"state": "succeeded", "result": "done"})
        );
    }

    #[test]
    fn terminal_record_keeps_raw_json_beside_the_typed_view() {
        let raw = serde_json::json!({
            "state": "succeeded",
            "result": "done",
            "result_truncated": true,
            "stopped_by": "token_budget",
            "usage": {
                "input_tokens": 7,
                "output_tokens": 3,
                "cache_read_tokens": 0,
                "cache_creation_tokens": 0
            },
            "newer_field": "kept",
        });
        let record = TerminalRecord::from_raw(raw.clone());

        assert_eq!(record.raw, raw);
        assert_eq!(
            record.terminal,
            Some(Terminal::Succeeded {
                result: "done".to_string()
            })
        );
        assert_eq!(record.stopped_by, Some(Bound::TokenBudget));
        assert_eq!(record.usage.expect("typed usage").input_tokens, 7);
        assert!(record.result_truncated);

        let unknown = serde_json::json!({"state": "paused_by_newer_basis", "detail": 1});
        let record = TerminalRecord::from_raw(unknown.clone());
        assert_eq!(record.raw, unknown);
        assert_eq!(record.terminal, None);
    }

    #[test]
    fn terminal_payloads_carry_truncation_and_bound_metadata() {
        let dir = tempfile::tempdir().unwrap();
        let paths = agent(&dir);
        let mut record = meta(&paths);
        record.pending_terminal = Some(Terminal::Failed {
            error: "took too long".to_string(),
        });
        record.stopped_by = Some(Bound::Deadline);
        assert_eq!(
            record.terminal_payload().unwrap(),
            serde_json::json!({"state": "failed", "error": "took too long", "stopped_by": "deadline"})
        );

        record.pending_terminal = Some(Terminal::Cancelled);
        assert_eq!(
            record.terminal_payload().unwrap(),
            serde_json::json!({"state": "cancelled"}),
            "a cancelled terminal never reports a bound"
        );
    }

    /// The terminal record is what `basis wait --json` and `basis list --json`
    /// read, and it is the only place a settled task's cost survives: the
    /// event journal can be capped, and the process that spent the tokens is
    /// gone.
    #[test]
    fn a_settled_task_records_what_it_spent_only_when_it_spent_something() {
        let dir = tempfile::tempdir().unwrap();
        let paths = agent(&dir);
        let mut record = meta(&paths);
        record.pending_terminal = Some(Terminal::Succeeded {
            result: "done".to_string(),
        });

        assert_eq!(
            record.terminal_payload().unwrap(),
            serde_json::json!({"state": "succeeded", "result": "done"}),
            "a task whose turns reported nothing claims no measurement"
        );

        record.usage = RunUsage {
            input_tokens: 900,
            output_tokens: 100,
            cache_read_tokens: 7,
            cache_creation_tokens: 3,
            reasoning_tokens: 0,
            thoughts_tokens: 0,
            model_responses: 1,
        };
        let payload = record.terminal_payload().unwrap();
        assert_eq!(payload["usage"]["input_tokens"], 900);
        assert_eq!(payload["usage"]["output_tokens"], 100);
        assert_eq!(payload["usage"]["cache_read_tokens"], 7);
        assert_eq!(payload["usage"]["cache_creation_tokens"], 3);
    }

    #[test]
    fn cancel_markers_signal_by_existence() {
        let dir = tempfile::tempdir().unwrap();
        let paths = agent(&dir);
        assert!(!cancel_requested(&paths));
        request_cancel(&paths, Some("caller/id")).unwrap();
        assert!(cancel_requested(&paths));
    }

    #[test]
    fn bounded_text_never_splits_utf8() {
        let (value, truncated) = bounded_text("a界b".to_string(), 2);
        assert_eq!(value, "a");
        assert!(truncated);
    }
}