| 1 | //! The stdio wiring: one handler per ACP method, each spawned off the dispatch loop when it |
| 2 | //! awaits anything (a model list, MCP servers, disk, or a turn that makes client requests). |
| 3 | |
| 4 | use std::sync::Arc; |
| 5 | |
| 6 | use agent_client_protocol::schema::v1::{ |
| 7 | AgentAuthCapabilities, AgentCapabilities, AuthenticateRequest, CancelNotification, |
| 8 | CloseSessionRequest, DeleteSessionRequest, Implementation, InitializeRequest, |
| 9 | InitializeResponse, ListSessionsRequest, LoadSessionRequest, LogoutCapabilities, LogoutRequest, |
| 10 | McpCapabilities, NewSessionRequest, PromptRequest, ResumeSessionRequest, |
| 11 | SessionAdditionalDirectoriesCapabilities, SessionCapabilities, SessionCloseCapabilities, |
| 12 | SessionDeleteCapabilities, SessionListCapabilities, SessionResumeCapabilities, |
| 13 | SetSessionConfigOptionRequest, |
| 14 | }; |
| 15 | use agent_client_protocol::{Agent, ConnectTo, Stdio}; |
| 16 | |
| 17 | use crate::agent::{self, AcpServer}; |
| 18 | use crate::{AgentInfo, Profile, cli}; |
| 19 | |
| 20 | /// How the server was started. |
| 21 | #[derive(Debug, Clone, Copy)] |
| 22 | pub struct ServeOptions { |
| 23 | /// Default for each session's auto-approve: `--yolo` or `<PREFIX>_YOLO=1`. |
| 24 | pub yolo: bool, |
| 25 | /// `tui` when launched by the terminal UI (it sets `<PREFIX>_SURFACE=tui`), else `acp`. |
| 26 | pub surface: &'static str, |
| 27 | } |
| 28 | |
| 29 | impl ServeOptions { |
| 30 | /// `<PREFIX>_YOLO` and `<PREFIX>_SURFACE` from the environment. |
| 31 | pub fn from_env(info: &AgentInfo) -> Self { |
| 32 | Self { |
| 33 | yolo: cli::env_flag(&info.env.var("YOLO")), |
| 34 | surface: cli::surface(&info.env), |
| 35 | } |
| 36 | } |
| 37 | } |
| 38 | |
| 39 | /// Serve ACP over stdin/stdout until the client disconnects, then stop every MCP server. |
| 40 | /// Logs must go to stderr: stdout carries JSON-RPC only. |
| 41 | pub async fn serve_stdio( |
| 42 | profile: Arc<dyn Profile>, |
| 43 | opts: ServeOptions, |
| 44 | ) -> agent_client_protocol::Result<()> { |
| 45 | serve(profile, opts, Stdio::new()).await |
| 46 | } |
| 47 | |
| 48 | /// Serve ACP over any transport until the client disconnects, then stop every MCP server. |
| 49 | /// |
| 50 | /// This is how a host embeds the agent in its own process instead of launching it as a |
| 51 | /// subprocess, for example an app in the macOS App Store sandbox: connect it to |
| 52 | /// [`Channel::duplex`](agent_client_protocol::Channel::duplex) and run an ACP client on the |
| 53 | /// other end. |
| 54 | /// |
| 55 | /// One call serves one connection: the transport is consumed, and the future resolves when |
| 56 | /// that client hangs up. To serve another client, call `serve` again with a new transport. |
| 57 | pub async fn serve( |
| 58 | profile: Arc<dyn Profile>, |
| 59 | opts: ServeOptions, |
| 60 | transport: impl ConnectTo<Agent> + 'static, |
| 61 | ) -> agent_client_protocol::Result<()> { |
| 62 | let info = profile.info(); |
| 63 | let prompt_caps = profile.prompt_capabilities(); |
| 64 | let agent = AcpServer::new(profile, opts.yolo, opts.surface); |
| 65 | tracing::info!("{} starting with model {}", info.name, agent.model()); |
| 66 | |
| 67 | let result = Agent |
| 68 | .builder() |
| 69 | .name(info.name) |
| 70 | .on_receive_request( |
| 71 | { |
| 72 | let agent = agent.clone(); |
| 73 | async move |req: InitializeRequest, responder, _cx| { |
| 74 | agent.set_client_caps(req.client_capabilities.clone()); |
| 75 | let methods = agent.auth_methods(&req.client_capabilities); |
| 76 | responder.respond( |
| 77 | InitializeResponse::new(agent::negotiate_version(req.protocol_version)) |
| 78 | .auth_methods(methods) |
| 79 | .agent_capabilities( |
| 80 | AgentCapabilities::new() |
| 81 | .prompt_capabilities(prompt_caps.clone()) |
| 82 | .session_capabilities( |
| 83 | SessionCapabilities::new() |
| 84 | .list(SessionListCapabilities::new()) |
| 85 | .additional_directories( |
| 86 | SessionAdditionalDirectoriesCapabilities::new(), |
| 87 | ) |
| 88 | .resume(SessionResumeCapabilities::new()) |
| 89 | .close(SessionCloseCapabilities::new()) |
| 90 | .delete(SessionDeleteCapabilities::new()), |
| 91 | ) |
| 92 | .auth( |
| 93 | AgentAuthCapabilities::new() |
| 94 | .logout(LogoutCapabilities::new()), |
| 95 | ) |
| 96 | .mcp_capabilities(McpCapabilities::new().http(true)) |
| 97 | .load_session(true), |
| 98 | ) |
| 99 | .agent_info(Implementation::new(info.name, info.version)), |
| 100 | ) |
| 101 | } |
| 102 | }, |
| 103 | agent_client_protocol::on_receive_request!(), |
| 104 | ) |
| 105 | .on_receive_request( |
| 106 | { |
| 107 | let agent = agent.clone(); |
| 108 | async move |req: AuthenticateRequest, responder, cx| { |
| 109 | // Spawned: a provider may wait on the client's answer to an elicitation, |
| 110 | // which the dispatch loop has to be free to deliver. |
| 111 | let agent = agent.clone(); |
| 112 | let connection = cx.clone(); |
| 113 | let request_id = responder.id().clone(); |
| 114 | cx.spawn(async move { |
| 115 | match agent.authenticate(req, &connection, request_id).await { |
| 116 | Ok(resp) => responder.respond(resp), |
| 117 | Err(e) => responder.respond_with_error(e), |
| 118 | } |
| 119 | }) |
| 120 | } |
| 121 | }, |
| 122 | agent_client_protocol::on_receive_request!(), |
| 123 | ) |
| 124 | .on_receive_request( |
| 125 | { |
| 126 | let agent = agent.clone(); |
| 127 | async move |_req: LogoutRequest, responder, cx| { |
| 128 | let agent = agent.clone(); |
| 129 | cx.spawn(async move { |
| 130 | match agent.logout().await { |
| 131 | Ok(resp) => responder.respond(resp), |
| 132 | Err(e) => responder.respond_with_error(e), |
| 133 | } |
| 134 | }) |
| 135 | } |
| 136 | }, |
| 137 | agent_client_protocol::on_receive_request!(), |
| 138 | ) |
| 139 | .on_receive_request( |
| 140 | { |
| 141 | let agent = agent.clone(); |
| 142 | async move |req: NewSessionRequest, responder, cx| { |
| 143 | // Listing models and starting MCP servers are slow; keep them off the |
| 144 | // dispatch loop. |
| 145 | let agent = agent.clone(); |
| 146 | let connection = cx.clone(); |
| 147 | cx.spawn(async move { |
| 148 | // Picks up a key stored by `--setup` since this process started. |
| 149 | if !agent.refresh_credentials().await { |
| 150 | return responder.respond_with_error(agent.auth_required()); |
| 151 | } |
| 152 | match agent.new_session(req).await { |
| 153 | Ok(resp) => { |
| 154 | let id = resp.session_id.clone(); |
| 155 | responder.respond(resp)?; |
| 156 | agent.send_commands(&connection, &id) |
| 157 | } |
| 158 | Err(e) => responder.respond_with_error(e), |
| 159 | } |
| 160 | }) |
| 161 | } |
| 162 | }, |
| 163 | agent_client_protocol::on_receive_request!(), |
| 164 | ) |
| 165 | .on_receive_request( |
| 166 | { |
| 167 | let agent = agent.clone(); |
| 168 | async move |req: ListSessionsRequest, responder, cx| { |
| 169 | // Reads the session directory from disk. |
| 170 | let agent = agent.clone(); |
| 171 | cx.spawn(async move { |
| 172 | match tokio::task::spawn_blocking(move || agent.list_sessions(req)).await { |
| 173 | Ok(Ok(resp)) => responder.respond(resp), |
| 174 | Ok(Err(e)) => responder.respond_with_error(e), |
| 175 | Err(e) => responder.respond_with_internal_error(e.to_string()), |
| 176 | } |
| 177 | }) |
| 178 | } |
| 179 | }, |
| 180 | agent_client_protocol::on_receive_request!(), |
| 181 | ) |
| 182 | .on_receive_request( |
| 183 | { |
| 184 | let agent = agent.clone(); |
| 185 | async move |req: LoadSessionRequest, responder, cx| { |
| 186 | let agent = agent.clone(); |
| 187 | let connection = cx.clone(); |
| 188 | cx.spawn(async move { |
| 189 | // Editors reopen a thread with session/load after terminal auth, so |
| 190 | // AUTH_REQUIRED here shows sign-in instead of a dead thread. |
| 191 | if !agent.refresh_credentials().await { |
| 192 | return responder.respond_with_error(agent.auth_required()); |
| 193 | } |
| 194 | let id = req.session_id.clone(); |
| 195 | match agent.load_session(req, &connection).await { |
| 196 | Ok(resp) => { |
| 197 | responder.respond(resp)?; |
| 198 | agent.send_commands(&connection, &id) |
| 199 | } |
| 200 | Err(e) => responder.respond_with_error(e), |
| 201 | } |
| 202 | }) |
| 203 | } |
| 204 | }, |
| 205 | agent_client_protocol::on_receive_request!(), |
| 206 | ) |
| 207 | .on_receive_request( |
| 208 | { |
| 209 | let agent = agent.clone(); |
| 210 | async move |req: ResumeSessionRequest, responder, cx| { |
| 211 | let agent = agent.clone(); |
| 212 | let connection = cx.clone(); |
| 213 | cx.spawn(async move { |
| 214 | if !agent.refresh_credentials().await { |
| 215 | return responder.respond_with_error(agent.auth_required()); |
| 216 | } |
| 217 | let id = req.session_id.clone(); |
| 218 | match agent.resume_session(req, &connection).await { |
| 219 | Ok(resp) => { |
| 220 | responder.respond(resp)?; |
| 221 | agent.send_commands(&connection, &id) |
| 222 | } |
| 223 | Err(e) => responder.respond_with_error(e), |
| 224 | } |
| 225 | }) |
| 226 | } |
| 227 | }, |
| 228 | agent_client_protocol::on_receive_request!(), |
| 229 | ) |
| 230 | .on_receive_request( |
| 231 | { |
| 232 | let agent = agent.clone(); |
| 233 | async move |req: CloseSessionRequest, responder, _cx| match agent.close_session(req) |
| 234 | { |
| 235 | Ok(resp) => responder.respond(resp), |
| 236 | Err(e) => responder.respond_with_error(e), |
| 237 | } |
| 238 | }, |
| 239 | agent_client_protocol::on_receive_request!(), |
| 240 | ) |
| 241 | .on_receive_request( |
| 242 | { |
| 243 | let agent = agent.clone(); |
| 244 | async move |req: DeleteSessionRequest, responder, _cx| match agent |
| 245 | .delete_session(req) |
| 246 | { |
| 247 | Ok(resp) => responder.respond(resp), |
| 248 | Err(e) => responder.respond_with_error(e), |
| 249 | } |
| 250 | }, |
| 251 | agent_client_protocol::on_receive_request!(), |
| 252 | ) |
| 253 | .on_receive_request( |
| 254 | { |
| 255 | let agent = agent.clone(); |
| 256 | async move |req: SetSessionConfigOptionRequest, responder, cx| { |
| 257 | let agent = agent.clone(); |
| 258 | let connection = cx.clone(); |
| 259 | cx.spawn(async move { |
| 260 | match agent.set_config_option(req, &connection).await { |
| 261 | Ok(resp) => responder.respond(resp), |
| 262 | Err(e) => responder.respond_with_error(e), |
| 263 | } |
| 264 | }) |
| 265 | } |
| 266 | }, |
| 267 | agent_client_protocol::on_receive_request!(), |
| 268 | ) |
| 269 | .on_receive_request( |
| 270 | { |
| 271 | let agent = agent.clone(); |
| 272 | async move |req: PromptRequest, responder, cx| { |
| 273 | // Run the turn off the dispatch loop so it can make requests to the client. |
| 274 | let agent = agent.clone(); |
| 275 | let connection = cx.clone(); |
| 276 | cx.spawn(async move { |
| 277 | if !agent.has_api_key() && !agent.refresh_credentials().await { |
| 278 | return responder.respond_with_error(agent.auth_required()); |
| 279 | } |
| 280 | agent.prompt(req, responder, connection).await |
| 281 | }) |
| 282 | } |
| 283 | }, |
| 284 | agent_client_protocol::on_receive_request!(), |
| 285 | ) |
| 286 | .on_receive_notification( |
| 287 | { |
| 288 | let agent = agent.clone(); |
| 289 | async move |n: CancelNotification, _cx| { |
| 290 | agent.cancel(&n.session_id); |
| 291 | Ok(()) |
| 292 | } |
| 293 | }, |
| 294 | agent_client_protocol::on_receive_notification!(), |
| 295 | ) |
| 296 | .connect_to(transport) |
| 297 | .await; |
| 298 | // The editor went away: stop MCP servers instead of leaving them to be orphaned. |
| 299 | agent.shutdown().await; |
| 300 | result |
| 301 | } |