tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

fix(goose-acp): heap allocations (#7322)

Browse Source
This commit is contained in:
Alex Hancock
2026-02-18 22:20:13 -05:00
committed by GitHub
parent b4b9f08d0b
commit c35b24975f
3 changed files with 107 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-acp/src/server.rs
+99 -86
View File
@@ -1193,112 +1193,125 @@ impl JrMessageHandler for GooseAcpHandler {
"goose-acp" "goose-acp"
} }
async fn handle_message( fn handle_message(
&mut self, &mut self,
message: MessageCx, message: MessageCx,
cx: JrConnectionCx<AgentToClient>, cx: JrConnectionCx<AgentToClient>,
) -> Result<Handled<MessageCx>, sacp::Error> { ) -> impl std::future::Future<Output = Result<Handled<MessageCx>, sacp::Error>> + Send {
use sacp::util::MatchMessageFrom; use sacp::util::MatchMessageFrom;
use sacp::JrRequestCx; use sacp::JrRequestCx;
MatchMessageFrom::new(message, &cx) let agent = self.agent.clone();
.if_request(
|req: InitializeRequest, req_cx: JrRequestCx<InitializeResponse>| async { // The MatchMessageFrom chain produces an ~85KB async state machine.
req_cx.respond(self.agent.on_initialize(req).await?) // Box::pin moves it to the heap so it doesn't overflow the tokio worker stack.
}, Box::pin(async move {
) MatchMessageFrom::new(message, &cx)
.await .if_request(
.if_request( |req: InitializeRequest, req_cx: JrRequestCx<InitializeResponse>| async {
|_req: AuthenticateRequest, req_cx: JrRequestCx<AuthenticateResponse>| async { req_cx.respond(agent.on_initialize(req).await?)
req_cx.respond(AuthenticateResponse::new()) },
}, )
) .await
.await .if_request(
.if_request( |_req: AuthenticateRequest, req_cx: JrRequestCx<AuthenticateResponse>| async {
|req: NewSessionRequest, req_cx: JrRequestCx<NewSessionResponse>| async { req_cx.respond(AuthenticateResponse::new())
req_cx.respond(self.agent.on_new_session(req).await?) },
}, )
) .await
.await .if_request(
.if_request( |req: NewSessionRequest, req_cx: JrRequestCx<NewSessionResponse>| async {
|req: LoadSessionRequest, req_cx: JrRequestCx<LoadSessionResponse>| async { req_cx.respond(agent.on_new_session(req).await?)
req_cx.respond(self.agent.on_load_session(req, &cx).await?) },
}, )
) .await
.await .if_request(
.if_request( |req: LoadSessionRequest, req_cx: JrRequestCx<LoadSessionResponse>| async {
|req: PromptRequest, req_cx: JrRequestCx<PromptResponse>| async { req_cx.respond(agent.on_load_session(req, &cx).await?)
let agent = self.agent.clone(); },
let cx_clone = cx.clone(); )
cx.spawn(async move { .await
match agent.on_prompt(req, &cx_clone).await { .if_request(
Ok(response) => { |req: PromptRequest, req_cx: JrRequestCx<PromptResponse>| async {
req_cx.respond(response)?; let agent = agent.clone();
let cx_clone = cx.clone();
cx.spawn(async move {
match agent.on_prompt(req, &cx_clone).await {
Ok(response) => {
req_cx.respond(response)?;
}
Err(e) => {
req_cx.respond_with_error(e)?;
}
} }
Err(e) => { Ok(())
req_cx.respond_with_error(e)?; })?;
}
}
Ok(()) Ok(())
})?; },
Ok(()) )
}, .await
) .if_notification(|notif: CancelNotification| async { agent.on_cancel(notif).await })
.await .await
.if_notification(|notif: CancelNotification| async { // Handle methods not yet in the sacp typed API.
self.agent.on_cancel(notif).await // - session/set_model: typed support pending in sacp
}) // - _<method>: custom requests that will eventually route to goose-server
.await .otherwise({
// Handle methods not yet in the sacp typed API. let agent = agent.clone();
// - session/set_model: typed support pending in sacp |message: MessageCx| async move {
// - _<method>: custom requests that will eventually route to goose-server match message {
.otherwise({ MessageCx::Request(req, request_cx)
let agent = self.agent.clone(); if req.method == "session/set_model" =>
|message: MessageCx| async move { {
match message { let params: SetSessionModelRequest =
MessageCx::Request(req, request_cx) serde_json::from_value(req.params).map_err(|e| {
if req.method == "session/set_model" => sacp::Error::invalid_params().data(e.to_string())
{ })?;
let params: SetSessionModelRequest = serde_json::from_value(req.params) let resp = agent
.map_err(|e| sacp::Error::invalid_params().data(e.to_string()))?; .on_set_model(&params.session_id.0, &params.model_id.0)
let resp = agent .await?;
.on_set_model(&params.session_id.0, &params.model_id.0) let json = serde_json::to_value(resp).map_err(|e| {
.await?; sacp::Error::internal_error().data(e.to_string())
let json = serde_json::to_value(resp) })?;
.map_err(|e| sacp::Error::internal_error().data(e.to_string()))?; request_cx.respond(json)?;
request_cx.respond(json)?; Ok(())
Ok(())
}
MessageCx::Request(req, request_cx) if req.method.starts_with('_') => {
match agent.handle_custom_request(&req.method, req.params).await {
Ok(json) => request_cx.respond(json)?,
Err(e) => request_cx.respond_with_error(e)?,
} }
Ok(()) MessageCx::Request(req, request_cx) if req.method.starts_with('_') => {
match agent.handle_custom_request(&req.method, req.params).await {
Ok(json) => request_cx.respond(json)?,
Err(e) => request_cx.respond_with_error(e)?,
}
Ok(())
}
_ => Err(sacp::Error::method_not_found()),
} }
_ => Err(sacp::Error::method_not_found()),
} }
} })
}) .await
.await .map(|()| Handled::Yes)
.map(|()| Handled::Yes) })
} }
} }
pub async fn serve<R, W>(agent: Arc<GooseAcpAgent>, read: R, write: W) -> Result<()> pub fn serve<R, W>(
agent: Arc<GooseAcpAgent>,
read: R,
write: W,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send>>
where where
R: futures::AsyncRead + Unpin + Send + 'static, R: futures::AsyncRead + Unpin + Send + 'static,
W: futures::AsyncWrite + Unpin + Send + 'static, W: futures::AsyncWrite + Unpin + Send + 'static,
{ {
let handler = GooseAcpHandler { agent }; Box::pin(async move {
let handler = GooseAcpHandler { agent };
AgentToClient::builder() AgentToClient::builder()
.name("goose-acp") .name("goose-acp")
.with_handler(handler) .with_handler(handler)
.serve(ByteStreams::new(write, read)) .serve(ByteStreams::new(write, read))
.await?; .await?;
Ok(()) Ok(())
})
} }
pub async fn run(builtins: Vec<String>) -> Result<()> { pub async fn run(builtins: Vec<String>) -> Result<()> {
crates/goose-acp/src/transport/http.rs
+4 -6
View File
@@ -41,13 +41,11 @@ impl HttpState {
let acp_session_id = uuid::Uuid::new_v4().to_string(); let acp_session_id = uuid::Uuid::new_v4().to_string();
let read_stream = ReceiverToAsyncRead::new(to_agent_rx);
let write_stream = SenderToAsyncWrite::new(from_agent_tx);
let fut = crate::server::serve(agent, read_stream.compat(), write_stream.compat_write());
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
let read_stream = ReceiverToAsyncRead::new(to_agent_rx); if let Err(e) = fut.await {
let write_stream = SenderToAsyncWrite::new(from_agent_tx);
if let Err(e) =
crate::server::serve(agent, read_stream.compat(), write_stream.compat_write()).await
{
error!("ACP session error: {}", e); error!("ACP session error: {}", e);
} }
}); });
crates/goose-acp/src/transport/websocket.rs
+4 -6
View File
@@ -36,13 +36,11 @@ impl WsState {
let acp_session_id = uuid::Uuid::new_v4().to_string(); let acp_session_id = uuid::Uuid::new_v4().to_string();
let read_stream = ReceiverToAsyncRead::new(to_agent_rx);
let write_stream = SenderToAsyncWrite::new(from_agent_tx);
let fut = crate::server::serve(agent, read_stream.compat(), write_stream.compat_write());
let handle = tokio::spawn(async move { let handle = tokio::spawn(async move {
let read_stream = ReceiverToAsyncRead::new(to_agent_rx); if let Err(e) = fut.await {
let write_stream = SenderToAsyncWrite::new(from_agent_tx);
if let Err(e) =
crate::server::serve(agent, read_stream.compat(), write_stream.compat_write()).await
{
error!("ACP WebSocket session error: {}", e); error!("ACP WebSocket session error: {}", e);
} }
}); });