polyc-a2a 2026.8.3

//! The A2A HTTP surface: the Agent Card discovery endpoint and the JSON-RPC
//! endpoint, plus the shared [`AppState`].
//!
//! Two routes:
//!
//! - `GET /.well-known/agent-card.json` serves the pre-signed Agent Card. This
//!   is discovery metadata, served unauthenticated (the card advertises the
//!   scheme a peer must then present, it never carries a secret itself).
//! - `POST /` is the A2A v1.0 JSON-RPC endpoint; the card's `supportedInterfaces`
//!   url advertises this address to peers. Every request must carry the
//!   configured `Authorization: Bearer <token>` (edge concern 5, see
//!   `polyc_rpc_client::edge`) — checked in `authenticate` before either
//!   dispatch path ever sees the body. Most methods return a unary JSON
//!   response (`rpc`'s `handle`); `SendStreamingMessage`/`TaskSubscription`
//!   (`#371`) instead return `text/event-stream` — `rpc`'s
//!   `is_streaming_method` peeks the body's `method` to decide which.

use std::sync::Arc;

use axum::{
    Json, Router,
    body::Bytes,
    extract::State,
    http::{HeaderMap, StatusCode, header},
    response::{
        IntoResponse, Response,
        sse::{Event, KeepAlive, Sse},
    },
    routing::{get, post},
};
use futures::StreamExt as _;
use polyc_rpc_client::Sensitive;
use polyc_runtime::admission::AdmissionGate;
use serde_json::Value;
use subtle::ConstantTimeEq;

use crate::{
    rpc,
    store::TaskStore,
    task::{ApprovalResponder, TurnRunner},
};

/// Shared handler state — cheaply clonable across requests.
#[derive(Clone)]
pub struct AppState {
    /// The fully-signed Agent Card, built once at startup and served verbatim.
    pub card: Arc<Value>,
    /// Runs a turn for each inbound `SendMessage`.
    pub runner: Arc<dyn TurnRunner>,
    /// Submits a human's decision for a task paused in `input-required`
    /// (`#792`), and re-drives the turn via `runner` once persisted.
    pub approvals: Arc<dyn ApprovalResponder>,
    /// Records produced tasks durably so `GetTask`/`CancelTask`/`ListTasks`
    /// can read them back — including from a different replica of this edge,
    /// and after a restart.
    pub store: Arc<dyn TaskStore>,
    /// Bounds how many turns this edge dials concurrently (`#795`). A burst
    /// beyond the bound sheds a `failed` task rather than dialing an
    /// already-loaded agent — the shared [`polyc_runtime::admission`] gate
    /// every edge fronts its dial with.
    pub turn_limit: AdmissionGate,
    /// Per-peer bearer credentials for authenticated JSON-RPC ingress.
    pub peers: PeerAuthenticator,
}

/// Deployment-configured mapping from bearer credentials to stable peer ids.
///
/// A credential belongs to exactly one peer. Construction refuses duplicate
/// peer ids, duplicate tokens, and empty components so two callers can never
/// share one ingress namespace accidentally.
#[derive(Clone, Default)]
pub struct PeerAuthenticator {
    entries: Arc<Vec<PeerCredential>>,
}

struct PeerCredential {
    peer_id: String,
    token: Sensitive<String>,
}

/// Errors parsing per-peer bearer credentials.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum PeerCredentialError {
    /// An entry was not `peer-id=token`.
    #[error("A2A peer credentials must use peer-id=token entries")]
    Malformed,
    /// A peer id or token was empty.
    #[error("A2A peer credential ids and tokens must not be empty")]
    Empty,
    /// A peer id appeared more than once.
    #[error("each A2A peer id must appear once")]
    DuplicatePeer,
    /// A bearer token was assigned to more than one peer.
    #[error("each A2A bearer token must identify exactly one peer")]
    DuplicateToken,
}

impl PeerAuthenticator {
    /// Parses whitespace-separated `peer-id=token` credentials.
    ///
    /// # Errors
    ///
    /// Returns [`PeerCredentialError`] for malformed, empty, or duplicate
    /// entries.
    pub fn parse(raw: &str) -> Result<Self, PeerCredentialError> {
        let mut entries: Vec<PeerCredential> = Vec::new();
        for item in raw.split_whitespace() {
            let (peer_id, token) = item.split_once('=').ok_or(PeerCredentialError::Malformed)?;
            if peer_id.is_empty() || token.is_empty() {
                return Err(PeerCredentialError::Empty);
            }
            if entries.iter().any(|entry| entry.peer_id == peer_id) {
                return Err(PeerCredentialError::DuplicatePeer);
            }
            if entries.iter().any(|entry| entry.token.expose() == token) {
                return Err(PeerCredentialError::DuplicateToken);
            }
            entries.push(PeerCredential {
                peer_id: peer_id.to_owned(),
                token: Sensitive::new(token.to_owned()),
            });
        }
        Ok(Self {
            entries: Arc::new(entries),
        })
    }

    /// Builds the one-peer form used by focused tests and local wiring.
    ///
    /// # Errors
    ///
    /// Returns [`PeerCredentialError::Empty`] when either value is empty.
    pub fn single(peer_id: &str, token: &str) -> Result<Self, PeerCredentialError> {
        Self::parse(&format!("{peer_id}={token}"))
    }

    /// Reports whether no peer credential is configured.
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    fn authenticate(&self, provided: &str) -> Option<String> {
        // Compare every configured candidate. The token is secret while the
        // peer id is only its public routing identity, so returning early on
        // the matching entry would turn registry order into a timing oracle.
        // `parse` rejects duplicate tokens, hence at most one assignment wins.
        let mut authenticated = None;
        for entry in self.entries.iter() {
            let matches: bool = provided
                .as_bytes()
                .ct_eq(entry.token.expose().as_bytes())
                .into();
            if matches {
                authenticated = Some(entry.peer_id.clone());
            }
        }
        authenticated
    }
}

/// Build the axum router for the A2A surface.
pub fn router(state: AppState) -> Router {
    Router::new()
        .route("/.well-known/agent-card.json", get(agent_card))
        .route("/", post(json_rpc))
        .with_state(state)
}

/// Serve the signed Agent Card.
async fn agent_card(State(state): State<AppState>) -> impl IntoResponse {
    (
        [(header::CONTENT_TYPE, "application/json")],
        state.card.to_string(),
    )
}

/// Handle one A2A JSON-RPC request.
///
/// Authenticates the caller before touching the body as JSON — a rejected
/// request never reaches [`rpc::handle`]/[`rpc::handle_streaming`], so an
/// unattributed caller can never run a turn. `SendStreamingMessage`/
/// `TaskSubscription` are routed to the SSE transport
/// ([`rpc::is_streaming_method`] peeks the method name); every other method
/// stays the unary JSON response.
async fn json_rpc(State(state): State<AppState>, headers: HeaderMap, body: Bytes) -> Response {
    let peer_id = match authenticate(&state, &headers) {
        Ok(peer_id) => peer_id,
        Err(rejection) => return *rejection,
    };
    if rpc::is_streaming_method(&body) {
        let mut stream = rpc::handle_streaming_for_peer(state, &body, peer_id);
        if !rpc::requires_durable_marker(&body) {
            let stream = stream.map(|value| {
                Ok::<_, std::convert::Infallible>(Event::default().data(value.to_string()))
            });
            return Sse::new(stream)
                .keep_alive(KeepAlive::default())
                .into_response();
        }
        let mut before_receipt = Vec::new();
        let mut received = false;
        while let Some(value) = stream.next().await {
            if rpc::is_durable_marker(&value) {
                received = true;
                break;
            }
            before_receipt.push(value);
        }
        if !received {
            let refusal = before_receipt.pop().unwrap_or_else(|| {
                serde_json::json!({
                    "jsonrpc": "2.0",
                    "id": null,
                    "error": { "code": -32603, "message": "ingress ended without a durable receipt" }
                })
            });
            return Json(refusal).into_response();
        }
        let stream = futures::stream::iter(before_receipt)
            .chain(stream)
            .filter(|value| futures::future::ready(!rpc::is_durable_marker(value)))
            .map(|value| {
                Ok::<_, std::convert::Infallible>(Event::default().data(value.to_string()))
            });
        return Sse::new(stream)
            .keep_alive(KeepAlive::default())
            .into_response();
    }
    Json(rpc::handle_for_peer(&state, &body, &peer_id).await).into_response()
}

/// Verify the `Authorization: Bearer <token>` header against
/// the configured per-peer bearer credentials in constant time.
///
/// Returns the stable peer id on success. Returns a rejection response on
/// failure: `503` when the deployment has no credentials configured (refuse
/// all traffic rather than accept it unauthenticated), or `401` for a
/// missing, malformed, or unknown credential.
fn authenticate(state: &AppState, headers: &HeaderMap) -> Result<String, Box<Response>> {
    if state.peers.is_empty() {
        return Err(Box::new(
            (
                StatusCode::SERVICE_UNAVAILABLE,
                "A2A peer credentials are not configured",
            )
                .into_response(),
        ));
    }

    let Some(provided) = headers
        .get(header::AUTHORIZATION)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| value.strip_prefix("Bearer "))
    else {
        return Err(Box::new(
            (StatusCode::UNAUTHORIZED, "missing bearer token").into_response(),
        ));
    };
    state
        .peers
        .authenticate(provided)
        .ok_or_else(|| Box::new((StatusCode::UNAUTHORIZED, "invalid bearer token").into_response()))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn peer_credentials_refuse_shared_tokens_and_duplicate_ids() {
        assert_eq!(
            PeerAuthenticator::parse("peer-a=one peer-b=one").err(),
            Some(PeerCredentialError::DuplicateToken)
        );
        assert_eq!(
            PeerAuthenticator::parse("peer-a=one peer-a=two").err(),
            Some(PeerCredentialError::DuplicatePeer)
        );
    }
}