diff --git a/CHANGELOG.md b/CHANGELOG.md index 54338d9ba..82f31dcf0 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### ⚠ BREAKING CHANGES +* **tracing:** removed the Agentex-native span processor and its `AgentexTracingProcessorConfig` (the Agentex server is retiring its Postgres spans API), along with `Trace.get_span` / `Trace.list_spans` and their async twins. `SGPTracingProcessorConfig` is the only processor config and registering any other type raises `ValueError`. The in-memory `Span` is now `agentex.lib.types.tracing.Span` (also exported as `agentex.lib.core.tracing.Span`); the generated `agentex.types.span.Span` disappears with the next client generation. Its `to_dict()` / `to_json()` return the full JSON-mode dump rather than only the fields that were set. * **harness:** removed the deprecated bespoke LangGraph tracing handler `create_langgraph_tracing_handler` (and its `AgentexLangGraphTracingHandler` class) from the public `agentex.lib.adk` surface. Span tracing is now derived from the canonical `StreamTaskMessage*` stream by `UnifiedEmitter` — wrap your run in the harness `*Turn` and drive `UnifiedEmitter.yield_turn` / `auto_send_turn`. The `agentex init` templates were migrated accordingly. * **harness:** removed the deprecated bespoke Pydantic-AI tracing handler `create_pydantic_ai_tracing_handler` (and its `AgentexPydanticAITracingHandler` class) from the public `agentex.lib.adk` surface. Span tracing is now derived from the canonical `StreamTaskMessage*` stream by `UnifiedEmitter` — wrap your run in `PydanticAITurn` and drive `UnifiedEmitter.yield_turn` / `auto_send_turn`. The `agentex init` templates were migrated accordingly. * **harness:** each harness now exposes exactly `__sync.py` + `__turn.py` under `agentex.lib.adk._modules`. The OpenAI harness `OpenAITurn` and `convert_openai_to_agentex_events` moved to `agentex.lib.adk._modules._openai_turn` / `_openai_sync`; back-compat shims remain at `agentex.lib.adk.providers._modules.{openai_turn,sync_provider}` for one release. Public facade names (`stream_pydantic_ai_events`, `stream_langgraph_events`, `emit_langgraph_messages`, etc.) are unchanged. diff --git a/examples/tutorials/10_async/10_temporal/020_state_machine/project/state_machines/deep_research.py b/examples/tutorials/10_async/10_temporal/020_state_machine/project/state_machines/deep_research.py index d1c4df00a..9277ff688 100644 --- a/examples/tutorials/10_async/10_temporal/020_state_machine/project/state_machines/deep_research.py +++ b/examples/tutorials/10_async/10_temporal/020_state_machine/project/state_machines/deep_research.py @@ -3,7 +3,7 @@ from pydantic import BaseModel -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.sdk.state_machine import StateMachine diff --git a/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py index 9b89d076e..119eefaa5 100644 --- a/src/agentex/lib/adk/_modules/tracing.py +++ b/src/agentex/lib/adk/_modules/tracing.py @@ -11,7 +11,6 @@ from temporalio.exceptions import ActivityError, TimeoutError as TemporalTimeoutError, is_cancelled_exception from agentex import AsyncAgentex # noqa: F401 -from agentex.lib.adk.utils._modules.client import create_async_agentex_client from agentex.lib.core.services.adk.tracing import TracingService from agentex.lib.core.temporal.activities.activity_helpers import ActivityHelpers from agentex.lib.core.temporal.activities.adk.tracing_activities import ( @@ -22,7 +21,7 @@ from agentex.lib.core.tracing.span_error import set_span_error from agentex.lib.core.tracing.tracer import AsyncTracer from agentex.lib.core.harness.types import TurnUsage -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.utils.logging import make_logger from agentex.lib.utils.model_utils import BaseModel from agentex.lib.utils.temporal import in_temporal_workflow @@ -145,46 +144,17 @@ def __init__(self, tracing_service: TracingService | None = None): Args: tracing_service (Optional[TracingService]): Optional pre-configured tracing service. - If None, will be lazily created on first use so the httpx client is - bound to the correct running event loop. + If None, one is created on first use. """ self._tracing_service_explicit = tracing_service self._tracing_service_lazy: TracingService | None = None - self._bound_loop_id: int | None = None @property def _tracing_service(self) -> TracingService: if self._tracing_service_explicit is not None: return self._tracing_service_explicit - - import asyncio - - # Determine the current event loop (if any). - try: - loop = asyncio.get_running_loop() - loop_id = id(loop) - except RuntimeError: - loop_id = None - - # Re-create the underlying httpx client when the event loop changes - # (e.g. between HTTP requests in a sync ASGI server) to avoid - # "Event loop is closed" / "bound to a different event loop" errors. - if self._tracing_service_lazy is None or (loop_id is not None and loop_id != self._bound_loop_id): - import httpx - - # Keepalive ON: connections are reused within a single event - # loop, eliminating the TLS-handshake-per-span penalty under - # load. Cross-loop safety is preserved by rebuilding the - # client whenever loop_id changes (the conditional above). - agentex_client = create_async_agentex_client( - http_client=httpx.AsyncClient( - limits=httpx.Limits(max_keepalive_connections=20), - ), - ) - tracer = AsyncTracer(agentex_client) - self._tracing_service_lazy = TracingService(tracer=tracer) - self._bound_loop_id = loop_id - + if self._tracing_service_lazy is None: + self._tracing_service_lazy = TracingService(tracer=AsyncTracer()) return self._tracing_service_lazy @asynccontextmanager diff --git a/src/agentex/lib/core/services/adk/tracing.py b/src/agentex/lib/core/services/adk/tracing.py index 77efffd9e..561f8cad3 100644 --- a/src/agentex/lib/core/services/adk/tracing.py +++ b/src/agentex/lib/core/services/adk/tracing.py @@ -2,7 +2,7 @@ from typing import Any -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.utils.logging import make_logger from agentex.lib.utils.temporal import heartbeat_if_in_workflow from agentex.lib.utils.model_utils import BaseModel diff --git a/src/agentex/lib/core/temporal/activities/adk/tracing_activities.py b/src/agentex/lib/core/temporal/activities/adk/tracing_activities.py index aec541afe..7955185f4 100644 --- a/src/agentex/lib/core/temporal/activities/adk/tracing_activities.py +++ b/src/agentex/lib/core/temporal/activities/adk/tracing_activities.py @@ -5,7 +5,7 @@ from temporalio import activity -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.utils.logging import make_logger from agentex.lib.utils.model_utils import BaseModel from agentex.lib.core.services.adk.tracing import TracingService diff --git a/src/agentex/lib/core/tracing/__init__.py b/src/agentex/lib/core/tracing/__init__.py index 580b53c20..2eadf79de 100644 --- a/src/agentex/lib/core/tracing/__init__.py +++ b/src/agentex/lib/core/tracing/__init__.py @@ -1,4 +1,4 @@ -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.tracing.trace import Trace, AsyncTrace from agentex.lib.core.tracing.tracer import Tracer, AsyncTracer from agentex.lib.core.tracing.span_error import ( diff --git a/src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py b/src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py deleted file mode 100644 index 448d013e9..000000000 --- a/src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py +++ /dev/null @@ -1,232 +0,0 @@ -import os -import asyncio -import weakref -from typing import TYPE_CHECKING, Any, Dict, override - -from agentex import Agentex -from agentex.types.span import Span -from agentex.lib.types.tracing import AgentexTracingProcessorConfig -from agentex.lib.utils.logging import make_logger -from agentex.lib.adk.utils._modules.client import create_async_agentex_client -from agentex.lib.core.tracing.processors.tracing_processor_interface import ( - SyncTracingProcessor, - AsyncTracingProcessor, -) - -if TYPE_CHECKING: - from agentex import AsyncAgentex - -logger = make_logger(__name__) - - -# NOTE: This is the Agentex-backend toggle (writes to the agentex `spans` -# table via the Agentex API). It is intentionally SEPARATE from the SGP/EGP -# processor's ``AGENTEX_TRACING_SKIP_SPAN_START`` so the two backends can be -# controlled independently. -_SKIP_SPAN_START_ENV = "AGENTEX_TRACING_SKIP_AGENTEX_SPAN_START" - - -def _skip_span_start_enabled() -> bool: - """Whether to skip the Agentex span-start write and persist each span only on end. - - The Agentex processor otherwise writes every span twice: a ``spans.create`` - on start (no ``end_time``/``output`` yet) and a ``spans.update`` on end. - The start row is overwritten by the end write moments later, so persisting - it doubles the per-span HTTP/DB write volume against the Agentex control - plane — the load that timed out span-start activities and pressured the - Agentex Postgres connection pool under load. - - When enabled (the default), the start write is skipped and the END write - becomes a single ``spans.create`` carrying the complete span — one INSERT - per span instead of an INSERT + UPDATE. (A plain ``spans.update`` on end - would 404 because the row was never created.) - - Default ON. Set ``AGENTEX_TRACING_SKIP_AGENTEX_SPAN_START`` to - ``0``/``false``/``no``/``off`` to restore the start write — e.g. if you - need in-flight spans visible before they complete, or spans that never end - (process crash) to still be persisted. - """ - raw = os.environ.get(_SKIP_SPAN_START_ENV, "1").strip().lower() - return raw not in ("0", "false", "no", "off") - - -def _create_kwargs(span: Span) -> Dict[str, Any]: - """Full-span kwargs for ``spans.create`` — used on start (skip disabled) and - on end (skip enabled, single-INSERT path).""" - return { - "name": span.name, - "start_time": span.start_time, - "end_time": span.end_time, - "id": span.id, - "trace_id": span.trace_id, - "parent_id": span.parent_id, - "input": span.input, - "output": span.output, - "data": span.data, - "task_id": span.task_id, - } - - -class AgentexSyncTracingProcessor(SyncTracingProcessor): - def __init__(self, config: AgentexTracingProcessorConfig): # noqa: ARG002 - self.client = Agentex() - # Capture the skip decision once at init: both halves of a span's - # lifecycle MUST agree, otherwise a start-skip + end-update lands on a - # non-existent row (404) — or the reverse double-creates. Re-reading the - # env per event would let a mid-span toggle (tests, config reload) split - # the decision. Deploy-time flag, so a single read is correct. - self._skip_span_start = _skip_span_start_enabled() - logger.info( - "Agentex tracing span-start write %s (%s)", - "disabled — end-only ingest" if self._skip_span_start else "enabled", - _SKIP_SPAN_START_ENV, - ) - - @override - def on_span_start(self, span: Span) -> None: - # End-only ingest: by default the start write is skipped (see - # _skip_span_start_enabled) so each span is persisted once, on end. - if self._skip_span_start: - return - self.client.spans.create(**_create_kwargs(span)) - - @override - def on_span_end(self, span: Span) -> None: - # End-only ingest: the start create was skipped, so persist the complete - # span as a single INSERT here (a bare spans.update would 404 — no row). - if self._skip_span_start: - self.client.spans.create(**_create_kwargs(span)) - return - - update: Dict[str, Any] = {} - if span.trace_id: - update["trace_id"] = span.trace_id - if span.name: - update["name"] = span.name - if span.parent_id: - update["parent_id"] = span.parent_id - if span.start_time: - update["start_time"] = span.start_time.isoformat() - if span.end_time is not None: - update["end_time"] = span.end_time.isoformat() - if span.input is not None: - update["input"] = span.input - if span.output is not None: - update["output"] = span.output - if span.data is not None: - update["data"] = span.data - - self.client.spans.update( - span.id, - **span.model_dump( - mode="json", - exclude={"id"}, - exclude_defaults=True, - exclude_none=True, - exclude_unset=True, - ), - ) - - @override - def shutdown(self) -> None: - pass - - -class AgentexAsyncTracingProcessor(AsyncTracingProcessor): - def __init__(self, config: AgentexTracingProcessorConfig): # noqa: ARG002 - # Per-event-loop client cache. httpx.AsyncClient is bound to the - # loop that created it, so in sync-ACP / streaming contexts (where - # the active loop can change between requests) we keep one client - # per loop instead of disabling keepalive entirely. The cache is a - # WeakKeyDictionary so a GC'd loop and its client are evicted - # automatically — using id() as a key would reuse entries when - # CPython recycles a freed loop's memory address. - self._clients_by_loop: weakref.WeakKeyDictionary[ - asyncio.AbstractEventLoop, "AsyncAgentex" - ] = weakref.WeakKeyDictionary() - # Capture the skip decision once at init: both halves of a span's - # lifecycle MUST agree, otherwise a start-skip + end-update lands on a - # non-existent row (404) — or the reverse double-creates. Re-reading the - # env per event would let a mid-span toggle (tests, config reload) split - # the decision. Deploy-time flag, so a single read is correct. - self._skip_span_start = _skip_span_start_enabled() - logger.info( - "Agentex tracing span-start write %s (%s)", - "disabled — end-only ingest" if self._skip_span_start else "enabled", - _SKIP_SPAN_START_ENV, - ) - - def _build_client(self) -> "AsyncAgentex": - import httpx - - # Keepalive ON: connections are reused within a single event loop, - # eliminating the TLS-handshake-per-span penalty under load. - return create_async_agentex_client( - http_client=httpx.AsyncClient( - limits=httpx.Limits(max_keepalive_connections=20), - ), - ) - - @property - def client(self) -> "AsyncAgentex": - try: - loop = asyncio.get_running_loop() - except RuntimeError: - return self._build_client() - client = self._clients_by_loop.get(loop) - if client is None: - client = self._build_client() - self._clients_by_loop[loop] = client - return client - - # TODO(AGX1-199): Add batch create/update endpoints to Agentex API and use - # them here instead of one HTTP call per span. - # https://linear.app/scale-epd/issue/AGX1-199/add-agentex-batch-endpoint-for-traces - @override - async def on_span_start(self, span: Span) -> None: - # End-only ingest: by default the start write is skipped (see - # _skip_span_start_enabled) so each span is persisted once, on end. - if self._skip_span_start: - return - await self.client.spans.create(**_create_kwargs(span)) - - @override - async def on_span_end(self, span: Span) -> None: - # End-only ingest: the start create was skipped, so persist the complete - # span as a single INSERT here (a bare spans.update would 404 — no row). - if self._skip_span_start: - await self.client.spans.create(**_create_kwargs(span)) - return - - update: Dict[str, Any] = {} - if span.trace_id: - update["trace_id"] = span.trace_id - if span.name: - update["name"] = span.name - if span.parent_id: - update["parent_id"] = span.parent_id - if span.start_time: - update["start_time"] = span.start_time.isoformat() - if span.end_time: - update["end_time"] = span.end_time.isoformat() - if span.input: - update["input"] = span.input - if span.output: - update["output"] = span.output - if span.data: - update["data"] = span.data - - await self.client.spans.update( - span.id, - **span.model_dump( - mode="json", - exclude={"id"}, - exclude_defaults=True, - exclude_none=True, - exclude_unset=True, - ), - ) - - @override - async def shutdown(self) -> None: - pass diff --git a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index 9ee269231..124e69b29 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -10,9 +10,8 @@ from scale_gp_beta.lib.tracing import create_span, flush_queue from scale_gp_beta.lib.tracing.span import Span as SGPSpan -from agentex.types.span import Span from agentex.lib.core.tracing import code_revision -from agentex.lib.types.tracing import SGPTracingProcessorConfig +from agentex.lib.types.tracing import Span, SGPTracingProcessorConfig from agentex.lib.utils.logging import make_logger from agentex.lib.core.observability import tracing_metrics_recording as _metrics from agentex.lib.environment_variables import EnvironmentVariables @@ -75,8 +74,8 @@ def _sgp_metadata(span: Span) -> Any: Returns a COPY rather than mutating ``span``. ``trace.py`` hands the same Span instance to every registered processor, so anything written onto - ``span.data`` here would also be serialized by the Agentex processor and - show up in caller-visible span data. ``__commit_sha__`` is opt-in and + ``span.data`` here would also reach every other processor and show up in + caller-visible span data. ``__commit_sha__`` is opt-in and SGP-scoped, so it must not leak that way. (The ``__source__`` / ``__agent_*`` keys set by ``_add_source_to_span`` do diff --git a/src/agentex/lib/core/tracing/processors/tracing_processor_interface.py b/src/agentex/lib/core/tracing/processors/tracing_processor_interface.py index f352f38c4..8da0e51c0 100644 --- a/src/agentex/lib/core/tracing/processors/tracing_processor_interface.py +++ b/src/agentex/lib/core/tracing/processors/tracing_processor_interface.py @@ -3,8 +3,7 @@ import asyncio from abc import ABC, abstractmethod -from agentex.types.span import Span -from agentex.lib.types.tracing import TracingProcessorConfig +from agentex.lib.types.tracing import Span, TracingProcessorConfig from agentex.lib.utils.logging import make_logger logger = make_logger(__name__) diff --git a/src/agentex/lib/core/tracing/span_error.py b/src/agentex/lib/core/tracing/span_error.py index f20eae471..81cb0edfd 100644 --- a/src/agentex/lib/core/tracing/span_error.py +++ b/src/agentex/lib/core/tracing/span_error.py @@ -9,15 +9,12 @@ ) from scale_gp_beta.lib.tracing.types import ErrorCategory -from agentex.types.span import Span +from agentex.lib.types.tracing import Span # Reserved key under ``Span.data`` carrying failure info for a span whose -# context-manager body raised. Mirrors the existing ``__span_type__`` / -# ``__source__`` reserved-key convention already read/written by the SGP -# processor. Stored in ``data`` because the Span model is generated from the -# OpenAPI spec and has no first-class status/error field; ``data`` is a real -# field, so it survives ``model_copy(deep=True)`` and round-trips to both the -# SGP and agentex-native span stores. +# context-manager body raised, alongside the ``__span_type__`` / ``__source__`` +# keys the SGP processor already reads. Kept in ``data`` so it survives +# ``model_copy(deep=True)`` and reaches the processors with the span. SPAN_ERROR_KEY = "__error__" ERROR_CATEGORY_UNKNOWN: ErrorCategory = "unknown" diff --git a/src/agentex/lib/core/tracing/span_queue.py b/src/agentex/lib/core/tracing/span_queue.py index d6ff7c1f6..20394269f 100644 --- a/src/agentex/lib/core/tracing/span_queue.py +++ b/src/agentex/lib/core/tracing/span_queue.py @@ -6,7 +6,7 @@ from enum import Enum from dataclasses import dataclass -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.utils.logging import make_logger from agentex.lib.core.observability import tracing_metrics_recording as _metrics from agentex.lib.core.tracing.processors.tracing_processor_interface import ( diff --git a/src/agentex/lib/core/tracing/trace.py b/src/agentex/lib/core/tracing/trace.py index 8f5260913..4e3fbd393 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -9,7 +9,7 @@ from pydantic import BaseModel from agentex import Agentex, AsyncAgentex -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.utils.logging import make_logger from agentex.lib.utils.model_utils import recursive_model_dump from agentex.lib.core.tracing.obs_ids import obs_correlation, warn_on_backend_drift @@ -218,23 +218,24 @@ def _begin_obs( class Trace: """ - Trace is a wrapper around the Agentex API for tracing. - It provides a context manager for spans and a way to start and end spans. - It also provides a way to get spans by ID and list all spans in a trace. + Trace groups the spans of one trace id and hands each span to the + registered processors. It provides a context manager for spans and a way + to start and end spans. """ def __init__( self, processors: list[SyncTracingProcessor], - client: Agentex, + client: Agentex | None = None, trace_id: str | None = None, ): """ Initialize a new trace with the specified trace ID. Args: - trace_id: Required trace ID to use for this trace. - processors: Optional list of tracing processors to use for this trace. + processors: Tracing processors every span is handed to. + client: Kept for backward compatibility, no longer used. + trace_id: Trace ID to use for this trace. """ self.processors = processors self.client = client @@ -252,7 +253,7 @@ def start_span( task_id: str | None = None, ) -> Span: """ - Start a new span and register it with the API. + Start a new span and hand it to the registered processors. Args: name: Name of the span. @@ -268,7 +269,6 @@ def start_span( if not self.trace_id: raise ValueError("Trace ID is required to start a span") - # Create a span using the client's spans resource start_time = datetime.now(UTC) serialized_input = recursive_model_dump(input) if input else None @@ -327,31 +327,6 @@ def end_span( return span - def get_span(self, span_id: str) -> Span: - """ - Get a span by ID. - - Args: - span_id: The ID of the span to get. - - Returns: - The requested span. - """ - # Query from Agentex API - span = self.client.spans.retrieve(span_id) - return span - - def list_spans(self) -> list[Span]: - """ - List all spans in this trace. - - Returns: - List of spans in this trace. - """ - # Query from Agentex API - spans = self.client.spans.list(trace_id=self.trace_id) - return spans - @contextmanager def span( self, @@ -380,15 +355,14 @@ def span( class AsyncTrace: """ - AsyncTrace is a wrapper around the Agentex API for tracing. - It provides a context manager for spans and a way to start and end spans. - It also provides a way to get spans by ID and list all spans in a trace. + AsyncTrace is the async version of Trace. It provides a context manager + for spans and a way to start and end spans. """ def __init__( self, processors: list[AsyncTracingProcessor], - client: AsyncAgentex, + client: AsyncAgentex | None = None, trace_id: str | None = None, span_queue: AsyncSpanQueue | None = None, ): @@ -417,7 +391,7 @@ async def start_span( task_id: str | None = None, ) -> Span: """ - Start a new span and register it with the API. + Start a new span and hand it to the registered processors. Args: name: Name of the span. @@ -432,7 +406,6 @@ async def start_span( if not self.trace_id: raise ValueError("Trace ID is required to start a span") - # Create a span using the client's spans resource start_time = datetime.now(UTC) serialized_input = recursive_model_dump(input) if input else None @@ -507,31 +480,6 @@ async def end_span( return span - async def get_span(self, span_id: str) -> Span: - """ - Get a span by ID. - - Args: - span_id: The ID of the span to get. - - Returns: - The requested span. - """ - # Query from Agentex API - span = await self.client.spans.retrieve(span_id) - return span - - async def list_spans(self) -> list[Span]: - """ - List all spans in this trace. - - Returns: - List of spans in this trace. - """ - # Query from Agentex API - spans = await self.client.spans.list(trace_id=self.trace_id) - return spans - @asynccontextmanager async def span( self, diff --git a/src/agentex/lib/core/tracing/tracer.py b/src/agentex/lib/core/tracing/tracer.py index 3af79977e..5dcdf3399 100644 --- a/src/agentex/lib/core/tracing/tracer.py +++ b/src/agentex/lib/core/tracing/tracer.py @@ -15,12 +15,12 @@ class Tracer: It manages the client connection and creates traces. """ - def __init__(self, client: Agentex): + def __init__(self, client: Agentex | None = None): """ - Initialize a new sync tracer with the provided client. + Initialize a new sync tracer. Args: - client: Agentex client instance used for API communication. + client: Kept for backward compatibility, no longer used. """ self.client = client @@ -47,12 +47,12 @@ class AsyncTracer: It manages the async client connection and creates async traces. """ - def __init__(self, client: AsyncAgentex): + def __init__(self, client: AsyncAgentex | None = None): """ - Initialize a new async tracer with the provided client. + Initialize a new async tracer. Args: - client: AsyncAgentex client instance used for API communication. + client: Kept for backward compatibility, no longer used. """ self.client = client diff --git a/src/agentex/lib/core/tracing/tracing_processor_manager.py b/src/agentex/lib/core/tracing/tracing_processor_manager.py index 07c440313..3e91e6c61 100644 --- a/src/agentex/lib/core/tracing/tracing_processor_manager.py +++ b/src/agentex/lib/core/tracing/tracing_processor_manager.py @@ -1,7 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING -from threading import Lock +from threading import RLock from agentex.lib.types.tracing import TracingProcessorConfig from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( @@ -13,17 +12,9 @@ AsyncTracingProcessor, ) -if TYPE_CHECKING: - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( # noqa: F401 - AgentexSyncTracingProcessor, - AgentexAsyncTracingProcessor, - ) - class TracingProcessorManager: def __init__(self): - # Mapping of processor config type to processor class - # Use lazy loading for agentex processors to avoid circular imports self.sync_config_registry: dict[str, type[SyncTracingProcessor]] = { "sgp": SGPSyncTracingProcessor, } @@ -33,23 +24,17 @@ def __init__(self): # Cache for processors self.sync_processors: list[SyncTracingProcessor] = [] self.async_processors: list[AsyncTracingProcessor] = [] - self.lock = Lock() - self._agentex_registered = False - - def _ensure_agentex_registered(self): - """Lazily register agentex processors to avoid circular imports.""" - if not self._agentex_registered: - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexSyncTracingProcessor, - AgentexAsyncTracingProcessor, - ) - self.sync_config_registry["agentex"] = AgentexSyncTracingProcessor - self.async_config_registry["agentex"] = AgentexAsyncTracingProcessor - self._agentex_registered = True + # Reentrant: set_processor_configs holds it while calling add_processor_config. + self.lock = RLock() def add_processor_config(self, processor_config: TracingProcessorConfig) -> None: with self.lock: - self._ensure_agentex_registered() + if processor_config.type not in self.sync_config_registry: + raise ValueError( + f"Unknown tracing processor type {processor_config.type!r}. " + f"Supported: {sorted(self.sync_config_registry)}. The Agentex span store " + "was removed, configure SGPTracingProcessorConfig instead." + ) sync_processor = self.sync_config_registry[processor_config.type] async_processor = self.async_config_registry[processor_config.type] self.sync_processors.append(sync_processor(processor_config)) @@ -73,8 +58,10 @@ def get_async_processors(self) -> list[AsyncTracingProcessor]: add_tracing_processor_config = GLOBAL_TRACING_PROCESSOR_MANAGER.add_processor_config set_tracing_processor_configs = GLOBAL_TRACING_PROCESSOR_MANAGER.set_processor_configs + def get_sync_tracing_processors(): return GLOBAL_TRACING_PROCESSOR_MANAGER.get_sync_processors() + def get_async_tracing_processors(): return GLOBAL_TRACING_PROCESSOR_MANAGER.get_async_processors() diff --git a/src/agentex/lib/types/tracing.py b/src/agentex/lib/types/tracing.py index 721d87794..d4d0eb4da 100644 --- a/src/agentex/lib/types/tracing.py +++ b/src/agentex/lib/types/tracing.py @@ -1,8 +1,9 @@ from __future__ import annotations -from typing import Literal, Annotated +from typing import Any, Literal +from datetime import datetime -from pydantic import Field +from pydantic import ConfigDict from agentex.lib.utils.model_utils import BaseModel @@ -20,8 +21,22 @@ class BaseModelWithTraceParams(BaseModel): parent_span_id: str | None = None -class AgentexTracingProcessorConfig(BaseModel): - type: Literal["agentex"] = "agentex" +class Span(BaseModel): + """In-memory span handed to tracing processors. Owned here, not by the generated client.""" + + # The generated model kept unknown keys, and custom processors may stash their own. + model_config = ConfigDict(extra="allow") + + id: str + name: str + start_time: datetime + trace_id: str + data: dict[str, Any] | list[dict[str, Any]] | None = None + end_time: datetime | None = None + input: dict[str, Any] | list[dict[str, Any]] | None = None + output: dict[str, Any] | list[dict[str, Any]] | None = None + parent_id: str | None = None + task_id: str | None = None class SGPTracingProcessorConfig(BaseModel): @@ -31,7 +46,4 @@ class SGPTracingProcessorConfig(BaseModel): sgp_base_url: str | None = None -TracingProcessorConfig = Annotated[ - AgentexTracingProcessorConfig | SGPTracingProcessorConfig, - Field(discriminator="type"), -] +TracingProcessorConfig = SGPTracingProcessorConfig diff --git a/tests/lib/adk/providers/test_litellm_usage.py b/tests/lib/adk/providers/test_litellm_usage.py index 5f5d480d9..bbce59f9e 100644 --- a/tests/lib/adk/providers/test_litellm_usage.py +++ b/tests/lib/adk/providers/test_litellm_usage.py @@ -12,7 +12,7 @@ from contextlib import asynccontextmanager from unittest.mock import AsyncMock, MagicMock -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.types.task_message import TaskMessage from agentex.lib.types.llm_messages import ( Delta, diff --git a/tests/lib/adk/test_langgraph_sync.py b/tests/lib/adk/test_langgraph_sync.py index 9e8c6e4f0..6bd331a78 100644 --- a/tests/lib/adk/test_langgraph_sync.py +++ b/tests/lib/adk/test_langgraph_sync.py @@ -251,7 +251,7 @@ class _FakeTracingBackend: spans_ended: list[str] = field(default_factory=list) async def start_span(self, **kw) -> Any: - from agentex.types.span import Span + from agentex.lib.types.tracing import Span sp = Span( id=f"span-{len(self.spans_started) + 1}", diff --git a/tests/lib/adk/test_tracing_activities.py b/tests/lib/adk/test_tracing_activities.py index 248ba94a7..b661e5009 100644 --- a/tests/lib/adk/test_tracing_activities.py +++ b/tests/lib/adk/test_tracing_activities.py @@ -5,7 +5,7 @@ from temporalio.testing import ActivityEnvironment -from agentex.types.span import Span +from agentex.lib.types.tracing import Span def _make_span(**overrides) -> Span: diff --git a/tests/lib/adk/test_tracing_module.py b/tests/lib/adk/test_tracing_module.py index c17ff5ff6..fe571ee36 100644 --- a/tests/lib/adk/test_tracing_module.py +++ b/tests/lib/adk/test_tracing_module.py @@ -7,7 +7,7 @@ from temporalio.exceptions import ActivityError import agentex.lib.adk._modules.tracing as _tracing_mod -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.harness.types import TurnUsage from agentex.lib.adk._modules.tracing import TurnSpan, TracingModule from agentex.lib.core.tracing.span_error import get_span_error diff --git a/tests/lib/adk/test_tracing_service.py b/tests/lib/adk/test_tracing_service.py index dceb000f5..ec41ad263 100644 --- a/tests/lib/adk/test_tracing_service.py +++ b/tests/lib/adk/test_tracing_service.py @@ -3,7 +3,7 @@ from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.services.adk.tracing import TracingService diff --git a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py index bf3dc8006..d6050224b 100644 --- a/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py +++ b/tests/lib/core/temporal/plugins/openai_agents/test_model_usage.py @@ -22,7 +22,7 @@ ) import agentex.lib.core.temporal.plugins.openai_agents.models.temporal_streaming_model as tsm -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.temporal.plugins.openai_agents.interceptors.context_interceptor import ( streaming_task_id, streaming_trace_id, diff --git a/tests/lib/core/tracing/processors/test_agentex_tracing_processor.py b/tests/lib/core/tracing/processors/test_agentex_tracing_processor.py deleted file mode 100644 index 84f37b495..000000000 --- a/tests/lib/core/tracing/processors/test_agentex_tracing_processor.py +++ /dev/null @@ -1,285 +0,0 @@ -from __future__ import annotations - -import asyncio -import weakref -from datetime import datetime, timezone -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -# AgentexAsyncTracingProcessor pulls in agentex.lib.adk via -# create_async_agentex_client, which in turn imports pydantic_ai at package -# init. Skip these tests cleanly when pydantic_ai isn't installed (the SDK -# dev venv state) so collection doesn't error out. -pytest.importorskip( - "pydantic_ai", - reason="agentex.lib.adk import chain requires pydantic_ai", -) - -# Import the processor module up front so unittest.mock.patch() can resolve -# attributes by string path. The tracing_processor_manager only loads this -# module lazily, so without this explicit import the patches below would fail -# with AttributeError at __enter__ time. -import agentex.lib.core.tracing.processors.agentex_tracing_processor # noqa: E402, F401 - -MODULE = "agentex.lib.core.tracing.processors.agentex_tracing_processor" - - -SKIP_ENV = "AGENTEX_TRACING_SKIP_AGENTEX_SPAN_START" - - -def _make_config() -> MagicMock: - """Empty config — AgentexTracingProcessorConfig is unused by __init__.""" - return MagicMock() - - -def _make_span(): - from agentex.types.span import Span - - now = datetime.now(timezone.utc) - return Span( - id="span-1", - trace_id="trace-1", - name="test-span", - start_time=now, - end_time=now, - input={"in": 1}, - output={"out": 2}, - ) - - -class TestAgentexSyncSkipSpanStart: - """The Agentex backend writes create-on-start + update-on-end by default. - End-only ingest (default) skips the start write and makes the END a single - create — verify the start is a no-op and end does an INSERT, not an UPDATE. - """ - - def test_start_skipped_and_end_creates_by_default(self, monkeypatch): - monkeypatch.delenv(SKIP_ENV, raising=False) # default ON - with patch(f"{MODULE}.Agentex") as MockAgentex: - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexSyncTracingProcessor, - ) - - processor = AgentexSyncTracingProcessor(_make_config()) - client = MockAgentex.return_value - span = _make_span() - - processor.on_span_start(span) - client.spans.create.assert_not_called() # start skipped - client.spans.update.assert_not_called() - - processor.on_span_end(span) - client.spans.create.assert_called_once() # single INSERT on end - client.spans.update.assert_not_called() # never a 404-prone UPDATE - - def test_start_creates_and_end_updates_when_skip_disabled(self, monkeypatch): - monkeypatch.setenv(SKIP_ENV, "0") - with patch(f"{MODULE}.Agentex") as MockAgentex: - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexSyncTracingProcessor, - ) - - processor = AgentexSyncTracingProcessor(_make_config()) - client = MockAgentex.return_value - span = _make_span() - - processor.on_span_start(span) - client.spans.create.assert_called_once() # start write restored - - processor.on_span_end(span) - client.spans.update.assert_called_once() # end is the UPDATE - - def test_skip_decision_captured_at_init_not_per_call(self, monkeypatch): - """The two halves of a span MUST use the same skip decision. A flag - toggled after construction must not split it (start-skip + end-update - would 404). The decision is captured once at init. - """ - monkeypatch.delenv(SKIP_ENV, raising=False) # construct with skip ON - with patch(f"{MODULE}.Agentex") as MockAgentex: - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexSyncTracingProcessor, - ) - - processor = AgentexSyncTracingProcessor(_make_config()) - client = MockAgentex.return_value - span = _make_span() - - processor.on_span_start(span) # skipped (cached ON) - monkeypatch.setenv(SKIP_ENV, "0") # toggle mid-span — must be ignored - processor.on_span_end(span) - - client.spans.create.assert_called_once() # still end-only INSERT - client.spans.update.assert_not_called() # NOT a 404-prone UPDATE - - -class TestAgentexAsyncSkipSpanStart: - async def test_start_skipped_and_end_creates_by_default(self, monkeypatch): - monkeypatch.delenv(SKIP_ENV, raising=False) # default ON - with patch(f"{MODULE}.create_async_agentex_client") as mock_factory: - client = MagicMock() - client.spans.create = AsyncMock() - client.spans.update = AsyncMock() - mock_factory.return_value = client - - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexAsyncTracingProcessor, - ) - - processor = AgentexAsyncTracingProcessor(_make_config()) - span = _make_span() - - await processor.on_span_start(span) - client.spans.create.assert_not_called() # start skipped - client.spans.update.assert_not_called() - - await processor.on_span_end(span) - client.spans.create.assert_awaited_once() # single INSERT on end - client.spans.update.assert_not_called() - - async def test_start_creates_and_end_updates_when_skip_disabled(self, monkeypatch): - monkeypatch.setenv(SKIP_ENV, "0") - with patch(f"{MODULE}.create_async_agentex_client") as mock_factory: - client = MagicMock() - client.spans.create = AsyncMock() - client.spans.update = AsyncMock() - mock_factory.return_value = client - - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexAsyncTracingProcessor, - ) - - processor = AgentexAsyncTracingProcessor(_make_config()) - span = _make_span() - - await processor.on_span_start(span) - client.spans.create.assert_awaited_once() # start write restored - - await processor.on_span_end(span) - client.spans.update.assert_awaited_once() # end is the UPDATE - - async def test_skip_decision_captured_at_init_not_per_call(self, monkeypatch): - """A flag toggled after construction must not split a span's lifecycle.""" - monkeypatch.delenv(SKIP_ENV, raising=False) # construct with skip ON - with patch(f"{MODULE}.create_async_agentex_client") as mock_factory: - client = MagicMock() - client.spans.create = AsyncMock() - client.spans.update = AsyncMock() - mock_factory.return_value = client - - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexAsyncTracingProcessor, - ) - - processor = AgentexAsyncTracingProcessor(_make_config()) - span = _make_span() - - await processor.on_span_start(span) # skipped (cached ON) - monkeypatch.setenv(SKIP_ENV, "0") # toggle mid-span — must be ignored - await processor.on_span_end(span) - - client.spans.create.assert_awaited_once() # still end-only INSERT - client.spans.update.assert_not_called() # NOT a 404-prone UPDATE - - -class TestAgentexAsyncTracingProcessor: - """Coverage for the per-event-loop client cache. The SGP processor has - matching tests; mirror them here so a regression in the Agentex side - (e.g. an accidental refactor that switches back to a plain dict, or - drops the lazy lookup) does not slip through unnoticed. - """ - - async def test_client_caches_per_event_loop(self): - """First access builds the client; subsequent accesses in the same - running loop must return the cached instance. - """ - with patch(f"{MODULE}.create_async_agentex_client") as mock_factory: - mock_factory.side_effect = lambda **kwargs: MagicMock() - - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexAsyncTracingProcessor, - ) - - processor = AgentexAsyncTracingProcessor(_make_config()) - - # Construction must not eagerly build the client (no running loop - # guarantee at module import time). - assert mock_factory.call_count == 0 - - c1 = processor.client - c2 = processor.client - c3 = processor.client - - assert mock_factory.call_count == 1, ( - f"Expected client to be built once per loop, but " - f"create_async_agentex_client was called {mock_factory.call_count} times" - ) - assert c1 is c2 is c3 - - async def test_client_keepalive_is_enabled(self): - """Regression guard: the per-loop client must use keepalive — the - whole reason for the per-loop cache. Verify max_keepalive_connections > 0. - """ - import httpx as _httpx - - captured_limits: list[_httpx.Limits] = [] - original_async_client = _httpx.AsyncClient - - def capture_limits(*args, **kwargs): - limits = kwargs.get("limits") - if limits is not None: - captured_limits.append(limits) - return original_async_client(*args, **kwargs) - - with patch(f"{MODULE}.create_async_agentex_client") as mock_factory, patch( - "httpx.AsyncClient", side_effect=capture_limits - ): - mock_factory.side_effect = lambda **kwargs: MagicMock() - - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexAsyncTracingProcessor, - ) - - processor = AgentexAsyncTracingProcessor(_make_config()) - _ = processor.client - - assert len(captured_limits) == 1 - max_keepalive = captured_limits[0].max_keepalive_connections - assert max_keepalive is not None and max_keepalive > 0, ( - f"Agentex async client should have keepalive enabled, got " - f"max_keepalive_connections={max_keepalive}" - ) - - def test_cache_is_weakkeydict_and_evicts_dead_loops(self): - """Regression guard for the id()-reuse bug: the per-loop cache must - be a WeakKeyDictionary so a GC'd loop's entry is evicted. Otherwise - a new loop landing at the same memory address would reuse the dead - loop's client, reintroducing the "bound to a different event loop" - error the per-loop cache was built to prevent. - """ - import gc - - with patch(f"{MODULE}.create_async_agentex_client"): - from agentex.lib.core.tracing.processors.agentex_tracing_processor import ( - AgentexAsyncTracingProcessor, - ) - - processor = AgentexAsyncTracingProcessor(_make_config()) - - # Storage type itself: WeakKeyDictionary, not plain dict. - assert isinstance(processor._clients_by_loop, weakref.WeakKeyDictionary) - - # End-to-end check: insert under a loop, drop the loop, the entry - # must vanish after GC. - loop = asyncio.new_event_loop() - try: - processor._clients_by_loop[loop] = MagicMock() - assert len(processor._clients_by_loop) == 1 - finally: - loop.close() - del loop - gc.collect() - assert len(processor._clients_by_loop) == 0, ( - "WeakKeyDictionary should have evicted the dead loop's entry; " - "remaining keys would cause stale-client reuse on id() recycling." - ) diff --git a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py index 6cd324f01..cc79a6054 100644 --- a/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py +++ b/tests/lib/core/tracing/processors/test_sgp_tracing_processor.py @@ -7,8 +7,7 @@ import pytest -from agentex.types.span import Span -from agentex.lib.types.tracing import SGPTracingProcessorConfig +from agentex.lib.types.tracing import Span, SGPTracingProcessorConfig MODULE = "agentex.lib.core.tracing.processors.sgp_tracing_processor" @@ -84,11 +83,10 @@ def test_commit_sha_is_stamped_after_opt_in(self, monkeypatch): def test_commit_sha_does_not_leak_onto_the_shared_span(self, monkeypatch): """trace.py hands ONE Span to every processor. If the commit SHA were - written onto span.data, a co-registered Agentex processor would - serialize it too, and it would surface in caller-visible span data.""" + written onto span.data, any co-registered processor would serialize it + too, and it would surface in caller-visible span data.""" from agentex.lib.core.tracing import code_revision from agentex.lib.core.tracing.processors.sgp_tracing_processor import _sgp_metadata - from agentex.lib.core.tracing.processors.agentex_tracing_processor import _create_kwargs monkeypatch.setenv("AGENT_COMMIT_SHA", self.SHA) code_revision.enable() @@ -96,7 +94,6 @@ def test_commit_sha_does_not_leak_onto_the_shared_span(self, monkeypatch): span = _make_span(); span.data = {} assert _sgp_metadata(span)["__commit_sha__"] == self.SHA # SGP sees it assert "__commit_sha__" not in span.data # the span does not - assert "__commit_sha__" not in (_create_kwargs(span)["data"] or {}) finally: code_revision.disable() diff --git a/tests/lib/core/tracing/processors/test_tracing_processor_interface.py b/tests/lib/core/tracing/processors/test_tracing_processor_interface.py index 12847b70d..dfa9d7b8b 100644 --- a/tests/lib/core/tracing/processors/test_tracing_processor_interface.py +++ b/tests/lib/core/tracing/processors/test_tracing_processor_interface.py @@ -5,8 +5,7 @@ from typing import override from datetime import UTC, datetime -from agentex.types.span import Span -from agentex.lib.types.tracing import TracingProcessorConfig +from agentex.lib.types.tracing import Span, TracingProcessorConfig from agentex.lib.core.tracing.processors.tracing_processor_interface import ( AsyncTracingProcessor, ) diff --git a/tests/lib/core/tracing/test_span_error.py b/tests/lib/core/tracing/test_span_error.py index 02e9645a4..ebbaf8fa5 100644 --- a/tests/lib/core/tracing/test_span_error.py +++ b/tests/lib/core/tracing/test_span_error.py @@ -12,7 +12,7 @@ CategorizedError as SGPCategorizedError, ) -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.tracing.trace import Trace, AsyncTrace from agentex.lib.core.tracing.span_error import ( SPAN_ERROR_KEY, diff --git a/tests/lib/core/tracing/test_span_model.py b/tests/lib/core/tracing/test_span_model.py new file mode 100644 index 000000000..6d356e9f1 --- /dev/null +++ b/tests/lib/core/tracing/test_span_model.py @@ -0,0 +1,21 @@ +from datetime import UTC, datetime + +from agentex.lib.types.tracing import Span + + +def _span(**extra) -> Span: + return Span(id="s1", name="n", trace_id="t1", start_time=datetime(2026, 1, 1, tzinfo=UTC), **extra) + + +def test_unknown_keys_survive_validation_and_a_json_round_trip(): + span = Span.model_validate({**_span().model_dump(), "extension": {"sampled": True}}) + + assert span.extension == {"sampled": True} # type: ignore[attr-defined] + assert Span.model_validate_json(span.model_dump_json()).model_dump()["extension"] == {"sampled": True} + + +def test_processors_can_attach_their_own_attributes(): + span = _span() + span.annotation = "custom" # type: ignore[attr-defined] + + assert span.model_copy(deep=True).model_dump()["annotation"] == "custom" diff --git a/tests/lib/core/tracing/test_span_queue.py b/tests/lib/core/tracing/test_span_queue.py index b8092daca..c5a4df248 100644 --- a/tests/lib/core/tracing/test_span_queue.py +++ b/tests/lib/core/tracing/test_span_queue.py @@ -7,7 +7,7 @@ from datetime import UTC, datetime from unittest.mock import AsyncMock, MagicMock, patch -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.tracing.span_queue import ( _DEFAULT_BATCH_SIZE, SpanEventType, diff --git a/tests/lib/core/tracing/test_span_queue_load.py b/tests/lib/core/tracing/test_span_queue_load.py index 652589881..edfe54b3a 100644 --- a/tests/lib/core/tracing/test_span_queue_load.py +++ b/tests/lib/core/tracing/test_span_queue_load.py @@ -41,7 +41,7 @@ import pytest -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.tracing.trace import AsyncTrace from agentex.lib.core.tracing.span_queue import AsyncSpanQueue diff --git a/tests/lib/core/tracing/test_tracing_processor_manager.py b/tests/lib/core/tracing/test_tracing_processor_manager.py new file mode 100644 index 000000000..158a70ca1 --- /dev/null +++ b/tests/lib/core/tracing/test_tracing_processor_manager.py @@ -0,0 +1,68 @@ +import threading +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from agentex.lib.types.tracing import SGPTracingProcessorConfig +from agentex.lib.core.tracing.tracing_processor_manager import TracingProcessorManager +from agentex.lib.core.tracing.processors.sgp_tracing_processor import ( + SGPSyncTracingProcessor, + SGPAsyncTracingProcessor, +) + +SGP_MODULE = "agentex.lib.core.tracing.processors.sgp_tracing_processor" + + +def _sgp_config() -> SGPTracingProcessorConfig: + return SGPTracingProcessorConfig(sgp_api_key="k", sgp_account_id="a", sgp_base_url="http://sgp.test") + + +def _patched_sgp(): + env = MagicMock() + env.refresh.return_value = MagicMock(ACP_TYPE=None, AGENT_NAME=None, AGENT_ID=None, AGENT_VERSION=None) + return ( + patch(f"{SGP_MODULE}.SGPClient"), + patch(f"{SGP_MODULE}.AsyncSGPClient"), + patch(f"{SGP_MODULE}.tracing.init"), + patch(f"{SGP_MODULE}.EnvironmentVariables", env), + ) + + +def test_unknown_processor_type_is_rejected_by_name(): + manager = TracingProcessorManager() + + with pytest.raises(ValueError, match="agentex.*sgp"): + manager.add_processor_config(SimpleNamespace(type="agentex")) # type: ignore[arg-type] + + assert manager.get_sync_processors() == [] + assert manager.get_async_processors() == [] + + +def test_sgp_config_registers_one_sync_and_one_async_processor(): + p1, p2, p3, p4 = _patched_sgp() + with p1, p2, p3, p4: + manager = TracingProcessorManager() + manager.add_processor_config(_sgp_config()) + + (sync_processor,) = manager.get_sync_processors() + (async_processor,) = manager.get_async_processors() + assert isinstance(sync_processor, SGPSyncTracingProcessor) + assert isinstance(async_processor, SGPAsyncTracingProcessor) + + +def test_set_processor_configs_registers_every_config_without_deadlocking(): + p1, p2, p3, p4 = _patched_sgp() + manager = TracingProcessorManager() + done = threading.Event() + + def register(): + with p1, p2, p3, p4: + manager.set_processor_configs([_sgp_config(), _sgp_config()]) + done.set() + + threading.Thread(target=register, daemon=True).start() + + assert done.wait(timeout=5), "set_processor_configs hung: the manager lock must be reentrant" + assert len(manager.get_sync_processors()) == 2 + assert len(manager.get_async_processors()) == 2 diff --git a/tests/test_adk_tracing_span_error.py b/tests/test_adk_tracing_span_error.py index c81015142..c8243cf41 100644 --- a/tests/test_adk_tracing_span_error.py +++ b/tests/test_adk_tracing_span_error.py @@ -21,7 +21,7 @@ import pytest -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.adk._modules.tracing import TracingModule from agentex.lib.core.tracing.span_error import get_span_error diff --git a/tests/test_obs_handle_registry.py b/tests/test_obs_handle_registry.py index 02d3adf0a..0c40f8068 100644 --- a/tests/test_obs_handle_registry.py +++ b/tests/test_obs_handle_registry.py @@ -25,7 +25,7 @@ ) import agentex.lib.core.tracing.trace as trace_mod -from agentex.types.span import Span +from agentex.lib.types.tracing import Span from agentex.lib.core.tracing.trace import _OBS_HANDLES, _OBS_HANDLES_MAX, Trace from agentex.lib.core.tracing.obs_span import ObsSpanHandle