prefactor-core 0.2.6__tar.gz → 0.3.0__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.6 → prefactor_core-0.3.0}/.gitignore +0 -1
  2. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/PKG-INFO +2 -2
  3. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/pyproject.toml +1 -1
  4. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/_version.py +1 -1
  5. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/client.py +49 -1
  6. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/managers/agent_instance.py +57 -2
  7. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/managers/span.py +21 -5
  8. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/models.py +2 -0
  9. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/operations.py +1 -0
  10. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/schema_registry.py +52 -0
  11. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/span_context.py +35 -6
  12. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_agent_instance_finish_status.py +3 -0
  13. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_agent_instance_register.py +3 -0
  14. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_failure_handling.py +31 -0
  15. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_sdk_header.py +8 -0
  16. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/README.md +0 -0
  17. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/examples/agent_e2e.py +0 -0
  18. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/__init__.py +0 -0
  19. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/config.py +0 -0
  20. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/context_stack.py +0 -0
  21. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/exceptions.py +0 -0
  22. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/managers/__init__.py +0 -0
  23. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/monitoring/__init__.py +0 -0
  24. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/monitoring/termination_monitor.py +0 -0
  25. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/__init__.py +0 -0
  26. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/base.py +0 -0
  27. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/executor.py +0 -0
  28. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/queue/memory.py +0 -0
  29. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/src/prefactor_core/utils.py +0 -0
  30. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/monitoring/__init__.py +0 -0
  31. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/monitoring/test_termination_monitor.py +0 -0
  32. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_client.py +0 -0
  33. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_imports.py +0 -0
  34. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_queue.py +0 -0
  35. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_span_context.py +0 -0
  36. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/tests/test_span_manager.py +0 -0
  37. {prefactor_core-0.2.6 → prefactor_core-0.3.0}/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.6
3
+ Version: 0.3.0
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.4
8
+ Requires-Dist: prefactor-http>=0.2.0
9
9
  Requires-Dist: pydantic>=2.0.0
10
10
  Description-Content-Type: text/markdown
11
11
 
@@ -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.4",
12
+ "prefactor-http>=0.2.0",
13
13
  "pydantic>=2.0.0",
14
14
  ]
15
15
 
@@ -3,5 +3,5 @@
3
3
  from __future__ import annotations
4
4
 
5
5
  PACKAGE_NAME = "prefactor-core"
6
- __version__ = "0.2.6"
6
+ __version__ = "0.3.0"
7
7
  PACKAGE_VERSION = __version__
@@ -11,6 +11,7 @@ import asyncio
11
11
  import logging
12
12
  import time
13
13
  from contextlib import asynccontextmanager
14
+ from datetime import datetime
14
15
  from typing import TYPE_CHECKING, Any
15
16
 
16
17
  from prefactor_http.client import PrefactorHttpClient
@@ -19,6 +20,7 @@ from prefactor_http.exceptions import (
19
20
  is_permanent_http_error,
20
21
  is_transient_http_error,
21
22
  )
23
+ from prefactor_http.models.types import InstancePurpose
22
24
 
23
25
  from ._version import PACKAGE_NAME as CORE_PACKAGE_NAME
24
26
  from ._version import PACKAGE_VERSION as CORE_PACKAGE_VERSION
@@ -142,6 +144,14 @@ class PrefactorCoreClient:
142
144
  )
143
145
  await self._http.__aenter__()
144
146
 
147
+ try:
148
+ await self._http.validate_token()
149
+ except BaseException as exc:
150
+ http = self._http
151
+ self._http = None
152
+ await http.__aexit__(type(exc), exc, exc.__traceback__)
153
+ raise
154
+
145
155
  # Initialize executor
146
156
  self._executor = TaskExecutor(
147
157
  queue=self._queue,
@@ -323,6 +333,11 @@ class PrefactorCoreClient:
323
333
  return
324
334
  raise
325
335
 
336
+ elif operation.type == OperationType.UPDATE_AGENT_INSTANCE:
337
+ await self._http.agent_instances.update(
338
+ agent_instance_id=operation.payload["instance_id"],
339
+ quality_payload=operation.payload.get("quality_payload"),
340
+ )
326
341
  elif operation.type == OperationType.CREATE_SPAN:
327
342
  await self._http.agent_spans.create(
328
343
  agent_instance_id=operation.payload["instance_id"],
@@ -331,6 +346,7 @@ class PrefactorCoreClient:
331
346
  id=operation.payload.get("span_id"),
332
347
  parent_span_id=operation.payload.get("parent_span_id"),
333
348
  payload=operation.payload.get("payload"),
349
+ started_at=operation.payload.get("started_at"),
334
350
  control_signal_callback=self._on_control_signal,
335
351
  )
336
352
 
@@ -393,6 +409,7 @@ class PrefactorCoreClient:
393
409
  instance_id: str | None = None,
394
410
  external_schema_version_id: str | None = None,
395
411
  environment_id: str | None = None,
412
+ purpose: InstancePurpose | None = None,
396
413
  ) -> "AgentInstanceHandle":
397
414
  """Create a new agent instance.
398
415
 
@@ -411,6 +428,8 @@ class PrefactorCoreClient:
411
428
  external_schema_version_id: Optional external identifier for the
412
429
  schema version. Defaults to "auto-generated" when using registry.
413
430
  environment_id: Optional environment ID used to scope the agent instance.
431
+ purpose: Why this instance ran — ``"live"``, ``"smoke_test"``,
432
+ or ``"eval"``. Omitted (None) lets the API default to ``"live"``.
414
433
 
415
434
  Returns:
416
435
  AgentInstanceHandle for the created instance.
@@ -452,6 +471,7 @@ class PrefactorCoreClient:
452
471
  agent_schema_version=final_schema_version,
453
472
  instance_id=instance_id,
454
473
  environment_id=environment_id,
474
+ purpose=purpose,
455
475
  )
456
476
 
457
477
  self._set_current_instance(instance_id)
@@ -467,6 +487,7 @@ class PrefactorCoreClient:
467
487
  schema_name: str,
468
488
  parent_span_id: str | None = None,
469
489
  payload: dict[str, Any] | None = None,
490
+ started_at: datetime | None = None,
470
491
  ) -> str:
471
492
  """Create a span and return its ID without finishing it.
472
493
 
@@ -478,6 +499,7 @@ class PrefactorCoreClient:
478
499
  schema_name: Name of the schema for this span.
479
500
  parent_span_id: Optional explicit parent span ID.
480
501
  payload: Optional initial payload (params/inputs) stored on creation.
502
+ started_at: Optional ISO 8601 start time (defaults to current time).
481
503
 
482
504
  Returns:
483
505
  The span ID.
@@ -494,23 +516,49 @@ class PrefactorCoreClient:
494
516
  schema_name=schema_name,
495
517
  parent_span_id=parent_span_id,
496
518
  payload=payload,
519
+ started_at=started_at,
497
520
  )
498
521
 
499
522
  async def finish_span(
500
523
  self,
501
524
  span_id: str,
502
525
  result_payload: dict[str, Any] | None = None,
526
+ timestamp: datetime | None = None,
503
527
  ) -> None:
504
528
  """Finish a previously created span.
505
529
 
506
530
  Args:
507
531
  span_id: The ID of the span to finish.
508
532
  result_payload: Optional result data to store on the span.
533
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
509
534
  """
510
535
  self._ensure_initialized()
511
536
  assert self._span_manager is not None
512
537
 
513
- await self._span_manager.finish(span_id, result_payload=result_payload)
538
+ await self._span_manager.finish(
539
+ span_id,
540
+ result_payload=result_payload,
541
+ timestamp=timestamp,
542
+ )
543
+
544
+ async def update_agent_instance(
545
+ self,
546
+ instance_id: str,
547
+ quality_payload: dict[str, Any] | None = None,
548
+ ) -> None:
549
+ """Update an agent instance (e.g., set quality payload).
550
+
551
+ Args:
552
+ instance_id: The ID of the instance to update.
553
+ quality_payload: Quality evaluation payload (None to clear).
554
+ """
555
+ self._ensure_initialized()
556
+ assert self._instance_manager is not None
557
+
558
+ await self._instance_manager.update(
559
+ instance_id,
560
+ quality_payload=quality_payload,
561
+ )
514
562
 
515
563
  @asynccontextmanager
516
564
  async def span(
@@ -14,7 +14,7 @@ from ..utils import generate_idempotency_key
14
14
 
15
15
  if TYPE_CHECKING:
16
16
  from prefactor_http.client import PrefactorHttpClient
17
- from prefactor_http.models.types import FinishStatus
17
+ from prefactor_http.models.types import FinishStatus, InstancePurpose
18
18
 
19
19
  from ..client import PrefactorCoreClient
20
20
 
@@ -64,6 +64,7 @@ class AgentInstanceManager:
64
64
  agent_id: str | None = None,
65
65
  instance_id: str | None = None,
66
66
  environment_id: str | None = None,
67
+ purpose: "InstancePurpose | None" = None,
67
68
  ) -> str:
68
69
  """Register a new agent instance.
69
70
 
@@ -79,6 +80,8 @@ class AgentInstanceManager:
79
80
  the API generates one.
80
81
  environment_id: Optional environment ID. Required when using an
81
82
  account-scoped token; omit when using a deployment-scoped token.
83
+ purpose: Why this instance ran — ``"live"``, ``"smoke_test"``,
84
+ or ``"eval"``. Omitted (None) lets the API default to ``"live"``.
82
85
 
83
86
  Returns:
84
87
  The instance ID (API-generated).
@@ -90,6 +93,7 @@ class AgentInstanceManager:
90
93
  environment_id=environment_id,
91
94
  id=instance_id,
92
95
  idempotency_key=generate_idempotency_key(),
96
+ purpose=purpose,
93
97
  )
94
98
  return result.id
95
99
 
@@ -158,6 +162,30 @@ class AgentInstanceManager:
158
162
 
159
163
  await self._enqueue(operation)
160
164
 
165
+ async def update(
166
+ self,
167
+ instance_id: str,
168
+ quality_payload: dict[str, Any] | None = None,
169
+ ) -> None:
170
+ """Update an agent instance (e.g., set quality payload).
171
+
172
+ Queues an update operation for the instance.
173
+
174
+ Args:
175
+ instance_id: The ID of the instance to update.
176
+ quality_payload: Quality evaluation payload (None to clear).
177
+ """
178
+ operation = Operation(
179
+ type=OperationType.UPDATE_AGENT_INSTANCE,
180
+ payload={
181
+ "instance_id": instance_id,
182
+ "quality_payload": quality_payload,
183
+ },
184
+ timestamp=datetime.now(timezone.utc),
185
+ )
186
+
187
+ await self._enqueue(operation)
188
+
161
189
 
162
190
  class AgentInstanceHandle:
163
191
  """Handle to an agent instance with convenience methods.
@@ -238,11 +266,30 @@ class AgentInstanceHandle:
238
266
  status=status,
239
267
  )
240
268
 
269
+ async def update(
270
+ self,
271
+ quality_payload: dict[str, Any] | None = None,
272
+ ) -> None:
273
+ """Update the instance (e.g., set quality payload).
274
+
275
+ This queues an update operation for the instance.
276
+
277
+ Args:
278
+ quality_payload: Quality evaluation payload (None to clear).
279
+ """
280
+ manager = self._client.instance_manager
281
+ assert manager is not None
282
+ await manager.update(
283
+ self._instance_id,
284
+ quality_payload=quality_payload,
285
+ )
286
+
241
287
  async def create_span(
242
288
  self,
243
289
  schema_name: str,
244
290
  parent_span_id: str | None = None,
245
291
  payload: dict[str, Any] | None = None,
292
+ started_at: datetime | None = None,
246
293
  ) -> str:
247
294
  """Create a span within this instance and return its ID.
248
295
 
@@ -252,6 +299,7 @@ class AgentInstanceHandle:
252
299
  schema_name: Name of the schema for this span.
253
300
  parent_span_id: Optional explicit parent span ID.
254
301
  payload: Optional initial payload (params/inputs) stored on creation.
302
+ started_at: Optional ISO 8601 start time (defaults to current time).
255
303
 
256
304
  Returns:
257
305
  The span ID.
@@ -262,20 +310,27 @@ class AgentInstanceHandle:
262
310
  schema_name=schema_name,
263
311
  parent_span_id=parent_span_id,
264
312
  payload=payload,
313
+ started_at=started_at,
265
314
  )
266
315
 
267
316
  async def finish_span(
268
317
  self,
269
318
  span_id: str,
270
319
  result_payload: dict[str, Any] | None = None,
320
+ timestamp: datetime | None = None,
271
321
  ) -> None:
272
322
  """Finish a previously created span.
273
323
 
274
324
  Args:
275
325
  span_id: The ID of the span to finish.
276
326
  result_payload: Optional result data to store on the span.
327
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
277
328
  """
278
- await self._client.finish_span(span_id, result_payload=result_payload)
329
+ await self._client.finish_span(
330
+ span_id,
331
+ result_payload=result_payload,
332
+ timestamp=timestamp,
333
+ )
279
334
 
280
335
  @asynccontextmanager
281
336
  async def span(
@@ -109,6 +109,7 @@ class SpanManager:
109
109
  self,
110
110
  temp_id: str,
111
111
  payload: dict[str, Any] | None = None,
112
+ started_at: datetime | None = None,
112
113
  ) -> str:
113
114
  """Post the span to the API as ``active`` and return the API-generated ID.
114
115
 
@@ -121,6 +122,7 @@ class SpanManager:
121
122
  Args:
122
123
  temp_id: The temporary ID returned by ``prepare()``.
123
124
  payload: Optional params/inputs to send with the span.
125
+ started_at: Optional ISO 8601 start time (defaults to current time).
124
126
 
125
127
  Returns:
126
128
  The API-generated span ID.
@@ -139,6 +141,7 @@ class SpanManager:
139
141
  status="active",
140
142
  payload=payload or {},
141
143
  parent_span_id=span.parent_span_id,
144
+ started_at=started_at,
142
145
  idempotency_key=generate_idempotency_key(),
143
146
  )
144
147
 
@@ -148,6 +151,7 @@ class SpanManager:
148
151
  span.id = api_id
149
152
  span.payload = payload or {}
150
153
  span.status = "active"
154
+ span.started_at = started_at or datetime.now(timezone.utc)
151
155
  del self._spans[temp_id]
152
156
  self._spans[api_id] = span
153
157
 
@@ -165,7 +169,11 @@ class SpanManager:
165
169
 
166
170
  return api_id
167
171
 
168
- async def cancel_unstarted(self, temp_id: str) -> None:
172
+ async def cancel_unstarted(
173
+ self,
174
+ temp_id: str,
175
+ timestamp: datetime | None = None,
176
+ ) -> None:
169
177
  """Cancel a span that was never started.
170
178
 
171
179
  When ``cancel()`` is called before ``start()``, the span has not yet
@@ -175,6 +183,7 @@ class SpanManager:
175
183
 
176
184
  Args:
177
185
  temp_id: The temporary ID returned by ``prepare()``.
186
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
178
187
 
179
188
  Raises:
180
189
  KeyError: If temp_id is not a known pending span.
@@ -197,11 +206,13 @@ class SpanManager:
197
206
  await self._http.agent_spans.finish(
198
207
  agent_span_id=api_id,
199
208
  status="cancelled",
209
+ timestamp=timestamp,
200
210
  idempotency_key=generate_idempotency_key(),
201
211
  )
202
212
 
213
+ effective_timestamp = timestamp or datetime.now(timezone.utc)
203
214
  span.status = "cancelled"
204
- span.finished_at = datetime.now(timezone.utc)
215
+ span.finished_at = effective_timestamp
205
216
 
206
217
  stack = SpanContextStack.get_stack()
207
218
  if temp_id in stack:
@@ -218,6 +229,7 @@ class SpanManager:
218
229
  parent_span_id: str | None = None,
219
230
  payload: dict[str, Any] | None = None,
220
231
  span_id: str | None = None,
232
+ started_at: datetime | None = None,
221
233
  ) -> str:
222
234
  """Create a span in one step (prepare + start).
223
235
 
@@ -230,6 +242,7 @@ class SpanManager:
230
242
  parent_span_id: Optional parent span ID (auto-detected if None).
231
243
  payload: Optional initial payload data.
232
244
  span_id: Ignored (API generates IDs).
245
+ started_at: Optional ISO 8601 start time (defaults to current time).
233
246
 
234
247
  Returns:
235
248
  The API-generated span ID.
@@ -239,7 +252,7 @@ class SpanManager:
239
252
  schema_name=schema_name,
240
253
  parent_span_id=parent_span_id,
241
254
  )
242
- return await self.start(temp_id, payload=payload)
255
+ return await self.start(temp_id, payload=payload, started_at=started_at)
243
256
 
244
257
  async def finish(
245
258
  self,
@@ -247,6 +260,7 @@ class SpanManager:
247
260
  result_payload: dict[str, Any] | None = None,
248
261
  status: "FinishStatus" = "complete",
249
262
  idempotency_key: str | None = None,
263
+ timestamp: datetime | None = None,
250
264
  ) -> None:
251
265
  """Mark a span as finished.
252
266
 
@@ -261,6 +275,7 @@ class SpanManager:
261
275
  to cancel a span that was never started.
262
276
  idempotency_key: Optional key to make repeated finish requests
263
277
  duplicate-safe. When omitted, a new key is generated.
278
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
264
279
 
265
280
  Raises:
266
281
  KeyError: If the span ID is not known.
@@ -268,8 +283,9 @@ class SpanManager:
268
283
  if span_id not in self._spans:
269
284
  raise KeyError(f"Unknown span: {span_id}")
270
285
 
286
+ effective_timestamp = timestamp or datetime.now(timezone.utc)
271
287
  self._spans[span_id].status = status
272
- self._spans[span_id].finished_at = datetime.now(timezone.utc)
288
+ self._spans[span_id].finished_at = effective_timestamp
273
289
 
274
290
  stack = SpanContextStack.get_stack()
275
291
  if span_id in stack:
@@ -288,7 +304,7 @@ class SpanManager:
288
304
  operation = Operation(
289
305
  type=OperationType.FINISH_SPAN,
290
306
  payload=op_payload,
291
- timestamp=datetime.now(timezone.utc),
307
+ timestamp=effective_timestamp,
292
308
  )
293
309
 
294
310
  await self._enqueue(operation)
@@ -28,10 +28,12 @@ class AgentInstance:
28
28
  id: str
29
29
  agent_id: str
30
30
  status: str = "pending"
31
+ purpose: str | None = None
31
32
  created_at: datetime = field(default_factory=datetime.now)
32
33
  started_at: datetime | None = None
33
34
  finished_at: datetime | None = None
34
35
  metadata: dict[str, Any] = field(default_factory=dict)
36
+ quality_payload: dict[str, Any] | None = None
35
37
 
36
38
 
37
39
  @dataclass
@@ -19,6 +19,7 @@ class OperationType(Enum):
19
19
  REGISTER_AGENT_INSTANCE = auto()
20
20
  START_AGENT_INSTANCE = auto()
21
21
  FINISH_AGENT_INSTANCE = auto()
22
+ UPDATE_AGENT_INSTANCE = auto()
22
23
  CREATE_SPAN = auto()
23
24
  FINISH_SPAN = auto()
24
25
 
@@ -67,6 +67,8 @@ class SchemaRegistry:
67
67
  self._span_result_schemas: dict[str, dict[str, Any]] = {}
68
68
  # span_type_schemas: name → full structured entry
69
69
  self._span_type_schemas: dict[str, dict[str, Any]] = {}
70
+ # quality_schema: optional quality schema for instance evaluations
71
+ self._quality_schema: dict[str, Any] | None = None
70
72
 
71
73
  def register(
72
74
  self,
@@ -188,6 +190,46 @@ class SchemaRegistry:
188
190
 
189
191
  self._span_type_schemas[name] = entry
190
192
 
193
+ def register_quality_schema(
194
+ self,
195
+ schema: dict[str, Any],
196
+ title: str | None = None,
197
+ description: str | None = None,
198
+ template: str | None = None,
199
+ data_risk: dict[str, Any] | None = None,
200
+ ) -> None:
201
+ """Register a quality schema for instance evaluations.
202
+
203
+ The quality schema defines the shape of quality payloads that can be
204
+ set on agent instances. It uses the same ``title``, ``description``,
205
+ ``template``, and ``data_risk`` fields as span type schemas.
206
+
207
+ Args:
208
+ schema: JSON Schema dict defining the quality payload structure.
209
+ title: Optional human-readable title (defaults to "quality" on API).
210
+ description: Optional description of the quality evaluation.
211
+ template: Optional display template using ``{{field}}`` interpolation.
212
+ data_risk: Optional data risk classification dict (same structure
213
+ as span type data_risk).
214
+
215
+ Raises:
216
+ ValueError: If a quality schema is already registered.
217
+ """
218
+ if self._quality_schema is not None:
219
+ raise ValueError("Quality schema is already registered")
220
+
221
+ entry: dict[str, Any] = {"schema": schema}
222
+ if title is not None:
223
+ entry["title"] = title
224
+ if description is not None:
225
+ entry["description"] = description
226
+ if template is not None:
227
+ entry["template"] = template
228
+ if data_risk is not None:
229
+ entry["data_risk"] = data_risk
230
+
231
+ self._quality_schema = entry
232
+
191
233
  def get(self, schema_name: str) -> dict[str, Any] | None:
192
234
  """Get a params schema by name.
193
235
 
@@ -244,6 +286,9 @@ class SchemaRegistry:
244
286
  if self._span_type_schemas:
245
287
  result["span_type_schemas"] = list(self._span_type_schemas.values())
246
288
 
289
+ if self._quality_schema is not None:
290
+ result["quality_schema"] = dict(self._quality_schema)
291
+
247
292
  return result
248
293
 
249
294
  def merge(self, other: "SchemaRegistry") -> None:
@@ -270,6 +315,10 @@ class SchemaRegistry:
270
315
  if name in self._span_type_schemas:
271
316
  conflicts.append(f"span_type_schemas/{name}")
272
317
 
318
+ if other._quality_schema is not None:
319
+ if self._quality_schema is not None:
320
+ conflicts.append("quality_schema")
321
+
273
322
  if conflicts:
274
323
  msg = f"Cannot merge registries - conflicting schemas: {conflicts}"
275
324
  raise ValueError(msg)
@@ -283,5 +332,8 @@ class SchemaRegistry:
283
332
  for name, entry in other._span_type_schemas.items():
284
333
  self._span_type_schemas[name] = entry.copy()
285
334
 
335
+ if other._quality_schema is not None:
336
+ self._quality_schema = other._quality_schema.copy()
337
+
286
338
 
287
339
  __all__ = ["SchemaRegistry"]
@@ -6,6 +6,7 @@ and ensures proper cleanup when the span completes.
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
+ from datetime import datetime
9
10
  from typing import TYPE_CHECKING, Any
10
11
 
11
12
  from .utils import generate_idempotency_key
@@ -80,6 +81,7 @@ class SpanContext:
80
81
  tuple[FinishStatus, tuple[tuple[str, Any], ...]] | None
81
82
  ) = None
82
83
  self._finish_idempotency_key: str | None = None
84
+ self._finish_timestamp: datetime | None = None
83
85
 
84
86
  @property
85
87
  def id(self) -> str:
@@ -93,7 +95,11 @@ class SpanContext:
93
95
  """
94
96
  return self._span_id
95
97
 
96
- async def start(self, payload: dict[str, Any] | None = None) -> None:
98
+ async def start(
99
+ self,
100
+ payload: dict[str, Any] | None = None,
101
+ started_at: datetime | None = None,
102
+ ) -> None:
97
103
  """Post the span to the API as ``active`` with the given params payload.
98
104
 
99
105
  This triggers ``POST /api/v1/agent_spans``. The span is created as
@@ -105,11 +111,14 @@ class SpanContext:
105
111
  payload: Optional params/inputs for the span (e.g. model name,
106
112
  prompt text, tool input). Stored as the span's ``payload``
107
113
  field in the API.
114
+ started_at: Optional ISO 8601 start time (defaults to current time).
108
115
  """
109
116
  if self._started:
110
117
  return
111
118
 
112
- api_id = await self._span_manager.start(self._span_id, payload=payload)
119
+ api_id = await self._span_manager.start(
120
+ self._span_id, payload=payload, started_at=started_at
121
+ )
113
122
  self._span_id = api_id
114
123
  self._started = True
115
124
 
@@ -124,37 +133,53 @@ class SpanContext:
124
133
  """
125
134
  self._result_payload.update(data)
126
135
 
127
- async def complete(self, result: dict[str, Any] | None = None) -> None:
136
+ async def complete(
137
+ self,
138
+ result: dict[str, Any] | None = None,
139
+ timestamp: datetime | None = None,
140
+ ) -> None:
128
141
  """Finish the span with ``complete`` status.
129
142
 
130
143
  Args:
131
144
  result: Optional result payload to attach to the span.
145
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
132
146
  """
133
147
  if result:
134
148
  self.set_result(result)
135
149
  self._finish_status = "complete"
150
+ self._finish_timestamp = timestamp
136
151
  await self._finish()
137
152
 
138
- async def fail(self, result: dict[str, Any] | None = None) -> None:
153
+ async def fail(
154
+ self,
155
+ result: dict[str, Any] | None = None,
156
+ timestamp: datetime | None = None,
157
+ ) -> None:
139
158
  """Finish the span with ``failed`` status.
140
159
 
141
160
  Args:
142
161
  result: Optional result payload (e.g. error details).
162
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
143
163
  """
144
164
  if result:
145
165
  self.set_result(result)
146
166
  self._finish_status = "failed"
167
+ self._finish_timestamp = timestamp
147
168
  await self._finish()
148
169
 
149
- async def cancel(self) -> None:
170
+ async def cancel(self, timestamp: datetime | None = None) -> None:
150
171
  """Finish the span with ``cancelled`` status.
151
172
 
152
173
  Can be called before or after ``start()``. If ``start()`` has not
153
174
  been called yet, the span is posted as ``pending`` then immediately
154
175
  cancelled — the API only accepts cancellation from the ``pending``
155
176
  state, so this is always a valid sequence.
177
+
178
+ Args:
179
+ timestamp: Optional ISO 8601 finish time (defaults to current time).
156
180
  """
157
181
  self._finish_status = "cancelled"
182
+ self._finish_timestamp = timestamp
158
183
  await self._finish()
159
184
 
160
185
  async def finish(self) -> None:
@@ -192,7 +217,10 @@ class SpanContext:
192
217
 
193
218
  try:
194
219
  if not self._started and self._finish_status == "cancelled":
195
- await self._span_manager.cancel_unstarted(self._span_id)
220
+ await self._span_manager.cancel_unstarted(
221
+ self._span_id,
222
+ timestamp=self._finish_timestamp,
223
+ )
196
224
  self._finished = True
197
225
  return
198
226
 
@@ -206,6 +234,7 @@ class SpanContext:
206
234
  result_payload=self._result_payload or None,
207
235
  status=self._finish_status,
208
236
  idempotency_key=self._finish_idempotency_key,
237
+ timestamp=self._finish_timestamp,
209
238
  )
210
239
  self._finished = True
211
240
  except Exception:
@@ -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(
@@ -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():
@@ -68,17 +68,28 @@ class _StubHttpClient:
68
68
  *_args,
69
69
  agent_instances: _StubAgentInstances | None = None,
70
70
  agent_spans: _StubAgentSpans | None = None,
71
+ validate_token_error: PrefactorAuthError | None = None,
71
72
  **_kwargs,
72
73
  ) -> None:
73
74
  self.agent_instances = agent_instances or _StubAgentInstances()
74
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
75
79
 
76
80
  async def __aenter__(self):
81
+ self.enter_calls += 1
77
82
  return self
78
83
 
79
84
  async def __aexit__(self, exc_type, exc, tb):
85
+ self.exit_calls += 1
80
86
  return None
81
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
+
82
93
 
83
94
  def _make_client_config(max_retries: int = 0) -> PrefactorCoreConfig:
84
95
  return PrefactorCoreConfig(
@@ -375,3 +386,23 @@ async def test_latched_failure_drops_already_queued_backlog():
375
386
  await client.close()
376
387
 
377
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