Cap concurrent calls per upstream server

docs/operations.md

@@ -794,6 +794,7 @@ policy = "deny"
 
 - `oauth = true` on a `url` server enables `oauth login` (optional `scopes`, `client_name`, or `client_id` and `client_secret` for a client registered ahead of time, and `authorization_params`).
 - Rarely used stdio servers: `lazy = true` starts the process on the first call and stops it after `idle_secs` (default 600) without calls; `/healthz` shows it as `idle` meanwhile and its tools stay searchable. A failed start fails that call only, and the next call tries again.
+- Parallel calls: the gateway sends a server at most `max_concurrent` calls at once (default 8), and the rest of a script's `asyncio.gather` waits its turn, so a burst of calls does not trip the server's rate limit. A call that gets no turn within `tool_timeout_secs` fails with a `TimeoutError` without being sent, and nothing is retried. A spoke calling one of the hub's servers waits out the hub's queue too, using the longer of its own `tool_timeout_secs` and the hub's. The setting takes 1 to 4096. When a tool's error reads as a rate limit (`rate limit`, `ratelimit` or `too many requests`), the script's error names this setting and the file to change it in, which is the included file for a server kept in one, so the agent can ask for it to be lowered, e.g. `max_concurrent = 4` under `[servers.linear]`. Builtin servers, such as `secrets`, answer inside the daemon and are not capped, so a script's finds still reach the owner as one prompt.
 - Keys stay in 1Password: a `headers`, `env` or `url` value can hold `{{ op://vault/item/field }}` references, as in `op inject`, e.g. `headers = { authorization = "Bearer {{ op://Dev/Context7/credential }}" }`. The daemon runs `op read` the first time the server connects and keeps the value in memory until it exits. A clean stop (`service restart`, `update`, a reboot) hands the values to the next daemon sealed, so it asks nothing again; to pick up a rotated key, run `codemode-mcp service restart --fresh`. With the 1Password app integration the Mac asks once when the daemon starts (a locked or unanswered prompt shows `secret_unavailable` in `/healthz` and retries every 10 minutes). `[secrets] op` names the `op` binary (default `op` on `PATH`). A `url` that carries the key, as Zoho's MCP URLs do, should embed the reference where the key sits: `url = "https://<server>-<id>.zohomcp.com/mcp/{{ op://Personal/Zoho MCP - InfraOps/key }}/message"`. References in a `url` may sit only in path segments and query names or values, never in the scheme, credentials, host, port or fragment, and the `url` needs a literal host; anything else is refused when the config loads. A key embedded in a `url` may hold only ASCII letters, digits and `- . _ ~ + / = : @ ! * , ;` (at least 16 bytes), which covers hex, base64, base64url, JWT-like and colon-separated keys; any other character, such as `\`, quotes, `<`, `>`, `^`, `{`, `}`, `|`, `[`, `]`, `#`, `?`, `%`, `&`, `$`, `(`, `)` or a space, refuses the server with a configuration error. It is encoded as data where it sits, so `/` or `+` in a key stays part of that one segment or query value. The secrets are exactly the values 1Password returns, in the encodings libraries print them in (percent-encoded in either case, JSON-escaped, HTML character references such as `&#43;` or `&plus;`, `+` read as a space), matched ignoring ASCII case, and any run of 8 or more of its bytes is redacted too, since rmcp cuts invalid responses to 256 bytes and can leave only part of a key; nothing else in the URL is guessed to be a key. A value read for an embedded reference that is itself an `http(s)` URL is refused with a configuration error naming the reference, and dropped here and from the phones' copies, so the next attempt asks again instead of replaying it. A `url` that is one whole reference (`url = "{{ op://Personal/Zoho MCP infraops/credential }}"`) is refused when the config loads, with or without OAuth, because codemode cannot tell which part of such a URL is the key; embed the reference instead, e.g. `url = "https://host/mcp/{{ op://Vault/Item/key }}"`. Log text is matched through nested escape layers (a client's `Debug` output escapes the JSON escapes of a response body), up to 4 passes. Only the configured template is ever shown (companion prompts name the reference, never the value), and the resolved URL is redacted from errors, `/healthz` and the daemon's logs. Whatever OAuth would store (credentials, pending state, the client registration) or print (the authorization link) is refused when it holds a key, and the link's `resource` parameter is left out only when it holds the key.
 - Per tool: `policy` (`allow`, `ask`, `deny`) and, for `ask`, `grant` (`session` or `per_call`). A `session` grant is given to the calling client, not to one MCP session: it covers every running session of that client name for 15 or 60 minutes.
 - The server is session-less: it issues no `mcp-session-id`, for any protocol version. Clients are told apart by the client name in the request `_meta`, the `initialize` handshake or the `user-agent` header. While `execute` or `wait` is waiting, it sends progress notifications every `limits.progress_interval_ms` (15000) to clients that sent a `progressToken`.

src/config.rs

@@ -147,6 +147,10 @@ pub struct Config {
     /// nothing can be written back
     #[serde(skip)]
     path: Option<PathBuf>,
+    /// For servers set up across included files, the file that sets each
+    /// one's `max_concurrent`, or else the file holding its table
+    #[serde(skip)]
+    server_files: BTreeMap<String, PathBuf>,
 }
 
 impl Default for Config {
@@ -169,6 +173,7 @@ impl Default for Config {
             hub: None,
             servers: BTreeMap::new(),
             path: None,
+            server_files: BTreeMap::new(),
         }
     }
 }
@@ -261,6 +266,15 @@ impl Config {
         self.path.as_deref()
     }
 
+    /// The file to change `servers.<server>.max_concurrent` in: the one
+    /// that sets it, or else the one holding the server's table
+    pub fn max_concurrent_file(&self, server: &str) -> Option<&Path> {
+        self.server_files
+            .get(server)
+            .map(PathBuf::as_path)
+            .or(self.path())
+    }
+
     /// Parses and validates configuration text; `path` is used in errors and
     /// to resolve relative `include` paths
     ///
@@ -309,21 +323,35 @@ impl Config {
                 }
                 merged.add(table, path)?;
 
-                merged
-                    .table
-                    .try_into()
-                    .map_err(|source| ConfigError::ParseIncluded {
-                        files: format!(
-                            "{} together with {}",
-                            path.display(),
-                            included_paths
-                                .iter()
-                                .map(|included| included.display().to_string())
-                                .collect::<Vec<_>>()
-                                .join(", ")
-                        ),
-                        source,
-                    })?
+                let Merged { table, origins } = merged;
+                let mut config: Self =
+                    table
+                        .try_into()
+                        .map_err(|source| ConfigError::ParseIncluded {
+                            files: format!(
+                                "{} together with {}",
+                                path.display(),
+                                included_paths
+                                    .iter()
+                                    .map(|included| included.display().to_string())
+                                    .collect::<Vec<_>>()
+                                    .join(", ")
+                            ),
+                            source,
+                        })?;
+                // a value already set must be changed where it is, or the
+                // files would clash
+                for name in config.servers.keys() {
+                    let key = ["servers", name, "max_concurrent"].map(str::to_owned);
+                    if let Some(file) = (1..=key.len())
+                        .rev()
+                        .find_map(|end| origins.get(&key[..end]))
+                    {
+                        config.server_files.insert(name.clone(), file.clone());
+                    }
+                }
+
+                config
             }
         };
         config.validate()?;
@@ -693,6 +721,14 @@ fn default_idle_secs() -> u64 {
     10 * 60
 }
 
+/// Highest `max_concurrent`; more is no limit in practice, and the
+/// semaphore that enforces it panics past its own maximum
+const MAX_CONCURRENT: usize = 4096;
+
+fn default_max_concurrent() -> usize {
+    8
+}
+
 fn default_op() -> String {
     "op".to_owned()
 }
@@ -1254,6 +1290,11 @@ pub struct ServerConfig {
     /// Seconds a `lazy` server may go without calls before it is stopped
     #[serde(default = "default_idle_secs")]
     pub idle_secs: u64,
+    /// Calls the gateway sends this server at once; the rest wait their
+    /// turn, so a script's `asyncio.gather` cannot trip the server's rate
+    /// limit. Builtin servers answer inside the daemon and are not limited
+    #[serde(default = "default_max_concurrent")]
+    pub max_concurrent: usize,
     /// Endpoint of a Streamable HTTP server; may hold `{{ op://... }}`
     /// references for a URL that carries a key, and is then never shown
     /// resolved
@@ -1485,6 +1526,11 @@ impl ServerConfig {
     }
 
     fn validate(&self, name: &str) -> Result<(), ConfigError> {
+        if !(1..=MAX_CONCURRENT).contains(&self.max_concurrent) {
+            return Err(ConfigError::Invalid(format!(
+                "servers.{name}.max_concurrent must be from 1 to {MAX_CONCURRENT}"
+            )));
+        }
         for (tool, tool_override) in &self.tools {
             tool_override.validate(name, tool)?;
         }

src/gateway.rs

@@ -119,7 +119,9 @@ impl ToolCallError {
         match self {
             Self::Unknown(_) => ExcType::NameError,
             Self::Denied(_) | Self::NotApproved { .. } => ExcType::PermissionError,
-            Self::Upstream(UpstreamError::TimedOut { .. }) => ExcType::TimeoutError,
+            Self::Upstream(UpstreamError::TimedOut { .. } | UpstreamError::Busy { .. }) => {
+                ExcType::TimeoutError
+            }
             Self::ResultTooLarge { .. } => ExcType::ValueError,
             Self::Hub(error) => match error.kind {
                 HubErrorKind::Unknown => ExcType::NameError,
@@ -140,7 +142,9 @@ impl ToolCallError {
         match self {
             Self::Unknown(_) => Outcome::NameError,
             Self::Denied(_) | Self::NotApproved { .. } => Outcome::Denied,
-            Self::Upstream(UpstreamError::TimedOut { .. }) => Outcome::Timeout,
+            Self::Upstream(UpstreamError::TimedOut { .. } | UpstreamError::Busy { .. }) => {
+                Outcome::Timeout
+            }
             Self::Hub(error) => match error.kind {
                 HubErrorKind::Unknown => Outcome::NameError,
                 HubErrorKind::Denied => Outcome::Denied,
@@ -1017,8 +1021,19 @@ impl Gateway {
                 approval: call.hub_request.clone(),
                 sensitive: call.sensitive,
             };
+            // the hub's own limit governs how long the call may queue and
+            // run there; an older hub does not say, so ours stands in
+            let timeout = self
+                .hub_view()
+                .servers
+                .iter()
+                .find(|server| server.name == call.server)
+                .and_then(|server| server.call_timeout_secs)
+                .map_or(self.config.tool_timeout(), |secs| {
+                    Duration::from_secs(secs).max(self.config.tool_timeout())
+                });
             let answer = hub
-                .call(&forwarded, self.config.tool_timeout())
+                .call(&forwarded, timeout)
                 .await
                 .map_err(|error| hub_unreachable(&error))?;
             let value = match answer {
@@ -1042,7 +1057,7 @@ impl Gateway {
             .map_err(|error| match error {
                 // the upstream's own error message can echo configured secrets
                 UpstreamError::Service(rmcp::service::ServiceError::McpError(_)) => {
-                    ToolCallError::ToolFailed(self.redactor.apply(&error.public_message()))
+                    self.tool_failed(&call.server, self.redactor.apply(&error.public_message()))
                 }
                 error => ToolCallError::Upstream(error),
             })?;
@@ -1183,6 +1198,10 @@ impl Gateway {
                         .map(|server| server.aliases.clone())
                         .unwrap_or_default(),
                     tools: export_tools(&self.config, snapshot),
+                    call_timeout_secs: self
+                        .upstreams
+                        .get(&snapshot.name)
+                        .map(|upstream| upstream.call_timeout().as_secs()),
                 })
             })
             .collect();
@@ -1298,6 +1317,33 @@ fn omit_none_arguments(entry: &ToolEntry, mut arguments: Map<String, Value>) ->
 }
 
 impl Gateway {
+    /// A tool's error, with where to lower `max_concurrent` when it reads
+    /// as the server's rate limit, so the agent can say what to change
+    fn tool_failed(&self, server: &str, message: String) -> ToolCallError {
+        let limit = self
+            .config
+            .servers
+            .get(server)
+            .filter(|config| config.builtin.is_none())
+            .map(|config| config.max_concurrent);
+        let Some(limit) = limit.filter(|_| looks_rate_limited(&message)) else {
+            return ToolCallError::ToolFailed(message);
+        };
+        let config = self.config.max_concurrent_file(server).map_or_else(
+            || "the codemode-mcp config".to_owned(),
+            |path| format!("`{}`", path.display()),
+        );
+
+        // the path is the owner's own, but anything shown passes the
+        // redactor, which a value read from 1Password may share text with
+        ToolCallError::ToolFailed(self.redactor.apply(&format!(
+            "{message}\n\n`{server}` looks rate limited. codemode sends it at most {limit} calls \
+             at once; make fewer calls at once, or ask the user to set a lower `max_concurrent` under \
+             `[servers.{server}]` in {config} on the machine that runs `{server}`, then run \
+             `codemode-mcp service restart`"
+        )))
+    }
+
     /// Converts an MCP tool result to the value a script receives
     ///
     /// Structured content wins. Without it, text that is a JSON object or
@@ -1335,7 +1381,7 @@ impl Gateway {
                 self.redactor.apply(&text.join("\n"))
             };
 
-            return Err(ToolCallError::ToolFailed(message));
+            return Err(self.tool_failed(&call.server, message));
         }
         if matches!(call.output, OutputSpec::Text) {
             // a text tool that sends only structured content still gives the
@@ -1361,6 +1407,22 @@ impl Gateway {
     }
 }
 
+/// Whether an upstream's error reads as a rate limit, such as HTTP 429's
+/// `Too Many Requests`
+fn looks_rate_limited(message: &str) -> bool {
+    let message = message.to_ascii_lowercase();
+
+    [
+        "rate limit",
+        "rate-limit",
+        "rate_limit",
+        "ratelimit",
+        "too many requests",
+    ]
+    .iter()
+    .any(|phrase| message.contains(phrase))
+}
+
 /// Parses text blocks as one JSON object or array; scalars stay text because
 /// a tool answering `123` or `true` as prose is more likely than as data
 fn json_from_text(blocks: &[String]) -> Option<Value> {

src/hub.rs

@@ -93,6 +93,10 @@ pub struct HubServer {
     pub aliases: Vec<String>,
     /// Its tools
     pub tools: Vec<ExportedTool>,
+    /// Seconds the hub lets a call of it wait for a turn under
+    /// `max_concurrent`, and then run; absent from older hubs
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub call_timeout_secs: Option<u64>,
 }
 
 /// Body of `POST /hub/v1/gate`: asks the hub whether a call may run now
@@ -195,9 +199,10 @@ impl From<&ToolCallError> for HubError {
         let kind = match error {
             ToolCallError::Unknown(_) => HubErrorKind::Unknown,
             ToolCallError::Denied(_) | ToolCallError::NotApproved { .. } => HubErrorKind::Denied,
-            ToolCallError::Upstream(crate::upstream::UpstreamError::TimedOut { .. }) => {
-                HubErrorKind::TimedOut
-            }
+            ToolCallError::Upstream(
+                crate::upstream::UpstreamError::TimedOut { .. }
+                | crate::upstream::UpstreamError::Busy { .. },
+            ) => HubErrorKind::TimedOut,
             ToolCallError::ResultTooLarge { .. } => HubErrorKind::TooLarge,
             ToolCallError::Hub(error) => error.kind,
             _ => HubErrorKind::Failed,

src/hub/client.rs

@@ -40,6 +40,8 @@ use crate::{
 pub const ID_FILE: &str = "spoke.id";
 /// How long a request that needs no person may take
 const QUICK: Duration = Duration::from_secs(30);
+/// The longest call timeout a forwarded call is waited out for
+const LONGEST_CALL: Duration = Duration::from_secs(24 * 60 * 60);
 /// How long a request that may wait for the owner may take: a long poll,
 /// or a read the owner approves on a device
 const PATIENT: Duration = Duration::from_secs(15 * 60);
@@ -277,14 +279,21 @@ impl HubClient {
         self.post("/gate", call, QUICK).await
     }
 
-    /// Runs a call of the hub's tool there; `tool_timeout` is the spoke's
-    /// own limit, which the hub's may exceed
+    /// Runs a call of the hub's tool there; `tool_timeout` is the longer of
+    /// the spoke's limit and the one the hub lists for the server
+    ///
+    /// The hub may first queue the call behind the server's
+    /// `max_concurrent` for up to that timeout, so the wait covers that too
     pub async fn call(
         &self,
         call: &ForwardedCall,
         tool_timeout: Duration,
     ) -> Result<CallAnswer, HubClientError> {
-        self.post("/call", call, tool_timeout + QUICK).await
+        // the hub names its own timeout, so a huge one is held to a day
+        // rather than overflowing the deadline
+        let wait = tool_timeout.min(LONGEST_CALL).saturating_mul(2) + QUICK;
+
+        self.post("/call", call, wait).await
     }
 
     /// The hub's verdict on request `id`, after a long poll; `None` while it