prefactor-core 0.2.4__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.4 → prefactor_core-0.2.6}/.gitignore +1 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/PKG-INFO +8 -4
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/README.md +6 -2
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/examples/agent_e2e.py +24 -6
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/pyproject.toml +1 -1
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/__init__.py +2 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/_version.py +1 -1
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/client.py +91 -10
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/exceptions.py +18 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/managers/agent_instance.py +13 -3
- 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.4 → prefactor_core-0.2.6}/tests/test_agent_instance_finish_status.py +67 -0
- prefactor_core-0.2.6/tests/test_agent_instance_register.py +123 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/tests/test_failure_handling.py +56 -1
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/tests/test_sdk_header.py +2 -2
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/config.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/context_stack.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/managers/__init__.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/managers/span.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/models.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/operations.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/queue/__init__.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/queue/base.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/queue/executor.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/queue/memory.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/schema_registry.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/span_context.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/src/prefactor_core/utils.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/tests/test_client.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/tests/test_imports.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/tests/test_queue.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/tests/test_span_context.py +0 -0
- {prefactor_core-0.2.4 → prefactor_core-0.2.6}/tests/test_span_manager.py +0 -0
- {prefactor_core-0.2.4 → 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
|
|
|
@@ -81,6 +81,10 @@ async def main():
|
|
|
81
81
|
asyncio.run(main())
|
|
82
82
|
```
|
|
83
83
|
|
|
84
|
+
`create_agent_instance()` supports two auth modes:
|
|
85
|
+
- Account-scoped token: pass `agent_id` and usually `environment_id`.
|
|
86
|
+
- Deployment-scoped token: omit `agent_id` and `environment_id`; the API derives both from the token.
|
|
87
|
+
|
|
84
88
|
## API Reference
|
|
85
89
|
|
|
86
90
|
### `PrefactorCoreClient`
|
|
@@ -98,9 +102,9 @@ await client.close()
|
|
|
98
102
|
|
|
99
103
|
```python
|
|
100
104
|
handle = await client.create_agent_instance(
|
|
101
|
-
agent_id="my-agent",
|
|
102
105
|
agent_version={"name": "My Agent", "external_identifier": "v1.0.0"},
|
|
103
106
|
agent_schema_version=None, # Optional: auto-generated if schema_registry is configured
|
|
107
|
+
agent_id="my-agent", # Optional for deployment-scoped tokens
|
|
104
108
|
external_schema_version_id=None, # Optional: reference an existing schema version
|
|
105
109
|
) -> AgentInstanceHandle
|
|
106
110
|
```
|
|
@@ -240,8 +244,8 @@ config = PrefactorCoreConfig(
|
|
|
240
244
|
async with PrefactorCoreClient(config) as client:
|
|
241
245
|
# agent_schema_version is generated automatically from the registry
|
|
242
246
|
instance = await client.create_agent_instance(
|
|
243
|
-
agent_id="my-agent",
|
|
244
247
|
agent_version={"name": "My Agent", "external_identifier": "v1.0.0"},
|
|
248
|
+
agent_id="my-agent", # Optional for deployment-scoped tokens
|
|
245
249
|
)
|
|
246
250
|
```
|
|
247
251
|
|
|
@@ -70,6 +70,10 @@ async def main():
|
|
|
70
70
|
asyncio.run(main())
|
|
71
71
|
```
|
|
72
72
|
|
|
73
|
+
`create_agent_instance()` supports two auth modes:
|
|
74
|
+
- Account-scoped token: pass `agent_id` and usually `environment_id`.
|
|
75
|
+
- Deployment-scoped token: omit `agent_id` and `environment_id`; the API derives both from the token.
|
|
76
|
+
|
|
73
77
|
## API Reference
|
|
74
78
|
|
|
75
79
|
### `PrefactorCoreClient`
|
|
@@ -87,9 +91,9 @@ await client.close()
|
|
|
87
91
|
|
|
88
92
|
```python
|
|
89
93
|
handle = await client.create_agent_instance(
|
|
90
|
-
agent_id="my-agent",
|
|
91
94
|
agent_version={"name": "My Agent", "external_identifier": "v1.0.0"},
|
|
92
95
|
agent_schema_version=None, # Optional: auto-generated if schema_registry is configured
|
|
96
|
+
agent_id="my-agent", # Optional for deployment-scoped tokens
|
|
93
97
|
external_schema_version_id=None, # Optional: reference an existing schema version
|
|
94
98
|
) -> AgentInstanceHandle
|
|
95
99
|
```
|
|
@@ -229,8 +233,8 @@ config = PrefactorCoreConfig(
|
|
|
229
233
|
async with PrefactorCoreClient(config) as client:
|
|
230
234
|
# agent_schema_version is generated automatically from the registry
|
|
231
235
|
instance = await client.create_agent_instance(
|
|
232
|
-
agent_id="my-agent",
|
|
233
236
|
agent_version={"name": "My Agent", "external_identifier": "v1.0.0"},
|
|
237
|
+
agent_id="my-agent", # Optional for deployment-scoped tokens
|
|
234
238
|
)
|
|
235
239
|
```
|
|
236
240
|
|
|
@@ -24,10 +24,15 @@ Run via mise (env vars set automatically):
|
|
|
24
24
|
Or manually:
|
|
25
25
|
PREFACTOR_API_URL=https://api.prefactor.ai \\
|
|
26
26
|
PREFACTOR_API_TOKEN=your-token \\
|
|
27
|
-
PREFACTOR_AGENT_ID=your-agent-id \\
|
|
28
27
|
python packages/core/examples/agent_e2e.py
|
|
28
|
+
|
|
29
|
+
Set `PREFACTOR_AGENT_ID` only when using an account-scoped token. With a
|
|
30
|
+
deployment-scoped token, leave it unset and the API derives the agent and
|
|
31
|
+
environment from the token.
|
|
29
32
|
"""
|
|
30
33
|
|
|
34
|
+
from __future__ import annotations
|
|
35
|
+
|
|
31
36
|
import asyncio
|
|
32
37
|
import os
|
|
33
38
|
|
|
@@ -182,7 +187,9 @@ async def simulate_retrieval(query: str) -> list[str]:
|
|
|
182
187
|
|
|
183
188
|
|
|
184
189
|
async def main() -> None:
|
|
185
|
-
agent_id = os.environ
|
|
190
|
+
agent_id = os.environ.get("PREFACTOR_AGENT_ID")
|
|
191
|
+
if agent_id is not None:
|
|
192
|
+
agent_id = agent_id.strip() or None
|
|
186
193
|
|
|
187
194
|
config = PrefactorCoreConfig(
|
|
188
195
|
http_config=HttpClientConfig(
|
|
@@ -194,10 +201,21 @@ async def main() -> None:
|
|
|
194
201
|
|
|
195
202
|
async with PrefactorCoreClient(config) as client:
|
|
196
203
|
# Register the agent instance — schema version derived from registry.
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
204
|
+
if agent_id is None:
|
|
205
|
+
instance = await client.create_agent_instance(
|
|
206
|
+
agent_version={
|
|
207
|
+
"name": "Example Agent",
|
|
208
|
+
"external_identifier": "v8.0.0",
|
|
209
|
+
}
|
|
210
|
+
)
|
|
211
|
+
else:
|
|
212
|
+
instance = await client.create_agent_instance(
|
|
213
|
+
agent_id=agent_id,
|
|
214
|
+
agent_version={
|
|
215
|
+
"name": "Example Agent",
|
|
216
|
+
"external_identifier": "v8.0.0",
|
|
217
|
+
},
|
|
218
|
+
)
|
|
201
219
|
|
|
202
220
|
await instance.start()
|
|
203
221
|
|
|
@@ -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."""
|
|
@@ -314,11 +387,12 @@ class PrefactorCoreClient:
|
|
|
314
387
|
|
|
315
388
|
async def create_agent_instance(
|
|
316
389
|
self,
|
|
317
|
-
agent_id: str,
|
|
318
390
|
agent_version: dict[str, Any],
|
|
319
391
|
agent_schema_version: dict[str, Any] | None = None,
|
|
392
|
+
agent_id: str | None = None,
|
|
320
393
|
instance_id: str | None = None,
|
|
321
394
|
external_schema_version_id: str | None = None,
|
|
395
|
+
environment_id: str | None = None,
|
|
322
396
|
) -> "AgentInstanceHandle":
|
|
323
397
|
"""Create a new agent instance.
|
|
324
398
|
|
|
@@ -329,13 +403,14 @@ class PrefactorCoreClient:
|
|
|
329
403
|
the registry's schemas will be used automatically.
|
|
330
404
|
|
|
331
405
|
Args:
|
|
332
|
-
agent_id: ID
|
|
406
|
+
agent_id: Agent ID. Omit when using a deployment-scoped token.
|
|
333
407
|
agent_version: Version information (name, etc.).
|
|
334
408
|
agent_schema_version: Schema version. Uses registry if not provided
|
|
335
409
|
and registry is configured.
|
|
336
410
|
instance_id: Optional custom ID for the instance.
|
|
337
411
|
external_schema_version_id: Optional external identifier for the
|
|
338
412
|
schema version. Defaults to "auto-generated" when using registry.
|
|
413
|
+
environment_id: Optional environment ID used to scope the agent instance.
|
|
339
414
|
|
|
340
415
|
Returns:
|
|
341
416
|
AgentInstanceHandle for the created instance.
|
|
@@ -356,7 +431,10 @@ class PrefactorCoreClient:
|
|
|
356
431
|
# Use registry to generate schema version
|
|
357
432
|
ext_id = external_schema_version_id
|
|
358
433
|
if ext_id is None:
|
|
359
|
-
|
|
434
|
+
if agent_id is None:
|
|
435
|
+
ext_id = f"auto-deployment-{time.time()}"
|
|
436
|
+
else:
|
|
437
|
+
ext_id = f"auto-{agent_id}-{time.time()}"
|
|
360
438
|
final_schema_version = self._config.schema_registry.to_agent_schema_version(
|
|
361
439
|
ext_id
|
|
362
440
|
)
|
|
@@ -373,8 +451,11 @@ class PrefactorCoreClient:
|
|
|
373
451
|
agent_version=agent_version,
|
|
374
452
|
agent_schema_version=final_schema_version,
|
|
375
453
|
instance_id=instance_id,
|
|
454
|
+
environment_id=environment_id,
|
|
376
455
|
)
|
|
377
456
|
|
|
457
|
+
self._set_current_instance(instance_id)
|
|
458
|
+
|
|
378
459
|
return AgentInstanceHandle(
|
|
379
460
|
instance_id=instance_id,
|
|
380
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
|
]
|
|
@@ -59,10 +59,11 @@ class AgentInstanceManager:
|
|
|
59
59
|
|
|
60
60
|
async def register(
|
|
61
61
|
self,
|
|
62
|
-
agent_id: str,
|
|
63
62
|
agent_version: dict[str, Any],
|
|
64
63
|
agent_schema_version: dict[str, Any],
|
|
64
|
+
agent_id: str | None = None,
|
|
65
65
|
instance_id: str | None = None,
|
|
66
|
+
environment_id: str | None = None,
|
|
66
67
|
) -> str:
|
|
67
68
|
"""Register a new agent instance.
|
|
68
69
|
|
|
@@ -70,12 +71,14 @@ class AgentInstanceManager:
|
|
|
70
71
|
the API-generated ID.
|
|
71
72
|
|
|
72
73
|
Args:
|
|
73
|
-
agent_id: ID
|
|
74
|
+
agent_id: Agent ID. Omit when using a deployment-scoped token.
|
|
74
75
|
agent_version: Version information (name, external_identifier, etc.).
|
|
75
76
|
agent_schema_version: Schema version information.
|
|
76
77
|
instance_id: Optional ID to forward to the API as ``id``. When
|
|
77
78
|
provided, the API uses it as the instance ID; when omitted,
|
|
78
79
|
the API generates one.
|
|
80
|
+
environment_id: Optional environment ID. Required when using an
|
|
81
|
+
account-scoped token; omit when using a deployment-scoped token.
|
|
79
82
|
|
|
80
83
|
Returns:
|
|
81
84
|
The instance ID (API-generated).
|
|
@@ -84,6 +87,7 @@ class AgentInstanceManager:
|
|
|
84
87
|
agent_id=agent_id,
|
|
85
88
|
agent_version=agent_version,
|
|
86
89
|
agent_schema_version=agent_schema_version,
|
|
90
|
+
environment_id=environment_id,
|
|
87
91
|
id=instance_id,
|
|
88
92
|
idempotency_key=generate_idempotency_key(),
|
|
89
93
|
)
|
|
@@ -214,12 +218,18 @@ class AgentInstanceHandle:
|
|
|
214
218
|
async def finish(self, status: "FinishStatus" = "complete") -> None:
|
|
215
219
|
"""Mark the instance as finished.
|
|
216
220
|
|
|
217
|
-
|
|
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.
|
|
218
224
|
|
|
219
225
|
Args:
|
|
220
226
|
status: Terminal status for the instance — one of ``"complete"``,
|
|
221
227
|
``"failed"``, or ``"cancelled"``. Defaults to ``"complete"``.
|
|
222
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)
|
|
223
233
|
manager = self._client.instance_manager
|
|
224
234
|
assert manager is not None
|
|
225
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()
|
|
@@ -0,0 +1,123 @@
|
|
|
1
|
+
"""Tests for environment_id threading through the core register flow."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from types import SimpleNamespace
|
|
6
|
+
from unittest.mock import patch
|
|
7
|
+
|
|
8
|
+
import pytest
|
|
9
|
+
from prefactor_core import PrefactorCoreClient
|
|
10
|
+
from prefactor_core.config import PrefactorCoreConfig, QueueConfig
|
|
11
|
+
from prefactor_http.config import HttpClientConfig
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def _make_config() -> PrefactorCoreConfig:
|
|
15
|
+
return PrefactorCoreConfig(
|
|
16
|
+
http_config=HttpClientConfig(
|
|
17
|
+
api_url="https://api.test.com",
|
|
18
|
+
api_token="test-token",
|
|
19
|
+
max_retries=0,
|
|
20
|
+
),
|
|
21
|
+
queue_config=QueueConfig(num_workers=1, max_retries=0),
|
|
22
|
+
)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
class _StubAgentInstances:
|
|
26
|
+
def __init__(self) -> None:
|
|
27
|
+
self.register_calls: list[dict] = []
|
|
28
|
+
|
|
29
|
+
async def register(self, **kwargs):
|
|
30
|
+
self.register_calls.append(kwargs)
|
|
31
|
+
return SimpleNamespace(id="inst-1")
|
|
32
|
+
|
|
33
|
+
async def start(self, **kwargs):
|
|
34
|
+
return SimpleNamespace(id=kwargs["agent_instance_id"])
|
|
35
|
+
|
|
36
|
+
async def finish(self, **kwargs):
|
|
37
|
+
return SimpleNamespace(id=kwargs["agent_instance_id"])
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class _StubAgentSpans:
|
|
41
|
+
async def create(self, **kwargs):
|
|
42
|
+
return SimpleNamespace(id="span-1")
|
|
43
|
+
|
|
44
|
+
async def finish(self, **kwargs):
|
|
45
|
+
return SimpleNamespace(id=kwargs["agent_span_id"])
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
class _StubHttpClient:
|
|
49
|
+
def __init__(self, *args, **kwargs) -> None:
|
|
50
|
+
self.agent_instances = _StubAgentInstances()
|
|
51
|
+
self.agent_spans = _StubAgentSpans()
|
|
52
|
+
|
|
53
|
+
async def __aenter__(self):
|
|
54
|
+
return self
|
|
55
|
+
|
|
56
|
+
async def __aexit__(self, exc_type, exc, tb):
|
|
57
|
+
return None
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
@pytest.mark.asyncio
|
|
61
|
+
async def test_create_agent_instance_passes_environment_id_to_register():
|
|
62
|
+
"""environment_id passed to create_agent_instance() reaches the HTTP layer."""
|
|
63
|
+
stub_http = _StubHttpClient()
|
|
64
|
+
with patch("prefactor_core.client.PrefactorHttpClient", return_value=stub_http):
|
|
65
|
+
async with PrefactorCoreClient(_make_config()) as client:
|
|
66
|
+
await client.create_agent_instance(
|
|
67
|
+
agent_id="agent-1",
|
|
68
|
+
agent_version={"name": "v1"},
|
|
69
|
+
agent_schema_version={"span_schemas": {}},
|
|
70
|
+
environment_id="env-1",
|
|
71
|
+
)
|
|
72
|
+
assert stub_http.agent_instances.register_calls[0]["environment_id"] == "env-1"
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
@pytest.mark.asyncio
|
|
76
|
+
async def test_create_agent_instance_without_environment_id_omits_field():
|
|
77
|
+
"""Omitting environment_id does not send it to the HTTP register call."""
|
|
78
|
+
stub_http = _StubHttpClient()
|
|
79
|
+
with patch("prefactor_core.client.PrefactorHttpClient", return_value=stub_http):
|
|
80
|
+
async with PrefactorCoreClient(_make_config()) as client:
|
|
81
|
+
await client.create_agent_instance(
|
|
82
|
+
agent_id="agent-1",
|
|
83
|
+
agent_version={"name": "v1"},
|
|
84
|
+
agent_schema_version={"span_schemas": {}},
|
|
85
|
+
)
|
|
86
|
+
call = stub_http.agent_instances.register_calls[0]
|
|
87
|
+
assert "environment_id" not in call or call["environment_id"] is None
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
@pytest.mark.asyncio
|
|
91
|
+
async def test_create_agent_instance_without_agent_id_omits_field():
|
|
92
|
+
"""Omitting agent_id does not send it to the HTTP register call."""
|
|
93
|
+
stub_http = _StubHttpClient()
|
|
94
|
+
with patch("prefactor_core.client.PrefactorHttpClient", return_value=stub_http):
|
|
95
|
+
async with PrefactorCoreClient(_make_config()) as client:
|
|
96
|
+
await client.create_agent_instance(
|
|
97
|
+
agent_version={"name": "v1"},
|
|
98
|
+
agent_schema_version={"span_schemas": {}},
|
|
99
|
+
)
|
|
100
|
+
call = stub_http.agent_instances.register_calls[0]
|
|
101
|
+
assert "agent_id" not in call or call["agent_id"] is None
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
@pytest.mark.asyncio
|
|
105
|
+
async def test_create_agent_instance_without_agent_id_uses_deployment_schema_fallback():
|
|
106
|
+
"""Registry fallback should not require agent_id to generate an external ID."""
|
|
107
|
+
stub_http = _StubHttpClient()
|
|
108
|
+
config = _make_config()
|
|
109
|
+
config.schema_registry = SimpleNamespace(
|
|
110
|
+
to_agent_schema_version=lambda external_id: {"external_identifier": external_id}
|
|
111
|
+
)
|
|
112
|
+
|
|
113
|
+
with patch("prefactor_core.client.PrefactorHttpClient", return_value=stub_http):
|
|
114
|
+
async with PrefactorCoreClient(config) as client:
|
|
115
|
+
await client.create_agent_instance(
|
|
116
|
+
agent_version={"name": "v1"},
|
|
117
|
+
agent_schema_version=None,
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
call = stub_http.agent_instances.register_calls[0]
|
|
121
|
+
assert call["agent_schema_version"]["external_identifier"].startswith(
|
|
122
|
+
"auto-deployment-"
|
|
123
|
+
)
|
|
@@ -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."""
|
|
@@ -47,7 +47,7 @@ class TestPrefactorCoreSdkHeader:
|
|
|
47
47
|
"""Core prepends the adapter entry ahead of the core entry."""
|
|
48
48
|
client = PrefactorCoreClient(
|
|
49
49
|
create_config(),
|
|
50
|
-
sdk_header_entry="prefactor-langchain@0.2.
|
|
50
|
+
sdk_header_entry="prefactor-langchain@0.2.5",
|
|
51
51
|
)
|
|
52
52
|
|
|
53
53
|
with (
|
|
@@ -63,7 +63,7 @@ class TestPrefactorCoreSdkHeader:
|
|
|
63
63
|
await client.initialize()
|
|
64
64
|
assert client._http is not None
|
|
65
65
|
assert client._http._sdk_header == (
|
|
66
|
-
f"prefactor-langchain@0.2.
|
|
66
|
+
f"prefactor-langchain@0.2.5 {CORE_SDK_HEADER_ENTRY}"
|
|
67
67
|
)
|
|
68
68
|
await client.close()
|
|
69
69
|
|
|
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
|