1use std::sync::Arc;
13
14use atd_protocol::{Response, ToolVisibility};
15use atd_runtime::capability::CapabilitySet;
16use atd_runtime::dispatch::ServerState;
17use atd_runtime::secrets::BearerIdentity;
18use atd_runtime::tracker::ReadTracker;
19use axum::Json;
20use axum::http::StatusCode;
21use axum::response::IntoResponse;
22use serde::{Deserialize, Serialize};
23use serde_json::{Value, json};
24
25#[derive(Debug, Deserialize)]
29pub struct JsonRpcRequest {
30 #[allow(dead_code)]
31 #[serde(default)]
32 pub jsonrpc: String,
33 #[serde(default)]
36 pub id: Option<Value>,
37 pub method: String,
38 #[serde(default)]
39 pub params: Value,
40}
41
42#[derive(Debug, Serialize)]
45pub struct JsonRpcResponse {
46 pub jsonrpc: &'static str,
47 #[serde(skip_serializing_if = "Option::is_none")]
48 pub id: Option<Value>,
49 #[serde(skip_serializing_if = "Option::is_none")]
50 pub result: Option<Value>,
51 #[serde(skip_serializing_if = "Option::is_none")]
52 pub error: Option<JsonRpcError>,
53}
54
55#[derive(Debug, Serialize)]
56pub struct JsonRpcError {
57 pub code: i32,
58 pub message: String,
59}
60
61pub fn error_response(
66 status: StatusCode,
67 id: Option<Value>,
68 code: i32,
69 message: impl Into<String>,
70) -> axum::response::Response {
71 error_response_with_headers(status, id, code, message, &[])
72}
73
74pub fn error_response_with_headers(
80 status: StatusCode,
81 id: Option<Value>,
82 code: i32,
83 message: impl Into<String>,
84 extra_headers: &[(&str, String)],
85) -> axum::response::Response {
86 let mut resp: axum::response::Response = (
87 status,
88 Json(JsonRpcResponse {
89 jsonrpc: "2.0",
90 id,
91 result: None,
92 error: Some(JsonRpcError {
93 code,
94 message: message.into(),
95 }),
96 }),
97 )
98 .into_response();
99
100 if !extra_headers.is_empty() {
101 let headers = resp.headers_mut();
102 for (name, value) in extra_headers {
103 if let (Ok(hn), Ok(hv)) = (
104 axum::http::HeaderName::from_bytes(name.as_bytes()),
105 axum::http::HeaderValue::from_str(value),
106 ) {
107 headers.insert(hn, hv);
108 }
109 }
110 }
111
112 resp
113}
114
115pub fn ok_response(id: Option<Value>, result: Value) -> axum::response::Response {
117 Json(JsonRpcResponse {
118 jsonrpc: "2.0",
119 id,
120 result: Some(result),
121 error: None,
122 })
123 .into_response()
124}
125
126pub fn handle_initialize(
135 id: Option<Value>,
136 server_version: &str,
137 token_broker: Option<&Arc<dyn atd_runtime::TokenBroker>>,
138) -> axum::response::Response {
139 let (name, version) = split_server_version(server_version);
140 let mut capabilities = json!({ "tools": {} });
141
142 if let Some(broker) = token_broker {
148 let formats = broker.accepted_token_formats();
149 if !formats.is_empty() {
150 let formats_value = serde_json::Value::Array(
151 formats
152 .iter()
153 .map(|f| serde_json::Value::String((*f).to_string()))
154 .collect(),
155 );
156 capabilities["experimental"] = json!({
157 "atd": {
158 "acceptedTokenFormats": formats_value
159 }
160 });
161 }
162 }
163
164 let info = json!({
165 "protocolVersion": "2025-06-18",
166 "serverInfo": {
167 "name": name,
168 "version": version,
169 },
170 "capabilities": capabilities,
171 });
172 ok_response(id, info)
173}
174
175fn split_server_version(s: &str) -> (String, String) {
176 if let Some(idx) = s.rfind(' ') {
181 let (name, ver) = s.split_at(idx);
182 (name.to_string(), ver.trim().to_string())
183 } else {
184 (s.to_string(), "0.0.0".to_string())
185 }
186}
187
188pub fn handle_initialized_notification(id: Option<Value>) -> axum::response::Response {
191 ok_response(id, json!({}))
192}
193
194pub fn handle_tools_list(id: Option<Value>, state: &ServerState) -> axum::response::Response {
201 let tools: Vec<Value> = state
202 .registry
203 .summaries()
204 .into_iter()
205 .filter(|s| !matches!(s.visibility, ToolVisibility::Hidden))
206 .map(|summary| {
207 let def_extras = state
214 .registry
215 .get(&summary.id)
216 .map(|t| {
217 let def = t.definition();
218 json!({
219 "version": def.version,
220 "outputSchema": def.output_schema,
221 })
222 })
223 .unwrap_or_else(|| json!({}));
224 let mut entry = json!({
225 "name": summary.id,
226 "description": summary.description,
227 "inputSchema": summary.input_schema.clone().unwrap_or_else(|| json!({})),
228 });
229 if let Some(obj) = entry.as_object_mut() {
230 if let Some(extra_obj) = def_extras.as_object() {
231 for (k, v) in extra_obj {
232 obj.insert(k.clone(), v.clone());
233 }
234 }
235 }
236 entry
237 })
238 .collect();
239 ok_response(id, json!({ "tools": tools }))
240}
241
242pub async fn handle_tools_call(
256 id: Option<Value>,
257 state: &Arc<ServerState>,
258 identity: Option<&BearerIdentity>,
259 params: Value,
260) -> axum::response::Response {
261 let tool_id = match params.get("name").and_then(|v| v.as_str()) {
262 Some(s) => s.to_string(),
263 None => {
264 return error_response(StatusCode::OK, id, -32602, "missing `name` parameter");
265 }
266 };
267 let mut args = params
268 .get("arguments")
269 .cloned()
270 .unwrap_or_else(|| Value::Object(Default::default()));
271
272 let cursor: Option<String> = match args.as_object_mut() {
276 Some(obj) => obj
277 .remove("__cursor")
278 .and_then(|v| v.as_str().map(String::from)),
279 None => None,
280 };
281
282 let caps = Arc::new(build_caps(state, identity));
286 let caller_id = identity.map(|i| i.caller_id.clone());
287 let tracker = Arc::new(ReadTracker::new());
291
292 let resp = if let Some(c) = cursor {
293 atd_runtime::dispatch::run_tool_continue(
294 state,
295 &tracker,
296 &caps,
297 caller_id.as_deref(),
298 tool_id.clone(),
299 c,
300 )
301 .await
302 } else {
303 atd_runtime::dispatch::run_tool(
304 state,
305 &tracker,
306 &caps,
307 caller_id.as_deref(),
308 tool_id.clone(),
309 args,
310 false, )
312 .await
313 };
314
315 wrap_tool_response(id, resp)
316}
317
318fn build_caps(state: &ServerState, identity: Option<&BearerIdentity>) -> CapabilitySet {
324 let Some(id) = identity else {
325 return CapabilitySet::empty();
326 };
327 let allow = CapabilitySet::from_iter(state.config.granted_capabilities.iter().cloned());
328 let (granted, _denied) = allow.intersect(&id.granted_capabilities);
329 let provenance = granted
335 .iter()
336 .map(|c| atd_runtime::audit::CapProvenance {
337 cap: c.clone(),
338 source: atd_runtime::audit::ProvSource::StringAllowList,
339 })
340 .collect();
341 CapabilitySet::with_provenance(granted, provenance)
342}
343
344pub fn wrap_tool_response(id: Option<Value>, resp: Response) -> axum::response::Response {
350 match resp {
351 Response::ToolResultResponse {
352 result,
353 success,
354 dry_run: _,
355 tool_id: _,
356 next_cursor,
357 } => {
358 let text = serde_json::to_string(&result).unwrap_or_else(|_| "{}".into());
359 let mut body = json!({
366 "content": [{ "type": "text", "text": text }],
367 "isError": !success,
368 });
369 if let Some(cur) = next_cursor {
370 if let Some(obj) = body.as_object_mut() {
371 obj.insert("nextCursor".into(), Value::String(cur));
372 }
373 }
374 ok_response(id, body)
375 }
376 Response::Error {
377 message,
378 code,
379 retryable,
380 details,
381 } => {
382 let jsonrpc_code = -32603i32;
389 let mut data = serde_json::Map::new();
390 if let Some(c) = code {
391 data.insert("atd_code".into(), Value::from(c));
392 }
393 if let Some(r) = retryable {
394 data.insert("retryable".into(), Value::Bool(r));
395 }
396 if let Some(d) = details {
397 data.insert("details".into(), d);
398 }
399 let err = JsonRpcError {
400 code: jsonrpc_code,
401 message,
402 };
403 let envelope = JsonRpcResponse {
407 jsonrpc: "2.0",
408 id,
409 result: None,
410 error: Some(err),
411 };
412 let mut value = serde_json::to_value(&envelope).unwrap_or_else(|_| {
413 json!({"jsonrpc":"2.0", "error": {"code": jsonrpc_code, "message": "serialise"}})
414 });
415 if !data.is_empty() {
416 if let Some(err_obj) = value.get_mut("error").and_then(|e| e.as_object_mut()) {
417 err_obj.insert("data".into(), Value::Object(data));
418 }
419 }
420 (StatusCode::OK, Json(value)).into_response()
421 }
422 other => error_response(
425 StatusCode::INTERNAL_SERVER_ERROR,
426 id,
427 -32603,
428 format!("unexpected dispatch variant: {other:?}"),
429 ),
430 }
431}
432
433#[cfg(test)]
434mod tests {
435 use super::*;
436 use atd_runtime::dispatch::SharedServerConfig;
437 use atd_runtime::registry::{CallFuture, Registry, Tool};
438 use atd_runtime::{CallContext, TierPolicy};
439 use http_body_util::BodyExt;
440 use std::sync::Arc;
441
442 fn stub_def(id: &str) -> atd_protocol::ToolDefinition {
443 use atd_protocol::{
444 BindingProtocol, SafetyLevel, ToolBinding, ToolCapability, ToolResources, ToolSafety,
445 ToolTier, ToolTrust, ToolVisibility, TrustLevel,
446 };
447 let _ = ToolTier::Warm; atd_protocol::ToolDefinition {
449 id: id.into(),
450 name: id.into(),
451 description: "stub".into(),
452 version: "0.0.0".into(),
453 capability: ToolCapability {
454 domain: "echo".into(),
455 actions: vec![],
456 tags: vec![],
457 intent_examples: vec![],
458 },
459 input_schema: json!({"type":"object"}),
460 output_schema: json!({"type":"object"}),
461 bindings: vec![ToolBinding {
462 protocol: BindingProtocol::Cli,
463 config: json!({}),
464 }],
465 safety: ToolSafety {
466 level: SafetyLevel::Read,
467 dry_run: false,
468 side_effects: vec![],
469 data_sensitivity: None,
470 },
471 resources: ToolResources {
472 timeout_ms: 1000,
473 max_concurrent: 1,
474 rate_limit_per_min: None,
475 estimated_tokens: None,
476 },
477 trust: ToolTrust {
478 publisher: "test".into(),
479 trust_level: TrustLevel::L0Unverified,
480 signature: None,
481 },
482 visibility: ToolVisibility::Read,
483 required_capabilities: vec![],
484 tier: None,
485 errors: vec![],
486 }
487 }
488
489 struct EchoStub {
490 def: atd_protocol::ToolDefinition,
491 }
492 impl EchoStub {
493 fn new() -> Self {
494 Self {
495 def: stub_def("ref:echo.say"),
496 }
497 }
498 }
499 impl Tool for EchoStub {
500 fn definition(&self) -> &atd_protocol::ToolDefinition {
501 &self.def
502 }
503 fn call<'a>(&'a self, args: serde_json::Value, _ctx: &'a CallContext) -> CallFuture<'a> {
504 Box::pin(async move { Ok(json!({"echoed": args})) })
505 }
506 }
507
508 fn state_with_echo() -> Arc<ServerState> {
509 let mut reg = Registry::new();
510 reg.register(Arc::new(EchoStub::new()));
511 Arc::new(ServerState {
512 registry: reg,
513 config: SharedServerConfig::for_test(),
514 tier_policy: TierPolicy::defaults(),
515 middleware: vec![],
516 metrics: Arc::new(atd_runtime::MetricsCounters::default()),
517 cursor_issuer: Arc::new(atd_runtime::cursor::CursorIssuer::new([0u8; 32])),
518 })
519 }
520
521 async fn body_to_json(resp: axum::response::Response) -> Value {
522 let bytes = resp.into_body().collect().await.unwrap().to_bytes();
523 serde_json::from_slice(&bytes).unwrap()
524 }
525
526 #[tokio::test]
527 async fn initialize_returns_server_info_with_tools_capability() {
528 let resp = handle_initialize(Some(json!(1)), "atd-test 0.3.0", None);
529 let body = body_to_json(resp).await;
530 assert_eq!(body["jsonrpc"], "2.0");
531 assert_eq!(body["id"], 1);
532 assert_eq!(body["result"]["serverInfo"]["name"], "atd-test");
533 assert_eq!(body["result"]["serverInfo"]["version"], "0.3.0");
534 assert!(body["result"]["capabilities"]["tools"].is_object());
535 assert_eq!(body["result"]["protocolVersion"], "2025-06-18");
536 assert!(body["result"]["capabilities"]["experimental"].is_null());
538 }
539
540 #[tokio::test]
541 async fn initialize_echoes_broker_accepted_token_formats_when_wired() {
542 use atd_runtime::secrets::{ResolveBearerFuture, ResolveFuture, TokenBroker};
543 #[derive(Debug)]
544 struct DeclaringBroker;
545 impl TokenBroker for DeclaringBroker {
546 fn resolve<'a>(&'a self, _caller_id: Option<&'a str>) -> ResolveFuture<'a> {
547 Box::pin(async { Ok(None) })
548 }
549 fn resolve_bearer<'a>(&'a self, _bearer: &'a str) -> ResolveBearerFuture<'a> {
550 Box::pin(async { Ok(None) })
551 }
552 fn accepted_token_formats(&self) -> &'static [&'static str] {
553 &["ucan-jwt", "opaque", "ce-pairing-code"]
554 }
555 }
556 let broker: Arc<dyn TokenBroker> = Arc::new(DeclaringBroker);
557 let resp = handle_initialize(Some(json!(1)), "atd-test 0.3.0", Some(&broker));
558 let body = body_to_json(resp).await;
559 let formats =
560 &body["result"]["capabilities"]["experimental"]["atd"]["acceptedTokenFormats"];
561 assert_eq!(
562 formats,
563 &json!(["ucan-jwt", "opaque", "ce-pairing-code"]),
564 "broker formats should appear in initialize.capabilities.experimental.atd"
565 );
566 }
567
568 #[tokio::test]
569 async fn initialize_omits_experimental_when_broker_declares_empty_formats() {
570 use atd_runtime::secrets::{ResolveBearerFuture, ResolveFuture, TokenBroker};
571 #[derive(Debug)]
572 struct SilentBroker;
573 impl TokenBroker for SilentBroker {
574 fn resolve<'a>(&'a self, _caller_id: Option<&'a str>) -> ResolveFuture<'a> {
575 Box::pin(async { Ok(None) })
576 }
577 fn resolve_bearer<'a>(&'a self, _bearer: &'a str) -> ResolveBearerFuture<'a> {
578 Box::pin(async { Ok(None) })
579 }
580 }
583 let broker: Arc<dyn TokenBroker> = Arc::new(SilentBroker);
584 let resp = handle_initialize(Some(json!(1)), "atd-test 0.3.0", Some(&broker));
585 let body = body_to_json(resp).await;
586 assert!(body["result"]["capabilities"]["experimental"].is_null());
587 }
588
589 #[tokio::test]
590 async fn tools_list_advertises_registered_tools() {
591 let state = state_with_echo();
592 let resp = handle_tools_list(Some(json!(2)), &state);
593 let body = body_to_json(resp).await;
594 let tools = body["result"]["tools"].as_array().unwrap();
595 assert_eq!(tools.len(), 1);
596 assert_eq!(tools[0]["name"], "ref:echo.say");
597 assert!(tools[0]["inputSchema"].is_object());
598 }
599
600 #[tokio::test]
601 async fn tools_call_echo_returns_is_error_false_with_serialised_result() {
602 let state = state_with_echo();
603 let params = json!({"name": "ref:echo.say", "arguments": {"hi": "x"}});
604 let resp = handle_tools_call(Some(json!(3)), &state, None, params).await;
605 let body = body_to_json(resp).await;
606 assert_eq!(body["result"]["isError"], false);
607 let text = body["result"]["content"][0]["text"].as_str().unwrap();
608 let parsed: Value = serde_json::from_str(text).unwrap();
609 assert_eq!(parsed["echoed"]["hi"], "x");
610 }
611
612 #[tokio::test]
613 async fn tools_call_missing_name_returns_invalid_params_error() {
614 let state = state_with_echo();
615 let params = json!({"arguments": {}});
616 let resp = handle_tools_call(Some(json!(4)), &state, None, params).await;
617 let body = body_to_json(resp).await;
618 assert_eq!(body["error"]["code"], -32602);
619 assert!(
620 body["error"]["message"]
621 .as_str()
622 .unwrap()
623 .contains("missing")
624 );
625 }
626
627 #[tokio::test]
628 async fn tools_call_unknown_tool_returns_error_envelope() {
629 let state = state_with_echo();
630 let params = json!({"name": "ref:does-not-exist", "arguments": {}});
631 let resp = handle_tools_call(Some(json!(5)), &state, None, params).await;
632 let body = body_to_json(resp).await;
633 assert_eq!(body["error"]["code"], -32603);
635 assert!(
636 body["error"]["message"]
637 .as_str()
638 .unwrap()
639 .contains("tool not found")
640 );
641 }
642
643 struct PageStubHttp {
650 def: atd_protocol::ToolDefinition,
651 }
652 impl PageStubHttp {
653 fn new() -> Self {
654 Self {
655 def: stub_def("ref:page_stub"),
656 }
657 }
658 }
659 impl Tool for PageStubHttp {
660 fn definition(&self) -> &atd_protocol::ToolDefinition {
661 &self.def
662 }
663 fn supports_pagination(&self) -> bool {
664 true
665 }
666 fn call<'a>(&'a self, _args: serde_json::Value, _ctx: &'a CallContext) -> CallFuture<'a> {
667 Box::pin(async { Ok(json!({"unused": true})) })
668 }
669 fn call_paginated<'a>(
670 &'a self,
671 _args: serde_json::Value,
672 ctx: &'a CallContext,
673 cursor: Option<&'a str>,
674 ) -> atd_runtime::PaginatedCallFuture<'a> {
675 let issuer = ctx.cursor_issuer().expect("issuer attached");
676 let page = match cursor {
677 None => 1u32,
678 Some(c) => {
679 let p = issuer.verify(c, 300).expect("pre-verified");
680 u32::from_be_bytes(p.opaque_state[..4].try_into().unwrap())
681 }
682 };
683 let payload = atd_runtime::cursor::CursorPayload {
684 tool_id: "ref:page_stub".into(),
685 caller_id: ctx.caller_id.clone(),
686 args_fingerprint: [0u8; 32],
687 page_index: page + 1,
688 issued_at_unix: std::time::SystemTime::now()
689 .duration_since(std::time::UNIX_EPOCH)
690 .unwrap()
691 .as_secs(),
692 server_session: issuer.session_nonce(),
693 opaque_state: (page + 1).to_be_bytes().to_vec(),
694 };
695 let next = issuer.issue(payload).ok();
696 Box::pin(async move {
697 Ok(atd_runtime::PaginatedResult {
698 value: json!({"page": page}),
699 next_cursor: next,
700 })
701 })
702 }
703 }
704
705 fn state_with_page_stub() -> Arc<ServerState> {
706 let mut reg = Registry::new();
707 reg.register(Arc::new(PageStubHttp::new()));
708 Arc::new(ServerState {
709 registry: reg,
710 config: SharedServerConfig::for_test(),
711 tier_policy: TierPolicy::defaults(),
712 middleware: vec![],
713 metrics: Arc::new(atd_runtime::MetricsCounters::default()),
714 cursor_issuer: Arc::new(atd_runtime::cursor::CursorIssuer::new([0u8; 32])),
715 })
716 }
717
718 #[tokio::test]
719 async fn tools_call_initial_paginated_call_surfaces_next_cursor_in_envelope() {
720 let state = state_with_page_stub();
721 let params = json!({"name": "ref:page_stub", "arguments": {}});
725 let resp = handle_tools_call(Some(json!(7)), &state, None, params).await;
726 let body = body_to_json(resp).await;
727 assert_eq!(body["result"]["isError"], false);
728 assert!(
729 body["result"]["nextCursor"].is_string(),
730 "paginating tool on first call must emit nextCursor; got: {body}"
731 );
732 }
733
734 #[tokio::test]
735 async fn tools_call_with_dunder_cursor_routes_to_run_tool_continue() {
736 let state = state_with_page_stub();
737 let payload = atd_runtime::cursor::CursorPayload {
740 tool_id: "ref:page_stub".into(),
741 caller_id: None,
742 args_fingerprint: [0u8; 32],
743 page_index: 2,
744 issued_at_unix: std::time::SystemTime::now()
745 .duration_since(std::time::UNIX_EPOCH)
746 .unwrap()
747 .as_secs(),
748 server_session: state.cursor_issuer.session_nonce(),
749 opaque_state: 2u32.to_be_bytes().to_vec(),
750 };
751 let cursor = state.cursor_issuer.issue(payload).unwrap();
752
753 let params = json!({
755 "name": "ref:page_stub",
756 "arguments": {"__cursor": cursor},
757 });
758 let resp = handle_tools_call(Some(json!(8)), &state, None, params).await;
759 let body = body_to_json(resp).await;
760 assert_eq!(body["result"]["isError"], false);
761 let text = body["result"]["content"][0]["text"].as_str().unwrap();
763 let parsed: Value = serde_json::from_str(text).unwrap();
764 assert_eq!(parsed["page"], 2);
765 assert!(
767 body["result"]["nextCursor"].is_string(),
768 "continuation must emit nextCursor; got: {body}"
769 );
770 }
771
772 #[tokio::test]
773 async fn tools_call_expired_cursor_returns_jsonrpc_error() {
774 let state = state_with_page_stub();
775 let payload = atd_runtime::cursor::CursorPayload {
777 tool_id: "ref:page_stub".into(),
778 caller_id: None,
779 args_fingerprint: [0u8; 32],
780 page_index: 2,
781 issued_at_unix: 1, server_session: state.cursor_issuer.session_nonce(),
783 opaque_state: 2u32.to_be_bytes().to_vec(),
784 };
785 let cursor = state.cursor_issuer.issue(payload).unwrap();
786 let params = json!({
787 "name": "ref:page_stub",
788 "arguments": {"__cursor": cursor},
789 });
790 let resp = handle_tools_call(Some(json!(9)), &state, None, params).await;
791 let body = body_to_json(resp).await;
792 assert_eq!(body["error"]["code"], -32603);
794 assert_eq!(body["error"]["data"]["atd_code"], 1020);
795 }
796}