Skip to main content

okbot_core/
agent.rs

1//! The agent port.
2
3use std::path::PathBuf;
4use std::time::Duration;
5
6use serde::{Deserialize, Serialize};
7
8use crate::async_trait;
9use crate::caps::{AgentCaps, ToolPolicy};
10use crate::secret::Secret;
11
12/// An MCP server the agent may call during a run.
13#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct McpServer {
15    /// Server name as the agent sees it (`kb`, `bot`).
16    pub name: String,
17    /// Streamable HTTP endpoint.
18    pub url: String,
19    /// Name of the environment variable holding the bearer token. The token itself is in
20    /// [`AgentTurn::env`], never in a command line or config file.
21    pub bearer_token_env: String,
22}
23
24/// Progress while an agent works.
25#[derive(Debug, Clone, PartialEq, Eq)]
26#[non_exhaustive]
27pub enum AgentProgress {
28    /// A piece of the answer.
29    TextDelta(String),
30    /// A whole message (backends that do not stream pieces).
31    Text(String),
32    /// The agent started using a tool.
33    ToolUse(String),
34}
35
36/// Receives [`AgentProgress`]; called from the backend's output reader, so keep it quick.
37#[derive(Clone)]
38pub struct ProgressFn(pub std::sync::Arc<dyn Fn(AgentProgress) + Send + Sync>);
39
40impl ProgressFn {
41    /// Wraps a callback.
42    pub fn new(f: impl Fn(AgentProgress) + Send + Sync + 'static) -> Self {
43        Self(std::sync::Arc::new(f))
44    }
45
46    /// Reports progress.
47    pub fn report(&self, progress: AgentProgress) {
48        (self.0)(progress);
49    }
50}
51
52impl std::fmt::Debug for ProgressFn {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        f.write_str("ProgressFn")
55    }
56}
57
58/// A native session to continue.
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct Resume {
61    /// The backend that owns the session ([`AgentBackend::name`] of the one that answered).
62    pub backend: String,
63    /// Session id from [`AgentReply::session_id`].
64    pub session: String,
65    /// The new messages only (the session already holds the earlier ones).
66    pub prompt: String,
67}
68
69/// One agent turn.
70#[derive(Debug, Clone)]
71pub struct AgentTurn {
72    /// Stable instructions: persona, rules, knowledge catalog, user profile.
73    pub system: String,
74    /// The conversation transcript and the new message.
75    pub prompt: String,
76    /// Tools the agent may use.
77    pub tools: ToolPolicy,
78    /// MCP servers for this run.
79    pub mcp: Vec<McpServer>,
80    /// Extra environment variables for the agent process (run tokens).
81    pub env: Vec<(String, Secret<String>)>,
82    /// Empty working directory for the agent process.
83    pub workdir: PathBuf,
84    /// Upper bound for the run.
85    pub timeout: Duration,
86    /// A session to continue instead of sending [`AgentTurn::prompt`]. Backends that cannot
87    /// resume it (another backend's session, expired) use the full prompt instead.
88    pub resume: Option<Resume>,
89    /// Progress callback.
90    pub progress: Option<ProgressFn>,
91}
92
93/// Token accounting, when the backend reports it.
94#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
95pub struct TokenUsage {
96    /// Input tokens.
97    pub input: u64,
98    /// Output tokens.
99    pub output: u64,
100    /// Reasoning tokens.
101    pub reasoning: u64,
102    /// Tokens read from the prompt cache.
103    pub cache_read: u64,
104    /// Tokens written to the prompt cache.
105    pub cache_write: u64,
106}
107
108/// The result of a successful turn.
109#[derive(Debug, Clone, PartialEq, Eq)]
110pub struct AgentReply {
111    /// The answer in Markdown.
112    pub text: String,
113    /// Which backend answered (`claude`, `codex`).
114    pub backend: String,
115    /// Model, when known.
116    pub model: Option<String>,
117    /// Native session id, for resuming.
118    pub session_id: Option<String>,
119    /// Token usage, when reported.
120    pub usage: Option<TokenUsage>,
121    /// Wall-clock duration.
122    pub duration: Duration,
123}
124
125/// Why a turn failed.
126#[derive(Debug, thiserror::Error)]
127#[non_exhaustive]
128pub enum AgentError {
129    /// The backend cannot run at all (CLI missing, logged out).
130    #[error("agent unavailable: {0}")]
131    Unavailable(String),
132    /// The backend cannot satisfy the turn (for example the tool policy); try another backend.
133    #[error("agent cannot handle this turn: {0}")]
134    Unsupported(String),
135    /// The run took longer than the turn's timeout.
136    #[error("agent timed out after {0:?}")]
137    Timeout(Duration),
138    /// The run failed. `retryable` says whether another backend (or a later retry) may succeed.
139    #[error("agent failed: {message}")]
140    Failed {
141        /// What went wrong, safe to log (never shown to chat users).
142        message: String,
143        /// Whether trying again elsewhere makes sense.
144        retryable: bool,
145    },
146}
147
148impl AgentError {
149    /// Whether a fallback backend should be tried.
150    pub fn should_fall_back(&self) -> bool {
151        match self {
152            Self::Unavailable(_) | Self::Unsupported(_) | Self::Timeout(_) => true,
153            Self::Failed { retryable, .. } => *retryable,
154        }
155    }
156}
157
158/// Runs agent turns (design ยง4.1).
159///
160/// Dropping the future returned by [`AgentBackend::run`] must stop the work (and kill any child
161/// process): the runtime cancels turns that way.
162#[async_trait]
163pub trait AgentBackend: Send + Sync {
164    /// Short name for logs and audit (`claude`, `fallback(claude,codex)`).
165    fn name(&self) -> &str;
166
167    /// What this backend supports natively.
168    fn caps(&self) -> AgentCaps;
169
170    /// Runs one turn.
171    async fn run(&self, turn: AgentTurn) -> Result<AgentReply, AgentError>;
172}