#![allow(clippy::unwrap_used, clippy::pedantic, missing_docs)]
use std::{
future::Future,
pin::Pin,
sync::{
Arc,
atomic::{AtomicUsize, Ordering},
},
};
use futures::Stream;
use axum::body::Body;
use http::{Request, StatusCode, header};
use http_body_util::BodyExt as _;
use polyc_a2a::{
AppState, InMemoryTaskStore, TurnOutcome, TurnRequest, TurnRunner, TurnStreamEvent,
UnconfiguredApprovalResponder, card::signed_card, router,
};
use polyc_crypto::Signer;
use polyc_runtime::admission::AdmissionGate;
use serde_json::{Value, json};
use tokio::sync::Notify;
use tower::ServiceExt as _;
const TOKEN: &str = "s3cr3t-a2a-token";
struct CountingRunner {
calls: Arc<AtomicUsize>,
outcome: TurnOutcome,
}
impl TurnRunner for CountingRunner {
fn run_turn<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<Box<dyn Future<Output = TurnOutcome> + Send + 'a>> {
self.calls.fetch_add(1, Ordering::SeqCst);
let outcome = self.outcome.clone();
Box::pin(async move { outcome })
}
}
fn app_with_token(token: &str, calls: Arc<AtomicUsize>) -> axum::Router {
let signer = Signer::from_seed(7);
let card = signed_card(
&polyc_a2a::card::CardConfig {
name: "Polychrome".to_owned(),
description: "test".to_owned(),
url: "https://agent.example/".to_owned(),
version: "0.1.3".to_owned(),
},
&signer,
);
router(AppState {
card: Arc::new(card),
runner: Arc::new(CountingRunner {
calls,
outcome: TurnOutcome::Completed {
text: "42".to_owned(),
},
}),
approvals: Arc::new(UnconfiguredApprovalResponder),
store: Arc::new(InMemoryTaskStore::new()),
turn_limit: AdmissionGate::new(64),
peers: if token.is_empty() {
polyc_a2a::PeerAuthenticator::default()
} else {
polyc_a2a::PeerAuthenticator::single("test-peer", token).unwrap()
},
})
}
fn send_message() -> Value {
json!({
"jsonrpc": "2.0",
"id": 1,
"method": "SendMessage",
"params": {
"message": {
"role": "ROLE_USER",
"messageId": "m1",
"contextId": "ctx-1",
"parts": [{ "text": "what is the weather?" }]
}
}
})
}
async fn post(app: axum::Router, auth: Option<&str>) -> http::Response<Body> {
post_value(app, auth, &send_message()).await
}
async fn post_value(
app: axum::Router,
auth: Option<&str>,
request: &Value,
) -> http::Response<Body> {
let mut builder = Request::builder()
.method("POST")
.uri("/")
.header(header::CONTENT_TYPE, "application/json");
if let Some(value) = auth {
builder = builder.header(header::AUTHORIZATION, value);
}
app.oneshot(
builder
.body(Body::from(serde_json::to_vec(request).unwrap()))
.unwrap(),
)
.await
.unwrap()
}
struct StalledReceiptRunner {
receipt: Arc<Notify>,
}
impl TurnRunner for StalledReceiptRunner {
fn receive_ingress<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<
Box<
dyn Future<
Output = Result<
polyc_a2a::task::IngressReceipt,
polyc_a2a::task::IngressReceiptError,
>,
> + Send
+ 'a,
>,
> {
Box::pin(async move {
self.receipt.notified().await;
Ok(polyc_a2a::task::IngressReceipt {
dispatch_id: "stalled-dispatch".to_owned(),
})
})
}
fn run_turn<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<Box<dyn Future<Output = TurnOutcome> + Send + 'a>> {
Box::pin(async {
TurnOutcome::Completed {
text: "done".to_owned(),
}
})
}
fn run_turn_streaming<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<Box<dyn Stream<Item = TurnStreamEvent> + Send + 'a>> {
Box::pin(async_stream::stream! {
yield TurnStreamEvent::DurablyReceived;
yield TurnStreamEvent::Outcome(TurnOutcome::Completed {
text: "done".to_owned(),
});
})
}
}
fn app_with_stalled_receipt(receipt: Arc<Notify>) -> axum::Router {
let signer = Signer::from_seed(7);
let card = signed_card(
&polyc_a2a::card::CardConfig {
name: "Polychrome".to_owned(),
description: "test".to_owned(),
url: "https://agent.example/".to_owned(),
version: "0.1.3".to_owned(),
},
&signer,
);
router(AppState {
card: Arc::new(card),
runner: Arc::new(StalledReceiptRunner { receipt }),
approvals: Arc::new(UnconfiguredApprovalResponder),
store: Arc::new(InMemoryTaskStore::new()),
turn_limit: AdmissionGate::new(64),
peers: polyc_a2a::PeerAuthenticator::single("test-peer", TOKEN).unwrap(),
})
}
#[tokio::test]
async fn missing_authorization_is_rejected_before_dispatch() {
let calls = Arc::new(AtomicUsize::new(0));
let app = app_with_token(TOKEN, calls.clone());
let response = post(app, None).await;
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(calls.load(Ordering::SeqCst), 0, "turn must never run");
}
#[tokio::test]
async fn wrong_authorization_is_rejected_before_dispatch() {
let calls = Arc::new(AtomicUsize::new(0));
let app = app_with_token(TOKEN, calls.clone());
let response = post(app, Some("Bearer not-the-token")).await;
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(calls.load(Ordering::SeqCst), 0, "turn must never run");
}
#[tokio::test]
async fn malformed_authorization_scheme_is_rejected_before_dispatch() {
let calls = Arc::new(AtomicUsize::new(0));
let app = app_with_token(TOKEN, calls.clone());
let response = post(app, Some(TOKEN)).await;
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(calls.load(Ordering::SeqCst), 0, "turn must never run");
}
#[tokio::test]
async fn correct_bearer_token_runs_the_turn() {
let calls = Arc::new(AtomicUsize::new(0));
let app = app_with_token(TOKEN, calls.clone());
let response = post(app, Some(&format!("Bearer {TOKEN}"))).await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"turn must run exactly once"
);
let body = response.into_body().collect().await.unwrap().to_bytes();
let value: Value = serde_json::from_slice(&body).unwrap();
assert_eq!(
value["result"]["task"]["status"]["state"],
"TASK_STATE_COMPLETED"
);
}
#[tokio::test]
async fn unconfigured_bearer_token_fails_every_request_closed() {
let calls = Arc::new(AtomicUsize::new(0));
let app = app_with_token("", calls.clone());
let response = post(app, Some("Bearer ")).await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(calls.load(Ordering::SeqCst), 0, "turn must never run");
}
#[tokio::test]
async fn agent_card_advertises_the_bearer_security_scheme() {
let calls = Arc::new(AtomicUsize::new(0));
let app = app_with_token(TOKEN, calls);
let response = app
.oneshot(
Request::builder()
.uri("/.well-known/agent-card.json")
.body(Body::empty())
.unwrap(),
)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::OK);
let body = response.into_body().collect().await.unwrap().to_bytes();
let card: Value = serde_json::from_slice(&body).unwrap();
let schemes = card["securitySchemes"]
.as_object()
.expect("card must advertise securitySchemes");
let bearer = schemes
.values()
.find(|scheme| scheme["scheme"] == "bearer")
.expect("card must advertise a bearer scheme");
assert_eq!(bearer["type"], "http");
assert!(
card["security"]
.as_array()
.is_some_and(|reqs| !reqs.is_empty()),
"card must require the advertised scheme"
);
}
#[tokio::test]
async fn unary_response_waits_for_the_durable_ingress_receipt() {
let receipt = Arc::new(Notify::new());
let app = app_with_stalled_receipt(receipt.clone());
let authorization = format!("Bearer {TOKEN}");
let request = send_message();
let response = post_value(app, Some(&authorization), &request);
tokio::pin!(response);
assert!(
tokio::time::timeout(std::time::Duration::from_millis(50), &mut response)
.await
.is_err(),
"unary success must remain pending while State has not acknowledged ingress"
);
receipt.notify_waiters();
assert_eq!(response.await.status(), StatusCode::OK);
}
#[tokio::test]
async fn streaming_headers_wait_for_the_durable_ingress_receipt() {
let receipt = Arc::new(Notify::new());
let app = app_with_stalled_receipt(receipt.clone());
let request = json!({
"jsonrpc": "2.0", "id": 8, "method": "SendStreamingMessage", "params": {
"message": {
"role": "ROLE_USER", "messageId": "m-stream", "parts": [{"text": "q"}]
}
}
});
let authorization = format!("Bearer {TOKEN}");
let response = post_value(app, Some(&authorization), &request);
tokio::pin!(response);
assert!(
tokio::time::timeout(std::time::Duration::from_millis(50), &mut response)
.await
.is_err(),
"SSE headers must remain pending while State has not acknowledged ingress"
);
receipt.notify_waiters();
let response = response.await;
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers()[header::CONTENT_TYPE],
"text/event-stream"
);
}