use std::sync::Arc;
use atd_protocol::{Response, ToolVisibility};
use atd_runtime::capability::CapabilitySet;
use atd_runtime::dispatch::ServerState;
use atd_runtime::secrets::BearerIdentity;
use atd_runtime::tracker::ReadTracker;
use axum::Json;
use axum::http::StatusCode;
use axum::response::IntoResponse;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
#[derive(Debug, Deserialize)]
pub struct JsonRpcRequest {
#[allow(dead_code)]
#[serde(default)]
pub jsonrpc: String,
#[serde(default)]
pub id: Option<Value>,
pub method: String,
#[serde(default)]
pub params: Value,
}
#[derive(Debug, Serialize)]
pub struct JsonRpcResponse {
pub jsonrpc: &'static str,
#[serde(skip_serializing_if = "Option::is_none")]
pub id: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<JsonRpcError>,
}
#[derive(Debug, Serialize)]
pub struct JsonRpcError {
pub code: i32,
pub message: String,
}
pub fn error_response(
status: StatusCode,
id: Option<Value>,
code: i32,
message: impl Into<String>,
) -> axum::response::Response {
error_response_with_headers(status, id, code, message, &[])
}
pub fn error_response_with_headers(
status: StatusCode,
id: Option<Value>,
code: i32,
message: impl Into<String>,
extra_headers: &[(&str, String)],
) -> axum::response::Response {
let mut resp: axum::response::Response = (
status,
Json(JsonRpcResponse {
jsonrpc: "2.0",
id,
result: None,
error: Some(JsonRpcError {
code,
message: message.into(),
}),
}),
)
.into_response();
if !extra_headers.is_empty() {
let headers = resp.headers_mut();
for (name, value) in extra_headers {
if let (Ok(hn), Ok(hv)) = (
axum::http::HeaderName::from_bytes(name.as_bytes()),
axum::http::HeaderValue::from_str(value),
) {
headers.insert(hn, hv);
}
}
}
resp
}
pub fn ok_response(id: Option<Value>, result: Value) -> axum::response::Response {
Json(JsonRpcResponse {
jsonrpc: "2.0",
id,
result: Some(result),
error: None,
})
.into_response()
}
pub fn handle_initialize(
id: Option<Value>,
server_version: &str,
token_broker: Option<&Arc<dyn atd_runtime::TokenBroker>>,
) -> axum::response::Response {
let (name, version) = split_server_version(server_version);
let mut capabilities = json!({ "tools": {} });
if let Some(broker) = token_broker {
let formats = broker.accepted_token_formats();
if !formats.is_empty() {
let formats_value = serde_json::Value::Array(
formats
.iter()
.map(|f| serde_json::Value::String((*f).to_string()))
.collect(),
);
capabilities["experimental"] = json!({
"atd": {
"acceptedTokenFormats": formats_value
}
});
}
}
let info = json!({
"protocolVersion": "2025-06-18",
"serverInfo": {
"name": name,
"version": version,
},
"capabilities": capabilities,
});
ok_response(id, info)
}
fn split_server_version(s: &str) -> (String, String) {
if let Some(idx) = s.rfind(' ') {
let (name, ver) = s.split_at(idx);
(name.to_string(), ver.trim().to_string())
} else {
(s.to_string(), "0.0.0".to_string())
}
}
pub fn handle_initialized_notification(id: Option<Value>) -> axum::response::Response {
ok_response(id, json!({}))
}
pub fn handle_tools_list(id: Option<Value>, state: &ServerState) -> axum::response::Response {
let tools: Vec<Value> = state
.registry
.summaries()
.into_iter()
.filter(|s| !matches!(s.visibility, ToolVisibility::Hidden))
.map(|summary| {
let def_extras = state
.registry
.get(&summary.id)
.map(|t| {
let def = t.definition();
json!({
"version": def.version,
"outputSchema": def.output_schema,
})
})
.unwrap_or_else(|| json!({}));
let mut entry = json!({
"name": summary.id,
"description": summary.description,
"inputSchema": summary.input_schema.clone().unwrap_or_else(|| json!({})),
});
if let Some(obj) = entry.as_object_mut() {
if let Some(extra_obj) = def_extras.as_object() {
for (k, v) in extra_obj {
obj.insert(k.clone(), v.clone());
}
}
}
entry
})
.collect();
ok_response(id, json!({ "tools": tools }))
}
pub async fn handle_tools_call(
id: Option<Value>,
state: &Arc<ServerState>,
identity: Option<&BearerIdentity>,
params: Value,
) -> axum::response::Response {
let tool_id = match params.get("name").and_then(|v| v.as_str()) {
Some(s) => s.to_string(),
None => {
return error_response(StatusCode::OK, id, -32602, "missing `name` parameter");
}
};
let mut args = params
.get("arguments")
.cloned()
.unwrap_or_else(|| Value::Object(Default::default()));
let cursor: Option<String> = match args.as_object_mut() {
Some(obj) => obj
.remove("__cursor")
.and_then(|v| v.as_str().map(String::from)),
None => None,
};
let caps = Arc::new(build_caps(state, identity));
let caller_id = identity.map(|i| i.caller_id.clone());
let tracker = Arc::new(ReadTracker::new());
let resp = if let Some(c) = cursor {
atd_runtime::dispatch::run_tool_continue(
state,
&tracker,
&caps,
caller_id.as_deref(),
tool_id.clone(),
c,
)
.await
} else {
atd_runtime::dispatch::run_tool(
state,
&tracker,
&caps,
caller_id.as_deref(),
tool_id.clone(),
args,
false, )
.await
};
wrap_tool_response(id, resp)
}
fn build_caps(state: &ServerState, identity: Option<&BearerIdentity>) -> CapabilitySet {
let Some(id) = identity else {
return CapabilitySet::empty();
};
let allow = CapabilitySet::from_iter(state.config.granted_capabilities.iter().cloned());
let (granted, _denied) = allow.intersect(&id.granted_capabilities);
let provenance = granted
.iter()
.map(|c| atd_runtime::audit::CapProvenance {
cap: c.clone(),
source: atd_runtime::audit::ProvSource::StringAllowList,
})
.collect();
CapabilitySet::with_provenance(granted, provenance)
}
pub fn wrap_tool_response(id: Option<Value>, resp: Response) -> axum::response::Response {
match resp {
Response::ToolResultResponse {
result,
success,
dry_run: _,
tool_id: _,
next_cursor,
} => {
let text = serde_json::to_string(&result).unwrap_or_else(|_| "{}".into());
let mut body = json!({
"content": [{ "type": "text", "text": text }],
"isError": !success,
});
if let Some(cur) = next_cursor {
if let Some(obj) = body.as_object_mut() {
obj.insert("nextCursor".into(), Value::String(cur));
}
}
ok_response(id, body)
}
Response::Error {
message,
code,
retryable,
details,
} => {
let jsonrpc_code = -32603i32;
let mut data = serde_json::Map::new();
if let Some(c) = code {
data.insert("atd_code".into(), Value::from(c));
}
if let Some(r) = retryable {
data.insert("retryable".into(), Value::Bool(r));
}
if let Some(d) = details {
data.insert("details".into(), d);
}
let err = JsonRpcError {
code: jsonrpc_code,
message,
};
let envelope = JsonRpcResponse {
jsonrpc: "2.0",
id,
result: None,
error: Some(err),
};
let mut value = serde_json::to_value(&envelope).unwrap_or_else(|_| {
json!({"jsonrpc":"2.0", "error": {"code": jsonrpc_code, "message": "serialise"}})
});
if !data.is_empty() {
if let Some(err_obj) = value.get_mut("error").and_then(|e| e.as_object_mut()) {
err_obj.insert("data".into(), Value::Object(data));
}
}
(StatusCode::OK, Json(value)).into_response()
}
other => error_response(
StatusCode::INTERNAL_SERVER_ERROR,
id,
-32603,
format!("unexpected dispatch variant: {other:?}"),
),
}
}
#[cfg(test)]
mod tests {
use super::*;
use atd_runtime::dispatch::SharedServerConfig;
use atd_runtime::registry::{CallFuture, Registry, Tool};
use atd_runtime::{CallContext, TierPolicy};
use http_body_util::BodyExt;
use std::sync::Arc;
fn stub_def(id: &str) -> atd_protocol::ToolDefinition {
use atd_protocol::{
BindingProtocol, SafetyLevel, ToolBinding, ToolCapability, ToolResources, ToolSafety,
ToolTier, ToolTrust, ToolVisibility, TrustLevel,
};
let _ = ToolTier::Warm; atd_protocol::ToolDefinition {
id: id.into(),
name: id.into(),
description: "stub".into(),
version: "0.0.0".into(),
capability: ToolCapability {
domain: "echo".into(),
actions: vec![],
tags: vec![],
intent_examples: vec![],
},
input_schema: json!({"type":"object"}),
output_schema: json!({"type":"object"}),
bindings: vec![ToolBinding {
protocol: BindingProtocol::Cli,
config: json!({}),
}],
safety: ToolSafety {
level: SafetyLevel::Read,
dry_run: false,
side_effects: vec![],
data_sensitivity: None,
},
resources: ToolResources {
timeout_ms: 1000,
max_concurrent: 1,
rate_limit_per_min: None,
estimated_tokens: None,
},
trust: ToolTrust {
publisher: "test".into(),
trust_level: TrustLevel::L0Unverified,
signature: None,
},
visibility: ToolVisibility::Read,
required_capabilities: vec![],
tier: None,
errors: vec![],
}
}
struct EchoStub {
def: atd_protocol::ToolDefinition,
}
impl EchoStub {
fn new() -> Self {
Self {
def: stub_def("ref:echo.say"),
}
}
}
impl Tool for EchoStub {
fn definition(&self) -> &atd_protocol::ToolDefinition {
&self.def
}
fn call<'a>(&'a self, args: serde_json::Value, _ctx: &'a CallContext) -> CallFuture<'a> {
Box::pin(async move { Ok(json!({"echoed": args})) })
}
}
fn state_with_echo() -> Arc<ServerState> {
let mut reg = Registry::new();
reg.register(Arc::new(EchoStub::new()));
Arc::new(ServerState {
registry: reg,
config: SharedServerConfig::for_test(),
tier_policy: TierPolicy::defaults(),
middleware: vec![],
metrics: Arc::new(atd_runtime::MetricsCounters::default()),
cursor_issuer: Arc::new(atd_runtime::cursor::CursorIssuer::new([0u8; 32])),
})
}
async fn body_to_json(resp: axum::response::Response) -> Value {
let bytes = resp.into_body().collect().await.unwrap().to_bytes();
serde_json::from_slice(&bytes).unwrap()
}
#[tokio::test]
async fn initialize_returns_server_info_with_tools_capability() {
let resp = handle_initialize(Some(json!(1)), "atd-test 0.3.0", None);
let body = body_to_json(resp).await;
assert_eq!(body["jsonrpc"], "2.0");
assert_eq!(body["id"], 1);
assert_eq!(body["result"]["serverInfo"]["name"], "atd-test");
assert_eq!(body["result"]["serverInfo"]["version"], "0.3.0");
assert!(body["result"]["capabilities"]["tools"].is_object());
assert_eq!(body["result"]["protocolVersion"], "2025-06-18");
assert!(body["result"]["capabilities"]["experimental"].is_null());
}
#[tokio::test]
async fn initialize_echoes_broker_accepted_token_formats_when_wired() {
use atd_runtime::secrets::{ResolveBearerFuture, ResolveFuture, TokenBroker};
#[derive(Debug)]
struct DeclaringBroker;
impl TokenBroker for DeclaringBroker {
fn resolve<'a>(&'a self, _caller_id: Option<&'a str>) -> ResolveFuture<'a> {
Box::pin(async { Ok(None) })
}
fn resolve_bearer<'a>(&'a self, _bearer: &'a str) -> ResolveBearerFuture<'a> {
Box::pin(async { Ok(None) })
}
fn accepted_token_formats(&self) -> &'static [&'static str] {
&["ucan-jwt", "opaque", "ce-pairing-code"]
}
}
let broker: Arc<dyn TokenBroker> = Arc::new(DeclaringBroker);
let resp = handle_initialize(Some(json!(1)), "atd-test 0.3.0", Some(&broker));
let body = body_to_json(resp).await;
let formats =
&body["result"]["capabilities"]["experimental"]["atd"]["acceptedTokenFormats"];
assert_eq!(
formats,
&json!(["ucan-jwt", "opaque", "ce-pairing-code"]),
"broker formats should appear in initialize.capabilities.experimental.atd"
);
}
#[tokio::test]
async fn initialize_omits_experimental_when_broker_declares_empty_formats() {
use atd_runtime::secrets::{ResolveBearerFuture, ResolveFuture, TokenBroker};
#[derive(Debug)]
struct SilentBroker;
impl TokenBroker for SilentBroker {
fn resolve<'a>(&'a self, _caller_id: Option<&'a str>) -> ResolveFuture<'a> {
Box::pin(async { Ok(None) })
}
fn resolve_bearer<'a>(&'a self, _bearer: &'a str) -> ResolveBearerFuture<'a> {
Box::pin(async { Ok(None) })
}
}
let broker: Arc<dyn TokenBroker> = Arc::new(SilentBroker);
let resp = handle_initialize(Some(json!(1)), "atd-test 0.3.0", Some(&broker));
let body = body_to_json(resp).await;
assert!(body["result"]["capabilities"]["experimental"].is_null());
}
#[tokio::test]
async fn tools_list_advertises_registered_tools() {
let state = state_with_echo();
let resp = handle_tools_list(Some(json!(2)), &state);
let body = body_to_json(resp).await;
let tools = body["result"]["tools"].as_array().unwrap();
assert_eq!(tools.len(), 1);
assert_eq!(tools[0]["name"], "ref:echo.say");
assert!(tools[0]["inputSchema"].is_object());
}
#[tokio::test]
async fn tools_call_echo_returns_is_error_false_with_serialised_result() {
let state = state_with_echo();
let params = json!({"name": "ref:echo.say", "arguments": {"hi": "x"}});
let resp = handle_tools_call(Some(json!(3)), &state, None, params).await;
let body = body_to_json(resp).await;
assert_eq!(body["result"]["isError"], false);
let text = body["result"]["content"][0]["text"].as_str().unwrap();
let parsed: Value = serde_json::from_str(text).unwrap();
assert_eq!(parsed["echoed"]["hi"], "x");
}
#[tokio::test]
async fn tools_call_missing_name_returns_invalid_params_error() {
let state = state_with_echo();
let params = json!({"arguments": {}});
let resp = handle_tools_call(Some(json!(4)), &state, None, params).await;
let body = body_to_json(resp).await;
assert_eq!(body["error"]["code"], -32602);
assert!(
body["error"]["message"]
.as_str()
.unwrap()
.contains("missing")
);
}
#[tokio::test]
async fn tools_call_unknown_tool_returns_error_envelope() {
let state = state_with_echo();
let params = json!({"name": "ref:does-not-exist", "arguments": {}});
let resp = handle_tools_call(Some(json!(5)), &state, None, params).await;
let body = body_to_json(resp).await;
assert_eq!(body["error"]["code"], -32603);
assert!(
body["error"]["message"]
.as_str()
.unwrap()
.contains("tool not found")
);
}
struct PageStubHttp {
def: atd_protocol::ToolDefinition,
}
impl PageStubHttp {
fn new() -> Self {
Self {
def: stub_def("ref:page_stub"),
}
}
}
impl Tool for PageStubHttp {
fn definition(&self) -> &atd_protocol::ToolDefinition {
&self.def
}
fn supports_pagination(&self) -> bool {
true
}
fn call<'a>(&'a self, _args: serde_json::Value, _ctx: &'a CallContext) -> CallFuture<'a> {
Box::pin(async { Ok(json!({"unused": true})) })
}
fn call_paginated<'a>(
&'a self,
_args: serde_json::Value,
ctx: &'a CallContext,
cursor: Option<&'a str>,
) -> atd_runtime::PaginatedCallFuture<'a> {
let issuer = ctx.cursor_issuer().expect("issuer attached");
let page = match cursor {
None => 1u32,
Some(c) => {
let p = issuer.verify(c, 300).expect("pre-verified");
u32::from_be_bytes(p.opaque_state[..4].try_into().unwrap())
}
};
let payload = atd_runtime::cursor::CursorPayload {
tool_id: "ref:page_stub".into(),
caller_id: ctx.caller_id.clone(),
args_fingerprint: [0u8; 32],
page_index: page + 1,
issued_at_unix: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs(),
server_session: issuer.session_nonce(),
opaque_state: (page + 1).to_be_bytes().to_vec(),
};
let next = issuer.issue(payload).ok();
Box::pin(async move {
Ok(atd_runtime::PaginatedResult {
value: json!({"page": page}),
next_cursor: next,
})
})
}
}
fn state_with_page_stub() -> Arc<ServerState> {
let mut reg = Registry::new();
reg.register(Arc::new(PageStubHttp::new()));
Arc::new(ServerState {
registry: reg,
config: SharedServerConfig::for_test(),
tier_policy: TierPolicy::defaults(),
middleware: vec![],
metrics: Arc::new(atd_runtime::MetricsCounters::default()),
cursor_issuer: Arc::new(atd_runtime::cursor::CursorIssuer::new([0u8; 32])),
})
}
#[tokio::test]
async fn tools_call_initial_paginated_call_surfaces_next_cursor_in_envelope() {
let state = state_with_page_stub();
let params = json!({"name": "ref:page_stub", "arguments": {}});
let resp = handle_tools_call(Some(json!(7)), &state, None, params).await;
let body = body_to_json(resp).await;
assert_eq!(body["result"]["isError"], false);
assert!(
body["result"]["nextCursor"].is_string(),
"paginating tool on first call must emit nextCursor; got: {body}"
);
}
#[tokio::test]
async fn tools_call_with_dunder_cursor_routes_to_run_tool_continue() {
let state = state_with_page_stub();
let payload = atd_runtime::cursor::CursorPayload {
tool_id: "ref:page_stub".into(),
caller_id: None,
args_fingerprint: [0u8; 32],
page_index: 2,
issued_at_unix: std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_secs(),
server_session: state.cursor_issuer.session_nonce(),
opaque_state: 2u32.to_be_bytes().to_vec(),
};
let cursor = state.cursor_issuer.issue(payload).unwrap();
let params = json!({
"name": "ref:page_stub",
"arguments": {"__cursor": cursor},
});
let resp = handle_tools_call(Some(json!(8)), &state, None, params).await;
let body = body_to_json(resp).await;
assert_eq!(body["result"]["isError"], false);
let text = body["result"]["content"][0]["text"].as_str().unwrap();
let parsed: Value = serde_json::from_str(text).unwrap();
assert_eq!(parsed["page"], 2);
assert!(
body["result"]["nextCursor"].is_string(),
"continuation must emit nextCursor; got: {body}"
);
}
#[tokio::test]
async fn tools_call_expired_cursor_returns_jsonrpc_error() {
let state = state_with_page_stub();
let payload = atd_runtime::cursor::CursorPayload {
tool_id: "ref:page_stub".into(),
caller_id: None,
args_fingerprint: [0u8; 32],
page_index: 2,
issued_at_unix: 1, server_session: state.cursor_issuer.session_nonce(),
opaque_state: 2u32.to_be_bytes().to_vec(),
};
let cursor = state.cursor_issuer.issue(payload).unwrap();
let params = json!({
"name": "ref:page_stub",
"arguments": {"__cursor": cursor},
});
let resp = handle_tools_call(Some(json!(9)), &state, None, params).await;
let body = body_to_json(resp).await;
assert_eq!(body["error"]["code"], -32603);
assert_eq!(body["error"]["data"]["atd_code"], 1020);
}
}