feat(acp-nats-agent): add server-side ACP agent framework

Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>

Author: yordis

rsworkspace/crates/acp-nats-agent/Cargo.toml

@@ -0,0 +1,36 @@
+[package]
+name = "acp-nats-agent"
+version = "0.0.1"
+edition = "2024"
+
+[lints]
+workspace = true
+
+[dependencies]
+acp-nats = { workspace = true }
+agent-client-protocol = { workspace = true, features = [
+    "unstable_auth_methods",
+    "unstable_boolean_config",
+    "unstable_cancel_request",
+    "unstable_message_id",
+    "unstable_session_close",
+    "unstable_session_fork",
+    "unstable_session_model",
+    "unstable_session_resume",
+    "unstable_session_usage",
+] }
+async-nats = { workspace = true }
+async-trait = { workspace = true }
+bytes = { workspace = true }
+futures = { workspace = true }
+serde = { workspace = true }
+serde_json = { workspace = true }
+tokio = { workspace = true, features = ["rt", "macros", "sync"] }
+tracing = { workspace = true }
+trogon-nats = { workspace = true }
+trogon-std = { workspace = true }
+
+[dev-dependencies]
+tokio = { workspace = true, features = ["test-util"] }
+trogon-nats = { workspace = true, features = ["test-support"] }
+trogon-std = { workspace = true, features = ["test-support"] }

rsworkspace/crates/acp-nats-agent/src/connection.rs

@@ -0,0 +1,282 @@
+use crate::dispatch::{AgentMethod, parse_agent_subject};
+use acp_nats::acp_prefix::AcpPrefix;
+use acp_nats::client_proxy::{NatsClientProxy, NatsClientProxyConfig};
+use acp_nats::session_id::AcpSessionId;
+use agent_client_protocol::{
+    Agent, AuthenticateRequest, CancelNotification, CloseSessionRequest, ExtRequest,
+    ForkSessionRequest, InitializeRequest, ListSessionsRequest, LoadSessionRequest,
+    NewSessionRequest, PromptRequest, ResumeSessionRequest, SetSessionConfigOptionRequest,
+    SetSessionModeRequest, SetSessionModelRequest,
+};
+use async_nats::Message;
+use bytes::Bytes;
+use futures::future::LocalBoxFuture;
+use futures::StreamExt;
+use std::rc::Rc;
+use std::time::Duration;
+use tracing::{error, info, warn};
+use trogon_nats::{FlushClient, PublishClient, RequestClient, SubscribeClient};
+
+pub struct AgentSideNatsConnectionConfig {
+    pub acp_prefix: AcpPrefix,
+    pub operation_timeout: Duration,
+}
+
+impl AgentSideNatsConnectionConfig {
+    pub fn new(acp_prefix: AcpPrefix) -> Self {
+        Self {
+            acp_prefix,
+            operation_timeout: Duration::from_secs(30),
+        }
+    }
+
+    pub fn with_operation_timeout(mut self, timeout: Duration) -> Self {
+        self.operation_timeout = timeout;
+        self
+    }
+}
+
+pub struct AgentSideNatsConnection<N> {
+    nats: N,
+    config: AgentSideNatsConnectionConfig,
+}
+
+impl<N> AgentSideNatsConnection<N>
+where
+    N: SubscribeClient + RequestClient + PublishClient + FlushClient + Clone + 'static,
+{
+    /// Creates a new agent-side NATS connection.
+    ///
+    /// Mirrors the ACP SDK's `AgentSideConnection::new` pattern:
+    /// - Takes an `Agent` impl to handle incoming requests
+    /// - Takes a NATS client instead of byte streams
+    /// - Takes a `spawn` function for concurrency control
+    /// - Returns `(Self, io_task)` where `Self` implements `Client`-creation per session
+    ///
+    /// The returned future subscribes to NATS agent subjects, dispatches incoming
+    /// messages to the agent via the spawn function, and runs until the subscriptions end.
+    pub fn new(
+        agent: impl Agent + 'static,
+        nats: N,
+        config: AgentSideNatsConnectionConfig,
+        spawn: impl Fn(LocalBoxFuture<'static, ()>) + 'static,
+    ) -> (Self, impl std::future::Future<Output = ()>) {
+        let nats_for_serve = nats.clone();
+        let prefix = config.acp_prefix.as_str().to_string();
+
+        let io_task = async move {
+            serve(agent, nats_for_serve, &prefix, spawn).await;
+        };
+
+        let conn = Self { nats, config };
+        (conn, io_task)
+    }
+
+    pub fn client_for_session(&self, session_id: AcpSessionId) -> NatsClientProxy<N> {
+        NatsClientProxy::new(
+            self.nats.clone(),
+            NatsClientProxyConfig {
+                session_id,
+                prefix: self.config.acp_prefix.clone(),
+                timeout: self.config.operation_timeout,
+            },
+        )
+    }
+}
+
+async fn serve<N, A>(
+    agent: A,
+    nats: N,
+    prefix: &str,
+    spawn: impl Fn(LocalBoxFuture<'static, ()>) + 'static,
+) where
+    N: SubscribeClient + PublishClient + Clone + 'static,
+    A: Agent + 'static,
+{
+    let global_wildcard = acp_nats::nats::agent::wildcards::all(prefix);
+    let session_wildcard = acp_nats::nats::agent::wildcards::all_sessions(prefix);
+
+    info!(
+        global = %global_wildcard,
+        session = %session_wildcard,
+        "Starting agent-side NATS connection"
+    );
+
+    let global_sub = match nats.subscribe(global_wildcard).await {
+        Ok(sub) => sub,
+        Err(e) => {
+            error!(error = %e, "Failed to subscribe to global agent subjects");
+            return;
+        }
+    };
+
+    let session_sub = match nats.subscribe(session_wildcard).await {
+        Ok(sub) => sub,
+        Err(e) => {
+            error!(error = %e, "Failed to subscribe to session agent subjects");
+            return;
+        }
+    };
+
+    let mut merged = futures::stream::select(global_sub, session_sub);
+    let agent = Rc::new(agent);
+    let nats = Rc::new(nats);
+
+    while let Some(msg) = merged.next().await {
+        let agent = agent.clone();
+        let nats = nats.clone();
+        spawn(Box::pin(async move {
+            dispatch_message(msg, agent.as_ref(), nats.as_ref()).await;
+        }));
+    }
+
+    info!("Agent-side NATS connection ended");
+}
+
+async fn dispatch_message<N: PublishClient, A: Agent>(msg: Message, agent: &A, nats: &N) {
+    let subject = msg.subject.as_str();
+
+    let parsed = match parse_agent_subject(subject) {
+        Some(p) => p,
+        None => {
+            warn!(subject, "Failed to parse agent subject");
+            return;
+        }
+    };
+
+    let result = match parsed.method {
+        AgentMethod::Initialize => {
+            handle_request(&msg, nats, |req: InitializeRequest| agent.initialize(req)).await
+        }
+        AgentMethod::Authenticate => {
+            handle_request(&msg, nats, |req: AuthenticateRequest| {
+                agent.authenticate(req)
+            })
+            .await
+        }
+        AgentMethod::SessionNew => {
+            handle_request(&msg, nats, |req: NewSessionRequest| agent.new_session(req)).await
+        }
+        AgentMethod::SessionList => {
+            handle_request(&msg, nats, |req: ListSessionsRequest| {
+                agent.list_sessions(req)
+            })
+            .await
+        }
+        AgentMethod::SessionLoad => {
+            handle_request(&msg, nats, |req: LoadSessionRequest| {
+                agent.load_session(req)
+            })
+            .await
+        }
+        AgentMethod::SessionPrompt => {
+            handle_request(&msg, nats, |req: PromptRequest| agent.prompt(req)).await
+        }
+        AgentMethod::SessionCancel => {
+            handle_notification(&msg, |req: CancelNotification| agent.cancel(req)).await
+        }
+        AgentMethod::SessionSetMode => {
+            handle_request(&msg, nats, |req: SetSessionModeRequest| {
+                agent.set_session_mode(req)
+            })
+            .await
+        }
+        AgentMethod::SessionSetConfigOption => {
+            handle_request(&msg, nats, |req: SetSessionConfigOptionRequest| {
+                agent.set_session_config_option(req)
+            })
+            .await
+        }
+        AgentMethod::SessionSetModel => {
+            handle_request(&msg, nats, |req: SetSessionModelRequest| {
+                agent.set_session_model(req)
+            })
+            .await
+        }
+        AgentMethod::SessionFork => {
+            handle_request(&msg, nats, |req: ForkSessionRequest| {
+                agent.fork_session(req)
+            })
+            .await
+        }
+        AgentMethod::SessionResume => {
+            handle_request(&msg, nats, |req: ResumeSessionRequest| {
+                agent.resume_session(req)
+            })
+            .await
+        }
+        AgentMethod::SessionClose => {
+            handle_request(&msg, nats, |req: CloseSessionRequest| {
+                agent.close_session(req)
+            })
+            .await
+        }
+        AgentMethod::Ext(_) => {
+            handle_request(&msg, nats, |req: ExtRequest| agent.ext_method(req)).await
+        }
+    };
+
+    if let Err(e) = result {
+        warn!(
+            subject,
+            session_id = parsed.session_id.as_ref().map(|s| s.as_str()).unwrap_or("-"),
+            error = %e,
+            "Error handling agent request"
+        );
+    }
+}
+
+async fn handle_request<N, Resp, ReqT, F>(
+    msg: &Message,
+    nats: &N,
+    handler: impl FnOnce(ReqT) -> F,
+) -> std::result::Result<(), String>
+where
+    N: PublishClient,
+    ReqT: serde::de::DeserializeOwned,
+    F: std::future::Future<Output = agent_client_protocol::Result<Resp>>,
+    Resp: serde::Serialize,
+{
+    let reply_to = msg
+        .reply
+        .as_ref()
+        .ok_or_else(|| "no reply subject".to_string())?;
+
+    let request: ReqT =
+        serde_json::from_slice(&msg.payload).map_err(|e| format!("deserialize request: {}", e))?;
+
+    let response = handler(request).await;
+
+    let response_bytes: Bytes = match response {
+        Ok(resp) => serde_json::to_vec(&resp)
+            .map_err(|e| format!("serialize response: {}", e))?
+            .into(),
+        Err(err) => serde_json::to_vec(&err)
+            .map_err(|e| format!("serialize error: {}", e))?
+            .into(),
+    };
+
+    nats.publish_with_headers(
+        reply_to.to_string(),
+        async_nats::HeaderMap::new(),
+        response_bytes,
+    )
+    .await
+    .map_err(|e| format!("respond: {}", e))
+}
+
+async fn handle_notification<ReqT, F>(
+    msg: &Message,
+    handler: impl FnOnce(ReqT) -> F,
+) -> std::result::Result<(), String>
+where
+    ReqT: serde::de::DeserializeOwned,
+    F: std::future::Future<Output = agent_client_protocol::Result<()>>,
+{
+    let request: ReqT = serde_json::from_slice(&msg.payload)
+        .map_err(|e| format!("deserialize notification: {}", e))?;
+
+    handler(request)
+        .await
+        .map_err(|e| format!("notification handler: {}", e))
+}