perf: improve daemon worker and gateway concurrency

crates/cli/src/daemon/broker/registry.rs

@@ -226,6 +226,22 @@ impl Registry {
         fingerprint: Fingerprint,
         session_id: &McpSessionId,
     ) -> Result<BrokerDirective, RegistryError> {
+        // Most directive polls observe stable ready/pass-through state. Keep
+        // those polls on the shared read lock; recovering-with-target remains
+        // intentionally mutable because it transitions to Ready below.
+        let inner = self.read();
+        let route = inner
+            .routes
+            .get(&fingerprint)
+            .ok_or(RegistryError::UnknownRoute)?;
+        if !route.refs.contains_key(session_id) {
+            return Err(RegistryError::UnknownMcpSession);
+        }
+        if let Some(directive) = route.stable_current_directive(session_id, self.retry_after_ms) {
+            return Ok(directive);
+        }
+        drop(inner);
+
         let mut inner = self.write();
         let route = inner
             .routes
@@ -886,6 +902,38 @@ impl RouteEntry {
         }
     }
 
+    // Return a directive without mutating the route. `Recovering` with an
+    // existing target deliberately returns None because the write-path
+    // promotes it back to Ready as part of serving the directive.
+    fn stable_current_directive(
+        &self,
+        session_id: &McpSessionId,
+        retry_after_ms: u64,
+    ) -> Option<BrokerDirective> {
+        Some(match &self.state {
+            RouteState::Empty => BrokerDirective::WaitForWorker { retry_after_ms },
+            RouteState::Activating {
+                owner,
+                launch: active_launch,
+            } if owner == session_id => active_launch.clone().into_directive(),
+            RouteState::Activating { .. }
+            | RouteState::Draining { .. }
+            | RouteState::Recovering { target: None, .. } => {
+                BrokerDirective::WaitForWorker { retry_after_ms }
+            }
+            RouteState::Ready { target } if !target.control_available() => {
+                BrokerDirective::WaitForWorker { retry_after_ms }
+            }
+            RouteState::Ready { target } => BrokerDirective::ReuseWorker {
+                endpoint: target.endpoint().to_owned(),
+            },
+            RouteState::PassThrough { .. } => BrokerDirective::UsePassThrough,
+            RouteState::Recovering {
+                target: Some(_), ..
+            } => return None,
+        })
+    }
+
     fn after_reference_removed(
         &mut self,
         removed_session: &McpSessionId,

crates/cli/src/daemon/broker/server/socket.rs

@@ -310,6 +310,13 @@ async fn run(state: Arc<DaemonState>, role: ComponentRole, local: bool, socket:
                     if send.try_send(Message::Text(text.into())).is_err() { break; }
                     continue;
                 }
+                // ACKs affect only the peer's bounded directive queue. Challenge
+                // issuance has no routing effect either, so neither needs to
+                // wake every established control connection.
+                let local_control_command = matches!(
+                    &request.command,
+                    Command::Acknowledge { .. } | Command::Challenge(_)
+                );
                 let readiness = if let Command::Ready(payload) = &request.command {
                     Some(socket_ready(&state, role, id.as_deref(), &generation, payload).await)
                 } else {
@@ -328,7 +335,7 @@ async fn run(state: Arc<DaemonState>, role: ComponentRole, local: bool, socket:
                     None => dispatch(&state, role, &mut id, &mut challenge, request.command).await,
                 };
                 let success = response.status().is_success();
-                if success && let Some(id) = &id {
+                if success && !local_control_command && let Some(id) = &id {
                     let mut peers = lock(&state.sockets.peers);
                     let entry = peers.entry(key(role, id));
                     use std::collections::hash_map::Entry;
@@ -362,7 +369,9 @@ async fn run(state: Arc<DaemonState>, role: ComponentRole, local: bool, socket:
                 last_reply = Some((request.request_id, text.to_string(), event.clone()));
                 let Ok(text) = serde_json::to_string(&event) else { break };
                 if send.try_send(Message::Text(text.into())).is_err() { break; }
-                state.sockets.changed.notify_waiters();
+                if !local_control_command {
+                    state.sockets.changed.notify_waiters();
+                }
             }
         }
     }

crates/cli/src/daemon/worker/managed.rs

@@ -106,6 +106,11 @@ const INTERNAL_DISPATCH_ROUTE_HEADER: &str = "x-nemo-relay-internal-dispatch-rou
 const INTERNAL_DISPATCH_BACKEND_HEADER: &str = "x-nemo-relay-internal-dispatch-backend";
 const INTERNAL_RETRY_AWARE_HEADER: &str = "x-nemo-relay-internal-retry-aware";
 
+#[derive(Clone, Copy)]
+pub(super) struct ProviderMiddlewareRequirements {
+    request_body_decode_required: bool,
+}
+
 /// Runtime-owned plugin activation, hook sessions, and response observation.
 pub(super) struct ManagedRuntime {
     config: GatewayConfig,
@@ -125,7 +130,7 @@ impl ManagedRuntime {
         let activation =
             crate::server::initialize_plugin_host(config.plugin_config.clone(), dynamic_plugins)
                 .await?;
-        if let Err(error) = reject_incompatible_execution_middleware() {
+        if let Err(error) = provider_middleware_requirements() {
             if let Some(activation) = activation {
                 let _ = activation.clear();
             }
@@ -144,8 +149,18 @@ impl ManagedRuntime {
 
     /// Rechecks the transport contract before a provider body is polled. Plugin activation is
     /// normally static, but this also fails closed if a component installs middleware later.
+    #[cfg(test)]
     pub(super) fn ensure_streaming_transport_compatible(&self) -> Result<(), CliError> {
-        reject_incompatible_execution_middleware()
+        self.provider_middleware_requirements().map(|_| ())
+    }
+
+    /// Snapshots the middleware contract once for an incoming provider request.
+    /// This keeps late registration fail-closed while avoiding separate global
+    /// registry enumeration for raw-delivery compatibility and body decoding.
+    pub(super) fn provider_middleware_requirements(
+        &self,
+    ) -> Result<ProviderMiddlewareRequirements, CliError> {
+        provider_middleware_requirements()
     }
 
     pub(super) async fn close(&self) -> Result<(), CliError> {
@@ -256,16 +271,33 @@ impl ManagedRuntime {
         }
     }
 
+    #[cfg(test)]
     pub(super) async fn proxy_provider(
+        &self,
+        upstream: PooledClient,
+        request: Request<Body>,
+        route: ProviderRoute,
+    ) -> Result<Response<RelayBody>, CliError> {
+        self.proxy_provider_with_requirements(
+            upstream,
+            request,
+            route,
+            self.provider_middleware_requirements()?,
+        )
+        .await
+    }
+
+    pub(super) async fn proxy_provider_with_requirements(
         &self,
         upstream: PooledClient,
         mut request: Request<Body>,
         route: ProviderRoute,
+        middleware: ProviderMiddlewareRequirements,
     ) -> Result<Response<RelayBody>, CliError> {
         let Some(surface) = provider_surface(request.uri().path()) else {
             return dispatch_unmanaged(upstream, request, route, &self.config).await;
         };
-        if !request_body_decode_required()? {
+        if !middleware.request_body_decode_required {
             strip_worker_headers(request.headers_mut());
             strip_untrusted_dispatch_headers(request.headers_mut());
             let streaming_hint = request_streaming_hint(request.headers());
@@ -1015,14 +1047,30 @@ fn prepared_streaming(request: &LlmRequest) -> bool {
     stream_mode(request)
 }
 
-fn request_body_decode_required() -> Result<bool, CliError> {
+fn provider_middleware_requirements() -> Result<ProviderMiddlewareRequirements, CliError> {
     let kinds = BTreeSet::from([
         RuntimeRegistrationKind::LlmSanitizeRequestGuardrail,
         RuntimeRegistrationKind::LlmConditionalExecutionGuardrail,
         RuntimeRegistrationKind::LlmRequestIntercept,
+        RuntimeRegistrationKind::LlmExecutionIntercept,
+        RuntimeRegistrationKind::LlmStreamExecutionIntercept,
     ]);
     let registrations = list_runtime_registrations(Some(&kinds)).map_err(CliError::from)?;
-    Ok(registrations.iter().any(registration_reads_request_body))
+    let incompatible = incompatible_registration_names(&registrations);
+    if !incompatible.is_empty() {
+        return Err(CliError::Config(format!(
+            "daemon worker raw delivery is incompatible with LLM execution middleware: {}",
+            incompatible.join(", ")
+        )));
+    }
+    Ok(ProviderMiddlewareRequirements {
+        request_body_decode_required: registrations.iter().any(registration_reads_request_body),
+    })
+}
+
+#[cfg(test)]
+fn request_body_decode_required() -> Result<bool, CliError> {
+    provider_middleware_requirements().map(|requirements| requirements.request_body_decode_required)
 }
 
 fn registration_reads_request_body(registration: &RuntimeRegistrationIdentity) -> bool {
@@ -1034,24 +1082,9 @@ fn registration_reads_request_body(registration: &RuntimeRegistrationIdentity) -
     )
 }
 
+#[cfg(test)]
 fn reject_incompatible_execution_middleware() -> Result<(), CliError> {
-    // Execution intercepts own the provider callback and may replace, suppress, retry, or mutate
-    // its result. The raw worker transport cannot safely invoke that contract while also returning
-    // the provider's response head and frames unchanged. Request intercepts and conditional
-    // execution guardrails remain supported above the transport boundary.
-    let kinds = BTreeSet::from([
-        RuntimeRegistrationKind::LlmExecutionIntercept,
-        RuntimeRegistrationKind::LlmStreamExecutionIntercept,
-    ]);
-    let registrations = list_runtime_registrations(Some(&kinds)).map_err(CliError::from)?;
-    let incompatible = incompatible_registration_names(&registrations);
-    if incompatible.is_empty() {
-        return Ok(());
-    }
-    Err(CliError::Config(format!(
-        "daemon worker raw delivery is incompatible with LLM execution middleware: {}",
-        incompatible.join(", ")
-    )))
+    provider_middleware_requirements().map(|_| ())
 }
 
 fn incompatible_registration_names(registrations: &[RuntimeRegistrationIdentity]) -> Vec<String> {

crates/cli/src/daemon/worker/runtime.rs

@@ -565,12 +565,19 @@ async fn proxy(State(state): State<Arc<WorkerState>>, request: Request<Body>) ->
     let Some(route) = PublicRoute::from_path(request.uri().path()) else {
         return StatusCode::NOT_FOUND.into_response();
     };
-    if matches!(route, PublicRoute::Provider(_))
-        && let Some(managed) = state.managed.as_ref()
-        && let Err(error) = managed.ensure_streaming_transport_compatible()
-    {
-        return route_failure_response(error);
-    }
+    let middleware = if matches!(route, PublicRoute::Provider(_)) {
+        match state
+            .managed
+            .as_ref()
+            .map(|managed| managed.provider_middleware_requirements())
+        {
+            Some(Ok(requirements)) => Some(requirements),
+            Some(Err(error)) => return route_failure_response(error),
+            None => None,
+        }
+    } else {
+        None
+    };
     let Some(in_flight) = state.admit() else {
         let mut response = message(
             StatusCode::SERVICE_UNAVAILABLE,
@@ -601,7 +608,12 @@ async fn proxy(State(state): State<Arc<WorkerState>>, request: Request<Body>) ->
         PublicRoute::Provider(provider) => {
             if let Some(managed) = state.managed.as_ref() {
                 let response = managed
-                    .proxy_provider(state.upstream.clone(), request, provider)
+                    .proxy_provider_with_requirements(
+                        state.upstream.clone(),
+                        request,
+                        provider,
+                        middleware.expect("managed provider requests have middleware requirements"),
+                    )
                     .await;
                 return match response {
                     Ok(response) => {

crates/cli/src/gateway/mod.rs

@@ -161,10 +161,11 @@ pub(crate) async fn passthrough(
 ) -> Result<Response<Body>, CliError> {
     state.touch();
     let authorization = state.authorize_provider_request(request.headers_mut())?;
-    let prepared = prepare_gateway_request(&state.config, request, authorization).await?;
+    let mut prepared = prepare_gateway_request(&state.config, request, authorization).await?;
+    let start = take_llm_gateway_start(&mut prepared);
     let prep = state
         .sessions
-        .prepare_gateway_call(&prepared.headers, build_llm_gateway_start(&prepared))
+        .prepare_gateway_call(&prepared.headers, start)
         .await?;
     run_managed_gateway(state, prepared, prep).await
 }
@@ -747,8 +748,7 @@ fn client_sse_body(
         while let Some(item) = json_stream.next().await {
             match item {
                 Ok(event_json) => {
-                    let frame = encode_sse_frame(&event_json, route);
-                    yield Ok::<Bytes, CliError>(Bytes::from(frame));
+                    yield Ok::<Bytes, CliError>(encode_sse_frame(&event_json, route));
                 }
                 Err(error) => {
                     guard.finish().await;
@@ -864,19 +864,25 @@ impl Drop for GatewayCallGuard {
 // Formats one SSE frame from a parsed event payload. Anthropic and OpenAI Responses events carry
 // the event name in the `type` field, so it is mirrored back onto the `event:` line; OpenAI Chat
 // chunks have no event name and emit only `data:`.
-fn encode_sse_frame(event_json: &Value, route: ProviderRoute) -> String {
-    let serialized = serde_json::to_string(event_json).unwrap_or_else(|_| "null".to_string());
+fn encode_sse_frame(event_json: &Value, route: ProviderRoute) -> Bytes {
     let event_name = match route {
-        ProviderRoute::AnthropicMessages | ProviderRoute::OpenAiResponses => event_json
-            .get("type")
-            .and_then(Value::as_str)
-            .map(ToOwned::to_owned),
+        ProviderRoute::AnthropicMessages | ProviderRoute::OpenAiResponses => {
+            event_json.get("type").and_then(Value::as_str)
+        }
         _ => None,
     };
-    match event_name {
-        Some(name) => format!("event: {name}\ndata: {serialized}\n\n"),
-        None => format!("data: {serialized}\n\n"),
+    let mut frame = Vec::with_capacity(64);
+    if let Some(name) = event_name {
+        frame.extend_from_slice(b"event: ");
+        frame.extend_from_slice(name.as_bytes());
+        frame.push(b'\n');
+    }
+    frame.extend_from_slice(b"data: ");
+    if serde_json::to_writer(&mut frame, event_json).is_err() {
+        frame.extend_from_slice(b"null");
     }
+    frame.extend_from_slice(b"\n\n");
+    Bytes::from(frame)
 }
 
 // Forwards the buffered request to the upstream provider with only the safe request headers. This

crates/cli/src/gateway/request.rs

@@ -211,6 +211,7 @@ fn passthrough_body_error(error: axum::Error) -> CliError {
     }
 }
 
+#[cfg(test)]
 pub(super) fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGatewayStart {
     build_llm_gateway_start_from_parts(
         &request.headers,
@@ -221,6 +222,22 @@ pub(super) fn build_llm_gateway_start(request: &PreparedGatewayRequest) -> LlmGa
     )
 }
 
+/// Transfers the already-parsed request JSON into the session start event.
+///
+/// The gateway only needs this representation once after request preparation, so moving it avoids
+/// a second full copy for large provider prompts. The borrowed builder above remains available for
+/// callers that must retain a prepared request.
+pub(super) fn take_llm_gateway_start(request: &mut PreparedGatewayRequest) -> LlmGatewayStart {
+    let request_json = std::mem::take(&mut request.request_json);
+    build_llm_gateway_start_from_parts(
+        &request.headers,
+        &request.path,
+        request.provider,
+        request_json,
+        request.streaming,
+    )
+}
+
 pub(super) fn build_llm_gateway_start_from_parts(
     headers: &HeaderMap,
     path: &str,