diff --git a/docs/CORTEX.md b/docs/CORTEX.md index 954da31b3..7b4596e07 100644 --- a/docs/CORTEX.md +++ b/docs/CORTEX.md @@ -118,6 +118,18 @@ The initial spawn request (first line of file, written by client). } ``` +### cancel +Inbound request to stop a running talent use. Cortex queues cancellation work off +the receive thread and terminalizes the use with an error carrying `reason_code`. +```json +{ + "event": "cancel", + "ts": 1234567890123, + "use_id": "1234567890123", + "reason_code": "chat_watchdog_cancelled" +} +``` + ### start Emitted when a talent run begins. ```json @@ -171,6 +183,18 @@ Emitted when the model produces reasoning/thinking content (model-dependent, pri } ``` +### progress +Emitted by synchronous generator runs as a liveness heartbeat while provider work +is still in flight. It intentionally carries no `summary`. +```json +{ + "event": "progress", + "ts": 1234567890123, + "use_id": "1234567890123", + "phase": "generate" +} +``` + ### talent_updated Emitted when control is handed off to a different agent (multi-agent scenarios). ```json @@ -189,7 +213,8 @@ Emitted when the talent run completes successfully. "event": "finish", "ts": 1234567890123, "use_id": "1234567890123", - "result": "Final response text to the owner" + "result": "Final response text to the owner", + "generate_progress_count": 3 } ``` @@ -201,7 +226,8 @@ Emitted when an error occurs during execution. "ts": 1234567890123, "use_id": "1234567890123", "error": "Error message", - "trace": "Full stack trace..." + "trace": "Full stack trace...", + "generate_progress_count": 3 } ``` diff --git a/solstone/think/cortex.py b/solstone/think/cortex.py index a762d0805..18d95b4ac 100644 --- a/solstone/think/cortex.py +++ b/solstone/think/cortex.py @@ -44,6 +44,8 @@ from solstone.think.talent import get_output_path from solstone.think.talents import TALENT_EXECUTION_MODULE from solstone.think.utils import get_journal, get_rev, now_ms +_CANCEL_REASON_CODE = "chat_watchdog_cancelled" + class TalentProcess: """Manages a running talent subprocess.""" @@ -657,8 +659,10 @@ class CortexService: self.stop_event = threading.Event() self.shutdown_requested = threading.Event() self.spawn_queue: queue.Queue = queue.Queue() + self.cancel_queue: queue.Queue = queue.Queue() self._pending_spawns: int = 0 self._spawn_worker: threading.Thread | None = None + self._cancel_worker: threading.Thread | None = None self._clock = clock or (lambda: datetime.now(timezone.utc)) self._wait = wait or self.stop_event.wait self._spp_renewal_controller: SppRenewalController | None = None @@ -810,6 +814,12 @@ class CortexService: daemon=True, ) self._spawn_worker.start() + self._cancel_worker = threading.Thread( + target=self._run_cancel_worker, + name="cortex-cancel-worker", + daemon=True, + ) + self._cancel_worker.start() self._start_spp_renewal_controller() self.logger.info("Cortex service started, listening for talent requests") @@ -864,15 +874,30 @@ class CortexService: """Handle incoming Callosum messages (callback).""" if self._spp_renewal_controller is not None: self._spp_renewal_controller.handle_supervisor_message(message) - # Filter for cortex tract and request event - if message.get("tract") != "cortex" or message.get("event") != "request": + # Filter for cortex tract and inbound work events. + event = message.get("event") + if message.get("tract") != "cortex" or event not in {"request", "cancel"}: return # Handle the request try: - self._handle_request(message) + if event == "request": + self._handle_request(message) + else: + self._handle_cancel(message) except Exception as e: - self.logger.exception(f"Error handling request: {e}") + self.logger.exception(f"Error handling {event}: {e}") + + def _handle_cancel(self, message: Dict[str, Any]) -> None: + """Queue cancellation for a running talent use without blocking callosum.""" + use_id = message.get("use_id") + if not isinstance(use_id, str) or not use_id: + self.logger.debug("Ignoring cortex cancel without use_id") + return + reason_code = message.get("reason_code") + if not isinstance(reason_code, str) or not reason_code: + reason_code = _CANCEL_REASON_CODE + self.cancel_queue.put({"use_id": use_id, "reason_code": reason_code}) def _handle_request(self, request: Dict[str, Any]) -> None: """Handle a new talent request from Callosum. @@ -951,6 +976,23 @@ class CortexService: self._pending_spawns -= 1 self.spawn_queue.task_done() + def _run_cancel_worker(self) -> None: + """Drain the cancel queue on a single dedicated thread (FIFO).""" + while not self.stop_event.is_set(): + try: + item = self.cancel_queue.get(timeout=0.5) + except queue.Empty: + continue + try: + self._cancel_talent_use( + str(item.get("use_id") or ""), + str(item.get("reason_code") or _CANCEL_REASON_CODE), + ) + except Exception as e: + self.logger.exception(f"Cancel worker error: {e}") + finally: + self.cancel_queue.task_done() + def _is_idle(self) -> bool: """True when nothing is running, queued, or mid-spawn on the worker.""" with self.lock: @@ -1106,6 +1148,35 @@ class CortexService: self._complete_use_file(use_id, agent.log_path) self._clear_request(use_id) + def _cancel_talent_use(self, use_id: str, reason_code: str) -> None: + """Terminalize and stop a running talent use cancelled by a caller.""" + if not use_id: + return + with self.lock: + agent = self.running_uses.get(use_id) + if agent is None or not self._claim_finalize(use_id): + return + + self.logger.warning("Talent %s cancelled by cortex cancel event", use_id) + error_event = self._create_error_event( + use_id, + "Talent cancelled by chat watchdog", + ) + error_event["reason_code"] = reason_code + self._append_use_event(use_id, agent.log_path, error_event) + + # Broadcast to callosum so wait_for_uses detects immediately + try: + event_copy = error_event.copy() + event_type = event_copy.pop("event", "error") + self.callosum.emit("cortex", event_type, **event_copy) + except Exception: + pass + + agent.stop() + self._complete_use_file(use_id, agent.log_path) + self._clear_request(use_id) + def _monitor_stdout(self, agent: TalentProcess) -> None: """Monitor talent stdout and append events to the JSONL file.""" if not agent.process.stdout: @@ -1484,6 +1555,11 @@ class CortexService: if self._spawn_worker is not None: self._spawn_worker.join(timeout=2.0) + # Let the cancel worker finish its current item and exit (~2s bound; the + # worker's 0.5s get-timeout guarantees it observes stop_event promptly). + if self._cancel_worker is not None: + self._cancel_worker.join(timeout=2.0) + # Terminalize any claims still queued. The worker has exited, so no other # thread pulls these concurrently. while True: diff --git a/tests/test_cortex.py b/tests/test_cortex.py index 4862f63a6..3f7aca765 100644 --- a/tests/test_cortex.py +++ b/tests/test_cortex.py @@ -69,6 +69,33 @@ def _completed_path(journal_path: Path, use_id: str, name: str = "chat") -> Path return journal_path / "talents" / name.replace(":", "--") / f"{use_id}.jsonl" +def _make_live_agent( + cortex_service, journal_path: Path, use_id: str, *, name: str = "chat" +): + from solstone.think.cortex import TalentProcess + + talent_dir = journal_path / "talents" / name.replace(":", "--") + talent_dir.mkdir(exist_ok=True) + active_path = talent_dir / f"{use_id}_active.jsonl" + request = { + "event": "request", + "use_id": use_id, + "ts": 1000, + "name": name, + "day": "20260410", + } + active_path.write_text(json.dumps(request) + "\n", encoding="utf-8") + mock_process = MagicMock() + mock_process.pid = 24680 + mock_process.wait.return_value = 0 + mock_process.stdout = MockPipe([]) + mock_process.stderr = MockPipe([]) + agent = TalentProcess(use_id, mock_process, active_path) + cortex_service.running_uses[use_id] = agent + cortex_service.use_requests[use_id] = request + return agent, active_path + + class _FakeClock: def __init__(self, now: datetime): self.now = now @@ -1599,6 +1626,42 @@ def test_slow_spawn_does_not_block_request_callback(cortex_service, mock_journal cortex_service._spawn_worker.join(timeout=1) +def test_cancel_callback_queues_without_blocking_stop(cortex_service): + block = threading.Event() + started = threading.Event() + cancelled: list[str] = [] + + def slow_cancel(use_id: str, _reason_code: str) -> None: + cancelled.append(use_id) + started.set() + block.wait(timeout=5) + + with patch.object(cortex_service, "_cancel_talent_use", side_effect=slow_cancel): + cortex_service._cancel_worker = threading.Thread( + target=cortex_service._run_cancel_worker, + daemon=True, + ) + cortex_service._cancel_worker.start() + + before = time.monotonic() + cortex_service._handle_callosum_message( + { + "tract": "cortex", + "event": "cancel", + "use_id": "cancel_slow_1", + } + ) + elapsed = time.monotonic() - before + + assert elapsed < 0.5 + assert started.wait(1) + assert cancelled == ["cancel_slow_1"] + + cortex_service.stop_event.set() + block.set() + cortex_service._cancel_worker.join(timeout=1) + + def test_claim_latency_independent_of_queue_depth(cortex_service, mock_journal): """AC2: request handling claims and queues without spawning inline.""" with patch.object(cortex_service, "_spawn_subprocess") as mock_spawn: @@ -2222,6 +2285,200 @@ def test_timeout_finalize_claim_beats_late_stdout_cleanup(cortex_service, mock_j assert [summary["use_id"] for summary in summaries].count(use_id) == 1 +def test_cancel_edge_cases_are_idempotent_without_spawning( + cortex_service, mock_journal +): + cortex_service.callosum = MagicMock() + with patch.object(cortex_service, "_spawn_subprocess") as mock_spawn: + cortex_service._cancel_talent_use("unknown_use", "chat_watchdog_cancelled") + cortex_service._handle_callosum_message({"tract": "cortex", "event": "cancel"}) + cortex_service._handle_callosum_message( + {"tract": "cortex", "event": "cancel", "use_id": {"bad": "shape"}} + ) + + completed_dir = mock_journal / "talents" / "chat" + completed_dir.mkdir() + finalized = completed_dir / "already_done.jsonl" + finalized.write_text('{"event":"finish","use_id":"already_done"}\n') + before_finalized = finalized.read_bytes() + cortex_service._cancel_talent_use("already_done", "chat_watchdog_cancelled") + + agent, active_path = _make_live_agent( + cortex_service, mock_journal, "live_cancel_twice" + ) + with patch.object(agent, "_signal_process_group"): + cortex_service._cancel_talent_use( + "live_cancel_twice", "chat_watchdog_cancelled" + ) + cortex_service._cancel_talent_use( + "live_cancel_twice", "chat_watchdog_cancelled" + ) + + completed_path = _completed_path(mock_journal, "live_cancel_twice") + rows = _read_jsonl(completed_path) + terminal_errors = [row for row in rows if row.get("event") == "error"] + assert not _active_path(mock_journal, "unknown_use").exists() + assert finalized.read_bytes() == before_finalized + assert not active_path.exists() + assert len(terminal_errors) == 1 + assert terminal_errors[0]["reason_code"] == "chat_watchdog_cancelled" + mock_spawn.assert_not_called() + + +def test_cancel_reason_code_does_not_change_timeout_reason_shape( + cortex_service, mock_journal +): + cortex_service.callosum = MagicMock() + cancel_agent, _cancel_active_path = _make_live_agent( + cortex_service, mock_journal, "reason_cancel" + ) + with patch.object(cancel_agent, "_signal_process_group"): + cortex_service._cancel_talent_use("reason_cancel", "chat_watchdog_cancelled") + + timeout_agent, _timeout_active_path = _make_live_agent( + cortex_service, mock_journal, "reason_timeout" + ) + with patch.object(timeout_agent, "_signal_process_group"): + cortex_service._timeout_talent("reason_timeout", timeout_agent, 1) + + cancel_error = [ + row + for row in _read_jsonl(_completed_path(mock_journal, "reason_cancel")) + if row.get("event") == "error" + ][0] + timeout_error = [ + row + for row in _read_jsonl(_completed_path(mock_journal, "reason_timeout")) + if row.get("event") == "error" + ][0] + assert cancel_error["reason_code"] == "chat_watchdog_cancelled" + assert "reason_code" not in timeout_error + + +def test_cancel_racing_real_finish_completes_once(cortex_service, mock_journal): + cortex_service.callosum = MagicMock() + use_id = "cancel_race_finish" + agent, active_path = _make_live_agent(cortex_service, mock_journal, use_id) + agent.process.stdout = MockPipe( + [ + json.dumps( + { + "event": "finish", + "use_id": use_id, + "ts": 1001, + "result": "already emitted", + } + ) + + "\n" + ] + ) + finish_appended = threading.Event() + release_finish = threading.Event() + original_append = cortex_service._append_use_event + + def append_spy(use_id_arg, active_path_arg, event): + appended = original_append(use_id_arg, active_path_arg, event) + if event.get("event") == "finish": + finish_appended.set() + release_finish.wait(timeout=1) + return appended + + monitor = threading.Thread(target=cortex_service._monitor_stdout, args=(agent,)) + with patch.object(cortex_service, "_append_use_event", side_effect=append_spy): + monitor.start() + assert finish_appended.wait(1) + with patch.object(agent, "_signal_process_group"): + cortex_service._cancel_talent_use(use_id, "chat_watchdog_cancelled") + release_finish.set() + monitor.join(timeout=1) + + assert not monitor.is_alive() + + completed_path = _completed_path(mock_journal, use_id) + rows = _read_jsonl(completed_path) + day_index = mock_journal / "talents" / "20260410.jsonl" + assert completed_path.exists() + assert not active_path.exists() + assert sum(1 for row in rows if row.get("event") == "finish") == 1 + terminal_errors = [row for row in rows if row.get("event") == "error"] + assert len(terminal_errors) == 1 + assert terminal_errors[0]["reason_code"] == "chat_watchdog_cancelled" + assert len(day_index.read_text(encoding="utf-8").splitlines()) == 1 + + +def test_cancel_claim_first_then_late_finish_appends_before_complete( + cortex_service, mock_journal +): + from solstone.think.cortex_client import get_use_end_state + + cortex_service.callosum = MagicMock() + use_id = "cancel_claim_first_late_finish" + agent, active_path = _make_live_agent(cortex_service, mock_journal, use_id) + agent.process.stdout = MockPipe( + [ + json.dumps( + { + "event": "finish", + "use_id": use_id, + "ts": 1002, + "result": "late emitted", + } + ) + + "\n" + ] + ) + cancel_error_appended = threading.Event() + allow_cancel_complete = threading.Event() + cancel_done = threading.Event() + original_append = cortex_service._append_use_event + + def append_spy(use_id_arg, active_path_arg, event): + appended = original_append(use_id_arg, active_path_arg, event) + if ( + event.get("event") == "error" + and event.get("reason_code") == "chat_watchdog_cancelled" + ): + cancel_error_appended.set() + allow_cancel_complete.wait(timeout=1) + return appended + + def cancel_use(): + try: + cortex_service._cancel_talent_use(use_id, "chat_watchdog_cancelled") + finally: + cancel_done.set() + + cancel_thread = threading.Thread(target=cancel_use) + with patch.object(cortex_service, "_append_use_event", side_effect=append_spy): + with patch.object(agent, "_signal_process_group"): + cancel_thread.start() + try: + assert cancel_error_appended.wait(1) + assert use_id not in cortex_service.running_uses + assert active_path.exists() + cortex_service._monitor_stdout(agent) + finally: + allow_cancel_complete.set() + cancel_thread.join(timeout=1) + + assert cancel_done.is_set() + assert not cancel_thread.is_alive() + + completed_path = _completed_path(mock_journal, use_id) + rows = _read_jsonl(completed_path) + day_index = mock_journal / "talents" / "20260410.jsonl" + day_rows = _read_jsonl(day_index) + assert completed_path.exists() + assert not active_path.exists() + assert [row.get("event") for row in rows] == ["request", "error", "finish"] + assert rows[1]["reason_code"] == "chat_watchdog_cancelled" + assert rows[2]["result"] == "late emitted" + assert cortex_service._has_finish_event(completed_path) is True + assert get_use_end_state(use_id) == "finish" + assert len(day_rows) == 1 + assert day_rows[0]["status"] == "completed" + + def test_monitor_stderr_after_completion_does_not_resurrect_active_log( cortex_service, mock_journal ):