prefactor-core 0.2.5__tar.gz → 0.2.6__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.
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/PKG-INFO +2 -2
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/examples/agent_e2e.py +4 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/pyproject.toml +1 -1
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/__init__.py +2 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/_version.py +1 -1
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/client.py +82 -7
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/exceptions.py +18 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/managers/agent_instance.py +7 -1
- prefactor_core-0.2.6/src/prefactor_core/monitoring/__init__.py +5 -0
- prefactor_core-0.2.6/src/prefactor_core/monitoring/termination_monitor.py +151 -0
- prefactor_core-0.2.6/tests/monitoring/__init__.py +0 -0
- prefactor_core-0.2.6/tests/monitoring/test_termination_monitor.py +259 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_agent_instance_finish_status.py +67 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_failure_handling.py +56 -1
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/.gitignore +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/README.md +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/config.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/context_stack.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/managers/__init__.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/managers/span.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/models.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/operations.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/queue/__init__.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/queue/base.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/queue/executor.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/queue/memory.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/schema_registry.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/span_context.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/src/prefactor_core/utils.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_agent_instance_register.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_client.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_imports.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_queue.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_sdk_header.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_span_context.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_span_manager.py +0 -0
- {prefactor_core-0.2.5 → prefactor_core-0.2.6}/tests/test_utils.py +0 -0
|
@@ -1,11 +1,11 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: prefactor-core
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.6
|
|
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.
|
|
8
|
+
Requires-Dist: prefactor-http>=0.1.4
|
|
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(
|
|
@@ -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",
|
|
@@ -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
|
|
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."""
|
|
@@ -155,6 +164,13 @@ class PrefactorCoreClient:
|
|
|
155
164
|
|
|
156
165
|
self._initialized = True
|
|
157
166
|
|
|
167
|
+
self._termination_monitor = TerminationMonitor(
|
|
168
|
+
fetch_instance=self._fetch_instance_for_poll,
|
|
169
|
+
)
|
|
170
|
+
self._sync_task = asyncio.create_task(
|
|
171
|
+
self._run_sync_loop(), name="prefactor-termination-sync"
|
|
172
|
+
)
|
|
173
|
+
|
|
158
174
|
async def close(self) -> None:
|
|
159
175
|
"""Close the client and cleanup resources.
|
|
160
176
|
|
|
@@ -164,6 +180,24 @@ class PrefactorCoreClient:
|
|
|
164
180
|
if not self._initialized:
|
|
165
181
|
return
|
|
166
182
|
|
|
183
|
+
if self._sync_task is not None:
|
|
184
|
+
if not self._sync_task.done():
|
|
185
|
+
self._sync_task.cancel()
|
|
186
|
+
try:
|
|
187
|
+
await self._sync_task
|
|
188
|
+
except asyncio.CancelledError:
|
|
189
|
+
pass
|
|
190
|
+
except Exception:
|
|
191
|
+
logger.exception(
|
|
192
|
+
"Termination sync loop exited with error during close()"
|
|
193
|
+
)
|
|
194
|
+
finally:
|
|
195
|
+
self._sync_task = None
|
|
196
|
+
|
|
197
|
+
if self._termination_monitor is not None:
|
|
198
|
+
self._termination_monitor.destroy()
|
|
199
|
+
self._termination_monitor = None
|
|
200
|
+
|
|
167
201
|
# Stop executor
|
|
168
202
|
if self._executor:
|
|
169
203
|
await self._executor.stop()
|
|
@@ -272,12 +306,23 @@ class PrefactorCoreClient:
|
|
|
272
306
|
)
|
|
273
307
|
|
|
274
308
|
elif operation.type == OperationType.FINISH_AGENT_INSTANCE:
|
|
275
|
-
|
|
276
|
-
|
|
277
|
-
|
|
278
|
-
|
|
279
|
-
|
|
280
|
-
|
|
309
|
+
try:
|
|
310
|
+
await self._http.agent_instances.finish(
|
|
311
|
+
agent_instance_id=operation.payload["instance_id"],
|
|
312
|
+
status=operation.payload.get("status", "complete"),
|
|
313
|
+
timestamp=operation.timestamp,
|
|
314
|
+
idempotency_key=operation.payload.get("idempotency_key"),
|
|
315
|
+
)
|
|
316
|
+
except PrefactorApiError as finish_err:
|
|
317
|
+
if finish_err.status_code == 409:
|
|
318
|
+
logger.debug(
|
|
319
|
+
"[prefactor:http] Agent instance %s already in"
|
|
320
|
+
" terminal state; skipping finish.",
|
|
321
|
+
operation.payload["instance_id"],
|
|
322
|
+
)
|
|
323
|
+
return
|
|
324
|
+
raise
|
|
325
|
+
|
|
281
326
|
elif operation.type == OperationType.CREATE_SPAN:
|
|
282
327
|
await self._http.agent_spans.create(
|
|
283
328
|
agent_instance_id=operation.payload["instance_id"],
|
|
@@ -286,6 +331,7 @@ class PrefactorCoreClient:
|
|
|
286
331
|
id=operation.payload.get("span_id"),
|
|
287
332
|
parent_span_id=operation.payload.get("parent_span_id"),
|
|
288
333
|
payload=operation.payload.get("payload"),
|
|
334
|
+
control_signal_callback=self._on_control_signal,
|
|
289
335
|
)
|
|
290
336
|
|
|
291
337
|
elif operation.type == OperationType.FINISH_SPAN:
|
|
@@ -295,6 +341,7 @@ class PrefactorCoreClient:
|
|
|
295
341
|
result_payload=operation.payload.get("result_payload"),
|
|
296
342
|
timestamp=operation.timestamp,
|
|
297
343
|
idempotency_key=operation.payload.get("idempotency_key"),
|
|
344
|
+
control_signal_callback=self._on_control_signal,
|
|
298
345
|
)
|
|
299
346
|
|
|
300
347
|
except Exception as e:
|
|
@@ -307,6 +354,32 @@ class PrefactorCoreClient:
|
|
|
307
354
|
)
|
|
308
355
|
raise
|
|
309
356
|
|
|
357
|
+
async def _fetch_instance_for_poll(self, instance_id: str):
|
|
358
|
+
if self._http is None:
|
|
359
|
+
return None
|
|
360
|
+
return await self._http.agent_instances.get(instance_id)
|
|
361
|
+
|
|
362
|
+
async def _run_sync_loop(self) -> None:
|
|
363
|
+
while True:
|
|
364
|
+
await asyncio.sleep(1)
|
|
365
|
+
if self._termination_monitor is None:
|
|
366
|
+
continue
|
|
367
|
+
try:
|
|
368
|
+
self._termination_monitor.sync(self._current_instance_id)
|
|
369
|
+
except Exception:
|
|
370
|
+
logger.exception("Termination sync iteration failed")
|
|
371
|
+
|
|
372
|
+
def _on_control_signal(self, reason: str | None) -> None:
|
|
373
|
+
if self._termination_monitor is not None:
|
|
374
|
+
self._termination_monitor.detect_termination(reason)
|
|
375
|
+
|
|
376
|
+
def _set_current_instance(self, instance_id: str | None) -> None:
|
|
377
|
+
self._current_instance_id = instance_id
|
|
378
|
+
|
|
379
|
+
def _clear_current_instance(self, instance_id: str) -> None:
|
|
380
|
+
if self._current_instance_id == instance_id:
|
|
381
|
+
self._current_instance_id = None
|
|
382
|
+
|
|
310
383
|
@property
|
|
311
384
|
def instance_manager(self) -> AgentInstanceManager | None:
|
|
312
385
|
"""Public accessor for the agent instance manager."""
|
|
@@ -381,6 +454,8 @@ class PrefactorCoreClient:
|
|
|
381
454
|
environment_id=environment_id,
|
|
382
455
|
)
|
|
383
456
|
|
|
457
|
+
self._set_current_instance(instance_id)
|
|
458
|
+
|
|
384
459
|
return AgentInstanceHandle(
|
|
385
460
|
instance_id=instance_id,
|
|
386
461
|
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
|
-
|
|
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,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 == []
|
|
@@ -102,3 +102,70 @@ async def test_finish_forwards_explicit_status(status):
|
|
|
102
102
|
|
|
103
103
|
assert len(stub_http.agent_instances.finish_calls) == 1
|
|
104
104
|
assert stub_http.agent_instances.finish_calls[0]["status"] == status
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
class TestFinishAgentInstance409Handling:
|
|
108
|
+
async def test_409_on_finish_treated_as_success(self):
|
|
109
|
+
"""FINISH_AGENT_INSTANCE with 409 response should not raise."""
|
|
110
|
+
from datetime import datetime, timezone
|
|
111
|
+
from unittest.mock import AsyncMock, MagicMock
|
|
112
|
+
|
|
113
|
+
from prefactor_core.operations import Operation, OperationType
|
|
114
|
+
from prefactor_http.exceptions import PrefactorApiError
|
|
115
|
+
|
|
116
|
+
mock_http = MagicMock()
|
|
117
|
+
mock_http.agent_instances = MagicMock()
|
|
118
|
+
mock_http.agent_instances.finish = AsyncMock(
|
|
119
|
+
side_effect=PrefactorApiError("already terminated", "conflict", 409)
|
|
120
|
+
)
|
|
121
|
+
mock_http.agent_spans = MagicMock()
|
|
122
|
+
mock_http.agent_spans.create = AsyncMock()
|
|
123
|
+
mock_http.agent_spans.finish = AsyncMock()
|
|
124
|
+
|
|
125
|
+
config = PrefactorCoreConfig(
|
|
126
|
+
http_config=HttpClientConfig(api_url="http://fake", api_token="tok")
|
|
127
|
+
)
|
|
128
|
+
client = PrefactorCoreClient(config)
|
|
129
|
+
client._http = mock_http
|
|
130
|
+
client._initialized = True
|
|
131
|
+
|
|
132
|
+
op = Operation(
|
|
133
|
+
type=OperationType.FINISH_AGENT_INSTANCE,
|
|
134
|
+
payload={
|
|
135
|
+
"instance_id": "inst-1",
|
|
136
|
+
"idempotency_key": "key-1",
|
|
137
|
+
"status": "complete",
|
|
138
|
+
},
|
|
139
|
+
timestamp=datetime.now(timezone.utc),
|
|
140
|
+
)
|
|
141
|
+
# Should not raise
|
|
142
|
+
await client._process_operation(op)
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
class TestAgentInstanceHandleFinishResetsMonitor:
|
|
146
|
+
async def test_finish_resets_termination_monitor(self):
|
|
147
|
+
"""handle.finish() should reset the termination monitor before enqueueing."""
|
|
148
|
+
from unittest.mock import AsyncMock, MagicMock
|
|
149
|
+
|
|
150
|
+
from prefactor_core.managers.agent_instance import AgentInstanceHandle
|
|
151
|
+
from prefactor_core.monitoring.termination_monitor import TerminationMonitor
|
|
152
|
+
|
|
153
|
+
fetch = AsyncMock(
|
|
154
|
+
return_value=MagicMock(status="active", termination_reason=None)
|
|
155
|
+
)
|
|
156
|
+
monitor = TerminationMonitor(fetch_instance=fetch)
|
|
157
|
+
monitor.detect_termination("run 1")
|
|
158
|
+
assert monitor.get_termination_event().is_set()
|
|
159
|
+
|
|
160
|
+
mock_client = MagicMock()
|
|
161
|
+
mock_client._termination_monitor = monitor
|
|
162
|
+
mock_client.instance_manager = MagicMock()
|
|
163
|
+
mock_client.instance_manager.finish_with_idempotency_key = AsyncMock()
|
|
164
|
+
|
|
165
|
+
handle = AgentInstanceHandle(instance_id="inst-1", client=mock_client)
|
|
166
|
+
await handle.finish()
|
|
167
|
+
|
|
168
|
+
# Monitor should be reset (new unset event)
|
|
169
|
+
assert not monitor.get_termination_event().is_set()
|
|
170
|
+
# finish_with_idempotency_key should still be called
|
|
171
|
+
mock_client.instance_manager.finish_with_idempotency_key.assert_called_once()
|
|
@@ -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
|
|
@@ -103,6 +104,60 @@ async def _wait_until(
|
|
|
103
104
|
raise AssertionError("Timed out waiting for expected condition")
|
|
104
105
|
|
|
105
106
|
|
|
107
|
+
@pytest.mark.asyncio
|
|
108
|
+
async def test_termination_sync_loop_survives_iteration_failures(caplog):
|
|
109
|
+
"""A failed sync iteration should be logged without stopping future syncs."""
|
|
110
|
+
client = PrefactorCoreClient(_make_client_config())
|
|
111
|
+
monitor = Mock()
|
|
112
|
+
calls = 0
|
|
113
|
+
|
|
114
|
+
def sync(_instance_id):
|
|
115
|
+
nonlocal calls
|
|
116
|
+
calls += 1
|
|
117
|
+
if calls == 1:
|
|
118
|
+
raise RuntimeError("sync failed")
|
|
119
|
+
|
|
120
|
+
monitor.sync.side_effect = sync
|
|
121
|
+
client._termination_monitor = monitor
|
|
122
|
+
original_sleep = asyncio.sleep
|
|
123
|
+
|
|
124
|
+
async def fast_sleep(_delay):
|
|
125
|
+
await original_sleep(0)
|
|
126
|
+
|
|
127
|
+
with (
|
|
128
|
+
caplog.at_level(logging.ERROR, logger="prefactor_core.client"),
|
|
129
|
+
patch("prefactor_core.client.asyncio.sleep", side_effect=fast_sleep),
|
|
130
|
+
):
|
|
131
|
+
task = asyncio.create_task(client._run_sync_loop())
|
|
132
|
+
try:
|
|
133
|
+
await _wait_until(lambda: monitor.sync.call_count >= 2)
|
|
134
|
+
finally:
|
|
135
|
+
task.cancel()
|
|
136
|
+
with pytest.raises(asyncio.CancelledError):
|
|
137
|
+
await task
|
|
138
|
+
|
|
139
|
+
assert "Termination sync iteration failed" in caplog.text
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
@pytest.mark.asyncio
|
|
143
|
+
async def test_close_observes_completed_sync_task_exceptions(caplog):
|
|
144
|
+
"""close() should await already-failed sync tasks and clear the task handle."""
|
|
145
|
+
client = PrefactorCoreClient(_make_client_config())
|
|
146
|
+
client._initialized = True
|
|
147
|
+
|
|
148
|
+
async def fail_sync_loop():
|
|
149
|
+
raise RuntimeError("sync task failed")
|
|
150
|
+
|
|
151
|
+
client._sync_task = asyncio.create_task(fail_sync_loop())
|
|
152
|
+
await asyncio.sleep(0)
|
|
153
|
+
|
|
154
|
+
with caplog.at_level(logging.ERROR, logger="prefactor_core.client"):
|
|
155
|
+
await client.close()
|
|
156
|
+
|
|
157
|
+
assert client._sync_task is None
|
|
158
|
+
assert "Termination sync loop exited with error during close()" in caplog.text
|
|
159
|
+
|
|
160
|
+
|
|
106
161
|
@pytest.mark.asyncio
|
|
107
162
|
async def test_permanent_worker_failure_latches_and_rejects_future_operations():
|
|
108
163
|
"""Permanent failures should latch and reject later queued operations."""
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|