fix(codemode): harden saved snippet execution

config/config.example.toml

@@ -191,7 +191,7 @@
 # `callTool(id, params)`, and typed `codemode.<upstream>.<tool>(params)`
 # helpers generated from connected servers' inputSchemas.
 # timeout_ms = 30000                        # valid range: 1..=60000  (wall-clock budget)
-# max_source_bytes = 131072                 # valid range: 1024..=1048576  (JavaScript source)
+# max_source_bytes = 1048576                # valid range: 1024..=1048576  (JavaScript source)
 # max_response_bytes = 24576                # valid range: 1024..=1048576
 # max_response_tokens = 6000                # valid range: 256..=256000
 # token_estimate_divisor = 4                # valid range: 1..=64  (lower is more conservative)

crates/labby-codemode/src/config.rs

@@ -27,7 +27,10 @@ pub(crate) const MAX_SNIPPET_RESOLVES_PER_RUN: usize = 32;
 pub(crate) const MAX_INTERNAL_CALLS_PER_RUN: usize = 32;
 
 /// Maximum total bytes of resolved snippet source allowed in a single run.
-pub(crate) const MAX_SNIPPET_RESOLVED_BYTES_PER_RUN: usize = 256 * 1024;
+/// Keep nested/composed snippets on the same hard byte budget as a direct Code
+/// Mode source so reusable helpers can carry normal agent context without
+/// inheriting a smaller legacy ceiling.
+pub(crate) const MAX_SNIPPET_RESOLVED_BYTES_PER_RUN: usize = MAX_SOURCE_BYTES;
 
 /// Default per-run `callTool` fan-out budget.
 const DEFAULT_MAX_CALLTOOL_PER_RUN: u64 = 512;
@@ -129,4 +132,9 @@ mod tests {
     fn max_source_bytes_is_stable() {
         assert_eq!(MAX_SOURCE_BYTES, 1024 * 1024);
     }
+
+    #[test]
+    fn composed_snippet_budget_matches_code_mode_hard_source_budget() {
+        assert_eq!(MAX_SNIPPET_RESOLVED_BYTES_PER_RUN, MAX_SOURCE_BYTES);
+    }
 }

crates/labby-codemode/src/execute.rs

@@ -12,6 +12,7 @@ use crate::host::{CodeModeHost, ExecCtx, ToolCallOutcome, ToolsRender};
 use labby_runtime::{CodeModeConfig, CodeModeResultShapePolicy};
 
 use super::CodeModeBroker;
+use super::config::MAX_SOURCE_BYTES;
 use super::normalize_user_code;
 use super::shape::shape_final_result;
 use super::truncate::{response_within_budget, truncate_execution_response};
@@ -101,11 +102,20 @@ impl<H: CodeModeHost> CodeModeBroker<'_, H> {
             }
             .into());
         }
+        let max_source_bytes = config.max_source_bytes.min(MAX_SOURCE_BYTES);
+        if code.len() > max_source_bytes {
+            return Err(ToolError::InvalidParam {
+                message: format!("code exceeds max length {max_source_bytes} bytes"),
+                param: "code".to_string(),
+            }
+            .into());
+        }
         let started = std::time::Instant::now();
         let mut response = self
             .execute_sandboxed(
                 code,
                 execution_timeout(config.timeout_ms),
+                max_source_bytes,
                 caller,
                 surface,
                 config.max_log_entries,
@@ -258,6 +268,7 @@ impl<H: CodeModeHost> CodeModeBroker<'_, H> {
         &self,
         code: &str,
         timeout: Duration,
+        snippet_max_bytes: usize,
         caller: CodeModeCaller,
         surface: CodeModeSurface,
         max_log_entries: usize,
@@ -325,6 +336,7 @@ impl<H: CodeModeHost> CodeModeBroker<'_, H> {
             max_log_bytes,
             trace_params,
             scope,
+            snippet_max_bytes,
             execution_id,
         )
         .await
@@ -877,6 +889,31 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn broker_enforces_configured_source_limit_before_runner_start() {
+        let host = NoopHost::default();
+        let broker = CodeModeBroker::new(Some(&host));
+        let config = CodeModeConfig {
+            max_source_bytes: 1024,
+            ..CodeModeConfig::default()
+        };
+        let code = format!("async () => \"{}\"", "x".repeat(2048));
+
+        let error = broker
+            .execute_with_raw_response(
+                &code,
+                CodeModeCaller::TrustedLocal,
+                CodeModeSurface::Cli,
+                config,
+                ToolScope::default(),
+                None,
+            )
+            .await
+            .expect_err("configured lower source limit must reject before runner start");
+
+        assert!(format!("{error}").contains("code exceeds max length 1024 bytes"));
+    }
+
     #[test]
     fn execution_timeout_reserves_response_delivery_margin() {
         assert_eq!(execution_timeout(30_000), Duration::from_millis(29_500));
@@ -1150,6 +1187,59 @@ mod tests {
         }
     }
 
+    #[tokio::test]
+    async fn raw_tool_call_boundary_rejects_undeclared_and_out_of_route_tools() {
+        let host = FixtureHost::new(vec![
+            CatalogDescriptor::tool("alpha", "tool1", "allowed", None, None),
+            CatalogDescriptor::tool("alpha", "other_tool", "undeclared sibling", None, None),
+            CatalogDescriptor::tool("beta", "tool2", "out of route", None, None),
+        ]);
+        let broker = CodeModeBroker::new(Some(&host));
+        let scope = ToolScope::scoped_namespaces(
+            vec!["alpha".to_string()],
+            vec!["alpha::tool1".to_string()],
+        );
+
+        for id in ["alpha::other_tool", "beta::tool2"] {
+            let error = broker
+                .call_tool_id(
+                    id,
+                    json!({}),
+                    CodeModeCaller::TrustedLocal,
+                    CodeModeSurface::Cli,
+                    &scope,
+                    ExecCtx::none(),
+                )
+                .await
+                .expect_err("raw callTool target outside the effective snippet scope must fail");
+            assert_eq!(error.kind(), "unknown_tool");
+            assert!(
+                error
+                    .to_string()
+                    .contains("outside this Code Mode execution capability set"),
+                "scope rejection must happen before host dispatch: {error:?}"
+            );
+        }
+
+        let allowed = broker
+            .call_tool_id(
+                "alpha::tool1",
+                json!({}),
+                CodeModeCaller::TrustedLocal,
+                CodeModeSurface::Cli,
+                &scope,
+                ExecCtx::none(),
+            )
+            .await
+            .expect_err("fixture host intentionally rejects real dispatch after scope admission");
+        assert!(
+            allowed
+                .to_string()
+                .contains("FixtureHost does not dispatch real tool calls"),
+            "the declared in-route tool must pass the scope boundary and reach the host"
+        );
+    }
+
     #[tokio::test]
     async fn dispatch_internal_call_resource_discovery_round_trips_uri() {
         let host = FixtureHost::new(Vec::new());

crates/labby-codemode/src/runner.rs

@@ -8,6 +8,7 @@ use serde_json::Value;
 
 use crate::CodeModeCallError;
 
+use super::config::MAX_SNIPPET_RESOLVED_BYTES_PER_RUN;
 use super::protocol::CODE_MODE_STACK_SIZE_LIMIT;
 use super::protocol::{
     CodeModeRunnerInput, CodeModeRunnerOutput, CodeModeRunnerResult, CodeModeRunnerState,
@@ -444,7 +445,7 @@ globalThis.__labSnippetResolveCount = 0;
 globalThis.__labSnippetResolvedBytes = 0;
 globalThis.__labSnippetMaxDepth = 8;
 globalThis.__labSnippetMaxResolves = 32;
-globalThis.__labSnippetMaxBytes = 262144;
+globalThis.__labSnippetMaxBytes = {snippet_max_bytes};
 {codec}
 globalThis.callTool = (id, params = {{}}) => {{
   if (typeof id !== "string" || id.trim() === "") {{
@@ -599,6 +600,7 @@ globalThis.__labMainPromise = (async () => {{
         codec = CODE_MODE_VALUE_CODEC_JS,
         invoker = invoker,
         proxy = proxy,
+        snippet_max_bytes = MAX_SNIPPET_RESOLVED_BYTES_PER_RUN,
     )
 }
 
@@ -956,3 +958,16 @@ fn runner_read_input() -> Result<CodeModeRunnerInput, RunnerReadError> {
         serde_json::from_str(&line).map_err(|err| RunnerReadError::Other(err.to_string()))
     })
 }
+
+#[cfg(test)]
+mod wrapper_tests {
+    use super::*;
+
+    #[test]
+    fn generated_wrapper_uses_the_shared_composed_snippet_budget() {
+        let wrapped = wrap_code_mode("async () => ({ ok: true })", "");
+        assert!(wrapped.contains(&format!(
+            "globalThis.__labSnippetMaxBytes = {MAX_SNIPPET_RESOLVED_BYTES_PER_RUN};"
+        )));
+    }
+}

crates/labby-codemode/src/runner_drive.rs

@@ -191,6 +191,8 @@ pub(crate) struct RunnerConfig {
     pub max_log_bytes: usize,
     pub trace_params: bool,
     pub capability_filter: ToolScope,
+    /// Effective total byte budget for source resolved through `codemode.run`.
+    pub snippet_max_bytes: usize,
     /// Durable-run execution id, minted by the caller (binary/gateway). `None`
     /// on the write-free/standalone path; flows into every [`ExecCtx`] so the
     /// host's `record_step` can key its per-execution journal buffer.
@@ -283,6 +285,7 @@ impl<H: CodeModeHost> CodeModeBroker<'_, H> {
         max_log_bytes: usize,
         trace_params: bool,
         capability_filter: ToolScope,
+        snippet_max_bytes: usize,
         execution_id: Option<Arc<str>>,
     ) -> Result<CodeModeExecutionResponse, CodeModeExecutionError> {
         // Read the openapi registry/client from the host at the config-build site
@@ -318,6 +321,7 @@ impl<H: CodeModeHost> CodeModeBroker<'_, H> {
             max_log_bytes,
             trace_params,
             capability_filter,
+            snippet_max_bytes: snippet_max_bytes.min(MAX_SNIPPET_RESOLVED_BYTES_PER_RUN),
             execution_id,
             openapi_registry,
             openapi_http_client,
@@ -1468,6 +1472,7 @@ mod tests {
             max_log_bytes: 4096,
             trace_params: false,
             capability_filter: ToolScope::default(),
+            snippet_max_bytes: MAX_SNIPPET_RESOLVED_BYTES_PER_RUN,
             execution_id: None,
             openapi_registry: labby_openapi::OpenApiRegistry::default(),
             openapi_http_client: labby_openapi::http::build_dispatch_client()

crates/labby-codemode/src/runner_drive/artifacts.rs

@@ -108,6 +108,10 @@ pub(super) async fn handle_snippet_resolve_event<H: CodeModeHost>(
     }
 }
 
+fn snippet_resolution_scope_allowed(caller: &CodeModeCaller, scope: &ToolScope) -> bool {
+    !scope.is_scoped() || matches!(caller, CodeModeCaller::TrustedLocal)
+}
+
 async fn resolve_snippet_for_runner<H: CodeModeHost>(
     broker: &CodeModeBroker<'_, H>,
     name: &str,
@@ -121,7 +125,7 @@ async fn resolve_snippet_for_runner<H: CodeModeHost>(
             required_scopes: vec!["lab:admin".to_string()],
         });
     }
-    if cfg.capability_filter.is_scoped() {
+    if !snippet_resolution_scope_allowed(&cfg.caller, &cfg.capability_filter) {
         return Err(ToolError::Forbidden {
             message: "codemode.run is not available on route-scoped Code Mode surfaces".to_string(),
             required_scopes: vec!["lab:admin".to_string()],
@@ -144,7 +148,7 @@ async fn resolve_snippet_for_runner<H: CodeModeHost>(
     let resolved = host.resolve_snippet(name, input).await?;
     let (name, code, input) = (resolved.name, resolved.code, resolved.input);
     state.snippet_resolved_bytes = state.snippet_resolved_bytes.saturating_add(code.len());
-    if state.snippet_resolved_bytes > MAX_SNIPPET_RESOLVED_BYTES_PER_RUN {
+    if state.snippet_resolved_bytes > cfg.snippet_max_bytes {
         return Err(ToolError::Sdk {
             sdk_kind: "snippet_budget_exceeded".to_string(),
             message: "resolved snippet code budget exceeded".to_string(),
@@ -262,8 +266,35 @@ fn artifact_call(
 
 #[cfg(test)]
 mod tests {
-    use super::artifact_writes_allowed;
+    use super::{artifact_writes_allowed, snippet_resolution_scope_allowed};
     use crate::ToolScope;
+    use crate::types::{CodeModeCaller, CodeModeCallerCapabilities};
+
+    #[test]
+    fn trusted_local_saved_snippets_may_compose_inside_declared_tool_scope() {
+        let scope = ToolScope::scoped_namespaces(
+            vec!["claude-macpoo".to_string()],
+            vec!["claude-macpoo::Bash".to_string()],
+        );
+        assert!(snippet_resolution_scope_allowed(
+            &CodeModeCaller::TrustedLocal,
+            &scope
+        ));
+
+        let route_scoped_admin = CodeModeCaller::Scoped {
+            capabilities: CodeModeCallerCapabilities {
+                can_read: true,
+                can_execute: true,
+                can_use_snippets: true,
+                is_admin: true,
+            },
+            sub: Some("admin".to_string()),
+        };
+        assert!(
+            !snippet_resolution_scope_allowed(&route_scoped_admin, &scope),
+            "route-scoped callers must not use nested snippet resolution to widen authority"
+        );
+    }
 
     #[test]
     fn artifact_writes_are_blocked_for_read_only_runs() {