Consume rejected editor request bodies before responding

concordia/utils/simulation_server.py

@@ -341,23 +341,26 @@ def _serve_server_sent_events(self) -> None:
 
       def _handle_set_component_state(self) -> None:
         """Handle POST /cmd/set_component_state for dynamic editing."""
-        if not server.step_controller.is_paused:
-          self._send_json({
-              'status': 'error',
-              'message': 'Simulation must be paused to edit state.',
-          })
-          return
-
-        if server.simulation is None:
-          self._send_json({
-              'status': 'error',
-              'message': 'No simulation instance available.',
-          })
-          return
-
         try:
           content_length = int(self.headers.get('Content-Length', 0))
           body = self.rfile.read(content_length)
+          # Consume the framed request before replying, even when editing is
+          # unavailable. Closing with unread body bytes can reset the socket
+          # before the client has received the JSON error response.
+          if not server.step_controller.is_paused:
+            self._send_json({
+                'status': 'error',
+                'message': 'Simulation must be paused to edit state.',
+            })
+            return
+
+          if server.simulation is None:
+            self._send_json({
+                'status': 'error',
+                'message': 'No simulation instance available.',
+            })
+            return
+
           data = json.loads(body.decode('utf-8'))
 
           entity_name = data['entity_name']

concordia/utils/simulation_server_test.py

@@ -177,6 +177,8 @@ def test_unknown_path_returns_404(self):
     self.assertEqual(cm.exception.code, 404)
 
   def test_set_component_state_requires_paused(self):
+    fake_sim = _FakeSimulation()
+    self.server.set_simulation(fake_sim)
     self.server.step_controller.play()
     status, payload = _request(
         self.base_url + '/cmd/set_component_state',
@@ -191,6 +193,45 @@ def test_set_component_state_requires_paused(self):
     self.assertEqual(status, 200)
     self.assertEqual(payload['status'], 'error')
     self.assertIn('paused', payload['message'])
+    self.assertEmpty(fake_sim.set_calls)
+    self.assertIsNone(self.server.cached_entity_info)
+
+  def test_rejected_malformed_body_preserves_state_error_and_recovers(self):
+    fake_sim = _FakeSimulation()
+    # Read both rejection responses fully, including when the body is too
+    # malformed to decode. State guards still take precedence over parsing.
+    for message in ('No simulation', 'paused'):
+      with self.subTest(message=message):
+        if message == 'paused':
+          self.server.set_simulation(fake_sim)
+          self.server.step_controller.play()
+        req = urllib.request.Request(
+            self.base_url + '/cmd/set_component_state',
+            data=b'not JSON: ' + bytes([255]) * 65536,
+            headers={'Content-Type': 'application/json'},
+        )
+        with urllib.request.urlopen(req, timeout=5) as response:
+          self.assertEqual(response.status, 200)
+          payload = json.loads(response.read().decode('utf-8'))
+        self.assertEqual(payload['status'], 'error')
+        self.assertIn(message, payload['message'])
+        self.assertEmpty(fake_sim.set_calls)
+        self.assertIsNone(self.server.cached_entity_info)
+
+    self.server.step_controller.pause()
+    status, payload = _request(
+        self.base_url + '/cmd/set_component_state',
+        method='POST',
+        data={
+            'entity_name': 'alice',
+            'component_name': 'memory',
+            'key': 'foo',
+            'value': 'bar',
+        },
+    )
+    self.assertEqual(status, 200)
+    self.assertEqual(payload['status'], 'ok')
+    self.assertEqual(fake_sim.set_calls, [('alice', 'memory', 'foo', 'bar')])
 
   def test_set_component_state_requires_simulation(self):
     self.assertTrue(self.server.step_controller.is_paused)