Skip to main content

acp_utils/
notifications.rs

1//! Typed wire-format types for Aether's custom ACP extension requests and
2//! notifications.
3use std::path::PathBuf;
4
5use agent_client_protocol::schema::v1::{AuthMethod, Meta};
6use agent_client_protocol::{JsonRpcNotification, JsonRpcRequest, JsonRpcResponse};
7pub use mcp_utils::display_meta::{ToolDisplayMeta, ToolResultMeta};
8use serde::{Deserialize, Serialize, de::DeserializeOwned};
9
10pub use mcp_utils::status::{McpServerAuthCapability, McpServerStatus, McpServerStatusEntry};
11
12pub const AETHER_META_NAMESPACE: &str = "contextbridge/aether";
13
14/// Context/token usage reported after an LLM call.
15///
16/// Per-call fields (`input_tokens`, `output_tokens`, `cache_read_tokens`,
17/// `cache_creation_tokens`, `reasoning_tokens`) come from the most recent
18/// API response. The `total_*` fields are cumulative across the agent's
19/// lifetime. The optional fields are `None` when the provider doesn't
20/// expose that dimension; this is semantically distinct from `Some(0)`.
21#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
22pub struct ContextUsage {
23    /// Current usage ratio (0.0 - 1.0), if context window is known.
24    pub usage_ratio: Option<f64>,
25    /// Maximum context limit, if known.
26    pub context_limit: Option<u32>,
27    /// Input tokens on the most recent API call (the current context size).
28    pub input_tokens: u32,
29    /// Output tokens on the most recent API call.
30    #[serde(default)]
31    pub output_tokens: u32,
32    /// Prompt tokens served from cache on the most recent API call.
33    #[serde(default, skip_serializing_if = "Option::is_none")]
34    pub cache_read_tokens: Option<u32>,
35    /// Prompt tokens written to cache on the most recent API call.
36    #[serde(default, skip_serializing_if = "Option::is_none")]
37    pub cache_creation_tokens: Option<u32>,
38    /// Reasoning tokens spent on the most recent API call.
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub reasoning_tokens: Option<u32>,
41    /// Cumulative input tokens since the agent started.
42    #[serde(default)]
43    pub total_input_tokens: u64,
44    /// Cumulative output tokens since the agent started.
45    #[serde(default)]
46    pub total_output_tokens: u64,
47    /// Cumulative cache-read tokens since the agent started.
48    #[serde(default)]
49    pub total_cache_read_tokens: u64,
50    /// Cumulative cache-creation tokens since the agent started.
51    #[serde(default)]
52    pub total_cache_creation_tokens: u64,
53    /// Cumulative reasoning tokens since the agent started.
54    #[serde(default)]
55    pub total_reasoning_tokens: u64,
56}
57
58impl ContextUsage {
59    /// Sum of cumulative input + output tokens.
60    pub fn total_tokens(&self) -> u64 {
61        self.total_input_tokens + self.total_output_tokens
62    }
63}
64
65/// Parameters for `_aether/context_compaction` notifications.
66#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
67#[notification(method = "_aether/context_compaction")]
68pub struct ContextCompactionParams {
69    pub active: bool,
70}
71
72/// Parameters for `_aether/context_cleared` notifications.
73#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, JsonRpcNotification)]
74#[notification(method = "_aether/context_cleared")]
75pub struct ContextClearedParams {}
76
77/// Parameters for `_aether/auth_methods_updated` notifications.
78#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
79#[notification(method = "_aether/auth_methods_updated")]
80pub struct AuthMethodsUpdatedParams {
81    pub auth_methods: Vec<AuthMethod>,
82}
83
84/// Parameters for the `_aether/prompt_search` request.
85#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
86#[request(method = "_aether/prompt_search", response = PromptSearchResponse)]
87#[serde(rename_all = "camelCase")]
88pub struct PromptSearchParams {
89    pub query: String,
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub limit: Option<usize>,
92}
93
94/// Response for the `_aether/prompt_search` request.
95#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
96#[serde(rename_all = "camelCase")]
97pub struct PromptSearchResponse {
98    pub query: String,
99    pub results: Vec<PromptSearchResult>,
100    pub truncated: bool,
101}
102
103/// A single prompt-history search hit.
104///
105/// `match_start` and `match_end` are UTF-8 byte offsets into `prompt` and are
106/// guaranteed to fall on char boundaries.
107#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
108#[serde(rename_all = "camelCase")]
109pub struct PromptSearchResult {
110    pub session_id: String,
111    pub cwd: PathBuf,
112    pub session_created_at: String,
113    pub prompt: String,
114    pub match_start: usize,
115    pub match_end: usize,
116}
117
118#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
119#[request(method = "_aether/session_preview", response = SessionPreviewResponse)]
120#[serde(rename_all = "camelCase")]
121pub struct SessionPreviewParams {
122    pub session_id: String,
123}
124
125#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
126#[serde(rename_all = "camelCase")]
127pub struct SessionPreviewResponse {
128    pub session_id: String,
129    pub cwd: PathBuf,
130    pub created_at: String,
131    pub model: String,
132    #[serde(default, skip_serializing_if = "Option::is_none")]
133    pub selected_mode: Option<String>,
134    pub transcript: Vec<SessionPreviewTurn>,
135    pub tool_call_count: usize,
136    pub truncated: bool,
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
140#[serde(rename_all = "camelCase")]
141pub struct SessionPreviewTurn {
142    pub role: SessionPreviewRole,
143    pub text: String,
144}
145
146#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
147#[serde(rename_all = "camelCase")]
148pub enum SessionPreviewRole {
149    User,
150    Assistant,
151}
152
153/// Parameters for the `_aether/workspace_list` request.
154#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
155#[request(method = "_aether/workspace_list", response = WorkspaceListResponse)]
156#[serde(rename_all = "camelCase")]
157pub struct WorkspaceListParams {
158    pub session_id: String,
159}
160
161/// Response for the `_aether/workspace_list` request: every managed workspace
162/// originating from the same git repository as the session's working directory.
163#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
164#[serde(rename_all = "camelCase")]
165pub struct WorkspaceListResponse {
166    pub workspaces: Vec<WorkspaceEntry>,
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
170#[serde(rename_all = "camelCase")]
171pub struct WorkspaceEntry {
172    pub path: PathBuf,
173    pub is_current: bool,
174}
175
176/// Parameters for the `_aether/workspace_move` request.
177#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
178#[request(method = "_aether/workspace_move", response = WorkspaceMoveResponse)]
179#[serde(rename_all = "camelCase")]
180pub struct WorkspaceMoveParams {
181    pub session_id: String,
182    pub target: WorkspaceMoveTarget,
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
186#[serde(tag = "kind", rename_all = "camelCase")]
187pub enum WorkspaceMoveTarget {
188    Existing { path: PathBuf },
189    New { name: String },
190}
191
192#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
193#[serde(rename_all = "camelCase")]
194pub struct WorkspaceMoveResponse {
195    pub new_cwd: PathBuf,
196}
197
198#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
199#[serde(rename_all = "camelCase")]
200pub struct SessionDisplayMeta {
201    #[serde(default, skip_serializing_if = "Option::is_none")]
202    pub model: Option<String>,
203    #[serde(default, skip_serializing_if = "Option::is_none")]
204    pub selected_mode: Option<String>,
205}
206
207impl SessionDisplayMeta {
208    #[must_use]
209    pub fn new(model: impl Into<String>, selected_mode: Option<String>) -> Self {
210        Self { model: Some(model.into()), selected_mode }
211    }
212
213    #[must_use]
214    pub fn to_meta(&self) -> Meta {
215        to_aether_meta(self)
216    }
217
218    #[must_use]
219    pub fn from_meta(meta: Option<&Meta>) -> Self {
220        from_aether_meta(meta)
221    }
222}
223
224#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
225#[serde(rename_all = "camelCase")]
226pub struct AetherCapabilities {
227    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
228    pub prompt_search: bool,
229    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
230    pub session_preview: bool,
231    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
232    pub workspace_move: bool,
233}
234
235impl AetherCapabilities {
236    #[must_use]
237    pub fn to_meta(self) -> Meta {
238        to_aether_meta(&self)
239    }
240
241    #[must_use]
242    pub fn from_meta(meta: Option<&Meta>) -> Self {
243        from_aether_meta(meta)
244    }
245}
246
247fn to_aether_meta<T: Serialize>(value: &T) -> Meta {
248    let mut meta = Meta::new();
249    meta.insert(AETHER_META_NAMESPACE.to_string(), serde_json::json!(value));
250    meta
251}
252
253fn from_aether_meta<T: DeserializeOwned + Default>(meta: Option<&Meta>) -> T {
254    meta.and_then(|m| m.get(AETHER_META_NAMESPACE))
255        .cloned()
256        .and_then(|value| serde_json::from_value(value).ok())
257        .unwrap_or_default()
258}
259
260/// Server→client MCP extension notifications (relay → wisp).
261#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
262#[notification(method = "_aether/mcp_event")]
263pub enum McpNotification {
264    ServerStatus { servers: Vec<McpServerStatusEntry> },
265}
266
267/// Client→server MCP extension requests (wisp → relay).
268#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
269#[notification(method = "_aether/mcp_request")]
270pub enum McpRequest {
271    Authenticate { session_id: String, server_name: String },
272}
273
274/// Parameters for `_aether/sub_agent_progress` notifications.
275///
276/// This is the wire format sent from the ACP server (`aether-cli`) to clients like `wisp`.
277#[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)]
278#[notification(method = "_aether/sub_agent_progress")]
279pub struct SubAgentProgressParams {
280    pub parent_tool_id: String,
281    pub task_id: String,
282    pub agent_name: String,
283    pub event: SubAgentEvent,
284}
285
286/// Subset of agent message variants relevant for sub-agent status display.
287///
288/// The ACP server (`aether-cli`) converts `AgentEvent` to this type before
289/// serializing, so the wire format only contains these known variants.
290#[derive(Debug, Clone, Serialize, Deserialize)]
291pub enum SubAgentEvent {
292    ToolCall { request: SubAgentToolRequest },
293    ToolCallUpdate { update: SubAgentToolCallUpdate },
294    ToolResult { result: SubAgentToolResult },
295    ToolError { error: SubAgentToolError },
296    Done,
297    Other,
298}
299
300#[derive(Debug, Clone, Serialize, Deserialize)]
301pub struct SubAgentToolRequest {
302    pub id: String,
303    pub name: String,
304    pub arguments: String,
305}
306
307#[derive(Debug, Clone, Serialize, Deserialize)]
308pub struct SubAgentToolCallUpdate {
309    pub id: String,
310    pub chunk: String,
311}
312
313#[derive(Debug, Clone, Serialize, Deserialize)]
314pub struct SubAgentToolResult {
315    pub id: String,
316    pub name: String,
317    pub result_meta: Option<ToolResultMeta>,
318}
319
320#[derive(Debug, Clone, Serialize, Deserialize)]
321pub struct SubAgentToolError {
322    pub id: String,
323    pub name: String,
324}
325
326#[cfg(test)]
327mod tests {
328    use agent_client_protocol::JsonRpcMessage;
329    use agent_client_protocol::schema::v1::AuthMethodAgent;
330
331    use super::*;
332
333    #[test]
334    fn wire_method_names_are_prefixed() {
335        assert_eq!(ContextClearedParams::default().method(), "_aether/context_cleared");
336        assert_eq!(AuthMethodsUpdatedParams { auth_methods: vec![] }.method(), "_aether/auth_methods_updated");
337        assert_eq!(McpNotification::ServerStatus { servers: vec![] }.method(), "_aether/mcp_event");
338        assert_eq!(
339            McpRequest::Authenticate { session_id: String::new(), server_name: String::new() }.method(),
340            "_aether/mcp_request"
341        );
342        assert_eq!(PromptSearchParams { query: String::new(), limit: None }.method(), "_aether/prompt_search");
343        assert_eq!(SessionPreviewParams { session_id: String::new() }.method(), "_aether/session_preview");
344        assert_eq!(WorkspaceListParams { session_id: String::new() }.method(), "_aether/workspace_list");
345        let move_params =
346            WorkspaceMoveParams { session_id: String::new(), target: WorkspaceMoveTarget::New { name: String::new() } };
347        assert_eq!(move_params.method(), "_aether/workspace_move");
348    }
349
350    #[test]
351    fn context_compaction_params_roundtrip() {
352        for active in [true, false] {
353            let params = ContextCompactionParams { active };
354            let untyped = params.to_untyped_message().expect("serializable");
355            assert_eq!(untyped.method(), "_aether/context_compaction");
356            let parsed = ContextCompactionParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
357            assert_eq!(parsed, params);
358        }
359    }
360
361    #[test]
362    fn context_cleared_params_roundtrip() {
363        let params = ContextClearedParams::default();
364        let untyped = params.to_untyped_message().expect("serializable");
365        assert_eq!(untyped.method(), "_aether/context_cleared");
366        let parsed = ContextClearedParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
367        assert_eq!(parsed, params);
368    }
369
370    #[test]
371    fn auth_methods_updated_roundtrip() {
372        let params = AuthMethodsUpdatedParams {
373            auth_methods: vec![
374                AuthMethod::Agent(AuthMethodAgent::new("anthropic", "Anthropic").description("authenticated")),
375                AuthMethod::Agent(AuthMethodAgent::new("openrouter", "OpenRouter")),
376            ],
377        };
378
379        let untyped = params.to_untyped_message().expect("serializable");
380        assert_eq!(untyped.method(), "_aether/auth_methods_updated");
381        let parsed = AuthMethodsUpdatedParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
382        assert_eq!(parsed, params);
383    }
384
385    #[test]
386    fn mcp_request_authenticate_roundtrip() {
387        let msg = McpRequest::Authenticate {
388            session_id: "session-0".to_string(),
389            server_name: "my oauth server".to_string(),
390        };
391
392        let untyped = msg.to_untyped_message().expect("serializable");
393        assert_eq!(untyped.method(), "_aether/mcp_request");
394        let parsed = McpRequest::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
395        assert_eq!(parsed, msg);
396    }
397
398    #[test]
399    fn mcp_notification_server_status_roundtrip() {
400        let msg = McpNotification::ServerStatus {
401            servers: vec![
402                McpServerStatusEntry::new("github", McpServerStatus::Connected { tool_count: 5 }),
403                McpServerStatusEntry::new("linear", McpServerStatus::NeedsOAuth)
404                    .with_auth_capability(McpServerAuthCapability::OAuth),
405                McpServerStatusEntry::new("slack", McpServerStatus::Failed { error: "connection timeout".to_string() }),
406            ],
407        };
408
409        let untyped = msg.to_untyped_message().expect("serializable");
410        assert_eq!(untyped.method(), "_aether/mcp_event");
411        let parsed = McpNotification::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
412        assert_eq!(parsed, msg);
413    }
414
415    #[test]
416    fn sub_agent_progress_params_roundtrip() {
417        let params = SubAgentProgressParams {
418            parent_tool_id: "call_123".to_string(),
419            task_id: "task_abc".to_string(),
420            agent_name: "explorer".to_string(),
421            event: SubAgentEvent::Done,
422        };
423
424        let untyped = params.to_untyped_message().expect("serializable");
425        assert_eq!(untyped.method(), "_aether/sub_agent_progress");
426    }
427
428    #[test]
429    fn mcp_server_status_entry_serde_roundtrip() {
430        let entry = McpServerStatusEntry::new("test-server", McpServerStatus::Connected { tool_count: 3 })
431            .with_auth_capability(McpServerAuthCapability::OAuth);
432
433        let json = serde_json::to_string(&entry).unwrap();
434        assert!(json.contains("\"auth_capability\":\"OAuth\""));
435        assert!(json.contains("\"deferTools\":false"));
436        let parsed: McpServerStatusEntry = serde_json::from_str(&json).unwrap();
437        assert_eq!(parsed, entry);
438        assert!(!parsed.deferred_tools);
439        assert!(parsed.can_authenticate());
440    }
441
442    #[test]
443    fn mcp_server_status_entry_deferred_tools_serde_roundtrip() {
444        let entry = McpServerStatusEntry::new("math", McpServerStatus::NeedsOAuth)
445            .with_auth_capability(McpServerAuthCapability::OAuth)
446            .with_deferred_tools(true);
447
448        let json = serde_json::to_string(&entry).unwrap();
449        assert!(json.contains("\"deferTools\":true"));
450        let parsed: McpServerStatusEntry = serde_json::from_str(&json).unwrap();
451        assert_eq!(parsed, entry);
452    }
453
454    #[test]
455    fn deserialize_tool_call_event() {
456        let json = r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{\"pattern\":\"test\"}"},"model_name":"m"}}"#;
457        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
458        assert!(matches!(event, SubAgentEvent::ToolCall { .. }));
459    }
460
461    #[test]
462    fn deserialize_tool_call_update_event() {
463        let json = r#"{"ToolCallUpdate":{"update":{"id":"c1","chunk":"{\"pattern\":\"test\"}"},"model_name":"m"}}"#;
464        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
465        assert!(matches!(event, SubAgentEvent::ToolCallUpdate { .. }));
466    }
467
468    #[test]
469    fn deserialize_tool_result_event() {
470        let json = r#"{"ToolResult":{"result":{"id":"c1","name":"grep","result_meta":{"display":{"title":"Grep","value":"'test' in src (3 matches)"}}}}}"#;
471        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
472        match event {
473            SubAgentEvent::ToolResult { result } => {
474                let result_meta = result.result_meta.expect("expected result_meta");
475                assert_eq!(result_meta.display.title, "Grep");
476            }
477            other => panic!("Expected ToolResult, got {other:?}"),
478        }
479    }
480
481    #[test]
482    fn deserialize_tool_error_event() {
483        let json = r#"{"ToolError":{"error":{"id":"c1","name":"grep"}}}"#;
484        let event: SubAgentEvent = serde_json::from_str(json).unwrap();
485        assert!(matches!(event, SubAgentEvent::ToolError { .. }));
486    }
487
488    #[test]
489    fn deserialize_done_event() {
490        let event: SubAgentEvent = serde_json::from_str(r#""Done""#).unwrap();
491        assert!(matches!(event, SubAgentEvent::Done));
492    }
493
494    #[test]
495    fn deserialize_other_variant() {
496        let event: SubAgentEvent = serde_json::from_str(r#""Other""#).unwrap();
497        assert!(matches!(event, SubAgentEvent::Other));
498    }
499
500    #[test]
501    fn tool_result_meta_map_roundtrip() {
502        let meta: ToolResultMeta = ToolDisplayMeta::new("Read file", "Cargo.toml, 156 lines").into();
503        let map = meta.clone().into_map();
504        let parsed = ToolResultMeta::from_map(&map).expect("should deserialize ToolResultMeta");
505        assert_eq!(parsed, meta);
506    }
507}