1use std::path::PathBuf;
4
5use agent_client_protocol::schema::AuthMethod;
6use agent_client_protocol::{JsonRpcNotification, JsonRpcRequest, JsonRpcResponse};
7pub use mcp_utils::display_meta::{ToolDisplayMeta, ToolResultMeta};
8pub use rmcp::model::CreateElicitationRequestParams;
9use serde::{Deserialize, Serialize, de::DeserializeOwned};
10
11pub use mcp_utils::status::{McpServerAuthCapability, McpServerStatus, McpServerStatusEntry};
12
13pub const AETHER_META_NAMESPACE: &str = "contextbridge/aether";
14
15#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonRpcNotification)]
23#[notification(method = "_aether/context_usage")]
24pub struct ContextUsageParams {
25 pub usage_ratio: Option<f64>,
26 pub context_limit: Option<u32>,
27 pub input_tokens: u32,
28 #[serde(default)]
29 pub output_tokens: u32,
30 #[serde(default, skip_serializing_if = "Option::is_none")]
31 pub cache_read_tokens: Option<u32>,
32 #[serde(default, skip_serializing_if = "Option::is_none")]
33 pub cache_creation_tokens: Option<u32>,
34 #[serde(default, skip_serializing_if = "Option::is_none")]
35 pub reasoning_tokens: Option<u32>,
36 #[serde(default)]
37 pub total_input_tokens: u64,
38 #[serde(default)]
39 pub total_output_tokens: u64,
40 #[serde(default)]
41 pub total_cache_read_tokens: u64,
42 #[serde(default)]
43 pub total_cache_creation_tokens: u64,
44 #[serde(default)]
45 pub total_reasoning_tokens: u64,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default, JsonRpcNotification)]
50#[notification(method = "_aether/context_cleared")]
51pub struct ContextClearedParams {}
52
53#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
55#[notification(method = "_aether/auth_methods_updated")]
56pub struct AuthMethodsUpdatedParams {
57 pub auth_methods: Vec<AuthMethod>,
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonRpcRequest)]
66#[request(method = "_aether/elicitation", response = ElicitationResponse)]
67pub struct ElicitationParams {
68 pub server_name: String,
69 pub request: CreateElicitationRequestParams,
70}
71
72pub use rmcp::model::ElicitationAction;
73
74#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
76#[request(method = "_aether/prompt_search", response = PromptSearchResponse)]
77#[serde(rename_all = "camelCase")]
78pub struct PromptSearchParams {
79 pub query: String,
80 #[serde(default, skip_serializing_if = "Option::is_none")]
81 pub limit: Option<usize>,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
86#[serde(rename_all = "camelCase")]
87pub struct PromptSearchResponse {
88 pub query: String,
89 pub results: Vec<PromptSearchResult>,
90 pub truncated: bool,
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
98#[serde(rename_all = "camelCase")]
99pub struct PromptSearchResult {
100 pub session_id: String,
101 pub cwd: PathBuf,
102 pub session_created_at: String,
103 pub prompt: String,
104 pub match_start: usize,
105 pub match_end: usize,
106}
107
108#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
109#[request(method = "_aether/session_preview", response = SessionPreviewResponse)]
110#[serde(rename_all = "camelCase")]
111pub struct SessionPreviewParams {
112 pub session_id: String,
113}
114
115#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
116#[serde(rename_all = "camelCase")]
117pub struct SessionPreviewResponse {
118 pub session_id: String,
119 pub cwd: PathBuf,
120 pub created_at: String,
121 pub model: String,
122 #[serde(default, skip_serializing_if = "Option::is_none")]
123 pub selected_mode: Option<String>,
124 pub transcript: Vec<SessionPreviewTurn>,
125 pub tool_call_count: usize,
126 pub truncated: bool,
127}
128
129#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
130#[serde(rename_all = "camelCase")]
131pub struct SessionPreviewTurn {
132 pub role: SessionPreviewRole,
133 pub text: String,
134}
135
136#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
137#[serde(rename_all = "camelCase")]
138pub enum SessionPreviewRole {
139 User,
140 Assistant,
141}
142
143#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
145#[request(method = "_aether/workspace_list", response = WorkspaceListResponse)]
146#[serde(rename_all = "camelCase")]
147pub struct WorkspaceListParams {
148 pub session_id: String,
149}
150
151#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
154#[serde(rename_all = "camelCase")]
155pub struct WorkspaceListResponse {
156 pub workspaces: Vec<WorkspaceEntry>,
157}
158
159#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
160#[serde(rename_all = "camelCase")]
161pub struct WorkspaceEntry {
162 pub path: PathBuf,
163 pub is_current: bool,
164}
165
166#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcRequest)]
168#[request(method = "_aether/workspace_move", response = WorkspaceMoveResponse)]
169#[serde(rename_all = "camelCase")]
170pub struct WorkspaceMoveParams {
171 pub session_id: String,
172 pub target: WorkspaceMoveTarget,
173}
174
175#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
176#[serde(tag = "kind", rename_all = "camelCase")]
177pub enum WorkspaceMoveTarget {
178 Existing { path: PathBuf },
179 New { name: String },
180}
181
182#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcResponse)]
183#[serde(rename_all = "camelCase")]
184pub struct WorkspaceMoveResponse {
185 pub new_cwd: PathBuf,
186}
187
188#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
189#[serde(rename_all = "camelCase")]
190pub struct SessionDisplayMeta {
191 #[serde(default, skip_serializing_if = "Option::is_none")]
192 pub model: Option<String>,
193 #[serde(default, skip_serializing_if = "Option::is_none")]
194 pub selected_mode: Option<String>,
195}
196
197impl SessionDisplayMeta {
198 #[must_use]
199 pub fn new(model: impl Into<String>, selected_mode: Option<String>) -> Self {
200 Self { model: Some(model.into()), selected_mode }
201 }
202
203 #[must_use]
204 pub fn to_meta(&self) -> agent_client_protocol::schema::Meta {
205 to_aether_meta(self)
206 }
207
208 #[must_use]
209 pub fn from_meta(meta: Option<&agent_client_protocol::schema::Meta>) -> Self {
210 from_aether_meta(meta)
211 }
212}
213
214#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
215#[serde(rename_all = "camelCase")]
216pub struct AetherCapabilities {
217 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
218 pub prompt_search: bool,
219 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
220 pub session_preview: bool,
221 #[serde(default, skip_serializing_if = "std::ops::Not::not")]
222 pub workspace_move: bool,
223}
224
225impl AetherCapabilities {
226 #[must_use]
227 pub fn to_meta(self) -> agent_client_protocol::schema::Meta {
228 to_aether_meta(&self)
229 }
230
231 #[must_use]
232 pub fn from_meta(meta: Option<&agent_client_protocol::schema::Meta>) -> Self {
233 from_aether_meta(meta)
234 }
235}
236
237fn to_aether_meta<T: Serialize>(value: &T) -> agent_client_protocol::schema::Meta {
238 let mut meta = agent_client_protocol::schema::Meta::new();
239 meta.insert(AETHER_META_NAMESPACE.to_string(), serde_json::json!(value));
240 meta
241}
242
243fn from_aether_meta<T: DeserializeOwned + Default>(meta: Option<&agent_client_protocol::schema::Meta>) -> T {
244 meta.and_then(|m| m.get(AETHER_META_NAMESPACE))
245 .cloned()
246 .and_then(|value| serde_json::from_value(value).ok())
247 .unwrap_or_default()
248}
249
250#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, JsonRpcResponse)]
252pub struct ElicitationResponse {
253 pub action: ElicitationAction,
254 pub content: Option<serde_json::Value>,
256}
257
258pub use mcp_utils::client::UrlElicitationCompleteParams;
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 UrlElicitationComplete(UrlElicitationCompleteParams),
266}
267
268#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, JsonRpcNotification)]
270#[notification(method = "_aether/mcp_request")]
271pub enum McpRequest {
272 Authenticate { session_id: String, server_name: String },
273}
274
275#[derive(Debug, Clone, Serialize, Deserialize, JsonRpcNotification)]
279#[notification(method = "_aether/sub_agent_progress")]
280pub struct SubAgentProgressParams {
281 pub parent_tool_id: String,
282 pub task_id: String,
283 pub agent_name: String,
284 pub event: SubAgentEvent,
285}
286
287#[derive(Debug, Clone, Serialize, Deserialize)]
292pub enum SubAgentEvent {
293 ToolCall { request: SubAgentToolRequest },
294 ToolCallUpdate { update: SubAgentToolCallUpdate },
295 ToolResult { result: SubAgentToolResult },
296 ToolError { error: SubAgentToolError },
297 Done,
298 Other,
299}
300
301#[derive(Debug, Clone, Serialize, Deserialize)]
302pub struct SubAgentToolRequest {
303 pub id: String,
304 pub name: String,
305 pub arguments: String,
306}
307
308#[derive(Debug, Clone, Serialize, Deserialize)]
309pub struct SubAgentToolCallUpdate {
310 pub id: String,
311 pub chunk: String,
312}
313
314#[derive(Debug, Clone, Serialize, Deserialize)]
315pub struct SubAgentToolResult {
316 pub id: String,
317 pub name: String,
318 pub result_meta: Option<ToolResultMeta>,
319}
320
321#[derive(Debug, Clone, Serialize, Deserialize)]
322pub struct SubAgentToolError {
323 pub id: String,
324 pub name: String,
325}
326
327#[cfg(test)]
328mod tests {
329 use agent_client_protocol::JsonRpcMessage;
330 use agent_client_protocol::schema::AuthMethodAgent;
331
332 use super::*;
333
334 #[test]
335 fn wire_method_names_are_prefixed() {
336 assert_eq!(ContextClearedParams::default().method(), "_aether/context_cleared");
337 assert!(AuthMethodsUpdatedParams { auth_methods: vec![] }.method() == "_aether/auth_methods_updated");
338 assert!(McpNotification::ServerStatus { servers: vec![] }.method() == "_aether/mcp_event");
339 assert!(
340 McpRequest::Authenticate { session_id: String::new(), server_name: String::new() }.method()
341 == "_aether/mcp_request"
342 );
343 assert_eq!(PromptSearchParams { query: String::new(), limit: None }.method(), "_aether/prompt_search");
344 assert_eq!(SessionPreviewParams { session_id: String::new() }.method(), "_aether/session_preview");
345 assert_eq!(WorkspaceListParams { session_id: String::new() }.method(), "_aether/workspace_list");
346 let move_params =
347 WorkspaceMoveParams { session_id: String::new(), target: WorkspaceMoveTarget::New { name: String::new() } };
348 assert_eq!(move_params.method(), "_aether/workspace_move");
349 }
350
351 #[test]
352 fn context_usage_params_roundtrip() {
353 let params = ContextUsageParams {
354 usage_ratio: Some(0.75),
355 context_limit: Some(100_000),
356 input_tokens: 75_000,
357 output_tokens: 1_200,
358 cache_read_tokens: Some(40_000),
359 cache_creation_tokens: Some(2_000),
360 reasoning_tokens: Some(500),
361 total_input_tokens: 200_000,
362 total_output_tokens: 8_000,
363 total_cache_read_tokens: 90_000,
364 total_cache_creation_tokens: 5_000,
365 total_reasoning_tokens: 1_500,
366 };
367
368 let untyped = params.to_untyped_message().expect("serializable");
369 assert_eq!(untyped.method(), "_aether/context_usage");
370 let parsed = ContextUsageParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
371 assert_eq!(parsed, params);
372 }
373
374 #[test]
375 fn context_usage_params_omits_unset_optional_token_fields() {
376 let params = ContextUsageParams {
377 usage_ratio: Some(0.1),
378 context_limit: Some(1_000),
379 input_tokens: 100,
380 output_tokens: 0,
381 cache_read_tokens: None,
382 cache_creation_tokens: None,
383 reasoning_tokens: None,
384 total_input_tokens: 0,
385 total_output_tokens: 0,
386 total_cache_read_tokens: 0,
387 total_cache_creation_tokens: 0,
388 total_reasoning_tokens: 0,
389 };
390
391 let raw = serde_json::to_string(¶ms).unwrap();
392 assert!(!raw.contains("\"cache_read_tokens\""));
393 assert!(!raw.contains("\"cache_creation_tokens\""));
394 assert!(!raw.contains("\"reasoning_tokens\""));
395 }
396
397 #[test]
398 fn context_cleared_params_roundtrip() {
399 let params = ContextClearedParams::default();
400 let untyped = params.to_untyped_message().expect("serializable");
401 assert_eq!(untyped.method(), "_aether/context_cleared");
402 let parsed = ContextClearedParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
403 assert_eq!(parsed, params);
404 }
405
406 #[test]
407 fn auth_methods_updated_roundtrip() {
408 let params = AuthMethodsUpdatedParams {
409 auth_methods: vec![
410 AuthMethod::Agent(AuthMethodAgent::new("anthropic", "Anthropic").description("authenticated")),
411 AuthMethod::Agent(AuthMethodAgent::new("openrouter", "OpenRouter")),
412 ],
413 };
414
415 let untyped = params.to_untyped_message().expect("serializable");
416 assert_eq!(untyped.method(), "_aether/auth_methods_updated");
417 let parsed = AuthMethodsUpdatedParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
418 assert_eq!(parsed, params);
419 }
420
421 #[test]
422 fn mcp_request_authenticate_roundtrip() {
423 let msg = McpRequest::Authenticate {
424 session_id: "session-0".to_string(),
425 server_name: "my oauth server".to_string(),
426 };
427
428 let untyped = msg.to_untyped_message().expect("serializable");
429 assert_eq!(untyped.method(), "_aether/mcp_request");
430 let parsed = McpRequest::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
431 assert_eq!(parsed, msg);
432 }
433
434 #[test]
435 fn mcp_notification_server_status_roundtrip() {
436 let msg = McpNotification::ServerStatus {
437 servers: vec![
438 McpServerStatusEntry::new("github", McpServerStatus::Connected { tool_count: 5 }),
439 McpServerStatusEntry::new("linear", McpServerStatus::NeedsOAuth)
440 .with_auth_capability(McpServerAuthCapability::OAuth),
441 McpServerStatusEntry::new("slack", McpServerStatus::Failed { error: "connection timeout".to_string() }),
442 ],
443 };
444
445 let untyped = msg.to_untyped_message().expect("serializable");
446 assert_eq!(untyped.method(), "_aether/mcp_event");
447 let parsed = McpNotification::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
448 assert_eq!(parsed, msg);
449 }
450
451 #[test]
452 fn mcp_notification_url_elicitation_complete_roundtrip() {
453 let msg = McpNotification::UrlElicitationComplete(UrlElicitationCompleteParams {
454 server_name: "github".to_string(),
455 elicitation_id: "el-456".to_string(),
456 });
457
458 let untyped = msg.to_untyped_message().expect("serializable");
459 let parsed = McpNotification::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
460 assert_eq!(parsed, msg);
461 }
462
463 #[test]
464 fn sub_agent_progress_params_roundtrip() {
465 let params = SubAgentProgressParams {
466 parent_tool_id: "call_123".to_string(),
467 task_id: "task_abc".to_string(),
468 agent_name: "explorer".to_string(),
469 event: SubAgentEvent::Done,
470 };
471
472 let untyped = params.to_untyped_message().expect("serializable");
473 assert_eq!(untyped.method(), "_aether/sub_agent_progress");
474 }
475
476 #[test]
477 fn elicitation_params_roundtrip() {
478 use rmcp::model::{ElicitationSchema, EnumSchema};
479
480 let params = ElicitationParams {
481 server_name: "github".to_string(),
482 request: CreateElicitationRequestParams::FormElicitationParams {
483 meta: None,
484 message: "Pick a color".to_string(),
485 requested_schema: ElicitationSchema::builder()
486 .required_enum_schema(
487 "color",
488 EnumSchema::builder(vec!["red".into(), "green".into(), "blue".into()]).untitled().build(),
489 )
490 .build()
491 .unwrap(),
492 },
493 };
494
495 let untyped = params.to_untyped_message().expect("serializable");
496 assert_eq!(untyped.method(), "_aether/elicitation");
497 let parsed = ElicitationParams::parse_message(untyped.method(), untyped.params()).expect("roundtrip");
498 assert_eq!(parsed, params);
499 }
500
501 #[test]
502 fn elicitation_params_url_variant_has_mode_field() {
503 let params = ElicitationParams {
504 server_name: "github".to_string(),
505 request: CreateElicitationRequestParams::UrlElicitationParams {
506 meta: None,
507 message: "Authorize GitHub".to_string(),
508 url: "https://github.com/login/oauth".to_string(),
509 elicitation_id: "el-123".to_string(),
510 },
511 };
512
513 let json = serde_json::to_string(¶ms).unwrap();
514 assert!(json.contains("\"mode\":\"url\""));
515 assert!(json.contains("\"server_name\":\"github\""));
516 }
517
518 #[test]
519 fn mcp_server_status_entry_serde_roundtrip() {
520 let entry = McpServerStatusEntry::new("test-server", McpServerStatus::Connected { tool_count: 3 })
521 .with_auth_capability(McpServerAuthCapability::OAuth);
522
523 let json = serde_json::to_string(&entry).unwrap();
524 assert!(json.contains("\"auth_capability\":\"OAuth\""));
525 assert!(json.contains("\"proxied\":false"));
526 let parsed: McpServerStatusEntry = serde_json::from_str(&json).unwrap();
527 assert_eq!(parsed, entry);
528 assert!(!parsed.proxied);
529 assert!(parsed.can_authenticate());
530 }
531
532 #[test]
533 fn mcp_server_status_entry_proxied_serde_roundtrip() {
534 let entry = McpServerStatusEntry::new("math", McpServerStatus::NeedsOAuth)
535 .with_auth_capability(McpServerAuthCapability::OAuth)
536 .with_proxied(true);
537
538 let json = serde_json::to_string(&entry).unwrap();
539 assert!(json.contains("\"proxied\":true"));
540 let parsed: McpServerStatusEntry = serde_json::from_str(&json).unwrap();
541 assert_eq!(parsed, entry);
542 }
543
544 #[test]
545 fn deserialize_tool_call_event() {
546 let json = r#"{"ToolCall":{"request":{"id":"c1","name":"grep","arguments":"{\"pattern\":\"test\"}"},"model_name":"m"}}"#;
547 let event: SubAgentEvent = serde_json::from_str(json).unwrap();
548 assert!(matches!(event, SubAgentEvent::ToolCall { .. }));
549 }
550
551 #[test]
552 fn deserialize_tool_call_update_event() {
553 let json = r#"{"ToolCallUpdate":{"update":{"id":"c1","chunk":"{\"pattern\":\"test\"}"},"model_name":"m"}}"#;
554 let event: SubAgentEvent = serde_json::from_str(json).unwrap();
555 assert!(matches!(event, SubAgentEvent::ToolCallUpdate { .. }));
556 }
557
558 #[test]
559 fn deserialize_tool_result_event() {
560 let json = r#"{"ToolResult":{"result":{"id":"c1","name":"grep","result_meta":{"display":{"title":"Grep","value":"'test' in src (3 matches)"}}}}}"#;
561 let event: SubAgentEvent = serde_json::from_str(json).unwrap();
562 match event {
563 SubAgentEvent::ToolResult { result } => {
564 let result_meta = result.result_meta.expect("expected result_meta");
565 assert_eq!(result_meta.display.title, "Grep");
566 }
567 other => panic!("Expected ToolResult, got {other:?}"),
568 }
569 }
570
571 #[test]
572 fn deserialize_tool_error_event() {
573 let json = r#"{"ToolError":{"error":{"id":"c1","name":"grep"}}}"#;
574 let event: SubAgentEvent = serde_json::from_str(json).unwrap();
575 assert!(matches!(event, SubAgentEvent::ToolError { .. }));
576 }
577
578 #[test]
579 fn deserialize_done_event() {
580 let event: SubAgentEvent = serde_json::from_str(r#""Done""#).unwrap();
581 assert!(matches!(event, SubAgentEvent::Done));
582 }
583
584 #[test]
585 fn deserialize_other_variant() {
586 let event: SubAgentEvent = serde_json::from_str(r#""Other""#).unwrap();
587 assert!(matches!(event, SubAgentEvent::Other));
588 }
589
590 #[test]
591 fn tool_result_meta_map_roundtrip() {
592 let meta: ToolResultMeta = ToolDisplayMeta::new("Read file", "Cargo.toml, 156 lines").into();
593 let map = meta.clone().into_map();
594 let parsed = ToolResultMeta::from_map(&map).expect("should deserialize ToolResultMeta");
595 assert_eq!(parsed, meta);
596 }
597}