fix(runtime): bound worker request payloads

Author: Aitomates

maf_starter/loop_worker_cli.py

@@ -14,6 +14,8 @@
     run_bounded_specialists,
 )
 
+MAX_REQUEST_BYTES = 1_000_000
+
 
 @dataclass(frozen=True)
 class WorkerEnvelope:
@@ -68,7 +70,15 @@ def main() -> int:
     parser.add_argument("--request", help="JSON worker request; stdin is used when omitted")
     args = parser.parse_args()
     try:
-        payload = json.loads(args.request if args.request is not None else sys.stdin.read())
+        if args.request is not None:
+            request_text = args.request
+            request_size = len(request_text.encode("utf-8"))
+        else:
+            request_text = sys.stdin.read(MAX_REQUEST_BYTES + 1)
+            request_size = len(request_text.encode("utf-8"))
+        if request_size > MAX_REQUEST_BYTES:
+            raise ValueError(f"Worker request exceeds {MAX_REQUEST_BYTES} bytes")
+        payload = json.loads(request_text)
         result = asyncio.run(execute_request(payload))
     except (ValueError, json.JSONDecodeError) as error:
         print(json.dumps({"error": str(error)}), file=sys.stderr)

tests/test_loop_worker_cli.py

@@ -1,7 +1,12 @@
 import asyncio
+import io
+import json
+import sys
 import unittest
+from contextlib import redirect_stderr
+from unittest.mock import patch
 
-from maf_starter.loop_worker_cli import execute_request
+from maf_starter.loop_worker_cli import MAX_REQUEST_BYTES, execute_request, main
 
 
 class LoopWorkerCliTests(unittest.TestCase):
@@ -22,3 +27,17 @@ def test_cli_request_executes_real_bounded_specialist_runtime(self) -> None:
         self.assertEqual(3, result.peakConcurrency)
         self.assertEqual(("research", "architecture", "security", "test"), result.roles)
         self.assertEqual(4, len(result.evidenceUris))
+
+    def test_cli_rejects_oversized_stdin_before_json_parsing(self) -> None:
+        oversized = json.dumps({"padding": "x" * MAX_REQUEST_BYTES})
+        stderr = io.StringIO()
+
+        with (
+            patch.object(sys, "argv", ["loop_worker_cli"]),
+            patch.object(sys, "stdin", io.StringIO(oversized)),
+            redirect_stderr(stderr),
+        ):
+            exit_code = main()
+
+        self.assertEqual(2, exit_code)
+        self.assertIn("exceeds", stderr.getvalue())