okbot_peer/
a2a.rs

1//! The A2A subset okbot speaks (JSON-RPC binding): sending a message, getting and cancelling a
2//! task, and the Agent Card, plus the okbot delegation extension.
3//!
4//! okbot sends A2A 1.0 (`SendMessage`, `GetTask`, `CancelTask`, `TASK_STATE_*`, parts without
5//! `kind`, `A2A-Version: 1.0`) and answers 0.3 requests (`message/send`, `tasks/get`,
6//! `tasks/cancel`) in 0.3's shapes, so older A2A clients keep working. Only this module knows A2A's
7//! shapes; the rest of the crate uses okbot's types.
8
9use okbot_core::{DelegationContext, PeerTaskState};
10use serde::{Deserialize, Serialize};
11use serde_json::{Map, Value, json};
12
13/// The okbot delegation extension (required on every message between okbots).
14pub const EXTENSION: &str =
15    "https://github.com/tidusvn05/okbot/blob/main/docs/peers.md#delegation-v1";
16/// Where the delegation context sits in a message's metadata.
17pub const METADATA_KEY: &str = "okbot/delegation";
18/// Where a finished task's token use sits in its status message's metadata.
19pub const USAGE_KEY: &str = "okbot/usage";
20/// The header naming the A2A version of a request.
21pub const VERSION_HEADER: &str = "a2a-version";
22/// The header naming the extensions a request uses.
23pub const EXTENSIONS_HEADER: &str = "a2a-extensions";
24
25/// An A2A protocol version.
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
27pub enum Version {
28    /// 0.3.0: `message/send`, lowercase states, `kind` on parts.
29    V03,
30    /// 1.0: `SendMessage`, `TASK_STATE_*`, parts without `kind`.
31    V10,
32}
33
34impl Version {
35    /// The `A2A-Version` header value.
36    pub fn header(self) -> &'static str {
37        match self {
38            Self::V03 => "0.3",
39            Self::V10 => "1.0",
40        }
41    }
42
43    /// From an `A2A-Version` header or a configuration value; missing means 0.3 (the spec's rule).
44    pub fn parse(value: Option<&str>) -> Option<Self> {
45        match value.map(str::trim) {
46            None | Some("" | "0.3" | "0.3.0") => Some(Self::V03),
47            Some("1.0" | "1.0.0" | "1") => Some(Self::V10),
48            Some(_) => None,
49        }
50    }
51}
52
53/// An operation okbot offers.
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub enum Op {
56    /// Send a message (a new task, or an answer to a task waiting for input).
57    Send,
58    /// Get a task.
59    Get,
60    /// Cancel a task.
61    Cancel,
62}
63
64/// The JSON-RPC method of `op` in `version`.
65pub fn method(op: Op, version: Version) -> &'static str {
66    match (op, version) {
67        (Op::Send, Version::V03) => "message/send",
68        (Op::Get, Version::V03) => "tasks/get",
69        (Op::Cancel, Version::V03) => "tasks/cancel",
70        (Op::Send, Version::V10) => "SendMessage",
71        (Op::Get, Version::V10) => "GetTask",
72        (Op::Cancel, Version::V10) => "CancelTask",
73    }
74}
75
76/// The operation and version a JSON-RPC method names.
77pub fn op_of(method: &str) -> Option<(Op, Version)> {
78    Some(match method {
79        "message/send" => (Op::Send, Version::V03),
80        "tasks/get" => (Op::Get, Version::V03),
81        "tasks/cancel" => (Op::Cancel, Version::V03),
82        "SendMessage" => (Op::Send, Version::V10),
83        "GetTask" => (Op::Get, Version::V10),
84        "CancelTask" => (Op::Cancel, Version::V10),
85        _ => return None,
86    })
87}
88
89/// JSON-RPC error codes.
90pub mod code {
91    /// Invalid JSON.
92    pub const PARSE: i64 = -32700;
93    /// Not a valid request.
94    pub const INVALID_REQUEST: i64 = -32600;
95    /// Unknown method.
96    pub const METHOD_NOT_FOUND: i64 = -32601;
97    /// Invalid parameters.
98    pub const INVALID_PARAMS: i64 = -32602;
99    /// A2A: no such task.
100    pub const TASK_NOT_FOUND: i64 = -32001;
101    /// A2A: the task cannot be cancelled.
102    pub const NOT_CANCELABLE: i64 = -32002;
103    /// A2A: the version is not supported.
104    pub const VERSION_NOT_SUPPORTED: i64 = -32009;
105    /// okbot: the request is refused (links, depth, deadline, limits).
106    pub const REFUSED: i64 = -32010;
107}
108
109/// A JSON-RPC request.
110#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
111pub struct Request {
112    /// `2.0`.
113    pub jsonrpc: String,
114    /// Request id.
115    pub id: Value,
116    /// Method.
117    pub method: String,
118    /// Parameters.
119    #[serde(default)]
120    pub params: Value,
121}
122
123impl Request {
124    /// A request with a random id.
125    pub fn new(method: &str, params: Value) -> Self {
126        Self {
127            jsonrpc: "2.0".into(),
128            id: Value::String(hex::encode(rand::random::<[u8; 8]>())),
129            method: method.into(),
130            params,
131        }
132    }
133}
134
135/// A successful response.
136pub fn result(id: &Value, result: &Value) -> Value {
137    json!({"jsonrpc": "2.0", "id": id, "result": result})
138}
139
140/// An error response.
141pub fn error(id: &Value, code: i64, message: &str) -> Value {
142    json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}})
143}
144
145/// A message part.
146#[derive(Debug, Clone, PartialEq)]
147pub enum Part {
148    /// Text.
149    Text(String),
150    /// Structured data.
151    Data(Value),
152}
153
154impl Part {
155    fn to_wire(&self, version: Version) -> Value {
156        match (self, version) {
157            (Self::Text(text), Version::V03) => json!({"kind": "text", "text": text}),
158            (Self::Data(data), Version::V03) => json!({"kind": "data", "data": data}),
159            (Self::Text(text), Version::V10) => json!({"text": text}),
160            (Self::Data(data), Version::V10) => json!({"data": data}),
161        }
162    }
163
164    /// Either version: the member present decides (1.0), which 0.3's `kind` agrees with.
165    fn from_wire(value: &Value) -> Option<Self> {
166        if let Some(text) = value.get("text").and_then(Value::as_str) {
167            return Some(Self::Text(text.to_string()));
168        }
169        value.get("data").map(|d| Self::Data(d.clone()))
170    }
171}
172
173/// An A2A message.
174#[derive(Debug, Clone, PartialEq)]
175pub struct Message {
176    /// Unique per message (the idempotency key).
177    pub message_id: String,
178    /// The conversation.
179    pub context_id: Option<String>,
180    /// Continues this task (an answer to `input-required`).
181    pub task_id: Option<String>,
182    /// From the caller (`user`) rather than the agent.
183    pub from_user: bool,
184    /// Content.
185    pub parts: Vec<Part>,
186    /// Extension data.
187    pub metadata: Option<Value>,
188    /// Extensions this message uses.
189    pub extensions: Vec<String>,
190}
191
192impl Message {
193    /// A message from the caller.
194    pub fn user(message_id: &str, context_id: &str, parts: Vec<Part>) -> Self {
195        Self {
196            message_id: message_id.into(),
197            context_id: Some(context_id.into()),
198            task_id: None,
199            from_user: true,
200            parts,
201            metadata: None,
202            extensions: vec![EXTENSION.into()],
203        }
204    }
205
206    /// A message from the agent (the callee).
207    pub fn agent(parts: Vec<Part>) -> Self {
208        Self {
209            message_id: hex::encode(rand::random::<[u8; 8]>()),
210            context_id: None,
211            task_id: None,
212            from_user: false,
213            parts,
214            metadata: None,
215            extensions: Vec::new(),
216        }
217    }
218
219    /// The text parts, joined.
220    pub fn text(&self) -> String {
221        self.parts
222            .iter()
223            .filter_map(|p| match p {
224                Part::Text(text) => Some(text.as_str()),
225                Part::Data(_) => None,
226            })
227            .collect::<Vec<_>>()
228            .join("\n")
229    }
230
231    /// The first data part.
232    pub fn data(&self) -> Option<&Value> {
233        self.parts.iter().find_map(|p| match p {
234            Part::Data(data) => Some(data),
235            Part::Text(_) => None,
236        })
237    }
238
239    /// Attaches the delegation context.
240    pub fn with_context(mut self, context: &DelegationContext) -> Self {
241        self.metadata = Some(json!({ METADATA_KEY: context }));
242        self
243    }
244
245    /// The delegation context, when present and valid.
246    pub fn context(&self) -> Option<DelegationContext> {
247        let value = self.metadata.as_ref()?.get(METADATA_KEY)?;
248        serde_json::from_value(value.clone()).ok()
249    }
250
251    /// The message in `version`'s shape.
252    pub fn to_wire(&self, version: Version) -> Value {
253        let mut m = Map::new();
254        if version == Version::V03 {
255            m.insert("kind".into(), json!("message"));
256        }
257        let role = match (version, self.from_user) {
258            (Version::V03, true) => "user",
259            (Version::V03, false) => "agent",
260            (Version::V10, true) => "ROLE_USER",
261            (Version::V10, false) => "ROLE_AGENT",
262        };
263        m.insert("role".into(), json!(role));
264        m.insert("messageId".into(), json!(self.message_id));
265        if let Some(c) = &self.context_id {
266            m.insert("contextId".into(), json!(c));
267        }
268        if let Some(t) = &self.task_id {
269            m.insert("taskId".into(), json!(t));
270        }
271        m.insert(
272            "parts".into(),
273            Value::Array(self.parts.iter().map(|p| p.to_wire(version)).collect()),
274        );
275        if let Some(metadata) = &self.metadata {
276            m.insert("metadata".into(), metadata.clone());
277        }
278        if !self.extensions.is_empty() {
279            m.insert("extensions".into(), json!(self.extensions));
280        }
281        Value::Object(m)
282    }
283
284    /// A message in either version's shape.
285    pub fn from_wire(value: &Value) -> Option<Self> {
286        let text = |key: &str| value.get(key).and_then(Value::as_str).map(str::to_string);
287        Some(Self {
288            message_id: text("messageId")?,
289            context_id: text("contextId"),
290            task_id: text("taskId"),
291            from_user: matches!(text("role").as_deref(), Some("user" | "ROLE_USER")),
292            parts: value
293                .get("parts")?
294                .as_array()?
295                .iter()
296                .filter_map(Part::from_wire)
297                .collect(),
298            metadata: value.get("metadata").cloned(),
299            extensions: value
300                .get("extensions")
301                .and_then(Value::as_array)
302                .map(|e| {
303                    e.iter()
304                        .filter_map(Value::as_str)
305                        .map(str::to_string)
306                        .collect()
307                })
308                .unwrap_or_default(),
309        })
310    }
311}
312
313/// A task state in `version`'s spelling.
314pub fn state_name(state: PeerTaskState, version: Version) -> &'static str {
315    match version {
316        Version::V03 => state.as_str(),
317        Version::V10 => match state {
318            PeerTaskState::Submitted => "TASK_STATE_SUBMITTED",
319            PeerTaskState::Working => "TASK_STATE_WORKING",
320            PeerTaskState::InputRequired => "TASK_STATE_INPUT_REQUIRED",
321            PeerTaskState::Completed => "TASK_STATE_COMPLETED",
322            PeerTaskState::Failed => "TASK_STATE_FAILED",
323            PeerTaskState::Canceled => "TASK_STATE_CANCELED",
324            PeerTaskState::Rejected => "TASK_STATE_REJECTED",
325        },
326    }
327}
328
329/// A task state in either spelling; an unknown one (`auth-required`, …) counts as failed.
330pub fn parse_state(name: &str) -> PeerTaskState {
331    let name = match name.strip_prefix("TASK_STATE_") {
332        Some(rest) => rest.to_ascii_lowercase().replace('_', "-"),
333        None => name.to_string(),
334    };
335    name.parse().unwrap_or(PeerTaskState::Failed)
336}
337
338/// An A2A task.
339#[derive(Debug, Clone, PartialEq)]
340pub struct Task {
341    /// Task id.
342    pub id: String,
343    /// The conversation.
344    pub context_id: String,
345    /// Where it stands.
346    pub state: PeerTaskState,
347    /// The answer, the reason, the question, or progress.
348    pub message: Option<Message>,
349    /// When it changed (RFC 3339).
350    pub timestamp: Option<String>,
351}
352
353impl Task {
354    /// The task in `version`'s shape.
355    pub fn to_wire(&self, version: Version) -> Value {
356        let mut status = Map::new();
357        status.insert("state".into(), json!(state_name(self.state, version)));
358        if let Some(message) = &self.message {
359            status.insert("message".into(), message.to_wire(version));
360        }
361        if let Some(t) = &self.timestamp {
362            status.insert("timestamp".into(), json!(t));
363        }
364        let mut task = Map::new();
365        if version == Version::V03 {
366            task.insert("kind".into(), json!("task"));
367        }
368        task.insert("id".into(), json!(self.id));
369        task.insert("contextId".into(), json!(self.context_id));
370        task.insert("status".into(), Value::Object(status));
371        Value::Object(task)
372    }
373
374    /// A task in either version's shape.
375    pub fn from_wire(value: &Value) -> Option<Self> {
376        let status = value.get("status")?;
377        Some(Self {
378            id: value.get("id")?.as_str()?.to_string(),
379            context_id: value
380                .get("contextId")
381                .and_then(Value::as_str)
382                .unwrap_or_default()
383                .to_string(),
384            state: parse_state(status.get("state")?.as_str()?),
385            message: status.get("message").and_then(Message::from_wire),
386            timestamp: status
387                .get("timestamp")
388                .and_then(Value::as_str)
389                .map(str::to_string),
390        })
391    }
392}
393
394/// The result of a send: 1.0 wraps the task (`{"task": …}`), 0.3 returns it.
395pub fn send_result(task: &Task, version: Version) -> Value {
396    match version {
397        Version::V03 => task.to_wire(version),
398        Version::V10 => json!({"task": task.to_wire(version)}),
399    }
400}
401
402/// The task in a send result of either version.
403pub fn unwrap_send_result(result: &Value) -> &Value {
404    result.get("task").unwrap_or(result)
405}
406
407/// The Agent Card of a member for its team (A2A 1.0, with the 0.3 interface too).
408pub fn agent_card(name: &str, description: &str, url: &str, skills: &[String]) -> Value {
409    json!({
410        "name": name,
411        "description": description,
412        "version": env!("CARGO_PKG_VERSION"),
413        "supportedInterfaces": [
414            {"url": url, "protocolBinding": "JSONRPC", "protocolVersion": "1.0"},
415            {"url": url, "protocolBinding": "JSONRPC", "protocolVersion": "0.3"},
416        ],
417        "capabilities": {
418            "streaming": false,
419            "pushNotifications": false,
420            "extensions": [{"uri": EXTENSION, "required": true,
421                            "description": "okbot delegation context (origin, chain, deadline, budget, trace); requests are signed with OkbotPeer tokens"}],
422        },
423        "defaultInputModes": ["text/plain"],
424        "defaultOutputModes": ["text/plain"],
425        "skills": skills.iter().map(|s| json!({
426            "id": s, "name": s, "description": s, "tags": [s],
427        })).collect::<Vec<_>>(),
428    })
429}
430
431#[cfg(test)]
432mod tests {
433    use super::*;
434    use okbot_core::{Budget, Origin};
435
436    fn context() -> DelegationContext {
437        DelegationContext {
438            origin: Origin {
439                bot: "assistant".into(),
440                conversation: "zalo:shop|-|123|-".into(),
441                person: Some("lan".into()),
442                name: "Lan".into(),
443            },
444            chain: vec!["assistant".into()],
445            deadline: "2026-10-24T09:15:00Z".parse().unwrap(),
446            budget: Budget {
447                turns: 4,
448                tokens: None,
449            },
450            trace: "4f1c2a9e0b7d3e61".into(),
451        }
452    }
453
454    fn sample_task() -> Task {
455        Task {
456            id: "t-1".into(),
457            context_id: "c-1".into(),
458            state: PeerTaskState::InputRequired,
459            message: Some(Message {
460                message_id: "a-1".into(),
461                ..Message::agent(vec![
462                    Part::Text("Restart nginx?".into()),
463                    Part::Data(
464                        json!({"prompt": "p-1", "buttons": [{"id": "allow", "label": "Allow"}]}),
465                    ),
466                ])
467            }),
468            timestamp: Some("2026-10-24T09:00:00Z".into()),
469        }
470    }
471
472    #[test]
473    fn wire_format_in_both_versions() {
474        let message = Message::user(
475            "m-1",
476            "c-1",
477            vec![Part::Text("Find the release notes".into())],
478        )
479        .with_context(&context());
480        for version in [Version::V10, Version::V03] {
481            let mut request = Request::new(
482                method(Op::Send, version),
483                json!({"message": message.to_wire(version)}),
484            );
485            request.id = json!(1);
486            let task = sample_task();
487            let name = version.header().replace('.', "_");
488            let result = send_result(&task, version);
489            insta::with_settings!({sort_maps => true}, {
490                insta::assert_json_snapshot!(format!("send_{name}"), request);
491                insta::assert_json_snapshot!(format!("send_result_{name}"), result);
492            });
493            // Read back whatever the version.
494            let parsed = Message::from_wire(&request.params["message"]).unwrap();
495            assert_eq!(parsed, message);
496            assert_eq!(parsed.context(), Some(context()));
497            let back = Task::from_wire(unwrap_send_result(&result)).unwrap();
498            assert_eq!(back, task);
499        }
500    }
501
502    #[test]
503    fn methods_states_and_versions() {
504        for op in [Op::Send, Op::Get, Op::Cancel] {
505            for version in [Version::V03, Version::V10] {
506                assert_eq!(op_of(method(op, version)), Some((op, version)));
507            }
508        }
509        assert_eq!(op_of("tasks/resubscribe"), None);
510        for state in [
511            PeerTaskState::Submitted,
512            PeerTaskState::Working,
513            PeerTaskState::InputRequired,
514            PeerTaskState::Completed,
515            PeerTaskState::Failed,
516            PeerTaskState::Canceled,
517            PeerTaskState::Rejected,
518        ] {
519            for version in [Version::V03, Version::V10] {
520                assert_eq!(parse_state(state_name(state, version)), state);
521            }
522        }
523        assert_eq!(
524            parse_state("TASK_STATE_AUTH_REQUIRED"),
525            PeerTaskState::Failed
526        );
527        assert_eq!(
528            Version::parse(None),
529            Some(Version::V03),
530            "missing means 0.3"
531        );
532        assert_eq!(Version::parse(Some("1.0")), Some(Version::V10));
533        assert_eq!(Version::parse(Some("2.0")), None);
534    }
535
536    #[test]
537    fn agent_card() {
538        let mut card = super::agent_card(
539            "ops",
540            "Operations",
541            "https://ops.example.com/a2a/acme",
542            &["infrastructure".into()],
543        );
544        card["version"] = json!("[version]");
545        insta::with_settings!({sort_maps => true}, {
546            insta::assert_json_snapshot!("agent_card", card);
547        });
548    }
549
550    #[test]
551    fn a_message_without_context_has_none() {
552        let message = Message::user("m", "c", Vec::new());
553        assert_eq!(message.context(), None);
554        let mut bad = message;
555        bad.metadata = Some(json!({METADATA_KEY: {"chain": "x"}}));
556        assert_eq!(bad.context(), None);
557    }
558}
okbot_peer/a2a.rs

okbot_peer/
a2a.rs
