Make invalid states unrepresentable: ScopedPath, MemoryConfig, TruncationReason, NonZero caps

src/args.rs

@@ -1,14 +1,15 @@
 use crate::tools::{DEFAULT_MAX_BYTES, DEFAULT_MAX_LINES, DEFAULT_MAX_RESULTS, OutputMode};
 use rmcp::schemars::{self, JsonSchema};
 use serde::{Deserialize, Serialize};
+use std::num::NonZeroUsize;
 
-fn default_max_results() -> usize {
+fn default_max_results() -> NonZeroUsize {
     DEFAULT_MAX_RESULTS
 }
-fn default_max_bytes() -> usize {
+fn default_max_bytes() -> NonZeroUsize {
     DEFAULT_MAX_BYTES
 }
-fn default_max_lines() -> usize {
+fn default_max_lines() -> NonZeroUsize {
     DEFAULT_MAX_LINES
 }
 const fn default_true() -> bool {
@@ -64,7 +65,7 @@ pub struct GrepArgs {
     pub after_context: usize,
     #[serde(default = "default_max_results")]
     #[schemars(description = "Maximum number of results to return (default 100)")]
-    pub max_results: usize,
+    pub max_results: NonZeroUsize,
     #[serde(default)]
     #[schemars(
         description = "Case-insensitive search (default false). Equivalent to prefixing pattern with (?i)."
@@ -88,7 +89,7 @@ pub struct GrepArgs {
     #[schemars(
         description = "Hard cap on total response size in bytes (default ~5 MiB). Truncates with a marker."
     )]
-    pub max_bytes: usize,
+    pub max_bytes: NonZeroUsize,
     #[serde(default)]
     #[schemars(
         description = "Output mode: 'files_with_matches' (default — list file paths only), 'content' (matching lines with line numbers), 'count' (per-file match tallies as path: N)."
@@ -107,7 +108,7 @@ pub struct FindArgs {
     pub pattern: String,
     #[serde(default = "default_max_results")]
     #[schemars(description = "Maximum number of results to return (default 100)")]
-    pub max_results: usize,
+    pub max_results: NonZeroUsize,
     #[serde(default)]
     #[schemars(description = "Include hidden files and directories (default false)")]
     pub include_hidden: bool,
@@ -133,10 +134,10 @@ pub struct CatArgs {
     pub offset: usize,
     #[serde(default = "default_max_lines")]
     #[schemars(description = "Maximum number of lines to return (default 2000)")]
-    pub max_lines: usize,
+    pub max_lines: NonZeroUsize,
     #[serde(default = "default_max_bytes")]
     #[schemars(description = "Maximum number of bytes to return (default ~5 MiB)")]
-    pub max_bytes: usize,
+    pub max_bytes: NonZeroUsize,
 }
 
 #[derive(Debug, Serialize, Deserialize, JsonSchema)]

src/cli.rs

@@ -1,7 +1,10 @@
 use clap::Parser;
 use std::net::SocketAddr;
+use std::num::NonZeroU32;
 use std::path::PathBuf;
 
+const DEFAULT_INITIALIZE_RATE: NonZeroU32 = NonZeroU32::new(12).unwrap();
+
 /// Command-line arguments for `code-mcp`.
 ///
 /// Parsed via clap. All fields are `pub(crate)` because they're only consumed
@@ -42,8 +45,8 @@ pub struct Args {
     /// per-minute rate (token bucket of capacity = rate, refilling over
     /// 60s). When exhausted, new initializes from that peer get 429 +
     /// Retry-After. Existing-session traffic is unaffected.
-    #[arg(long, default_value_t = 12)]
-    pub(crate) initialize_rate_per_min: u32,
+    #[arg(long, default_value_t = DEFAULT_INITIALIZE_RATE)]
+    pub(crate) initialize_rate_per_min: NonZeroU32,
 
     /// Trust the rightmost entry of `X-Forwarded-For` as the peer IP
     /// instead of the TCP socket address. Assumes a single trusted proxy

src/gate.rs

@@ -117,6 +117,12 @@ mod tests {
     use axum::Router;
     use axum::routing::any;
     use std::net::Ipv4Addr;
+    use std::num::NonZeroU32;
+
+    fn nz(n: u32) -> NonZeroU32 {
+        NonZeroU32::new(n).unwrap()
+    }
+
     use tower::ServiceExt;
 
     fn dummy_addr() -> SocketAddr {
@@ -148,7 +154,7 @@ mod tests {
         let ctx = Arc::new(GateCtx {
             sessions: Arc::new(LocalSessionManager::default()),
             max_sessions: 0, // would block POSTs
-            limiter: PeerLimiter::per_minute(1),
+            limiter: PeerLimiter::per_minute(nz(1)),
             trust_forwarded_for: false,
             activity: Arc::new(ActivityTracker::new()),
         });
@@ -166,7 +172,7 @@ mod tests {
         let ctx = Arc::new(GateCtx {
             sessions: Arc::new(LocalSessionManager::default()),
             max_sessions: 0,
-            limiter: PeerLimiter::per_minute(1),
+            limiter: PeerLimiter::per_minute(nz(1)),
             trust_forwarded_for: false,
             activity: Arc::new(ActivityTracker::new()),
         });
@@ -183,7 +189,7 @@ mod tests {
         let ctx = Arc::new(GateCtx {
             sessions: Arc::new(LocalSessionManager::default()),
             max_sessions: 0,
-            limiter: PeerLimiter::per_minute(100),
+            limiter: PeerLimiter::per_minute(nz(100)),
             trust_forwarded_for: false,
             activity: Arc::new(ActivityTracker::new()),
         });
@@ -201,7 +207,7 @@ mod tests {
         let ctx = Arc::new(GateCtx {
             sessions: Arc::new(LocalSessionManager::default()),
             max_sessions: 1000,
-            limiter: PeerLimiter::new(2.0, 0.001, 1024), // ~no refill
+            limiter: PeerLimiter::new(2.0, 0.001, 1024).unwrap(), // ~no refill
             trust_forwarded_for: false,
             activity: Arc::new(ActivityTracker::new()),
         });

src/limiter.rs

@@ -1,7 +1,18 @@
 use std::collections::HashMap;
 use std::net::IpAddr;
+use std::num::NonZeroU32;
 use std::sync::Mutex;
 use std::time::{Duration, Instant};
+use thiserror::Error;
+
+/// Returned by [`PeerLimiter::new`] when `capacity` or `refill_per_sec` is
+/// zero, negative, or NaN. `try_consume` divides by `refill_per_sec` and
+/// feeds the quotient to `Duration::from_secs_f64`, which panics on
+/// non-finite input — this error makes that state unconstructible instead
+/// of merely `debug_assert`-ed.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Error)]
+#[error("capacity and refill_per_sec must be positive and finite")]
+pub struct InvalidRate;
 
 /// Per-peer token-bucket limiter. Each peer (keyed by `IpAddr`) gets a
 /// bucket that refills continuously toward `capacity` at `refill_per_sec`.
@@ -22,20 +33,27 @@ struct Bucket {
 }
 
 impl PeerLimiter {
-    pub fn new(capacity: f64, refill_per_sec: f64, evict_threshold: usize) -> Self {
-        debug_assert!(capacity > 0.0 && refill_per_sec > 0.0);
-        Self {
+    pub fn new(
+        capacity: f64,
+        refill_per_sec: f64,
+        evict_threshold: usize,
+    ) -> Result<Self, InvalidRate> {
+        // NaN fails both comparisons, so it lands here too.
+        if !(capacity > 0.0 && refill_per_sec > 0.0) {
+            return Err(InvalidRate);
+        }
+        Ok(Self {
             capacity,
             refill_per_sec,
             buckets: Mutex::new(HashMap::new()),
             evict_threshold,
-        }
+        })
     }
 
     /// Convenience: a per-minute rate (capacity = `rate`, refill = `rate`/60s).
-    pub fn per_minute(rate: u32) -> Self {
-        let cap = f64::from(rate.max(1));
-        Self::new(cap, cap / 60.0, 4096)
+    pub fn per_minute(rate: NonZeroU32) -> Self {
+        let cap = f64::from(rate.get());
+        Self::new(cap, cap / 60.0, 4096).expect("derived from a NonZeroU32 rate")
     }
 
     /// Try to consume one token. Returns the time until the next token would
@@ -91,15 +109,15 @@ mod tests {
 
     #[test]
     fn allows_within_capacity() {
-        let l = PeerLimiter::new(3.0, 0.001, 1024); // refill effectively zero
+        let l = PeerLimiter::new(3.0, 0.001, 1024).unwrap(); // refill effectively zero
         for _ in 0..3 {
             assert!(l.try_consume(ip(1)).is_ok());
         }
     }
 
     #[test]
     fn blocks_after_exhaustion() {
-        let l = PeerLimiter::new(2.0, 0.001, 1024);
+        let l = PeerLimiter::new(2.0, 0.001, 1024).unwrap();
         l.try_consume(ip(1)).unwrap();
         l.try_consume(ip(1)).unwrap();
         let err = l
@@ -110,16 +128,25 @@ mod tests {
 
     #[test]
     fn separate_ips_get_separate_buckets() {
-        let l = PeerLimiter::new(1.0, 0.001, 1024);
+        let l = PeerLimiter::new(1.0, 0.001, 1024).unwrap();
         assert!(l.try_consume(ip(1)).is_ok());
         assert!(l.try_consume(ip(2)).is_ok());
         assert!(l.try_consume(ip(1)).is_err());
     }
 
+    #[test]
+    fn new_rejects_non_positive_rates() {
+        assert!(PeerLimiter::new(0.0, 1.0, 8).is_err());
+        assert!(PeerLimiter::new(1.0, 0.0, 8).is_err());
+        assert!(PeerLimiter::new(-1.0, 1.0, 8).is_err());
+        assert!(PeerLimiter::new(1.0, f64::NAN, 8).is_err());
+        assert!(PeerLimiter::new(1.0, 1.0, 8).is_ok());
+    }
+
     #[test]
     fn refills_over_time() {
         // 10 tokens/sec → one token per 100ms.
-        let l = PeerLimiter::new(1.0, 10.0, 1024);
+        let l = PeerLimiter::new(1.0, 10.0, 1024).unwrap();
         l.try_consume(ip(1)).unwrap();
         assert!(l.try_consume(ip(1)).is_err());
         std::thread::sleep(Duration::from_millis(150));
@@ -130,7 +157,7 @@ mod tests {
     fn evicts_stale_entries_when_threshold_exceeded() {
         // Tiny threshold + fast refill so the prior entry is "stale" by the
         // time we hit the threshold.
-        let l = PeerLimiter::new(1.0, 1000.0, 4);
+        let l = PeerLimiter::new(1.0, 1000.0, 4).unwrap();
         for n in 0..4 {
             l.try_consume(ip(n)).unwrap();
         }

src/main.rs

@@ -64,7 +64,7 @@ use crate::cli::Args;
 use crate::gate::{GateCtx, gate};
 use crate::limiter::PeerLimiter;
 use crate::scope::Scope;
-use crate::server::CodeMcpServer;
+use crate::server::{CodeMcpServer, MemoryConfig};
 
 use crate::error::AppError;
 
@@ -79,29 +79,34 @@ async fn main() -> Result<(), AppError> {
 
     let args = Args::parse();
 
-    // If a memory dir is configured, load <dir>/instructions.md once at startup.
-    // It's appended to the InitializeResult.instructions payload.
-    let extra_instructions = if let Some(dir) = args.memory_dir.as_ref() {
-        let path = dir.join("instructions.md");
-        match tokio::fs::read_to_string(&path).await {
-            Ok(s) => {
-                tracing::info!(path = %path.display(), "loaded extra instructions");
-                Some(s)
-            }
-            Err(e) if e.kind() == std::io::ErrorKind::NotFound => None,
-            Err(e) => {
-                tracing::warn!(path = %path.display(), error = %e, "could not read instructions.md");
-                None
+    // If a memory dir is configured, load <dir>/instructions.md once at
+    // startup. It's appended to the InitializeResult.instructions payload,
+    // so the dir and its extra instructions travel together in one type.
+    let memory = match args.memory_dir.as_ref() {
+        Some(dir) => {
+            let path = dir.join("instructions.md");
+            let extra = match tokio::fs::read_to_string(&path).await {
+                Ok(s) => {
+                    tracing::info!(path = %path.display(), "loaded extra instructions");
+                    Some(s)
+                }
+                Err(e) if e.kind() == std::io::ErrorKind::NotFound => None,
+                Err(e) => {
+                    tracing::warn!(path = %path.display(), error = %e, "could not read instructions.md");
+                    None
+                }
+            };
+            MemoryConfig::Enabled {
+                dir: dir.clone(),
+                extra,
             }
         }
-    } else {
-        None
+        None => MemoryConfig::Disabled,
     };
 
     let scope = Scope::new(args.project.clone())?;
     tracing::info!(root = %scope.root().display(), "project scope active");
 
-    let memory_dir = args.memory_dir.clone();
     let cancel = CancellationToken::new();
     let session_manager = Arc::new(LocalSessionManager::default());
     let sessions_for_gate = session_manager.clone();
@@ -118,13 +123,7 @@ async fn main() -> Result<(), AppError> {
         ..Default::default()
     };
     let service = StreamableHttpService::new(
-        move || {
-            Ok(CodeMcpServer::new(
-                memory_dir.clone(),
-                extra_instructions.clone(),
-                scope.clone(),
-            ))
-        },
+        move || Ok(CodeMcpServer::new(memory.clone(), scope.clone())),
         session_manager,
         config,
     );
@@ -138,7 +137,7 @@ async fn main() -> Result<(), AppError> {
     });
     tracing::info!(
         max_sessions = args.max_sessions,
-        initialize_rate_per_min = args.initialize_rate_per_min,
+        initialize_rate_per_min = args.initialize_rate_per_min.get(),
         trust_forwarded_for = args.trust_forwarded_for,
         session_idle_timeout_secs = args.session_idle_timeout_secs,
         session_sweep_interval_secs = args.session_sweep_interval_secs,

src/scope.rs

@@ -30,11 +30,12 @@ impl Scope {
         &self.root
     }
 
-    /// Validate that `input` is within the scope, returning its canonical
-    /// path. Symlinks in `input` are resolved before the containment
-    /// check, so a symlink inside the project that points outside it is
-    /// rejected.
-    pub fn check<P: AsRef<Path>>(&self, input: P) -> Result<PathBuf, AppError> {
+    /// Validate that `input` is within the scope, returning a [`ScopedPath`]
+    /// — the canonical path, with the in-scope proof carried by the type so
+    /// downstream tool entry points can require it. Symlinks in `input` are
+    /// resolved before the containment check, so a symlink inside the
+    /// project that points outside it is rejected.
+    pub fn check<P: AsRef<Path>>(&self, input: P) -> Result<ScopedPath, AppError> {
         let input = input.as_ref();
         let canon = input
             .canonicalize()
@@ -46,7 +47,20 @@ impl Scope {
                 self.root.display()
             )));
         }
-        Ok(canon)
+        Ok(ScopedPath(canon))
+    }
+}
+
+/// A path proven — via [`Scope::check`] — to canonicalize inside a scope's
+/// root. The only constructor is `Scope::check`, and tool entry points take
+/// `&ScopedPath`, so a path that skipped scope validation cannot reach them:
+/// the "forgot to call `check`" bug is a compile error, not a security hole.
+#[derive(Debug, Clone)]
+pub struct ScopedPath(PathBuf);
+
+impl AsRef<Path> for ScopedPath {
+    fn as_ref(&self) -> &Path {
+        &self.0
     }
 }
 
@@ -64,7 +78,11 @@ mod tests {
         fs::write(td.path().join("a.txt"), "x")?;
         let s = Scope::new(td.path())?;
         let canon = s.check(td.path().join("a.txt"))?;
-        assert!(canon.ends_with("a.txt"), "got {}", canon.display());
+        assert!(
+            canon.as_ref().ends_with("a.txt"),
+            "got {}",
+            canon.as_ref().display()
+        );
         Ok(())
     }
 
@@ -109,10 +127,10 @@ mod tests {
             // Either way, reading /etc/passwd via traversal should not succeed.
             Err(AppError::NotFound(_)) => Ok(()),
             Ok(p) => {
-                if p.starts_with(td.path()) {
+                if p.as_ref().starts_with(td.path()) {
                     Ok(())
                 } else {
-                    Err(format!("escaped root: {}", p.display()).into())
+                    Err(format!("escaped root: {}", p.as_ref().display()).into())
                 }
             }
             other => Err(format!("unexpected: {:?}", other).into()),