diff --git a/packages/core/examples/agent_e2e.py b/packages/core/examples/agent_e2e.py index 5262566..5db653d 100644 --- a/packages/core/examples/agent_e2e.py +++ b/packages/core/examples/agent_e2e.py @@ -31,6 +31,8 @@ environment from the token. """ +from __future__ import annotations + import asyncio import os @@ -186,6 +188,8 @@ async def simulate_retrieval(query: str) -> list[str]: async def main() -> None: agent_id = os.environ.get("PREFACTOR_AGENT_ID") + if agent_id is not None: + agent_id = agent_id.strip() or None config = PrefactorCoreConfig( http_config=HttpClientConfig( diff --git a/packages/core/pyproject.toml b/packages/core/pyproject.toml index 328c715..9df986d 100644 --- a/packages/core/pyproject.toml +++ b/packages/core/pyproject.toml @@ -9,7 +9,7 @@ authors = [ ] requires-python = ">=3.11.0, <4.0.0" dependencies = [ - "prefactor-http>=0.1.3", + "prefactor-http>=0.1.4", "pydantic>=2.0.0", ] diff --git a/packages/core/src/prefactor_core/__init__.py b/packages/core/src/prefactor_core/__init__.py index b9c69af..bccc494 100644 --- a/packages/core/src/prefactor_core/__init__.py +++ b/packages/core/src/prefactor_core/__init__.py @@ -16,6 +16,7 @@ OperationError, PrefactorCoreError, PrefactorTelemetryFailureError, + PrefactorTerminatedError, SpanNotFoundError, ) from .managers.agent_instance import AgentInstanceHandle @@ -43,6 +44,7 @@ "InstanceNotFoundError", "SpanNotFoundError", "PrefactorTelemetryFailureError", + "PrefactorTerminatedError", # Models "AgentInstance", "Span", diff --git a/packages/core/src/prefactor_core/_version.py b/packages/core/src/prefactor_core/_version.py index 170234f..34964bc 100644 --- a/packages/core/src/prefactor_core/_version.py +++ b/packages/core/src/prefactor_core/_version.py @@ -3,5 +3,5 @@ from __future__ import annotations PACKAGE_NAME = "prefactor-core" -__version__ = "0.2.5" +__version__ = "0.2.6" PACKAGE_VERSION = __version__ diff --git a/packages/core/src/prefactor_core/client.py b/packages/core/src/prefactor_core/client.py index 15393d8..414f077 100644 --- a/packages/core/src/prefactor_core/client.py +++ b/packages/core/src/prefactor_core/client.py @@ -7,13 +7,18 @@ from __future__ import annotations +import asyncio import logging import time from contextlib import asynccontextmanager from typing import TYPE_CHECKING, Any from prefactor_http.client import PrefactorHttpClient -from prefactor_http.exceptions import is_permanent_http_error, is_transient_http_error +from prefactor_http.exceptions import ( + PrefactorApiError, + is_permanent_http_error, + is_transient_http_error, +) from ._version import PACKAGE_NAME as CORE_PACKAGE_NAME from ._version import PACKAGE_VERSION as CORE_PACKAGE_VERSION @@ -26,6 +31,7 @@ ) from .managers.agent_instance import AgentInstanceManager from .managers.span import SpanManager +from .monitoring.termination_monitor import TerminationMonitor from .operations import Operation, OperationType from .queue.base import Queue from .queue.executor import TaskExecutor @@ -86,6 +92,9 @@ def __init__( self._initialized = False self._telemetry_failure: PrefactorTelemetryFailureError | None = None self._telemetry_failure_observed = False + self._termination_monitor: TerminationMonitor | None = None + self._sync_task: asyncio.Task | None = None + self._current_instance_id: str | None = None def _build_http_sdk_header(self) -> str: """Build the effective SDK header for HTTP requests.""" @@ -155,6 +164,13 @@ async def initialize(self) -> None: self._initialized = True + self._termination_monitor = TerminationMonitor( + fetch_instance=self._fetch_instance_for_poll, + ) + self._sync_task = asyncio.create_task( + self._run_sync_loop(), name="prefactor-termination-sync" + ) + async def close(self) -> None: """Close the client and cleanup resources. @@ -164,6 +180,24 @@ async def close(self) -> None: if not self._initialized: return + if self._sync_task is not None: + if not self._sync_task.done(): + self._sync_task.cancel() + try: + await self._sync_task + except asyncio.CancelledError: + pass + except Exception: + logger.exception( + "Termination sync loop exited with error during close()" + ) + finally: + self._sync_task = None + + if self._termination_monitor is not None: + self._termination_monitor.destroy() + self._termination_monitor = None + # Stop executor if self._executor: await self._executor.stop() @@ -272,12 +306,23 @@ async def _process_operation(self, operation: Operation) -> None: ) elif operation.type == OperationType.FINISH_AGENT_INSTANCE: - await self._http.agent_instances.finish( - agent_instance_id=operation.payload["instance_id"], - status=operation.payload.get("status", "complete"), - timestamp=operation.timestamp, - idempotency_key=operation.payload.get("idempotency_key"), - ) + try: + await self._http.agent_instances.finish( + agent_instance_id=operation.payload["instance_id"], + status=operation.payload.get("status", "complete"), + timestamp=operation.timestamp, + idempotency_key=operation.payload.get("idempotency_key"), + ) + except PrefactorApiError as finish_err: + if finish_err.status_code == 409: + logger.debug( + "[prefactor:http] Agent instance %s already in" + " terminal state; skipping finish.", + operation.payload["instance_id"], + ) + return + raise + elif operation.type == OperationType.CREATE_SPAN: await self._http.agent_spans.create( agent_instance_id=operation.payload["instance_id"], @@ -286,6 +331,7 @@ async def _process_operation(self, operation: Operation) -> None: id=operation.payload.get("span_id"), parent_span_id=operation.payload.get("parent_span_id"), payload=operation.payload.get("payload"), + control_signal_callback=self._on_control_signal, ) elif operation.type == OperationType.FINISH_SPAN: @@ -295,6 +341,7 @@ async def _process_operation(self, operation: Operation) -> None: result_payload=operation.payload.get("result_payload"), timestamp=operation.timestamp, idempotency_key=operation.payload.get("idempotency_key"), + control_signal_callback=self._on_control_signal, ) except Exception as e: @@ -307,6 +354,32 @@ async def _process_operation(self, operation: Operation) -> None: ) raise + async def _fetch_instance_for_poll(self, instance_id: str): + if self._http is None: + return None + return await self._http.agent_instances.get(instance_id) + + async def _run_sync_loop(self) -> None: + while True: + await asyncio.sleep(1) + if self._termination_monitor is None: + continue + try: + self._termination_monitor.sync(self._current_instance_id) + except Exception: + logger.exception("Termination sync iteration failed") + + def _on_control_signal(self, reason: str | None) -> None: + if self._termination_monitor is not None: + self._termination_monitor.detect_termination(reason) + + def _set_current_instance(self, instance_id: str | None) -> None: + self._current_instance_id = instance_id + + def _clear_current_instance(self, instance_id: str) -> None: + if self._current_instance_id == instance_id: + self._current_instance_id = None + @property def instance_manager(self) -> AgentInstanceManager | None: """Public accessor for the agent instance manager.""" @@ -381,6 +454,8 @@ async def create_agent_instance( environment_id=environment_id, ) + self._set_current_instance(instance_id) + return AgentInstanceHandle( instance_id=instance_id, client=self, diff --git a/packages/core/src/prefactor_core/exceptions.py b/packages/core/src/prefactor_core/exceptions.py index b225132..9c88558 100644 --- a/packages/core/src/prefactor_core/exceptions.py +++ b/packages/core/src/prefactor_core/exceptions.py @@ -58,6 +58,23 @@ def __init__( self.dropped_operations = dropped_operations +class PrefactorTerminatedError(PrefactorCoreError): + """Raised when the agent instance has been terminated by p2. + + Args: + reason: Optional reason reported by p2 for the termination. + """ + + def __init__(self, reason: str | None = None) -> None: + msg = ( + f"Agent instance terminated by p2: {reason}" + if reason + else "Agent instance terminated by p2" + ) + super().__init__(msg) + self.reason = reason + + __all__ = [ "PrefactorCoreError", "ClientNotInitializedError", @@ -66,4 +83,5 @@ def __init__( "InstanceNotFoundError", "SpanNotFoundError", "PrefactorTelemetryFailureError", + "PrefactorTerminatedError", ] diff --git a/packages/core/src/prefactor_core/managers/agent_instance.py b/packages/core/src/prefactor_core/managers/agent_instance.py index 2453ef0..bb0a964 100644 --- a/packages/core/src/prefactor_core/managers/agent_instance.py +++ b/packages/core/src/prefactor_core/managers/agent_instance.py @@ -218,12 +218,18 @@ async def start(self) -> None: async def finish(self, status: "FinishStatus" = "complete") -> None: """Mark the instance as finished. - This queues a finish operation for the instance. + Resets the termination monitor (fence + new event) before enqueueing + the HTTP finish so stale span responses from the dying run cannot + trigger termination on the next run. Args: status: Terminal status for the instance — one of ``"complete"``, ``"failed"``, or ``"cancelled"``. Defaults to ``"complete"``. """ + monitor = getattr(self._client, "_termination_monitor", None) + if monitor is not None: + monitor.reset() + self._client._clear_current_instance(self._instance_id) manager = self._client.instance_manager assert manager is not None await manager.finish_with_idempotency_key( diff --git a/packages/core/src/prefactor_core/monitoring/__init__.py b/packages/core/src/prefactor_core/monitoring/__init__.py new file mode 100644 index 0000000..f5990b5 --- /dev/null +++ b/packages/core/src/prefactor_core/monitoring/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from prefactor_core.monitoring.termination_monitor import TerminationMonitor + +__all__ = ["TerminationMonitor"] diff --git a/packages/core/src/prefactor_core/monitoring/termination_monitor.py b/packages/core/src/prefactor_core/monitoring/termination_monitor.py new file mode 100644 index 0000000..b63c0da --- /dev/null +++ b/packages/core/src/prefactor_core/monitoring/termination_monitor.py @@ -0,0 +1,151 @@ +"""TerminationMonitor — detects agent instance termination via two paths: +1. Fast: control signal in span API responses (pushed via callback) +2. Slow: polling the instance status endpoint every `poll_interval` seconds +""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Callable + +logger = logging.getLogger(__name__) + + +class TerminationMonitor: + """Monitors an agent instance for termination. + + Thread-safety note: `detect_termination` and `get_termination_event().is_set()` + are safe to call from sync worker threads — they only read/write a bool under + the GIL (asyncio.Event._value is a plain bool). + """ + + def __init__( + self, + fetch_instance: Callable, + poll_interval: float = 30.0, + ) -> None: + self._fetch_instance = fetch_instance + self._poll_interval = poll_interval + + self._event = asyncio.Event() + self._termination_reason: str | None = None + self._callbacks: list[Callable[[], None]] = [] + + self._current_instance_id: str | None = None + self._poll_task: asyncio.Task | None = None + self._destroyed = False + + # Fence blocks stale span callbacks from a previous run after reset() + self._fenced = False + # Generation increments on each reset so stale polls self-discard + self._generation = 0 + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + @property + def termination_reason(self) -> str | None: + return self._termination_reason + + def get_termination_event(self) -> asyncio.Event: + return self._event + + def detect_termination(self, reason: str | None) -> None: + """Signal that the instance has been terminated. + + No-op if already terminated, destroyed, or fenced (post-reset stale call). + """ + if self._destroyed or self._event.is_set() or self._fenced: + return + self._termination_reason = reason + self._event.set() + self._stop_poll() + for cb in list(self._callbacks): + try: + cb() + except Exception: + logger.exception("Termination callback failed: %r", cb) + + def sync(self, instance_id: str | None) -> None: + """Update the tracked instance ID and start/stop the fallback poll. + + Idempotent: calling with the same non-None ID when a poll is already + running does not restart it (preserving the sleep interval). + """ + if instance_id is not None: + self._fenced = False + if instance_id == self._current_instance_id: + # Same ID — no change needed; poll (if any) keeps running + return + self._current_instance_id = instance_id + self._stop_poll() + if instance_id is not None and not self._event.is_set() and not self._destroyed: + self._start_poll(instance_id, self._generation) + + def reset(self) -> None: + """Prepare monitor for the next agent run. + + - Creates a fresh (unset) event + - Clears the termination reason + - Cancels any in-flight poll + - Sets fence so stale callbacks from the dying run are ignored + - Increments generation so stale polls self-discard + """ + self._stop_poll() + self._event = asyncio.Event() + self._termination_reason = None + self._current_instance_id = None + self._generation += 1 + self._fenced = True + + def subscribe(self, callback: Callable[[], None]) -> Callable[[], None]: + """Register a callback invoked on termination. Returns an unsubscribe fn.""" + self._callbacks.append(callback) + + def unsubscribe() -> None: + try: + self._callbacks.remove(callback) + except ValueError: + pass + + return unsubscribe + + def destroy(self) -> None: + """Permanently shut down the monitor (no further events will fire).""" + self._destroyed = True + self._stop_poll() + self._callbacks.clear() + + # ------------------------------------------------------------------ + # Internal + # ------------------------------------------------------------------ + + def _start_poll(self, instance_id: str, generation: int) -> None: + self._poll_task = asyncio.get_event_loop().create_task( + self._poll(instance_id, generation) + ) + + def _stop_poll(self) -> None: + if self._poll_task is not None and not self._poll_task.done(): + self._poll_task.cancel() + self._poll_task = None + + async def _poll(self, instance_id: str, generation: int) -> None: + while True: + await asyncio.sleep(self._poll_interval) + # Guard: generation changed means reset() was called + if self._generation != generation: + return + try: + instance = await self._fetch_instance(instance_id) + except Exception: + logger.debug("Termination poll error (transient)", exc_info=True) + continue + # Guard again after await + if self._generation != generation: + return + if instance.status == "terminated": + self.detect_termination(instance.termination_reason) + return diff --git a/packages/core/tests/monitoring/__init__.py b/packages/core/tests/monitoring/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/packages/core/tests/monitoring/test_termination_monitor.py b/packages/core/tests/monitoring/test_termination_monitor.py new file mode 100644 index 0000000..185da9d --- /dev/null +++ b/packages/core/tests/monitoring/test_termination_monitor.py @@ -0,0 +1,259 @@ +"""Tests for TerminationMonitor — 19 tests covering primary path, fallback poll, +reset(), and callback lifecycle.""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +from prefactor_core.monitoring.termination_monitor import TerminationMonitor + + +def _make_monitor(fetch_instance=None) -> TerminationMonitor: + if fetch_instance is None: + fetch_instance = AsyncMock( + return_value=MagicMock(status="active", termination_reason=None) + ) + return TerminationMonitor(fetch_instance=fetch_instance) + + +def _terminated_instance(reason: str | None = "test reason"): + inst = MagicMock() + inst.status = "terminated" + inst.termination_reason = reason + return inst + + +# --------------------------------------------------------------------------- +# Primary path (5 tests) +# --------------------------------------------------------------------------- + + +class TestPrimaryPath: + async def test_detect_termination_sets_event(self): + monitor = _make_monitor() + assert not monitor.get_termination_event().is_set() + monitor.detect_termination("reason") + assert monitor.get_termination_event().is_set() + + async def test_reason_propagates(self): + monitor = _make_monitor() + monitor.detect_termination("my reason") + assert monitor.termination_reason == "my reason" + + async def test_null_reason_accepted(self): + monitor = _make_monitor() + monitor.detect_termination(None) + assert monitor.get_termination_event().is_set() + assert monitor.termination_reason is None + + async def test_second_detect_termination_is_idempotent(self): + callback = MagicMock() + monitor = _make_monitor() + monitor.subscribe(callback) + monitor.detect_termination("first") + monitor.detect_termination("second") + callback.assert_called_once() + assert monitor.termination_reason == "first" + + async def test_detect_termination_noop_after_destroy(self): + monitor = _make_monitor() + monitor.destroy() + monitor.detect_termination("reason") + assert not monitor.get_termination_event().is_set() + + +# --------------------------------------------------------------------------- +# Fallback poll (5 tests) +# --------------------------------------------------------------------------- + + +class TestFallbackPoll: + async def test_poll_starts_when_instance_id_arrives(self): + monitor = _make_monitor() + monitor.sync("inst-1") + await asyncio.sleep(0) # yield to let task start + assert monitor._poll_task is not None + assert not monitor._poll_task.done() + monitor.destroy() + + async def test_poll_stops_when_sync_called_with_none(self): + monitor = _make_monitor() + monitor.sync("inst-1") + await asyncio.sleep(0) + poll_task = monitor._poll_task + assert poll_task is not None + monitor.sync(None) + await asyncio.sleep(0) + assert poll_task.cancelled() or poll_task.done() + monitor.destroy() + + async def test_no_poll_without_instance_id(self): + monitor = _make_monitor() + monitor.sync(None) + await asyncio.sleep(0) + assert monitor._poll_task is None + monitor.destroy() + + async def test_poll_stops_after_termination_detected(self): + monitor = _make_monitor() + monitor.sync("inst-1") + await asyncio.sleep(0) + poll_task = monitor._poll_task + assert poll_task is not None + monitor.detect_termination("reason") + await asyncio.sleep(0) + assert poll_task.cancelled() or poll_task.done() + monitor.destroy() + + async def test_poll_survives_transient_http_errors(self): + fetch = AsyncMock(side_effect=Exception("network error")) + monitor = TerminationMonitor(fetch_instance=fetch, poll_interval=0.05) + monitor.sync("inst-1") + await asyncio.sleep(0.2) # let poll fire a couple times + # monitor should not be terminated — error was swallowed + assert not monitor.get_termination_event().is_set() + monitor.destroy() + + +# --------------------------------------------------------------------------- +# reset() (7 tests) +# --------------------------------------------------------------------------- + + +class TestReset: + async def test_reset_creates_fresh_event(self): + monitor = _make_monitor() + monitor.detect_termination("reason") + old_event = monitor.get_termination_event() + assert old_event.is_set() + monitor.reset() + new_event = monitor.get_termination_event() + assert not new_event.is_set() + assert new_event is not old_event + + async def test_reset_allows_new_termination_after_sync_with_new_id(self): + monitor = _make_monitor() + monitor.detect_termination("run 1") + monitor.reset() + # fenced — detect_termination should be blocked + monitor.detect_termination("stale") + assert not monitor.get_termination_event().is_set() + # sync with new id lifts fence + monitor.sync("inst-2") + monitor.detect_termination("run 2") + assert monitor.get_termination_event().is_set() + assert monitor.termination_reason == "run 2" + monitor.destroy() + + async def test_get_termination_event_returns_new_event_after_reset(self): + monitor = _make_monitor() + getter = monitor.get_termination_event + monitor.detect_termination("reason") + monitor.reset() + # getter returns new (unset) event + assert not getter().is_set() + + async def test_reset_cancels_poll_task(self): + monitor = _make_monitor() + monitor.sync("inst-1") + await asyncio.sleep(0) + poll_task = monitor._poll_task + assert poll_task is not None + monitor.reset() + await asyncio.sleep(0) + assert poll_task.cancelled() or poll_task.done() + + async def test_reset_preserves_callbacks(self): + callback = MagicMock() + monitor = _make_monitor() + monitor.subscribe(callback) + monitor.reset() + # sync to lift fence, then detect + monitor.sync("inst-2") + monitor.detect_termination("after reset") + callback.assert_called_once() + + async def test_fence_blocks_detect_termination_until_sync_with_new_id(self): + monitor = _make_monitor() + monitor.detect_termination("run 1") + monitor.reset() # fenced = True + + # stale span response fires during queue drain + monitor.detect_termination("stale span response") + assert not monitor.get_termination_event().is_set() + + # sync with None doesn't lift fence + monitor.sync(None) + monitor.detect_termination("still stale") + assert not monitor.get_termination_event().is_set() + + # sync with new instance id lifts fence + monitor.sync("instance-2") + monitor.detect_termination("run 2") + assert monitor.get_termination_event().is_set() + assert monitor.termination_reason == "run 2" + monitor.destroy() + + async def test_stale_poll_response_discarded_after_reset(self): + """Poll fires for old instance after reset — generation check discards it.""" + slow_fetch = AsyncMock(return_value=_terminated_instance("old run")) + monitor = TerminationMonitor(fetch_instance=slow_fetch, poll_interval=0.05) + monitor.sync("inst-1") + await asyncio.sleep(0) # poll task started + + # reset before poll completes + monitor.reset() + await asyncio.sleep(0.2) # let poll fire with old generation + + # monitor should NOT be terminated + assert not monitor.get_termination_event().is_set() + monitor.destroy() + + +# --------------------------------------------------------------------------- +# Callback lifecycle (2 tests) +# --------------------------------------------------------------------------- + + +class TestCallbackLifecycle: + async def test_unsubscribe_removes_callback(self): + callback = MagicMock() + monitor = _make_monitor() + unsubscribe = monitor.subscribe(callback) + unsubscribe() + monitor.detect_termination("reason") + callback.assert_not_called() + + async def test_callbacks_fire_in_registration_order(self): + order = [] + monitor = _make_monitor() + monitor.subscribe(lambda: order.append(1)) + monitor.subscribe(lambda: order.append(2)) + monitor.subscribe(lambda: order.append(3)) + monitor.detect_termination("reason") + assert order == [1, 2, 3] + + async def test_bad_callback_does_not_block_later_callbacks(self): + order = [] + monitor = _make_monitor() + + def failing_callback(): + order.append("bad") + raise RuntimeError("callback failed") + + monitor.subscribe(failing_callback) + monitor.subscribe(lambda: order.append("good")) + + monitor.detect_termination("reason") + + assert order == ["bad", "good"] + assert monitor.get_termination_event().is_set() + + async def test_destroy_clears_callbacks(self): + monitor = _make_monitor() + monitor.subscribe(MagicMock()) + + monitor.destroy() + + assert monitor._callbacks == [] diff --git a/packages/core/tests/test_agent_instance_finish_status.py b/packages/core/tests/test_agent_instance_finish_status.py index 3ef2bbe..9f77ed5 100644 --- a/packages/core/tests/test_agent_instance_finish_status.py +++ b/packages/core/tests/test_agent_instance_finish_status.py @@ -102,3 +102,70 @@ async def test_finish_forwards_explicit_status(status): assert len(stub_http.agent_instances.finish_calls) == 1 assert stub_http.agent_instances.finish_calls[0]["status"] == status + + +class TestFinishAgentInstance409Handling: + async def test_409_on_finish_treated_as_success(self): + """FINISH_AGENT_INSTANCE with 409 response should not raise.""" + from datetime import datetime, timezone + from unittest.mock import AsyncMock, MagicMock + + from prefactor_core.operations import Operation, OperationType + from prefactor_http.exceptions import PrefactorApiError + + mock_http = MagicMock() + mock_http.agent_instances = MagicMock() + mock_http.agent_instances.finish = AsyncMock( + side_effect=PrefactorApiError("already terminated", "conflict", 409) + ) + mock_http.agent_spans = MagicMock() + mock_http.agent_spans.create = AsyncMock() + mock_http.agent_spans.finish = AsyncMock() + + config = PrefactorCoreConfig( + http_config=HttpClientConfig(api_url="http://fake", api_token="tok") + ) + client = PrefactorCoreClient(config) + client._http = mock_http + client._initialized = True + + op = Operation( + type=OperationType.FINISH_AGENT_INSTANCE, + payload={ + "instance_id": "inst-1", + "idempotency_key": "key-1", + "status": "complete", + }, + timestamp=datetime.now(timezone.utc), + ) + # Should not raise + await client._process_operation(op) + + +class TestAgentInstanceHandleFinishResetsMonitor: + async def test_finish_resets_termination_monitor(self): + """handle.finish() should reset the termination monitor before enqueueing.""" + from unittest.mock import AsyncMock, MagicMock + + from prefactor_core.managers.agent_instance import AgentInstanceHandle + from prefactor_core.monitoring.termination_monitor import TerminationMonitor + + fetch = AsyncMock( + return_value=MagicMock(status="active", termination_reason=None) + ) + monitor = TerminationMonitor(fetch_instance=fetch) + monitor.detect_termination("run 1") + assert monitor.get_termination_event().is_set() + + mock_client = MagicMock() + mock_client._termination_monitor = monitor + mock_client.instance_manager = MagicMock() + mock_client.instance_manager.finish_with_idempotency_key = AsyncMock() + + handle = AgentInstanceHandle(instance_id="inst-1", client=mock_client) + await handle.finish() + + # Monitor should be reset (new unset event) + assert not monitor.get_termination_event().is_set() + # finish_with_idempotency_key should still be called + mock_client.instance_manager.finish_with_idempotency_key.assert_called_once() diff --git a/packages/core/tests/test_failure_handling.py b/packages/core/tests/test_failure_handling.py index 646cf47..ec074ec 100644 --- a/packages/core/tests/test_failure_handling.py +++ b/packages/core/tests/test_failure_handling.py @@ -3,8 +3,9 @@ from __future__ import annotations import asyncio +import logging from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import aiohttp import pytest @@ -103,6 +104,60 @@ async def _wait_until( raise AssertionError("Timed out waiting for expected condition") +@pytest.mark.asyncio +async def test_termination_sync_loop_survives_iteration_failures(caplog): + """A failed sync iteration should be logged without stopping future syncs.""" + client = PrefactorCoreClient(_make_client_config()) + monitor = Mock() + calls = 0 + + def sync(_instance_id): + nonlocal calls + calls += 1 + if calls == 1: + raise RuntimeError("sync failed") + + monitor.sync.side_effect = sync + client._termination_monitor = monitor + original_sleep = asyncio.sleep + + async def fast_sleep(_delay): + await original_sleep(0) + + with ( + caplog.at_level(logging.ERROR, logger="prefactor_core.client"), + patch("prefactor_core.client.asyncio.sleep", side_effect=fast_sleep), + ): + task = asyncio.create_task(client._run_sync_loop()) + try: + await _wait_until(lambda: monitor.sync.call_count >= 2) + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert "Termination sync iteration failed" in caplog.text + + +@pytest.mark.asyncio +async def test_close_observes_completed_sync_task_exceptions(caplog): + """close() should await already-failed sync tasks and clear the task handle.""" + client = PrefactorCoreClient(_make_client_config()) + client._initialized = True + + async def fail_sync_loop(): + raise RuntimeError("sync task failed") + + client._sync_task = asyncio.create_task(fail_sync_loop()) + await asyncio.sleep(0) + + with caplog.at_level(logging.ERROR, logger="prefactor_core.client"): + await client.close() + + assert client._sync_task is None + assert "Termination sync loop exited with error during close()" in caplog.text + + @pytest.mark.asyncio async def test_permanent_worker_failure_latches_and_rejects_future_operations(): """Permanent failures should latch and reject later queued operations.""" diff --git a/packages/http/README.md b/packages/http/README.md index b575224..074395a 100644 --- a/packages/http/README.md +++ b/packages/http/README.md @@ -267,7 +267,7 @@ config = HttpClientConfig( ```python from prefactor_http import AgentStatus, FinishStatus -# AgentStatus = Literal["pending", "active", "complete", "failed", "cancelled"] +# AgentStatus = Literal["pending", "active", "complete", "failed", "cancelled", "terminated"] # FinishStatus = Literal["complete", "failed", "cancelled"] ``` diff --git a/packages/http/src/prefactor_http/_version.py b/packages/http/src/prefactor_http/_version.py index f0baca2..23c6bb7 100644 --- a/packages/http/src/prefactor_http/_version.py +++ b/packages/http/src/prefactor_http/_version.py @@ -3,5 +3,5 @@ from __future__ import annotations PACKAGE_NAME = "prefactor-http" -__version__ = "0.1.3" +__version__ = "0.1.4" PACKAGE_VERSION = __version__ diff --git a/packages/http/src/prefactor_http/endpoints/agent_instance.py b/packages/http/src/prefactor_http/endpoints/agent_instance.py index 6b057e6..158889b 100644 --- a/packages/http/src/prefactor_http/endpoints/agent_instance.py +++ b/packages/http/src/prefactor_http/endpoints/agent_instance.py @@ -196,3 +196,24 @@ async def finish( ) return self._parse_response(response, "agent_instances.finish") + + async def get(self, agent_instance_id: str) -> AgentInstance: + """Fetch an agent instance by ID. + + GET /api/v1/agent_instance/{agent_instance_id} + + Args: + agent_instance_id: The instance ID to fetch. + + Returns: + The agent instance. + + Raises: + PrefactorNotFoundError: If instance not found. + PrefactorApiError: On other errors. + """ + response = await self._client.request( + "GET", + f"/api/v1/agent_instance/{agent_instance_id}", + ) + return self._parse_response(response, "agent_instances.get") diff --git a/packages/http/src/prefactor_http/endpoints/agent_span.py b/packages/http/src/prefactor_http/endpoints/agent_span.py index 8b816ee..a446b4c 100644 --- a/packages/http/src/prefactor_http/endpoints/agent_span.py +++ b/packages/http/src/prefactor_http/endpoints/agent_span.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Callable from pydantic import ValidationError @@ -20,6 +20,15 @@ from prefactor_http.client import PrefactorHttpClient +def _check_control_signal( + response: dict, + callback: Callable[[str | None], None], +) -> None: + control = response.get("control") + if isinstance(control, dict) and control.get("terminate"): + callback(control.get("reason")) + + def _validate_idempotency_key(key: str) -> None: """Validate that an idempotency key is at most 64 characters. @@ -74,6 +83,7 @@ async def create( started_at: datetime | None = None, finished_at: datetime | None = None, idempotency_key: str | None = None, + control_signal_callback: Callable[[str | None], None] | None = None, ) -> AgentSpan: """Create a new agent span. @@ -123,6 +133,9 @@ async def create( json_data=body, ) + if control_signal_callback is not None: + _check_control_signal(response, control_signal_callback) + return self._parse_response(response, "agent_spans.create") async def finish( @@ -132,6 +145,7 @@ async def finish( result_payload: dict | None = None, timestamp: datetime | None = None, idempotency_key: str | None = None, + control_signal_callback: Callable[[str | None], None] | None = None, ) -> AgentSpan: """Finish an agent span. @@ -167,4 +181,7 @@ async def finish( json_data=finish_request.model_dump(exclude_none=True), ) + if control_signal_callback is not None: + _check_control_signal(response, control_signal_callback) + return self._parse_response(response, "agent_spans.finish") diff --git a/packages/http/src/prefactor_http/models/agent_instance.py b/packages/http/src/prefactor_http/models/agent_instance.py index f08aee8..f376fba 100644 --- a/packages/http/src/prefactor_http/models/agent_instance.py +++ b/packages/http/src/prefactor_http/models/agent_instance.py @@ -233,11 +233,13 @@ class AgentInstance(BaseModel): agent_id: Agent ID agent_version_id: Agent version ID environment_id: Environment ID + agent_deployment_id: Agent deployment ID status: Instance status inserted_at: When the instance was created updated_at: When the instance was last updated started_at: When the instance started (null if not started) finished_at: When the instance finished (null if not finished) + termination_reason: Reason for termination (null if not terminated) span_counts: Span counts for this instance """ @@ -253,4 +255,5 @@ class AgentInstance(BaseModel): updated_at: datetime started_at: datetime | None = None finished_at: datetime | None = None + termination_reason: str | None = None span_counts: AgentInstanceSpanCounts | None = None diff --git a/packages/http/src/prefactor_http/models/types.py b/packages/http/src/prefactor_http/models/types.py index 8a1fb61..6ea9b1c 100644 --- a/packages/http/src/prefactor_http/models/types.py +++ b/packages/http/src/prefactor_http/models/types.py @@ -1,6 +1,10 @@ """Shared type definitions for Prefactor API models.""" +from __future__ import annotations + from typing import Literal -AgentStatus = Literal["pending", "active", "complete", "failed", "cancelled"] +AgentStatus = Literal[ + "pending", "active", "complete", "failed", "cancelled", "terminated" +] FinishStatus = Literal["complete", "failed", "cancelled"] diff --git a/packages/http/tests/test_endpoints.py b/packages/http/tests/test_endpoints.py index f1e47d6..0f9a05f 100644 --- a/packages/http/tests/test_endpoints.py +++ b/packages/http/tests/test_endpoints.py @@ -443,3 +443,145 @@ async def test_exactly_64_chars_is_accepted(self, config): idempotency_key="a" * 64, ) assert result.id == "span-1" + + +class TestAgentInstanceGet: + async def test_get_returns_agent_instance(self, config): + instance_data = { + "status": "success", + "details": { + "type": "agent_instance", + "id": "inst-123", + "account_id": "acc-1", + "agent_id": "agent-1", + "agent_version_id": "ver-1", + "environment_id": "env-1", + "agent_deployment_id": "dep-1", + "status": "terminated", + "termination_reason": "admin terminated", + "inserted_at": NOW, + "updated_at": NOW, + }, + } + + with aioresponses() as m: + m.get( + "https://api.test.com/api/v1/agent_instance/inst-123", + payload=instance_data, + ) + async with PrefactorHttpClient(config) as client: + result = await client.agent_instances.get("inst-123") + + assert result.id == "inst-123" + assert result.status == "terminated" + assert result.termination_reason == "admin terminated" + + +class TestAgentSpanControlSignal: + async def test_create_calls_callback_when_control_signal_present(self, config): + from unittest.mock import MagicMock + + span_data = { + "status": "success", + "details": {**MOCK_SPAN}, + "control": {"terminate": True, "reason": "demo termination"}, + } + + callback = MagicMock() + with aioresponses() as m: + m.post("https://api.test.com/api/v1/agent_spans", payload=span_data) + async with PrefactorHttpClient(config) as client: + await client.agent_spans.create( + agent_instance_id="inst-1", + schema_name="langchain:llm", + status="pending", + control_signal_callback=callback, + ) + + callback.assert_called_once_with("demo termination") + + async def test_create_no_callback_when_control_absent(self, config): + from unittest.mock import MagicMock + + span_data = {"status": "success", "details": {**MOCK_SPAN}} + + callback = MagicMock() + with aioresponses() as m: + m.post("https://api.test.com/api/v1/agent_spans", payload=span_data) + async with PrefactorHttpClient(config) as client: + await client.agent_spans.create( + agent_instance_id="inst-1", + schema_name="langchain:llm", + status="pending", + control_signal_callback=callback, + ) + + callback.assert_not_called() + + async def test_create_ignores_non_dict_control_payload(self, config): + from unittest.mock import MagicMock + + span_data = { + "status": "success", + "details": {**MOCK_SPAN}, + "control": "terminate", + } + + callback = MagicMock() + with aioresponses() as m: + m.post("https://api.test.com/api/v1/agent_spans", payload=span_data) + async with PrefactorHttpClient(config) as client: + await client.agent_spans.create( + agent_instance_id="inst-1", + schema_name="langchain:llm", + status="pending", + control_signal_callback=callback, + ) + + callback.assert_not_called() + + async def test_finish_calls_callback_when_control_signal_present(self, config): + from unittest.mock import MagicMock + + span_data = { + "status": "success", + "details": {**MOCK_SPAN}, + "control": {"terminate": True, "reason": None}, + } + + callback = MagicMock() + with aioresponses() as m: + m.post( + "https://api.test.com/api/v1/agent_spans/span-1/finish", + payload=span_data, + ) + async with PrefactorHttpClient(config) as client: + await client.agent_spans.finish( + agent_span_id="span-1", + control_signal_callback=callback, + ) + + callback.assert_called_once_with(None) + + async def test_finish_ignores_non_dict_control_payload(self, config): + from unittest.mock import MagicMock + + span_data = { + "status": "success", + "details": {**MOCK_SPAN}, + "control": ["terminate"], + } + + callback = MagicMock() + with aioresponses() as m: + m.post( + "https://api.test.com/api/v1/agent_spans/span-1/finish", + payload=span_data, + ) + async with PrefactorHttpClient(config) as client: + await client.agent_spans.finish( + agent_span_id="span-1", + control_signal_callback=callback, + ) + + callback.assert_not_called() diff --git a/packages/http/tests/test_models.py b/packages/http/tests/test_models.py index fe29458..00cd5c2 100644 --- a/packages/http/tests/test_models.py +++ b/packages/http/tests/test_models.py @@ -352,3 +352,37 @@ def test_summary_optional(self): def test_summary_present(self): span = self._make_span(summary="Completed tool call") assert span.summary == "Completed tool call" + + +class TestAgentInstanceTerminatedReason: + """Tests parsing of AgentInstance termination reasons.""" + + def _make_instance(self, status="active", **kwargs): + """Create an AgentInstance with default required fields.""" + from datetime import datetime, timezone + + return AgentInstance( + type="agent_instance", + id="inst-1", + account_id="acc-1", + agent_id="agent-1", + agent_version_id="ver-1", + environment_id="env-1", + agent_deployment_id="dep-1", + status=status, + inserted_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + **kwargs, + ) + + def test_termination_reason_defaults_none(self): + """Termination reason defaults to None for non-terminated instances.""" + instance = self._make_instance() + assert instance.termination_reason is None + + def test_termination_reason_parsed(self): + """Termination reason is parsed when present on a terminated instance.""" + instance = self._make_instance( + status="terminated", termination_reason="admin action" + ) + assert instance.termination_reason == "admin action" diff --git a/packages/langchain/examples/termination_demo.py b/packages/langchain/examples/termination_demo.py new file mode 100644 index 0000000..de0cf3f --- /dev/null +++ b/packages/langchain/examples/termination_demo.py @@ -0,0 +1,203 @@ +"""Termination demo — runs a LangChain agent in a service loop and demonstrates +automatic detection of p2-initiated termination. + +Required env vars: + PREFACTOR_API_URL e.g. http://localhost:4000 + PREFACTOR_AGENT_ID agent ID on the target p2 instance + PREFACTOR_API_TOKEN deployment/BA token for SDK (span creation etc.) + PREFACTOR_BA_TOKEN BA token for terminate API (falls back to + PREFACTOR_API_TOKEN) + +Optional env vars: + PREFACTOR_AUTO_TERMINATE_DELAY seconds before demo calls terminate API (default: 6) + PREFACTOR_RESTART_DELAY seconds to wait between runs (default: 75) + +Usage: + PREFACTOR_API_URL=http://localhost:4000 \\ + PREFACTOR_AGENT_ID= \\ + PREFACTOR_API_TOKEN= \\ + PREFACTOR_AUTO_TERMINATE_DELAY=6 \\ + PREFACTOR_RESTART_DELAY=75 \\ + python packages/langchain/examples/termination_demo.py +""" + +from __future__ import annotations + +import asyncio +import logging +import os + +import aiohttp +from langchain.agents import create_agent +from langchain_anthropic import ChatAnthropic +from langchain_core.tools import tool +from prefactor_core import PrefactorTerminatedError +from prefactor_langchain.middleware import PrefactorMiddleware + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s" +) +logger = logging.getLogger("termination-demo") + + +@tool +def get_current_time() -> str: + """Return the current UTC time as a string.""" + from datetime import datetime, timezone + + return datetime.now(timezone.utc).isoformat() + + +async def terminate_after_delay( + api_url: str, + ba_token: str, + instance_id: str, + delay: float, +) -> None: + """Call the terminate endpoint after a configured delay. + + Args: + api_url: Base Prefactor API URL. + ba_token: Bearer token for terminate API authorization. + instance_id: Agent instance identifier to terminate. + delay: Seconds to wait before posting the terminate request. + + Returns: + None. + + Raises: + aiohttp.ClientResponseError: If the terminate API returns a non-2xx + response. + aiohttp.ClientError: If the terminate request fails before a response. + """ + await asyncio.sleep(delay) + url = f"{api_url.rstrip('/')}/api/v1/agent_instance/{instance_id}/terminate" + logger.info("Calling terminate API: POST %s", url) + timeout = aiohttp.ClientTimeout(total=10) + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.post( + url, + headers={ + "Authorization": f"Bearer {ba_token}", + "Content-Type": "application/json", + }, + json={"reason": "demo termination"}, + ) as resp: + body = await resp.text() + logger.info("Terminate API: status=%s body=%s", resp.status, body) + resp.raise_for_status() + + +async def _cancel_and_await_terminate_task(task: asyncio.Task[None]) -> None: + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + except Exception: + logger.exception("Auto-terminate task failed") + raise + + +async def run_once( + run_number: int, + api_url: str, + api_token: str, + agent_id: str, + auto_terminate_delay: float, +) -> None: + """Run one termination-demo agent session. + + Args: + run_number: Sequential run counter used in log messages. + api_url: Base Prefactor API URL. + api_token: SDK API token for Prefactor telemetry. + agent_id: Agent identifier used to initialize the middleware. + auto_terminate_delay: Seconds before the demo calls the terminate API. + + Returns: + None. + + Raises: + PrefactorTerminatedError: If Prefactor signals agent termination. + Exception: If agent invocation, termination cleanup, or middleware + cleanup fails. + """ + environment_id = os.environ.get("PREFACTOR_ENVIRONMENT_ID") + middleware = PrefactorMiddleware.from_config( + api_url=api_url, + api_token=api_token, + agent_id=agent_id, + agent_name="termination-demo-agent", + environment_id=environment_id, + ) + + model = ChatAnthropic(model_name="claude-haiku-4-5-20251001") + agent = create_agent( + model, tools=[get_current_time], middleware=[middleware], checkpointer=None + ) + + instance = await middleware.ensure_initialized() + logger.info("Run #%d — Agent instance: %s", run_number, instance.id) + logger.info("Auto-terminate in %.0fs...", auto_terminate_delay) + + ba_token = os.environ.get("PREFACTOR_BA_TOKEN", api_token) + terminate_task = asyncio.create_task( + terminate_after_delay(api_url, ba_token, instance.id, auto_terminate_delay) + ) + + try: + result = await agent.ainvoke( + { + "messages": [ + { + "role": "user", + "content": "What time is it? Then keep asking every 2 seconds.", + } + ] + }, + ) + logger.info("Run #%d completed normally: %s", run_number, result) + except PrefactorTerminatedError as e: + logger.info("Run #%d terminated: %s", run_number, e) + raise + finally: + try: + await _cancel_and_await_terminate_task(terminate_task) + finally: + await middleware.close() + + +async def main() -> None: + """Run the termination demo restart loop. + + Returns: + None. + """ + api_url = os.environ["PREFACTOR_API_URL"] + api_token = os.environ["PREFACTOR_API_TOKEN"] + agent_id = os.environ["PREFACTOR_AGENT_ID"] + auto_terminate_delay = float(os.environ.get("PREFACTOR_AUTO_TERMINATE_DELAY", "6")) + restart_delay = float(os.environ.get("PREFACTOR_RESTART_DELAY", "75")) + + run_number = 0 + while True: + run_number += 1 + try: + await run_once( + run_number, api_url, api_token, agent_id, auto_terminate_delay + ) + except PrefactorTerminatedError: + logger.info("Service continues — next run in %.0fs.", restart_delay) + await asyncio.sleep(restart_delay) + except (KeyboardInterrupt, asyncio.CancelledError): + logger.info("Stopped by user.") + break + except Exception as e: + logger.exception("Unexpected error in run #%d: %s", run_number, e) + await asyncio.sleep(restart_delay) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/packages/langchain/pyproject.toml b/packages/langchain/pyproject.toml index 155ebc1..fcfca6c 100644 --- a/packages/langchain/pyproject.toml +++ b/packages/langchain/pyproject.toml @@ -9,7 +9,7 @@ authors = [ ] requires-python = ">=3.11.0, <4.0.0" dependencies = [ - "prefactor-core>=0.2.5", + "prefactor-core>=0.2.6", "langchain-core>=1.0.0", ] diff --git a/packages/langchain/src/prefactor_langchain/_version.py b/packages/langchain/src/prefactor_langchain/_version.py index d915904..d84d7e1 100644 --- a/packages/langchain/src/prefactor_langchain/_version.py +++ b/packages/langchain/src/prefactor_langchain/_version.py @@ -1,5 +1,5 @@ """Package version for prefactor-langchain.""" PACKAGE_NAME = "prefactor-langchain" -__version__ = "0.2.5" +__version__ = "0.2.6" PACKAGE_VERSION = __version__ diff --git a/packages/langchain/src/prefactor_langchain/middleware.py b/packages/langchain/src/prefactor_langchain/middleware.py index 4734384..275c26c 100644 --- a/packages/langchain/src/prefactor_langchain/middleware.py +++ b/packages/langchain/src/prefactor_langchain/middleware.py @@ -15,6 +15,7 @@ PrefactorCoreClient, PrefactorCoreConfig, PrefactorTelemetryFailureError, + PrefactorTerminatedError, SchemaRegistry, SpanContext, ) @@ -167,6 +168,7 @@ def __init__( self._client = None self._agent_id = agent_id self._agent_name = agent_name + self._environment_id: str | None = None self._instance = instance self._owns_instance = False self._owns_client = False @@ -180,6 +182,7 @@ def __init__( self._tool_span_types = ( self._register_tool_schemas(None, tool_schemas) if tool_schemas else {} ) + self._get_termination_event = None return if client is None: @@ -201,6 +204,7 @@ def __init__( self._client = client self._agent_id = agent_id self._agent_name = agent_name + self._environment_id = None self._instance: AgentInstanceHandle | None = None self._owns_instance = True @@ -214,6 +218,7 @@ def __init__( self._pending_emit_futures: list[asyncio.Task[None]] = [] self._pending_emit_error: Exception | None = None self._tool_span_types = {} + self._get_termination_event = None if tool_schemas: self._tool_span_types = self._register_tool_schemas(client, tool_schemas) @@ -237,6 +242,19 @@ def _prefer_shutdown_error( return new_error return current + def _throw_if_terminated(self) -> None: + if self._get_termination_event is None: + return + event = self._get_termination_event() + if event.is_set(): + monitor = ( + self._client._termination_monitor + if self._client and hasattr(self._client, "_termination_monitor") + else None + ) + reason = monitor.termination_reason if monitor is not None else None + raise PrefactorTerminatedError(reason) + @classmethod def from_config( cls, @@ -244,6 +262,7 @@ def from_config( api_token: str, agent_id: str | None = None, agent_name: str | None = None, + environment_id: str | None = None, schema_registry: SchemaRegistry | None = None, include_langchain_schemas: bool = True, tool_schemas: Mapping[str, LangChainToolSchemaConfig | Mapping[str, Any]] @@ -259,6 +278,7 @@ def from_config( api_token: The API token for authentication. agent_id: Optional agent identifier for categorization. agent_name: Optional human-readable agent name. + environment_id: Optional environment identifier for scoping the agent. schema_registry: Optional SchemaRegistry for registering span schemas. include_langchain_schemas: If True and schema_registry is provided, automatically register LangChain-specific schemas. @@ -303,6 +323,7 @@ def from_config( middleware._client = client middleware._agent_id = agent_id middleware._agent_name = agent_name + middleware._environment_id = environment_id middleware._instance = None middleware._owns_instance = True middleware._owns_client = True @@ -314,6 +335,7 @@ def from_config( middleware._pending_emit_futures = [] middleware._pending_emit_error = None middleware._tool_span_types = tool_span_types + middleware._get_termination_event = None logger.debug("PrefactorMiddleware created via from_config()") return middleware @@ -393,10 +415,19 @@ async def _ensure_initialized(self) -> AgentInstanceHandle: }, agent_schema_version=None, # Will use registry if available external_schema_version_id=schema_version_id, + environment_id=self._environment_id, ) self._owns_instance = True await self._instance.start() + if ( + self._client is not None + and hasattr(self._client, "_termination_monitor") + and self._client._termination_monitor is not None + ): + self._get_termination_event = ( + self._client._termination_monitor.get_termination_event + ) logger.debug("Initialized agent instance %s", self._instance.id) return self._instance @@ -772,6 +803,7 @@ def before_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None: Optional state updates. """ try: + self._throw_if_terminated() if self._instance is None: return None @@ -796,6 +828,8 @@ def before_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None: except Exception as e: _raise_if_telemetry_failure(e) + if isinstance(e, PrefactorTerminatedError): + raise logger.error("Error in before_agent: %s", e, exc_info=True) return None @@ -879,6 +913,7 @@ async def abefore_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None Optional state updates. """ try: + self._throw_if_terminated() instance = await self._ensure_initialized() messages = [] @@ -918,6 +953,8 @@ async def abefore_agent(self, state: Any, runtime: Any) -> dict[str, Any] | None except Exception as e: _raise_if_telemetry_failure(e) + if isinstance(e, PrefactorTerminatedError): + raise logger.error("Error in abefore_agent: %s", e, exc_info=True) return None @@ -976,6 +1013,7 @@ def wrap_model_call( Returns: The model response. """ + self._throw_if_terminated() inputs = self._extract_model_inputs(request) span_data = LLMSpan( name=self._get_name_from_request(request), @@ -1014,6 +1052,7 @@ async def awrap_model_call( Returns: The model response. """ + self._throw_if_terminated() instance = await self._ensure_initialized() inputs = self._extract_model_inputs(request) @@ -1058,6 +1097,7 @@ def wrap_tool_call( Returns: The tool response. """ + self._throw_if_terminated() inputs = self._extract_tool_inputs(request) tool_name = inputs.get("tool_name", "unknown_tool") @@ -1100,6 +1140,7 @@ async def awrap_tool_call( Returns: The tool response. """ + self._throw_if_terminated() instance = await self._ensure_initialized() inputs = self._extract_tool_inputs(request) diff --git a/packages/langchain/tests/test_middleware.py b/packages/langchain/tests/test_middleware.py index 75b5643..f2f71de 100644 --- a/packages/langchain/tests/test_middleware.py +++ b/packages/langchain/tests/test_middleware.py @@ -820,4 +820,101 @@ def test_langchain_tool_schema(self): """Test LANGCHAIN_TOOL_SCHEMA is exported.""" assert LANGCHAIN_TOOL_SCHEMA is not None assert LANGCHAIN_TOOL_SCHEMA.get("type") == "object" - assert "properties" in LANGCHAIN_TOOL_SCHEMA + + +class TestMiddlewareThrowIfTerminated: + """Tests that middleware raises PrefactorTerminatedError when monitor is set.""" + + def _make_middleware_with_monitor( + self, + terminated: bool = False, + reason: str | None = "test reason", + ): + from prefactor_core.monitoring.termination_monitor import TerminationMonitor + from prefactor_langchain.middleware import PrefactorMiddleware + + fetch = AsyncMock(return_value=Mock(status="active", termination_reason=None)) + monitor = TerminationMonitor(fetch_instance=fetch) + if terminated: + monitor.detect_termination(reason) + + mock_client = Mock() + mock_client._initialized = True + mock_client._termination_monitor = monitor + + middleware = PrefactorMiddleware.__new__(PrefactorMiddleware) + middleware._client = mock_client + middleware._agent_id = None + middleware._agent_name = None + middleware._instance = Mock() + middleware._owns_instance = False + middleware._owns_client = False + middleware._agent_span_cm = None + middleware._agent_span_context = None + middleware._agent_span_id = None + middleware._current_parent_span_id = None + middleware._loop = asyncio.get_event_loop() + middleware._pending_emit_futures = [] + middleware._pending_emit_error = None + middleware._tool_span_types = {} + middleware._get_termination_event = monitor.get_termination_event + + return middleware, monitor + + async def test_throw_if_terminated_raises_when_event_set(self): + from prefactor_core import PrefactorTerminatedError + + middleware, _ = self._make_middleware_with_monitor( + terminated=True, reason="test" + ) + with pytest.raises(PrefactorTerminatedError) as exc_info: + middleware._throw_if_terminated() + assert exc_info.value.reason == "test" + + async def test_throw_if_terminated_noop_when_not_terminated(self): + middleware, _ = self._make_middleware_with_monitor(terminated=False) + middleware._throw_if_terminated() # should not raise + + async def test_throw_if_terminated_noop_when_getter_is_none(self): + from prefactor_langchain.middleware import PrefactorMiddleware + + middleware = PrefactorMiddleware.__new__(PrefactorMiddleware) + middleware._get_termination_event = None + middleware._client = None + middleware._throw_if_terminated() # should not raise + + async def test_before_agent_raises_when_terminated(self): + from prefactor_core import PrefactorTerminatedError + + middleware, _ = self._make_middleware_with_monitor( + terminated=True, reason="terminated" + ) + with pytest.raises(PrefactorTerminatedError): + middleware.before_agent({}, Mock()) + + async def test_abefore_agent_raises_when_terminated(self): + from prefactor_core import PrefactorTerminatedError + + middleware, _ = self._make_middleware_with_monitor( + terminated=True, reason="terminated" + ) + with pytest.raises(PrefactorTerminatedError): + await middleware.abefore_agent({}, Mock()) + + async def test_awrap_model_call_raises_when_terminated(self): + from prefactor_core import PrefactorTerminatedError + + middleware, _ = self._make_middleware_with_monitor( + terminated=True, reason="terminated" + ) + with pytest.raises(PrefactorTerminatedError): + await middleware.awrap_model_call(Mock(), AsyncMock()) + + async def test_awrap_tool_call_raises_when_terminated(self): + from prefactor_core import PrefactorTerminatedError + + middleware, _ = self._make_middleware_with_monitor( + terminated=True, reason="terminated" + ) + with pytest.raises(PrefactorTerminatedError): + await middleware.awrap_tool_call(Mock(), AsyncMock()) diff --git a/packages/langchain/tests/test_termination_demo.py b/packages/langchain/tests/test_termination_demo.py new file mode 100644 index 0000000..84d6bbe --- /dev/null +++ b/packages/langchain/tests/test_termination_demo.py @@ -0,0 +1,73 @@ +"""Tests for the LangChain termination demo.""" + +from __future__ import annotations + +import asyncio +import importlib.util +import sys +from pathlib import Path + +import aiohttp +import pytest +from aioresponses import aioresponses + + +def _load_termination_demo(): + module_path = ( + Path(__file__).parents[1] / "examples" / "termination_demo.py" + ).resolve() + spec = importlib.util.spec_from_file_location("termination_demo", module_path) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +@pytest.mark.asyncio +async def test_terminate_after_delay_raises_for_non_success_response(): + """Terminate API failures should be visible to the demo caller.""" + demo = _load_termination_demo() + url = "https://api.test.com/api/v1/agent_instance/inst-1/terminate" + + with aioresponses() as responses: + responses.post(url, status=500, payload={"error": "failed"}) + + with pytest.raises(aiohttp.ClientResponseError): + await demo.terminate_after_delay( + api_url="https://api.test.com", + ba_token="token", + instance_id="inst-1", + delay=0, + ) + + +@pytest.mark.asyncio +async def test_cancel_and_await_terminate_task_reraises_completed_failures(): + """Completed terminate task failures should not be swallowed.""" + demo = _load_termination_demo() + + async def fail(): + raise RuntimeError("terminate failed") + + task = asyncio.create_task(fail()) + await asyncio.sleep(0) + + with pytest.raises(RuntimeError, match="terminate failed"): + await demo._cancel_and_await_terminate_task(task) + + +@pytest.mark.asyncio +async def test_cancel_and_await_terminate_task_ignores_normal_cancellation(): + """Cancelling an unfinished terminate task should not fail cleanup.""" + demo = _load_termination_demo() + + async def wait_forever(): + await asyncio.Event().wait() + + task = asyncio.create_task(wait_forever()) + + await demo._cancel_and_await_terminate_task(task) + + assert task.cancelled()