Skip to content

feat(tracing): add W3C distributed tracing and incremental handler batching - #225

Open
pradystar wants to merge 7 commits into
mainfrom
feat/w3c-trace-context-migration
Open

feat(tracing): add W3C distributed tracing and incremental handler batching#225
pradystar wants to merge 7 commits into
mainfrom
feat/w3c-trace-context-migration

Conversation

@pradystar

Copy link
Copy Markdown
Collaborator

Summary

This PR completes the SDK’s W3C distributed-tracing migration and satisfies both distributed-tracing requirements:

  • DT.A — Distributed context propagation: SDK-native telemetry, native OpenTelemetry spans, and supported upstream OTel HTTP instrumentation can participate in the same trace across service boundaries.
  • DT.B — Incremental span delivery: Completed native-handler operations enter the existing OpenTelemetry BatchSpanProcessor when their callbacks end without waiting for the entire logical trace to complete.

Both automatic and explicit propagation remain supported.

Detailed requirements mapping, lifecycle behavior, compatibility analysis, deletion rationale, and validation evidence are available in the detailed review context PR_REVIEW_DETAILS.md

What changed

Standard W3C propagation

Distributed context now uses:

  • traceparent
  • tracestate
  • W3C baggage

This replaces the proprietary Splunk-AO-Trace-ID and Splunk-AO-Parent-ID propagation format and allows Splunk AO telemetry to interoperate with upstream OpenTelemetry instrumentation and non-Python services.

Automatic distributed tracing

Applications can configure supported automatic instrumentation with:

from fastapi import FastAPI
from splunk_ao import configure_distributed_tracing

app = FastAPI()
tracer_provider = configure_distributed_tracing(app=app)

This:

  • Creates or accepts a concrete TracerProvider.
  • Registers the Splunk AO span processor once.
  • Enables installed supported HTTP instrumentors.
  • Instruments the supplied FastAPI or Starlette application.
  • Returns the application-owned provider for normal shutdown.
  • Does not replace the process-global tracer provider.

Install the optional dependencies with:

pip install "splunk-ao[distributed-tracing]"

Automatic incoming extraction is supported for FastAPI and Starlette.

Automatic outgoing injection is supported for:

  • Requests
  • HTTPX sync and async clients
  • aiohttp

Applications using other frameworks or custom transports can register the corresponding upstream OTel instrumentor or continue using explicit propagation.

Explicit propagation remains supported

Existing explicit integrations can continue to use:

from splunk_ao import get_tracing_headers

Incoming context can continue to use extract_tracing_context() or TracingMiddleware.

These APIs now use the same W3C context as automatic instrumentation. The explicit approach is not deprecated.

Incremental native-handler completion

LangChain, CrewAI, Google ADK, and OpenAI Agents operations now enter the existing span-processing pipeline when their individual callbacks end.

A completed child can therefore become eligible for scheduled or size-based BSP export while its parent agent or workflow remains active.

This does not mean every callback causes an immediate network request. The existing BatchSpanProcessor still determines network-export timing according to standard BSP configuration.

No per-trace flush() is required, and flush() does not end active work.

Control-plane HTTP suppression

SDK-owned authentication, health-check, routing, CRUD, token-refresh, and related HTTP operations are scoped under standard OTel HTTP suppression.

This prevents SDK control-plane calls from appearing as unrelated application GET or POST traces when automatic HTTP instrumentation is enabled.

Suppression is scoped to the SDK request and does not suppress application HTTP traffic.

Compatibility guarantees

This change preserves:

  • Existing single-service parent-child relationships.
  • Existing local trace structure with and without distributed tracing.
  • Logger, @log, and native handler instrumentation.
  • SDK-native and caller-owned OTel instrumentation.
  • Existing standalone and O11y deployment routing.
  • Caller ownership of caller-created providers.
  • The rule that the internal trace envelope is not exported as a span.
  • The behavior that flush() drains completed work without ending active operations.
  • ingestion_hook whole-tree compatibility.
  • flush_on_chain_end compatibility.

Distributed tracing adds remote ancestry without otherwise reshaping the application’s local telemetry tree.

start_new_trace remains a handler ownership option, not a distributed-tracing switch. Users do not need to change its default value to enable distributed tracing.

Session propagation and privacy

An explicit session is propagated through W3C baggage as:

gen_ai.conversation.id

The implementation does not propagate project, Agent Stream, agent, experiment, application, routing, endpoint, authentication, prompt, response, or embedding data through baggage.

Intentional breaking changes

The migration removes the proprietary propagation surface:

  • trace_id= and span_id= arguments on SplunkAOLogger. Note that all example might not be updated and examples might be stale. The examples will be fixed in a future PR.
  • The logger instance method get_tracing_headers().
  • Splunk-AO-Trace-ID and Splunk-AO-Parent-ID.

Applications should use the module-level W3C helper:

from splunk_ao import get_tracing_headers

or enable automatic propagation through configure_distributed_tracing().

Removed delivery path

The previous distributed-mode REST task handler maintained a second queue, dependency-ordering system, retry lifecycle, and shutdown path.

Normally exported telemetry now follows one processing path:

completed operation
→ SpanSink
→ OpenTelemetry BatchSpanProcessor
→ deployment-aware OTLP exporter

The obsolete task handler and its dedicated tests were removed. This does not remove any logger, decorator, or native framework-handler support.

Validation summary

Completed validation includes:

  • Full root suite: 2,163 passed, 4 skipped
  • LangChain and CrewAI compatibility: 202 passed
  • Google ADK suite: 252 passed
  • A2A suite: 68 passed
  • Final focused root regression set: 212 passed
  • Automatic instrumentation tests: 20 passed
  • Control-plane suppression and O11y HTTP tests: 41 passed
  • Ruff lint and formatting checks: passed
  • Mypy: 118 source files passed
  • Poetry lock consistency: passed
  • git diff --check: passed

Live standalone validation succeeded for both automatic and explicit cross-service distributed tracing.

@pradystar pradystar changed the title feat(tracing): add automatic W3C distributed tracing and incremental handler batching feat(tracing): add W3C distributed tracing and incremental handler batching Aug 13, 2026

@fercor-cisco fercor-cisco left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 This review was generated by the Astra agent (claude-sonnet-5). It may contain mistakes.

Verdict: request_changes — A concrete id-reuse aliasing bug in the new automatic-instrumentation bookkeeping (http_instrumentation.py) can cause silent double-instrumentation, false conflict errors, or unbounded growth of module-level state; there are also stale/incorrect docs left over from the removed proprietary header scheme.

General Comments

  • 🟠 major (bug): http_instrumentation.py tracks instrumentation state using id(app) / id(tracer_provider) as keys in plain module-level dict/set structures (_client_provider_ids, _instrumented_apps) without holding any reference (strong or weak) to the underlying objects. _configured_providers correctly uses a WeakSet to avoid this problem, but the other two do not. Once an app/tracer_provider is garbage-collected, CPython can and does reuse its id() for an unrelated object; a subsequent unrelated app/provider can then collide with a stale entry, causing _instrument_app/_validate_client_ownership to silently skip instrumentation (treating the new object as 'already instrumented') or to raise a false 'already instrumented through Splunk AO with another tracer provider' RuntimeError for objects that never actually conflicted. This is realistic in any process that creates many short-lived FastAPI apps or TracerProviders (tests, multi-tenant setups, hot-reload dev servers) and is also a slow memory leak since entries are never evicted. Recommend keying by id(obj) only while also holding a weakref to the object (or using WeakValueDictionary/an id->weakref map with a weakref.finalize cleanup callback) so aliasing cannot occur and stale entries are pruned automatically, matching the pattern already used for _configured_providers.
  • 🟡 minor (documentation): splunk-ao-migration-tool/README.md (section 6, 'HTTP Tracing Headers') still instructs users to migrate X-Galileo-Trace-ID/X-Galileo-Parent-ID to Splunk-AO-Trace-ID/Splunk-AO-Parent-ID, and says 'The get_tracing_headers() function return value now uses the new header names.' This PR removes Splunk-AO-Trace-ID/Splunk-AO-Parent-ID entirely in favor of W3C traceparent/tracestate, so this guidance is now wrong and will mislead migrating users into propagating headers the SDK no longer understands. Please update this section to point at the W3C headers (or get_tracing_headers()'s new W3C-based output) in the same change, per AGENTS.md's requirement to update docs when propagation/telemetry paths change.

Follow-ups

Suggested follow-up work that could be tracked as Jira tickets:

  • src/splunk_ao/decorator.py:1343-1373: _session_id_context (now defined in session_context.py) is a single global ContextVar shared across all SplunkAOLogger instances within a context, rather than being scoped per-logger. _set_active_session_id/set_session_context therefore make one logger's set_session/clear_session call affect get_effective_session_id() for any other logger sharing the same async/thread context (e.g. two loggers for different agent streams created in the same request). This may be intentional per the 'one explicit session, request-local' design, but is worth a design discussion/doc note since it's a behavior change from the previous per-instance self.session_id semantics.
  • src/splunk_ao/middleware/tracing.py:16-17: The module docstring's usage example calls logger.conclude(output=str(result)) twice in a row. If unintentional this is a confusing copy-paste artifact in documentation; if intentional (concluding the workflow span then the trace) it deserves a comment explaining why, since readers copying the example verbatim may not realize two distinct steps are being concluded.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

src/splunk_ao/http_instrumentation.py:2441-2443 (line not in diff)

🟠 major (bug): _client_provider_ids and _instrumented_apps key on id(tracer_provider)/id(app) without keeping the objects alive or otherwise validating identity. After the underlying object is garbage collected, Python may reuse the same id for a new, unrelated app/provider, causing silent skip-instrumentation or false ownership-conflict errors for that new object. Track a weakref.ref (or use weakref.finalize to prune the entry when the original object dies) alongside the id, or switch to WeakValueDictionary/WeakSet keyed by the object itself the way _configured_providers already does.

Suggested change
_client_provider_ids: dict[str, tuple[int, "weakref.ReferenceType[Any]"]] = {}
_instrumented_apps: WeakSet[Any] = WeakSet() # store (app, provider, framework) via a wrapper holding weakrefs

🤖 Generated by the Astra agent

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 major (bug): Module-level _client_provider_ids: dict[str, int] and _instrumented_apps: set[tuple[int, int, str]] store raw id() values with no reference back to the objects, unlike _configured_providers which is a WeakSet. This is an id-reuse aliasing hazard and an unbounded-growth leak; see PR-level comment for details and a suggested fix (weak references / finalizers keyed alongside the ids).

🤖 Generated by the Astra agent

@fercor-cisco fercor-cisco left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 This review was generated by the Astra agent (claude-opus-5). It may contain mistakes.

Verdict: request_changes — Several verified defects in new code: the distributed-tracing extra can't actually import its own instrumentors, one bad span aborts the whole handler trace (reproducible from the shipped ADK path), clear_session() cannot clear an inbound baggage session, unfinished callback spans leak OTel context activations, and two flush tests no longer exercise the code they name.

General Comments

  • 🟠 major (question): Was _activate_handler_step validated against the async LangChain callback dispatch path?

The new activation mechanism (logger.py:644-665) exists so that the handler span becomes the active W3C parent for outbound HTTP made while a callback is in flight. That requires the otel_context.attach() to take effect in the application's execution context.

But SplunkAOAsyncCallback does not set run_inline, so langchain-core dispatches async callbacks through asyncio.gather(*coros). gather wraps each coroutine in a Task, which copies the current contextvars.Context; mutations inside the Task do not propagate back to the caller. The same applies to sync handlers invoked from ainvoke, which langchain-core runs via run_in_executor(copy_context().run, ...).

If that is right, then for the async paths (a) the attach is invisible to the application, so automatic outbound propagation from inside a handler callback silently does nothing, and (b) the matching detach at end-callback time runs in a different Context, which is what triggers the ERROR-noise problem flagged on logger.py:659. Note _sync_otel_context_impl already anticipates exactly this class of failure (logger.py:496-499) — the new activation path does not.

Exported parent/child topology is unaffected (it comes from _otel_ids, not live context), so this is specifically about the propagation guarantee. Please confirm with a test that drives the real async LangChain callbacks and asserts the injected traceparent parent-id matches the active handler span — the live validation described in the PR body appears to have used @log/openai, not the async framework handlers.

  • 🟠 major (testing): The +183 lines of new OpenAI Agents lifecycle code are never driven through the public TracingProcessor interface. grep -rn "on_span_start" tests/ is empty. tests/test_openai_agents.py:46-101 calls _start_owned_root / _start_incremental_span / _finish_incremental_span directly with hand-built Nodes, test_simple_agent uses ingestion_hook (legacy branch), and test_complex_agent is @pytest.mark.skip.

That is why shape mismatches like the output=None case flagged on span_lifecycle.py aren't caught: the tests construct nodes that are known-good rather than nodes the framework actually produces. Please add at least one test that feeds real Span/Trace objects through on_trace_starton_span_starton_span_endon_trace_end on the non-hook path and asserts the exported span set and that _active_steps is empty afterwards.

  • 🟡 minor (testing): install_session_propagator() mutates the process-global textmap propagator, and _client_providers / _instrumented_apps / _configured_providers are module-level. These are reset only by the autouse fixture local to tests/test_http_instrumentation.py:119-133; tests/conftest.py has no guard.

Any future test (or an existing one that transitively calls configure_distributed_tracing) outside that file will permanently wrap the global propagator for the rest of the session and pin real TracerProviders in the strong _client_providers dict. Given -n auto xdist and the project's own rule to "reset global OTel context, providers/processors, SDK configuration ... after tests", this belongs in conftest.py as an autouse fixture rather than in one test module.

  • 🟡 minor (documentation): The new examples live under examples/logging-samples/DT_2.0/. DT_2.0 is an internal requirement codename that means nothing to a public contributor, which conflicts with AGENTS.md's "Do not document unavailable internal context. Repository documentation must stand alone."

It also duplicates the pre-existing examples/logging-samples/distributed-tracing/, which still ships SPLUNK_AO_MODE=distributed in its .env.example and README — so the repo now contains two DT guides that contradict each other on configuration. Suggest naming the new directories by capability (e.g. distributed-tracing-w3c/ and distributed-tracing-auto/) and either updating or deleting the old sample in this PR, since the PR body's "examples will be fixed in a future PR" note doesn't cover a duplicate that this PR itself creates.

Follow-ups

Suggested follow-up work that could be tracked as Jira tickets:

  • src/splunk_ao/handlers/base_async_handler.py:31-95: async_commit/async_end_node duplicate ~55 lines of commit/end_node from the sync base class, differing only in awaiting async_flush. The duplication has already drifted: the _owned_root is not None branch uses serialize_to_str(root_output) for the envelope output (line 66) where the sync path uses SplunkAOLogger._coerce_output(root_output) (base_handler.py:103), so a root output that is a list of ContentBlocks is stringified on the async path and preserved on the sync path. The divergence predates this PR, but this PR copy-pasted the new branch structure into both, doubling the surface. Consider extracting the shared body into a template method taking a flush callable, and aligning the output coercion.
  • src/splunk_ao/handlers/span_lifecycle.py:42-48: _step_number is now implemented twice with different behaviour: span_lifecycle._step_number swallows silently, SplunkAOBaseHandler._step_number (base_handler.py:283-292) logs a warning, and log_node_tree (base_handler.py:379-384) inlines a third copy. Consolidate on one helper so the LangGraph step-number rule can't drift between the incremental and legacy paths.
  • src/splunk_ao/utils/singleton.py:87-120: Removing trace_id/span_id from the singleton cache key also removed the only per-request differentiation for the decorator path — the deleted docstring said the key existed "for proper isolation of concurrent requests in async web servers." @log on an async endpoint now shares one cached SplunkAOLogger across concurrent requests on the same key. Per-instance ContextVar parent stacks limit the damage, but self.traces and self.session_id remain shared instance state, and the shipped examples/logging-samples/distributed-tracing/retrieval_service.py:56 uses exactly that pattern. Worth a design note or a request-scoped key that doesn't depend on the removed proprietary IDs.
  • src/splunk_ao/handlers/openai_agents/handler.py:60-64: SplunkAOTracingProcessor is installed process-wide but keeps single-valued per-trace state (_owned_trace, _caller_parent, _owned_root, _owned_root_node_id, _nodes, _active_steps). Partly pre-existing for _nodes/_owned_trace, but this PR adds three more fields with the same shape. Keying all per-trace state by trace_id would make concurrent Runner.run(...) calls safe and is a prerequisite for the on_trace_end cleanup fix flagged inline.

Comment thread pyproject.toml
crewai = ["crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4); python_version < '3.14'", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)"]
middleware = ["starlette"]
all = ["langchain-core", "langchain", "langsmith (>=0.8.0)", "openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "crewai (>=0.152.0,<2.0.0); python_version < '3.14'", "starlette", "litellm (>=1.83.14,<2.0.0); python_version < '3.14'", "uv (>=0.9.6); python_version < '3.14'", "aiohttp (>=3.14.1,<4); python_version < '3.14'", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)"]
distributed-tracing = ["aiohttp (>=3.14.1,<4)", "opentelemetry-instrumentation-fastapi (==0.59b0)", "opentelemetry-instrumentation-starlette (==0.59b0)", "opentelemetry-instrumentation-requests (==0.59b0)", "opentelemetry-instrumentation-httpx (==0.59b0)", "opentelemetry-instrumentation-aiohttp-client (==0.59b0)"]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 major (bug): The distributed-tracing extra cannot import its own instrumentors. _load_instrumentors() (src/splunk_ao/http_instrumentation.py:27-43) eagerly imports all five upstream instrumentors. opentelemetry-instrumentation-fastapi imports fastapi at module scope (it defines _InstrumentedFastAPI(fastapi.FastAPI)), which is why it declares instruments = ["fastapi (>=0.92,<1.0)"] (poetry.lock:3229) rather than a hard dependency. But fastapi is only a dev dependency here (pyproject.toml:103) and appears in no extra.

Failure scenario: a Flask or CLI application that only wants outbound propagation runs pip install "splunk-ao[distributed-tracing]", then configure_distributed_tracing()ImportError("Automatic distributed tracing requires optional dependencies. Install them with: pip install 'splunk-ao[distributed-tracing]'") — telling the user to install the extra they just installed. [all] is affected too (it carries starlette but never fastapi). CI doesn't catch it because fastapi is a dev dep and every test patches _load_instrumentors.

Add fastapi to both extras, and see the companion note on http_instrumentation.py:27 for making the import lazy so instrument_requests=True, instrument_httpx=False, instrument_aiohttp_client=False doesn't require every framework either.

🤖 Generated by the Astra agent

Comment on lines +27 to +43
def _load_instrumentors() -> dict[str, type]:
try:
from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.requests import RequestsInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.starlette import StarletteInstrumentor # noqa: PLC0415
except ImportError as exc:
raise ImportError(_INSTALL_MESSAGE) from exc

return {
"fastapi": FastAPIInstrumentor,
"starlette": StarletteInstrumentor,
"requests": RequestsInstrumentor,
"httpx": HTTPXClientInstrumentor,
"aiohttp-client": AioHttpClientInstrumentor,
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 major (design): _load_instrumentors() is all-or-nothing: it imports FastAPI, Starlette, HTTPX, Requests and aiohttp instrumentors regardless of what the caller asked for. Combined with the missing fastapi in the extra (see pyproject.toml:33), a caller who explicitly opts out of a transport still pays for its import.

Failure scenario: configure_distributed_tracing(instrument_httpx=False, instrument_aiohttp_client=False) in a process without aiohttp installed raises ImportError even though aiohttp instrumentation was explicitly declined.

Suggest resolving each instrumentor lazily for only the components actually requested, keeping the "resolve imports before constructing a processor" property that the comment at line 197-198 relies on:

Suggested change
def _load_instrumentors() -> dict[str, type]:
try:
from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.requests import RequestsInstrumentor # noqa: PLC0415
from opentelemetry.instrumentation.starlette import StarletteInstrumentor # noqa: PLC0415
except ImportError as exc:
raise ImportError(_INSTALL_MESSAGE) from exc
return {
"fastapi": FastAPIInstrumentor,
"starlette": StarletteInstrumentor,
"requests": RequestsInstrumentor,
"httpx": HTTPXClientInstrumentor,
"aiohttp-client": AioHttpClientInstrumentor,
}
_INSTRUMENTOR_IMPORTS: dict[str, tuple[str, str]] = {
"fastapi": ("opentelemetry.instrumentation.fastapi", "FastAPIInstrumentor"),
"starlette": ("opentelemetry.instrumentation.starlette", "StarletteInstrumentor"),
"requests": ("opentelemetry.instrumentation.requests", "RequestsInstrumentor"),
"httpx": ("opentelemetry.instrumentation.httpx", "HTTPXClientInstrumentor"),
"aiohttp-client": ("opentelemetry.instrumentation.aiohttp_client", "AioHttpClientInstrumentor"),
}
def _load_instrumentors(names: tuple[str, ...]) -> dict[str, type]:
"""Import only the instrumentors needed for the requested components."""
from importlib import import_module # noqa: PLC0415
resolved: dict[str, type] = {}
for name in names:
module_path, attribute = _INSTRUMENTOR_IMPORTS[name]
try:
resolved[name] = getattr(import_module(module_path), attribute)
except ImportError as exc:
raise ImportError(f"{_INSTALL_MESSAGE} (missing support for {name})") from exc
return resolved

🤖 Generated by the Astra agent

Comment on lines +140 to +147
def _reset_handler_state(self) -> None:
"""Release one completed callback tree without touching caller-owned state."""
self._nodes.clear()
self._active_steps.clear()
self._root_node = None
self._owned_trace = None
self._owned_root = None
self._owned_parent = None

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 major (bug): _reset_handler_state() drops unfinished _active_steps without detaching their OTel context activations or releasing their OTel identities. Compare handlers/openai_agents/handler.py:110-112, which was written to do exactly this cleanup — the base handler is missing it.

Failure scenario: a LangChain root chain fails, so on_chain_errorend_node(root)_finish_incremental_node(root)finally: _reset_handler_state(). If a child LLM run never received on_llm_end/on_llm_error (cancelled task, .astream() closed early, a framework that doesn't emit an error callback for that node type), its HandlerStepContext.token is never detached. That child's SpanContext therefore remains the current OTel span for the rest of the execution context: subsequent instrumented outbound HTTP and the next handler run get parented to a span that was never exported, and logger._otel_ids retains its entry.

Mirror the OpenAI Agents cleanup here (and reverse the iteration order — see the note on that file):

Suggested change
def _reset_handler_state(self) -> None:
"""Release one completed callback tree without touching caller-owned state."""
self._nodes.clear()
self._active_steps.clear()
self._root_node = None
self._owned_trace = None
self._owned_root = None
self._owned_parent = None
def _reset_handler_state(self) -> None:
"""Release one completed callback tree without touching caller-owned state."""
for state in reversed(list(self._active_steps.values())):
self._splunk_ao_logger._restore_handler_step_context(state.activation)
state.activation = None
self._splunk_ao_logger._release_otel_context(state.step)
self._nodes.clear()
self._active_steps.clear()
self._root_node = None
self._owned_trace = None
self._owned_root = None
self._owned_parent = None

🤖 Generated by the Astra agent

Comment on lines +98 to +118
if node.node_type in ("llm", "chat"):
return LoggedLlmSpan(
**common,
input=input_value,
output=output,
metrics=LlmMetrics.model_validate(
{
"duration_ns": duration_ns,
"num_input_tokens": params.get("num_input_tokens"),
"num_output_tokens": params.get("num_output_tokens"),
"num_total_tokens": params.get("num_total_tokens", params.get("total_tokens")),
"time_to_first_token_ns": params.get("time_to_first_token_ns"),
"num_reasoning_tokens": params.get("num_reasoning_tokens"),
"num_cached_input_tokens": params.get("num_cached_input_tokens"),
}
),
tools=params.get("tools"),
events=params.get("events"),
model=params.get("model"),
temperature=params.get("temperature"),
)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 major (bug): build_handler_step raises on output=None for an LLM node, and a shipped integration produces exactly that. LoggedLlmSpan.output is a non-optional LoggedMessage whose mode="before" validator ends in raise ValueError("LLM span output must be a Message, a string, or a dict.") for None.

Failure scenario: splunk-ao-adk/src/splunk_ao_adk/span_manager.py:132-134 normalizes an empty LLM output list to None and passes it straight to end_node(output=normalized_output). params.get("output", "") returns None (the key exists), so LoggedLlmSpan(output=None) raises. On the old path add_llm_span swallowed this and lost one span; on the new path the except at base_handler.py:206 calls _conclude_owned_state_on_failure(), so an ADK turn with no model output now concludes the whole handler-owned root and envelope with status 500 mid-run. The ADK test added in this PR (splunk-ao-adk/tests/test_span_manager.py:29) only covers the non-empty case.

Coerce at this seam rather than validating strictly. _created_at (lines 27-34) has the same shape: datetime.fromisoformat(start_time_iso) on a malformed string raises ValueError with the same blast radius — it should fall back to datetime.now(tz=UTC) the way _step_number already degrades gracefully.

Suggested change
if node.node_type in ("llm", "chat"):
return LoggedLlmSpan(
**common,
input=input_value,
output=output,
metrics=LlmMetrics.model_validate(
{
"duration_ns": duration_ns,
"num_input_tokens": params.get("num_input_tokens"),
"num_output_tokens": params.get("num_output_tokens"),
"num_total_tokens": params.get("num_total_tokens", params.get("total_tokens")),
"time_to_first_token_ns": params.get("time_to_first_token_ns"),
"num_reasoning_tokens": params.get("num_reasoning_tokens"),
"num_cached_input_tokens": params.get("num_cached_input_tokens"),
}
),
tools=params.get("tools"),
events=params.get("events"),
model=params.get("model"),
temperature=params.get("temperature"),
)
if node.node_type in ("llm", "chat"):
return LoggedLlmSpan(
**common,
input=input_value,
output=output if output is not None else "",

🤖 Generated by the Astra agent

Comment on lines +206 to +213
except Exception:
self._splunk_ao_logger._restore_handler_step_context(state.activation)
state.activation = None
self._conclude_owned_state_on_failure()
_logger.warning("Failed to complete handler telemetry for node %s", node_id, exc_info=True)
finally:
if is_root:
self._reset_handler_state()

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟠 major (bug): A single failing span tears down the entire handler trace. _conclude_owned_state_on_failure() walks the parent stack concluding operations and then concludes _owned_trace with status 500 — for what may be one malformed child node (see the output=None case on span_lifecycle.py, or any pydantic rejection of framework-supplied content).

Failure scenario: an ADK/LangChain run with 20 spans hits one unserializable output on span 3. Instead of losing that one span, the whole owned root and envelope are concluded and emitted with status 500 while the framework keeps firing callbacks. Every later _start_incremental_step then fails in _record_otel_ids (the parent identity was released by the conclude) and logs a warning, so the remaining 17 spans are lost too — and the user sees a trace falsely marked as a 500.

A per-node failure should drop only that node: pop it from _active_steps, restore its activation, release its identity, and leave the owned trace intact. Reserve _conclude_owned_state_on_failure() for failures of the root/envelope itself.

🤖 Generated by the Astra agent

Comment on lines +187 to +193
restored_logger = self.get_logger_instance(
project=_project_context.get(),
agent_stream=_agent_stream_context.get(),
experiment_id=_experiment_id_context.get(),
)
if isinstance(restored_logger, SplunkAOLogger):
restored_logger._set_active_session_id(restored_session_id)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 minor (design): __exit__ now constructs a second logger — keyed on the restored outer context, which is a different singleton key from the one flushed at line 170 — purely to publish a session id. get_logger_instance creates and caches a real SplunkAOLogger when none exists, which means deployment/config resolution, exporter and sink construction, a background BatchSpanProcessor thread, and an atexit registration.

Failure scenario: with splunk_ao_context(project="A"): ... used at top level where no logger exists for the default context — exiting the block now builds a logger for the default context (and any ConfigurationError raised there escapes from __exit__, replacing an exception from the body). Repeating this over several distinct outer contexts accumulates loggers and BSP threads that were never used for telemetry.

Only publish the session id when a logger for the restored context already exists, rather than forcing one into existence.

🤖 Generated by the Astra agent

Comment thread src/splunk_ao/config.py
super().set_validated_api_client()
with _suppress_control_plane_http():
super().set_validated_api_client()
assert self.validated_api_client is not None

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 nit (bug): assert used for control flow: under python -O the statement is stripped and client = self.validated_api_client / client.host raises a bare AttributeError: 'NoneType' object has no attribute 'host' instead of a useful error. Prefer an explicit check so the failure mode is the same in optimized builds:

Suggested change
assert self.validated_api_client is not None
client = self.validated_api_client
if client is None:
raise ConfigurationError("Validated API client was not initialized.")

🤖 Generated by the Astra agent

Comment on lines +188 to +189
if state.step is self._owned_root or node_id == str(getattr(self._root_node, "run_id", "")):
self._owned_root = final if self._owned_root is not None else self._owned_root

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 nit (other): This is hard to read for no benefit: the right-hand side reduces to "assign final only if _owned_root is already non-None", and node_id == str(getattr(self._root_node, "run_id", "")) restates is_root, which was already computed on line 179.

Suggested change
if state.step is self._owned_root or node_id == str(getattr(self._root_node, "run_id", "")):
self._owned_root = final if self._owned_root is not None else self._owned_root
if self._owned_root is not None and (state.step is self._owned_root or is_root):
self._owned_root = final

🤖 Generated by the Astra agent

Comment on lines +126 to +133
def test_middleware_detaches_context_on_exception(app_factory) -> None:
make_app, _, _ = app_factory
with pytest.raises(RuntimeError, match="request failed"):
TestClient(make_app(fail=True)).get(
"/test", headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"}
)

@patch("splunk_ao.logger.logger.Projects")
@patch("splunk_ao.logger.logger.AgentStreams")
def test_invalid_uuid_headers_raise_exception(
mock_logstreams_client: Mock, mock_projects_client: Mock, app: FastAPI, client: TestClient
):
"""Test that invalid UUID headers raise SplunkAOLoggerException."""
from splunk_ao.exceptions import SplunkAOLoggerException
assert not trace.get_current_span().get_span_context().is_valid

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 minor (testing): This assertion cannot fail, so the one test named for the detach invariant proves nothing. otel_context.attach happens inside the ASGI task's copied contextvars.Context; a ContextVar set there can never leak back into the test thread, so assert not trace.get_current_span()...is_valid is true whether or not middleware/tracing.py:71 detaches. Deleting the finally: otel_context.detach(token) keeps this test green.

To actually pin the invariant, have the endpoint capture otel_context.get_current() (or the active span context) before and after call_next from inside the middleware's context — e.g. wrap the middleware and assert the token is consumed — or assert that a second request in the same task does not inherit the first request's parent.

🤖 Generated by the Astra agent

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

tests/test_base_handler.py:349-391 (line not in diff)

🟠 major (testing): These two tests no longer exercise commit(), and they pass while the span-build path throws. _ingestion_hook is an undeclared pydantic private attribute assigned only in SplunkAOLogger.__init__ (logger.py:285), so it is not in dir(SplunkAOLogger). Mock(spec=SplunkAOLogger) therefore raises AttributeError on access, getattr(..., None) yields None, and end_node takes the new incremental branch at base_handler.py:588commit() is never called, despite both docstrings saying "commit() calls flush()".

Worse, inside that branch finalize_handler_step receives step_id=<Mock>.id, pydantic rejects it, the ValidationError is swallowed at base_handler.py:206, and flush() is still invoked at line 592 — so mock_logger.flush.assert_called_once() is green while the entire new span pipeline is failing. These are the only two tests covering flush_on_chain_end on the incremental path, so its real flush behaviour is unverified.

Rewrite them against the real splunk_ao_logger fixture used elsewhere in this file (or a recording sink) and assert the exported span set in addition to the flush call, so a build failure cannot pass.

🤖 Generated by the Astra agent

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants