Skip to content

fix(goose-acp): heap allocations #7322

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
alexhancock merged 1 commit into from
Feb 19, 2026
Merged

fix(goose-acp): heap allocations #7322

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
185 changes: 99 additions & 86 deletions crates/goose-acp/src/server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1193,112 +1193,125 @@ impl JrMessageHandler for GooseAcpHandler {
"goose-acp"
}

async fn handle_message(
fn handle_message(
&mut self,
message: MessageCx,
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::JrRequestCx;

MatchMessageFrom::new(message, &cx)
.if_request(
|req: InitializeRequest, req_cx: JrRequestCx<InitializeResponse>| async {
req_cx.respond(self.agent.on_initialize(req).await?)
},
)
.await
.if_request(
|_req: AuthenticateRequest, req_cx: JrRequestCx<AuthenticateResponse>| async {
req_cx.respond(AuthenticateResponse::new())
},
)
.await
.if_request(
|req: NewSessionRequest, req_cx: JrRequestCx<NewSessionResponse>| async {
req_cx.respond(self.agent.on_new_session(req).await?)
},
)
.await
.if_request(
|req: LoadSessionRequest, req_cx: JrRequestCx<LoadSessionResponse>| async {
req_cx.respond(self.agent.on_load_session(req, &cx).await?)
},
)
.await
.if_request(
|req: PromptRequest, req_cx: JrRequestCx<PromptResponse>| async {
let agent = self.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)?;
let agent = self.agent.clone();

// The MatchMessageFrom chain produces an ~85KB async state machine.
// Box::pin moves it to the heap so it doesn't overflow the tokio worker stack.
Box::pin(async move {
MatchMessageFrom::new(message, &cx)
.if_request(
|req: InitializeRequest, req_cx: JrRequestCx<InitializeResponse>| async {
req_cx.respond(agent.on_initialize(req).await?)
},
)
.await
.if_request(
|_req: AuthenticateRequest, req_cx: JrRequestCx<AuthenticateResponse>| async {
req_cx.respond(AuthenticateResponse::new())
},
)
.await
.if_request(
|req: NewSessionRequest, req_cx: JrRequestCx<NewSessionResponse>| async {
req_cx.respond(agent.on_new_session(req).await?)
},
)
.await
.if_request(
|req: LoadSessionRequest, req_cx: JrRequestCx<LoadSessionResponse>| async {
req_cx.respond(agent.on_load_session(req, &cx).await?)
},
)
.await
.if_request(
|req: PromptRequest, req_cx: JrRequestCx<PromptResponse>| async {
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)?;
}
}
}
Ok(())
})?;
Ok(())
},
)
.await
.if_notification(|notif: CancelNotification| async {
self.agent.on_cancel(notif).await
})
.await
// Handle methods not yet in the sacp typed API.
// - session/set_model: typed support pending in sacp
// - _<method>: custom requests that will eventually route to goose-server
.otherwise({
let agent = self.agent.clone();
|message: MessageCx| async move {
match message {
MessageCx::Request(req, request_cx)
if req.method == "session/set_model" =>
{
let params: SetSessionModelRequest = serde_json::from_value(req.params)
.map_err(|e| sacp::Error::invalid_params().data(e.to_string()))?;
let resp = agent
.on_set_model(&params.session_id.0, &params.model_id.0)
.await?;
let json = serde_json::to_value(resp)
.map_err(|e| sacp::Error::internal_error().data(e.to_string()))?;
request_cx.respond(json)?;
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(())
},
)
.await
.if_notification(|notif: CancelNotification| async { agent.on_cancel(notif).await })
.await
// Handle methods not yet in the sacp typed API.
// - session/set_model: typed support pending in sacp
// - _<method>: custom requests that will eventually route to goose-server
.otherwise({
let agent = agent.clone();
|message: MessageCx| async move {
match message {
MessageCx::Request(req, request_cx)
if req.method == "session/set_model" =>
{
let params: SetSessionModelRequest =
serde_json::from_value(req.params).map_err(|e| {
sacp::Error::invalid_params().data(e.to_string())
})?;
let resp = agent
.on_set_model(&params.session_id.0, &params.model_id.0)
.await?;
let json = serde_json::to_value(resp).map_err(|e| {
sacp::Error::internal_error().data(e.to_string())
})?;
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(())
}
_ => Err(sacp::Error::method_not_found()),
}
_ => Err(sacp::Error::method_not_found()),
}
}
})
.await
.map(|()| Handled::Yes)
})
.await
.map(|()| Handled::Yes)
})
Copy link
Collaborator

@jamadeo jamadeo Feb 19, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The copilot explanation is convincing. That is quite a lot of .awaits, each one growing the size of the state machine. 85kb is quite a lot less than the default 2mb though.

}
}

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
R: futures::AsyncRead + Unpin + Send + 'static,
W: futures::AsyncWrite + Unpin + Send + 'static,
{
let handler = GooseAcpHandler { agent };
Box::pin(async move {
let handler = GooseAcpHandler { agent };

AgentToClient::builder()
.name("goose-acp")
.with_handler(handler)
.serve(ByteStreams::new(write, read))
.await?;
AgentToClient::builder()
.name("goose-acp")
.with_handler(handler)
.serve(ByteStreams::new(write, read))
.await?;

Ok(())
Ok(())
})
}

pub async fn run(builtins: Vec<String>) -> Result<()> {
Expand Down
10 changes: 4 additions & 6 deletions crates/goose-acp/src/transport/http.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,13 +41,11 @@ impl HttpState {

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 read_stream = ReceiverToAsyncRead::new(to_agent_rx);
let write_stream = SenderToAsyncWrite::new(from_agent_tx);

if let Err(e) =
crate::server::serve(agent, read_stream.compat(), write_stream.compat_write()).await
{
if let Err(e) = fut.await {
error!("ACP session error: {}", e);
}
});
Expand Down
10 changes: 4 additions & 6 deletions crates/goose-acp/src/transport/websocket.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,11 @@ impl WsState {

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 read_stream = ReceiverToAsyncRead::new(to_agent_rx);
let write_stream = SenderToAsyncWrite::new(from_agent_tx);

if let Err(e) =
crate::server::serve(agent, read_stream.compat(), write_stream.compat_write()).await
{
if let Err(e) = fut.await {
error!("ACP WebSocket session error: {}", e);
}
});
Expand Down
Loading