Pass connection handle into ACP request handlers

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<GetExtensionsResponse, agent_client_protocol::Error> { .. }
+/// async fn on_get_extensions(
+///     &self,
+///     cx: &ConnectionTo<Client>,
+/// ) -> Result<GetExtensionsResponse, agent_client_protocol::Error> { .. }
 ///
 /// // Typed params — JSON params auto-deserialized
 /// #[custom_method(GetSessionRequest)]
-/// async fn on_get_session(&self, req: GetSessionRequest) -> Result<GetSessionResponse, agent_client_protocol::Error> { .. }
+/// async fn on_get_session(
+///     &self,
+///     cx: &ConnectionTo<Client>,
+///     req: GetSessionRequest,
+/// ) -> Result<GetSessionResponse, agent_client_protocol::Error> { .. }
 /// ```
 ///
 /// The return type must be `Result<T, agent_client_protocol::Error>` 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<agent_client_protocol::Client>,
             method: &str,
             params: serde_json::Value,
         ) -> Result<serde_json::Value, agent_client_protocol::Error> {
@@ -221,8 +229,10 @@ struct Route {
     ok_type: Option<Type>,
 }
 
-/// 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<Type> {
+    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<Type> {
                     continue;
                 }
             }
+
+            if !saw_connection {
+                saw_connection = true;
+                continue;
+            }
+
             return Some((*pat_type.ty).clone());
         }
     }

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<Client>,
         method: &str,
         params: serde_json::Value,
     ) -> Result<serde_json::Value, agent_client_protocol::Error> {
-        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<Client>,
         req: AddExtensionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_add_extension(req).await
@@ -22,6 +24,7 @@ impl GooseAcpAgent {
     #[custom_method(RemoveExtensionRequest)]
     async fn dispatch_remove_extension(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: RemoveExtensionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_remove_extension(req).await
@@ -30,6 +33,7 @@ impl GooseAcpAgent {
     #[custom_method(GetToolsRequest)]
     async fn dispatch_get_tools(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: GetToolsRequest,
     ) -> Result<GetToolsResponse, agent_client_protocol::Error> {
         self.on_get_tools(req).await
@@ -38,6 +42,7 @@ impl GooseAcpAgent {
     #[custom_method(GooseToolCallRequest)]
     async fn dispatch_call_tool(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: GooseToolCallRequest,
     ) -> Result<GooseToolCallResponse, agent_client_protocol::Error> {
         self.on_call_tool(req).await
@@ -46,6 +51,7 @@ impl GooseAcpAgent {
     #[custom_method(ReadResourceRequest)]
     async fn dispatch_read_resource(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ReadResourceRequest,
     ) -> Result<ReadResourceResponse, agent_client_protocol::Error> {
         self.on_read_resource(req).await
@@ -54,6 +60,7 @@ impl GooseAcpAgent {
     #[custom_method(UpdateWorkingDirRequest)]
     async fn dispatch_update_working_dir(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: UpdateWorkingDirRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         req: SetSessionSystemPromptRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         req: DeleteSessionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_delete_session(req).await
@@ -78,6 +87,7 @@ impl GooseAcpAgent {
     #[custom_method(GetExtensionsRequest)]
     async fn dispatch_get_extensions(
         &self,
+        _cx: &ConnectionTo<Client>,
     ) -> Result<GetExtensionsResponse, agent_client_protocol::Error> {
         self.on_get_extensions().await
     }
@@ -85,6 +95,7 @@ impl GooseAcpAgent {
     #[custom_method(AddConfigExtensionRequest)]
     async fn dispatch_add_config_extension(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: AddConfigExtensionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         req: RemoveConfigExtensionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ToggleConfigExtensionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         req: GetSessionExtensionsRequest,
     ) -> Result<GetSessionExtensionsResponse, agent_client_protocol::Error> {
         self.on_get_session_extensions(req).await
@@ -117,6 +131,7 @@ impl GooseAcpAgent {
     #[custom_method(ListProvidersRequest)]
     async fn dispatch_list_providers(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ListProvidersRequest,
     ) -> Result<ListProvidersResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderSupportedModelsListRequest,
     ) -> Result<ProviderSupportedModelsListResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderCatalogListRequest,
     ) -> Result<ProviderCatalogListResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderSetupCatalogListRequest,
     ) -> Result<ProviderSetupCatalogListResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderCatalogTemplateRequest,
     ) -> Result<ProviderCatalogTemplateResponse, agent_client_protocol::Error> {
         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<Client>,
         req: CustomProviderCreateRequest,
     ) -> Result<CustomProviderCreateResponse, agent_client_protocol::Error> {
         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<Client>,
         req: CustomProviderReadRequest,
     ) -> Result<CustomProviderReadResponse, agent_client_protocol::Error> {
         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<Client>,
         req: CustomProviderUpdateRequest,
     ) -> Result<CustomProviderUpdateResponse, agent_client_protocol::Error> {
         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<Client>,
         req: CustomProviderDeleteRequest,
     ) -> Result<CustomProviderDeleteResponse, agent_client_protocol::Error> {
         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<Client>,
         req: RefreshProviderInventoryRequest,
     ) -> Result<RefreshProviderInventoryResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderConfigReadRequest,
     ) -> Result<ProviderConfigReadResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderConfigStatusRequest,
     ) -> Result<ProviderConfigStatusResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderConfigSaveRequest,
     ) -> Result<ProviderConfigChangeResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderConfigDeleteRequest,
     ) -> Result<ProviderConfigChangeResponse, agent_client_protocol::Error> {
         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<Client>,
         req: ProviderConfigAuthenticateRequest,
     ) -> Result<ProviderConfigChangeResponse, agent_client_protocol::Error> {
         self.on_authenticate_provider_config(req).await
@@ -237,6 +266,7 @@ impl GooseAcpAgent {
     #[custom_method(PreferencesReadRequest)]
     async fn dispatch_preferences_read(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: PreferencesReadRequest,
     ) -> Result<PreferencesReadResponse, agent_client_protocol::Error> {
         self.on_preferences_read(req).await
@@ -245,6 +275,7 @@ impl GooseAcpAgent {
     #[custom_method(PreferencesSaveRequest)]
     async fn dispatch_preferences_save(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: PreferencesSaveRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_preferences_save(req).await
@@ -253,6 +284,7 @@ impl GooseAcpAgent {
     #[custom_method(PreferencesRemoveRequest)]
     async fn dispatch_preferences_remove(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: PreferencesRemoveRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_preferences_remove(req).await
@@ -261,6 +293,7 @@ impl GooseAcpAgent {
     #[custom_method(DefaultsReadRequest)]
     async fn dispatch_defaults_read(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: DefaultsReadRequest,
     ) -> Result<DefaultsReadResponse, agent_client_protocol::Error> {
         self.on_defaults_read(req).await
@@ -269,6 +302,7 @@ impl GooseAcpAgent {
     #[custom_method(DefaultsSaveRequest)]
     async fn dispatch_defaults_save(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: DefaultsSaveRequest,
     ) -> Result<DefaultsReadResponse, agent_client_protocol::Error> {
         self.on_defaults_save(req).await
@@ -277,6 +311,7 @@ impl GooseAcpAgent {
     #[custom_method(OnboardingImportScanRequest)]
     async fn dispatch_onboarding_import_scan(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: OnboardingImportScanRequest,
     ) -> Result<OnboardingImportScanResponse, agent_client_protocol::Error> {
         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<Client>,
         req: OnboardingImportApplyRequest,
     ) -> Result<OnboardingImportApplyResponse, agent_client_protocol::Error> {
         self.on_onboarding_import_apply(req).await
@@ -293,6 +329,7 @@ impl GooseAcpAgent {
     #[custom_method(ExportSessionRequest)]
     async fn dispatch_export_session(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ExportSessionRequest,
     ) -> Result<ExportSessionResponse, agent_client_protocol::Error> {
         self.on_export_session(req).await
@@ -301,6 +338,7 @@ impl GooseAcpAgent {
     #[custom_method(ImportSessionRequest)]
     async fn dispatch_import_session(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ImportSessionRequest,
     ) -> Result<ImportSessionResponse, agent_client_protocol::Error> {
         self.on_import_session(req).await
@@ -309,6 +347,7 @@ impl GooseAcpAgent {
     #[custom_method(UpdateSessionProjectRequest)]
     async fn dispatch_update_session_project(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: UpdateSessionProjectRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_update_session_project(req).await
@@ -317,6 +356,7 @@ impl GooseAcpAgent {
     #[custom_method(RenameSessionRequest)]
     async fn dispatch_rename_session(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: RenameSessionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_rename_session(req).await
@@ -325,6 +365,7 @@ impl GooseAcpAgent {
     #[custom_method(ArchiveSessionRequest)]
     async fn dispatch_archive_session(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ArchiveSessionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_archive_session(req).await
@@ -333,6 +374,7 @@ impl GooseAcpAgent {
     #[custom_method(UnarchiveSessionRequest)]
     async fn dispatch_unarchive_session(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: UnarchiveSessionRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_unarchive_session(req).await
@@ -341,6 +383,7 @@ impl GooseAcpAgent {
     #[custom_method(CreateSourceRequest)]
     async fn dispatch_create_source(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: CreateSourceRequest,
     ) -> Result<CreateSourceResponse, agent_client_protocol::Error> {
         self.on_create_source(req).await
@@ -349,6 +392,7 @@ impl GooseAcpAgent {
     #[custom_method(ListSourcesRequest)]
     async fn dispatch_list_sources(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ListSourcesRequest,
     ) -> Result<ListSourcesResponse, agent_client_protocol::Error> {
         self.on_list_sources(req).await
@@ -357,6 +401,7 @@ impl GooseAcpAgent {
     #[custom_method(UpdateSourceRequest)]
     async fn dispatch_update_source(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: UpdateSourceRequest,
     ) -> Result<UpdateSourceResponse, agent_client_protocol::Error> {
         self.on_update_source(req).await
@@ -365,6 +410,7 @@ impl GooseAcpAgent {
     #[custom_method(DeleteSourceRequest)]
     async fn dispatch_delete_source(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: DeleteSourceRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_delete_source(req).await
@@ -373,6 +419,7 @@ impl GooseAcpAgent {
     #[custom_method(ExportSourceRequest)]
     async fn dispatch_export_source(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ExportSourceRequest,
     ) -> Result<ExportSourceResponse, agent_client_protocol::Error> {
         self.on_export_source(req).await
@@ -381,6 +428,7 @@ impl GooseAcpAgent {
     #[custom_method(ImportSourcesRequest)]
     async fn dispatch_import_sources(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: ImportSourcesRequest,
     ) -> Result<ImportSourcesResponse, agent_client_protocol::Error> {
         self.on_import_sources(req).await
@@ -389,6 +437,7 @@ impl GooseAcpAgent {
     #[custom_method(DictationTranscribeRequest)]
     async fn dispatch_dictation_transcribe(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: DictationTranscribeRequest,
     ) -> Result<DictationTranscribeResponse, agent_client_protocol::Error> {
         self.on_dictation_transcribe(req).await
@@ -397,6 +446,7 @@ impl GooseAcpAgent {
     #[custom_method(DictationConfigRequest)]
     async fn dispatch_dictation_config(
         &self,
+        _cx: &ConnectionTo<Client>,
         _req: DictationConfigRequest,
     ) -> Result<DictationConfigResponse, agent_client_protocol::Error> {
         self.on_dictation_config(_req).await
@@ -405,6 +455,7 @@ impl GooseAcpAgent {
     #[custom_method(DictationSecretSaveRequest)]
     async fn dispatch_dictation_secret_save(
         &self,
+        _cx: &ConnectionTo<Client>,
         req: DictationSecretSaveRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         req: DictationSecretDeleteRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         _req: DictationModelsListRequest,
     ) -> Result<DictationModelsListResponse, agent_client_protocol::Error> {
         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<Client>,
         _req: DictationModelDownloadRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         _req: DictationModelDownloadProgressRequest,
     ) -> Result<DictationModelDownloadProgressResponse, agent_client_protocol::Error> {
         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<Client>,
         _req: DictationModelCancelRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         _req: DictationModelDeleteRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         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<Client>,
         req: DictationModelSelectRequest,
     ) -> Result<EmptyResponse, agent_client_protocol::Error> {
         self.on_dictation_model_select(req).await

crates/goose/src/acp/server/dispatch.rs

@@ -367,30 +367,33 @@ impl HandleDispatchFrom<Client> 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(())
                         }
                     }
                 })