agent-protocols 0.6.0

use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine as _;
use ed25519_dalek::{Signature, Signer, SigningKey, VerifyingKey};
use serde::de::Error as DeError;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use serde_json::Value;
use sha3::{Digest, Sha3_256};
use std::collections::{BTreeMap, HashMap};
use std::fmt;
use std::str::FromStr;
use std::time::{SystemTime, UNIX_EPOCH};

use crate::error::{Result, SdkError};

pub const AGENT_ID_PREFIX: &str = "did:agent:";
pub const DEFAULT_LIVE_WRITE_WINDOW_MS: i64 = 300_000;
pub const DEFAULT_NONCE_TTL_MS: i64 = 300_000;
pub const DEFAULT_REQUEST_JWT_TTL_SECS: i64 = 300;
pub const MAX_NONCE_HEADER: &str = "Max-Seen-Nonce";
pub const MAX_SAFE_NONCE: u64 = 0x1FFFFFFFFFFFFF;

#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
pub struct AgentId(String);

impl AgentId {
    pub fn from_public_key(public_key: &[u8; 32]) -> Self {
        Self(format!(
            "{}{}",
            AGENT_ID_PREFIX,
            URL_SAFE_NO_PAD.encode(public_key)
        ))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }

    pub fn public_key_bytes(&self) -> Result<[u8; 32]> {
        let encoded = self
            .0
            .strip_prefix(AGENT_ID_PREFIX)
            .ok_or(SdkError::InvalidAgentIdPrefix)?;
        // URL_SAFE_NO_PAD rejects padding characters and non-zero trailing
        // bits, so non-canonical encodings of the same key are rejected.
        let bytes = URL_SAFE_NO_PAD.decode(encoded)?;
        bytes
            .try_into()
            .map_err(|bytes: Vec<u8>| SdkError::InvalidPublicKeyLength(bytes.len()))
    }

    pub fn verifying_key(&self) -> Result<VerifyingKey> {
        Ok(VerifyingKey::from_bytes(&self.public_key_bytes()?)?)
    }
}

impl fmt::Debug for AgentId {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.debug_tuple("AgentId").field(&self.0).finish()
    }
}

impl fmt::Display for AgentId {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.0)
    }
}

impl FromStr for AgentId {
    type Err = SdkError;

    fn from_str(value: &str) -> Result<Self> {
        if !value.starts_with(AGENT_ID_PREFIX) {
            return Err(SdkError::InvalidAgentIdPrefix);
        }
        let agent_id = Self(value.to_owned());
        agent_id.public_key_bytes()?;
        Ok(agent_id)
    }
}

impl TryFrom<String> for AgentId {
    type Error = SdkError;

    fn try_from(value: String) -> Result<Self> {
        Self::from_str(&value)
    }
}

impl Serialize for AgentId {
    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_str(&self.0)
    }
}

impl<'de> Deserialize<'de> for AgentId {
    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let value = String::deserialize(deserializer)?;
        Self::from_str(&value).map_err(D::Error::custom)
    }
}

pub struct AgentSigner {
    signing_key: SigningKey,
}

impl AgentSigner {
    pub fn generate() -> Self {
        let mut seed = [0_u8; 32];
        rand::fill(&mut seed);
        Self::from_seed(seed)
    }

    pub fn from_seed(seed: [u8; 32]) -> Self {
        Self {
            signing_key: SigningKey::from_bytes(&seed),
        }
    }

    pub fn agent_id(&self) -> AgentId {
        AgentId::from_public_key(&self.signing_key.verifying_key().to_bytes())
    }

    pub fn verifying_key(&self) -> VerifyingKey {
        self.signing_key.verifying_key()
    }

    pub fn sign_event<P>(&self, event: Event<P>) -> Result<Envelope<P>>
    where
        P: Serialize,
    {
        let event_hash = event_hash_bytes(&event)?;
        let hash = URL_SAFE_NO_PAD.encode(&event_hash);
        let signature = sign_event_hash(&self.signing_key, &event_hash)
            .expect("event hashes are always 32 bytes");
        Ok(Envelope {
            hash,
            event,
            signature,
        })
    }

    pub fn sign_request_jwt(&self, claims: &RequestJwtClaims) -> Result<String> {
        let agent_id = self.agent_id();
        if claims.iss != agent_id || claims.sub != agent_id {
            return Err(SdkError::InvalidJwtClaim("iss/sub"));
        }

        let header = RequestJwtHeader {
            alg: "EdDSA".to_owned(),
            typ: "JWT".to_owned(),
            kid: agent_id,
        };
        let header =
            URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header).expect("JWT header is JSON-safe"));
        let payload =
            URL_SAFE_NO_PAD.encode(serde_json::to_vec(claims).expect("JWT claims are JSON-safe"));
        let signing_input = format!("{header}.{payload}");
        let signature = self.signing_key.sign(signing_input.as_bytes());
        Ok(format!(
            "{signing_input}.{}",
            URL_SAFE_NO_PAD.encode(signature.to_bytes())
        ))
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct Event<P = Value> {
    pub protocol: String,
    #[serde(rename = "type")]
    pub kind: String,
    pub actor: AgentId,
    pub created_at: i64,
    pub nonce: u64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub room_id: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub base_seq: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub base_hash: Option<String>,
    #[serde(default, skip_serializing_if = "Vec::is_empty")]
    pub mentions: Vec<AgentId>,
    pub payload: P,
    #[serde(default, flatten)]
    pub extra: BTreeMap<String, Value>,
}

impl<P> Event<P> {
    pub fn new(
        protocol: impl Into<String>,
        kind: impl Into<String>,
        actor: AgentId,
        created_at: i64,
        nonce: u64,
        payload: P,
    ) -> Self {
        Self {
            protocol: protocol.into(),
            kind: kind.into(),
            actor,
            created_at,
            nonce,
            room_id: None,
            base_seq: None,
            base_hash: None,
            mentions: Vec::new(),
            payload,
            extra: BTreeMap::new(),
        }
    }

    pub fn with_room_id(mut self, room_id: impl Into<String>) -> Self {
        self.room_id = Some(room_id.into());
        self
    }

    pub fn with_room_head(mut self, base_seq: u64, base_hash: impl Into<String>) -> Self {
        self.base_seq = Some(base_seq);
        self.base_hash = Some(base_hash.into());
        self
    }

    pub fn with_mentions(mut self, mentions: Vec<AgentId>) -> Self {
        self.mentions = mentions;
        self
    }

    pub fn with_mention(mut self, agent_id: AgentId) -> Self {
        self.mentions.push(agent_id);
        self
    }
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct Envelope<P = Value> {
    pub hash: String,
    pub event: Event<P>,
    pub signature: String,
}

pub fn canonical_event_bytes<P>(event: &Event<P>) -> Result<Vec<u8>>
where
    P: Serialize,
{
    serde_jcs::to_vec(event).map_err(|err| SdkError::CanonicalJson(err.to_string()))
}

pub fn event_hash<P>(event: &Event<P>) -> Result<String>
where
    P: Serialize,
{
    Ok(URL_SAFE_NO_PAD.encode(event_hash_bytes(event)?))
}

pub fn event_hash_bytes<P>(event: &Event<P>) -> Result<[u8; 32]>
where
    P: Serialize,
{
    validate_nonce(event.nonce)?;
    let digest = Sha3_256::digest(canonical_event_bytes(event)?);
    let mut bytes = [0_u8; 32];
    bytes.copy_from_slice(&digest);
    Ok(bytes)
}

pub fn sign_event<P>(signing_key: &SigningKey, event: &Event<P>) -> Result<String>
where
    P: Serialize,
{
    let event_hash = event_hash_bytes(event)?;
    Ok(sign_event_hash(signing_key, &event_hash).expect("event hashes are always 32 bytes"))
}

/// Signs a precomputed 32-byte event hash. Signing a digest supplied by
/// another component without seeing the event it commits to (blind signing) is
/// NOT RECOMMENDED: `actor` and all event content are inside the digest, so a
/// blind signer can be tricked into signing arbitrary events attributed to its
/// key. Prefer [`sign_event`], which canonicalizes and hashes the event.
pub fn sign_event_hash(signing_key: &SigningKey, event_hash: &[u8]) -> Result<String> {
    let signature = signing_key.sign(valid_event_hash_bytes(event_hash)?);
    Ok(URL_SAFE_NO_PAD.encode(signature.to_bytes()))
}

pub fn verify_event_hash<P>(envelope: &Envelope<P>) -> Result<()>
where
    P: Serialize,
{
    let expected = event_hash(&envelope.event)?;
    if expected == envelope.hash {
        Ok(())
    } else {
        Err(SdkError::InvalidEventHash {
            expected,
            actual: envelope.hash.clone(),
        })
    }
}

pub fn verify_signature<P>(envelope: &Envelope<P>) -> Result<()>
where
    P: Serialize,
{
    let event_hash = event_hash_bytes(&envelope.event)?;
    let verifying_key = envelope.event.actor.verifying_key()?;
    verify_event_hash_signature(&verifying_key, &event_hash, &envelope.signature)
}

pub fn verify_event_hash_signature(
    verifying_key: &VerifyingKey,
    event_hash: &[u8],
    encoded_signature: &str,
) -> Result<()> {
    let signature_bytes = URL_SAFE_NO_PAD.decode(encoded_signature)?;
    let signature_bytes: [u8; 64] = signature_bytes
        .try_into()
        .map_err(|bytes: Vec<u8>| SdkError::InvalidSignatureLength(bytes.len()))?;
    let signature = Signature::from_bytes(&signature_bytes);
    verifying_key.verify_strict(valid_event_hash_bytes(event_hash)?, &signature)?;
    Ok(())
}

fn valid_event_hash_bytes(event_hash: &[u8]) -> Result<&[u8; 32]> {
    if event_hash.len() != 32 {
        return Err(SdkError::InvalidEventHashLength(event_hash.len()));
    }
    Ok(event_hash
        .try_into()
        .expect("event hash length was checked above"))
}

pub fn verify_envelope<P>(envelope: &Envelope<P>) -> Result<()>
where
    P: Serialize,
{
    verify_event_hash(envelope)?;
    verify_signature(envelope)
}

pub fn verify_timestamp(created_at: i64, now_ms: i64, window_ms: i64) -> Result<()> {
    if window_ms < 0 || created_at.abs_diff(now_ms) > window_ms as u64 {
        Err(SdkError::TimestampOutOfWindow)
    } else {
        Ok(())
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct NonceRecord {
    pub max_nonce: u64,
    pub expires_at: i64,
}

pub trait NonceStore {
    fn check_and_update(
        &mut self,
        actor: &AgentId,
        nonce: u64,
        now_ms: i64,
        ttl_ms: i64,
    ) -> Result<u64>;

    fn max_nonce(&self, actor: &AgentId, now_ms: i64) -> Option<u64>;
}

#[derive(Clone, Debug, Default)]
pub struct MemoryNonceStore {
    records: HashMap<AgentId, NonceRecord>,
}

impl MemoryNonceStore {
    pub fn new() -> Self {
        Self::default()
    }
}

impl NonceStore for MemoryNonceStore {
    fn check_and_update(
        &mut self,
        actor: &AgentId,
        nonce: u64,
        now_ms: i64,
        ttl_ms: i64,
    ) -> Result<u64> {
        validate_nonce(nonce)?;
        if ttl_ms < 0 {
            return Err(SdkError::InvalidNonce(
                "nonce cache ttl must be non-negative".to_owned(),
            ));
        }

        if let Some(record) = self.records.get(actor) {
            if record.expires_at > now_ms && nonce <= record.max_nonce {
                return Err(SdkError::NonceNotGreater {
                    max_nonce: record.max_nonce,
                });
            }
        }

        let record = NonceRecord {
            max_nonce: nonce,
            expires_at: now_ms.saturating_add(ttl_ms),
        };
        self.records.insert(actor.clone(), record);
        Ok(record.max_nonce)
    }

    fn max_nonce(&self, actor: &AgentId, now_ms: i64) -> Option<u64> {
        self.records
            .get(actor)
            .filter(|record| record.expires_at > now_ms)
            .map(|record| record.max_nonce)
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ClientNonceManager {
    next_nonce: u64,
}

impl Default for ClientNonceManager {
    fn default() -> Self {
        Self { next_nonce: 1 }
    }
}

impl ClientNonceManager {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn with_next(next_nonce: u64) -> Result<Self> {
        validate_nonce(next_nonce)?;
        Ok(Self { next_nonce })
    }

    pub fn peek(&self) -> u64 {
        self.next_nonce
    }

    pub fn next_nonce(&mut self) -> Result<u64> {
        let nonce = self.next_nonce;
        validate_nonce(nonce)?;
        self.next_nonce += 1;
        Ok(nonce)
    }

    pub fn observe_max_nonce(&mut self, max_nonce: u64) {
        if max_nonce >= self.next_nonce {
            self.next_nonce = max_nonce.saturating_add(1);
        }
    }

    pub fn observe_max_nonce_header(&mut self, value: &str) -> Result<()> {
        let max_nonce = value
            .parse::<u64>()
            .map_err(|_| SdkError::InvalidNonce("invalid max nonce header".to_owned()))?;
        validate_nonce(max_nonce)?;
        self.observe_max_nonce(max_nonce);
        Ok(())
    }
}

pub fn validate_nonce(nonce: u64) -> Result<()> {
    if nonce == 0 || nonce > MAX_SAFE_NONCE {
        // Number.MAX_SAFE_INTEGER in JavaScript, to prevent interoperability issues with JS clients using Number for nonces
        Err(SdkError::InvalidNonce(
            "nonce must be a positive integer less than or equal to 9007199254740991".to_owned(),
        ))
    } else {
        Ok(())
    }
}

#[derive(Clone, Debug)]
pub struct LiveWriteOptions {
    pub now_ms: i64,
    pub window_ms: i64,
    pub nonce_ttl_ms: i64,
}

impl Default for LiveWriteOptions {
    fn default() -> Self {
        Self {
            now_ms: unix_ms(),
            window_ms: DEFAULT_LIVE_WRITE_WINDOW_MS,
            nonce_ttl_ms: DEFAULT_NONCE_TTL_MS,
        }
    }
}

pub fn verify_live_envelope<P, S>(
    envelope: &Envelope<P>,
    options: &LiveWriteOptions,
    nonce_store: &mut S,
) -> Result<u64>
where
    P: Serialize,
    S: NonceStore,
{
    verify_envelope(envelope)?;
    verify_timestamp(envelope.event.created_at, options.now_ms, options.window_ms)?;
    nonce_store.check_and_update(
        &envelope.event.actor,
        envelope.event.nonce,
        options.now_ms,
        options.nonce_ttl_ms,
    )
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct RequestJwtHeader {
    pub alg: String,
    pub typ: String,
    pub kid: AgentId,
}

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct RequestJwtClaims {
    pub iss: AgentId,
    pub sub: AgentId,
    /// Origin of the receiving service: scheme, host, and non-default port
    /// with no path. All Agent Protocols endpoints on one origin share this
    /// value; derive it with [`service_origin`].
    pub aud: String,
    pub iat: i64,
    pub exp: i64,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RequestBinding {
    pub audience: String,
}

impl RequestBinding {
    pub fn new(audience: impl Into<String>) -> Self {
        Self {
            audience: audience.into(),
        }
    }
}

impl RequestJwtClaims {
    pub fn new(agent_id: AgentId, binding: RequestBinding, issued_at: i64, ttl_secs: i64) -> Self {
        Self {
            iss: agent_id.clone(),
            sub: agent_id,
            aud: binding.audience,
            iat: issued_at,
            exp: issued_at + ttl_secs,
        }
    }
}

#[derive(Clone, Debug)]
pub struct RequestAuthContext {
    pub audience: String,
    pub now_secs: i64,
    pub max_ttl_secs: i64,
}

impl RequestAuthContext {
    pub fn new(audience: impl Into<String>) -> Self {
        let binding = RequestBinding::new(audience);
        Self {
            audience: binding.audience,
            now_secs: unix_secs(),
            max_ttl_secs: DEFAULT_REQUEST_JWT_TTL_SECS,
        }
    }
}

pub fn verify_request_jwt(token: &str, context: &RequestAuthContext) -> Result<RequestJwtClaims> {
    let parts: Vec<&str> = token.split('.').collect();
    if parts.len() != 3 {
        return Err(SdkError::InvalidJwt(
            "expected three compact JWS parts".into(),
        ));
    }

    let header: RequestJwtHeader = serde_json::from_slice(&URL_SAFE_NO_PAD.decode(parts[0])?)?;
    let claims: RequestJwtClaims = serde_json::from_slice(&URL_SAFE_NO_PAD.decode(parts[1])?)?;
    let signature_bytes = URL_SAFE_NO_PAD.decode(parts[2])?;
    let signature_bytes: [u8; 64] = signature_bytes
        .try_into()
        .map_err(|bytes: Vec<u8>| SdkError::InvalidSignatureLength(bytes.len()))?;
    let signature = Signature::from_bytes(&signature_bytes);
    let signing_input = format!("{}.{}", parts[0], parts[1]);

    if header.alg != "EdDSA" {
        return Err(SdkError::InvalidJwtClaim("alg"));
    }
    if header.typ != "JWT" {
        return Err(SdkError::InvalidJwtClaim("typ"));
    }
    if header.kid != claims.iss || claims.iss != claims.sub {
        return Err(SdkError::InvalidJwtClaim("kid/iss/sub"));
    }

    header
        .kid
        .verifying_key()?
        .verify_strict(signing_input.as_bytes(), &signature)?;

    if claims.aud != context.audience {
        return Err(SdkError::InvalidJwtClaim("aud"));
    }
    if claims.exp <= claims.iat {
        return Err(SdkError::InvalidJwtClaim("exp/iat"));
    }
    if claims.iat > context.now_secs || claims.exp < context.now_secs {
        return Err(SdkError::InvalidJwtClaim("iat/exp"));
    }
    if claims.exp - claims.iat > context.max_ttl_secs {
        return Err(SdkError::InvalidJwtClaim("ttl"));
    }

    Ok(claims)
}

/// Derives the request JWT `aud` from a request URL: the service origin —
/// scheme, host, and non-default port, with no path (Agent Identity Section 8).
pub fn service_origin(url: &str) -> Result<String> {
    let parsed = url::Url::parse(url)
        .map_err(|_| SdkError::InvalidPayload(format!("not a valid URL: {url}")))?;
    if parsed.scheme() != "https" && parsed.scheme() != "http" {
        return Err(SdkError::InvalidPayload(format!(
            "not an HTTP(S) URL: {url}"
        )));
    }
    Ok(parsed.origin().ascii_serialization())
}

pub fn unix_ms() -> i64 {
    let duration = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default();
    duration.as_millis().min(i64::MAX as u128) as i64
}

pub fn unix_secs() -> i64 {
    unix_ms() / 1000
}

#[cfg(test)]
mod tests {
    use serde::ser;
    use serde_json::json;

    use super::*;

    struct FailingPayload;

    impl Serialize for FailingPayload {
        fn serialize<S>(&self, _serializer: S) -> std::result::Result<S::Ok, S::Error>
        where
            S: Serializer,
        {
            Err(ser::Error::custom("payload cannot be serialized"))
        }
    }

    #[test]
    fn signs_and_verifies_event_envelope() {
        let signer = AgentSigner::from_seed([7; 32]);
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1_779_753_600_000,
            1,
            json!({"id": signer.agent_id(), "name": "ResearchAgent"}),
        );

        let envelope = signer.sign_event(event).unwrap();

        assert!(!envelope.hash.starts_with("evt_"));
        assert_eq!(envelope.hash.len(), 43);
        verify_envelope(&envelope).unwrap();
    }

    #[test]
    fn signs_and_verifies_raw_event_hash_bytes() {
        let seed = [18; 32];
        let signer = AgentSigner::from_seed(seed);
        let signing_key = SigningKey::from_bytes(&seed);
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1_779_753_600_000,
            1,
            json!({"id": signer.agent_id(), "name": "ResearchAgent"}),
        );

        let digest = event_hash_bytes(&event).unwrap();
        let signature = sign_event_hash(&signing_key, &digest).unwrap();
        let envelope = signer.sign_event(event).unwrap();

        assert_eq!(envelope.signature, signature);
        verify_event_hash_signature(&signing_key.verifying_key(), &digest, &signature).unwrap();
    }

    #[test]
    fn rejects_tampered_event_payload() {
        let signer = AgentSigner::from_seed([8; 32]);
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1_779_753_600_000,
            1,
            json!({"name": "before"}),
        );
        let mut envelope = signer.sign_event(event).unwrap();
        envelope.event.payload = json!({"name": "after"});

        assert!(verify_envelope(&envelope).is_err());
    }

    #[test]
    fn rejects_non_canonical_event_payloads() {
        let signer = AgentSigner::from_seed([52; 32]);
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1000,
            1,
            FailingPayload,
        );
        assert!(matches!(
            canonical_event_bytes(&event),
            Err(SdkError::CanonicalJson(_))
        ));
        assert!(matches!(
            event_hash(&event),
            Err(SdkError::CanonicalJson(_))
        ));
        let envelope = Envelope {
            hash: "h".to_owned(),
            event,
            signature: "s".to_owned(),
        };
        assert!(matches!(
            verify_event_hash(&envelope),
            Err(SdkError::CanonicalJson(_))
        ));
        assert!(matches!(
            verify_signature(&envelope),
            Err(SdkError::CanonicalJson(_))
        ));
        assert!(matches!(
            verify_envelope(&envelope),
            Err(SdkError::CanonicalJson(_))
        ));
    }

    #[test]
    fn rejects_nonce_reuse() {
        let signer = AgentSigner::from_seed([9; 32]);
        let options = LiveWriteOptions {
            now_ms: 1000,
            window_ms: 1000,
            nonce_ttl_ms: 1000,
        };
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1000,
            1,
            json!({"name": "ResearchAgent"}),
        );
        let envelope = signer.sign_event(event).unwrap();
        let mut store = MemoryNonceStore::new();

        let max_nonce = verify_live_envelope(&envelope, &options, &mut store).unwrap();
        assert_eq!(max_nonce, 1);
        assert!(matches!(
            verify_live_envelope(&envelope, &options, &mut store),
            Err(SdkError::NonceNotGreater { max_nonce: 1 })
        ));
    }

    #[test]
    fn client_nonce_manager_observes_server_max() {
        let mut manager = ClientNonceManager::new();

        assert_eq!(manager.next_nonce().unwrap(), 1);
        manager.observe_max_nonce(5);

        assert_eq!(manager.peek(), 6);
        manager.observe_max_nonce(4);
        assert_eq!(manager.peek(), 6);
        assert_eq!(manager.next_nonce().unwrap(), 6);

        let mut invalid = ClientNonceManager {
            next_nonce: MAX_SAFE_NONCE + 1,
        };
        assert!(matches!(
            invalid.next_nonce(),
            Err(SdkError::InvalidNonce(_))
        ));
    }

    #[test]
    fn rejects_nonce_values_outside_safe_json_integer_range() {
        let signer = AgentSigner::from_seed([16; 32]);
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1000,
            MAX_SAFE_NONCE + 1,
            json!({"id": signer.agent_id(), "name": "ResearchAgent"}),
        );

        assert!(matches!(
            signer.sign_event(event),
            Err(SdkError::InvalidNonce(_))
        ));
    }

    #[test]
    fn signs_and_verifies_request_jwt() {
        let signer = AgentSigner::from_seed([10; 32]);
        let claims = RequestJwtClaims::new(
            signer.agent_id(),
            RequestBinding::new("https://api.example.com"),
            100,
            300,
        );
        let token = signer.sign_request_jwt(&claims).unwrap();
        let context = RequestAuthContext {
            audience: "https://api.example.com".to_owned(),
            now_secs: 120,
            max_ttl_secs: 300,
        };
        let verified = verify_request_jwt(&token, &context).unwrap();

        assert_eq!(verified.iss, signer.agent_id());
    }

    #[test]
    fn rejects_request_jwts_with_non_positive_ttl() {
        let signer = AgentSigner::from_seed([17; 32]);
        let claims = RequestJwtClaims::new(
            signer.agent_id(),
            RequestBinding::new("https://api.example.com"),
            100,
            0,
        );
        let token = signer.sign_request_jwt(&claims).unwrap();
        let context = RequestAuthContext {
            audience: "https://api.example.com".to_owned(),
            now_secs: 100,
            max_ttl_secs: 300,
        };

        assert!(matches!(
            verify_request_jwt(&token, &context),
            Err(SdkError::InvalidJwtClaim("exp/iat"))
        ));
    }

    #[test]
    fn agent_id_string_conversions() {
        let signer = AgentSigner::from_seed([20; 32]);
        let agent_id = signer.agent_id();
        let text = agent_id.to_string();

        assert_eq!(agent_id.as_str(), text);
        assert!(format!("{agent_id:?}").starts_with("AgentId("));
        assert_eq!(AgentId::try_from(text.clone()).unwrap(), agent_id);
        assert_eq!(text.parse::<AgentId>().unwrap(), agent_id);
        assert!(matches!(
            "not-a-did".parse::<AgentId>(),
            Err(SdkError::InvalidAgentIdPrefix)
        ));
        assert!(matches!(
            AgentId("not-a-did".to_owned()).public_key_bytes(),
            Err(SdkError::InvalidAgentIdPrefix)
        ));
        assert!(AgentId(format!("{AGENT_ID_PREFIX}%"))
            .public_key_bytes()
            .is_err());
        let short = AgentId(format!(
            "{}{}",
            AGENT_ID_PREFIX,
            URL_SAFE_NO_PAD.encode([1_u8; 31])
        ));
        assert!(matches!(
            short.public_key_bytes(),
            Err(SdkError::InvalidPublicKeyLength(31))
        ));
        assert!(short.verifying_key().is_err());
        assert!(format!("{}%", AGENT_ID_PREFIX).parse::<AgentId>().is_err());
    }

    #[test]
    fn agent_signer_exposes_keys_and_generates() {
        let signer = AgentSigner::from_seed([21; 32]);
        assert_eq!(
            signer.verifying_key().to_bytes(),
            signer.agent_id().public_key_bytes().unwrap()
        );

        let generated = AgentSigner::generate();
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            generated.agent_id(),
            1000,
            1,
            json!({"name": "Generated"}),
        );
        verify_envelope(&generated.sign_event(event).unwrap()).unwrap();
    }

    #[test]
    fn sign_request_jwt_rejects_foreign_claims() {
        let signer = AgentSigner::from_seed([22; 32]);
        let other = AgentSigner::from_seed([23; 32]);
        let claims = RequestJwtClaims {
            iss: other.agent_id(),
            sub: other.agent_id(),
            aud: "https://api.example.com".to_owned(),
            iat: 100,
            exp: 400,
        };
        assert!(matches!(
            signer.sign_request_jwt(&claims),
            Err(SdkError::InvalidJwtClaim("iss/sub"))
        ));
    }

    #[test]
    fn free_sign_event_matches_signer() {
        let seed = [24; 32];
        let signer = AgentSigner::from_seed(seed);
        let signing_key = SigningKey::from_bytes(&seed);
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1000,
            1,
            json!({"name": "Agent"}),
        );

        let signature = sign_event(&signing_key, &event).unwrap();
        assert_eq!(signer.sign_event(event).unwrap().signature, signature);
    }

    #[test]
    fn rejects_event_hash_with_wrong_length() {
        let signing_key = SigningKey::from_bytes(&[25; 32]);
        assert!(matches!(
            sign_event_hash(&signing_key, &[0_u8; 10]),
            Err(SdkError::InvalidEventHashLength(10))
        ));
        assert!(matches!(
            verify_event_hash_signature(&signing_key.verifying_key(), &[0_u8; 10], "AAAA"),
            Err(SdkError::InvalidSignatureLength(_))
        ));
        let signature = URL_SAFE_NO_PAD.encode([0_u8; 64]);
        assert!(matches!(
            verify_event_hash_signature(&signing_key.verifying_key(), &[0_u8; 10], &signature),
            Err(SdkError::InvalidEventHashLength(10))
        ));
    }

    #[test]
    fn verify_timestamp_rejects_drift_and_negative_window() {
        verify_timestamp(1000, 1200, 1000).unwrap();
        assert!(matches!(
            verify_timestamp(0, 1_000_000, 1000),
            Err(SdkError::TimestampOutOfWindow)
        ));
        assert!(matches!(
            verify_timestamp(100, 100, -1),
            Err(SdkError::TimestampOutOfWindow)
        ));
    }

    #[test]
    fn verify_live_envelope_propagates_envelope_and_timestamp_errors() {
        let signer = AgentSigner::from_seed([53; 32]);
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1000,
            1,
            json!({"name": "ResearchAgent"}),
        );
        let mut envelope = signer.sign_event(event).unwrap();
        let mut store = MemoryNonceStore::new();
        let options = LiveWriteOptions {
            now_ms: 1000,
            window_ms: 1000,
            nonce_ttl_ms: 1000,
        };

        envelope.signature = "invalid".to_owned();
        assert!(verify_live_envelope(&envelope, &options, &mut store).is_err());

        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            0,
            2,
            json!({"name": "ResearchAgent"}),
        );
        let envelope = signer.sign_event(event).unwrap();
        let expired_options = LiveWriteOptions {
            now_ms: 3000,
            ..options
        };
        assert!(matches!(
            verify_live_envelope(&envelope, &expired_options, &mut store),
            Err(SdkError::TimestampOutOfWindow)
        ));
    }

    #[test]
    fn memory_nonce_store_rejects_negative_ttl_and_expires() {
        let actor = AgentSigner::from_seed([26; 32]).agent_id();
        let mut store = MemoryNonceStore::new();
        assert!(matches!(
            store.check_and_update(&actor, 1, 1000, -1),
            Err(SdkError::InvalidNonce(_))
        ));

        store.check_and_update(&actor, 4, 1000, 1000).unwrap();
        // a strictly greater nonce is accepted while the record is still live
        assert_eq!(store.check_and_update(&actor, 5, 1100, 1000).unwrap(), 5);
        assert_eq!(store.max_nonce(&actor, 1500), Some(5));
        // record has expired by now_ms past expires_at
        assert_eq!(store.max_nonce(&actor, 5000), None);
        assert_eq!(store.max_nonce(&actor, 1500), Some(5));
        // a fresh actor has no record at all
        let stranger = AgentSigner::from_seed([27; 32]).agent_id();
        assert_eq!(store.max_nonce(&stranger, 1500), None);
    }

    #[test]
    fn client_nonce_manager_with_next_and_header() {
        assert!(ClientNonceManager::with_next(0).is_err());
        let mut manager = ClientNonceManager::with_next(5).unwrap();
        assert_eq!(manager.peek(), 5);

        manager.observe_max_nonce_header("10").unwrap();
        assert_eq!(manager.peek(), 11);
        assert!(matches!(
            manager.observe_max_nonce_header("not-a-number"),
            Err(SdkError::InvalidNonce(_))
        ));
        assert!(matches!(
            manager.observe_max_nonce_header("9007199254740992"),
            Err(SdkError::InvalidNonce(_))
        ));
    }

    #[test]
    fn default_options_and_context_use_real_clock() {
        let options = LiveWriteOptions::default();
        assert_eq!(options.window_ms, DEFAULT_LIVE_WRITE_WINDOW_MS);
        assert_eq!(options.nonce_ttl_ms, DEFAULT_NONCE_TTL_MS);
        assert!(options.now_ms > 0);

        let context = RequestAuthContext::new("https://api.example.com");
        assert_eq!(context.audience, "https://api.example.com");
        assert_eq!(context.max_ttl_secs, DEFAULT_REQUEST_JWT_TTL_SECS);
        assert_eq!(context.now_secs, unix_secs());
        assert!(unix_ms() >= context.now_secs * 1000);
    }

    fn encode_token(header: Value, claims: Value) -> String {
        let header = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header).unwrap());
        let claims = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&claims).unwrap());
        let signature = URL_SAFE_NO_PAD.encode([0_u8; 64]);
        format!("{header}.{claims}.{signature}")
    }

    #[test]
    fn verify_request_jwt_rejects_malformed_tokens() {
        let signer = AgentSigner::from_seed([28; 32]);
        let agent_id = signer.agent_id();
        let other = AgentSigner::from_seed([29; 32]).agent_id();
        let context = RequestAuthContext {
            audience: "https://api.example.com".to_owned(),
            now_secs: 200,
            max_ttl_secs: 300,
        };
        let valid_claims = json!({
            "iss": agent_id,
            "sub": agent_id,
            "aud": "https://api.example.com",
            "iat": 100,
            "exp": 400,
        });

        assert!(matches!(
            verify_request_jwt("only.two", &context),
            Err(SdkError::InvalidJwt(_))
        ));
        assert!(verify_request_jwt("%.e30.AAAA", &context).is_err());
        let invalid_header = format!(
            "{}.{}.{}",
            URL_SAFE_NO_PAD.encode(b"not-json"),
            URL_SAFE_NO_PAD.encode(serde_json::to_vec(&valid_claims).unwrap()),
            URL_SAFE_NO_PAD.encode([0_u8; 64])
        );
        assert!(verify_request_jwt(&invalid_header, &context).is_err());
        let header = URL_SAFE_NO_PAD.encode(
            serde_json::to_vec(&json!({"alg": "EdDSA", "typ": "JWT", "kid": agent_id})).unwrap(),
        );
        let claims = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&valid_claims).unwrap());
        assert!(verify_request_jwt(&format!("{header}.%.AAAA"), &context).is_err());
        assert!(verify_request_jwt(&format!("{header}.{claims}.%"), &context).is_err());
        assert!(matches!(
            verify_request_jwt(&format!("{header}.{claims}.AAAA"), &context),
            Err(SdkError::InvalidSignatureLength(_))
        ));
        assert!(verify_request_jwt(
            &format!("{header}.{claims}.{}", URL_SAFE_NO_PAD.encode([0_u8; 64])),
            &context
        )
        .is_err());
        let invalid_key = AgentId::from_public_key(&[255_u8; 32]);
        assert!(verify_request_jwt(
            &encode_token(
                json!({"alg": "EdDSA", "typ": "JWT", "kid": invalid_key}),
                json!({
                    "iss": invalid_key,
                    "sub": invalid_key,
                    "aud": "https://api.example.com",
                    "iat": 100,
                    "exp": 400,
                }),
            ),
            &context,
        )
        .is_err());
        assert!(matches!(
            verify_request_jwt(
                &encode_token(
                    json!({"alg": "HS256", "typ": "JWT", "kid": agent_id}),
                    valid_claims.clone(),
                ),
                &context,
            ),
            Err(SdkError::InvalidJwtClaim("alg"))
        ));
        assert!(matches!(
            verify_request_jwt(
                &encode_token(
                    json!({"alg": "EdDSA", "typ": "JWS", "kid": agent_id}),
                    valid_claims.clone(),
                ),
                &context,
            ),
            Err(SdkError::InvalidJwtClaim("typ"))
        ));
        assert!(matches!(
            verify_request_jwt(
                &encode_token(
                    json!({"alg": "EdDSA", "typ": "JWT", "kid": other}),
                    valid_claims,
                ),
                &context,
            ),
            Err(SdkError::InvalidJwtClaim("kid/iss/sub"))
        ));
    }

    #[test]
    fn rejects_non_canonical_base64url_encodings() {
        let signer = AgentSigner::from_seed([62; 32]);
        let agent_id = signer.agent_id();
        let suffix = agent_id
            .as_str()
            .strip_prefix(AGENT_ID_PREFIX)
            .unwrap()
            .to_owned();

        // Padding characters are rejected even when the decoded bytes match.
        let padded = AgentId(format!("{AGENT_ID_PREFIX}{suffix}=="));
        assert!(padded.public_key_bytes().is_err());

        // Non-zero trailing bits give the same key a second string form:
        // replace the final character with one sharing its used bits but
        // different trailing bits (the last char of a 32-byte value uses 4 of
        // its 6 bits).
        const ALPHABET: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
        let last = suffix.as_bytes()[suffix.len() - 1];
        let index = ALPHABET.iter().position(|&c| c == last).unwrap();
        let non_canonical = ALPHABET[(index & !3) | ((index & 3) ^ 1)] as char;
        let mut tweaked = suffix[..suffix.len() - 1].to_owned();
        tweaked.push(non_canonical);
        let tweaked = AgentId(format!("{AGENT_ID_PREFIX}{tweaked}"));
        assert!(tweaked.public_key_bytes().is_err());

        // Signatures must be canonical base64url too.
        let event = Event::new(
            "agent-profile/1.0",
            "profile.update",
            signer.agent_id(),
            1000,
            1,
            json!({"name": "Agent"}),
        );
        let mut envelope = signer.sign_event(event).unwrap();
        envelope.signature.push('=');
        assert!(verify_signature(&envelope).is_err());
    }

    #[test]
    fn service_origin_derives_request_jwt_audience() {
        assert_eq!(
            service_origin("https://api.example.com/v1/rooms/room1").unwrap(),
            "https://api.example.com"
        );
        assert_eq!(
            service_origin("https://api.example.com:8443/path?q=1").unwrap(),
            "https://api.example.com:8443"
        );
        assert_eq!(
            service_origin("https://API.Example.com:443/").unwrap(),
            "https://api.example.com"
        );
        assert!(service_origin("not a url").is_err());
        assert!(service_origin("ftp://example.com").is_err());
    }

    #[test]
    fn verify_request_jwt_enforces_audience_and_window() {
        let signer = AgentSigner::from_seed([30; 32]);

        let wrong_audience = signer
            .sign_request_jwt(&RequestJwtClaims::new(
                signer.agent_id(),
                RequestBinding::new("https://api.example.com"),
                100,
                300,
            ))
            .unwrap();
        let context = RequestAuthContext {
            audience: "https://other.example.com".to_owned(),
            now_secs: 200,
            max_ttl_secs: 300,
        };
        assert!(matches!(
            verify_request_jwt(&wrong_audience, &context),
            Err(SdkError::InvalidJwtClaim("aud"))
        ));

        let expired = signer
            .sign_request_jwt(&RequestJwtClaims::new(
                signer.agent_id(),
                RequestBinding::new("https://api.example.com"),
                1000,
                300,
            ))
            .unwrap();
        let later = RequestAuthContext {
            audience: "https://api.example.com".to_owned(),
            now_secs: 5000,
            max_ttl_secs: 300,
        };
        assert!(matches!(
            verify_request_jwt(&expired, &later),
            Err(SdkError::InvalidJwtClaim("iat/exp"))
        ));

        let long_ttl = signer
            .sign_request_jwt(&RequestJwtClaims::new(
                signer.agent_id(),
                RequestBinding::new("https://api.example.com"),
                100,
                400,
            ))
            .unwrap();
        let strict = RequestAuthContext {
            audience: "https://api.example.com".to_owned(),
            now_secs: 200,
            max_ttl_secs: 300,
        };
        assert!(matches!(
            verify_request_jwt(&long_ttl, &strict),
            Err(SdkError::InvalidJwtClaim("ttl"))
        ));
    }
}