From d7d87206022d0b4c4eb0a69e92a80b96a0d97fa8 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Fri, 11 Sep 2026 09:09:44 -0400 Subject: [PATCH 1/2] refactor(tracing): drop the Agentex Postgres span processor and own the Span model in lib The Agentex server is retiring its Postgres-backed /spans API, so the processor that wrote to it goes, along with the Trace helpers that read spans back through the client. The SGP processor is the only sink. The in-memory Span handed to processors was the generated client type for that API. It now lives in agentex.lib.types.tracing and is still exported from agentex.lib.core.tracing, so it survives the generated surface disappearing when the server spec drops the routes. Co-Authored-By: Claude Fable 5.1 --- .../project/state_machines/deep_research.py | 2 +- src/agentex/lib/adk/_modules/tracing.py | 2 +- src/agentex/lib/core/services/adk/tracing.py | 2 +- .../activities/adk/tracing_activities.py | 2 +- src/agentex/lib/core/tracing/__init__.py | 2 +- .../processors/agentex_tracing_processor.py | 232 -------------- .../processors/sgp_tracing_processor.py | 3 +- .../processors/tracing_processor_interface.py | 3 +- src/agentex/lib/core/tracing/span_error.py | 2 +- src/agentex/lib/core/tracing/span_queue.py | 2 +- src/agentex/lib/core/tracing/trace.py | 52 +--- .../core/tracing/tracing_processor_manager.py | 30 +- src/agentex/lib/types/tracing.py | 25 +- tests/lib/adk/providers/test_litellm_usage.py | 2 +- tests/lib/adk/test_langgraph_sync.py | 2 +- tests/lib/adk/test_tracing_activities.py | 2 +- tests/lib/adk/test_tracing_module.py | 2 +- tests/lib/adk/test_tracing_service.py | 2 +- .../plugins/openai_agents/test_model_usage.py | 2 +- .../test_agentex_tracing_processor.py | 285 ------------------ .../processors/test_sgp_tracing_processor.py | 9 +- .../test_tracing_processor_interface.py | 3 +- tests/lib/core/tracing/test_span_error.py | 2 +- tests/lib/core/tracing/test_span_queue.py | 2 +- .../lib/core/tracing/test_span_queue_load.py | 2 +- .../tracing/test_tracing_processor_manager.py | 15 + tests/test_adk_tracing_span_error.py | 2 +- tests/test_obs_handle_registry.py | 2 +- 28 files changed, 64 insertions(+), 629 deletions(-) delete mode 100644 src/agentex/lib/core/tracing/processors/agentex_tracing_processor.py delete mode 100644 tests/lib/core/tracing/processors/test_agentex_tracing_processor.py create mode 100644 tests/lib/core/tracing/test_tracing_processor_manager.py 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..7b6c5dcf4 100644 --- a/src/agentex/lib/adk/_modules/tracing.py +++ b/src/agentex/lib/adk/_modules/tracing.py @@ -22,7 +22,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 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..b0546d1b1 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 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..376bcca35 100644 --- a/src/agentex/lib/core/tracing/span_error.py +++ b/src/agentex/lib/core/tracing/span_error.py @@ -9,7 +9,7 @@ ) 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__`` / 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..dddf5643a 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 @@ -220,7 +220,6 @@ 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. """ def __init__( @@ -268,7 +267,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,30 +325,7 @@ 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( @@ -382,7 +357,6 @@ 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. """ def __init__( @@ -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,30 +480,7 @@ 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( diff --git a/src/agentex/lib/core/tracing/tracing_processor_manager.py b/src/agentex/lib/core/tracing/tracing_processor_manager.py index 07c440313..6ac80a032 100644 --- a/src/agentex/lib/core/tracing/tracing_processor_manager.py +++ b/src/agentex/lib/core/tracing/tracing_processor_manager.py @@ -1,6 +1,5 @@ from __future__ import annotations -from typing import TYPE_CHECKING from threading import Lock from agentex.lib.types.tracing import TracingProcessorConfig @@ -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, } @@ -34,22 +25,15 @@ def __init__(self): 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 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 +57,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..36f0137e8 100644 --- a/src/agentex/lib/types/tracing.py +++ b/src/agentex/lib/types/tracing.py @@ -1,8 +1,7 @@ from __future__ import annotations -from typing import Literal, Annotated - -from pydantic import Field +from typing import Any, Literal +from datetime import datetime from agentex.lib.utils.model_utils import BaseModel @@ -20,8 +19,19 @@ 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.""" + + 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 +41,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_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..e48cfd345 --- /dev/null +++ b/tests/lib/core/tracing/test_tracing_processor_manager.py @@ -0,0 +1,15 @@ +from types import SimpleNamespace + +import pytest + +from agentex.lib.core.tracing.tracing_processor_manager import TracingProcessorManager + + +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() == [] 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 From 62b97be67792f3eda1610d66d5cb60acc6eb9659 Mon Sep 17 00:00:00 2001 From: Mohammad Atallah Date: Fri, 11 Sep 2026 10:05:54 -0400 Subject: [PATCH 2/2] fix(tracing): keep extra span keys, make the processor lock reentrant, drop the tracer's dead client Review follow-ups on the span-processor removal: - Span keeps unknown keys like the generated model it replaces, so a custom processor's own attributes and older worker payloads survive. - set_processor_configs re-enters add_processor_config under the manager lock, which deadlocked on a plain Lock. Reentrant now, with tests for the SGP happy path and the batch registration. - Trace and Tracer no longer use the client they receive, so it is optional and the ADK tracing module stops building an httpx client per event loop. - Changelog entry for the breaking removals, stale prose updated. Co-Authored-By: Claude Fable 5.1 --- CHANGELOG.md | 1 + src/agentex/lib/adk/_modules/tracing.py | 36 ++----------- .../processors/sgp_tracing_processor.py | 4 +- src/agentex/lib/core/tracing/span_error.py | 9 ++-- src/agentex/lib/core/tracing/trace.py | 26 +++++---- src/agentex/lib/core/tracing/tracer.py | 12 ++--- .../core/tracing/tracing_processor_manager.py | 5 +- src/agentex/lib/types/tracing.py | 5 ++ tests/lib/core/tracing/test_span_model.py | 21 ++++++++ .../tracing/test_tracing_processor_manager.py | 53 +++++++++++++++++++ 10 files changed, 109 insertions(+), 63 deletions(-) create mode 100644 tests/lib/core/tracing/test_span_model.py 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/src/agentex/lib/adk/_modules/tracing.py b/src/agentex/lib/adk/_modules/tracing.py index 7b6c5dcf4..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 ( @@ -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/tracing/processors/sgp_tracing_processor.py b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py index b0546d1b1..124e69b29 100644 --- a/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py +++ b/src/agentex/lib/core/tracing/processors/sgp_tracing_processor.py @@ -74,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/span_error.py b/src/agentex/lib/core/tracing/span_error.py index 376bcca35..81cb0edfd 100644 --- a/src/agentex/lib/core/tracing/span_error.py +++ b/src/agentex/lib/core/tracing/span_error.py @@ -12,12 +12,9 @@ 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/trace.py b/src/agentex/lib/core/tracing/trace.py index dddf5643a..4e3fbd393 100644 --- a/src/agentex/lib/core/tracing/trace.py +++ b/src/agentex/lib/core/tracing/trace.py @@ -218,22 +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. + 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 @@ -251,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. @@ -325,8 +327,6 @@ def end_span( return span - - @contextmanager def span( self, @@ -355,14 +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. + 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, ): @@ -391,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. @@ -480,8 +480,6 @@ async def end_span( return span - - @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 6ac80a032..3e91e6c61 100644 --- a/src/agentex/lib/core/tracing/tracing_processor_manager.py +++ b/src/agentex/lib/core/tracing/tracing_processor_manager.py @@ -1,6 +1,6 @@ from __future__ import annotations -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 ( @@ -24,7 +24,8 @@ def __init__(self): # Cache for processors self.sync_processors: list[SyncTracingProcessor] = [] self.async_processors: list[AsyncTracingProcessor] = [] - self.lock = Lock() + # 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: diff --git a/src/agentex/lib/types/tracing.py b/src/agentex/lib/types/tracing.py index 36f0137e8..d4d0eb4da 100644 --- a/src/agentex/lib/types/tracing.py +++ b/src/agentex/lib/types/tracing.py @@ -3,6 +3,8 @@ from typing import Any, Literal from datetime import datetime +from pydantic import ConfigDict + from agentex.lib.utils.model_utils import BaseModel @@ -22,6 +24,9 @@ class BaseModelWithTraceParams(BaseModel): 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 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_tracing_processor_manager.py b/tests/lib/core/tracing/test_tracing_processor_manager.py index e48cfd345..158a70ca1 100644 --- a/tests/lib/core/tracing/test_tracing_processor_manager.py +++ b/tests/lib/core/tracing/test_tracing_processor_manager.py @@ -1,8 +1,32 @@ +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(): @@ -13,3 +37,32 @@ def test_unknown_processor_type_is_rejected_by_name(): 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