klieo-a2a 3.3.0

//! `A2aServer` end-to-end dispatch over `MemoryBus`.

use bytes::Bytes;
use klieo_a2a::envelope::{JsonRpcRequest, JsonRpcResponse};
use klieo_a2a::handler::EchoHandler;
use klieo_a2a::server::A2aServer;
use klieo_auth_common::AllowAnonymous;
use klieo_bus_memory::MemoryBus;
use klieo_core::{DurableName, Headers};
use serde_json::json;
use std::sync::Arc;
use std::time::Duration;
use tokio_stream::StreamExt;

const APP: &str = "test";
const AGENT: &str = "agent-x";

fn rpc_subject() -> String {
    format!("{APP}.a2a.{AGENT}.rpc")
}

async fn run_server(bus: Arc<MemoryBus>) {
    let server = A2aServer::new(
        APP.into(),
        AGENT.into(),
        Arc::new(EchoHandler::default()),
        Arc::new(AllowAnonymous),
        bus.pubsub.clone(),
    );
    tokio::spawn(async move {
        server.run().await.expect("server.run");
    });
    // Give the server time to subscribe before any publish.
    tokio::time::sleep(Duration::from_millis(50)).await;
}

async fn round_trip(bus: &Arc<MemoryBus>, payload: Bytes, reply_to: &str) -> JsonRpcResponse {
    let mut reply_stream = bus
        .pubsub
        .subscribe(reply_to, DurableName::new("test-reply"))
        .await
        .unwrap();
    let mut headers = Headers::new();
    headers.insert("Reply-To".into(), reply_to.into());
    headers.insert("A2A-Version".into(), "1.0".into());
    bus.pubsub
        .publish(&rpc_subject(), payload, headers)
        .await
        .unwrap();
    let msg = tokio::time::timeout(Duration::from_secs(2), reply_stream.next())
        .await
        .expect("reply timeout")
        .expect("stream ended")
        .expect("bus error");
    serde_json::from_slice(&msg.payload).expect("decode response")
}

#[tokio::test]
async fn send_message_round_trips() {
    use klieo_a2a::client::A2aClient;
    use klieo_a2a::types::{Message, Part, Role, SendMessageResult, TaskStatus};
    let bus = Arc::new(MemoryBus::new());
    run_server(bus.clone()).await;
    let message = Message {
        messageId: "m-1".into(),
        contextId: Some("c-1".into()),
        taskId: None,
        role: Role::User,
        parts: vec![Part::Text {
            content: "hi".into(),
            metadata: None,
            mediaType: None,
            filename: None,
        }],
        metadata: None,
        extensions: vec![],
        referenceTaskIds: vec![],
    };
    let client = A2aClient::new(bus.pubsub.clone(), APP);
    let result = client
        .send_message(AGENT, message, Duration::from_secs(2))
        .await
        .expect("send_message");
    match result {
        SendMessageResult::Task(task) => assert!(matches!(task.status, TaskStatus::Completed)),
        SendMessageResult::Message(_) => panic!("EchoHandler returns a Task"),
    }
}

#[tokio::test]
async fn unknown_method_returns_minus_32601() {
    let bus = Arc::new(MemoryBus::new());
    run_server(bus.clone()).await;
    let req = JsonRpcRequest {
        jsonrpc: "2.0".into(),
        id: json!(2),
        method: "Frobnicate".into(),
        params: json!({}),
    };
    let resp = round_trip(
        &bus,
        Bytes::from(serde_json::to_vec(&req).unwrap()),
        "test.reply.unknown",
    )
    .await;
    assert_eq!(resp.error.unwrap().code, -32601);
}

#[tokio::test]
async fn malformed_json_returns_minus_32700() {
    let bus = Arc::new(MemoryBus::new());
    run_server(bus.clone()).await;
    let resp = round_trip(&bus, Bytes::from_static(b"not json {"), "test.reply.bad").await;
    assert_eq!(resp.error.unwrap().code, -32700);
}

mod task_store_tests {
    use klieo_a2a::task_store::A2aTaskStore;
    use klieo_a2a::types::{Task, TaskStatus};
    use klieo_bus_memory::MemoryBus;
    use std::sync::Arc;

    fn make_task(id: &str, ctx: &str) -> Task {
        Task {
            id: id.into(),
            contextId: ctx.into(),
            status: TaskStatus::Submitted,
            artifacts: vec![],
            history: vec![],
            metadata: None,
        }
    }

    #[tokio::test]
    async fn round_trip_a_task() {
        let bus = Arc::new(MemoryBus::new());
        let store = A2aTaskStore::new(bus.kv.clone(), "a2a.tasks".into());
        let t = make_task("t-1", "c-1");
        store.put(&t).await.unwrap();
        let back = store.get("t-1").await.unwrap().unwrap();
        assert_eq!(back.id, "t-1");
    }

    #[tokio::test]
    async fn list_tasks_filters_by_context() {
        let bus = Arc::new(MemoryBus::new());
        let store = A2aTaskStore::new(bus.kv.clone(), "a2a.tasks".into());
        store.put(&make_task("t-1", "c-1")).await.unwrap();
        store.put(&make_task("t-2", "c-1")).await.unwrap();
        store.put(&make_task("t-3", "c-2")).await.unwrap();
        let in_c1 = store.list(Some("c-1")).await.unwrap();
        assert_eq!(in_c1.len(), 2);
    }

    #[tokio::test]
    async fn delete_removes_from_index() {
        let bus = Arc::new(MemoryBus::new());
        let store = A2aTaskStore::new(bus.kv.clone(), "a2a.tasks".into());
        store.put(&make_task("t-1", "c-1")).await.unwrap();
        store.delete("t-1").await.unwrap();
        let in_c1 = store.list(Some("c-1")).await.unwrap();
        assert!(in_c1.is_empty());
        assert!(store.get("t-1").await.unwrap().is_none());
    }
}