diff --git a/concordia/utils/simulation_server.py b/concordia/utils/simulation_server.py index 6c1f94f84..b9a9f4de0 100644 --- a/concordia/utils/simulation_server.py +++ b/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'] diff --git a/concordia/utils/simulation_server_test.py b/concordia/utils/simulation_server_test.py index 029625592..49e85f3d0 100644 --- a/concordia/utils/simulation_server_test.py +++ b/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)