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}