feat(unstable): Add support for session/close methods (#77)

Author: benbrandt

Cargo.lock

@@ -17,7 +17,7 @@ tokio = { version = "1.48", features = ["full"] }
 tokio-util = { version = "0.7", features = ["compat"] }
 
 # Protocol
-agent-client-protocol-schema = { version = "=0.11.1" }
+agent-client-protocol-schema = { version = "=0.11.2" }
 
 # Serialization
 serde = { version = "1.0", features = ["derive", "rc"] }

Cargo.toml

@@ -7,6 +7,8 @@ use agent_client_protocol_schema::{
     NewSessionResponse, PromptRequest, PromptResponse, Result, SetSessionConfigOptionRequest,
     SetSessionConfigOptionResponse, SetSessionModeRequest, SetSessionModeResponse,
 };
+#[cfg(feature = "unstable_session_close")]
+use agent_client_protocol_schema::{CloseSessionRequest, CloseSessionResponse};
 #[cfg(feature = "unstable_session_fork")]
 use agent_client_protocol_schema::{ForkSessionRequest, ForkSessionResponse};
 #[cfg(feature = "unstable_session_resume")]
@@ -183,6 +185,21 @@ pub trait Agent {
         Err(Error::method_not_found())
     }
 
+    /// **UNSTABLE**
+    ///
+    /// This capability is not part of the spec yet, and may be removed or changed at any point.
+    ///
+    /// Closes an active session, freeing up any resources associated with it.
+    ///
+    /// The agent must cancel any ongoing work (as if `session/cancel` was called)
+    /// and then free up any resources associated with the session.
+    ///
+    /// Only available if the Agent supports the `session.close` capability.
+    #[cfg(feature = "unstable_session_close")]
+    async fn close_session(&self, _args: CloseSessionRequest) -> Result<CloseSessionResponse> {
+        Err(Error::method_not_found())
+    }
+
     /// Handles extension method requests from the client.
     ///
     /// Extension methods provide a way to add custom functionality while maintaining
@@ -255,6 +272,10 @@ impl<T: Agent> Agent for Rc<T> {
     async fn resume_session(&self, args: ResumeSessionRequest) -> Result<ResumeSessionResponse> {
         self.as_ref().resume_session(args).await
     }
+    #[cfg(feature = "unstable_session_close")]
+    async fn close_session(&self, args: CloseSessionRequest) -> Result<CloseSessionResponse> {
+        self.as_ref().close_session(args).await
+    }
     async fn ext_method(&self, args: ExtRequest) -> Result<ExtResponse> {
         self.as_ref().ext_method(args).await
     }
@@ -313,6 +334,10 @@ impl<T: Agent> Agent for Arc<T> {
     async fn resume_session(&self, args: ResumeSessionRequest) -> Result<ResumeSessionResponse> {
         self.as_ref().resume_session(args).await
     }
+    #[cfg(feature = "unstable_session_close")]
+    async fn close_session(&self, args: CloseSessionRequest) -> Result<CloseSessionResponse> {
+        self.as_ref().close_session(args).await
+    }
     async fn ext_method(&self, args: ExtRequest) -> Result<ExtResponse> {
         self.as_ref().ext_method(args).await
     }

src/agent-client-protocol/src/agent.rs

@@ -184,6 +184,17 @@ impl Agent for ClientSideConnection {
             .await
     }
 
+    #[cfg(feature = "unstable_session_close")]
+    async fn close_session(&self, args: CloseSessionRequest) -> Result<CloseSessionResponse> {
+        self.conn
+            .request::<Option<_>>(
+                AGENT_METHOD_NAMES.session_close,
+                Some(ClientRequest::CloseSessionRequest(args)),
+            )
+            .await
+            .map(Option::unwrap_or_default)
+    }
+
     async fn set_session_config_option(
         &self,
         args: SetSessionConfigOptionRequest,
@@ -567,6 +578,10 @@ impl Side for AgentSide {
             m if m == AGENT_METHOD_NAMES.session_resume => serde_json::from_str(params.get())
                 .map(ClientRequest::ResumeSessionRequest)
                 .map_err(Into::into),
+            #[cfg(feature = "unstable_session_close")]
+            m if m == AGENT_METHOD_NAMES.session_close => serde_json::from_str(params.get())
+                .map(ClientRequest::CloseSessionRequest)
+                .map_err(Into::into),
             m if m == AGENT_METHOD_NAMES.session_set_config_option => {
                 serde_json::from_str(params.get())
                     .map(ClientRequest::SetSessionConfigOptionRequest)
@@ -655,6 +670,11 @@ impl<T: Agent> MessageHandler<AgentSide> for T {
                 let response = self.resume_session(args).await?;
                 Ok(AgentResponse::ResumeSessionResponse(response))
             }
+            #[cfg(feature = "unstable_session_close")]
+            ClientRequest::CloseSessionRequest(args) => {
+                let response = self.close_session(args).await?;
+                Ok(AgentResponse::CloseSessionResponse(response))
+            }
             ClientRequest::SetSessionConfigOptionRequest(args) => {
                 let response = self.set_session_config_option(args).await?;
                 Ok(AgentResponse::SetSessionConfigOptionResponse(response))