prefactor-core 0.2.5__tar.gz → 0.2.7__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (37) hide show
  1. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/.gitignore +0 -1
  2. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/PKG-INFO +2 -2
  3. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/examples/agent_e2e.py +4 -0
  4. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/pyproject.toml +1 -1
  5. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/__init__.py +2 -0
  6. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/_version.py +1 -1
  7. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/client.py +90 -7
  8. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/exceptions.py +18 -0
  9. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/managers/agent_instance.py +7 -1
  10. prefactor_core-0.2.7/src/prefactor_core/monitoring/__init__.py +5 -0
  11. prefactor_core-0.2.7/src/prefactor_core/monitoring/termination_monitor.py +151 -0
  12. prefactor_core-0.2.7/tests/monitoring/__init__.py +0 -0
  13. prefactor_core-0.2.7/tests/monitoring/test_termination_monitor.py +259 -0
  14. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_agent_instance_finish_status.py +70 -0
  15. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_agent_instance_register.py +3 -0
  16. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_failure_handling.py +87 -1
  17. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_sdk_header.py +8 -0
  18. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/README.md +0 -0
  19. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/config.py +0 -0
  20. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/context_stack.py +0 -0
  21. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/managers/__init__.py +0 -0
  22. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/managers/span.py +0 -0
  23. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/models.py +0 -0
  24. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/operations.py +0 -0
  25. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/queue/__init__.py +0 -0
  26. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/queue/base.py +0 -0
  27. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/queue/executor.py +0 -0
  28. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/queue/memory.py +0 -0
  29. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/schema_registry.py +0 -0
  30. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/span_context.py +0 -0
  31. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/src/prefactor_core/utils.py +0 -0
  32. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_client.py +0 -0
  33. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_imports.py +0 -0
  34. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_queue.py +0 -0
  35. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_span_context.py +0 -0
  36. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_span_manager.py +0 -0
  37. {prefactor_core-0.2.5 → prefactor_core-0.2.7}/tests/test_utils.py +0 -0
@@ -41,7 +41,6 @@ env/
41
41
 
42
42
  # uv
43
43
  .uv/
44
- uv.lock
45
44
 
46
45
  # Type checkers
47
46
  .mypy_cache/
@@ -1,11 +1,11 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: prefactor-core
3
- Version: 0.2.5
3
+ Version: 0.2.7
4
4
  Summary: Core Prefactor SDK with async queue-based operations
5
5
  Author-email: Prefactor Pty Ltd <josh@prefactor.tech>
6
6
  License: MIT
7
7
  Requires-Python: <4.0.0,>=3.11.0
8
- Requires-Dist: prefactor-http>=0.1.3
8
+ Requires-Dist: prefactor-http>=0.1.5
9
9
  Requires-Dist: pydantic>=2.0.0
10
10
  Description-Content-Type: text/markdown
11
11
 
@@ -31,6 +31,8 @@ deployment-scoped token, leave it unset and the API derives the agent and
31
31
  environment from the token.
32
32
  """
33
33
 
34
+ from __future__ import annotations
35
+
34
36
  import asyncio
35
37
  import os
36
38
 
@@ -186,6 +188,8 @@ async def simulate_retrieval(query: str) -> list[str]:
186
188
 
187
189
  async def main() -> None:
188
190
  agent_id = os.environ.get("PREFACTOR_AGENT_ID")
191
+ if agent_id is not None:
192
+ agent_id = agent_id.strip() or None
189
193
 
190
194
  config = PrefactorCoreConfig(
191
195
  http_config=HttpClientConfig(
@@ -9,7 +9,7 @@ authors = [
9
9
  ]
10
10
  requires-python = ">=3.11.0, <4.0.0"
11
11
  dependencies = [
12
- "prefactor-http>=0.1.3",
12
+ "prefactor-http>=0.1.5",
13
13
  "pydantic>=2.0.0",
14
14
  ]
15
15
 
@@ -16,6 +16,7 @@ from .exceptions import (
16
16
  OperationError,
17
17
  PrefactorCoreError,
18
18
  PrefactorTelemetryFailureError,
19
+ PrefactorTerminatedError,
19
20
  SpanNotFoundError,
20
21
  )
21
22
  from .managers.agent_instance import AgentInstanceHandle
@@ -43,6 +44,7 @@ __all__ = [
43
44
  "InstanceNotFoundError",
44
45
  "SpanNotFoundError",
45
46
  "PrefactorTelemetryFailureError",
47
+ "PrefactorTerminatedError",
46
48
  # Models
47
49
  "AgentInstance",
48
50
  "Span",
@@ -3,5 +3,5 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  PACKAGE_NAME = "prefactor-core"
6
- __version__ = "0.2.5"
6
+ __version__ = "0.2.7"
7
7
  PACKAGE_VERSION = __version__
@@ -7,13 +7,18 @@ through an async queue-based architecture.
7
7
 
8
8
  from __future__ import annotations
9
9
 
10
+ import asyncio
10
11
  import logging
11
12
  import time
12
13
  from contextlib import asynccontextmanager
13
14
  from typing import TYPE_CHECKING, Any
14
15
 
15
16
  from prefactor_http.client import PrefactorHttpClient
16
- from prefactor_http.exceptions import is_permanent_http_error, is_transient_http_error
17
+ from prefactor_http.exceptions import (
18
+ PrefactorApiError,
19
+ is_permanent_http_error,
20
+ is_transient_http_error,
21
+ )
17
22
 
18
23
  from ._version import PACKAGE_NAME as CORE_PACKAGE_NAME
19
24
  from ._version import PACKAGE_VERSION as CORE_PACKAGE_VERSION
@@ -26,6 +31,7 @@ from .exceptions import (
26
31
  )
27
32
  from .managers.agent_instance import AgentInstanceManager
28
33
  from .managers.span import SpanManager
34
+ from .monitoring.termination_monitor import TerminationMonitor
29
35
  from .operations import Operation, OperationType
30
36
  from .queue.base import Queue
31
37
  from .queue.executor import TaskExecutor
@@ -86,6 +92,9 @@ class PrefactorCoreClient:
86
92
  self._initialized = False
87
93
  self._telemetry_failure: PrefactorTelemetryFailureError | None = None
88
94
  self._telemetry_failure_observed = False
95
+ self._termination_monitor: TerminationMonitor | None = None
96
+ self._sync_task: asyncio.Task | None = None
97
+ self._current_instance_id: str | None = None
89
98
 
90
99
  def _build_http_sdk_header(self) -> str:
91
100
  """Build the effective SDK header for HTTP requests."""
@@ -133,6 +142,14 @@ class PrefactorCoreClient:
133
142
  )
134
143
  await self._http.__aenter__()
135
144
 
145
+ try:
146
+ await self._http.validate_token()
147
+ except BaseException as exc:
148
+ http = self._http
149
+ self._http = None
150
+ await http.__aexit__(type(exc), exc, exc.__traceback__)
151
+ raise
152
+
136
153
  # Initialize executor
137
154
  self._executor = TaskExecutor(
138
155
  queue=self._queue,
@@ -155,6 +172,13 @@ class PrefactorCoreClient:
155
172
 
156
173
  self._initialized = True
157
174
 
175
+ self._termination_monitor = TerminationMonitor(
176
+ fetch_instance=self._fetch_instance_for_poll,
177
+ )
178
+ self._sync_task = asyncio.create_task(
179
+ self._run_sync_loop(), name="prefactor-termination-sync"
180
+ )
181
+
158
182
  async def close(self) -> None:
159
183
  """Close the client and cleanup resources.
160
184
 
@@ -164,6 +188,24 @@ class PrefactorCoreClient:
164
188
  if not self._initialized:
165
189
  return
166
190
 
191
+ if self._sync_task is not None:
192
+ if not self._sync_task.done():
193
+ self._sync_task.cancel()
194
+ try:
195
+ await self._sync_task
196
+ except asyncio.CancelledError:
197
+ pass
198
+ except Exception:
199
+ logger.exception(
200
+ "Termination sync loop exited with error during close()"
201
+ )
202
+ finally:
203
+ self._sync_task = None
204
+
205
+ if self._termination_monitor is not None:
206
+ self._termination_monitor.destroy()
207
+ self._termination_monitor = None
208
+
167
209
  # Stop executor
168
210
  if self._executor:
169
211
  await self._executor.stop()
@@ -272,12 +314,23 @@ class PrefactorCoreClient:
272
314
  )
273
315
 
274
316
  elif operation.type == OperationType.FINISH_AGENT_INSTANCE:
275
- await self._http.agent_instances.finish(
276
- agent_instance_id=operation.payload["instance_id"],
277
- status=operation.payload.get("status", "complete"),
278
- timestamp=operation.timestamp,
279
- idempotency_key=operation.payload.get("idempotency_key"),
280
- )
317
+ try:
318
+ await self._http.agent_instances.finish(
319
+ agent_instance_id=operation.payload["instance_id"],
320
+ status=operation.payload.get("status", "complete"),
321
+ timestamp=operation.timestamp,
322
+ idempotency_key=operation.payload.get("idempotency_key"),
323
+ )
324
+ except PrefactorApiError as finish_err:
325
+ if finish_err.status_code == 409:
326
+ logger.debug(
327
+ "[prefactor:http] Agent instance %s already in"
328
+ " terminal state; skipping finish.",
329
+ operation.payload["instance_id"],
330
+ )
331
+ return
332
+ raise
333
+
281
334
  elif operation.type == OperationType.CREATE_SPAN:
282
335
  await self._http.agent_spans.create(
283
336
  agent_instance_id=operation.payload["instance_id"],
@@ -286,6 +339,7 @@ class PrefactorCoreClient:
286
339
  id=operation.payload.get("span_id"),
287
340
  parent_span_id=operation.payload.get("parent_span_id"),
288
341
  payload=operation.payload.get("payload"),
342
+ control_signal_callback=self._on_control_signal,
289
343
  )
290
344
 
291
345
  elif operation.type == OperationType.FINISH_SPAN:
@@ -295,6 +349,7 @@ class PrefactorCoreClient:
295
349
  result_payload=operation.payload.get("result_payload"),
296
350
  timestamp=operation.timestamp,
297
351
  idempotency_key=operation.payload.get("idempotency_key"),
352
+ control_signal_callback=self._on_control_signal,
298
353
  )
299
354
 
300
355
  except Exception as e:
@@ -307,6 +362,32 @@ class PrefactorCoreClient:
307
362
  )
308
363
  raise
309
364
 
365
+ async def _fetch_instance_for_poll(self, instance_id: str):
366
+ if self._http is None:
367
+ return None
368
+ return await self._http.agent_instances.get(instance_id)
369
+
370
+ async def _run_sync_loop(self) -> None:
371
+ while True:
372
+ await asyncio.sleep(1)
373
+ if self._termination_monitor is None:
374
+ continue
375
+ try:
376
+ self._termination_monitor.sync(self._current_instance_id)
377
+ except Exception:
378
+ logger.exception("Termination sync iteration failed")
379
+
380
+ def _on_control_signal(self, reason: str | None) -> None:
381
+ if self._termination_monitor is not None:
382
+ self._termination_monitor.detect_termination(reason)
383
+
384
+ def _set_current_instance(self, instance_id: str | None) -> None:
385
+ self._current_instance_id = instance_id
386
+
387
+ def _clear_current_instance(self, instance_id: str) -> None:
388
+ if self._current_instance_id == instance_id:
389
+ self._current_instance_id = None
390
+
310
391
  @property
311
392
  def instance_manager(self) -> AgentInstanceManager | None:
312
393
  """Public accessor for the agent instance manager."""
@@ -381,6 +462,8 @@ class PrefactorCoreClient:
381
462
  environment_id=environment_id,
382
463
  )
383
464
 
465
+ self._set_current_instance(instance_id)
466
+
384
467
  return AgentInstanceHandle(
385
468
  instance_id=instance_id,
386
469
  client=self,
@@ -58,6 +58,23 @@ class PrefactorTelemetryFailureError(PrefactorCoreError):
58
58
  self.dropped_operations = dropped_operations
59
59
 
60
60
 
61
+ class PrefactorTerminatedError(PrefactorCoreError):
62
+ """Raised when the agent instance has been terminated by p2.
63
+
64
+ Args:
65
+ reason: Optional reason reported by p2 for the termination.
66
+ """
67
+
68
+ def __init__(self, reason: str | None = None) -> None:
69
+ msg = (
70
+ f"Agent instance terminated by p2: {reason}"
71
+ if reason
72
+ else "Agent instance terminated by p2"
73
+ )
74
+ super().__init__(msg)
75
+ self.reason = reason
76
+
77
+
61
78
  __all__ = [
62
79
  "PrefactorCoreError",
63
80
  "ClientNotInitializedError",
@@ -66,4 +83,5 @@ __all__ = [
66
83
  "InstanceNotFoundError",
67
84
  "SpanNotFoundError",
68
85
  "PrefactorTelemetryFailureError",
86
+ "PrefactorTerminatedError",
69
87
  ]
@@ -218,12 +218,18 @@ class AgentInstanceHandle:
218
218
  async def finish(self, status: "FinishStatus" = "complete") -> None:
219
219
  """Mark the instance as finished.
220
220
 
221
- This queues a finish operation for the instance.
221
+ Resets the termination monitor (fence + new event) before enqueueing
222
+ the HTTP finish so stale span responses from the dying run cannot
223
+ trigger termination on the next run.
222
224
 
223
225
  Args:
224
226
  status: Terminal status for the instance — one of ``"complete"``,
225
227
  ``"failed"``, or ``"cancelled"``. Defaults to ``"complete"``.
226
228
  """
229
+ monitor = getattr(self._client, "_termination_monitor", None)
230
+ if monitor is not None:
231
+ monitor.reset()
232
+ self._client._clear_current_instance(self._instance_id)
227
233
  manager = self._client.instance_manager
228
234
  assert manager is not None
229
235
  await manager.finish_with_idempotency_key(
@@ -0,0 +1,5 @@
1
+ from __future__ import annotations
2
+
3
+ from prefactor_core.monitoring.termination_monitor import TerminationMonitor
4
+
5
+ __all__ = ["TerminationMonitor"]
@@ -0,0 +1,151 @@
1
+ """TerminationMonitor — detects agent instance termination via two paths:
2
+ 1. Fast: control signal in span API responses (pushed via callback)
3
+ 2. Slow: polling the instance status endpoint every `poll_interval` seconds
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ import asyncio
9
+ import logging
10
+ from typing import Callable
11
+
12
+ logger = logging.getLogger(__name__)
13
+
14
+
15
+ class TerminationMonitor:
16
+ """Monitors an agent instance for termination.
17
+
18
+ Thread-safety note: `detect_termination` and `get_termination_event().is_set()`
19
+ are safe to call from sync worker threads — they only read/write a bool under
20
+ the GIL (asyncio.Event._value is a plain bool).
21
+ """
22
+
23
+ def __init__(
24
+ self,
25
+ fetch_instance: Callable,
26
+ poll_interval: float = 30.0,
27
+ ) -> None:
28
+ self._fetch_instance = fetch_instance
29
+ self._poll_interval = poll_interval
30
+
31
+ self._event = asyncio.Event()
32
+ self._termination_reason: str | None = None
33
+ self._callbacks: list[Callable[[], None]] = []
34
+
35
+ self._current_instance_id: str | None = None
36
+ self._poll_task: asyncio.Task | None = None
37
+ self._destroyed = False
38
+
39
+ # Fence blocks stale span callbacks from a previous run after reset()
40
+ self._fenced = False
41
+ # Generation increments on each reset so stale polls self-discard
42
+ self._generation = 0
43
+
44
+ # ------------------------------------------------------------------
45
+ # Public API
46
+ # ------------------------------------------------------------------
47
+
48
+ @property
49
+ def termination_reason(self) -> str | None:
50
+ return self._termination_reason
51
+
52
+ def get_termination_event(self) -> asyncio.Event:
53
+ return self._event
54
+
55
+ def detect_termination(self, reason: str | None) -> None:
56
+ """Signal that the instance has been terminated.
57
+
58
+ No-op if already terminated, destroyed, or fenced (post-reset stale call).
59
+ """
60
+ if self._destroyed or self._event.is_set() or self._fenced:
61
+ return
62
+ self._termination_reason = reason
63
+ self._event.set()
64
+ self._stop_poll()
65
+ for cb in list(self._callbacks):
66
+ try:
67
+ cb()
68
+ except Exception:
69
+ logger.exception("Termination callback failed: %r", cb)
70
+
71
+ def sync(self, instance_id: str | None) -> None:
72
+ """Update the tracked instance ID and start/stop the fallback poll.
73
+
74
+ Idempotent: calling with the same non-None ID when a poll is already
75
+ running does not restart it (preserving the sleep interval).
76
+ """
77
+ if instance_id is not None:
78
+ self._fenced = False
79
+ if instance_id == self._current_instance_id:
80
+ # Same ID — no change needed; poll (if any) keeps running
81
+ return
82
+ self._current_instance_id = instance_id
83
+ self._stop_poll()
84
+ if instance_id is not None and not self._event.is_set() and not self._destroyed:
85
+ self._start_poll(instance_id, self._generation)
86
+
87
+ def reset(self) -> None:
88
+ """Prepare monitor for the next agent run.
89
+
90
+ - Creates a fresh (unset) event
91
+ - Clears the termination reason
92
+ - Cancels any in-flight poll
93
+ - Sets fence so stale callbacks from the dying run are ignored
94
+ - Increments generation so stale polls self-discard
95
+ """
96
+ self._stop_poll()
97
+ self._event = asyncio.Event()
98
+ self._termination_reason = None
99
+ self._current_instance_id = None
100
+ self._generation += 1
101
+ self._fenced = True
102
+
103
+ def subscribe(self, callback: Callable[[], None]) -> Callable[[], None]:
104
+ """Register a callback invoked on termination. Returns an unsubscribe fn."""
105
+ self._callbacks.append(callback)
106
+
107
+ def unsubscribe() -> None:
108
+ try:
109
+ self._callbacks.remove(callback)
110
+ except ValueError:
111
+ pass
112
+
113
+ return unsubscribe
114
+
115
+ def destroy(self) -> None:
116
+ """Permanently shut down the monitor (no further events will fire)."""
117
+ self._destroyed = True
118
+ self._stop_poll()
119
+ self._callbacks.clear()
120
+
121
+ # ------------------------------------------------------------------
122
+ # Internal
123
+ # ------------------------------------------------------------------
124
+
125
+ def _start_poll(self, instance_id: str, generation: int) -> None:
126
+ self._poll_task = asyncio.get_event_loop().create_task(
127
+ self._poll(instance_id, generation)
128
+ )
129
+
130
+ def _stop_poll(self) -> None:
131
+ if self._poll_task is not None and not self._poll_task.done():
132
+ self._poll_task.cancel()
133
+ self._poll_task = None
134
+
135
+ async def _poll(self, instance_id: str, generation: int) -> None:
136
+ while True:
137
+ await asyncio.sleep(self._poll_interval)
138
+ # Guard: generation changed means reset() was called
139
+ if self._generation != generation:
140
+ return
141
+ try:
142
+ instance = await self._fetch_instance(instance_id)
143
+ except Exception:
144
+ logger.debug("Termination poll error (transient)", exc_info=True)
145
+ continue
146
+ # Guard again after await
147
+ if self._generation != generation:
148
+ return
149
+ if instance.status == "terminated":
150
+ self.detect_termination(instance.termination_reason)
151
+ return
File without changes
@@ -0,0 +1,259 @@
1
+ """Tests for TerminationMonitor — 19 tests covering primary path, fallback poll,
2
+ reset(), and callback lifecycle."""
3
+
4
+ from __future__ import annotations
5
+
6
+ import asyncio
7
+ from unittest.mock import AsyncMock, MagicMock
8
+
9
+ from prefactor_core.monitoring.termination_monitor import TerminationMonitor
10
+
11
+
12
+ def _make_monitor(fetch_instance=None) -> TerminationMonitor:
13
+ if fetch_instance is None:
14
+ fetch_instance = AsyncMock(
15
+ return_value=MagicMock(status="active", termination_reason=None)
16
+ )
17
+ return TerminationMonitor(fetch_instance=fetch_instance)
18
+
19
+
20
+ def _terminated_instance(reason: str | None = "test reason"):
21
+ inst = MagicMock()
22
+ inst.status = "terminated"
23
+ inst.termination_reason = reason
24
+ return inst
25
+
26
+
27
+ # ---------------------------------------------------------------------------
28
+ # Primary path (5 tests)
29
+ # ---------------------------------------------------------------------------
30
+
31
+
32
+ class TestPrimaryPath:
33
+ async def test_detect_termination_sets_event(self):
34
+ monitor = _make_monitor()
35
+ assert not monitor.get_termination_event().is_set()
36
+ monitor.detect_termination("reason")
37
+ assert monitor.get_termination_event().is_set()
38
+
39
+ async def test_reason_propagates(self):
40
+ monitor = _make_monitor()
41
+ monitor.detect_termination("my reason")
42
+ assert monitor.termination_reason == "my reason"
43
+
44
+ async def test_null_reason_accepted(self):
45
+ monitor = _make_monitor()
46
+ monitor.detect_termination(None)
47
+ assert monitor.get_termination_event().is_set()
48
+ assert monitor.termination_reason is None
49
+
50
+ async def test_second_detect_termination_is_idempotent(self):
51
+ callback = MagicMock()
52
+ monitor = _make_monitor()
53
+ monitor.subscribe(callback)
54
+ monitor.detect_termination("first")
55
+ monitor.detect_termination("second")
56
+ callback.assert_called_once()
57
+ assert monitor.termination_reason == "first"
58
+
59
+ async def test_detect_termination_noop_after_destroy(self):
60
+ monitor = _make_monitor()
61
+ monitor.destroy()
62
+ monitor.detect_termination("reason")
63
+ assert not monitor.get_termination_event().is_set()
64
+
65
+
66
+ # ---------------------------------------------------------------------------
67
+ # Fallback poll (5 tests)
68
+ # ---------------------------------------------------------------------------
69
+
70
+
71
+ class TestFallbackPoll:
72
+ async def test_poll_starts_when_instance_id_arrives(self):
73
+ monitor = _make_monitor()
74
+ monitor.sync("inst-1")
75
+ await asyncio.sleep(0) # yield to let task start
76
+ assert monitor._poll_task is not None
77
+ assert not monitor._poll_task.done()
78
+ monitor.destroy()
79
+
80
+ async def test_poll_stops_when_sync_called_with_none(self):
81
+ monitor = _make_monitor()
82
+ monitor.sync("inst-1")
83
+ await asyncio.sleep(0)
84
+ poll_task = monitor._poll_task
85
+ assert poll_task is not None
86
+ monitor.sync(None)
87
+ await asyncio.sleep(0)
88
+ assert poll_task.cancelled() or poll_task.done()
89
+ monitor.destroy()
90
+
91
+ async def test_no_poll_without_instance_id(self):
92
+ monitor = _make_monitor()
93
+ monitor.sync(None)
94
+ await asyncio.sleep(0)
95
+ assert monitor._poll_task is None
96
+ monitor.destroy()
97
+
98
+ async def test_poll_stops_after_termination_detected(self):
99
+ monitor = _make_monitor()
100
+ monitor.sync("inst-1")
101
+ await asyncio.sleep(0)
102
+ poll_task = monitor._poll_task
103
+ assert poll_task is not None
104
+ monitor.detect_termination("reason")
105
+ await asyncio.sleep(0)
106
+ assert poll_task.cancelled() or poll_task.done()
107
+ monitor.destroy()
108
+
109
+ async def test_poll_survives_transient_http_errors(self):
110
+ fetch = AsyncMock(side_effect=Exception("network error"))
111
+ monitor = TerminationMonitor(fetch_instance=fetch, poll_interval=0.05)
112
+ monitor.sync("inst-1")
113
+ await asyncio.sleep(0.2) # let poll fire a couple times
114
+ # monitor should not be terminated — error was swallowed
115
+ assert not monitor.get_termination_event().is_set()
116
+ monitor.destroy()
117
+
118
+
119
+ # ---------------------------------------------------------------------------
120
+ # reset() (7 tests)
121
+ # ---------------------------------------------------------------------------
122
+
123
+
124
+ class TestReset:
125
+ async def test_reset_creates_fresh_event(self):
126
+ monitor = _make_monitor()
127
+ monitor.detect_termination("reason")
128
+ old_event = monitor.get_termination_event()
129
+ assert old_event.is_set()
130
+ monitor.reset()
131
+ new_event = monitor.get_termination_event()
132
+ assert not new_event.is_set()
133
+ assert new_event is not old_event
134
+
135
+ async def test_reset_allows_new_termination_after_sync_with_new_id(self):
136
+ monitor = _make_monitor()
137
+ monitor.detect_termination("run 1")
138
+ monitor.reset()
139
+ # fenced — detect_termination should be blocked
140
+ monitor.detect_termination("stale")
141
+ assert not monitor.get_termination_event().is_set()
142
+ # sync with new id lifts fence
143
+ monitor.sync("inst-2")
144
+ monitor.detect_termination("run 2")
145
+ assert monitor.get_termination_event().is_set()
146
+ assert monitor.termination_reason == "run 2"
147
+ monitor.destroy()
148
+
149
+ async def test_get_termination_event_returns_new_event_after_reset(self):
150
+ monitor = _make_monitor()
151
+ getter = monitor.get_termination_event
152
+ monitor.detect_termination("reason")
153
+ monitor.reset()
154
+ # getter returns new (unset) event
155
+ assert not getter().is_set()
156
+
157
+ async def test_reset_cancels_poll_task(self):
158
+ monitor = _make_monitor()
159
+ monitor.sync("inst-1")
160
+ await asyncio.sleep(0)
161
+ poll_task = monitor._poll_task
162
+ assert poll_task is not None
163
+ monitor.reset()
164
+ await asyncio.sleep(0)
165
+ assert poll_task.cancelled() or poll_task.done()
166
+
167
+ async def test_reset_preserves_callbacks(self):
168
+ callback = MagicMock()
169
+ monitor = _make_monitor()
170
+ monitor.subscribe(callback)
171
+ monitor.reset()
172
+ # sync to lift fence, then detect
173
+ monitor.sync("inst-2")
174
+ monitor.detect_termination("after reset")
175
+ callback.assert_called_once()
176
+
177
+ async def test_fence_blocks_detect_termination_until_sync_with_new_id(self):
178
+ monitor = _make_monitor()
179
+ monitor.detect_termination("run 1")
180
+ monitor.reset() # fenced = True
181
+
182
+ # stale span response fires during queue drain
183
+ monitor.detect_termination("stale span response")
184
+ assert not monitor.get_termination_event().is_set()
185
+
186
+ # sync with None doesn't lift fence
187
+ monitor.sync(None)
188
+ monitor.detect_termination("still stale")
189
+ assert not monitor.get_termination_event().is_set()
190
+
191
+ # sync with new instance id lifts fence
192
+ monitor.sync("instance-2")
193
+ monitor.detect_termination("run 2")
194
+ assert monitor.get_termination_event().is_set()
195
+ assert monitor.termination_reason == "run 2"
196
+ monitor.destroy()
197
+
198
+ async def test_stale_poll_response_discarded_after_reset(self):
199
+ """Poll fires for old instance after reset — generation check discards it."""
200
+ slow_fetch = AsyncMock(return_value=_terminated_instance("old run"))
201
+ monitor = TerminationMonitor(fetch_instance=slow_fetch, poll_interval=0.05)
202
+ monitor.sync("inst-1")
203
+ await asyncio.sleep(0) # poll task started
204
+
205
+ # reset before poll completes
206
+ monitor.reset()
207
+ await asyncio.sleep(0.2) # let poll fire with old generation
208
+
209
+ # monitor should NOT be terminated
210
+ assert not monitor.get_termination_event().is_set()
211
+ monitor.destroy()
212
+
213
+
214
+ # ---------------------------------------------------------------------------
215
+ # Callback lifecycle (2 tests)
216
+ # ---------------------------------------------------------------------------
217
+
218
+
219
+ class TestCallbackLifecycle:
220
+ async def test_unsubscribe_removes_callback(self):
221
+ callback = MagicMock()
222
+ monitor = _make_monitor()
223
+ unsubscribe = monitor.subscribe(callback)
224
+ unsubscribe()
225
+ monitor.detect_termination("reason")
226
+ callback.assert_not_called()
227
+
228
+ async def test_callbacks_fire_in_registration_order(self):
229
+ order = []
230
+ monitor = _make_monitor()
231
+ monitor.subscribe(lambda: order.append(1))
232
+ monitor.subscribe(lambda: order.append(2))
233
+ monitor.subscribe(lambda: order.append(3))
234
+ monitor.detect_termination("reason")
235
+ assert order == [1, 2, 3]
236
+
237
+ async def test_bad_callback_does_not_block_later_callbacks(self):
238
+ order = []
239
+ monitor = _make_monitor()
240
+
241
+ def failing_callback():
242
+ order.append("bad")
243
+ raise RuntimeError("callback failed")
244
+
245
+ monitor.subscribe(failing_callback)
246
+ monitor.subscribe(lambda: order.append("good"))
247
+
248
+ monitor.detect_termination("reason")
249
+
250
+ assert order == ["bad", "good"]
251
+ assert monitor.get_termination_event().is_set()
252
+
253
+ async def test_destroy_clears_callbacks(self):
254
+ monitor = _make_monitor()
255
+ monitor.subscribe(MagicMock())
256
+
257
+ monitor.destroy()
258
+
259
+ assert monitor._callbacks == []
@@ -53,6 +53,9 @@ class _StubHttpClient:
53
53
  async def __aexit__(self, exc_type, exc, tb):
54
54
  return None
55
55
 
56
+ async def validate_token(self):
57
+ return {"status": "success"}
58
+
56
59
 
57
60
  def _make_client_config() -> PrefactorCoreConfig:
58
61
  return PrefactorCoreConfig(
@@ -102,3 +105,70 @@ async def test_finish_forwards_explicit_status(status):
102
105
 
103
106
  assert len(stub_http.agent_instances.finish_calls) == 1
104
107
  assert stub_http.agent_instances.finish_calls[0]["status"] == status
108
+
109
+
110
+ class TestFinishAgentInstance409Handling:
111
+ async def test_409_on_finish_treated_as_success(self):
112
+ """FINISH_AGENT_INSTANCE with 409 response should not raise."""
113
+ from datetime import datetime, timezone
114
+ from unittest.mock import AsyncMock, MagicMock
115
+
116
+ from prefactor_core.operations import Operation, OperationType
117
+ from prefactor_http.exceptions import PrefactorApiError
118
+
119
+ mock_http = MagicMock()
120
+ mock_http.agent_instances = MagicMock()
121
+ mock_http.agent_instances.finish = AsyncMock(
122
+ side_effect=PrefactorApiError("already terminated", "conflict", 409)
123
+ )
124
+ mock_http.agent_spans = MagicMock()
125
+ mock_http.agent_spans.create = AsyncMock()
126
+ mock_http.agent_spans.finish = AsyncMock()
127
+
128
+ config = PrefactorCoreConfig(
129
+ http_config=HttpClientConfig(api_url="http://fake", api_token="tok")
130
+ )
131
+ client = PrefactorCoreClient(config)
132
+ client._http = mock_http
133
+ client._initialized = True
134
+
135
+ op = Operation(
136
+ type=OperationType.FINISH_AGENT_INSTANCE,
137
+ payload={
138
+ "instance_id": "inst-1",
139
+ "idempotency_key": "key-1",
140
+ "status": "complete",
141
+ },
142
+ timestamp=datetime.now(timezone.utc),
143
+ )
144
+ # Should not raise
145
+ await client._process_operation(op)
146
+
147
+
148
+ class TestAgentInstanceHandleFinishResetsMonitor:
149
+ async def test_finish_resets_termination_monitor(self):
150
+ """handle.finish() should reset the termination monitor before enqueueing."""
151
+ from unittest.mock import AsyncMock, MagicMock
152
+
153
+ from prefactor_core.managers.agent_instance import AgentInstanceHandle
154
+ from prefactor_core.monitoring.termination_monitor import TerminationMonitor
155
+
156
+ fetch = AsyncMock(
157
+ return_value=MagicMock(status="active", termination_reason=None)
158
+ )
159
+ monitor = TerminationMonitor(fetch_instance=fetch)
160
+ monitor.detect_termination("run 1")
161
+ assert monitor.get_termination_event().is_set()
162
+
163
+ mock_client = MagicMock()
164
+ mock_client._termination_monitor = monitor
165
+ mock_client.instance_manager = MagicMock()
166
+ mock_client.instance_manager.finish_with_idempotency_key = AsyncMock()
167
+
168
+ handle = AgentInstanceHandle(instance_id="inst-1", client=mock_client)
169
+ await handle.finish()
170
+
171
+ # Monitor should be reset (new unset event)
172
+ assert not monitor.get_termination_event().is_set()
173
+ # finish_with_idempotency_key should still be called
174
+ mock_client.instance_manager.finish_with_idempotency_key.assert_called_once()
@@ -56,6 +56,9 @@ class _StubHttpClient:
56
56
  async def __aexit__(self, exc_type, exc, tb):
57
57
  return None
58
58
 
59
+ async def validate_token(self):
60
+ return {"status": "success"}
61
+
59
62
 
60
63
  @pytest.mark.asyncio
61
64
  async def test_create_agent_instance_passes_environment_id_to_register():
@@ -3,8 +3,9 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  import asyncio
6
+ import logging
6
7
  from types import SimpleNamespace
7
- from unittest.mock import patch
8
+ from unittest.mock import Mock, patch
8
9
 
9
10
  import aiohttp
10
11
  import pytest
@@ -67,17 +68,28 @@ class _StubHttpClient:
67
68
  *_args,
68
69
  agent_instances: _StubAgentInstances | None = None,
69
70
  agent_spans: _StubAgentSpans | None = None,
71
+ validate_token_error: PrefactorAuthError | None = None,
70
72
  **_kwargs,
71
73
  ) -> None:
72
74
  self.agent_instances = agent_instances or _StubAgentInstances()
73
75
  self.agent_spans = agent_spans or _StubAgentSpans()
76
+ self._validate_token_error = validate_token_error
77
+ self.enter_calls = 0
78
+ self.exit_calls = 0
74
79
 
75
80
  async def __aenter__(self):
81
+ self.enter_calls += 1
76
82
  return self
77
83
 
78
84
  async def __aexit__(self, exc_type, exc, tb):
85
+ self.exit_calls += 1
79
86
  return None
80
87
 
88
+ async def validate_token(self):
89
+ if self._validate_token_error is not None:
90
+ raise self._validate_token_error
91
+ return {"status": "success"}
92
+
81
93
 
82
94
  def _make_client_config(max_retries: int = 0) -> PrefactorCoreConfig:
83
95
  return PrefactorCoreConfig(
@@ -103,6 +115,60 @@ async def _wait_until(
103
115
  raise AssertionError("Timed out waiting for expected condition")
104
116
 
105
117
 
118
+ @pytest.mark.asyncio
119
+ async def test_termination_sync_loop_survives_iteration_failures(caplog):
120
+ """A failed sync iteration should be logged without stopping future syncs."""
121
+ client = PrefactorCoreClient(_make_client_config())
122
+ monitor = Mock()
123
+ calls = 0
124
+
125
+ def sync(_instance_id):
126
+ nonlocal calls
127
+ calls += 1
128
+ if calls == 1:
129
+ raise RuntimeError("sync failed")
130
+
131
+ monitor.sync.side_effect = sync
132
+ client._termination_monitor = monitor
133
+ original_sleep = asyncio.sleep
134
+
135
+ async def fast_sleep(_delay):
136
+ await original_sleep(0)
137
+
138
+ with (
139
+ caplog.at_level(logging.ERROR, logger="prefactor_core.client"),
140
+ patch("prefactor_core.client.asyncio.sleep", side_effect=fast_sleep),
141
+ ):
142
+ task = asyncio.create_task(client._run_sync_loop())
143
+ try:
144
+ await _wait_until(lambda: monitor.sync.call_count >= 2)
145
+ finally:
146
+ task.cancel()
147
+ with pytest.raises(asyncio.CancelledError):
148
+ await task
149
+
150
+ assert "Termination sync iteration failed" in caplog.text
151
+
152
+
153
+ @pytest.mark.asyncio
154
+ async def test_close_observes_completed_sync_task_exceptions(caplog):
155
+ """close() should await already-failed sync tasks and clear the task handle."""
156
+ client = PrefactorCoreClient(_make_client_config())
157
+ client._initialized = True
158
+
159
+ async def fail_sync_loop():
160
+ raise RuntimeError("sync task failed")
161
+
162
+ client._sync_task = asyncio.create_task(fail_sync_loop())
163
+ await asyncio.sleep(0)
164
+
165
+ with caplog.at_level(logging.ERROR, logger="prefactor_core.client"):
166
+ await client.close()
167
+
168
+ assert client._sync_task is None
169
+ assert "Termination sync loop exited with error during close()" in caplog.text
170
+
171
+
106
172
  @pytest.mark.asyncio
107
173
  async def test_permanent_worker_failure_latches_and_rejects_future_operations():
108
174
  """Permanent failures should latch and reject later queued operations."""
@@ -320,3 +386,23 @@ async def test_latched_failure_drops_already_queued_backlog():
320
386
  await client.close()
321
387
 
322
388
  assert exc_info.value.dropped_operations == 1
389
+
390
+
391
+ @pytest.mark.asyncio
392
+ async def test_initialize_raises_auth_error_when_token_validation_fails():
393
+ """Invalid tokens should fail during initialize() before workers start."""
394
+ stub_http = _StubHttpClient(
395
+ validate_token_error=PrefactorAuthError("Token expired", "bad_authtoken", 401)
396
+ )
397
+
398
+ with patch("prefactor_core.client.PrefactorHttpClient", return_value=stub_http):
399
+ client = PrefactorCoreClient(_make_client_config())
400
+ with pytest.raises(PrefactorAuthError) as exc_info:
401
+ await client.initialize()
402
+
403
+ assert exc_info.value.code == "bad_authtoken"
404
+ assert stub_http.enter_calls == 1
405
+ assert stub_http.exit_calls == 1
406
+ assert client._http is None
407
+ assert client._executor is None
408
+ await client.close()
@@ -37,6 +37,10 @@ class TestPrefactorCoreSdkHeader:
37
37
  "prefactor_http.client.PrefactorHttpClient.__aexit__",
38
38
  AsyncMock(return_value=None),
39
39
  ),
40
+ patch(
41
+ "prefactor_http.client.PrefactorHttpClient.validate_token",
42
+ AsyncMock(return_value={"status": "success"}),
43
+ ),
40
44
  ):
41
45
  await client.initialize()
42
46
  assert client._http is not None
@@ -59,6 +63,10 @@ class TestPrefactorCoreSdkHeader:
59
63
  "prefactor_http.client.PrefactorHttpClient.__aexit__",
60
64
  AsyncMock(return_value=None),
61
65
  ),
66
+ patch(
67
+ "prefactor_http.client.PrefactorHttpClient.validate_token",
68
+ AsyncMock(return_value={"status": "success"}),
69
+ ),
62
70
  ):
63
71
  await client.initialize()
64
72
  assert client._http is not None
File without changes