From 8292b01a47789dba7a1d8dfd792a0a72a1a647cc Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Fri, 29 May 2026 14:06:35 -0400 Subject: [PATCH] Pass connection handle into ACP request handlers --- crates/goose-acp-macros/src/lib.rs | 28 +++++++-- .../goose/src/acp/server/custom_dispatch.rs | 60 ++++++++++++++++++- crates/goose/src/acp/server/dispatch.rs | 49 ++++++++------- 3 files changed, 107 insertions(+), 30 deletions(-) diff --git a/crates/goose-acp-macros/src/lib.rs b/crates/goose-acp-macros/src/lib.rs index beaa3e62bb..3b47737deb 100644 --- a/crates/goose-acp-macros/src/lib.rs +++ b/crates/goose-acp-macros/src/lib.rs @@ -25,16 +25,23 @@ use syn::{ /// /// # Handler signatures /// -/// Handlers may take zero or one parameter (beyond `&self`): +/// Handlers receive the connection, followed by zero or one request parameter: /// /// ```ignore /// // No params — called for requests with no/empty params /// #[custom_method(GetExtensionsRequest)] -/// async fn on_get_extensions(&self) -> Result { .. } +/// async fn on_get_extensions( +/// &self, +/// cx: &ConnectionTo, +/// ) -> Result { .. } /// /// // Typed params — JSON params auto-deserialized /// #[custom_method(GetSessionRequest)] -/// async fn on_get_session(&self, req: GetSessionRequest) -> Result { .. } +/// async fn on_get_session( +/// &self, +/// cx: &ConnectionTo, +/// req: GetSessionRequest, +/// ) -> Result { .. } /// ``` /// /// The return type must be `Result` where `T: Serialize`. @@ -92,7 +99,7 @@ pub fn custom_methods(_attr: TokenStream, item: TokenStream) -> TokenStream { if <#req_type as agent_client_protocol::JsonRpcMessage>::matches_method(method) { let req = serde_json::from_value(params) .map_err(|e| agent_client_protocol::Error::invalid_params().data(e.to_string()))?; - let result = self.#fn_ident(req).await?; + let result = self.#fn_ident(cx, req).await?; return serde_json::to_value(&result) .map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string())); } @@ -101,7 +108,7 @@ pub fn custom_methods(_attr: TokenStream, item: TokenStream) -> TokenStream { None => { quote! { if <#req_type as agent_client_protocol::JsonRpcMessage>::matches_method(method) { - let result = self.#fn_ident().await?; + let result = self.#fn_ident(cx).await?; return serde_json::to_value(&result) .map_err(|e| agent_client_protocol::Error::internal_error().data(e.to_string())); } @@ -184,6 +191,7 @@ pub fn custom_methods(_attr: TokenStream, item: TokenStream) -> TokenStream { let dispatcher = quote! { async fn handle_custom_request( &self, + cx: &agent_client_protocol::ConnectionTo, method: &str, params: serde_json::Value, ) -> Result { @@ -221,8 +229,10 @@ struct Route { ok_type: Option, } -/// Extract the type of the first non-self parameter, if any. +/// Extract the request parameter after `&self` and connection, if any. fn extract_param_type(sig: &syn::Signature) -> Option { + let mut saw_connection = false; + for input in &sig.inputs { if let FnArg::Typed(pat_type) = input { if let Pat::Ident(pat_ident) = &*pat_type.pat { @@ -230,6 +240,12 @@ fn extract_param_type(sig: &syn::Signature) -> Option { continue; } } + + if !saw_connection { + saw_connection = true; + continue; + } + return Some((*pat_type.ty).clone()); } } diff --git a/crates/goose/src/acp/server/custom_dispatch.rs b/crates/goose/src/acp/server/custom_dispatch.rs index 81c6e67107..c2bde54a5f 100644 --- a/crates/goose/src/acp/server/custom_dispatch.rs +++ b/crates/goose/src/acp/server/custom_dispatch.rs @@ -5,15 +5,17 @@ use goose_acp_macros::custom_methods; impl GooseAcpAgent { pub async fn dispatch_custom_request( &self, + cx: &ConnectionTo, method: &str, params: serde_json::Value, ) -> Result { - self.handle_custom_request(method, params).await + self.handle_custom_request(cx, method, params).await } #[custom_method(AddExtensionRequest)] async fn dispatch_add_extension( &self, + _cx: &ConnectionTo, req: AddExtensionRequest, ) -> Result { self.on_add_extension(req).await @@ -22,6 +24,7 @@ impl GooseAcpAgent { #[custom_method(RemoveExtensionRequest)] async fn dispatch_remove_extension( &self, + _cx: &ConnectionTo, req: RemoveExtensionRequest, ) -> Result { self.on_remove_extension(req).await @@ -30,6 +33,7 @@ impl GooseAcpAgent { #[custom_method(GetToolsRequest)] async fn dispatch_get_tools( &self, + _cx: &ConnectionTo, req: GetToolsRequest, ) -> Result { self.on_get_tools(req).await @@ -38,6 +42,7 @@ impl GooseAcpAgent { #[custom_method(GooseToolCallRequest)] async fn dispatch_call_tool( &self, + _cx: &ConnectionTo, req: GooseToolCallRequest, ) -> Result { self.on_call_tool(req).await @@ -46,6 +51,7 @@ impl GooseAcpAgent { #[custom_method(ReadResourceRequest)] async fn dispatch_read_resource( &self, + _cx: &ConnectionTo, req: ReadResourceRequest, ) -> Result { self.on_read_resource(req).await @@ -54,6 +60,7 @@ impl GooseAcpAgent { #[custom_method(UpdateWorkingDirRequest)] async fn dispatch_update_working_dir( &self, + _cx: &ConnectionTo, req: UpdateWorkingDirRequest, ) -> Result { self.on_update_working_dir(req).await @@ -62,6 +69,7 @@ impl GooseAcpAgent { #[custom_method(SetSessionSystemPromptRequest)] async fn dispatch_set_session_system_prompt( &self, + _cx: &ConnectionTo, req: SetSessionSystemPromptRequest, ) -> Result { self.on_set_session_system_prompt(req).await @@ -70,6 +78,7 @@ impl GooseAcpAgent { #[custom_method(DeleteSessionRequest)] async fn dispatch_delete_session( &self, + _cx: &ConnectionTo, req: DeleteSessionRequest, ) -> Result { self.on_delete_session(req).await @@ -78,6 +87,7 @@ impl GooseAcpAgent { #[custom_method(GetExtensionsRequest)] async fn dispatch_get_extensions( &self, + _cx: &ConnectionTo, ) -> Result { self.on_get_extensions().await } @@ -85,6 +95,7 @@ impl GooseAcpAgent { #[custom_method(AddConfigExtensionRequest)] async fn dispatch_add_config_extension( &self, + _cx: &ConnectionTo, req: AddConfigExtensionRequest, ) -> Result { self.on_add_config_extension(req).await @@ -93,6 +104,7 @@ impl GooseAcpAgent { #[custom_method(RemoveConfigExtensionRequest)] async fn dispatch_remove_config_extension( &self, + _cx: &ConnectionTo, req: RemoveConfigExtensionRequest, ) -> Result { self.on_remove_config_extension(req).await @@ -101,6 +113,7 @@ impl GooseAcpAgent { #[custom_method(ToggleConfigExtensionRequest)] async fn dispatch_toggle_config_extension( &self, + _cx: &ConnectionTo, req: ToggleConfigExtensionRequest, ) -> Result { self.on_toggle_config_extension(req).await @@ -109,6 +122,7 @@ impl GooseAcpAgent { #[custom_method(GetSessionExtensionsRequest)] async fn dispatch_get_session_extensions( &self, + _cx: &ConnectionTo, req: GetSessionExtensionsRequest, ) -> Result { self.on_get_session_extensions(req).await @@ -117,6 +131,7 @@ impl GooseAcpAgent { #[custom_method(ListProvidersRequest)] async fn dispatch_list_providers( &self, + _cx: &ConnectionTo, req: ListProvidersRequest, ) -> Result { self.on_list_providers(req).await @@ -125,6 +140,7 @@ impl GooseAcpAgent { #[custom_method(ProviderSupportedModelsListRequest)] async fn dispatch_list_provider_supported_models( &self, + _cx: &ConnectionTo, req: ProviderSupportedModelsListRequest, ) -> Result { self.on_list_provider_supported_models(req).await @@ -133,6 +149,7 @@ impl GooseAcpAgent { #[custom_method(ProviderCatalogListRequest)] async fn dispatch_list_provider_catalog( &self, + _cx: &ConnectionTo, req: ProviderCatalogListRequest, ) -> Result { self.on_list_provider_catalog(req).await @@ -141,6 +158,7 @@ impl GooseAcpAgent { #[custom_method(ProviderSetupCatalogListRequest)] async fn dispatch_list_provider_setup_catalog( &self, + _cx: &ConnectionTo, req: ProviderSetupCatalogListRequest, ) -> Result { self.on_list_provider_setup_catalog(req).await @@ -149,6 +167,7 @@ impl GooseAcpAgent { #[custom_method(ProviderCatalogTemplateRequest)] async fn dispatch_get_provider_catalog_template( &self, + _cx: &ConnectionTo, req: ProviderCatalogTemplateRequest, ) -> Result { self.on_get_provider_catalog_template(req).await @@ -157,6 +176,7 @@ impl GooseAcpAgent { #[custom_method(CustomProviderCreateRequest)] async fn dispatch_create_custom_provider( &self, + _cx: &ConnectionTo, req: CustomProviderCreateRequest, ) -> Result { self.on_create_custom_provider(req).await @@ -165,6 +185,7 @@ impl GooseAcpAgent { #[custom_method(CustomProviderReadRequest)] async fn dispatch_read_custom_provider( &self, + _cx: &ConnectionTo, req: CustomProviderReadRequest, ) -> Result { self.on_read_custom_provider(req).await @@ -173,6 +194,7 @@ impl GooseAcpAgent { #[custom_method(CustomProviderUpdateRequest)] async fn dispatch_update_custom_provider( &self, + _cx: &ConnectionTo, req: CustomProviderUpdateRequest, ) -> Result { self.on_update_custom_provider(req).await @@ -181,6 +203,7 @@ impl GooseAcpAgent { #[custom_method(CustomProviderDeleteRequest)] async fn dispatch_delete_custom_provider( &self, + _cx: &ConnectionTo, req: CustomProviderDeleteRequest, ) -> Result { self.on_delete_custom_provider(req).await @@ -189,6 +212,7 @@ impl GooseAcpAgent { #[custom_method(RefreshProviderInventoryRequest)] async fn dispatch_refresh_provider_inventory( &self, + _cx: &ConnectionTo, req: RefreshProviderInventoryRequest, ) -> Result { self.on_refresh_provider_inventory(req).await @@ -197,6 +221,7 @@ impl GooseAcpAgent { #[custom_method(ProviderConfigReadRequest)] async fn dispatch_read_provider_config( &self, + _cx: &ConnectionTo, req: ProviderConfigReadRequest, ) -> Result { self.on_read_provider_config(req).await @@ -205,6 +230,7 @@ impl GooseAcpAgent { #[custom_method(ProviderConfigStatusRequest)] async fn dispatch_provider_config_status( &self, + _cx: &ConnectionTo, req: ProviderConfigStatusRequest, ) -> Result { self.on_provider_config_status(req).await @@ -213,6 +239,7 @@ impl GooseAcpAgent { #[custom_method(ProviderConfigSaveRequest)] async fn dispatch_save_provider_config( &self, + _cx: &ConnectionTo, req: ProviderConfigSaveRequest, ) -> Result { self.on_save_provider_config(req).await @@ -221,6 +248,7 @@ impl GooseAcpAgent { #[custom_method(ProviderConfigDeleteRequest)] async fn dispatch_delete_provider_config( &self, + _cx: &ConnectionTo, req: ProviderConfigDeleteRequest, ) -> Result { self.on_delete_provider_config(req).await @@ -229,6 +257,7 @@ impl GooseAcpAgent { #[custom_method(ProviderConfigAuthenticateRequest)] async fn dispatch_authenticate_provider_config( &self, + _cx: &ConnectionTo, req: ProviderConfigAuthenticateRequest, ) -> Result { self.on_authenticate_provider_config(req).await @@ -237,6 +266,7 @@ impl GooseAcpAgent { #[custom_method(PreferencesReadRequest)] async fn dispatch_preferences_read( &self, + _cx: &ConnectionTo, req: PreferencesReadRequest, ) -> Result { self.on_preferences_read(req).await @@ -245,6 +275,7 @@ impl GooseAcpAgent { #[custom_method(PreferencesSaveRequest)] async fn dispatch_preferences_save( &self, + _cx: &ConnectionTo, req: PreferencesSaveRequest, ) -> Result { self.on_preferences_save(req).await @@ -253,6 +284,7 @@ impl GooseAcpAgent { #[custom_method(PreferencesRemoveRequest)] async fn dispatch_preferences_remove( &self, + _cx: &ConnectionTo, req: PreferencesRemoveRequest, ) -> Result { self.on_preferences_remove(req).await @@ -261,6 +293,7 @@ impl GooseAcpAgent { #[custom_method(DefaultsReadRequest)] async fn dispatch_defaults_read( &self, + _cx: &ConnectionTo, req: DefaultsReadRequest, ) -> Result { self.on_defaults_read(req).await @@ -269,6 +302,7 @@ impl GooseAcpAgent { #[custom_method(DefaultsSaveRequest)] async fn dispatch_defaults_save( &self, + _cx: &ConnectionTo, req: DefaultsSaveRequest, ) -> Result { self.on_defaults_save(req).await @@ -277,6 +311,7 @@ impl GooseAcpAgent { #[custom_method(OnboardingImportScanRequest)] async fn dispatch_onboarding_import_scan( &self, + _cx: &ConnectionTo, req: OnboardingImportScanRequest, ) -> Result { self.on_onboarding_import_scan(req).await @@ -285,6 +320,7 @@ impl GooseAcpAgent { #[custom_method(OnboardingImportApplyRequest)] async fn dispatch_onboarding_import_apply( &self, + _cx: &ConnectionTo, req: OnboardingImportApplyRequest, ) -> Result { self.on_onboarding_import_apply(req).await @@ -293,6 +329,7 @@ impl GooseAcpAgent { #[custom_method(ExportSessionRequest)] async fn dispatch_export_session( &self, + _cx: &ConnectionTo, req: ExportSessionRequest, ) -> Result { self.on_export_session(req).await @@ -301,6 +338,7 @@ impl GooseAcpAgent { #[custom_method(ImportSessionRequest)] async fn dispatch_import_session( &self, + _cx: &ConnectionTo, req: ImportSessionRequest, ) -> Result { self.on_import_session(req).await @@ -309,6 +347,7 @@ impl GooseAcpAgent { #[custom_method(UpdateSessionProjectRequest)] async fn dispatch_update_session_project( &self, + _cx: &ConnectionTo, req: UpdateSessionProjectRequest, ) -> Result { self.on_update_session_project(req).await @@ -317,6 +356,7 @@ impl GooseAcpAgent { #[custom_method(RenameSessionRequest)] async fn dispatch_rename_session( &self, + _cx: &ConnectionTo, req: RenameSessionRequest, ) -> Result { self.on_rename_session(req).await @@ -325,6 +365,7 @@ impl GooseAcpAgent { #[custom_method(ArchiveSessionRequest)] async fn dispatch_archive_session( &self, + _cx: &ConnectionTo, req: ArchiveSessionRequest, ) -> Result { self.on_archive_session(req).await @@ -333,6 +374,7 @@ impl GooseAcpAgent { #[custom_method(UnarchiveSessionRequest)] async fn dispatch_unarchive_session( &self, + _cx: &ConnectionTo, req: UnarchiveSessionRequest, ) -> Result { self.on_unarchive_session(req).await @@ -341,6 +383,7 @@ impl GooseAcpAgent { #[custom_method(CreateSourceRequest)] async fn dispatch_create_source( &self, + _cx: &ConnectionTo, req: CreateSourceRequest, ) -> Result { self.on_create_source(req).await @@ -349,6 +392,7 @@ impl GooseAcpAgent { #[custom_method(ListSourcesRequest)] async fn dispatch_list_sources( &self, + _cx: &ConnectionTo, req: ListSourcesRequest, ) -> Result { self.on_list_sources(req).await @@ -357,6 +401,7 @@ impl GooseAcpAgent { #[custom_method(UpdateSourceRequest)] async fn dispatch_update_source( &self, + _cx: &ConnectionTo, req: UpdateSourceRequest, ) -> Result { self.on_update_source(req).await @@ -365,6 +410,7 @@ impl GooseAcpAgent { #[custom_method(DeleteSourceRequest)] async fn dispatch_delete_source( &self, + _cx: &ConnectionTo, req: DeleteSourceRequest, ) -> Result { self.on_delete_source(req).await @@ -373,6 +419,7 @@ impl GooseAcpAgent { #[custom_method(ExportSourceRequest)] async fn dispatch_export_source( &self, + _cx: &ConnectionTo, req: ExportSourceRequest, ) -> Result { self.on_export_source(req).await @@ -381,6 +428,7 @@ impl GooseAcpAgent { #[custom_method(ImportSourcesRequest)] async fn dispatch_import_sources( &self, + _cx: &ConnectionTo, req: ImportSourcesRequest, ) -> Result { self.on_import_sources(req).await @@ -389,6 +437,7 @@ impl GooseAcpAgent { #[custom_method(DictationTranscribeRequest)] async fn dispatch_dictation_transcribe( &self, + _cx: &ConnectionTo, req: DictationTranscribeRequest, ) -> Result { self.on_dictation_transcribe(req).await @@ -397,6 +446,7 @@ impl GooseAcpAgent { #[custom_method(DictationConfigRequest)] async fn dispatch_dictation_config( &self, + _cx: &ConnectionTo, _req: DictationConfigRequest, ) -> Result { self.on_dictation_config(_req).await @@ -405,6 +455,7 @@ impl GooseAcpAgent { #[custom_method(DictationSecretSaveRequest)] async fn dispatch_dictation_secret_save( &self, + _cx: &ConnectionTo, req: DictationSecretSaveRequest, ) -> Result { self.on_dictation_secret_save(req).await @@ -413,6 +464,7 @@ impl GooseAcpAgent { #[custom_method(DictationSecretDeleteRequest)] async fn dispatch_dictation_secret_delete( &self, + _cx: &ConnectionTo, req: DictationSecretDeleteRequest, ) -> Result { self.on_dictation_secret_delete(req).await @@ -421,6 +473,7 @@ impl GooseAcpAgent { #[custom_method(DictationModelsListRequest)] async fn dispatch_dictation_models_list( &self, + _cx: &ConnectionTo, _req: DictationModelsListRequest, ) -> Result { self.on_dictation_models_list(_req).await @@ -429,6 +482,7 @@ impl GooseAcpAgent { #[custom_method(DictationModelDownloadRequest)] async fn dispatch_dictation_model_download( &self, + _cx: &ConnectionTo, _req: DictationModelDownloadRequest, ) -> Result { self.on_dictation_model_download(_req).await @@ -437,6 +491,7 @@ impl GooseAcpAgent { #[custom_method(DictationModelDownloadProgressRequest)] async fn dispatch_dictation_model_download_progress( &self, + _cx: &ConnectionTo, _req: DictationModelDownloadProgressRequest, ) -> Result { self.on_dictation_model_download_progress(_req).await @@ -445,6 +500,7 @@ impl GooseAcpAgent { #[custom_method(DictationModelCancelRequest)] async fn dispatch_dictation_model_cancel( &self, + _cx: &ConnectionTo, _req: DictationModelCancelRequest, ) -> Result { self.on_dictation_model_cancel(_req).await @@ -453,6 +509,7 @@ impl GooseAcpAgent { #[custom_method(DictationModelDeleteRequest)] async fn dispatch_dictation_model_delete( &self, + _cx: &ConnectionTo, _req: DictationModelDeleteRequest, ) -> Result { self.on_dictation_model_delete(_req).await @@ -461,6 +518,7 @@ impl GooseAcpAgent { #[custom_method(DictationModelSelectRequest)] async fn dispatch_dictation_model_select( &self, + _cx: &ConnectionTo, req: DictationModelSelectRequest, ) -> Result { self.on_dictation_model_select(req).await diff --git a/crates/goose/src/acp/server/dispatch.rs b/crates/goose/src/acp/server/dispatch.rs index 7e23333e25..ea7e255a52 100644 --- a/crates/goose/src/acp/server/dispatch.rs +++ b/crates/goose/src/acp/server/dispatch.rs @@ -367,30 +367,33 @@ impl HandleDispatchFrom for GooseAcpHandler { } }) .await - .otherwise({ - let agent = agent.clone(); - let cx = cx.clone(); - |message: Dispatch| async move { - match message { - Dispatch::Request(req, responder) => { - cx.spawn(async move { - match agent.dispatch_custom_request(&req.method, req.params).await { - Ok(json) => responder.respond(json)?, - Err(e) => responder.respond_with_error(e)?, - } - Ok(()) - })?; + .otherwise(|message: Dispatch| async { + match message { + Dispatch::Request(req, responder) => { + let agent = agent.clone(); + let cx_spawn = cx.clone(); + let method = req.method; + let params = req.params; + cx.spawn(async move { + match agent + .dispatch_custom_request(&cx_spawn, &method, params) + .await + { + Ok(json) => responder.respond(json)?, + Err(e) => responder.respond_with_error(e)?, + } Ok(()) - } - Dispatch::Response(result, router) => { - debug!(method = %router.method(), id = %router.id(), ok = result.is_ok(), "routing response"); - router.respond_with_result(result)?; - Ok(()) - } - Dispatch::Notification(notif) => { - debug!(method = %notif.method, "unhandled notification"); - Ok(()) - } + })?; + Ok(()) + } + Dispatch::Response(result, router) => { + debug!(method = %router.method(), id = %router.id(), ok = result.is_ok(), "routing response"); + router.respond_with_result(result)?; + Ok(()) + } + Dispatch::Notification(notif) => { + debug!(method = %notif.method, "unhandled notification"); + Ok(()) } } })