1use 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#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
22pub struct ContextUsage {
23 pub usage_ratio: Option<f64>,
25 pub context_limit: Option<u32>,
27 pub input_tokens: u32,
29 #[serde(default)]
31 pub output_tokens: u32,
32 #[serde(default, skip_serializing_if = "Option::is_none")]
34 pub cache_read_tokens: Option<u32>,
35 #[serde(default, skip_serializing_if = "Option::is_none")]
37 pub cache_creation_tokens: Option<u32>,
38 #[serde(default, skip_serializing_if = "Option::is_none")]
40 pub reasoning_tokens: Option<u32>,
41 #[serde(default)]
43 pub total_input_tokens: u64,
44 #[serde(default)]
46 pub total_output_tokens: u64,
47 #[serde(default)]
49 pub total_cache_read_tokens: u64,
50 #[serde(default)]
52 pub total_cache_creation_tokens: u64,
53 #[serde(default)]
55 pub total_reasoning_tokens: u64,
56}
57
58impl ContextUsage {
59 pub fn total_tokens(&self) -> u64 {
61 self.total_input_tokens + self.total_output_tokens
62 }
63}
64
65#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
67#[notification(method = "_aether/context_compaction")]
68pub struct ContextCompactionParams {
69 pub active: bool,
70}
71
72#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, JsonRpcNotification)]
74#[notification(method = "_aether/context_cleared")]
75pub struct ContextClearedParams {}
76
77#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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#[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}