diff --git a/AGENTS.md b/AGENTS.md index 7a6f7509..a3087621 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -134,7 +134,13 @@ CI supports Python 3.11–3.14; root CI also spans Linux, macOS, and Windows. - Handler/decorator/OpenAI/ADK telemetry uses the internal logged-step path and converts completed steps to immutable OTel spans. The internal trace envelope is never exported as a span. - `start_splunk_ao_span()` is SDK-native OTel. `add_splunk_ao_span_processor()` and A2A instrument caller-owned OTel. + `configure_distributed_tracing()` is the combined automatic HTTP setup: it creates or accepts an application-owned + provider, registers Splunk AO export once per provider, configures supported upstream instrumentors, and returns the + provider for application shutdown. `instrument_distributed_tracing()` remains transport-only. - Never replace the process-global tracer provider. Register processors on the provided provider; respect ownership. +- Keep SDK-owned authentication, validation, CRUD, routing-resolution, and other control-plane HTTP calls out of + application traces with scoped OTel HTTP suppression. Never use broad backend URL exclusions that can suppress user + traffic, and always restore instrumentation immediately after the SDK request. - Treat ended `ReadableSpan` objects as immutable. Normalize by copying at export, never by mutating private fields. - Completed spans enqueue immediately. `flush()` drains completed work without ending active work; `terminate()` drains, shuts down SDK-owned resources, and discards unfinished state. Caller-owned providers use `shutdown()`. diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index faac9a7a..808621b4 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -86,10 +86,40 @@ client/server wrapping and message-context propagation. The SDK may add processors, but it must never silently replace the global tracer provider. The application owns the provider and calls `shutdown()`; SDK-owned logger/export resources use their own termination path. +### Distributed context and HTTP transports + +Supported automatic propagation reuses upstream OpenTelemetry instrumentors rather than custom SDK HTTP wrappers. +`configure_distributed_tracing()` creates or accepts an OpenTelemetry SDK provider, registers Splunk AO export, and +configures FastAPI/Starlette inbound instrumentation plus Requests, HTTPX sync/async, and aiohttp-client outbound +instrumentation. It returns the application-owned provider for shutdown and never sets or replaces the global tracer +provider. `instrument_distributed_tracing()` remains the transport-only entry point for applications that manage +processor registration separately. Manual `get_tracing_headers()` and `TracingMiddleware` remain fallbacks for +unsupported transports; supported automatic instrumentation does not require either one on individual requests. + +High-level setup is idempotent per provider for Splunk AO processor registration. Each process-wide client instrumentor +and each server application must still have one instrumentation owner; do not combine the helper with direct upstream +instrumentation of the same component. + +SDK-owned authentication, health-check, CRUD, routing-resolution, token-refresh, and streaming control-plane requests +execute inside OTel HTTP-instrumentation suppression. This boundary is scoped to the SDK request and must restore the +caller's context afterward; application traffic to the same host must remain instrumentable. Apply suppression at the +shared configuration/API-client boundary, never by excluding deployment URLs globally. + +An explicit SDK session propagates as the standard W3C baggage member `gen_ai.conversation.id`. Export normalization +may derive the local compatibility attribute `splunk_ao.session.id`, but that attribute is not propagated as a second +baggage member. The SDK must not propagate authentication, deployment, Project, Agent Stream/log stream, experiment, +endpoint, model, workflow, or agent identity in baggage. + ## Span Lifecycle and Export Ownership `SpanSink` owns the SDK's private provider and `BatchSpanProcessor` for internal telemetry. A completed operation is enqueued immediately, allowing scheduled export without an explicit flush. +Without an explicit internal `BatchConfig`, both SDK-owned and caller-owned processors defer to standard OTel +`OTEL_BSP_*` environment configuration. + +Framework handlers register a stable operation identity at each start callback and enqueue that operation at its +matching end callback, so a child can enter `BatchSpanProcessor` while its parent remains active. The deprecated +`ingestion_hook` is the compatibility exception and retains whole-tree payload construction at trace completion. - `flush()` / `async_flush()` drain completed work and do not end an active trace. - `terminate()` drains completed work, shuts down SDK-owned telemetry resources, and releases unfinished state. @@ -156,6 +186,9 @@ receiver failures, and expose bounded acknowledgement health. Do not log secrets These handlers adapt framework events into the internal Path 1 model. Keep optional dependencies lazy and imports free of surprising side effects. CrewAI is excluded on Python 3.14; code and tests must support the unavailable path. +With `start_new_trace=True`, a handler owns and concludes its local envelope and root while still inheriting an active +W3C parent. With `start_new_trace=False`, it attaches below caller-owned active state, emits only its subtree, and leaves +the caller open. ### OpenAI wrapper diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b750ee4..71521672 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,50 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Added the `distributed-tracing` extra and high-level + `configure_distributed_tracing()` setup for Splunk AO export plus supported + upstream FastAPI/Starlette, Requests, HTTPX, and aiohttp-client instrumentation. + The lower-level `instrument_distributed_tracing()` transport helper remains + available for applications that configure their provider separately. +- Explicit SDK sessions now propagate across supported services as standard + W3C `gen_ai.conversation.id` baggage. SDK routing, authentication, and + application identity are not propagated. +- Added module-level `get_tracing_headers()` and + `extract_tracing_context()` helpers for standard W3C `traceparent` and + `tracestate` propagation. `TracingMiddleware` now extracts and scopes that + OpenTelemetry context for Starlette/FastAPI requests. +- SDK-owned authentication, validation, and CRUD HTTP requests are suppressed + from automatic application HTTP instrumentation, preventing control-plane + calls from appearing as unrelated application traces. + +### Changed + +- The temporarily retained `mode="batch"` and `mode="distributed"` values now + use identical scheduled OTLP batch export and W3C propagation. Concluding an + operation ends and queues it; a per-operation `flush()` is not required. +- Callback handlers keep a live real root during framework execution so + outbound W3C context uses the same identity and visible hierarchy that is + exported at commit. +- Normal LangChain, CrewAI, Google ADK, and OpenAI Agents callbacks now enqueue + each completed operation into the existing `BatchSpanProcessor` at that + operation's end callback. The deprecated `ingestion_hook` retains its + whole-tree compatibility behavior. +- Default logger-owned batching now honors standard OpenTelemetry + `OTEL_BSP_*` configuration, matching caller-owned OTel paths. Explicit + internal `BatchConfig` values remain authoritative when supplied. + +### Removed + +- **Breaking:** Removed `trace_id=` and `span_id=` from `SplunkAOLogger`, the + logger-level proprietary `get_tracing_headers()` method, and custom + `Splunk-AO-Trace-ID` / `Splunk-AO-Parent-ID` continuation. Use the new + module-level W3C helpers instead. +- Removed the obsolete distributed-only REST streaming worker and task queue; + normal telemetry in both retained modes uses the existing + `BatchSpanProcessor` path. + ## [0.2.1] - 2026-08-07 ### Fixed diff --git a/README.md b/README.md index 7c4e7d77..d687eda9 100644 --- a/README.md +++ b/README.md @@ -32,6 +32,7 @@ Install the optional integration dependencies used by your application: ```shell pip install "splunk-ao[openai]" pip install "splunk-ao[langchain]" langchain-openai +pip install "splunk-ao[distributed-tracing]" ``` Other available extras include `crewai`, `middleware`, and `all`. @@ -347,7 +348,8 @@ tracer_provider.shutdown() `add_splunk_ao_span_processor()` reads the same deployment and routing configuration described in [Setup](#setup). It registers a `SplunkAOSpanProcessor`, which uses a standard OpenTelemetry -`BatchSpanProcessor`. +`BatchSpanProcessor`. Default logger-owned and caller-owned processors honor +the standard `OTEL_BSP_*` batch processor environment variables. If your application needs another processor or exporter, register it through the standard OpenTelemetry extension points. For example, to send the same @@ -364,6 +366,49 @@ The application owns a `TracerProvider` that it constructs, so it must call `tracer_provider.shutdown()` during teardown. A `SplunkAOLogger` created elsewhere owns a separate provider and should be terminated separately. +##### Automatic HTTP distributed tracing + +The `distributed-tracing` extra provides a supported one-time setup for +upstream OpenTelemetry FastAPI/Starlette, Requests, HTTPX sync/async, and +aiohttp-client instrumentation. The high-level helper creates an application-owned +provider, registers Splunk AO export, and configures those transports together: + +```python +from fastapi import FastAPI +from splunk_ao import configure_distributed_tracing + +app = FastAPI() +provider = configure_distributed_tracing(app=app) +``` + +The application owns the returned `provider` and calls `provider.shutdown()` +during process teardown. The helper never replaces the process-global provider. +Pass an existing OpenTelemetry SDK `TracerProvider` through the +`tracer_provider` argument when other OpenTelemetry/OpenInference agent or model +instrumentations should use that same provider and join the trace. + +For advanced composition, `add_splunk_ao_span_processor()` remains the +export-only API and `instrument_distributed_tracing()` remains the supported +transport-only API. Do not also directly instrument a client or application +component that is owned by the high-level helper; each component should have one +instrumentation owner. + +After this startup call, supported inbound requests extract W3C context and +supported outbound clients inject it automatically. Application code does not +call `get_tracing_headers()` for each request and does not install +`TracingMiddleware`. Those APIs remain supported for transports and frameworks +outside the automatic support matrix. + +SDK-owned authentication, health-check, Project, Agent Stream, and other +control-plane HTTP requests are scoped out of automatic client instrumentation. +They therefore do not appear as application traces, while HTTP requests made by +the application remain instrumented normally. + +An explicit SDK session propagates as the standard +`gen_ai.conversation.id` W3C baggage member. Project, Agent Stream, +experiment, deployment, authentication, model, workflow, and agent identity +remain local and are never added to baggage by the SDK. + #### Export diagnostics The SDK does not change your application's global logging configuration. @@ -416,6 +461,19 @@ not close a trace that was started by the caller. In short-lived jobs, ensure the owning logger or provider is terminated or shut down after framework work finishes. +For normal OTLP export, LangChain, CrewAI, Google ADK, and OpenAI Agents +handlers enqueue each operation into the existing `BatchSpanProcessor` when +that operation's end callback runs. A child can therefore enter a scheduled +or size-based export batch while its parent remains active. This is immediate +enqueue, not one network request per span, and no per-trace `flush()` is +required. The deprecated `ingestion_hook` compatibility path still constructs +one mutable whole-tree payload at trace completion. + +`start_new_trace` controls ownership, not whether distributed tracing works. +Keep its default `True` for a standalone handler. Set it to `False` only when +the handler is intentionally placed beneath caller-owned active logger state; +the handler emits its own subtree and does not conclude the caller. + #### Datasets Create a dataset: diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example b/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example new file mode 100644 index 00000000..90a5f99f --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/.env.example @@ -0,0 +1,23 @@ +# Select exactly one deployment block. Do not mix o11y and standalone variables. +# O11y Cloud +SPLUNK_AO_REALM="your-o11y-realm" +SPLUNK_AO_O11Y_TOKEN="your-o11y-ingest-or-ingest+api-token" +# SPLUNK_AO_O11Y_API_TOKEN="your-o11y-api-token" + +# Standalone +SPLUNK_AO_API_KEY="your-splunk-ao-api-key" +SPLUNK_AO_CONSOLE_URL="https://app.galileo.ai" + +# Configure project and agentstream +SPLUNK_AO_PROJECT="your-splunk-ao-project" +SPLUNK_AO_AGENT_STREAM="distributed-tracing-example" + +# OpenAI configuration +OPENAI_API_KEY="your-openai-api-key" +OPENAI_MODEL=gpt-5-nano + +OTEL_SERVICE_NAME=distributed-tracing-manual + +# Optional delay to see partial traces. +OTEL_BSP_SCHEDULE_DELAY="1000" +BATCH_DEMO_SLEEP_SECONDS="20" diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/README.md b/examples/logging-samples/DT_2.0/distributed-tracing-auto/README.md new file mode 100644 index 00000000..08939358 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/README.md @@ -0,0 +1,39 @@ +## Configure + +Copy `.env.example` to `.env` and set the credentials and project and agentstream. + +```bash +python -m venv venv +source venv/bin/activate +pip install -r requirements.txt +cp .env.example .env +``` + +Start the retrieval service: +```bash +uvicorn retrieval_service:app --port 8000 +``` + +Then run the orchestrator: +```bash +python main_run.py +``` + +### Observe batching before the root ends + +To demonstrate incremental export, add these values to an existing `.env`: + +```dotenv +OTEL_BSP_SCHEDULE_DELAY="1000" +BATCH_DEMO_SLEEP_SECONDS="20" +``` + +`OTEL_BSP_SCHEDULE_DELAY` is in milliseconds and is read when each batch span +processor is created. After the distributed retrieval call returns, the example +prints a batching-demo message and keeps the orchestrator root active for 20 +seconds. + +Set`BATCH_DEMO_SLEEP_SECONDS="0"` to disable the pause. + +The local tree remains the same as the baseline. Completed children enter the +batch processor at their own end callbacks, so no per-trace flush is needed. diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/main_run.py b/examples/logging-samples/DT_2.0/distributed-tracing-auto/main_run.py new file mode 100644 index 00000000..237a1326 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/main_run.py @@ -0,0 +1,90 @@ +"""Orchestrator using automatic W3C HTTP client instrumentation.""" + +import asyncio +import os +from uuid import uuid4 + +import httpx +from dotenv import load_dotenv + +from splunk_ao import configure_distributed_tracing, log, openai, splunk_ao_context + +load_dotenv() + +provider = configure_distributed_tracing() + +openai_client = openai.OpenAI() +OPENAI_MODEL = os.environ["OPENAI_MODEL"] +RETRIEVAL_SERVICE_URL = "http://localhost:8000" +BATCH_DEMO_SLEEP_SECONDS = float(os.getenv("BATCH_DEMO_SLEEP_SECONDS", "0")) + + +@log +async def orchestrator_agent(question: str) -> str: + """Run one independently traced RAG request.""" + analysis = analyze_question(question) + + # HTTPX creates the client span and injects the active W3C context. + async with httpx.AsyncClient(base_url=RETRIEVAL_SERVICE_URL, timeout=100.0) as client: + try: + response = await client.post("/retrieve", json={"query": question}) + response.raise_for_status() + retrieved_docs = response.json()["results"] + except httpx.HTTPError: + retrieved_docs = [] + + if BATCH_DEMO_SLEEP_SECONDS > 0: + print( + "Batching demo: analysis and distributed retrieval spans have ended; " + f"the orchestrator root remains active for {BATCH_DEMO_SLEEP_SECONDS:g} seconds." + ) + await asyncio.sleep(BATCH_DEMO_SLEEP_SECONDS) + + context = format_context(analysis, retrieved_docs) + response = openai_client.chat.completions.create( + messages=[ + { + "role": "system", + "content": ( + "Answer the user's question using only the provided context. " + f"If the context is insufficient, say so.\n\nContext:\n{context}" + ), + }, + {"role": "user", "content": question}, + ], + model=OPENAI_MODEL, + ) + return response.choices[0].message.content or "" + + +@log +def analyze_question(question: str) -> dict[str, object]: + question_lower = question.lower() + return { + "needs_company_info": any(word in question_lower for word in ("company", "work", "employer")), + "needs_location_info": any(word in question_lower for word in ("location", "where", "city", "live")), + "question_type": "factual", + } + + +@log +def format_context(analysis: dict[str, object], documents: list[str]) -> str: + rendered = [f"Analysis: {analysis}", "", "Retrieved Documents:"] + rendered.extend(f"{index}. {document}" for index, document in enumerate(documents, 1)) + return "\n".join(rendered) + + +async def main() -> None: + # One app-owned session is propagated as gen_ai.conversation.id baggage. + # set_session() is local-only; unlike start_session(), it performs no CRUD request. + splunk_ao_context.set_session(str(uuid4())) + for question in ("What did Galileo Galilei research?", "Where did Galileo Galilei work?"): + answer = await orchestrator_agent(question) + print(f"Question: {question}\nAnswer: {answer}\n") + + +if __name__ == "__main__": + try: + asyncio.run(main()) + finally: + provider.shutdown() diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/requirements.txt b/examples/logging-samples/DT_2.0/distributed-tracing-auto/requirements.txt new file mode 100644 index 00000000..d351de56 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/requirements.txt @@ -0,0 +1,6 @@ +python-dotenv==1.2.2 +pydantic==2.11.9 +httpx==0.28.1 +fastapi==0.117.1 +uvicorn==0.37.0 +-e ../../../..[openai,distributed-tracing] diff --git a/examples/logging-samples/DT_2.0/distributed-tracing-auto/retrieval_service.py b/examples/logging-samples/DT_2.0/distributed-tracing-auto/retrieval_service.py new file mode 100644 index 00000000..a0a276bd --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing-auto/retrieval_service.py @@ -0,0 +1,56 @@ +"""FastAPI service using automatic W3C server instrumentation.""" + +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from dotenv import load_dotenv +from fastapi import FastAPI +from pydantic import BaseModel + +from splunk_ao import configure_distributed_tracing, log + +load_dotenv() + + +@asynccontextmanager +async def lifespan(app: FastAPI) -> AsyncIterator[None]: + yield + provider.shutdown() + + +app = FastAPI(title="Retrieval Service", lifespan=lifespan) +provider = configure_distributed_tracing(app=app) + + +class RetrievalRequest(BaseModel): + query: str + + +class RetrievalResponse(BaseModel): + results: list[str] + + +@log(span_type="retriever") +def retrieval_service(query: str) -> list[str]: + knowledge_base = { + "birthplace": ["Galileo Galilei was born in Pisa, Italy in 1564."], + "profession": ["Galileo taught geometry, mechanics, and astronomy at the University of Padua."], + "research": ["Galileo's telescopic observations transformed our understanding of the universe."], + } + query_lower = query.lower() + results: list[str] = [] + for category, facts in knowledge_base.items(): + if category in query_lower or any(word in query_lower for word in ("work", "location", "education")): + results.extend(facts) + return results[:3] + + +@app.post("/retrieve", response_model=RetrievalResponse) +@log +async def retrieve_endpoint(request: RetrievalRequest) -> RetrievalResponse: + return RetrievalResponse(results=retrieval_service(request.query)) + + +@app.get("/health") +async def health() -> dict[str, str]: + return {"status": "healthy"} diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/.env.example b/examples/logging-samples/DT_2.0/distributed-tracing/.env.example new file mode 100644 index 00000000..7ddde639 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/.env.example @@ -0,0 +1,19 @@ +# Select exactly one deployment block. Do not mix o11y and standalone variables. +# O11y Cloud +SPLUNK_AO_REALM="your-o11y-realm" +SPLUNK_AO_O11Y_TOKEN="your-o11y-ingest-or-ingest+api-token" +# SPLUNK_AO_O11Y_API_TOKEN="your-o11y-api-token" + +# Standalone +SPLUNK_AO_API_KEY="your-splunk-ao-api-key" +SPLUNK_AO_CONSOLE_URL="https://app.galileo.ai" + +# Configure project and agentstream +SPLUNK_AO_PROJECT="your-splunk-ao-project" +SPLUNK_AO_AGENT_STREAM="distributed-tracing-example" + +# OpenAI configuration +OPENAI_API_KEY="your-openai-api-key" +OPENAI_MODEL=gpt-5-nano + +OTEL_SERVICE_NAME=distributed-tracing-manual diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/README.md b/examples/logging-samples/DT_2.0/distributed-tracing/README.md new file mode 100644 index 00000000..cb93a933 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/README.md @@ -0,0 +1,25 @@ +## Configure + +Copy `.env.example` to `.env` and set the credentials and project and agentstream. + +```bash +python -m venv venv +source venv/bin/activate +pip install -r requirements.txt +``` + +## Run + +Start the retrieval service in one terminal: +```bash +uvicorn retrieval_service:app --reload --port 8000 +``` + +Run the orchestrator in another: +```bash +python main_run.py +``` + +Each top level `orchestrator_agent()` call creates an independent trace. Its +outbound request carries the active orchestrator workflow span ID, and the +retrieval endpoint's workflow becomes its direct downstream child. diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/main_run.py b/examples/logging-samples/DT_2.0/distributed-tracing/main_run.py new file mode 100644 index 00000000..e2d44dd4 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/main_run.py @@ -0,0 +1,77 @@ +"""Orchestrator process for the W3C distributed-tracing example.""" + +import asyncio +from uuid import uuid4 + +import httpx +from dotenv import load_dotenv + +from splunk_ao import get_tracing_headers, log, openai, splunk_ao_context + +load_dotenv() + +openai_client = openai.OpenAI() +RETRIEVAL_SERVICE_URL = "http://localhost:8000" + + +@log +async def orchestrator_agent(question: str) -> str: + """Run one independently traced RAG request.""" + analysis = analyze_question(question) + + # The active @log workflow is a real exportable operation. The returned + # carrier contains W3C traceparent and optional tracestate fields. + headers = get_tracing_headers() + async with httpx.AsyncClient(base_url=RETRIEVAL_SERVICE_URL, timeout=100.0) as client: + try: + response = await client.post("/retrieve", json={"query": question}, headers=headers) + response.raise_for_status() + retrieved_docs = response.json()["results"] + except httpx.HTTPError: + retrieved_docs = [] + + context = format_context(analysis, retrieved_docs) + response = openai_client.chat.completions.create( + messages=[ + { + "role": "system", + "content": ( + "Answer the user's question using only the provided context. " + f"If the context is insufficient, say so.\n\nContext:\n{context}" + ), + }, + {"role": "user", "content": question}, + ], + model="gpt-5-mini", + ) + return response.choices[0].message.content or "" + + +@log +def analyze_question(question: str) -> dict[str, object]: + question_lower = question.lower() + return { + "needs_company_info": any(word in question_lower for word in ("company", "work", "employer")), + "needs_location_info": any(word in question_lower for word in ("location", "where", "city", "live")), + "question_type": "factual", + } + + +@log +def format_context(analysis: dict[str, object], documents: list[str]) -> str: + rendered = [f"Analysis: {analysis}", "", "Retrieved Documents:"] + rendered.extend(f"{index}. {document}" for index, document in enumerate(documents, 1)) + return "\n".join(rendered) + + +async def main() -> None: + # One app-owned session is propagated as gen_ai.conversation.id baggage. + # set_session() is local-only; unlike start_session(), it performs no CRUD request. + splunk_ao_context.set_session(str(uuid4())) + for question in ("What did Galileo Galilei research?", "Where did Galileo Galilei work?"): + answer = await orchestrator_agent(question) + print(f"Question: {question}\nAnswer: {answer}\n") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/requirements.txt b/examples/logging-samples/DT_2.0/distributed-tracing/requirements.txt new file mode 100644 index 00000000..5e0a8f23 --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/requirements.txt @@ -0,0 +1,6 @@ +python-dotenv==1.2.2 +pydantic==2.11.9 +httpx==0.28.1 +fastapi==0.117.1 +uvicorn==0.37.0 +-e ../../../..[openai,middleware] diff --git a/examples/logging-samples/DT_2.0/distributed-tracing/retrieval_service.py b/examples/logging-samples/DT_2.0/distributed-tracing/retrieval_service.py new file mode 100644 index 00000000..b402514e --- /dev/null +++ b/examples/logging-samples/DT_2.0/distributed-tracing/retrieval_service.py @@ -0,0 +1,49 @@ +"""FastAPI retrieval process that continues incoming W3C context.""" + +from dotenv import load_dotenv +from fastapi import FastAPI +from pydantic import BaseModel + +from splunk_ao import log +from splunk_ao.middleware import TracingMiddleware + +load_dotenv() + +app = FastAPI(title="Retrieval Service") +app.add_middleware(TracingMiddleware) + + +class RetrievalRequest(BaseModel): + query: str + + +class RetrievalResponse(BaseModel): + results: list[str] + + +@log(span_type="retriever") +def retrieval_service(query: str) -> list[str]: + knowledge_base = { + "birthplace": ["Galileo Galilei was born in Pisa, Italy in 1564."], + "profession": ["Galileo taught geometry, mechanics, and astronomy at the University of Padua."], + "research": ["Galileo's telescopic observations transformed our understanding of the universe."], + } + query_lower = query.lower() + results: list[str] = [] + for category, facts in knowledge_base.items(): + if category in query_lower or any(word in query_lower for word in ("work", "location", "education")): + results.extend(facts) + return results[:3] + + +@app.post("/retrieve", response_model=RetrievalResponse) +@log +async def retrieve_endpoint(request: RetrievalRequest) -> RetrievalResponse: + # @log creates the real downstream workflow beneath the middleware's + # extracted remote parent; the retriever is its local child. + return RetrievalResponse(results=retrieval_service(request.query)) + + +@app.get("/health") +async def health() -> dict[str, str]: + return {"status": "healthy"} diff --git a/poetry.lock b/poetry.lock index 5edd9823..68a4dc3a 100644 --- a/poetry.lock +++ b/poetry.lock @@ -7,7 +7,7 @@ description = "Happy Eyeballs for asyncio" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "aiohappyeyeballs-2.6.1-py3-none-any.whl", hash = "sha256:f349ba8f4b75cb25c99c5c2d84e997e485204d2902a9597802b0371f09331fb8"}, {file = "aiohappyeyeballs-2.6.1.tar.gz", hash = "sha256:c3f9d0113123803ccadfdf3f0faa505bc78e6a72d1cc4806cbd719826e943558"}, @@ -20,7 +20,7 @@ description = "Async http client/server framework (asyncio)" optional = true python-versions = ">=3.10" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "aiohttp-3.14.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:eb0495d778817619273c108784292be161a924b9f5ae5cbbc70a2caa6838250b"}, {file = "aiohttp-3.14.3-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:c3c200cf9757edd785051dc699c7ecbec22110dbfcb3fefc7a9f9695eda8ea7a"}, @@ -163,7 +163,7 @@ description = "aiosignal: a list of registered asynchronous callbacks" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "aiosignal-1.4.0-py3-none-any.whl", hash = "sha256:053243f8b92b990551949e63930a839ff0cf0b0ebbe0597b0f3fb19e1a0fe82e"}, {file = "aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7"}, @@ -246,6 +246,23 @@ files = [ [package.extras] test = ["coverage", "mypy", "pexpect", "ruff", "wheel"] +[[package]] +name = "asgiref" +version = "3.12.1" +description = "ASGI specs, helper code, and adapters" +optional = true +python-versions = ">=3.10" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "asgiref-3.12.1-py3-none-any.whl", hash = "sha256:fe386d1c2bff7259ea95929266d12a8cf9a8b5a1c2598402967d8792e7a7c094"}, + {file = "asgiref-3.12.1.tar.gz", hash = "sha256:59dcb51c272ad209d59bed5708a64a333083e86017d7fcdd67498eeab7784340"}, +] + +[package.extras] +mypy = ["mypy (>=1.14.0)"] +tests = ["pytest", "pytest-asyncio"] + [[package]] name = "attrs" version = "25.3.0" @@ -682,7 +699,7 @@ files = [ {file = "colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6"}, {file = "colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44"}, ] -markers = {main = "(platform_system == \"Windows\" or python_version <= \"3.13\") and (extra == \"crewai\" or extra == \"all\" or platform_system == \"Windows\") and (os_name == \"nt\" or platform_system == \"Windows\" or sys_platform == \"win32\")", dev = "platform_system == \"Windows\" or sys_platform == \"win32\"", test = "sys_platform == \"win32\""} +markers = {main = "(python_version <= \"3.13\" or platform_system == \"Windows\") and (extra == \"crewai\" or extra == \"all\" or platform_system == \"Windows\") and (os_name == \"nt\" or platform_system == \"Windows\" or sys_platform == \"win32\")", dev = "platform_system == \"Windows\" or sys_platform == \"win32\"", test = "sys_platform == \"win32\""} [[package]] name = "coloredlogs" @@ -948,7 +965,7 @@ files = [ {file = "distro-1.9.0-py3-none-any.whl", hash = "sha256:7bffd925d65168f85027d8da9af6bddab658135b840670a223589bc0c8ef02b2"}, {file = "distro-1.9.0.tar.gz", hash = "sha256:2fa77c6fd8940f116ee1d6b94a2f90b13b5ea8d019b98bc8bafdcabcdd9bdbed"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or extra == \"langchain\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\" or extra == \"langchain\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\" or extra == \"langchain\") or extra == \"openai\" or extra == \"all\" or extra == \"langchain\""} [[package]] name = "docstring-parser" @@ -1154,7 +1171,7 @@ description = "A list-like structure which implements collections.abc.MutableSeq optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "frozenlist-1.7.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:cc4df77d638aa2ed703b878dd093725b72a824c3c546c076e8fdf276f78ee84a"}, {file = "frozenlist-1.7.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:716a9973a2cc963160394f701964fe25012600f3d311f60c790400b00e568b61"}, @@ -1918,7 +1935,7 @@ files = [ {file = "jiter-0.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:1b28302349dc65703a9e4ead16f163b1c339efffbe1049c30a44b001a2a4fff9"}, {file = "jiter-0.10.0.tar.gz", hash = "sha256:07a7142c38aacc85194391108dc91b5b57093c978a9932bd86a36862759d9500"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\") or extra == \"openai\" or extra == \"all\""} [[package]] name = "json-repair" @@ -2614,7 +2631,7 @@ description = "multidict implementation" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "multidict-6.6.4-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:b8aa6f0bd8125ddd04a6593437bad6a7e70f300ff4180a531654aa2ab3f6d58f"}, {file = "multidict-6.6.4-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:b9e5853bbd7264baca42ffc53391b490d65fe62849bf2c690fa3f6273dbcd0cb"}, @@ -2944,7 +2961,7 @@ files = [ {file = "openai-2.32.0-py3-none-any.whl", hash = "sha256:4dcc9badeb4bf54ad0d187453742f290226d30150890b7890711bda4f32f192f"}, {file = "openai-2.32.0.tar.gz", hash = "sha256:c54b27a9e4cb8d51f0dd94972ffd1a04437efeb259a9e60d8922b8bd26fe55e0"}, ] -markers = {main = "extra == \"openai\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"openai\") or extra == \"openai\" or extra == \"all\""} [package.dependencies] anyio = ">=3.5.0,<5" @@ -3123,6 +3140,162 @@ opentelemetry-sdk = ">=1.38.0,<1.39.0" requests = ">=2.7,<3.0" typing-extensions = ">=4.5.0" +[[package]] +name = "opentelemetry-instrumentation" +version = "0.59b0" +description = "Instrumentation Tools & Auto Instrumentation for OpenTelemetry Python" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation-0.59b0-py3-none-any.whl", hash = "sha256:44082cc8fe56b0186e87ee8f7c17c327c4c2ce93bdbe86496e600985d74368ee"}, + {file = "opentelemetry_instrumentation-0.59b0.tar.gz", hash = "sha256:6010f0faaacdaf7c4dff8aac84e226d23437b331dcda7e70367f6d73a7db1adc"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.4,<2.0" +opentelemetry-semantic-conventions = "0.59b0" +packaging = ">=18.0" +wrapt = ">=1.0.0,<2.0.0" + +[[package]] +name = "opentelemetry-instrumentation-aiohttp-client" +version = "0.59b0" +description = "OpenTelemetry aiohttp client instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_aiohttp_client-0.59b0-py3-none-any.whl", hash = "sha256:8e7f234d2b6b385d5b1757ba7f13baa71c45ec020e1a08ba7f7649da7958470a"}, + {file = "opentelemetry_instrumentation_aiohttp_client-0.59b0.tar.gz", hash = "sha256:665ee520f2fb5f44d6d600918eb7bbd2c29b355e0d0deda49991db857adee51f"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" +wrapt = ">=1.0.0,<2.0.0" + +[package.extras] +instruments = ["aiohttp (>=3.0,<4.0)"] + +[[package]] +name = "opentelemetry-instrumentation-asgi" +version = "0.59b0" +description = "ASGI instrumentation for OpenTelemetry" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_asgi-0.59b0-py3-none-any.whl", hash = "sha256:ba9703e09d2c33c52fa798171f344c8123488fcd45017887981df088452d3c53"}, + {file = "opentelemetry_instrumentation_asgi-0.59b0.tar.gz", hash = "sha256:2509d6fe9fd829399ce3536e3a00426c7e3aa359fc1ed9ceee1628b56da40e7a"}, +] + +[package.dependencies] +asgiref = ">=3.0,<4.0" +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["asgiref (>=3.0,<4.0)"] + +[[package]] +name = "opentelemetry-instrumentation-fastapi" +version = "0.59b0" +description = "OpenTelemetry FastAPI Instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_fastapi-0.59b0-py3-none-any.whl", hash = "sha256:0d8d00ff7d25cca40a4b2356d1d40a8f001e0668f60c102f5aa6bb721d660c4f"}, + {file = "opentelemetry_instrumentation_fastapi-0.59b0.tar.gz", hash = "sha256:e8fe620cfcca96a7d634003df1bc36a42369dedcdd6893e13fb5903aeeb89b2b"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-instrumentation-asgi = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["fastapi (>=0.92,<1.0)"] + +[[package]] +name = "opentelemetry-instrumentation-httpx" +version = "0.59b0" +description = "OpenTelemetry HTTPX Instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_httpx-0.59b0-py3-none-any.whl", hash = "sha256:7dc9f66aef4ca3904d877f459a70c78eafd06131dc64d713b9b1b5a7d0a48f05"}, + {file = "opentelemetry_instrumentation_httpx-0.59b0.tar.gz", hash = "sha256:a1cb9b89d9f05a82701cc9ab9cfa3db54fd76932489449778b350bc1b9f0e872"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" +wrapt = ">=1.0.0,<2.0.0" + +[package.extras] +instruments = ["httpx (>=0.18.0)"] + +[[package]] +name = "opentelemetry-instrumentation-requests" +version = "0.59b0" +description = "OpenTelemetry requests instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_requests-0.59b0-py3-none-any.whl", hash = "sha256:d43121532877e31a46c48649279cec2504ee1e0ceb3c87b80fe5ccd7eafc14c1"}, + {file = "opentelemetry_instrumentation_requests-0.59b0.tar.gz", hash = "sha256:9af2ffe3317f03074d7f865919139e89170b6763a0251b68c25e8e64e04b3400"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["requests (>=2.0,<3.0)"] + +[[package]] +name = "opentelemetry-instrumentation-starlette" +version = "0.59b0" +description = "OpenTelemetry Starlette Instrumentation" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_instrumentation_starlette-0.59b0-py3-none-any.whl", hash = "sha256:a833b97d297e4b2aaf58041612663dbabb8380e3993c76a4e32b3e351470f321"}, + {file = "opentelemetry_instrumentation_starlette-0.59b0.tar.gz", hash = "sha256:3f033fd92d6a8e4122ebcb3d83afc5c64d6be7930e9094876eb02b8afbd08ba5"}, +] + +[package.dependencies] +opentelemetry-api = ">=1.12,<2.0" +opentelemetry-instrumentation = "0.59b0" +opentelemetry-instrumentation-asgi = "0.59b0" +opentelemetry-semantic-conventions = "0.59b0" +opentelemetry-util-http = "0.59b0" + +[package.extras] +instruments = ["starlette (>=0.13)"] + [[package]] name = "opentelemetry-proto" version = "1.38.0" @@ -3171,6 +3344,19 @@ files = [ opentelemetry-api = "1.38.0" typing-extensions = ">=4.5.0" +[[package]] +name = "opentelemetry-util-http" +version = "0.59b0" +description = "Web util for OpenTelemetry" +optional = true +python-versions = ">=3.9" +groups = ["main"] +markers = "extra == \"distributed-tracing\" or extra == \"all\"" +files = [ + {file = "opentelemetry_util_http-0.59b0-py3-none-any.whl", hash = "sha256:6d036a07563bce87bf521839c0671b507a02a0d39d7ea61b88efa14c6e25355d"}, + {file = "opentelemetry_util_http-0.59b0.tar.gz", hash = "sha256:ae66ee91be31938d832f3b4bc4eb8a911f6eddd38969c4a871b1230db2a0a560"}, +] + [[package]] name = "orjson" version = "3.11.9" @@ -3254,7 +3440,7 @@ files = [ {file = "orjson-3.11.9-cp314-cp314-win_arm64.whl", hash = "sha256:eebdbdeef0094e4f5aefa20dcd4eb2368ab5e7a3b4edea27f1e7b2892e009cf9"}, {file = "orjson-3.11.9.tar.gz", hash = "sha256:4fef17e1f8722c11587a6ef18e35902450221da0028e65dbaaa543619e68e48f"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation != \"PyPy\""} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") or extra == \"langchain\" or extra == \"all\"", test = "platform_python_implementation != \"PyPy\""} [[package]] name = "ormsgpack" @@ -3339,7 +3525,7 @@ files = [ {file = "packaging-24.2-py3-none-any.whl", hash = "sha256:09abb1bccd265c01f4a3aa3f7a7db064b36514d2cba19a2f694fe6150451a759"}, {file = "packaging-24.2.tar.gz", hash = "sha256:c228a6dc5e932d346bc5739379109d49e8853dd8223571c7c5b55260edc0b97f"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\" or extra == \"openai\" or extra == \"distributed-tracing\") or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"distributed-tracing\""} [[package]] name = "pathspec" @@ -3602,7 +3788,7 @@ description = "Accelerated property cache" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "propcache-0.3.2-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:22d9962a358aedbb7a2e36187ff273adeaab9743373a272976d2e348d08c7770"}, {file = "propcache-0.3.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:0d0fda578d1dc3f77b6b5a5dce3b9ad69a8250a891760a548df850a5e8da87f3"}, @@ -3983,7 +4169,7 @@ files = [ {file = "pycparser-2.22-py3-none-any.whl", hash = "sha256:c3702b6d3dd8c7abc1afa565d7e63d53a1d0bd86cdc24edd75470f4de499cfcc"}, {file = "pycparser-2.22.tar.gz", hash = "sha256:491c8be9c040f5390f5bf44a5b07752bd07f56edf992381b05c701439eec10f6"}, ] -markers = {main = "(python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"crewai\") and (platform_python_implementation != \"PyPy\" or extra == \"langchain\" or extra == \"all\") and implementation_name != \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\" or extra == \"openai\") and (platform_python_implementation == \"PyPy\" or python_version <= \"3.13\" or extra == \"openai\" or extra == \"all\" or extra == \"crewai\") and (platform_python_implementation == \"PyPy\" or extra == \"crewai\" or extra == \"all\" or extra == \"openai\")", test = "platform_python_implementation == \"PyPy\" and implementation_name != \"PyPy\""} +markers = {main = "(python_version <= \"3.13\" or extra == \"langchain\" or extra == \"all\" or extra == \"openai\" or extra == \"crewai\") and (platform_python_implementation != \"PyPy\" or extra == \"langchain\" or extra == \"all\") and implementation_name != \"PyPy\" and (extra == \"langchain\" or extra == \"all\" or extra == \"crewai\" or extra == \"openai\") and (platform_python_implementation == \"PyPy\" or extra == \"openai\" or extra == \"all\" or extra == \"crewai\")", test = "platform_python_implementation == \"PyPy\" and implementation_name != \"PyPy\""} [[package]] name = "pydantic" @@ -4575,7 +4761,7 @@ files = [ {file = "PyYAML-6.0.2-cp39-cp39-win_amd64.whl", hash = "sha256:39693e1f8320ae4f43943590b49779ffb98acb81f788220ea932a6b6c51004d8"}, {file = "pyyaml-6.0.2.tar.gz", hash = "sha256:d584d9ec91ad65861cc08d42e834324ef890a082e591037abe114850ff7bbc3e"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") or extra == \"langchain\" or extra == \"all\""} [[package]] name = "referencing" @@ -5199,7 +5385,7 @@ files = [ {file = "tenacity-9.1.2-py3-none-any.whl", hash = "sha256:f77bf36710d8b73a50b2dd155c97b870017ad21afe6ab300326b0371b3b05138"}, {file = "tenacity-9.1.2.tar.gz", hash = "sha256:1169d376c297e7de388d18b4481760d478b0e99a777cad3a9c86e556f4b697cb"}, ] -markers = {main = "extra == \"langchain\" or extra == \"all\" or python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\")"} +markers = {main = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"langchain\") or extra == \"langchain\" or extra == \"all\""} [package.extras] doc = ["reno", "sphinx"] @@ -6241,7 +6427,7 @@ description = "Yet another URL library" optional = true python-versions = ">=3.9" groups = ["main"] -markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\")" +markers = "python_version <= \"3.13\" and (extra == \"crewai\" or extra == \"all\" or extra == \"distributed-tracing\") or python_version < \"3.14\" and (extra == \"distributed-tracing\" or extra == \"all\" or extra == \"crewai\") or extra == \"distributed-tracing\" or extra == \"all\"" files = [ {file = "yarl-1.20.1-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:6032e6da6abd41e4acda34d75a816012717000fa6839f37124a47fcefc49bec4"}, {file = "yarl-1.20.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:2c7b34d804b8cf9b214f05015c4fee2ebe7ed05cf581e7192c06555c71f4446a"}, @@ -6510,8 +6696,9 @@ cffi = {version = ">=1.11", markers = "platform_python_implementation == \"PyPy\ cffi = ["cffi (>=1.11)"] [extras] -all = ["aiohttp", "crewai", "cryptography", "langchain", "langchain-core", "langsmith", "litellm", "mcp", "openai", "openai-agents", "packaging", "pdfminer-six", "starlette", "uv"] +all = ["aiohttp", "crewai", "cryptography", "langchain", "langchain-core", "langsmith", "litellm", "mcp", "openai", "openai-agents", "opentelemetry-instrumentation-aiohttp-client", "opentelemetry-instrumentation-fastapi", "opentelemetry-instrumentation-httpx", "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-starlette", "packaging", "pdfminer-six", "starlette", "uv"] crewai = ["aiohttp", "crewai", "cryptography", "litellm", "mcp", "pdfminer-six", "uv"] +distributed-tracing = ["aiohttp", "opentelemetry-instrumentation-aiohttp-client", "opentelemetry-instrumentation-fastapi", "opentelemetry-instrumentation-httpx", "opentelemetry-instrumentation-requests", "opentelemetry-instrumentation-starlette"] langchain = ["langchain", "langchain-core", "langsmith"] middleware = ["starlette"] openai = ["cryptography", "mcp", "openai", "openai-agents", "packaging"] @@ -6519,4 +6706,4 @@ openai = ["cryptography", "mcp", "openai", "openai-agents", "packaging"] [metadata] lock-version = "2.1" python-versions = "^3.11,<3.15" -content-hash = "5472c5c6216450621d5aae8e3dbdb21a3a071c4c2e3fb383f3552b4dc365d63c" +content-hash = "8219ff5a4bc2c3be82d45473a92cb251482d8b3b5c02ca3a41ef68432f52deb1" diff --git a/pyproject.toml b/pyproject.toml index c9e1aa6d..fd32d71f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,13 +30,14 @@ langchain = ["langchain-core", "langchain", "langsmith (>=0.8.0)"] openai = ["openai (>=2.8.0,<3.0.0)", "packaging (>=24.2,<25.0)", "openai-agents (>=0.4.0,<1.0.0)", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)"] 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)"] +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)", "cryptography (>=50.0.0)", "mcp (>=1.27.2,<2)", "pdfminer-six (>=20251107)", "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)"] [tool.poetry.dependencies] python = "^3.11,<3.15" -aiohttp = { version = ">=3.14.1,<4", optional = true, python = ">=3.11,<3.14" } +aiohttp = { version = ">=3.14.1,<4", optional = true } pydantic = "^2.11.9" pyjwt = ">=2.13.0,<3.0.0" wrapt = "^1.14" @@ -58,6 +59,11 @@ typing-extensions = { version = ">=4.5.0" } opentelemetry-sdk = "^1.38.0" opentelemetry-api = "^1.38.0" opentelemetry-exporter-otlp-proto-http = "^1.38.0" +opentelemetry-instrumentation-fastapi = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-starlette = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-requests = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-httpx = { version = "0.59b0", optional = true } +opentelemetry-instrumentation-aiohttp-client = { version = "0.59b0", optional = true } filelock = ">=3.20.1" idna = ">=3.15,<4" python-dotenv = ">=1.2.2" diff --git a/splunk-ao-adk/tests/test_span_manager.py b/splunk-ao-adk/tests/test_span_manager.py index 27eda6c2..0c18f254 100644 --- a/splunk-ao-adk/tests/test_span_manager.py +++ b/splunk-ao-adk/tests/test_span_manager.py @@ -4,10 +4,61 @@ from uuid import uuid4 import pytest +from opentelemetry.sdk.trace import ReadableSpan +from splunk_ao.handlers.base_handler import SplunkAOBaseHandler +from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao_adk.span_manager import INTEGRATION_TAG, SpanManager +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_adk_span_manager_enqueues_llm_at_callback_end_without_flush(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("SPLUNK_AO_API_KEY", "test-api-key") + monkeypatch.setenv("SPLUNK_AO_CONSOLE_URL", "https://console.test") + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, integration="google_adk", flush_on_chain_end=False) + manager = SpanManager(handler) + root_id = uuid4() + llm_id = uuid4() + try: + manager.start_run(root_id, "question", agent_name="researcher") + manager.start_llm(llm_id, root_id, "prompt", model="model") + + manager.end_llm(llm_id, "answer", num_input_tokens=2, num_output_tokens=1, total_tokens=3) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["chat"] + assert str(root_id) in handler._active_steps + assert sink.force_flush_calls == 0 + + manager.end_run(root_id, "done") + + child, root = sink.spans + assert child.parent == root.context + assert [span.name for span in sink.spans] == [ + "chat model", + "invoke_workflow invocation [researcher]", + ] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class TestSpanManagerRunSpans: """Tests for run (invocation) span lifecycle.""" diff --git a/src/splunk_ao/__init__.py b/src/splunk_ao/__init__.py index 0582d1fc..8b4648f3 100644 --- a/src/splunk_ao/__init__.py +++ b/src/splunk_ao/__init__.py @@ -58,6 +58,7 @@ ) from splunk_ao.experiment import Experiment from splunk_ao.handlers.agent_control import SplunkAOAgentControlBridge, setup_agent_control_bridge +from splunk_ao.http_instrumentation import configure_distributed_tracing, instrument_distributed_tracing from splunk_ao.integration import Integration from splunk_ao.logger import SplunkAOLogger from splunk_ao.logger.control import ControlAppliesTo, ControlCheckStage, ControlResult, ControlSpan @@ -80,7 +81,7 @@ SplunkAOFutureError, ValidationError, ) -from splunk_ao.tracing import get_tracing_headers +from splunk_ao.tracing import extract_tracing_context, get_tracing_headers from splunk_ao.types import MetricSpec from splunk_ao.utils.log_config import enable_console_logging @@ -160,6 +161,7 @@ "ValidationError", "WorkflowSpan", "add_records_to_annotation_queue", + "configure_distributed_tracing", "create_annotation_queue", "create_annotation_queue_field", "create_api_key", @@ -167,10 +169,12 @@ "delete_annotation_queue_field", "delete_api_key", "enable_console_logging", + "extract_tracing_context", "get_agent_control_target", "get_annotation_queue", "get_annotation_queue_records", "get_tracing_headers", + "instrument_distributed_tracing", "is_dependency_available", "list_annotation_queue_fields", "list_annotation_queue_users", diff --git a/src/splunk_ao/config.py b/src/splunk_ao/config.py index bf002819..8b2c3a3c 100644 --- a/src/splunk_ao/config.py +++ b/src/splunk_ao/config.py @@ -2,6 +2,7 @@ # We need to ignore syntax errors until https://github.com/python/mypy/issues/17535 is resolved. import os from collections.abc import Iterator +from contextlib import contextmanager from typing import Any, ClassVar, Optional from httpx import Response @@ -17,7 +18,33 @@ from splunk_ao.shared.exceptions import ConfigurationError, MissingConfigurationError -class O11yApiClient(ApiClient): +@contextmanager +def _suppress_control_plane_http() -> Iterator[None]: + """Suppress optional OTel HTTP instrumentation around SDK-owned requests.""" + try: + from opentelemetry.instrumentation.utils import suppress_http_instrumentation # noqa: PLC0415 + except ImportError: + yield + return + + with suppress_http_instrumentation(): + yield + + +class _ControlPlaneApiClient(ApiClient): + """Keep SDK-owned HTTP requests out of application telemetry.""" + + async def arequest(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Any: + with _suppress_control_plane_http(): + return await super().arequest(method, path, *args, **kwargs) + + @contextmanager + def stream_request(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Iterator[Response]: + with _suppress_control_plane_http(), super().stream_request(method, path, *args, **kwargs) as response: + yield response + + +class O11yApiClient(_ControlPlaneApiClient): """API client for Splunk Observability Cloud AO endpoints.""" o11y_token: SecretStr @@ -41,8 +68,10 @@ def _prefixed(self, path: str) -> str: async def arequest(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Any: return await super().arequest(method, self._prefixed(path), *args, **kwargs) + @contextmanager def stream_request(self, method: RequestMethod, path: str, *args: Any, **kwargs: Any) -> Iterator[Response]: - return super().stream_request(method, self._prefixed(path), *args, **kwargs) + with super().stream_request(method, self._prefixed(path), *args, **kwargs) as response: + yield response # Mapping of SPLUNK_AO_* → GALILEO_* env var pairs used by the bridge. @@ -105,7 +134,8 @@ def set_api_url(cls, api_url: str | Url | None, info: ValidationInfo) -> Url: """Derive the O11y API URL from its realm and preserve standalone validation.""" if cls._is_o11y_env(): return Url(O11yConfig.from_env().require_api_url()) - return super().set_api_url(api_url, info) + with _suppress_control_plane_http(): + return super().set_api_url(api_url, info) @model_validator(mode="after") def set_jwt_token(self) -> "SplunkAOConfig": @@ -114,7 +144,8 @@ def set_jwt_token(self) -> "SplunkAOConfig": self.jwt_token = None self.refresh_token = None return self - super().set_jwt_token() + with _suppress_control_plane_http(): + super().set_jwt_token() return self @model_validator(mode="after") @@ -126,7 +157,17 @@ def set_validated_api_client(self) -> "SplunkAOConfig": host=o11y.api_root, o11y_token=o11y.crud_token, jwt_token=SecretStr(""), ssl_context=self.ssl_context ) return self - super().set_validated_api_client() + with _suppress_control_plane_http(): + super().set_validated_api_client() + assert self.validated_api_client is not None + client = self.validated_api_client + self.validated_api_client = _ControlPlaneApiClient( + host=client.host, + jwt_token=client.jwt_token, + raise_on_unexpected_status=client.raise_on_unexpected_status, + ssl_context=client.ssl_context, + thread_local=client.thread_local, + ) return self def _uses_o11y_api_client(self) -> bool: @@ -136,7 +177,8 @@ def refresh_jwt_token(self) -> None: """Skip JWT refresh when authenticating directly with an O11y token.""" if self._uses_o11y_api_client(): return - super().refresh_jwt_token() + with _suppress_control_plane_http(): + super().refresh_jwt_token() @classmethod def get(cls, **kwargs: Any) -> "SplunkAOConfig": diff --git a/src/splunk_ao/constants/__init__.py b/src/splunk_ao/constants/__init__.py index c8c96f7b..4040a835 100644 --- a/src/splunk_ao/constants/__init__.py +++ b/src/splunk_ao/constants/__init__.py @@ -7,18 +7,14 @@ DEFAULT_API_URL = "https://api.galileo.ai/" DEFAULT_CONSOLE_URL = "https://app.galileo.ai/" -# HTTP header prefix for all Splunk AO headers -SPLUNK_AO_HEADER_PREFIX = "Splunk-AO" - # Type definitions LoggerModeType = Literal["batch", "distributed"] __all__ = ( + "DEFAULT_AGENT_STREAM_NAME", "DEFAULT_API_URL", "DEFAULT_CONSOLE_URL", - "DEFAULT_AGENT_STREAM_NAME", "DEFAULT_MODE", "DEFAULT_PROJECT_NAME", - "SPLUNK_AO_HEADER_PREFIX", "LoggerModeType", ) diff --git a/src/splunk_ao/constants/tracing.py b/src/splunk_ao/constants/tracing.py deleted file mode 100644 index c62539bb..00000000 --- a/src/splunk_ao/constants/tracing.py +++ /dev/null @@ -1,8 +0,0 @@ -"""Constants for distributed tracing.""" - -from splunk_ao.constants import SPLUNK_AO_HEADER_PREFIX - -# HTTP header names for propagating distributed tracing context -# These headers follow the pattern of namespaced custom headers (Splunk-AO-*) -TRACE_ID_HEADER = f"{SPLUNK_AO_HEADER_PREFIX}-Trace-ID" -PARENT_ID_HEADER = f"{SPLUNK_AO_HEADER_PREFIX}-Parent-ID" diff --git a/src/splunk_ao/decorator.py b/src/splunk_ao/decorator.py index fb02472b..354f3d84 100644 --- a/src/splunk_ao/decorator.py +++ b/src/splunk_ao/decorator.py @@ -61,11 +61,11 @@ def call_llm(prompt, temperature=0.7): from galileo_core.schemas.logging.trace import Trace from splunk_ao.constants import LoggerModeType from splunk_ao.logger import SplunkAOLogger -from splunk_ao.logger.logger import STUB_TRACE_NAME from splunk_ao.schema.content_blocks import is_content_block_list from splunk_ao.schema.datasets import DatasetRecord from splunk_ao.schema.metrics import LocalMetricConfig from splunk_ao.schema.trace import SPAN_TYPE +from splunk_ao.session_context import _session_id_context, set_session_context from splunk_ao.shared.exceptions import ConfigurationError from splunk_ao.utils import _get_timestamp from splunk_ao.utils.env_helpers import _get_mode_or_default @@ -92,12 +92,6 @@ def call_llm(prompt, temperature=0.7): _experiment_id_context: ContextVar[str | None] = ContextVar("experiment_id_context", default=None) _span_stack_context: ContextVar[list[WorkflowSpan] | None] = ContextVar("span_stack_context", default=None) _mode_context: ContextVar[LoggerModeType | None] = ContextVar("mode_context", default=None) -_session_id_context: ContextVar[str | None] = ContextVar("session_id_context", default=None) - -# Distributed tracing context variables (for middleware) -_trace_id_context: ContextVar[str | None] = ContextVar("trace_id_context", default=None) -_parent_id_context: ContextVar[str | None] = ContextVar("parent_id_context", default=None) - # Context variables for dataset fields (ground truth/reference output) # These allow setting ground truth data that will be attached to all spans # created within the context, enabling scorers that require reference output. @@ -188,7 +182,15 @@ def __exit__( _trace_context.set(_get_or_init_list(_trace_stack).pop()) _mode_context.set(_get_or_init_list(_mode_stack).pop()) _span_stack_context.set(_get_or_init_list(_span_stack_stack).pop()) - _session_id_context.set(_get_or_init_list(_session_id_stack).pop()) + restored_session_id = _get_or_init_list(_session_id_stack).pop() + _session_id_context.set(restored_session_id) + 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) def __call__( self, @@ -1024,34 +1026,6 @@ def _handle_call_result( status_code = span_params.get("status_code") logger.conclude(output=output, duration_ns=span_params["duration_ns"], status_code=status_code) - # In distributed mode, update parent trace output after concluding a top-level workflow - # This ensures the trace shows the latest workflow's output (last workflow wins) - # Skip stub traces (created from distributed tracing headers - they're managed by the client) - if logger.mode == "distributed" and not stack: - current_parent = logger.current_parent() - if current_parent is not None and isinstance(current_parent, Trace): - is_stub_trace = current_parent.name == STUB_TRACE_NAME - - if not is_stub_trace: - # _coerce_output preserves str and List[ContentBlock], - # serializes everything else (Message, List[Document], etc.) to JSON string. - if output is not None: - current_parent.output = SplunkAOLogger._coerce_output(output) - if redacted_output is not None: - current_parent.redacted_output = SplunkAOLogger._coerce_output(redacted_output) - - # Update trace duration - # Note: In distributed mode, trace.created_at may be set by the server - # Using max() to ensure parent is never shorter than its children. - if current_parent.created_at: - elapsed_ns = convert_time_delta_to_ns(_get_timestamp() - current_parent.created_at) - workflow_ns = span_params.get("duration_ns", 0) - prev_ns = current_parent.metrics.duration_ns or 0 - current_parent.metrics.duration_ns = max(elapsed_ns, workflow_ns, prev_ns) - - if status_code is not None: - current_parent.status_code = status_code - else: # Non-concludable spans (llm, tool, retriever) are added to the parent span_methods = {"llm": "add_llm_span", "tool": "add_tool_span", "retriever": "add_retriever_span"} @@ -1234,12 +1208,6 @@ def get_logger_instance( "experiment_id": experiment_id or _experiment_id_context.get(), "mode": _get_mode_or_default(mode) if mode is not None else _mode_context.get(), } - trace_id_from_context = _trace_id_context.get() - span_id_from_context = _parent_id_context.get() - if trace_id_from_context: - kwargs["trace_id"] = trace_id_from_context - if span_id_from_context: - kwargs["span_id"] = span_id_from_context if ingestion_hook is not None: kwargs["ingestion_hook"] = ingestion_hook @@ -1374,10 +1342,6 @@ def reset(self) -> None: _span_stack_context.set([]) _trace_context.set(None) _session_id_context.set(None) - # Reset distributed tracing context - _trace_id_context.set(None) - _parent_id_context.set(None) - # Clear all stacks _get_or_init_list(_project_stack).clear() _get_or_init_list(_agent_stream_stack).clear() @@ -1476,6 +1440,7 @@ def start_session( def clear_session(self) -> None: """Clear the session in the active context logger instance.""" self.get_logger_instance().clear_session() + set_session_context(None) def set_session(self, session_id: str) -> None: """ @@ -1487,6 +1452,7 @@ def set_session(self, session_id: str) -> None: The id of the session to set. """ self.get_logger_instance().set_session(session_id) + set_session_context(session_id) splunk_ao_context = SplunkAODecorator() diff --git a/src/splunk_ao/exporter/sink.py b/src/splunk_ao/exporter/sink.py index b32a6088..d23b7505 100644 --- a/src/splunk_ao/exporter/sink.py +++ b/src/splunk_ao/exporter/sink.py @@ -54,14 +54,16 @@ def shutdown(self) -> None: def build_batch_processor(exporter: SpanExporter, config: BatchConfig | None = None) -> BatchSpanProcessor: - """Build a BatchSpanProcessor with explicit SDK defaults.""" - batch_config = config or BatchConfig() + """Build a BatchSpanProcessor with OTel defaults or explicit SDK configuration.""" + if config is None: + return BatchSpanProcessor(exporter) + return BatchSpanProcessor( exporter, - max_queue_size=batch_config.max_queue_size, - schedule_delay_millis=batch_config.schedule_delay_millis, - export_timeout_millis=batch_config.export_timeout_millis, - max_export_batch_size=batch_config.max_export_batch_size, + max_queue_size=config.max_queue_size, + schedule_delay_millis=config.schedule_delay_millis, + export_timeout_millis=config.export_timeout_millis, + max_export_batch_size=config.max_export_batch_size, ) diff --git a/src/splunk_ao/handlers/base_async_handler.py b/src/splunk_ao/handlers/base_async_handler.py index 19d0f91f..fa7b45ba 100644 --- a/src/splunk_ao/handlers/base_async_handler.py +++ b/src/splunk_ao/handlers/base_async_handler.py @@ -21,7 +21,9 @@ class SplunkAOAsyncBaseHandler(SplunkAOBaseHandler): _nodes : dict[UUID, Node] A dictionary of nodes, where the key is the run_id and the value is the node. _start_new_trace : bool - Whether to start a new trace when a chain starts. Set this to `False` to continue using the current trace. + Whether the handler owns and concludes a local trace lifecycle. An + active W3C parent is still inherited when this is ``True``. Set it to + ``False`` only when attaching to caller-owned active logger state. _flush_on_chain_end : bool Whether to flush the trace when a chain ends. """ @@ -42,34 +44,55 @@ async def async_commit(self) -> None: _logger.warning("Unable to add nodes to trace: Root node does not exist") return - owned_trace = None + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + self._finish_incremental_node(root_node) + if self._flush_on_chain_end: + await self._splunk_ao_logger.async_flush() + return + try: - if self._start_new_trace: - owned_trace = self._splunk_ao_logger.start_trace( - input=serialize_to_str(root_node.span_params.get("input", "")), - name=root_node.span_params.get("name"), - metadata=root_node.span_params.get("metadata"), + if self._owned_root is not None: + self._update_owned_root(root_node) + self._log_node_children(root_node) + root_output = self._root_output(root_node) + self._splunk_ao_logger.conclude( + output=serialize_to_str(root_output), + duration_ns=root_node.span_params.get("duration_ns"), + status_code=root_node.span_params.get("status_code"), ) - - self.log_node_tree(root_node) - root_output = root_node.span_params.get("output", "") - - if self._start_new_trace: + if self._owned_trace is not None: + self._conclude_owned_trace( + self._owned_trace, + output=serialize_to_str(root_output), + status_code=root_node.span_params.get("status_code"), + ) + elif self._start_new_trace: + if self._owned_trace is None: + self._start_owned_root(root_node) + if self._owned_trace is None: + return + + self.log_node_tree(root_node) + root_output = self._root_output(root_node) self._conclude_owned_trace( - owned_trace, + self._owned_trace, output=serialize_to_str(root_output), status_code=root_node.span_params.get("status_code"), ) + elif self._has_reusable_caller_root(): + # Preserve the pre-DTA visible tree while leaving the live + # caller operation current and caller-owned. + self.log_node_tree(root_node) + else: + _logger.warning("Unable to commit async handler telemetry: no caller-owned active operation") if self._flush_on_chain_end: await self._splunk_ao_logger.async_flush() except Exception: - if owned_trace is not None: - self._conclude_owned_trace(owned_trace, output="", status_code=500) + self._conclude_owned_state_on_failure() _logger.warning("Failed to commit async handler telemetry", exc_info=True) finally: - self._nodes.clear() - self._root_node = None + self._reset_handler_state() async def async_end_node(self, run_id: UUID, **kwargs: Any) -> None: """ @@ -94,6 +117,13 @@ async def async_end_node(self, run_id: UUID, **kwargs: Any) -> None: # Update node parameters node.span_params.update(**kwargs) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + is_root = self._root_node is node + self._finish_incremental_node(node) + if is_root and self._flush_on_chain_end: + await self._splunk_ao_logger.async_flush() + return + # Check if this is the root node and commit if so root = self._root_node if root and node.run_id == root.run_id: diff --git a/src/splunk_ao/handlers/base_handler.py b/src/splunk_ao/handlers/base_handler.py index c095e717..1eca5887 100644 --- a/src/splunk_ao/handlers/base_handler.py +++ b/src/splunk_ao/handlers/base_handler.py @@ -6,8 +6,10 @@ from uuid import UUID from splunk_ao import splunk_ao_context +from splunk_ao.handlers.span_lifecycle import HandlerSpanState, build_handler_step, finalize_handler_step from splunk_ao.logger import SplunkAOLogger from splunk_ao.schema.handlers import INTEGRATION, NODE_TYPE, Node +from splunk_ao.schema.logged import LoggedAgentSpan, LoggedWorkflowSpan from splunk_ao.schema.trace import TracesIngestRequest from splunk_ao.utils.serialization import convert_to_string_dict, serialize_to_str @@ -25,7 +27,9 @@ class SplunkAOBaseHandler: _nodes : dict[UUID, Node] A dictionary of nodes, where the key is the run_id and the value is the node. _start_new_trace : bool - Whether to start a new trace when a chain starts. Set this to `False` to continue using the current trace. + Whether the handler owns and concludes a local trace lifecycle. An + active W3C parent is still inherited when this is ``True``. Set it to + ``False`` only when attaching to caller-owned active logger state. _flush_on_chain_end : bool Whether to flush the trace when a chain ends. _root_node : Optional[Node] @@ -56,6 +60,10 @@ def __init__( self._nodes: dict[str, Node] = {} self._root_node: Node | None = None self._integration: INTEGRATION = integration + self._owned_trace: Any = None + self._owned_root: Any = None + self._owned_parent: Any = None + self._active_steps: dict[str, HandlerSpanState] = {} def commit(self) -> None: """Commit the nodes to the trace using the Splunk AO Logger. Optionally flush the trace.""" @@ -73,36 +81,270 @@ def commit(self) -> None: _logger.warning("Unable to add nodes to trace: Root node does not exist") return - owned_trace = None + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + self._finish_incremental_node(root_node) + if self._flush_on_chain_end: + self._splunk_ao_logger.flush() + return + try: - if self._start_new_trace: - owned_trace = self._splunk_ao_logger.start_trace( - input=SplunkAOLogger._coerce_output(root_node.span_params.get("input", "")), - name=root_node.span_params.get("name"), - metadata=root_node.span_params.get("metadata"), + if self._owned_root is not None: + self._update_owned_root(root_node) + self._log_node_children(root_node) + root_output = self._root_output(root_node) + self._splunk_ao_logger.conclude( + output=serialize_to_str(root_output), + duration_ns=root_node.span_params.get("duration_ns"), + status_code=root_node.span_params.get("status_code"), ) - - self.log_node_tree(root_node) - - # Conclude the trace with the root node's output - root_output = root_node.span_params.get("output", "") - - if self._start_new_trace: + if self._owned_trace is not None: + self._conclude_owned_trace( + self._owned_trace, + output=SplunkAOLogger._coerce_output(root_output), + status_code=root_node.span_params.get("status_code"), + ) + elif self._start_new_trace: + if self._owned_trace is None: + self._start_owned_root(root_node) + if self._owned_trace is None: + return + + # Leaf-only roots cannot remain active in the Path 1 logger: + # leaf creation is an immediate completion operation. Preserve + # the pre-DTA direct leaf topology at commit while the owned + # envelope retains the incoming context. + self.log_node_tree(root_node) + root_output = self._root_output(root_node) self._conclude_owned_trace( - owned_trace, + self._owned_trace, output=SplunkAOLogger._coerce_output(root_output), status_code=root_node.span_params.get("status_code"), ) + elif self._has_reusable_caller_root(): + # Preserve the pre-DTA visible tree: the buffered handler root + # remains a child of the caller-owned live operation. The + # handler concludes only that buffered subtree and returns to + # the same caller operation. + self.log_node_tree(root_node) + else: + _logger.warning("Unable to commit handler telemetry: no caller-owned active operation") if self._flush_on_chain_end: self._splunk_ao_logger.flush() except Exception: - if owned_trace is not None: - self._conclude_owned_trace(owned_trace, output="", status_code=500) + self._conclude_owned_state_on_failure() _logger.warning("Failed to commit handler telemetry", exc_info=True) finally: - self._nodes.clear() - self._root_node = None + self._reset_handler_state() + + 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 _start_incremental_step(self, node: Node) -> None: + """Create and activate one real callback span for normal OTLP egress.""" + node_id = str(node.run_id) + if node is self._root_node and self._owned_root is not None: + self._active_steps[node_id] = HandlerSpanState(step=self._owned_root, activation=None) + return + + if node.parent_run_id is not None: + parent_state = self._active_steps.get(str(node.parent_run_id)) + parent = parent_state.step if parent_state is not None else None + else: + parent = self._owned_parent or self._owned_trace + if parent is None: + raise RuntimeError(f"No active parent is available for handler node {node_id}") + + step = build_handler_step(node) + self._splunk_ao_logger._register_handler_step(step, parent) + activation = self._splunk_ao_logger._activate_handler_step(step) + self._active_steps[node_id] = HandlerSpanState(step=step, activation=activation) + + def _finish_incremental_node(self, node: Node) -> None: + """Finalize and enqueue one callback span when that callback ends.""" + node_id = str(node.run_id) + state = self._active_steps.get(node_id) + if state is None: + _logger.warning("Unable to complete handler node %s: no active span state", node_id) + if node is self._root_node: + self._reset_handler_state() + return + + is_root = node is self._root_node + try: + if is_root: + node.span_params["output"] = self._root_output(node) + if state.step is self._owned_root: + node.span_params["input"] = serialize_to_str(node.span_params.get("input", "")) + final = finalize_handler_step(node, state) + final = self._splunk_ao_logger._replace_handler_step(state.step, final) + state.step = final + 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 + + self._splunk_ao_logger._restore_handler_step_context(state.activation) + state.activation = None + if self._splunk_ao_logger.current_parent() is final: + self._splunk_ao_logger._set_current_parent(final._parent) + self._splunk_ao_logger._complete_handler_step(final) + self._active_steps.pop(node_id, None) + + if is_root: + root_output = node.span_params.get("output", "") + if self._owned_trace is not None: + self._conclude_owned_trace( + self._owned_trace, + output=SplunkAOLogger._coerce_output(root_output), + status_code=node.span_params.get("status_code"), + ) + 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() + + def _start_owned_root(self, root_node: Node) -> None: + """Open the handler-owned envelope and real root at callback start.""" + if self._owned_root is not None: + return + + if not self._start_new_trace: + # Caller-owned mode preserves the caller's proprietary cursor. + # The incremental handler root is explicitly parented beneath it + # and becomes the active OTel operation during its callback. + self._owned_parent = self._splunk_ao_logger.current_parent() + return + + if self._owned_trace is None: + self._owned_trace = self._splunk_ao_logger.start_trace( + input=SplunkAOLogger._coerce_output(root_node.span_params.get("input", "")), + name=root_node.span_params.get("name"), + metadata=root_node.span_params.get("metadata"), + ) + parent = self._owned_trace + + if parent is None: + return + self._owned_parent = parent + + # The legacy ingestion-hook path logs a leaf root at commit. Normal + # OTLP egress creates and activates it through _start_incremental_step. + if root_node.node_type not in ("agent", "chain", "workflow"): + return + + metadata = root_node.span_params.get("metadata") + if metadata is not None: + metadata = convert_to_string_dict(metadata) + step_number = self._step_number(metadata) + root_kwargs = { + "input": serialize_to_str(root_node.span_params.get("input", "")), + "name": root_node.span_params.get("name"), + "metadata": metadata, + "tags": root_node.span_params.get("tags"), + "created_at": root_node.span_params.get("created_at"), + "status_code": root_node.span_params.get("status_code"), + "step_number": step_number, + } + if root_node.node_type == "agent": + self._owned_root = self._splunk_ao_logger.add_agent_span(**root_kwargs) + else: + self._owned_root = self._splunk_ao_logger.add_workflow_span(**root_kwargs) + + if self._owned_root is None: + self._conclude_owned_state_on_failure() + self._owned_trace = None + self._owned_parent = None + + def _update_owned_root(self, root_node: Node) -> None: + """Apply final callback data to the live root before it is completed.""" + if self._owned_root is None: + return + self._sync_owned_root_kind(root_node) + self._owned_root.input = serialize_to_str(root_node.span_params.get("input", "")) + if root_node.span_params.get("name") is not None: + self._owned_root.name = root_node.span_params["name"] + if root_node.span_params.get("tags") is not None: + self._owned_root.tags = root_node.span_params["tags"] + metadata = root_node.span_params.get("metadata") + if metadata is not None: + converted_metadata = convert_to_string_dict(metadata) + self._owned_root.user_metadata = converted_metadata + self._owned_root.step_number = self._step_number(converted_metadata) + + @staticmethod + def _step_number(metadata: dict[str, str] | None) -> int | None: + """Read the LangGraph step number using the pre-DTA conversion rule.""" + if not metadata or not (value := metadata.get("langgraph_step")): + return None + try: + return int(value) + except Exception as exc: + _logger.warning(f"Invalid step number: {value}, exception raised {exc}") + return None + + def _sync_owned_root_kind(self, root_node: Node) -> None: + """Preserve a late LangGraph chain-to-agent classification. + + LangGraph identifies some roots from child metadata after the root + callback has started. Replace only the proprietary model object while + retaining its UUID and preassigned OTel identity, so outbound W3C + context and the eventually exported agent span remain identical. + """ + if root_node.node_type != "agent" or not isinstance(self._owned_root, LoggedWorkflowSpan): + return + if self._splunk_ao_logger.current_parent() is not self._owned_root: + return + + previous = self._owned_root + parent = previous._parent + if parent is None: + return + replacement = LoggedAgentSpan.model_validate(previous.model_dump(exclude={"type"})) + replacement._parent = parent + replacement.spans = previous.spans + parent.spans = [replacement if child is previous else child for child in parent.spans] + self._owned_root = replacement + self._splunk_ao_logger._set_current_parent(replacement) + + def _log_node_children(self, root_node: Node) -> None: + """Log buffered descendants beneath the already-active real root.""" + for child_id in root_node.children: + child_node = self._nodes.get(child_id) + if child_node is not None: + self.log_node_tree(child_node) + else: + _logger.warning(f"Child node {child_id} not found") + + def _root_output(self, root_node: Node) -> Any: + output = root_node.span_params.get("output", "") + if output or not root_node.children: + return output + last_child = self._nodes.get(root_node.children[-1]) + return last_child.span_params.get("output", "") if last_child is not None else "" + + def _has_reusable_caller_root(self) -> bool: + """Return whether start_new_trace=False has caller-owned active state.""" + return self._splunk_ao_logger.current_parent() is not None + + def _conclude_owned_state_on_failure(self) -> None: + """Close only the handler root/envelope and return to caller state.""" + while ( + self._owned_parent is not None + and self._splunk_ao_logger.current_parent() is not None + and self._splunk_ao_logger.current_parent() is not self._owned_parent + ): + self._splunk_ao_logger.conclude(output="", status_code=500) + if self._owned_trace is not None: + self._conclude_owned_trace(self._owned_trace, output="", status_code=500) def _conclude_owned_trace(self, trace: Any, output: Any, status_code: int | None) -> None: current_parent = self._splunk_ao_logger.current_parent() @@ -142,7 +384,7 @@ def log_node_tree(self, node: Node) -> None: _logger.warning(f"Invalid step number: {metadata_step_number}, exception raised {e}") # Log the current node based on its type - if node.node_type == "chain": + if node.node_type in ("chain", "workflow"): self._splunk_ao_logger.add_workflow_span( input=input_, output=output, @@ -286,6 +528,22 @@ def start_node(self, node_type: NODE_TYPE, parent_run_id: UUID | None, run_id: U if not self._root_node: _logger.debug(f"Setting root node to {node_id}") self._root_node = node + try: + self._start_owned_root(node) + except Exception: + self._conclude_owned_state_on_failure() + self._owned_trace = None + self._owned_root = None + self._owned_parent = None + _logger.warning("Failed to start handler root telemetry", exc_info=True) + + # A LangGraph child may have reclassified the existing root from a + # chain to an agent immediately before this callback reached us. + if self._root_node is not None: + self._sync_owned_root_kind(self._root_node) + root_state = self._active_steps.get(str(self._root_node.run_id)) + if root_state is not None and self._owned_root is not None: + root_state.step = self._owned_root # Add to parent's children if parent exists if parent_run_id: @@ -295,6 +553,13 @@ def start_node(self, node_type: NODE_TYPE, parent_run_id: UUID | None, run_id: U else: _logger.debug(f"Parent node {parent_node_id} not found for {node_id}") + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + try: + self._start_incremental_step(node) + except Exception: + self._conclude_owned_state_on_failure() + _logger.warning("Failed to start handler span telemetry for node %s", node_id, exc_info=True) + return node def end_node(self, run_id: UUID, **kwargs: Any) -> None: @@ -320,6 +585,13 @@ def end_node(self, run_id: UUID, **kwargs: Any) -> None: # Update node parameters node.span_params.update(**kwargs) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + is_root = self._root_node is node + self._finish_incremental_node(node) + if is_root and self._flush_on_chain_end: + self._splunk_ao_logger.flush() + return + # Check if this is the root node and commit if so root = self._root_node if root and node.run_id == root.run_id: diff --git a/src/splunk_ao/handlers/openai_agents/handler.py b/src/splunk_ao/handlers/openai_agents/handler.py index 6826a880..d4236263 100644 --- a/src/splunk_ao/handlers/openai_agents/handler.py +++ b/src/splunk_ao/handlers/openai_agents/handler.py @@ -9,6 +9,7 @@ from galileo_core.schemas.logging.span import LlmMetrics, LlmSpan from galileo_core.schemas.logging.span import Span as SplunkAOSpan from splunk_ao import SplunkAOLogger, splunk_ao_context +from splunk_ao.handlers.span_lifecycle import HandlerSpanState, build_handler_step, finalize_handler_step from splunk_ao.schema.handlers import Node from splunk_ao.utils import _get_timestamp from splunk_ao.utils.openai_agents import ( @@ -59,6 +60,10 @@ def __init__(self, splunk_ao_logger: SplunkAOLogger | None = None, flush_on_trac self._last_status_code: int | None = None self._first_input: Any = None self._owned_trace: Any = None + self._owned_root: Any = None + self._owned_root_node_id: str | None = None + self._caller_parent: Any = None + self._active_steps: dict[str, HandlerSpanState] = {} def on_trace_start(self, trace: Trace) -> None: """Called when an OpenAI Agent trace starts.""" @@ -73,6 +78,17 @@ def on_trace_start(self, trace: Trace) -> None: }, ) self._nodes[trace.trace_id] = node + self._caller_parent = self._splunk_ao_logger.current_parent() + if self._caller_parent is None: + try: + self._owned_trace = self._splunk_ao_logger.start_trace( + input="Agent Workflow", + name="Trace", + created_at=datetime.fromisoformat(node.span_params["start_time_iso"]), + ) + except Exception: + self._owned_trace = None + _logger.warning("Failed to start OpenAI Agents trace telemetry", exc_info=True) def on_trace_end(self, trace: Trace) -> None: """Called when an OpenAI Agent trace ends.""" @@ -91,11 +107,18 @@ def on_trace_end(self, trace: Trace) -> None: self._conclude_current_trace_on_failure() _logger.warning("Failed to commit OpenAI Agents telemetry", exc_info=True) finally: + for state in self._active_steps.values(): + self._splunk_ao_logger._restore_handler_step_context(state.activation) + self._splunk_ao_logger._release_otel_context(state.step) self._nodes = {} self._last_output = None self._last_status_code = None self._first_input = None self._owned_trace = None + self._owned_root = None + self._owned_root_node_id = None + self._caller_parent = None + self._active_steps = {} def _commit_trace(self, trace: Trace) -> None: if not self._nodes: @@ -103,11 +126,48 @@ def _commit_trace(self, trace: Trace) -> None: return root_node = self._nodes.get(trace.trace_id) - if root_node: - self._log_node_tree(root_node, first_node=True) - else: + if root_node is None: _logger.warning(f"Root node {trace.trace_id} not found") - self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + return + + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + if self._active_steps: + _logger.warning( + "OpenAI Agents trace %s ended with %d unfinished callback spans", + trace.trace_id, + len(self._active_steps), + ) + if self._owned_trace is not None: + self._owned_trace.input = self._first_input or "Agent Workflow" + self._owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") + if self._splunk_ao_logger.current_parent() is self._owned_trace: + self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + return + + live_root_node = self._nodes.get(self._owned_root_node_id) if self._owned_root_node_id else None + if self._owned_trace is not None and self._owned_root is not None and live_root_node is not None: + self._owned_trace.input = self._first_input or "Agent Workflow" + self._owned_trace.metrics.duration_ns = root_node.span_params.get("duration_ns") + self._update_owned_root(live_root_node) + self._log_node_tree(live_root_node, reuse_current=True) + for child_id in root_node.children: + if child_id == self._owned_root_node_id: + continue + child = self._nodes.get(child_id) + if child is not None: + self._log_node_tree(child) + self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) + return + + # A caller-owned operation remains current; log framework operations + # beneath it without ever concluding the caller's parent. + for child_id in root_node.children: + child = self._nodes.get(child_id) + if child is not None: + self._log_node_tree(child) + + if self._owned_trace is not None: + self._splunk_ao_logger.conclude(output=self._last_output, status_code=self._last_status_code) def _conclude_current_trace_on_failure(self) -> None: if self._owned_trace is None: @@ -123,7 +183,7 @@ def _conclude_current_trace_on_failure(self) -> None: if root is self._owned_trace: self._splunk_ao_logger.conclude(output="", status_code=500, conclude_all=True) - def _log_node_tree(self, node: Node, first_node: bool = False) -> None: + def _log_node_tree(self, node: Node, reuse_current: bool = False) -> None: """ Log a node and its children recursively. @@ -131,8 +191,8 @@ def _log_node_tree(self, node: Node, first_node: bool = False) -> None: ---------- node : Node The node to log. - first_node : bool - Whether this is the root trace node. + reuse_current : bool + Whether this node is the live root already created at span start. """ is_workflow_span = False input = node.span_params.get("input", "") @@ -145,26 +205,18 @@ def _log_node_tree(self, node: Node, first_node: bool = False) -> None: # Convert metadata to a dict[str, str] if metadata is not None: metadata = convert_to_string_dict(metadata) - if first_node: - self._owned_trace = self._splunk_ao_logger.add_trace( - input=self._first_input or "Agent Workflow", - output=self._last_output, - duration_ns=node.span_params.get("duration_ns"), - created_at=start_time_iso, - name="Trace", - tags=tags, - ) # Log the current node based on its type - elif node.node_type in ("agent", "chain", "workflow"): - self._splunk_ao_logger.add_workflow_span( - input=input or node.node_type.capitalize() + " Step", - output=output, - name=name, - metadata=metadata, - tags=tags, - created_at=start_time_iso, - duration_ns=node.span_params.get("duration_ns"), - ) + if node.node_type in ("agent", "chain", "workflow"): + if not reuse_current: + self._splunk_ao_logger.add_workflow_span( + input=input or node.node_type.capitalize() + " Step", + output=output, + name=name, + metadata=metadata, + tags=tags, + created_at=start_time_iso, + duration_ns=node.span_params.get("duration_ns"), + ) is_workflow_span = True elif node.node_type in ("llm", "chat"): tools = node.span_params.get("tools") @@ -239,10 +291,97 @@ def _log_node_tree(self, node: Node, first_node: bool = False) -> None: if error: output = error status_code = 500 - self._splunk_ao_logger.conclude(output=serialize_to_str(output), status_code=status_code) + self._splunk_ao_logger.conclude( + output=serialize_to_str(output), + duration_ns=node.span_params.get("duration_ns"), + status_code=status_code, + ) self._last_status_code = status_code self._last_output = output + def _start_owned_root(self, node: Node) -> None: + """Create the first top-level real operation while framework work is active.""" + if self._owned_trace is None or self._owned_root is not None: + return + if node.node_type not in ("agent", "chain", "workflow"): + return + + metadata = node.span_params.get("metadata") + if metadata is not None: + metadata = convert_to_string_dict(metadata) + self._owned_root = self._splunk_ao_logger.add_workflow_span( + input=node.span_params.get("input") or node.node_type.capitalize() + " Step", + name=node.span_params.get("name"), + metadata=metadata, + tags=node.span_params.get("tags"), + created_at=datetime.fromisoformat(node.span_params["start_time_iso"]), + ) + if self._owned_root is not None: + self._owned_root_node_id = str(node.run_id) + + def _update_owned_root(self, node: Node) -> None: + """Apply final OpenAI callback data to the already-live root.""" + if self._owned_root is None: + return + self._owned_root.input = node.span_params.get("input") or node.node_type.capitalize() + " Step" + if node.span_params.get("name") is not None: + self._owned_root.name = node.span_params["name"] + if node.span_params.get("tags") is not None: + self._owned_root.tags = node.span_params["tags"] + metadata = node.span_params.get("metadata") + if metadata is not None: + self._owned_root.user_metadata = convert_to_string_dict(metadata) + + def _start_incremental_span(self, node: Node) -> None: + """Create and activate one real OpenAI Agents callback span.""" + node_id = str(node.run_id) + if node_id == self._owned_root_node_id and self._owned_root is not None: + self._active_steps[node_id] = HandlerSpanState(step=self._owned_root, activation=None) + return + + if str(node.parent_run_id) in self._active_steps: + parent = self._active_steps[str(node.parent_run_id)].step + else: + parent = self._owned_trace or self._caller_parent + if parent is None: + raise RuntimeError(f"No active parent is available for OpenAI Agents span {node_id}") + + step = build_handler_step(node, openai_agents=True) + self._splunk_ao_logger._register_handler_step(step, parent) + activation = self._splunk_ao_logger._activate_handler_step(step) + self._active_steps[node_id] = HandlerSpanState(step=step, activation=activation) + + def _finish_incremental_span(self, node: Node) -> None: + """Finalize and enqueue an OpenAI Agents span at callback completion.""" + node_id = str(node.run_id) + state = self._active_steps.get(node_id) + if state is None: + _logger.warning("Unable to complete OpenAI Agents span %s: no active state", node_id) + return + try: + if node.node_type in ("agent", "chain", "workflow") and not node.span_params.get("output"): + last_child = self._nodes.get(node.children[-1]) if node.children else None + if last_child is not None: + node.span_params["output"] = last_child.span_params.get("output", "") + final = finalize_handler_step(node, state, openai_agents=True) + final = self._splunk_ao_logger._replace_handler_step(state.step, final) + state.step = final + if node_id == self._owned_root_node_id: + self._owned_root = final + + self._splunk_ao_logger._restore_handler_step_context(state.activation) + state.activation = None + if self._splunk_ao_logger.current_parent() is final: + self._splunk_ao_logger._set_current_parent(final._parent) + self._splunk_ao_logger._complete_handler_step(final) + self._active_steps.pop(node_id, None) + self._last_output = node.span_params.get("output") + self._last_status_code = node.span_params.get("status_code", 200) + except Exception: + self._splunk_ao_logger._restore_handler_step_context(state.activation) + state.activation = None + _logger.warning("Failed to complete OpenAI Agents span %s", node_id, exc_info=True) + def on_span_start(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span starts.""" span_id = span.span_id @@ -320,6 +459,21 @@ def on_span_start(self, span: Span[Any]) -> None: _logger.warning(f"Parent node {parent_id} not found for span {span_id} in trace {trace_id}") return parent_node.children.append(span_id) + if parent_id == trace_id: + try: + self._start_owned_root(node) + except Exception: + self._conclude_current_trace_on_failure() + self._owned_trace = None + self._owned_root = None + self._owned_root_node_id = None + _logger.warning("Failed to start OpenAI Agents root telemetry", exc_info=True) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + try: + self._start_incremental_span(node) + except Exception: + self._conclude_current_trace_on_failure() + _logger.warning("Failed to start OpenAI Agents span telemetry for %s", span_id, exc_info=True) def on_span_end(self, span: Span[Any]) -> None: """Called when an OpenAI Agent span ends.""" @@ -404,6 +558,8 @@ def on_span_end(self, span: Span[Any]) -> None: # Update the node's parameters node.span_params.update(end_params) + if not getattr(self._splunk_ao_logger, "_ingestion_hook", None): + self._finish_incremental_span(node) def shutdown(self) -> None: """Called when the application stops. Flushes any remaining logs.""" diff --git a/src/splunk_ao/handlers/span_lifecycle.py b/src/splunk_ao/handlers/span_lifecycle.py new file mode 100644 index 00000000..2e76ef80 --- /dev/null +++ b/src/splunk_ao/handlers/span_lifecycle.py @@ -0,0 +1,143 @@ +"""Shared callback-to-span construction for incremental handler telemetry.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import UTC, datetime +from uuid import UUID, uuid4 + +from galileo_core.schemas.logging.agent import AgentType +from galileo_core.schemas.logging.span import LlmMetrics, RetrieverSpan, StepWithChildSpans, ToolSpan +from galileo_core.schemas.logging.step import BaseStep, Metrics +from splunk_ao.logger.logger import HandlerStepContext +from splunk_ao.schema.handlers import Node +from splunk_ao.schema.logged import LoggedAgentSpan, LoggedLlmSpan, LoggedWorkflowSpan +from splunk_ao.utils.retrievers import convert_to_documents +from splunk_ao.utils.serialization import convert_to_string_dict, serialize_to_str + + +@dataclass +class HandlerSpanState: + """Mutable ownership state for one in-flight framework callback.""" + + step: BaseStep + activation: HandlerStepContext | None + + +def _created_at(node: Node) -> datetime: + created_at = node.span_params.get("created_at") + if isinstance(created_at, datetime): + return created_at + start_time_iso = node.span_params.get("start_time_iso") + if isinstance(start_time_iso, str) and start_time_iso: + return datetime.fromisoformat(start_time_iso) + return datetime.now(tz=UTC) + + +def _metadata(node: Node) -> dict[str, str] | None: + metadata = node.span_params.get("metadata") + return convert_to_string_dict(metadata) if metadata is not None else None + + +def _step_number(metadata: dict[str, str] | None) -> int | None: + if not metadata or not (value := metadata.get("langgraph_step")): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def build_handler_step( + node: Node, *, step_id: UUID | None = None, children: list[BaseStep] | None = None, openai_agents: bool = False +) -> BaseStep: + """Build a schema-validated provisional or final span from a callback node.""" + params = node.span_params + metadata = _metadata(node) + step_number = _step_number(metadata) + input_value = params.get("input", "") + output = params.get("output", "") + status_code = params.get("status_code") + if params.get("error"): + output = params["error"] + status_code = 500 + if openai_agents and node.node_type in ("agent", "chain", "workflow", "tool"): + input_value = input_value or node.node_type.capitalize() + " Step" + + name = params.get("name") + if openai_agents and node.node_type in ("llm", "chat"): + name = name or "LLM Span" + common = { + "name": name, + "created_at": _created_at(node), + "user_metadata": metadata, + "tags": params.get("tags"), + "status_code": status_code, + "id": step_id or uuid4(), + "step_number": step_number, + } + duration_ns = params.get("duration_ns") + + if node.node_type in ("chain", "workflow") or (openai_agents and node.node_type == "agent"): + return LoggedWorkflowSpan( + **common, + input=input_value, + output=serialize_to_str(output), + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + ) + if node.node_type == "agent": + return LoggedAgentSpan( + **common, + input=input_value, + output=serialize_to_str(output), + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + agent_type=params.get("agent_type", AgentType.default), + ) + 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 == "retriever": + retriever_common = {**common, "status_code": None} + return RetrieverSpan( + **retriever_common, + input=serialize_to_str(input_value), + output=convert_to_documents(output, "output"), + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + ) + if node.node_type == "tool": + return ToolSpan( + **common, + input=serialize_to_str(input_value), + output=serialize_to_str(output) if output is not None else None, + metrics=Metrics(duration_ns=duration_ns), + spans=children or [], + tool_call_id=params.get("tool_call_id"), + ) + raise ValueError(f"Unsupported handler node type: {node.node_type}") + + +def finalize_handler_step(node: Node, state: HandlerSpanState, *, openai_agents: bool = False) -> BaseStep: + """Build the final validated form while retaining the provisional UUID and children.""" + children = list(state.step.spans) if isinstance(state.step, StepWithChildSpans) else None + return build_handler_step(node, step_id=state.step.id, children=children, openai_agents=openai_agents) diff --git a/src/splunk_ao/http_instrumentation.py b/src/splunk_ao/http_instrumentation.py new file mode 100644 index 00000000..07749a2c --- /dev/null +++ b/src/splunk_ao/http_instrumentation.py @@ -0,0 +1,207 @@ +"""Supported automatic OpenTelemetry export and HTTP instrumentation setup.""" + +from __future__ import annotations + +from typing import Any, cast +from weakref import WeakSet + +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider + +from splunk_ao.otel import TracerProvider, add_splunk_ao_span_processor +from splunk_ao.session_context import install_session_propagator + +_INSTALL_MESSAGE = ( + "Automatic distributed tracing requires optional dependencies. " + "Install them with: pip install 'splunk-ao[distributed-tracing]'" +) + + +_client_provider_ids: dict[str, int] = {} +_instrumented_apps: set[tuple[int, int, str]] = set() +_configured_providers: WeakSet[SDKTracerProvider] = WeakSet() + + +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, + } + + +def _validate_client_ownership(names: tuple[str, ...], provider_id: int) -> None: + """Reject provider conflicts before any instrumentation is changed.""" + for name in names: + configured_provider_id = _client_provider_ids.get(name) + if configured_provider_id is not None and configured_provider_id != provider_id: + raise RuntimeError( + f"{name} is already instrumented through Splunk AO with another tracer provider; " + "configure each process-wide client instrumentor through one owner" + ) + + +def _resolve_app_framework(app: Any) -> str: + try: + from fastapi import FastAPI # noqa: PLC0415 + from starlette.applications import Starlette # noqa: PLC0415 + except ImportError as exc: + raise ImportError(_INSTALL_MESSAGE) from exc + + if isinstance(app, FastAPI): + framework = "fastapi" + elif isinstance(app, Starlette): + framework = "starlette" + else: + raise TypeError("app must be a FastAPI or Starlette application") + return framework + + +def _instrument_app(app: Any, instrumentors: dict[str, Any], tracer_provider: TracerProvider) -> None: + framework = _resolve_app_framework(app) + + key = (id(app), id(tracer_provider), framework) + if key in _instrumented_apps: + return + instrumentors[framework].instrument_app(app, tracer_provider=tracer_provider) + _instrumented_apps.add(key) + + +def _instrument_supported_transports( + *, + instrumentors: dict[str, type], + tracer_provider: TracerProvider, + app: Any | None, + instrument_requests: bool, + instrument_httpx: bool, + instrument_aiohttp_client: bool, +) -> None: + clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} + enabled_clients = tuple(name for name, enabled in clients.items() if enabled) + provider_id = id(tracer_provider) + _validate_client_ownership(enabled_clients, provider_id) + + if app is not None: + _instrument_app(app, instrumentors, tracer_provider) + + for name in enabled_clients: + if _client_provider_ids.get(name) == provider_id: + continue + instrumentors[name]().instrument(tracer_provider=tracer_provider) + _client_provider_ids[name] = provider_id + + install_session_propagator() + + +def instrument_distributed_tracing( + *, + tracer_provider: TracerProvider | None = None, + app: Any | None = None, + instrument_requests: bool = True, + instrument_httpx: bool = True, + instrument_aiohttp_client: bool = True, +) -> None: + """Enable supported upstream OTel server and client instrumentation. + + Parameters + ---------- + tracer_provider : TracerProvider | None + Caller-owned provider used by the upstream instrumentors. The current + global provider is used when omitted; this function never replaces it. + app : Any | None + Optional FastAPI or Starlette application to instrument. + instrument_requests : bool + Whether to instrument process-wide Requests clients. + instrument_httpx : bool + Whether to instrument process-wide HTTPX sync and async clients. + instrument_aiohttp_client : bool + Whether to instrument process-wide aiohttp clients. + + Notes + ----- + Client instrumentation is process-wide. Configure each supported client + through one owner. The caller also owns provider shutdown. + """ + instrumentors = _load_instrumentors() + resolved_provider = tracer_provider or cast(TracerProvider, trace.get_tracer_provider()) + _instrument_supported_transports( + instrumentors=instrumentors, + tracer_provider=resolved_provider, + app=app, + instrument_requests=instrument_requests, + instrument_httpx=instrument_httpx, + instrument_aiohttp_client=instrument_aiohttp_client, + ) + + +def configure_distributed_tracing( + *, + tracer_provider: SDKTracerProvider | None = None, + app: Any | None = None, + instrument_requests: bool = True, + instrument_httpx: bool = True, + instrument_aiohttp_client: bool = True, +) -> SDKTracerProvider: + """Configure Splunk AO export and supported automatic HTTP propagation. + + Parameters + ---------- + tracer_provider : SDKTracerProvider | None + Optional application-owned OpenTelemetry SDK provider. A new provider + is created when omitted. This function never replaces the process-global + provider. + app : Any | None + Optional FastAPI or Starlette application to instrument. + instrument_requests : bool + Whether to instrument process-wide Requests clients. + instrument_httpx : bool + Whether to instrument process-wide HTTPX sync and async clients. + instrument_aiohttp_client : bool + Whether to instrument process-wide aiohttp clients. + + Returns + ------- + SDKTracerProvider + The configured provider. The caller owns it and must call ``shutdown()`` + during process teardown. + + Notes + ----- + Repeated calls with the same provider do not attach another Splunk AO span + processor. Use :func:`instrument_distributed_tracing` instead when export + has already been configured separately on the provider. + """ + # Resolve optional imports before constructing a processor and its worker so + # a missing extra cannot leave partially configured background resources. + instrumentors = _load_instrumentors() + resolved_provider = tracer_provider or SDKTracerProvider() + clients = {"requests": instrument_requests, "httpx": instrument_httpx, "aiohttp-client": instrument_aiohttp_client} + enabled_clients = tuple(name for name, enabled in clients.items() if enabled) + _validate_client_ownership(enabled_clients, id(resolved_provider)) + if app is not None: + _resolve_app_framework(app) + + if resolved_provider not in _configured_providers: + add_splunk_ao_span_processor(resolved_provider) + _configured_providers.add(resolved_provider) + + _instrument_supported_transports( + instrumentors=instrumentors, + tracer_provider=resolved_provider, + app=app, + instrument_requests=instrument_requests, + instrument_httpx=instrument_httpx, + instrument_aiohttp_client=instrument_aiohttp_client, + ) + return resolved_provider diff --git a/src/splunk_ao/logger/logger.py b/src/splunk_ao/logger/logger.py index 8ebd8937..68575824 100644 --- a/src/splunk_ao/logger/logger.py +++ b/src/splunk_ao/logger/logger.py @@ -1,7 +1,6 @@ import asyncio import atexit import contextlib -import copy import inspect import json import logging @@ -13,10 +12,6 @@ from datetime import datetime from typing import TYPE_CHECKING, Any, TypeVar, Union -if TYPE_CHECKING: - from splunk_ao.handlers.agent_control import SplunkAOAgentControlBridge - -import backoff from opentelemetry import context as otel_context from opentelemetry import trace as otel_trace from opentelemetry.sdk.trace.id_generator import RandomIdGenerator @@ -32,16 +27,14 @@ LlmSpanAllowedInputType, LlmSpanAllowedOutputType, RetrieverSpan, - Span, StepWithChildSpans, ToolSpan, ) -from galileo_core.schemas.logging.step import BaseStep, Metrics, StepType +from galileo_core.schemas.logging.step import BaseStep, Metrics from galileo_core.schemas.logging.trace import Trace from galileo_core.schemas.shared.traces_logger import TracesLogger from splunk_ao.agent_streams import AgentStreams from splunk_ao.constants import LoggerModeType -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER from splunk_ao.converter import SpanConverter from splunk_ao.deployment import DeploymentMode, O11yConfig, StandaloneConfig, resolve_deployment from splunk_ao.exceptions import SplunkAOLoggerException @@ -56,7 +49,6 @@ ) from splunk_ao.exporter.diagnostics import get_export_health from splunk_ao.logger.control import ControlAppliesTo, ControlCheckStage, ControlResult -from splunk_ao.logger.task_handler import ThreadPoolTaskHandler from splunk_ao.projects import Projects from splunk_ao.schema.content_blocks import ( DataContentBlock, @@ -81,42 +73,29 @@ LogRecordsSearchRequest, RetrieverSpanAllowedOutputType, SessionCreateRequest, - SpansIngestRequest, - SpanUpdateRequest, TracesIngestRequest, - TraceUpdateRequest, ) +from splunk_ao.session_context import get_effective_session_id, set_session_context from splunk_ao.traces import Traces -from splunk_ao.utils.decorators import ( - async_warn_catch_exception, - nop_async, - nop_sync, - retry_on_transient_http_error, - warn_catch_exception, -) +from splunk_ao.utils.decorators import async_warn_catch_exception, nop_async, nop_sync, warn_catch_exception from splunk_ao.utils.env_helpers import _get_mode_or_default from splunk_ao.utils.metrics import populate_local_metrics from splunk_ao.utils.retrievers import convert_to_documents from splunk_ao.utils.serialization import serialize_to_str +if TYPE_CHECKING: + from splunk_ao.handlers.agent_control import SplunkAOAgentControlBridge + # Type alias for metadata values that can be auto-converted to strings MetadataValue = str | bool | int | float | None StepT = TypeVar("StepT", bound=BaseStep) -STREAMING_MAX_RETRIES = 5 -STREAMING_MAX_TIME_SECONDS = 70 # Maximum time to spend retrying a single request -DISTRIBUTED_FLUSH_TIMEOUT_SECONDS = 90 # Timeout for waiting on background trace/span update tasks during flush() -# Default upper bound for shutdown wait in terminate(). Independent from -# DISTRIBUTED_FLUSH_TIMEOUT_SECONDS (which guards live flushes); kept aligned -# at 90s so explicit terminate() calls behave like flush() by default. DEFAULT_TERMINATE_TIMEOUT_SECONDS = 90 # Absolute threshold above which a slow SplunkAOLogger.terminate() shutdown is # logged as a warning. Fast-path shutdowns are sub-millisecond; >1s always # indicates a real anomaly (busy-poll, stuck task, in-flight HTTP retry) and # should be visible in CI logs regardless of the configured timeout. _SLOW_SHUTDOWN_WARN_THRESHOLD_SECONDS = 1.0 -STUB_TRACE_NAME = "stub_trace" # Name for stub traces created from distributed tracing headers - _logger = logging.getLogger("splunk_ao.logger") _otel_id_generator = RandomIdGenerator() @@ -127,6 +106,8 @@ class OtelIds: span_context: SpanContext parent_span_context: SpanContext | None + exportable: bool + session_id: str | None @dataclass(frozen=True) @@ -136,9 +117,19 @@ class ActiveOtelContext: logger_id: int step_id: uuid.UUID span_context: SpanContext + exportable: bool token: Token +@dataclass(frozen=True) +class HandlerStepContext: + """OTel activation owned by one framework callback operation.""" + + prior_context: otel_context.Context + token: Token + span_context: SpanContext + + @dataclass(frozen=True) class OtelContextState: """Request-local state for all active proprietary logger contexts.""" @@ -150,6 +141,22 @@ class OtelContextState: _otel_context_state: ContextVar[OtelContextState | None] = ContextVar("_otel_context_state", default=None) +def _has_active_exportable_span_context() -> bool: + """Return whether the current OTel context can be propagated as a wire parent.""" + current = otel_trace.get_current_span().get_span_context() + if not current.is_valid: + return False + + state = _otel_context_state.get() + if state is None or not state.active_contexts: + return True + + active = state.active_contexts[-1] + if active.span_context != current: + return True + return active.exportable + + class SplunkAOLogger(TracesLogger): """ This class can be used to upload traces to Splunk AO. @@ -218,16 +225,12 @@ class SplunkAOLogger(TracesLogger): agent_stream_id: str | None = None experiment_id: str | None = None session_id: str | None = None - trace_id: str | None = None - span_id: str | None = None local_metrics: list[LocalMetricConfig] | None = None mode: LoggerModeType | None = None _session_external_id: str | None = None _logger = logging.getLogger("splunk_ao.logger") _traces_client: Union["Traces", None] = None - _task_handler: ThreadPoolTaskHandler - _trace_completion_submitted: bool _otel_ids: dict[uuid.UUID, OtelIds] = PrivateAttr(default_factory=dict) _pending_otel_steps: set[uuid.UUID] = PrivateAttr(default_factory=set) @@ -238,8 +241,6 @@ def __init__( agent_stream: str | None = None, agent_stream_id: str | None = None, experiment_id: str | None = None, - trace_id: str | None = None, - span_id: str | None = None, local_metrics: list[LocalMetricConfig] | None = None, mode: str | None = None, ingestion_hook: Callable[[TracesIngestRequest], None] | None = None, @@ -261,25 +262,12 @@ def __init__( Agent stream ID. experiment_id: Optional[str] Experiment ID. Used by the experiment runner. - trace_id: Optional[str] - Trace ID for distributed tracing. This can only be used in "distributed" mode. - - When provided, creates a local stub trace without fetching from the backend. - This allows downstream services to continue a distributed trace without waiting - for backend ingestion. - span_id: Optional[str] - Parent span ID for distributed tracing. This can only be used in "distributed" mode. - - When provided, creates a local stub span without fetching from the backend. - This allows downstream services to continue a distributed trace without waiting - for backend ingestion. local_metrics: Optional[list[LocalMetricConfig]] Local metrics mode: Optional[str] Logger mode: "batch" or "distributed". Defaults to SPLUNK_AO_MODE env var, or "batch" if not set. - Both modes enqueue completed spans for scheduled OTLP batch export, and - flush() only drains the export queue. "distributed" additionally enables - the legacy trace_id/span_id continuation parameters. + Both accepted values enqueue completed spans for the same scheduled OTLP + batch export. The value is retained temporarily for ingestion-hook compatibility. ingestion_hook: Optional[Callable[[TracesIngestRequest], None]] A callable that intercepts trace data before ingestion. This hook is called when the logger is flushed and can be a @@ -290,7 +278,6 @@ def __init__( super().__init__() mode = _get_mode_or_default(mode) self.mode: LoggerModeType = mode - self._task_counter = 0 self._terminated = False self._traces_client = None self._pending_otel_steps = set() @@ -313,31 +300,6 @@ def __init__( self._auto_enable_agent_control_if_available() return - if trace_id or span_id: - if self.mode != "distributed": - raise SplunkAOLoggerException("trace_id or span_id can only be used in distributed mode") - if span_id and not trace_id: - raise SplunkAOLoggerException( - "trace_id is required when span_id is provided. " - "In distributed tracing, both trace_id and span_id must be propagated together." - ) - - # Validate UUIDs to prevent crashes from malformed input - if trace_id: - try: - uuid.UUID(trace_id) - except (ValueError, AttributeError, TypeError) as e: - raise SplunkAOLoggerException(f"Invalid trace_id: '{trace_id}' is not a valid UUID. Error: {e}") - - if span_id: - try: - uuid.UUID(span_id) - except (ValueError, AttributeError, TypeError) as e: - raise SplunkAOLoggerException(f"Invalid span_id: '{span_id}' is not a valid UUID. Error: {e}") - - self.trace_id = trace_id - self.span_id = span_id - if (agent_stream or agent_stream_id) and experiment_id: raise SplunkAOLoggerException("User cannot specify both an agent stream and an experiment.") @@ -368,7 +330,9 @@ def __init__( "User must provide project_name or project_id to SplunkAOLogger, or set it as an environment variable." ) if self.experiment_id is None and self.agent_stream_name is None and self.agent_stream_id is None: - raise SplunkAOLoggerException("agent_stream or agent_stream_id is required to initialize SplunkAOLogger.") + raise SplunkAOLoggerException( + "agent_stream or agent_stream_id is required to initialize SplunkAOLogger." + ) if local_metrics: self.local_metrics = local_metrics @@ -393,12 +357,6 @@ def __init__( else: self._sink = build_span_sink(build_standalone_exporter(StandaloneConfig.from_env(), routing)) - # If continuing an existing distributed trace, create local stubs instead of - # fetching from the backend to avoid race conditions with eventual consistency. - # Note: trace_id/span_id can ONLY be provided in distributed mode for distributed tracing - if self.trace_id: - self._init_distributed_trace_stubs() - # cleans up when the python interpreter closes atexit.register(self.terminate) self._auto_enable_agent_control_if_available() @@ -424,13 +382,13 @@ def reset_parent_tracking(self) -> None: if root is not None: self._discard_otel_subtree(root) - def _assign_otel_context(self, otel_trace_id: int, trace_state: TraceState) -> SpanContext: - """Create a sampled local SpanContext without changing the active context.""" + def _assign_otel_context(self, otel_trace_id: int, trace_flags: TraceFlags, trace_state: TraceState) -> SpanContext: + """Create a local SpanContext without changing the active context.""" return SpanContext( trace_id=otel_trace_id, span_id=_otel_id_generator.generate_span_id(), is_remote=False, - trace_flags=TraceFlags(TraceFlags.SAMPLED), + trace_flags=trace_flags, trace_state=trace_state, ) @@ -450,14 +408,18 @@ def _record_otel_ids( if parent_span_context is None: otel_trace_id = _otel_id_generator.generate_trace_id() + trace_flags = TraceFlags(TraceFlags.SAMPLED) trace_state = TraceState() else: otel_trace_id = parent_span_context.trace_id + trace_flags = parent_span_context.trace_flags trace_state = parent_span_context.trace_state ids = OtelIds( - span_context=self._assign_otel_context(otel_trace_id, trace_state), + span_context=self._assign_otel_context(otel_trace_id, trace_flags, trace_state), parent_span_context=parent_span_context, + exportable=not isinstance(step, Trace), + session_id=get_effective_session_id(self.session_id), ) self._otel_ids[step.id] = ids return ids @@ -516,10 +478,14 @@ def _sync_otel_context_impl(self, current_parent: StepWithChildSpans | None) -> base_context = state.base_context if state is not None else otel_context.get_current() remaining_contexts = tuple(active for active in active_contexts if active.logger_id != logger_id) desired_contexts = tuple( - (logger_id, step_id, self._otel_ids[step_id].span_context) for step_id in desired_step_ids + (logger_id, step_id, self._otel_ids[step_id].span_context, self._otel_ids[step_id].exportable) + for step_id in desired_step_ids ) contexts_to_restore = ( - tuple((active.logger_id, active.step_id, active.span_context) for active in remaining_contexts) + tuple( + (active.logger_id, active.step_id, active.span_context, active.exportable) + for active in remaining_contexts + ) + desired_contexts ) @@ -536,11 +502,15 @@ def _sync_otel_context_impl(self, current_parent: StepWithChildSpans | None) -> otel_context.attach(base_context) rebuilt_contexts: list[ActiveOtelContext] = [] - for owner_id, step_id, span_context in contexts_to_restore: + for owner_id, step_id, span_context, exportable in contexts_to_restore: ctx = otel_trace.set_span_in_context(NonRecordingSpan(span_context)) rebuilt_contexts.append( ActiveOtelContext( - logger_id=owner_id, step_id=step_id, span_context=span_context, token=otel_context.attach(ctx) + logger_id=owner_id, + step_id=step_id, + span_context=span_context, + exportable=exportable, + token=otel_context.attach(ctx), ) ) @@ -605,12 +575,15 @@ def _emit_and_release(self, finished_step: BaseStep) -> None: if isinstance(finished_step, Trace): return + if not bool(ids.span_context.trace_flags & TraceFlags.SAMPLED): + return + self._sink.emit( self._converter.convert_span( span=finished_step, span_context=ids.span_context, parent_span_context=self._export_parent_context(finished_step, ids), - session_id=self.session_id, + session_id=ids.session_id, resource=self._resource, ) ) @@ -657,6 +630,66 @@ def _mark_potentially_parentable(self, span: BaseStep) -> None: if not self._ingestion_hook and span.id in self._otel_ids: self._pending_otel_steps.add(span.id) + def _register_handler_step(self, step: BaseStep, parent: StepWithChildSpans) -> None: + """Register a callback-owned step under an explicit proprietary parent.""" + step._parent = parent + step.dataset_input = parent.dataset_input + step.dataset_output = parent.dataset_output + step.dataset_metadata = parent.dataset_metadata + parent.add_child_span(step) + if isinstance(step, LoggedWorkflowSpan | LoggedAgentSpan) and isinstance(parent, LoggedTrace): + step.conversation_root = True + self._record_otel_ids(step, parent_step=parent) + + def _activate_handler_step(self, step: BaseStep) -> HandlerStepContext | None: + """Make a callback-owned stable identity current for transport propagation.""" + ids = self._otel_ids.get(step.id) + if ids is None: + return None + prior_context = otel_context.get_current() + active_context = otel_trace.set_span_in_context(NonRecordingSpan(ids.span_context), prior_context) + return HandlerStepContext( + prior_context=prior_context, token=otel_context.attach(active_context), span_context=ids.span_context + ) + + def _restore_handler_step_context(self, activation: HandlerStepContext | None) -> None: + """Restore context owned by one callback without disturbing caller state.""" + if activation is None: + return + try: + otel_context.detach(activation.token) + except (RuntimeError, ValueError): + current = otel_trace.get_current_span().get_span_context() + if current == activation.span_context: + otel_context.attach(activation.prior_context) + self._logger.warning("Failed to detach handler OTel context; restored the prior context when safe.") + + def _replace_handler_step(self, provisional: BaseStep, final: BaseStep) -> BaseStep: + """Replace a provisional callback model while preserving identity and topology.""" + parent = provisional._parent + final._parent = parent + final.dataset_input = provisional.dataset_input + final.dataset_output = provisional.dataset_output + final.dataset_metadata = provisional.dataset_metadata + if isinstance(provisional, LoggedWorkflowSpan | LoggedAgentSpan) and isinstance( + final, LoggedWorkflowSpan | LoggedAgentSpan + ): + final.conversation_root = provisional.conversation_root + if isinstance(provisional, StepWithChildSpans) and isinstance(final, StepWithChildSpans): + final.spans = provisional.spans + final._last_child_created_at = provisional._last_child_created_at + for child in final.spans: + child._parent = final + if parent is not None: + parent.spans = [final if child is provisional else child for child in parent.spans] + if self.current_parent() is provisional and isinstance(final, StepWithChildSpans): + self._set_current_parent(final) + return final + + def _complete_handler_step(self, step: BaseStep) -> None: + """Enqueue one completed handler operation through the shared completion seam.""" + self._complete_step(step) + def _current_span_id(self) -> uuid.UUID: """Return the current proprietary parent ID for internal lifecycle tests.""" current_parent = self.current_parent() @@ -746,49 +779,6 @@ def _ensure_session_crud_client(self) -> Traces: self._traces_client = self._create_traces_client() return self._traces_client - def _init_distributed_trace_stubs(self) -> None: - """ - Initialize local stub objects for distributed tracing. To only be used in distributed mode. - - When a downstream service receives trace_id/span_id via headers, we create - local stub objects instead of fetching from the backend. This avoids race - conditions as the parent trace/span may not have been ingested yet when the downstream service starts. - - The stubs are placeholders that allow: - 1. Adding new spans to the distributed trace - 2. Proper parent-child relationships via _parent pointers - 3. Correct trace_id in ingestion requests - - Note: trace_id and span_id are already validated as UUIDs in __init__ - """ - stub_trace = LoggedTrace( - input="", - name=STUB_TRACE_NAME, - created_at=datetime.now(), - id=uuid.UUID(self.trace_id), - metrics=Metrics(duration_ns=0), - ) - self.traces.append(stub_trace) - - # Set trace as current parent using parent pointers - stub_trace._parent = None # Root trace has no parent - self._record_otel_ids(stub_trace) - self._set_current_parent(stub_trace) - - if self.span_id: - # If span_id is provided, also add the span (it's the immediate parent) - stub_span = LoggedWorkflowSpan( - input="", - name="stub_parent_span", - created_at=datetime.now(), - id=uuid.UUID(self.span_id), - metrics=Metrics(duration_ns=0), - ) - # Set parent pointer and update current parent - stub_span._parent = stub_trace - self._record_otel_ids(stub_span, parent_step=stub_trace) - self._set_current_parent(stub_span) - def add_trace( self, input: str, @@ -968,228 +958,6 @@ def _get_last_output(node: BaseStep | None) -> tuple[IngestOutputType | None, In return None, None - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _ingest_trace_streaming(self, trace: Trace, is_complete: bool = False) -> None: - traces_ingest_request = TracesIngestRequest( - traces=[copy.deepcopy(trace)], session_id=self.session_id, is_complete=is_complete, reliable=True - ) - - task_id = f"trace-ingest-{trace.id}" - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info(f"Retry #{self._task_handler.get_retry_count(task_id)} for task {task_id}"), - ), - on_giveup=lambda details: self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", exc_info=False - ), - ) - @retry_on_transient_http_error - async def ingest_traces_with_backoff(request: Any) -> None: - return await self._traces_client.ingest_traces(request) - - self._task_handler.submit_task( - task_id, lambda: ingest_traces_with_backoff(traces_ingest_request), dependent_on_prev=False - ) - self._logger.info("ingested trace %s.", trace.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _ingest_span_streaming(self, span: Span) -> None: - parent_step: StepWithChildSpans | None = ( - self.current_parent() - if span.type - not in [ - StepType.trace, - StepType.workflow, - StepType.agent, - ] # TODO: change this to StepWithChildSpans once we fix tool and retriever spans in `core - else self.previous_parent() - ) - if parent_step is None: - raise ValueError("A trace needs to be created in order to add a span.") - - # Use IDs from the current trace and parent step - trace_id = self.traces[0].id - parent_id = parent_step.id - spans_ingest_request = SpansIngestRequest( - spans=[copy.deepcopy(span)], trace_id=trace_id, parent_id=parent_id, reliable=True - ) - - task_id = f"span-ingest-{span.id}" - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info(f"Retry #{self._task_handler.get_retry_count(task_id)} for task {task_id}"), - ), - on_giveup=lambda details: self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", exc_info=False - ), - ) - @retry_on_transient_http_error - async def ingest_spans_with_backoff(request: Any) -> None: - return await self._traces_client.ingest_spans(request) - - self._task_handler.submit_task( - task_id, lambda: ingest_spans_with_backoff(spans_ingest_request), dependent_on_prev=False - ) - self._logger.info("ingested span %s.", span.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _update_trace_streaming(self, trace: Trace, is_complete: bool = False) -> None: - output: str | None = None - if trace.output is not None: - output = trace.output if isinstance(trace.output, str) else serialize_to_str(trace.output) - trace_update_request = TraceUpdateRequest( - trace_id=trace.id, - log_stream_id=self.agent_stream_id, - experiment_id=self.experiment_id, - output=output, - status_code=trace.status_code, - tags=trace.tags, - is_complete=is_complete, - duration_ns=trace.metrics.duration_ns, - reliable=True, - ) - - # Use counter to make each update task unique (same trace can be updated multiple times) - self._task_counter += 1 - task_id = f"trace-update-{trace.id}-{self._task_counter}" - - # Find the most recent trace update task for this specific trace (if any) - # This ensures trace updates for the same trace happen in order - prev_trace_update_task = None - for existing_task_id in reversed(list(self._task_handler._tasks.keys())): - if existing_task_id.startswith(f"trace-update-{trace.id}-"): - prev_trace_update_task = existing_task_id - break - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info(f"Retry #{self._task_handler.get_retry_count(task_id)} for trace update {task_id}"), - ), - on_giveup=lambda details: ( - self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", - exc_info=False, - ), - ), - ) - @retry_on_transient_http_error - async def update_trace_with_backoff(request: Any) -> None: - return await self._traces_client.update_trace(request) - - # Submit with dependency on the previous trace update for this trace - if prev_trace_update_task: - self._task_handler.submit_task_with_parent( - task_id, lambda: update_trace_with_backoff(trace_update_request), parent_task_id=prev_trace_update_task - ) - else: - self._task_handler.submit_task( - task_id, lambda: update_trace_with_backoff(trace_update_request), dependent_on_prev=True - ) - - # Mark that we've submitted the trace completion update to prevent duplicates - if is_complete: - self._trace_completion_submitted = True - - self._logger.info("updated trace %s.", trace.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _update_span_streaming(self, span: Span) -> None: - span_update_request = SpanUpdateRequest( - span_id=span.id, - log_stream_id=self.agent_stream_id, - experiment_id=self.experiment_id, - output=span.output, - status_code=span.status_code, - tags=span.tags, - duration_ns=span.metrics.duration_ns, - reliable=True, - ) - - # Use counter to make each update task unique (same span can be updated multiple times) - self._task_counter += 1 - task_id = f"span-update-{span.id}-{self._task_counter}" - - # Find the most recent update/ingest task for this specific span - # This ensures span updates happen in order - parent_task_id = None - for existing_task_id in reversed(list(self._task_handler._tasks.keys())): - if existing_task_id.startswith(f"span-update-{span.id}-") or existing_task_id == f"span-ingest-{span.id}": - parent_task_id = existing_task_id - break - - # If no previous task found, depend on the span ingest - if not parent_task_id: - parent_task_id = f"span-ingest-{span.id}" - - @backoff.on_exception( - backoff.expo, - Exception, - max_tries=self._max_retries, - max_time=self._max_time, - base=2, - logger=None, - on_backoff=lambda details: ( - self._task_handler.increment_retry(task_id), - self._logger.info( - f"Retry #{self._task_handler.get_retry_count(task_id)} for task {task_id}, waiting {details['wait']:.1f}s" - ), - ), - on_giveup=lambda details: self._logger.error( - f"Task {task_id} failed after {details['tries']} attempts: {details.get('exception')}", exc_info=False - ), - ) - @retry_on_transient_http_error - async def update_span_with_backoff(request: Any) -> None: - return await self._traces_client.update_span(request) - - self._task_handler.submit_task_with_parent( - task_id, lambda: update_span_with_backoff(span_update_request), parent_task_id=parent_task_id - ) - self._logger.info("updated span %s.", span.id) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _ingest_step_streaming(self, step: StepWithChildSpans, is_complete: bool = False) -> None: - if isinstance(step, LoggedTrace): - self._ingest_trace_streaming(step, is_complete=is_complete) - else: - self._ingest_span_streaming(step) - - @nop_sync - @warn_catch_exception(exceptions=(Exception,)) - def _update_step_streaming(self, step: StepWithChildSpans, is_complete: bool = False) -> None: - if isinstance(step, LoggedTrace): - self._update_trace_streaming(step, is_complete=is_complete) - else: - self._update_span_streaming(step) - @nop_sync @warn_catch_exception(exceptions=(Exception,)) def previous_parent(self) -> StepWithChildSpans | None: @@ -1198,8 +966,6 @@ def previous_parent(self) -> StepWithChildSpans | None: @nop_sync @warn_catch_exception(exceptions=(Exception,)) def has_active_trace(self) -> bool: - if self.mode == "distributed" and (self.trace_id or self.span_id): - return True current_parent = self.current_parent() # Each logger has its own per-instance ContextVar for parent tracking. # The traces check is a sanity check to ensure consistency. @@ -1222,71 +988,6 @@ def disable_agent_control(self) -> None: if bridge is not None: bridge.unregister() - def get_tracing_headers(self) -> dict[str, str]: - """ - Get tracing headers for distributed tracing. - Returns headers that can be passed to downstream services to continue the distributed trace. - - Returns - ------- - dict[str, str] - Dictionary with the following headers: - - Splunk-AO-Trace-ID: The root trace ID - - Splunk-AO-Parent-ID: The ID of the current parent (trace or span) that downstream - spans should attach to - - Raises - ------ - SplunkAOLoggerException - If not in distributed mode or if no trace has been started. - - Examples - -------- - ```python - logger = SplunkAOLogger(mode="distributed") - logger.start_trace(input="question") - headers = logger.get_tracing_headers() - # headers = { - # "Splunk-AO-Trace-ID": "...", - # "Splunk-AO-Parent-ID": "...", # trace ID as parent - # } - - logger.add_workflow_span(input="workflow", name="orchestrator") - headers = logger.get_tracing_headers() - # headers = { - # "Splunk-AO-Trace-ID": "...", - # "Splunk-AO-Parent-ID": "...", # workflow span ID as parent - # } - - # Pass headers to HTTP request - response = httpx.post(url, headers=headers) - ``` - - Note: Project and log_stream are configured per service (via env vars or logger initialization), - not propagated via headers, following standard distributed tracing patterns. - """ - if self.mode != "distributed": - raise SplunkAOLoggerException( - "get_tracing_headers is only supported in distributed mode for distributed tracing." - ) - - if len(self.traces) == 0: - raise SplunkAOLoggerException("Start trace before getting tracing headers.") - - headers: dict[str, str] = {} - - root_trace = self.traces[-1] - headers[TRACE_ID_HEADER] = str(root_trace.id) - - current_parent = self.current_parent() - - if not current_parent: - raise SplunkAOLoggerException("No parent trace or span found.") - - headers[PARENT_ID_HEADER] = str(current_parent.id) - - return headers - @nop_sync @warn_catch_exception() def start_trace( @@ -2263,72 +1964,6 @@ async def async_flush(self) -> None: if not await asyncio.to_thread(self._sink.force_flush): raise RuntimeError("force_flush timed out; some spans may not have been exported") - @async_warn_catch_exception(exceptions=(Exception,)) - async def _wait_for_all_tasks_async(self, timeout_seconds: int) -> None: - """Wait for all background tasks to complete (async polling). - - Parameters - ---------- - timeout_seconds: int - Maximum time to wait for tasks to complete - """ - start_wait = time.time() - while not self._task_handler.all_tasks_completed(): - if time.time() - start_wait > timeout_seconds: - raise TimeoutError( - f"Flush timeout reached after {timeout_seconds}s. " - "Some trace/span update requests may still be in progress." - ) - await asyncio.sleep(0.1) - - @warn_catch_exception(exceptions=(Exception,)) - def _wait_for_all_tasks_sync(self, timeout_seconds: int) -> None: - """Wait for all background tasks to complete (synchronous polling). - - Parameters - ---------- - timeout_seconds: int - Maximum time to wait for tasks to complete - """ - start_wait = time.time() - while not self._task_handler.all_tasks_completed(): - if time.time() - start_wait > timeout_seconds: - self._logger.warning( - f"Terminate timeout reached after {timeout_seconds}s. " - "Some trace/span update requests may still be in progress." - ) - break - time.sleep(0.1) - - @warn_catch_exception(exceptions=(Exception,)) - def _wait_for_pending_span_ingests(self, timeout_seconds: int) -> None: - """Wait for all pending span ingest tasks to complete. - - Note: Uses time.sleep() for polling even though callers may have @nop_sync. - This briefly blocks but is acceptable since we're just polling task status. - - Parameters - ---------- - timeout_seconds: int - Maximum time to wait for span ingests to complete - """ - pending_span_tasks = [ - task_id - for task_id in self._task_handler._tasks - if task_id.startswith("span-ingest-") and self._task_handler.get_status(task_id) in ["pending", "running"] - ] - - if pending_span_tasks: - start_wait = time.time() - while pending_span_tasks and (time.time() - start_wait) < timeout_seconds: - pending_span_tasks = [ - task_id - for task_id in pending_span_tasks - if self._task_handler.get_status(task_id) in ["pending", "running"] - ] - if pending_span_tasks: - time.sleep(0.1) - @nop_sync @warn_catch_exception(exceptions=(Exception,)) def _auto_conclude_trace(self) -> None: @@ -2343,39 +1978,13 @@ def _auto_conclude_trace(self) -> None: # Use the last trace in self.traces (should be the only active trace) trace = self.traces[-1] - # Don't auto-conclude stub traces - they're owned by the upstream service - # Downstream services that receive distributed tracing headers create stubs - # but should not mark them as complete - if trace.name == STUB_TRACE_NAME: - return - # If there are unconcluded items in the stack, conclude them if self._parent_stack: self._logger.info("Concluding unconcluded spans before flush...") # Get output from last child span if trace has no explicit output output, redacted_output = SplunkAOLogger._get_last_output(trace) # conclude() with conclude_all=True will conclude all unconcluded items in _parent_stack - # This will mark the trace as complete in distributed mode self.conclude(output=output, redacted_output=redacted_output, conclude_all=True) - elif self.mode == "distributed": - if not self._trace_completion_submitted: - # Wait for all span ingests to complete before marking trace complete - self._wait_for_pending_span_ingests(timeout_seconds=DISTRIBUTED_FLUSH_TIMEOUT_SECONDS) - self._update_trace_streaming(trace, is_complete=True) - - async def _flush_distributed(self) -> list[LoggedTrace]: - """Legacy proprietary distributed flush path, unused by OTLP egress.""" - self._auto_conclude_trace() - - # Wait for all pending trace/span update requests to complete - self._logger.info("Waiting for all distributed tracing tasks to complete...") - await self._wait_for_all_tasks_async(timeout_seconds=DISTRIBUTED_FLUSH_TIMEOUT_SECONDS) - self._logger.info("All distributed tracing requests are complete.") - - self.traces = [] - self._set_current_parent(None) - - return [] async def _flush_batch(self) -> list[LoggedTrace]: """Flush in batch mode: conclude unconcluded traces and send all traces to backend.""" @@ -2490,7 +2099,7 @@ async def _start_or_get_session_async( ) -> str: self._session_external_id = external_id if self._ingestion_hook: - self.session_id = str(uuid.uuid4()) + self._set_active_session_id(str(uuid.uuid4())) self._logger.info("Session started: session_id=%s, external_id=%s", self.session_id, external_id) return self.session_id @@ -2515,7 +2124,7 @@ async def _start_or_get_session_async( if sessions and len(sessions["records"]) > 0: session_id = sessions["records"][0]["id"] self._logger.info(f"Session {session_id} with external ID {external_id} already exists; using it.") - self.session_id = session_id + self._set_active_session_id(session_id) return session_id except Exception: self._logger.error("Failed to search for session with external ID %s", external_id, exc_info=True) @@ -2529,9 +2138,14 @@ async def _start_or_get_session_async( ) self._logger.info("Session started with ID: %s", session["id"]) - self.session_id = str(session["id"]) + self._set_active_session_id(str(session["id"])) return self.session_id + def _set_active_session_id(self, session_id: str | None) -> None: + """Update compatibility and request-local session state together.""" + self.session_id = session_id + set_session_context(session_id) + @nop_async async def async_start_session( self, @@ -2603,11 +2217,15 @@ def start_session( str The ID of the session (existing or newly created). """ - return async_run( + session_id = async_run( self._start_or_get_session_async( name=name, previous_session_id=previous_session_id, external_id=external_id, metadata=metadata ) ) + # ``async_run`` may execute in another context; publish the resolved ID + # into the synchronous caller's request-local context as well. + self._set_active_session_id(session_id) + return session_id @nop_sync @warn_catch_exception(exceptions=(Exception,)) @@ -2625,14 +2243,14 @@ def set_session(self, session_id: str) -> None: None """ self._logger.info("Setting the current session to %s", session_id) - self.session_id = session_id + self._set_active_session_id(session_id) self._logger.info("Current session set to %s", session_id) @nop_sync @warn_catch_exception(exceptions=(Exception,)) def clear_session(self) -> None: self._logger.info("Clearing the current session from the logger...") - self.session_id = None + self._set_active_session_id(None) self._logger.info("Current session cleared.") @nop_async diff --git a/src/splunk_ao/logger/task_handler.py b/src/splunk_ao/logger/task_handler.py deleted file mode 100644 index 3864c158..00000000 --- a/src/splunk_ao/logger/task_handler.py +++ /dev/null @@ -1,200 +0,0 @@ -import time -from collections.abc import Awaitable, Callable, Coroutine -from concurrent.futures import Future -from typing import Any, Literal - -from galileo_core.helpers.event_loop_thread_pool import EventLoopThreadPool - -NUM_THREADS = 4 - -TaskStatus = Literal["not_found", "pending", "running", "completed", "failed"] - - -class ThreadPoolTaskHandler: - """A task handler that manages dependencies and executes tasks in a thread pool.""" - - _pool: EventLoopThreadPool - _tasks: dict[str, dict] - _retry_counts: dict[str, int] - - def __init__(self, num_threads: int = NUM_THREADS): - self._tasks = {} - self._retry_counts = {} - self._pool = EventLoopThreadPool(num_threads=num_threads) - - def _handle_task_completion(self, task_id: str) -> None: - """Handle the completion of a task, triggering any children.""" - # Find all child tasks that depend on this task - for _child_task_id, task in list(self._tasks.items()): - if task.get("parent_task_id") == task_id and task.get("callback"): - # Execute the callback which will submit the child task - task["callback"]() - - def _add_or_update_task( - self, - task_id: str, - future: Future | None = None, - start_time: float | None = None, - parent_task_id: str | None = None, - callback: Callable | None = None, - ) -> None: - """ - Track a submitted future. - - Parameters - ---------- - task_id: str - The ID of the task. - future: Optional[Future] - The future to track. - start_time: Optional[float] - The start time of the task. - parent_task_id: Optional[str] - The ID of the parent task. - callback: Optional[Callable] - The callback to run when the task is completed. - """ - self._tasks[task_id] = { - "future": future, - "start_time": start_time, - "parent_task_id": parent_task_id, - "callback": callback, - } - self._retry_counts[task_id] = 0 - - def submit_task( - self, task_id: str, async_fn: Callable[[], Awaitable[Any]] | Coroutine, dependent_on_prev: bool = False - ) -> None: - """ - Submit a task to the thread pool. - - Parameters - ---------- - task_id: str - The ID of the task. - async_fn: Union[Callable[[], Awaitable[Any]], Coroutine] - The async function to submit to the thread pool. - dependent_on_prev: bool - Whether the task depends on the previous task. - """ - - def _submit(*args) -> None: - future = self._pool.submit(async_fn, wait_for_result=False) - future.add_done_callback(lambda f: self._handle_task_completion(task_id)) - self._add_or_update_task(task_id=task_id, future=future, start_time=time.time(), parent_task_id=None) - - if dependent_on_prev: - if not self._tasks: - _submit() - else: - last_task_id = list(self._tasks.keys())[-1] - if self.get_status(last_task_id) == "completed": - _submit() - else: - self._add_or_update_task( - task_id=task_id, future=None, start_time=None, parent_task_id=last_task_id, callback=_submit - ) - else: - _submit() - - def submit_task_with_parent( - self, task_id: str, async_fn: Callable[[], Awaitable[Any]] | Coroutine, parent_task_id: str - ) -> None: - """ - Submit a task that depends on a specific parent task. - - Parameters - ---------- - task_id: str - The ID of the task. - async_fn: Union[Callable[[], Awaitable[Any]], Coroutine] - The async function to submit to the thread pool. - parent_task_id: str - The ID of the parent task this depends on. - """ - - def _submit(*args) -> None: - future = self._pool.submit(async_fn, wait_for_result=False) - future.add_done_callback(lambda f: self._handle_task_completion(task_id)) - self._add_or_update_task(task_id=task_id, future=future, start_time=time.time(), parent_task_id=None) - - if parent_task_id not in self._tasks or self.get_status(parent_task_id) == "completed": - _submit() - else: - self._add_or_update_task( - task_id=task_id, future=None, start_time=None, parent_task_id=parent_task_id, callback=_submit - ) - - def get_children(self, parent_task_id: str) -> list[dict]: - """Get the children of a task.""" - return [task for _, task in self._tasks.items() if task.get("parent_task_id") == parent_task_id] - - def increment_retry(self, task_id: str) -> None: - """ - Increment the retry count for a task. - - Parameters - ---------- - task_id: str - The ID of the task. - """ - self._retry_counts[task_id] = self._retry_counts.get(task_id, 0) + 1 - - def get_status(self, task_id: str) -> TaskStatus: - """ - Returns the status of a task. - - Parameters - ---------- - task_id: str - The ID of the task. - - Returns - ------- - TaskStatus - The status of the task. - """ - if task_id not in self._tasks: - return "not_found" - - task = self._tasks[task_id] - - if task.get("parent_task_id"): - return "pending" - - future = task["future"] - if not future: - return "not_found" - - if not future.done(): - return "running" - - try: - future.result() # This will raise exception if task failed - return "completed" - except Exception: - return "failed" - - def get_result(self, task_id: str) -> Any: - """Get result if task completed, otherwise raises exception.""" - if task_id not in self._tasks: - raise ValueError(f"Task {task_id} not found") - return self._tasks[task_id]["future"].result() - - def get_retry_count(self, task_id: str) -> int: - """Get the retry count for a task.""" - return self._retry_counts.get(task_id, 0) - - def all_tasks_completed(self) -> bool: - """ - Check if all tasks are completed. - - Returns - ------- - bool - True if all tasks are completed, False otherwise. - """ - return all(self.get_status(task_id) not in ["running", "pending"] for task_id in self._tasks) - - def terminate(self) -> None: - self._pool.stop() diff --git a/src/splunk_ao/middleware/tracing.py b/src/splunk_ao/middleware/tracing.py index 57b2725f..ceb677db 100644 --- a/src/splunk_ao/middleware/tracing.py +++ b/src/splunk_ao/middleware/tracing.py @@ -1,62 +1,37 @@ -""" -Distributed tracing middleware for Starlette-based applications. - -This middleware automatically extracts distributed tracing headers from incoming HTTP requests -and makes them available to the Splunk AO logger within request handlers. - -Works with any ASGI framework built on Starlette: -- FastAPI -- Starlette -- Any other Starlette-based framework +"""W3C trace-context middleware for Starlette-based applications. -Example usage with FastAPI: - ```python - from fastapi import FastAPI - from splunk_ao.middleware import TracingMiddleware, get_request_logger +``TracingMiddleware`` extracts ``traceparent`` and ``tracestate`` through the +configured OpenTelemetry propagator. Applications explicitly own the local +Splunk AO logger and trace lifecycle:: - app = FastAPI() app.add_middleware(TracingMiddleware) @app.post("/process") async def process_request(data: dict): - # Logger automatically continues the distributed trace logger = get_request_logger() - logger.add_workflow_span(input=str(data), name="process_workflow") - # ... process request ... - logger.conclude(output="done") - return {"status": "success"} - ``` - -Example usage with Starlette: - ```python - from starlette.applications import Starlette - from starlette.routing import Route - from splunk_ao.middleware import TracingMiddleware, get_request_logger - - async def homepage(request): - logger = get_request_logger() - logger.add_workflow_span(input="homepage", name="homepage_handler") - logger.conclude(output="success") - return {"status": "ok"} - - app = Starlette( - routes=[Route("/", homepage)], - middleware=[TracingMiddleware] - ) - ``` + try: + logger.start_trace(input=str(data), name="request") + logger.add_workflow_span(input=str(data), name="process") + result = await process(data) + logger.conclude(output=str(result)) + logger.conclude(output=str(result)) + return {"result": result} + finally: + logger.terminate() """ import logging from typing import Any, NoReturn -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER -from splunk_ao.decorator import _parent_id_context, _trace_id_context +from opentelemetry import context as otel_context + from splunk_ao.logger import SplunkAOLogger +from splunk_ao.tracing import extract_tracing_context _logger = logging.getLogger(__name__) INSTALL_ERR_MSG = ( - "Starlette is not installed. Install optional middleware dependencies with: pip install galileo[middleware]" + "Starlette is not installed. Install optional middleware dependencies with: pip install splunk-ao[middleware]" ) try: @@ -65,7 +40,7 @@ async def homepage(request): from starlette.responses import Response from starlette.types import ASGIApp except ImportError: - # Create stub classes if Starlette is not available + # Keep imports available without the optional Starlette dependency. class BaseHTTPMiddleware: # type: ignore[no-redef] def __init__(self, *args: Any, **kwargs: Any) -> NoReturn: raise ImportError(INSTALL_ERR_MSG) @@ -84,144 +59,27 @@ class ASGIApp: # type: ignore[no-redef] class TracingMiddleware(BaseHTTPMiddleware): - """ - Middleware that extracts distributed tracing headers from incoming requests. - - This middleware looks for the following headers in incoming HTTP requests: - - Splunk-AO-Trace-ID: The root trace ID - - Splunk-AO-Parent-ID: The parent span/trace ID to attach to - - These values are stored in context variables, making them available to request - handlers via the `get_request_logger()` function. - - The middleware is compatible with FastAPI and any Starlette-based framework. - - Note: Project and log_stream are configured per service via environment variables - (SPLUNK_AO_PROJECT and SPLUNK_AO_AGENT_STREAM). They are not propagated via headers, - following standard distributed tracing patterns. - """ + """Attach incoming W3C trace context for the duration of each request.""" def __init__(self, app: ASGIApp) -> None: - """ - Initialize the tracing middleware. - - Parameters - ---------- - app : ASGIApp - The ASGI application - """ super().__init__(app) async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response: - """ - Process the request and extract tracing headers. - - Parameters - ---------- - request : Request - The incoming HTTP request - call_next : RequestResponseEndpoint - The next middleware or route handler - - Returns - ------- - Response - The HTTP response - """ - # Extract tracing headers from request - trace_id = request.headers.get(TRACE_ID_HEADER) - parent_id = request.headers.get(PARENT_ID_HEADER) - - # Store in context variables (thread-safe for async) - trace_id_token = _trace_id_context.set(trace_id) - parent_id_token = _parent_id_context.set(parent_id) - + """Extract, attach, and reliably detach the request's trace context.""" + extracted = extract_tracing_context(request.headers) + token = otel_context.attach(extracted) try: - # Process the request return await call_next(request) finally: - # Clean up context variables - _trace_id_context.reset(trace_id_token) - _parent_id_context.reset(parent_id_token) + otel_context.detach(token) def get_request_logger() -> SplunkAOLogger: - """ - Get a request-scoped SplunkAOLogger configured for distributed mode. - - Distributed mode enables the legacy trace_id/span_id continuation parameters. - Completed spans use the same scheduled OTLP batch export as batch mode. - - This function should be called within a request handler after the TracingMiddleware has - been registered. It creates a new SplunkAOLogger instance per request that automatically - continues the distributed trace from the upstream service. - - The logger is configured using trace context extracted by the middleware: - - Splunk-AO-Trace-ID: Root trace ID - - Splunk-AO-Parent-ID: Parent span/trace ID to attach to - - Project and log_stream are configured per service via environment variables - (SPLUNK_AO_PROJECT and SPLUNK_AO_AGENT_STREAM), not propagated via headers, following - standard distributed tracing patterns. - - If no tracing headers were present in the request, a regular logger is returned - (using SPLUNK_AO_PROJECT and SPLUNK_AO_AGENT_STREAM env vars). - - Note: This creates a new logger per request, unlike the decorator's get_logger_instance() - which uses a singleton pattern. - - Returns - ------- - SplunkAOLogger - A logger instance configured for the current request's trace context - - Examples - -------- - ```python - @app.post("/process") - async def process_request(data: dict): - logger = get_request_logger() - - # This span will be attached to the distributed trace - logger.add_workflow_span(input=str(data), name="process_workflow") - result = await process(data) - logger.conclude(output=str(result)) - - logger.flush() - logger.terminate() - - return {"result": result} - ``` - - ```python - @app.post("/retrieve") - async def retrieve_endpoint(query: str): - # Get logger with trace context from upstream service - logger = get_request_logger() - - # If trace context exists, this creates a workflow span - # Otherwise, it starts a new trace - if logger.trace_id: - logger.add_workflow_span(input=query, name="retrieval_service") - else: - logger.start_trace(input=query, name="retrieval_service") - - results = retrieve(query, logger) - - logger.conclude(output=str(results)) - logger.flush() - logger.terminate() + """Create a request-scoped logger under the middleware's active context. - return {"results": results} - ``` + The caller must explicitly call ``start_trace()`` before adding spans and + should call ``terminate()`` in ``finally``. With an extracted W3C parent, + the first real local operation becomes its descendant; without one, it + starts a new trace. """ - # Get trace context from middleware - trace_id = _trace_id_context.get() - parent_id = _parent_id_context.get() - - # Create logger with trace context - # Project and log_stream come from env vars (SPLUNK_AO_PROJECT, SPLUNK_AO_AGENT_STREAM) - # If parent_id equals trace_id, it means the parent is the root trace itself, - # not a span. In this case, we should pass None as span_id to avoid - # SplunkAOLoggerException when it tries to look up a span with the trace_id. - return SplunkAOLogger(mode="distributed", trace_id=trace_id, span_id=parent_id if parent_id != trace_id else None) + return SplunkAOLogger() diff --git a/src/splunk_ao/otel.py b/src/splunk_ao/otel.py index 28bfba5b..d1155c14 100644 --- a/src/splunk_ao/otel.py +++ b/src/splunk_ao/otel.py @@ -18,13 +18,12 @@ from splunk_ao.config import SplunkAOConfig from splunk_ao.converter import build_span_attributes from splunk_ao.decorator import ( + _agent_stream_context, _dataset_input_context, _dataset_metadata_context, _dataset_output_context, _experiment_id_context, - _agent_stream_context, _project_context, - _session_id_context, ) from splunk_ao.deployment import DeploymentMode, O11yConfig, StandaloneConfig from splunk_ao.exporter import ( @@ -35,6 +34,7 @@ resolve_routing, ) from splunk_ao.exporter.diagnostics import get_export_health +from splunk_ao.session_context import get_effective_session_id logger = logging.getLogger(__name__) @@ -243,7 +243,7 @@ def __init__( def on_start(self, span: Span, parent_context: context.Context | None = None) -> None: """Handle span start events by delegating to the underlying processor.""" - session_id = _session_id_context.get(None) + session_id = get_effective_session_id(context=parent_context) if session_id: span.set_attribute("gen_ai.conversation.id", session_id) @@ -324,12 +324,13 @@ def start_splunk_ao_span(splunk_ao_span: SplunkAOSpan) -> Generator[trace.Span, is_conversation_root = not trace.get_current_span().get_span_context().is_valid and isinstance( splunk_ao_span, WorkflowSpan | AgentSpan ) + session_id = get_effective_session_id() with tracer.start_as_current_span(splunk_ao_span.name) as span: try: yield span finally: try: - attributes = build_span_attributes(splunk_ao_span, _session_id_context.get(None)) + attributes = build_span_attributes(splunk_ao_span, session_id) if is_conversation_root: attributes[GEN_AI_CONVERSATION_ROOT] = True for key, value in attributes.items(): diff --git a/src/splunk_ao/session_context.py b/src/splunk_ao/session_context.py new file mode 100644 index 00000000..bda49cbe --- /dev/null +++ b/src/splunk_ao/session_context.py @@ -0,0 +1,107 @@ +"""Request-local session context and W3C baggage propagation.""" + +from __future__ import annotations + +from contextvars import ContextVar +from typing import Any + +from opentelemetry import baggage, propagate +from opentelemetry import context as otel_context +from opentelemetry.baggage.propagation import W3CBaggagePropagator +from opentelemetry.context import Context +from opentelemetry.propagators import textmap + +GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id" + +_session_id_context: ContextVar[str | None] = ContextVar("session_id_context", default=None) + + +def set_session_context(session_id: str | None) -> None: + """Set the request-local explicit session selection.""" + _session_id_context.set(session_id) + + +def _baggage_session_id(context: Context | None = None) -> str | None: + """Resolve an opaque conversation ID from baggage.""" + value = baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=context) + return value.strip() if isinstance(value, str) and value.strip() else None + + +def get_effective_session_id(logger_session_id: str | None = None, context: Context | None = None) -> str | None: + """Resolve the request-local, inbound, or compatibility session ID.""" + local_session_id = _session_id_context.get(None) + if local_session_id: + return local_session_id + baggage_session_id = _baggage_session_id(context) + if baggage_session_id: + return baggage_session_id + return logger_session_id + + +def _context_with_conversation_id(context: Context, session_id: str | None) -> Context: + """Return a context with the conversation baggage normalized or removed.""" + normalized = baggage.remove_baggage(GEN_AI_CONVERSATION_ID, context=context) + if session_id: + normalized = baggage.set_baggage(GEN_AI_CONVERSATION_ID, session_id, context=normalized) + return normalized + + +class SplunkAOSessionPropagator(textmap.TextMapPropagator): + """Preserve the configured propagator while normalizing session baggage.""" + + def __init__(self, delegate: textmap.TextMapPropagator) -> None: + self._delegate = delegate + self._baggage = W3CBaggagePropagator() + + def extract( + self, + carrier: textmap.CarrierT, + context: Context | None = None, + getter: textmap.Getter[textmap.CarrierT] = textmap.default_getter, + ) -> Context: + extracted = self._delegate.extract(carrier, context=context, getter=getter) + extracted = self._baggage.extract(carrier, context=extracted, getter=getter) + return _context_with_conversation_id(extracted, _baggage_session_id(extracted)) + + def inject( + self, + carrier: textmap.CarrierT, + context: Context | None = None, + setter: textmap.Setter[textmap.CarrierT] = textmap.default_setter, + ) -> None: + active = context if context is not None else otel_context.get_current() + normalized = _context_with_conversation_id(active, get_effective_session_id(context=active)) + self._delegate.inject(carrier, context=normalized, setter=setter) + self._baggage.inject(carrier, context=normalized, setter=setter) + + @property + def fields(self) -> set[str]: + return set(self._delegate.fields) | set(self._baggage.fields) + + +def _session_propagator() -> SplunkAOSessionPropagator: + current = propagate.get_global_textmap() + if isinstance(current, SplunkAOSessionPropagator): + return current + return SplunkAOSessionPropagator(current) + + +def install_session_propagator() -> None: + """Install the session-aware adapter without replacing its delegate.""" + current = propagate.get_global_textmap() + if not isinstance(current, SplunkAOSessionPropagator): + propagate.set_global_textmap(SplunkAOSessionPropagator(current)) + + +def inject_session_context( + carrier: Any, context: Context | None = None, setter: textmap.Setter[Any] = textmap.default_setter +) -> None: + """Inject configured trace context and normalized session baggage.""" + _session_propagator().inject(carrier, context=context, setter=setter) + + +def extract_session_context( + carrier: Any, context: Context | None = None, getter: textmap.Getter[Any] = textmap.default_getter +) -> Context: + """Extract configured trace context and normalize session baggage.""" + return _session_propagator().extract(carrier, context=context, getter=getter) diff --git a/src/splunk_ao/tracing.py b/src/splunk_ao/tracing.py index a1092ec8..69352b67 100644 --- a/src/splunk_ao/tracing.py +++ b/src/splunk_ao/tracing.py @@ -1,45 +1,44 @@ -"""Utilities for distributed tracing with Splunk AO.""" +"""W3C distributed-tracing utilities for Splunk AO.""" -from splunk_ao.decorator import splunk_ao_context +from collections.abc import Mapping, MutableMapping +from opentelemetry.context import Context -def get_tracing_headers() -> dict[str, str]: - """ - Get tracing headers for distributed tracing (decorator usage). +from splunk_ao.exceptions import SplunkAOLoggerException +from splunk_ao.logger.logger import _has_active_exportable_span_context +from splunk_ao.session_context import extract_session_context, inject_session_context + + +def get_tracing_headers(carrier: MutableMapping[str, str] | None = None) -> MutableMapping[str, str]: + """Inject the active operation's W3C trace context into a carrier. - Returns headers from the singleton logger context that can be passed to downstream services. - For direct logger instances, use logger.get_tracing_headers() instead. + Parameters + ---------- + carrier : MutableMapping[str, str] | None + Existing carrier to populate. A new dictionary is created when omitted. Returns ------- - dict[str, str] - Dictionary with Splunk-AO-Trace-ID and Splunk-AO-Parent-ID headers + MutableMapping[str, str] + The supplied carrier, or a new dictionary, containing ``traceparent`` + and any other fields owned by the configured OTel propagator. Raises ------ SplunkAOLoggerException - If not in distributed mode or if no trace has been started - - Examples - -------- - Using with decorators to propagate trace context to downstream services: - - ```python - from splunk_ao import log, get_tracing_headers - import httpx - - @log() - async def orchestrator(): - # Get headers to pass to downstream service - headers = get_tracing_headers() - - # Call downstream service with trace context - async with httpx.AsyncClient() as client: - response = await client.post( - "http://service:8000/process", - headers=headers, - json={"data": "..."} - ) - ``` + If there is no active exportable operation span. The internal Splunk AO + trace envelope is not exportable and cannot be used as a wire parent. """ - return splunk_ao_context.get_logger_instance().get_tracing_headers() + if not _has_active_exportable_span_context(): + raise SplunkAOLoggerException("Distributed tracing requires an active exportable operation span") + + if carrier is None: + carrier = {} + inject_session_context(carrier) + return carrier + + +def extract_tracing_context(carrier: Mapping[str, str]) -> Context: + """Extract W3C trace context from an incoming text-map carrier.""" + normalized = {str(key).lower(): value for key, value in carrier.items()} + return extract_session_context(normalized) diff --git a/src/splunk_ao/utils/env_helpers.py b/src/splunk_ao/utils/env_helpers.py index 19132002..4b291a6d 100644 --- a/src/splunk_ao/utils/env_helpers.py +++ b/src/splunk_ao/utils/env_helpers.py @@ -20,8 +20,9 @@ def _get_mode_or_default(mode: str | None) -> LoggerModeType: ------- LoggerModeType The mode value to use: - - "batch": Uses scheduled OTLP batch export (default) - - "distributed": Also enables legacy trace_id/span_id continuation + Both "batch" and the temporarily retained "distributed" value use + scheduled OTLP batch export. The distinction remains only for the + deprecated ingestion-hook compatibility guard. """ if mode is None: mode = getenv("SPLUNK_AO_MODE", DEFAULT_MODE) diff --git a/src/splunk_ao/utils/singleton.py b/src/splunk_ao/utils/singleton.py index 99c00017..eb985dec 100644 --- a/src/splunk_ao/utils/singleton.py +++ b/src/splunk_ao/utils/singleton.py @@ -53,8 +53,6 @@ def _get_key( agent_stream_id: str | None, mode: str, experiment_id: str | None = None, - trace_id: str | None = None, - span_id: str | None = None, ingestion_hook_id: int | None = None, ) -> tuple[str, ...]: """ @@ -74,18 +72,13 @@ def _get_key( The experiment ID. mode: The logger mode. - trace_id: (Optional[str]) - The distributed trace ID. - span_id: (Optional[str]) - The distributed parent span ID. ingestion_hook_id: (Optional[int]) Identity of the temporary ingestion hook compatibility path. Returns ------- Tuple[str, ...] - A tuple key used for caching. Includes trace_id and span_id for proper - isolation of concurrent requests in async web servers. + A tuple key used for caching. """ _logger.debug("current thread is %s", threading.current_thread().name) @@ -121,11 +114,6 @@ def _get_key( destination_key = f"id:{routing.agent_stream_id or ''}" base_key = (*key, deployment.value, project_key, destination_key) - # Add trace_id and span_id to key if present (for distributed tracing) - if trace_id is not None: - base_key = (*base_key, trace_id) - if span_id is not None: - base_key = (*base_key, span_id) if ingestion_hook_id is not None: base_key = (*base_key, str(ingestion_hook_id)) @@ -158,8 +146,6 @@ def get( experiment_id: str | None = None, mode: str | None = None, local_metrics: list[LocalMetricConfig] | None = None, - trace_id: str | None = None, - span_id: str | None = None, ingestion_hook: Callable | None = None, ) -> SplunkAOLogger: """ @@ -197,8 +183,6 @@ def get( agent_stream_id, mode, experiment_id, - trace_id, - span_id, ingestion_hook_id=id(ingestion_hook) if ingestion_hook else None, ) @@ -221,8 +205,6 @@ def get( "experiment_id": experiment_id, "local_metrics": local_metrics, "mode": mode, - "trace_id": trace_id, - "span_id": span_id, "ingestion_hook": ingestion_hook, } # Create the logger with filtered kwargs. diff --git a/tests/conftest.py b/tests/conftest.py index cc3a8ba1..3f046681 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -62,11 +62,10 @@ from splunk_ao.collaborator import CollaboratorRole from splunk_ao.config import SplunkAOConfig from splunk_ao.configuration import _CONFIGURATION_KEYS, Configuration -from splunk_ao.decorator import _mode_context +from splunk_ao.decorator import _mode_context, _session_id_context from splunk_ao.resources.models import DatasetContent, DatasetRow, DatasetRowValuesDict from splunk_ao.resources.models.messages_list_item import MessagesListItem from splunk_ao.utils.singleton import SplunkAOLoggerSingleton -from tests.testutils.setup import setup_thread_pool_request_capture class _TestSpanSink: @@ -159,8 +158,10 @@ def _clear_otel_test_context() -> None: @pytest.fixture(autouse=True) def reset_otel_test_context() -> Generator[None, None, None]: _clear_otel_test_context() + _session_id_context.set(None) yield _clear_otel_test_context() + _session_id_context.set(None) @pytest.fixture(autouse=True) @@ -380,29 +381,6 @@ def dataset_content_150_rows(): return DatasetContent(rows=rows) -@pytest.fixture -def thread_pool_capture(): - """ - Pytest fixture that provides a function to capture thread pool requests from distributed tracing methods. - - Usage: - def test_distributed_method(thread_pool_capture): - logger = SplunkAOLogger(project="test", agent_stream="test", mode="distributed") - capture = thread_pool_capture(logger) - - logger._ingest_trace_streaming(trace) - - capture.mock_pool.assert_called_once() - request = capture.get_latest_request() - assert isinstance(request, TracesIngestRequest) - """ - - def _capture_factory(logger): - return setup_thread_pool_request_capture(logger) - - return _capture_factory - - @pytest.fixture def enable_galileo_logging(): """Temporarily enable SDK logging for tests that need to capture log output.""" diff --git a/tests/test_async_base_handler.py b/tests/test_async_base_handler.py index 10e4d2f9..cdc6dd10 100644 --- a/tests/test_async_base_handler.py +++ b/tests/test_async_base_handler.py @@ -1,14 +1,133 @@ +import asyncio import uuid from collections.abc import Generator from unittest.mock import Mock, patch import pytest +from opentelemetry.sdk.trace import ReadableSpan +from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_async_handler import SplunkAOAsyncBaseHandler from splunk_ao.logger.logger import SplunkAOLogger +from splunk_ao.session_context import set_session_context from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_async_child_enqueues_when_its_callback_ends() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOAsyncBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request" + ) + await handler.async_start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + + await handler.async_end_node(child_id, output="answer", model="model") + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["chat"] + assert str(root_id) in handler._active_steps + assert sink.force_flush_calls == 0 + + await handler.async_end_node(root_id, output="done") + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ + "chat", + "invoke_workflow", + ] + assert sink.spans[0].parent == sink.spans[1].context + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + +@pytest.mark.asyncio +async def test_concurrent_async_handlers_keep_parent_and_session_context_isolated() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + first_started = asyncio.Event() + second_started = asyncio.Event() + + async def run_handler(label: str, own_started: asyncio.Event, other_started: asyncio.Event) -> None: + set_session_context(label) + handler = SplunkAOAsyncBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name=f"{label}-root", input="request" + ) + own_started.set() + await other_started.wait() + await handler.async_start_node( + node_type="llm", + parent_run_id=root_id, + run_id=child_id, + name=f"{label}-child", + input="prompt", + model="model", + ) + await asyncio.sleep(0) + await handler.async_end_node(child_id, output="answer", model="model") + await handler.async_end_node(root_id, output="done") + finally: + set_session_context(None) + + try: + await asyncio.gather( + run_handler("conversation-a", first_started, second_started), + run_handler("conversation-b", second_started, first_started), + ) + + assert len(sink.spans) == 4 + for label in ("conversation-a", "conversation-b"): + request_spans = [ + span for span in sink.spans if (span.attributes or {}).get("gen_ai.conversation.id") == label + ] + root = next( + span + for span in request_spans + if (span.attributes or {}).get("gen_ai.operation.name") == "invoke_workflow" + ) + child = next( + span for span in request_spans if (span.attributes or {}).get("gen_ai.operation.name") == "chat" + ) + assert child.parent == root.context + assert child.context.trace_id == root.context.trace_id + assert (child.attributes or {}).get("gen_ai.conversation.id") == label + assert (root.attributes or {}).get("gen_ai.conversation.id") == label + root_trace_ids = { + span.context.trace_id + for span in sink.spans + if (span.attributes or {}).get("gen_ai.operation.name") == "invoke_workflow" + } + assert len(root_trace_ids) == 2 + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class TestSplunkAOAsyncBaseHandlerCallback: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") @@ -30,6 +149,7 @@ def handler(self, splunk_ao_logger: SplunkAOLogger) -> Generator[SplunkAOAsyncBa yield handler # Clean up after each test handler._root_node = None + splunk_ao_logger.terminate() @pytest.mark.asyncio async def test_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: @@ -82,6 +202,24 @@ async def test_start_node(self, handler: SplunkAOAsyncBaseHandler) -> None: assert handler._root_node assert handler._root_node.run_id == parent_id + @pytest.mark.asyncio + async def test_live_root_is_reused_by_async_commit( + self, handler: SplunkAOAsyncBaseHandler, splunk_ao_logger: SplunkAOLogger + ) -> None: + run_id = uuid.uuid4() + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=run_id, name="Root", input="request" + ) + live_root = handler._owned_root + live_context = splunk_ao_logger._otel_ids[live_root.id].span_context + + headers = get_tracing_headers() + await handler.async_end_node(run_id, output="done") + + assert headers["traceparent"].split("-")[2] == format(live_context.span_id, "016x") + assert splunk_ao_logger.traces[0].spans[0] is live_root + assert splunk_ao_logger.current_parent() is None + @pytest.mark.asyncio async def test_end_node(self, handler: SplunkAOAsyncBaseHandler, splunk_ao_logger: SplunkAOLogger) -> None: """Test ending a node and updating its parameters""" @@ -109,9 +247,42 @@ async def test_commit_failure_concludes_handler_owned_trace( run_id = uuid.uuid4() await handler.async_start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") - with patch.object(handler, "log_node_tree", side_effect=RuntimeError("conversion failed")): + with patch.object(handler, "_log_node_children", side_effect=RuntimeError("conversion failed")): await handler.async_end_node(run_id, output="result") assert splunk_ao_logger.current_parent() is None assert handler._nodes == {} assert handler._root_node is None + + @pytest.mark.asyncio + async def test_start_new_trace_false_preserves_handler_root_and_caller( + self, splunk_ao_logger: SplunkAOLogger + ) -> None: + splunk_ao_logger.start_trace(input="request", name="caller") + caller_operation = splunk_ao_logger.add_workflow_span(input="outer", name="outer") + handler = SplunkAOAsyncBaseHandler( + splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False + ) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + await handler.async_start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name="handler-root", input="work" + ) + await handler.async_start_node( + node_type="llm", + parent_run_id=root_id, + run_id=child_id, + name="handler-child", + input="prompt", + output="answer", + model="model", + ) + + await handler.async_end_node(root_id, output="done") + + assert splunk_ao_logger.current_parent() is caller_operation + [handler_root] = caller_operation.spans + assert handler_root.name == "handler-root" + assert [child.name for child in handler_root.spans] == ["handler-child"] + splunk_ao_logger.conclude(output="outer done") + splunk_ao_logger.conclude(output="done") diff --git a/tests/test_base_handler.py b/tests/test_base_handler.py index a730d652..1e53a11b 100644 --- a/tests/test_base_handler.py +++ b/tests/test_base_handler.py @@ -1,14 +1,225 @@ +import json import uuid from collections.abc import Generator from unittest.mock import Mock, patch import pytest +from opentelemetry import context as otel_context +from opentelemetry import propagate +from opentelemetry.context import Context +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.trace.status import StatusCode +from splunk_ao import get_tracing_headers from splunk_ao.handlers.base_handler import SplunkAOBaseHandler from splunk_ao.logger.logger import SplunkAOLogger +from splunk_ao.schema.logged import LoggedAgentSpan from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def operation_names(spans: list[ReadableSpan]) -> list[str | None]: + return [(span.attributes or {}).get("gen_ai.operation.name") for span in spans] + + +def test_normal_otel_child_enqueues_at_callback_end_without_flush() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + root_step = handler._active_steps[str(root_id)].step + handler.start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + + handler.end_node(child_id, output="answer", model="model") + + assert operation_names(sink.spans) == ["chat"] + assert str(root_id) in handler._active_steps + assert logger.current_parent() is root_step + assert sink.force_flush_calls == 0 + + handler.end_node(root_id, output="done") + + assert operation_names(sink.spans) == ["chat", "invoke_workflow"] + child_span, root_span = sink.spans + assert child_span.parent == root_span.context + assert root_span.parent is None + assert [child.name for child in root_step.spans] == ["child"] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + +def test_incremental_handler_root_is_active_beneath_caller_owned_operation() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + logger.start_trace(input="caller trace") + caller = logger.add_workflow_span(input="caller operation", name="caller") + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, start_new_trace=False, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="handler", input="request") + handler_root = handler._active_steps[str(root_id)].step + handler_context = logger._otel_ids[handler_root.id].span_context + headers = get_tracing_headers() + handler.start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + handler.end_node(child_id, output="answer", model="model") + handler.end_node(root_id, output="done") + + assert headers["traceparent"].split("-")[2] == format(handler_context.span_id, "016x") + assert logger.current_parent() is caller + child_span, handler_span = sink.spans + assert child_span.parent == handler_span.context + assert handler_span.parent == logger._otel_ids[caller.id].span_context + assert [step.name for step in caller.spans] == ["handler"] + finally: + logger.conclude(output="caller done") + logger.conclude(output="trace done") + logger.terminate() + + +BASELINE_HANDLER_TOPOLOGY = { + "invoke_workflow handler-root": ("invoke_workflow", None), + "execute_tool search": ("execute_tool", "invoke_workflow handler-root"), + "chat model": ("chat", "execute_tool search"), +} + + +def _handler_topology(carrier: dict[str, str] | None = None) -> dict[str, tuple[str | None, str | None]]: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + tool_id = uuid.uuid4() + llm_id = uuid.uuid4() + token = otel_context.attach(propagate.extract(carrier) if carrier is not None else Context()) + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="handler-root", input="request") + handler.start_node(node_type="tool", parent_run_id=root_id, run_id=tool_id, name="search", input="query") + handler.start_node( + node_type="llm", parent_run_id=tool_id, run_id=llm_id, name="model-call", input="prompt", model="model" + ) + handler.end_node(llm_id, output="answer", model="model") + handler.end_node(tool_id, output="result") + handler.end_node(root_id, output="done") + + names_by_context = {(span.context.trace_id, span.context.span_id): span.name for span in sink.spans} + return { + span.name: ( + (span.attributes or {}).get("gen_ai.operation.name"), + names_by_context.get((span.parent.trace_id, span.parent.span_id)) if span.parent is not None else None, + ) + for span in sink.spans + } + finally: + logger.terminate() + otel_context.detach(token) + + +def test_distributed_and_non_distributed_handlers_preserve_baseline_local_topology() -> None: + remote_carrier = { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "tracestate": "vendor=value", + } + + assert _handler_topology() == BASELINE_HANDLER_TOPOLOGY + assert _handler_topology(remote_carrier) == BASELINE_HANDLER_TOPOLOGY + + +def test_leaf_only_handler_root_is_active_exportable_and_enqueued_without_flush() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + try: + handler.start_node( + node_type="llm", parent_run_id=None, run_id=root_id, name="leaf-root", input="prompt", model="model" + ) + active_span_id = get_tracing_headers()["traceparent"].split("-")[2] + + handler.end_node(root_id, output="answer", model="model") + + [span] = sink.spans + assert span.name == "chat model" + assert (span.attributes or {}).get("gen_ai.operation.name") == "chat" + assert format(span.context.span_id, "016x") == active_span_id + assert span.parent is None + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + +def test_unsampled_remote_parent_suppresses_incremental_handler_export() -> None: + remote_context = propagate.extract({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"}) + token = otel_context.attach(remote_context) + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + handler.start_node( + node_type="llm", parent_run_id=root_id, run_id=child_id, name="child", input="prompt", model="model" + ) + handler.end_node(child_id, output="answer", model="model") + handler.end_node(root_id, output="done") + + assert sink.spans == [] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + otel_context.detach(token) + + +def test_handler_error_span_enqueues_with_error_status_before_root_end() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + handler = SplunkAOBaseHandler(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="root", input="request") + handler.start_node( + node_type="tool", parent_run_id=root_id, run_id=child_id, name="failing-tool", input="arguments" + ) + + handler.end_node(child_id, error="tool failed") + + [span] = sink.spans + assert span.name == "execute_tool failing-tool" + assert span.status.status_code is StatusCode.ERROR + assert json.loads(str((span.attributes or {}).get("gen_ai.tool.call.result"))) == {"value": "tool failed"} + assert str(root_id) in handler._active_steps + assert sink.force_flush_calls == 0 + + handler.end_node(root_id, output="failed", status_code=500) + finally: + logger.terminate() + + class TestSplunkAOBaseHandler: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") @@ -24,9 +235,9 @@ def splunk_ao_logger(self, mock_traces_client: Mock, mock_projects_client: Mock, @pytest.fixture def handler(self, splunk_ao_logger: SplunkAOLogger) -> Generator[SplunkAOBaseHandler, None, None]: """Creates a SplunkAOBaseHandler with a mock logger""" - return SplunkAOBaseHandler(splunk_ao_logger=splunk_ao_logger, flush_on_chain_end=False) - # Reset the root node before each test - # Clean up after each test + handler = SplunkAOBaseHandler(splunk_ao_logger=splunk_ao_logger, flush_on_chain_end=False) + yield handler + splunk_ao_logger.terminate() def test_initialization(self, splunk_ao_logger: SplunkAOLogger) -> None: """Test callback initialization with various parameters""" @@ -77,6 +288,45 @@ def test_start_node(self, handler: SplunkAOBaseHandler) -> None: assert handler._root_node assert handler._root_node.run_id == parent_id + def test_root_is_live_for_w3c_injection_and_reused_at_commit( + self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAOLogger + ) -> None: + run_id = uuid.uuid4() + handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Root", input="request") + live_root = handler._owned_root + live_context = splunk_ao_logger._otel_ids[live_root.id].span_context + + headers = get_tracing_headers() + handler.end_node(run_id, output="done") + + assert headers["traceparent"].split("-")[2] == format(live_context.span_id, "016x") + assert len(splunk_ao_logger.traces[0].spans) == 1 + assert splunk_ao_logger.traces[0].spans[0] is live_root + assert splunk_ao_logger.current_parent() is None + + def test_late_langgraph_agent_classification_preserves_live_otel_identity( + self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAOLogger + ) -> None: + root_id = uuid.uuid4() + child_id = uuid.uuid4() + root_node = handler.start_node( + node_type="chain", parent_run_id=None, run_id=root_id, name="Root", input="request" + ) + original_root = handler._owned_root + original_ids = splunk_ao_logger._otel_ids[original_root.id] + + root_node.node_type = "agent" + handler.start_node(node_type="llm", parent_run_id=root_id, run_id=child_id, name="Child", input="prompt") + + assert isinstance(handler._owned_root, LoggedAgentSpan) + assert handler._owned_root.id == original_root.id + assert splunk_ao_logger._otel_ids[handler._owned_root.id] is original_ids + assert splunk_ao_logger.current_parent() is handler._owned_root + + handler.end_node(child_id, output="answer", model="model") + handler.end_node(root_id, output="done") + assert isinstance(splunk_ao_logger.traces[0].spans[0], LoggedAgentSpan) + def test_end_node(self, handler: SplunkAOBaseHandler, splunk_ao_logger: SplunkAOLogger) -> None: """Test ending a node and updating its parameters""" # Create a node @@ -146,7 +396,7 @@ def test_commit_failure_concludes_only_handler_owned_trace( run_id = uuid.uuid4() handler.start_node(node_type="chain", parent_run_id=None, run_id=run_id, name="Test", input="test") - with patch.object(handler, "log_node_tree", side_effect=RuntimeError("conversion failed")): + with patch.object(handler, "_log_node_children", side_effect=RuntimeError("conversion failed")): handler.end_node(run_id, output="result") assert splunk_ao_logger.current_parent() is None @@ -154,7 +404,8 @@ def test_commit_failure_concludes_only_handler_owned_trace( assert handler._root_node is None def test_commit_failure_preserves_caller_owned_trace(self, splunk_ao_logger: SplunkAOLogger) -> None: - caller_trace = splunk_ao_logger.start_trace(input="request", name="caller") + splunk_ao_logger.start_trace(input="request", name="caller") + caller_operation = splunk_ao_logger.add_workflow_span(input="outer", name="outer") handler = SplunkAOBaseHandler( splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False ) @@ -164,5 +415,39 @@ def test_commit_failure_preserves_caller_owned_trace(self, splunk_ao_logger: Spl with patch.object(handler, "log_node_tree", side_effect=RuntimeError("conversion failed")): handler.end_node(run_id, output="result") - assert splunk_ao_logger.current_parent() is caller_trace + assert splunk_ao_logger.current_parent() is caller_operation + splunk_ao_logger.conclude(output="outer done") + splunk_ao_logger.conclude(output="done") + + def test_start_new_trace_false_preserves_handler_root_beneath_caller( + self, splunk_ao_logger: SplunkAOLogger + ) -> None: + splunk_ao_logger.start_trace(input="request", name="caller") + caller_operation = splunk_ao_logger.add_workflow_span(input="outer", name="outer") + handler = SplunkAOBaseHandler( + splunk_ao_logger=splunk_ao_logger, start_new_trace=False, flush_on_chain_end=False + ) + caller_context = splunk_ao_logger._otel_ids[caller_operation.id].span_context + root_id = uuid.uuid4() + child_id = uuid.uuid4() + handler.start_node(node_type="chain", parent_run_id=None, run_id=root_id, name="handler-root", input="work") + headers = get_tracing_headers() + handler.start_node( + node_type="llm", + parent_run_id=root_id, + run_id=child_id, + name="handler-child", + input="prompt", + output="answer", + model="model", + ) + + handler.end_node(root_id, output="done") + + assert splunk_ao_logger.current_parent() is caller_operation + assert headers["traceparent"].split("-")[2] == format(caller_context.span_id, "016x") + [handler_root] = caller_operation.spans + assert handler_root.name == "handler-root" + assert [child.name for child in handler_root.spans] == ["handler-child"] + splunk_ao_logger.conclude(output="outer done") splunk_ao_logger.conclude(output="done") diff --git a/tests/test_crewai_handler.py b/tests/test_crewai_handler.py index 61cd33e5..96d36e10 100644 --- a/tests/test_crewai_handler.py +++ b/tests/test_crewai_handler.py @@ -5,11 +5,13 @@ from unittest.mock import Mock, patch import pytest +from opentelemetry.sdk.trace import ReadableSpan # Skip all tests in this module on Python 3.14+ (crewai doesn't support it yet) pytestmark = pytest.mark.skipif(sys.version_info >= (3, 14), reason="crewai does not support Python 3.14+") from splunk_ao.handlers.crewai.handler import CrewAIEventListener # noqa: E402 +from splunk_ao.logger.logger import SplunkAOLogger # noqa: E402 from splunk_ao.schema.handlers import NodeType # noqa: E402 from tests.testutils.setup import ( # noqa: E402 setup_mock_logstreams_client, @@ -76,6 +78,65 @@ def __init__(self, raw="Test output"): self.raw = raw +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_crewai_adapter_enqueues_task_at_callback_end_without_flush() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + crew_id = uuid.uuid4() + task_id = uuid.uuid4() + crew = MockCrew(crew_id=crew_id) + agent = MockAgent(crew=crew) + task = MockTask(task_id=task_id, description="Research market trends", agent=agent) + try: + with ( + patch("splunk_ao.handlers.crewai.handler._crewai_imports_resolved", True), + patch("splunk_ao.handlers.crewai.handler.CREWAI_AVAILABLE", False), + patch("splunk_ao.handlers.crewai.handler.LITE_LLM_AVAILABLE", False), + ): + callback = CrewAIEventListener(splunk_ao_logger=logger, start_new_trace=True, flush_on_crew_completed=False) + + callback._handle_crew_kickoff_started( + MockSource(id=crew_id), MockEvent(crew_name="Test Crew", inputs={"question": "market trends"}) + ) + callback._handle_task_started(MockSource(id=task_id), MockEvent(task=task)) + + callback._handle_task_completed(MockSource(id=task_id), MockEvent(output=MockOutput("Done"))) + + assert len(sink.spans) == 1 + assert sink.spans[0].name == "invoke_workflow Research market trends" + assert str(crew_id) in callback._handler._active_steps + assert sink.force_flush_calls == 0 + + callback._handle_crew_kickoff_completed( + MockSource(id=crew_id), MockEvent(output=MockOutput("Crew completed successfully")) + ) + + child, root = sink.spans + assert child.parent == root.context + assert [span.name for span in sink.spans] == [ + "invoke_workflow Research market trends", + "invoke_workflow Test Crew", + ] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + @pytest.fixture def mock_splunk_ao_logger(): """Creates a mock Galileo logger for testing.""" diff --git a/tests/test_decorator_distributed.py b/tests/test_decorator_distributed.py index d01f74ef..3b6a2e77 100644 --- a/tests/test_decorator_distributed.py +++ b/tests/test_decorator_distributed.py @@ -3,12 +3,11 @@ from unittest.mock import Mock, patch import pytest +from opentelemetry import context as otel_context from galileo_core.schemas.shared.document import Document from galileo_core.schemas.shared.multimodal import ContentModality -from splunk_ao import Message, MessageRole, log, splunk_ao_context -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER -from splunk_ao.decorator import _parent_id_context, _trace_id_context +from splunk_ao import Message, MessageRole, extract_tracing_context, log, splunk_ao_context from splunk_ao.schema.content_blocks import DataContentBlock, TextContentBlock from splunk_ao.tracing import get_tracing_headers from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client @@ -49,33 +48,40 @@ def test_decorator_get_tracing_headers(reset_context: None, distributed_clients: @log(span_type="workflow") def orchestrator(query: str) -> dict: - trace = splunk_ao_context.get_current_trace() - return {"result": query, "headers": get_tracing_headers(), "trace_id": str(trace.id)} + return {"result": query, "headers": get_tracing_headers()} result = orchestrator("test input") headers = result["headers"] - assert headers[TRACE_ID_HEADER] == result["trace_id"] - assert PARENT_ID_HEADER in headers + assert "traceparent" in headers + assert not any(header.lower().startswith("splunk-ao-") for header in headers) + traceparent = headers["traceparent"].split("-") + [span] = logger._sink.spans + assert int(traceparent[1], 16) == span.context.trace_id + assert int(traceparent[2], 16) == span.context.span_id assert logger.current_parent() is None def test_decorator_respects_incoming_distributed_context(reset_context: None, distributed_clients: Mock) -> None: - trace_id = "12345678-1234-4678-9abc-123456789abc" - parent_id = "87654321-4321-4876-9cba-987654321cba" - _trace_id_context.set(trace_id) - _parent_id_context.set(parent_id) + trace_id = "4bf92f3577b34da6a3ce929d0e0e4736" + parent_id = "00f067aa0ba902b7" logger = init_logger() @log(span_type="workflow") def downstream_service(query: str) -> str: return f"processed: {query}" - assert downstream_service("test input") == "processed: test input" + token = otel_context.attach(extract_tracing_context({"traceparent": f"00-{trace_id}-{parent_id}-01"})) + try: + assert downstream_service("test input") == "processed: test input" + finally: + otel_context.detach(token) + assert logger.mode == "distributed" - assert str(logger.traces[0].id) == trace_id - assert logger.traces[0].name == "stub_trace" - assert (logger._sink.spans[-1].attributes or {})["gen_ai.operation.name"] == "invoke_workflow" + exported = logger._sink.spans[-1] + assert exported.context.trace_id == int(trace_id, 16) + assert exported.parent.span_id == int(parent_id, 16) + assert (exported.attributes or {})["gen_ai.operation.name"] == "invoke_workflow" def test_completed_workflow_is_enqueued_and_flush_does_not_change_ownership( diff --git a/tests/test_exporter_sink.py b/tests/test_exporter_sink.py index 6dd7cef0..8c5f6c92 100644 --- a/tests/test_exporter_sink.py +++ b/tests/test_exporter_sink.py @@ -2,6 +2,7 @@ from collections.abc import Generator, Sequence from typing import Any +from unittest.mock import patch import pytest from opentelemetry import trace @@ -58,6 +59,15 @@ def test_batch_processor_wraps_exporter(shutdown_workers: list[Any]) -> None: assert isinstance(processor, BatchSpanProcessor) +def test_default_batch_processor_delegates_configuration_to_otel() -> None: + exporter = RecordingExporter() + + with patch("splunk_ao.exporter.sink.BatchSpanProcessor") as processor_class: + build_batch_processor(exporter) + + processor_class.assert_called_once_with(exporter) + + def test_batch_processor_accepts_custom_config(shutdown_workers: list[Any]) -> None: exporter = RecordingExporter() processor = build_batch_processor(exporter, BatchConfig(max_export_batch_size=1)) diff --git a/tests/test_http_instrumentation.py b/tests/test_http_instrumentation.py new file mode 100644 index 00000000..0ecd7bfa --- /dev/null +++ b/tests/test_http_instrumentation.py @@ -0,0 +1,635 @@ +from __future__ import annotations + +from collections.abc import Generator +from dataclasses import dataclass +from typing import Any +from unittest.mock import MagicMock, patch + +import aiohttp +import httpx +import pytest +import requests +from aiohttp import TraceRequestExceptionParams, TraceRequestStartParams +from fastapi import FastAPI +from multidict import CIMultiDict +from opentelemetry import baggage, context, propagate, trace +from opentelemetry.context import Context +from opentelemetry.instrumentation.aiohttp_client import AioHttpClientInstrumentor +from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor +from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor +from opentelemetry.instrumentation.requests import RequestsInstrumentor +from opentelemetry.instrumentation.starlette import StarletteInstrumentor +from opentelemetry.propagators import textmap +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse +from starlette.routing import Route +from starlette.testclient import TestClient +from yarl import URL + +from splunk_ao import configure_distributed_tracing, instrument_distributed_tracing +from splunk_ao.http_instrumentation import ( + _client_provider_ids, + _configured_providers, + _instrumented_apps, + _load_instrumentors, +) +from splunk_ao.logger.logger import SplunkAOLogger +from splunk_ao.session_context import GEN_AI_CONVERSATION_ID, SplunkAOSessionPropagator, _session_id_context + + +class FakeTracerProvider: + def add_span_processor(self, span_processor: Any) -> None: + return None + + def get_tracer( + self, + instrumenting_module_name: str, + instrumenting_library_version: str | None = None, + schema_url: str | None = None, + attributes: Any | None = None, + ) -> MagicMock: + return MagicMock() + + +class CustomPropagator(textmap.TextMapPropagator): + def extract( + self, + carrier: textmap.CarrierT, + context: context.Context | None = None, + getter: textmap.Getter[textmap.CarrierT] = textmap.default_getter, + ) -> context.Context: + return context if context is not None else Context() + + def inject( + self, + carrier: textmap.CarrierT, + context: context.Context | None = None, + setter: textmap.Setter[textmap.CarrierT] = textmap.default_setter, + ) -> None: + setter.set(carrier, "x-custom-context", "preserved") + + @property + def fields(self) -> set[str]: + return {"x-custom-context"} + + +@dataclass +class RecordingInstrumentor: + instrument_calls: list[dict[str, Any]] + app_calls: list[tuple[Any, dict[str, Any]]] + + def instrument(self, **kwargs: Any) -> None: + self.instrument_calls.append(kwargs) + + def instrument_app(self, app: Any, **kwargs: Any) -> None: + self.app_calls.append((app, kwargs)) + + +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + return True + + def shutdown(self) -> None: + return None + + +def instrumentor_type(recorder: RecordingInstrumentor) -> type: + class Instrumentor: + def __new__(cls) -> RecordingInstrumentor: + return recorder + + @classmethod + def instrument_app(cls, app: Any, **kwargs: Any) -> None: + recorder.instrument_app(app, **kwargs) + + return Instrumentor + + +@pytest.fixture(autouse=True) +def reset_instrumentation_state() -> Generator[None, None, None]: + previous_propagator = propagate.get_global_textmap() + session_token = _session_id_context.set(None) + _client_provider_ids.clear() + _configured_providers.clear() + _instrumented_apps.clear() + try: + yield + finally: + _client_provider_ids.clear() + _configured_providers.clear() + _instrumented_apps.clear() + propagate.set_global_textmap(previous_propagator) + _session_id_context.reset(session_token) + + +@pytest.fixture +def instrumentors() -> tuple[dict[str, type], dict[str, RecordingInstrumentor]]: + recorders = { + name: RecordingInstrumentor(instrument_calls=[], app_calls=[]) + for name in ("fastapi", "starlette", "requests", "httpx", "aiohttp_client") + } + types = { + "fastapi": instrumentor_type(recorders["fastapi"]), + "starlette": instrumentor_type(recorders["starlette"]), + "requests": instrumentor_type(recorders["requests"]), + "httpx": instrumentor_type(recorders["httpx"]), + "aiohttp-client": instrumentor_type(recorders["aiohttp_client"]), + } + return types, recorders + + +def test_missing_extra_fails_before_instrumentation_or_propagator_change() -> None: + previous_propagator = propagate.get_global_textmap() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", side_effect=ImportError("install the extra")), + pytest.raises(ImportError, match="install the extra"), + ): + instrument_distributed_tracing(tracer_provider=FakeTracerProvider()) + + assert _client_provider_ids == {} + assert _instrumented_apps == set() + assert propagate.get_global_textmap() is previous_propagator + + +def test_configure_distributed_tracing_creates_and_returns_provider( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = TracerProvider() + processor = MagicMock() + app = FastAPI() + + with ( + patch("splunk_ao.http_instrumentation.SDKTracerProvider", return_value=provider), + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor", return_value=processor) as add_processor, + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + ): + configured_provider = configure_distributed_tracing(app=app) + + assert configured_provider is provider + add_processor.assert_called_once_with(provider) + assert recorders["fastapi"].app_calls == [(app, {"tracer_provider": provider})] + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + + +def test_configure_distributed_tracing_registers_processor_once_per_provider( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = TracerProvider() + processor = MagicMock() + + with ( + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor", return_value=processor) as add_processor, + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + ): + first_result = configure_distributed_tracing(tracer_provider=provider) + second_result = configure_distributed_tracing(tracer_provider=provider) + + assert first_result is provider + assert second_result is provider + add_processor.assert_called_once_with(provider) + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + + +def test_configure_distributed_tracing_missing_extra_does_not_register_processor() -> None: + provider = TracerProvider() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", side_effect=ImportError("install the extra")), + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor") as add_processor, + pytest.raises(ImportError, match="install the extra"), + ): + configure_distributed_tracing(tracer_provider=provider) + + add_processor.assert_not_called() + assert provider not in _configured_providers + + +def test_configure_distributed_tracing_invalid_app_does_not_register_processor( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, _ = instrumentors + provider = TracerProvider() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + patch("splunk_ao.http_instrumentation.add_splunk_ao_span_processor") as add_processor, + pytest.raises(TypeError, match="FastAPI or Starlette"), + ): + configure_distributed_tracing(tracer_provider=provider, app=object()) + + add_processor.assert_not_called() + assert provider not in _configured_providers + + +def test_distributed_tracing_extra_loads_every_supported_upstream_instrumentor() -> None: + assert set(_load_instrumentors()) == {"fastapi", "starlette", "requests", "httpx", "aiohttp-client"} + + +@pytest.mark.parametrize(("app", "expected"), [(FastAPI(), "fastapi"), (Starlette(), "starlette")]) +def test_instruments_matching_app_and_every_enabled_client_once( + app: Any, expected: str, instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]] +) -> None: + types, recorders = instrumentors + provider = FakeTracerProvider() + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing(tracer_provider=provider, app=app) + instrument_distributed_tracing(tracer_provider=provider, app=app) + + assert recorders[expected].app_calls == [(app, {"tracer_provider": provider})] + other_framework = "starlette" if expected == "fastapi" else "fastapi" + assert recorders[other_framework].app_calls == [] + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + assert recorders["httpx"].instrument_calls == [{"tracer_provider": provider}] + assert recorders["aiohttp_client"].instrument_calls == [{"tracer_provider": provider}] + assert isinstance(propagate.get_global_textmap(), SplunkAOSessionPropagator) + + +def test_disabled_clients_are_not_instrumented( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=FakeTracerProvider(), + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + assert all(not recorder.instrument_calls for recorder in recorders.values()) + + +def test_provider_conflict_fails_before_partial_instrumentation( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + first_provider = FakeTracerProvider() + second_provider = FakeTracerProvider() + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing(tracer_provider=first_provider) + with pytest.raises(RuntimeError, match="another tracer provider"): + instrument_distributed_tracing(tracer_provider=second_provider, app=FastAPI()) + + assert recorders["fastapi"].app_calls == [] + assert all(len(recorders[name].instrument_calls) == 1 for name in ("requests", "httpx", "aiohttp_client")) + + +def test_uses_current_provider_without_replacing_it( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, recorders = instrumentors + provider = FakeTracerProvider() + + with ( + patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types), + patch("splunk_ao.http_instrumentation.trace.get_tracer_provider", return_value=provider), + patch.object(trace, "set_tracer_provider") as set_provider, + ): + instrument_distributed_tracing() + + set_provider.assert_not_called() + assert recorders["requests"].instrument_calls == [{"tracer_provider": provider}] + + +def test_installed_global_propagator_injects_only_conversation_session( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, _ = instrumentors + span_context = SpanContext( + trace_id=0x4BF92F3577B34DA6A3CE929D0E0E4736, + span_id=0x00F067AA0BA902B7, + is_remote=False, + trace_flags=TraceFlags.SAMPLED, + ) + token = context.attach(trace.set_span_in_context(NonRecordingSpan(span_context))) + _session_id_context.set("conversation-123") + try: + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=FakeTracerProvider(), + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + carrier: dict[str, str] = {} + propagate.inject(carrier) + finally: + context.detach(token) + + extracted = propagate.extract(carrier) + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-123" + assert "splunk_ao.session.id" not in carrier.get("baggage", "") + + +def test_session_adapter_preserves_application_configured_propagator( + instrumentors: tuple[dict[str, type], dict[str, RecordingInstrumentor]], +) -> None: + types, _ = instrumentors + custom = CustomPropagator() + propagate.set_global_textmap(custom) + + with patch("splunk_ao.http_instrumentation._load_instrumentors", return_value=types): + instrument_distributed_tracing( + tracer_provider=FakeTracerProvider(), + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + + carrier: dict[str, str] = {} + propagate.inject(carrier) + installed = propagate.get_global_textmap() + assert carrier["x-custom-context"] == "preserved" + assert isinstance(installed, SplunkAOSessionPropagator) + assert installed._delegate is custom + + +def test_real_requests_instrumentor_injects_trace_and_conversation_without_manual_headers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured_headers: dict[str, str] = {} + + def capture_send(session: requests.Session, request: requests.PreparedRequest, **kwargs: Any) -> requests.Response: + del session, kwargs + captured_headers.update(request.headers) + response = requests.Response() + response.status_code = 200 + response.request = request + response.url = request.url + return response + + monkeypatch.setattr(requests.Session, "send", capture_send) + provider = TracerProvider() + requests_instrumentor = RequestsInstrumentor() + try: + instrument_distributed_tracing( + tracer_provider=provider, instrument_httpx=False, instrument_aiohttp_client=False + ) + _session_id_context.set("conversation-outbound") + with provider.get_tracer(__name__).start_as_current_span("caller") as caller: + response = requests.get("https://example.test/automatic-context", timeout=1) + expected_trace_id = f"{caller.get_span_context().trace_id:032x}" + + assert response.status_code == 200 + assert captured_headers["traceparent"].split("-")[1] == expected_trace_id + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-outbound" + assert "splunk_ao.session.id" not in captured_headers["baggage"] + finally: + requests_instrumentor.uninstrument() + provider.shutdown() + + +def test_real_fastapi_instrumentor_extracts_trace_and_conversation_without_sdk_middleware() -> None: + expected_trace_id = 0x4BF92F3577B34DA6A3CE929D0E0E4736 + observed: dict[str, Any] = {} + provider = TracerProvider() + app = FastAPI() + + @app.get("/automatic-context") + def automatic_context() -> dict[str, bool]: + observed["trace_id"] = trace.get_current_span().get_span_context().trace_id + observed["conversation_id"] = baggage.get_baggage(GEN_AI_CONVERSATION_ID) + return {"ok": True} + + try: + instrument_distributed_tracing( + tracer_provider=provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + with TestClient(app) as client: + response = client.get( + "/automatic-context", + headers={ + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=conversation-inbound", + }, + ) + + assert response.json() == {"ok": True} + assert observed == {"trace_id": expected_trace_id, "conversation_id": "conversation-inbound"} + finally: + FastAPIInstrumentor.uninstrument_app(app) + provider.shutdown() + + +def test_automatic_requests_to_fastapi_continues_active_path1_context_without_manual_headers( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured_headers: dict[str, str] = {} + observed: dict[str, Any] = {} + + def capture_send(session: requests.Session, request: requests.PreparedRequest, **kwargs: Any) -> requests.Response: + del session, kwargs + captured_headers.update(request.headers) + response = requests.Response() + response.status_code = 200 + response.request = request + response.url = request.url + return response + + monkeypatch.setattr(requests.Session, "send", capture_send) + provider = TracerProvider() + app = FastAPI() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + + @app.get("/downstream") + def downstream() -> dict[str, bool]: + observed["trace_id"] = trace.get_current_span().get_span_context().trace_id + observed["conversation_id"] = baggage.get_baggage(GEN_AI_CONVERSATION_ID) + return {"ok": True} + + try: + instrument_distributed_tracing( + tracer_provider=provider, app=app, instrument_httpx=False, instrument_aiohttp_client=False + ) + logger.set_session("conversation-cross-service") + logger.start_trace(input="request") + path1_operation = logger.add_workflow_span(input="agent work", name="agent") + path1_context = logger._otel_ids[path1_operation.id].span_context + + response = requests.get("https://example.test/downstream", timeout=1) + downstream_token = context.attach(Context()) + try: + with TestClient(app) as client: + downstream_response = client.get( + "/downstream", + headers={"traceparent": captured_headers["traceparent"], "baggage": captured_headers["baggage"]}, + ) + finally: + context.detach(downstream_token) + + assert response.status_code == 200 + assert downstream_response.json() == {"ok": True} + assert captured_headers["traceparent"].split("-")[1] == format(path1_context.trace_id, "032x") + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-cross-service" + assert observed == {"trace_id": path1_context.trace_id, "conversation_id": "conversation-cross-service"} + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + RequestsInstrumentor().uninstrument() + FastAPIInstrumentor.uninstrument_app(app) + provider.shutdown() + + +def test_real_httpx_sync_instrumentation_injects_active_path1_context(monkeypatch: pytest.MonkeyPatch) -> None: + captured_headers: dict[str, str] = {} + + def capture_sync(transport: httpx.HTTPTransport, request: httpx.Request) -> httpx.Response: + del transport + captured_headers.update(request.headers) + return httpx.Response(200, request=request) + + monkeypatch.setattr(httpx.HTTPTransport, "handle_request", capture_sync) + provider = TracerProvider() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + try: + instrument_distributed_tracing( + tracer_provider=provider, instrument_requests=False, instrument_aiohttp_client=False + ) + logger.set_session("conversation-httpx") + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="agent") + operation_context = logger._otel_ids[operation.id].span_context + + with httpx.Client() as client: + assert client.get("https://example.test/sync").status_code == 200 + + assert captured_headers["traceparent"].split("-")[1] == format(operation_context.trace_id, "032x") + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-httpx" + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + HTTPXClientInstrumentor().uninstrument() + provider.shutdown() + + +@pytest.mark.asyncio +async def test_real_httpx_async_instrumentation_injects_active_path1_context(monkeypatch: pytest.MonkeyPatch) -> None: + captured_headers: dict[str, str] = {} + + async def capture_async(transport: httpx.AsyncHTTPTransport, request: httpx.Request) -> httpx.Response: + del transport + captured_headers.update(request.headers) + return httpx.Response(200, request=request) + + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", capture_async) + provider = TracerProvider() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + try: + instrument_distributed_tracing( + tracer_provider=provider, instrument_requests=False, instrument_aiohttp_client=False + ) + logger.set_session("conversation-httpx") + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="agent") + operation_context = logger._otel_ids[operation.id].span_context + + async with httpx.AsyncClient() as client: + assert (await client.get("https://example.test/async")).status_code == 200 + + assert captured_headers["traceparent"].split("-")[1] == format(operation_context.trace_id, "032x") + assert captured_headers["baggage"] == "gen_ai.conversation.id=conversation-httpx" + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + HTTPXClientInstrumentor().uninstrument() + provider.shutdown() + + +@pytest.mark.asyncio +async def test_real_aiohttp_instrumentation_installs_hook_that_injects_active_path1_context() -> None: + provider = TracerProvider() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=RecordingSink()) + try: + instrument_distributed_tracing(tracer_provider=provider, instrument_requests=False, instrument_httpx=False) + logger.set_session("conversation-aiohttp") + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="agent") + operation_context = logger._otel_ids[operation.id].span_context + + async with aiohttp.ClientSession() as session: + trace_config = next( + config + for config in session._trace_configs + if getattr(config, "_is_instrumented_by_opentelemetry", False) + ) + trace_context = trace_config.trace_config_ctx() + headers: CIMultiDict[str] = CIMultiDict() + url = URL("https://example.test/aiohttp") + start_params = TraceRequestStartParams("GET", url, headers) + for callback in trace_config.on_request_start: + await callback(session, trace_context, start_params) + + assert headers["traceparent"].split("-")[1] == format(operation_context.trace_id, "032x") + assert headers["baggage"] == "gen_ai.conversation.id=conversation-aiohttp" + + exception_params = TraceRequestExceptionParams("GET", url, headers, RuntimeError("synthetic stop")) + for callback in trace_config.on_request_exception: + await callback(session, trace_context, exception_params) + finally: + logger.clear_session() + if logger.current_parent() is not None: + logger.conclude(output="done", conclude_all=True) + logger.terminate() + AioHttpClientInstrumentor().uninstrument() + provider.shutdown() + + +def test_real_starlette_instrumentor_extracts_remote_context_without_sdk_middleware() -> None: + expected_trace_id = 0x4BF92F3577B34DA6A3CE929D0E0E4736 + observed: dict[str, Any] = {} + provider = TracerProvider() + + async def automatic_context(request: Request) -> JSONResponse: + del request + observed["trace_id"] = trace.get_current_span().get_span_context().trace_id + observed["conversation_id"] = baggage.get_baggage(GEN_AI_CONVERSATION_ID) + return JSONResponse({"ok": True}) + + app = Starlette(routes=[Route("/automatic-context", automatic_context)]) + try: + instrument_distributed_tracing( + tracer_provider=provider, + app=app, + instrument_requests=False, + instrument_httpx=False, + instrument_aiohttp_client=False, + ) + with TestClient(app) as client: + response = client.get( + "/automatic-context", + headers={ + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=conversation-starlette", + }, + ) + + assert response.json() == {"ok": True} + assert observed == {"trace_id": expected_trace_id, "conversation_id": "conversation-starlette"} + finally: + StarletteInstrumentor.uninstrument_app(app) + provider.shutdown() diff --git a/tests/test_langchain.py b/tests/test_langchain.py index ca8b6144..f2bb85e9 100644 --- a/tests/test_langchain.py +++ b/tests/test_langchain.py @@ -9,6 +9,7 @@ from langchain_core.documents import Document from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage from langchain_core.outputs import ChatGeneration, LLMResult +from opentelemetry.sdk.trace import ReadableSpan from galileo_core.schemas.shared.document import Document as GalileoDocument from splunk_ao import Message, MessageRole, splunk_ao_context @@ -22,6 +23,51 @@ from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_langchain_adapter_enqueues_child_at_callback_end_without_flush() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + callback = SplunkAOCallback(splunk_ao_logger=logger, flush_on_chain_end=False) + root_id = uuid.uuid4() + child_id = uuid.uuid4() + try: + callback.on_chain_start(serialized={"name": "root"}, inputs={"query": "question"}, run_id=root_id) + callback.on_tool_start(serialized={"name": "search"}, input_str="query", run_id=child_id, parent_run_id=root_id) + + callback.on_tool_end(output="result", run_id=child_id, parent_run_id=root_id) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["execute_tool"] + assert str(root_id) in callback._handler._active_steps + assert sink.force_flush_calls == 0 + + callback.on_chain_end(outputs={"answer": "done"}, run_id=root_id) + + child, root = sink.spans + assert child.parent == root.context + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ + "execute_tool", + "invoke_workflow", + ] + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class TestSplunkAOCallback: @pytest.fixture @patch("splunk_ao.logger.logger.AgentStreams") diff --git a/tests/test_logger_distributed.py b/tests/test_logger_distributed.py index 647c1e5b..c1026c5c 100644 --- a/tests/test_logger_distributed.py +++ b/tests/test_logger_distributed.py @@ -1,10 +1,11 @@ from collections.abc import Generator -from unittest.mock import Mock -from uuid import uuid4 +from unittest.mock import Mock, patch import pytest +from opentelemetry import context from opentelemetry.sdk.trace import ReadableSpan +from splunk_ao import extract_tracing_context from splunk_ao.logger import SplunkAOLogger @@ -37,6 +38,23 @@ def operation_names(spans: list[ReadableSpan]) -> list[str | None]: return [(span.attributes or {}).get("gen_ai.operation.name") for span in spans] +@pytest.mark.parametrize("mode", ["batch", "distributed"]) +def test_both_mode_values_construct_the_shared_batch_span_sink(mode: str) -> None: + sink = RecordingSink() + exporter = Mock() + with ( + patch("splunk_ao.logger.logger.build_standalone_exporter", return_value=exporter), + patch("splunk_ao.logger.logger.build_span_sink", return_value=sink) as build_sink, + ): + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", mode=mode) + + try: + build_sink.assert_called_once_with(exporter) + assert not hasattr(logger, "_task_handler") + finally: + logger.terminate() + + def test_distributed_mode_uses_otlp_completion_queue(distributed_logger: tuple[SplunkAOLogger, RecordingSink]) -> None: logger, sink = distributed_logger root = logger.start_trace(input="question") @@ -125,46 +143,82 @@ def test_distributed_conclude_all_emits_inner_to_outer( assert len({span.context.trace_id for span in sink.spans if span.context is not None}) == 1 -def test_distributed_continuation_keeps_stub_parent_local() -> None: +@pytest.mark.parametrize("mode", ["batch", "distributed"]) +def test_retained_mode_values_use_identical_w3c_otlp_behavior(mode: str) -> None: sink = RecordingSink() - trace_id = str(uuid4()) - parent_id = str(uuid4()) - logger = SplunkAOLogger( - project_id="project-id", - agent_stream_id="log-stream-id", - mode="distributed", - trace_id=trace_id, - span_id=parent_id, - _sink=sink, - ) + trace_id = "4bf92f3577b34da6a3ce929d0e0e4736" + parent_id = "00f067aa0ba902b7" + token = context.attach(extract_tracing_context({"traceparent": f"00-{trace_id}-{parent_id}-01"})) + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="log-stream-id", mode=mode, _sink=sink) try: - assert str(logger.traces[0].id) == trace_id - assert str(logger.current_parent().id) == parent_id + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="work", name="operation") + operation_context = logger._otel_ids[operation.id] + logger.conclude(output="done") + logger.conclude(output="complete") + + assert operation_context.span_context.trace_id == int(trace_id, 16) + assert operation_context.parent_span_context is not None + assert operation_names(sink.spans) == ["invoke_workflow"] + assert sink.spans[0].parent is not None + assert sink.spans[0].parent.span_id == int(parent_id, 16) + assert sink.force_flush_calls == 0 + assert logger.current_parent() is None + finally: + logger.terminate() + context.detach(token) - logger.add_llm_span(input="prompt", output="answer", model="model") - assert operation_names(sink.spans) == ["chat"] - assert sink.spans[0].parent == logger._otel_ids[logger.current_parent().id].span_context +@pytest.mark.parametrize("mode", ["batch", "distributed"]) +def test_mode_does_not_change_parent_child_topology_or_require_flush(mode: str) -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="log-stream-id", mode=mode, _sink=sink) + try: + logger.start_trace(input="request") + logger.add_workflow_span(input="workflow") + logger.add_agent_span(input="agent") + logger.add_llm_span(input="prompt", output="answer", model="model") + logger.conclude(output="complete", conclude_all=True) + + llm, agent, workflow = sink.spans + assert operation_names(sink.spans) == ["chat", "invoke_agent", "invoke_workflow"] + assert llm.parent == agent.context + assert agent.parent == workflow.context + assert workflow.parent is None + assert sink.force_flush_calls == 0 + assert logger.current_parent() is None finally: logger.terminate() -def test_distributed_terminate_discards_unfinished_stubs() -> None: +def test_unsampled_remote_parent_is_preserved_and_not_exported() -> None: sink = RecordingSink() - logger = SplunkAOLogger( - project_id="project-id", - agent_stream_id="log-stream-id", - mode="distributed", - trace_id=str(uuid4()), - span_id=str(uuid4()), - _sink=sink, + token = context.attach( + extract_tracing_context({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"}) ) + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="log-stream-id", _sink=sink) + try: + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="work") + operation_context = logger._otel_ids[operation.id].span_context + logger.add_llm_span(input="prompt", output="answer", model="model") + logger.conclude(output="done") + logger.conclude(output="complete") - logger.terminate() - logger.terminate() + assert operation_context.is_valid + assert not bool(operation_context.trace_flags) + assert sink.spans == [] + finally: + logger.terminate() + context.detach(token) - assert sink.spans == [] - assert sink.force_flush_calls == 1 - assert sink.shutdown_calls == 1 - assert logger.current_parent() is None - assert logger._otel_ids == {} + +def test_custom_trace_and_span_id_constructor_arguments_are_removed() -> None: + with pytest.raises(TypeError, match="trace_id"): + SplunkAOLogger(trace_id="4bf92f35-77b3-4da6-a3ce-929d0e0e4736") + with pytest.raises(TypeError, match="span_id"): + SplunkAOLogger(span_id="00f067aa-0ba9-42b7-8000-000000000000") + + +def test_old_logger_header_method_is_removed() -> None: + assert not hasattr(SplunkAOLogger, "get_tracing_headers") diff --git a/tests/test_logger_otel_context.py b/tests/test_logger_otel_context.py index a9481b26..62a84d35 100644 --- a/tests/test_logger_otel_context.py +++ b/tests/test_logger_otel_context.py @@ -106,29 +106,24 @@ def test_completed_leaf_siblings_share_parent_and_do_not_become_active( logger.conclude(output="done") -def test_otel_identity_failure_does_not_interrupt_span_creation_or_legacy_streaming( +def test_otel_identity_failure_does_not_interrupt_span_creation( make_logger: Callable[[], SplunkAOLogger], monkeypatch: pytest.MonkeyPatch ) -> None: logger = make_logger() root = logger.start_trace(input="q") - ingest_step = Mock() warning = Mock() def fail_assignment(*args, **kwargs) -> SpanContext: raise RuntimeError("identity failure") - logger.mode = "distributed" with monkeypatch.context() as patch: patch.setattr(SplunkAOLogger, "_assign_otel_context", fail_assignment) - patch.setattr(SplunkAOLogger, "_ingest_step_streaming", ingest_step) patch.setattr(logger._logger, "warning", warning) span = logger.add_llm_span(input="prompt", output="answer", model="model") - logger.mode = "batch" assert span is not None assert span in root.spans assert span.id not in logger._otel_ids - ingest_step.assert_not_called() assert "Failed to assign OTel identity" in warning.call_args.args[0] logger.conclude(output="done") @@ -139,22 +134,17 @@ def test_otel_sync_failure_does_not_interrupt_parentable_span_creation( ) -> None: logger = make_logger() root = logger.start_trace(input="q") - ingest_step = Mock() warning = Mock() - logger.mode = "distributed" with monkeypatch.context() as patch: patch.setattr(SplunkAOLogger, "_sync_otel_context_impl", Mock(side_effect=RuntimeError("sync failure"))) - patch.setattr(SplunkAOLogger, "_ingest_step_streaming", ingest_step) patch.setattr(logger._logger, "warning", warning) workflow = logger.add_workflow_span(input="workflow") - logger.mode = "batch" assert workflow is not None assert workflow in root.spans assert logger.current_parent() is workflow assert workflow.id in logger._otel_ids - ingest_step.assert_not_called() assert "Failed to synchronize OTel context" in warning.call_args.args[0] logger._sync_otel_context(workflow) @@ -168,22 +158,17 @@ def test_otel_sync_and_release_failures_do_not_interrupt_hook_conclusion( logger = make_logger() root = logger.start_trace(input="q") workflow = logger.add_workflow_span(input="workflow") - update_step = Mock() warning = Mock() - logger.mode = "distributed" with monkeypatch.context() as patch: patch.setattr(SplunkAOLogger, "_sync_otel_context_impl", Mock(side_effect=RuntimeError("sync failure"))) patch.setattr(SplunkAOLogger, "_discard_otel_subtree", Mock(side_effect=RuntimeError("release failure"))) - patch.setattr(SplunkAOLogger, "_update_step_streaming", update_step) patch.setattr(logger._logger, "warning", warning) parent = logger.conclude(output="workflow-output") - logger.mode = "batch" assert parent is root assert logger.current_parent() is root assert workflow.output == "workflow-output" - update_step.assert_not_called() warning_messages = [call.args[0] for call in warning.call_args_list] assert any("Failed to synchronize OTel context" in message for message in warning_messages) assert any("Failed to release OTel context" in message for message in warning_messages) @@ -287,7 +272,7 @@ def test_start_trace_inherits_remote_parent_and_preserves_tracestate(make_logger context.detach(token) -def test_trace_flags_are_always_sampled_with_unsampled_remote_parent(make_logger: Callable[[], SplunkAOLogger]) -> None: +def test_trace_flags_preserve_unsampled_remote_parent(make_logger: Callable[[], SplunkAOLogger]) -> None: remote_context = propagate.extract({"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"}) token = context.attach(remote_context) try: @@ -295,8 +280,8 @@ def test_trace_flags_are_always_sampled_with_unsampled_remote_parent(make_logger root = logger.start_trace(input="q") child = logger.add_llm_span(input="prompt", output="answer", model="model") - assert logger._otel_ids[root.id].span_context.trace_flags == TraceFlags(TraceFlags.SAMPLED) - assert logger._otel_ids[child.id].span_context.trace_flags == TraceFlags(TraceFlags.SAMPLED) + assert logger._otel_ids[root.id].span_context.trace_flags == TraceFlags(TraceFlags.DEFAULT) + assert logger._otel_ids[child.id].span_context.trace_flags == TraceFlags(TraceFlags.DEFAULT) logger.conclude(output="a") finally: context.detach(token) diff --git a/tests/test_logger_otel_egress.py b/tests/test_logger_otel_egress.py index 2a3a976c..250dcc04 100644 --- a/tests/test_logger_otel_egress.py +++ b/tests/test_logger_otel_egress.py @@ -76,6 +76,23 @@ def test_complete_leaf_is_enqueued_before_flush(otlp_logger: SplunkAOLogger, rec assert recording_sink.force_flush_calls == 0 +def test_path1_session_is_captured_when_stable_span_identity_is_assigned( + otlp_logger: SplunkAOLogger, recording_sink: RecordingSink +) -> None: + otlp_logger.set_session("conversation-at-start") + otlp_logger.start_trace(input="question") + otlp_logger.add_workflow_span(input="work", name="workflow") + + otlp_logger.set_session("conversation-changed-later") + otlp_logger.conclude(output="done") + otlp_logger.conclude(output="trace done") + + [workflow_span] = recording_sink.spans + assert workflow_span.attributes["gen_ai.conversation.id"] == "conversation-at-start" + assert "splunk_ao.session.id" not in workflow_span.attributes + otlp_logger.clear_session() + + def test_single_llm_trace_emits_only_real_child(otlp_logger: SplunkAOLogger, recording_sink: RecordingSink) -> None: otlp_logger.add_single_llm_span_trace(input="question", output="answer", model="gpt-5") diff --git a/tests/test_middleware_tracing.py b/tests/test_middleware_tracing.py index 81ee7751..e6f5680a 100644 --- a/tests/test_middleware_tracing.py +++ b/tests/test_middleware_tracing.py @@ -1,288 +1,141 @@ -"""Tests for distributed tracing middleware.""" +"""Tests for W3C tracing middleware.""" -from unittest.mock import Mock, patch +from typing import Any +from unittest.mock import Mock import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from opentelemetry import trace +from opentelemetry.sdk.trace import ReadableSpan -from splunk_ao.constants.tracing import PARENT_ID_HEADER, TRACE_ID_HEADER -from splunk_ao.decorator import _parent_id_context, _trace_id_context +from splunk_ao import get_tracing_headers from splunk_ao.logger import SplunkAOLogger from splunk_ao.middleware import TracingMiddleware, get_request_logger -from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client -@pytest.fixture -def app(): - """Create a test FastAPI app with tracing middleware.""" - app = FastAPI() - app.add_middleware(TracingMiddleware) - - @app.get("/test") - async def test_endpoint(): - logger = get_request_logger() - return { - "trace_id": logger.trace_id, - "span_id": logger.span_id, - "is_logger": isinstance(logger, SplunkAOLogger), # Verify logger was created - } - - @app.get("/test-context") - async def test_context_endpoint(): - # Direct access to context variables - return {"trace_id": _trace_id_context.get(), "parent_id": _parent_id_context.get()} +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] - @app.get("/test-add-span") - async def test_add_span_endpoint(): - """Endpoint that tries to add a span to test distributed tracing stub creation. + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) - Note: In distributed tracing (distributed mode), when span_id is provided, - the logger creates local stub objects in __init__ without fetching from backend. - """ - try: - logger = get_request_logger() - # The validation should have happened during get_request_logger() call above - # If we get here without exception, validation passed (or didn't run) - # Try to add a span to confirm logger works - logger.add_workflow_span(input="test", name="test_span") - return { - "error": None, - "success": True, - "trace_id": str(logger.trace_id), - "span_id": str(logger.span_id) if logger.span_id else None, - } - except Exception as e: - return {"error": str(e), "error_type": type(e).__name__, "success": False} + def force_flush(self) -> bool: + return True - return app + def shutdown(self) -> None: + return None @pytest.fixture -def client(app): - """Create a test client.""" - return TestClient(app) - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_middleware_extracts_headers( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that middleware correctly extracts tracing headers.""" - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - # Use valid UUID4 values - trace_id = "6c4e3f7e-4a9a-4e7e-8c1f-3a9a3a9a3a9d" - parent_id = "6c4e3f7e-4a9a-4e7e-8c1f-3a9a3a9a3a9e" - - # Override get_span to return a span that belongs to the correct trace - # Also override get_trace to return a trace with the correct trace_id - import datetime - from unittest.mock import AsyncMock - from uuid import UUID - - now = datetime.datetime.now() - mock_span = { - "id": UUID(parent_id), - "trace_id": UUID(trace_id), # Same trace as header - "type": "workflow", - "name": "test-workflow-span", - "input": "test-input", - "output": None, - "created_at": now, - "updated_at": now, - "user_metadata": {}, - "metrics": {}, - "parent_id": None, - } - mock_trace = { - "id": UUID(trace_id), - "name": "test-trace", - "type": "trace", - "input": "test-input", - "output": None, - "created_at": now, - "updated_at": now, - "user_metadata": {}, - "spans": [], - } - mock_instance = mock_traces_client.return_value - mock_instance.get_span = AsyncMock(return_value=mock_span) - mock_instance.get_trace = AsyncMock(return_value=mock_trace) - - response = client.get("/test", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) - - assert response.status_code == 200 - data = response.json() - assert data["is_logger"] is True - assert data["trace_id"] == trace_id - assert data["span_id"] == parent_id - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_middleware_handles_missing_headers( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that middleware handles requests without tracing headers.""" - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - response = client.get("/test") - - assert response.status_code == 200 - data = response.json() - assert data["is_logger"] is True - assert data["trace_id"] is None - assert data["span_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_middleware_handles_partial_headers( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that middleware handles requests with only trace ID. - - Note: In normal operation, both TRACE_ID_HEADER and PARENT_ID_HEADER should be present. +def app_factory(monkeypatch: pytest.MonkeyPatch): + loggers: list[SplunkAOLogger] = [] + sinks: list[RecordingSink] = [] - However, when only TRACE_ID_HEADER is provided (no PARENT_ID_HEADER), the middleware - handles this by creating a logger with the trace_id but no span_id, - allowing the service to start a new root span within the existing trace. - """ - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) + def logger_factory() -> SplunkAOLogger: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + loggers.append(logger) + sinks.append(sink) + return logger - trace_id = "12345678-1234-4678-9abc-123456789abc" + monkeypatch.setattr("splunk_ao.middleware.tracing.SplunkAOLogger", logger_factory) - response = client.get("/test", headers={TRACE_ID_HEADER: trace_id}) + def make_app(*, fail: bool = False) -> FastAPI: + app = FastAPI() + app.add_middleware(TracingMiddleware) - assert response.status_code == 200 - data = response.json() - assert data["is_logger"] is True - assert data["trace_id"] == trace_id - assert data["span_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_context_cleanup_after_request( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that context variables are cleaned up after request.""" - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - trace_id = "12345678-1234-4678-9abc-123456789abc" - parent_id = "87654321-4321-4876-9cba-987654321cba" - - # First request with headers - response1 = client.get("/test-context", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) - assert response1.status_code == 200 - data1 = response1.json() - assert data1["trace_id"] == trace_id - assert data1["parent_id"] == parent_id - - # Second request without headers - should not see previous request's context - response2 = client.get("/test-context") - assert response2.status_code == 200 - data2 = response2.json() - assert data2["trace_id"] is None - assert data2["parent_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_get_request_logger_when_parent_id_equals_trace_id( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that get_request_logger handles the case when parent_id equals trace_id. - - When upstream services forward headers immediately after start_trace(), both - Splunk-AO-Trace-ID and Splunk-AO-Parent-ID are identical (the root trace id). - In this case, we should pass None as span_id to avoid SplunkAOLoggerException. - """ - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - trace_id = "12345678-1234-4678-9abc-123456789abc" - # parent_id equals trace_id (first hop after start_trace) - parent_id = trace_id - - response = client.get("/test", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) + @app.get("/test") + async def endpoint() -> dict[str, Any]: + logger = get_request_logger() + try: + logger.start_trace(input="request", name="request") + operation = logger.add_workflow_span(input="work", name="operation") + ids = logger._otel_ids[operation.id] + outbound = get_tracing_headers() + if fail: + raise RuntimeError("request failed") + logger.conclude(output="done") + logger.conclude(output="complete") + return { + "trace_id": format(ids.span_context.trace_id, "032x"), + "span_id": format(ids.span_context.span_id, "016x"), + "parent_span_id": ( + format(ids.parent_span_context.span_id, "016x") if ids.parent_span_context is not None else None + ), + "outbound": outbound, + } + finally: + logger.terminate() + + return app + + return make_app, loggers, sinks + + +def test_middleware_and_request_logger_continue_w3c_parent(app_factory) -> None: + make_app, _, sinks = app_factory + trace_id = "4bf92f3577b34da6a3ce929d0e0e4736" + parent_id = "00f067aa0ba902b7" + response = TestClient(make_app()).get( + "/test", headers={"traceparent": f"00-{trace_id}-{parent_id}-01", "tracestate": "vendor=value"} + ) assert response.status_code == 200 data = response.json() - - # When parent_id equals trace_id, span_id should be None (not the trace_id) - # This prevents SplunkAOLogger from trying to look up a span with the trace_id - assert data["is_logger"] is True assert data["trace_id"] == trace_id - assert data["span_id"] is None - - -@patch("splunk_ao.logger.logger.AgentStreams") -@patch("splunk_ao.logger.logger.Projects") -@patch("splunk_ao.logger.logger.Traces") -def test_mismatched_trace_and_span_ids( - mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, client: TestClient -): - """Test that get_request_logger creates stubs for distributed tracing. - - In distributed tracing (distributed mode), when trace_id and span_id are provided, - the logger creates local stub objects instead of fetching from the backend. - This allows distributed tracing to work without backend validation, with eventual - consistency handling any mismatches during ingestion retries. - """ - setup_mock_traces_client(mock_traces_client) - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - - trace_id = "12345678-1234-4678-9abc-123456789abc" - parent_id = "99999999-9999-4999-9999-999999999999" - - response = client.get("/test-add-span", headers={TRACE_ID_HEADER: trace_id, PARENT_ID_HEADER: parent_id}) + assert data["outbound"]["tracestate"] == "vendor=value" + [exported] = sinks[0].spans + assert exported.context.trace_id == int(trace_id, 16) + assert exported.parent.span_id == int(parent_id, 16) + + +def test_middleware_without_header_starts_new_trace_and_does_not_leak(app_factory) -> None: + make_app, _, sinks = app_factory + client = TestClient(make_app()) + first = client.get( + "/test", headers={"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"} + ).json() + second = client.get("/test").json() + + assert first["trace_id"] == "4bf92f3577b34da6a3ce929d0e0e4736" + assert second["trace_id"] != first["trace_id"] + [first_exported] = sinks[0].spans + [second_exported] = sinks[1].spans + assert first_exported.parent.span_id == int("00f067aa0ba902b7", 16) + assert second_exported.parent is None + + +def test_middleware_ignores_proprietary_headers(app_factory) -> None: + make_app, _, sinks = app_factory + response = TestClient(make_app()).get( + "/test", + headers={ + "Splunk-AO-Trace-ID": "4bf92f35-77b3-4da6-a3ce-929d0e0e4736", + "Splunk-AO-Parent-ID": "00f067aa-0ba9-42b7-8000-000000000000", + }, + ) assert response.status_code == 200 - data = response.json() + assert response.json()["trace_id"] != "4bf92f3577b34da6a3ce929d0e0e4736" + [exported] = sinks[0].spans + assert exported.parent is None - # With stubs, logger initialization succeeds without backend validation - # The stub objects are created locally, allowing distributed tracing to work - # even when the parent span hasn't been ingested yet (eventual consistency) - assert data["success"] is True, "Expected logger initialization to succeed with stubs" - assert data["error"] is None - assert data["trace_id"] == trace_id - assert data["span_id"] == parent_id +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 - setup_mock_projects_client(mock_projects_client) - setup_mock_logstreams_client(mock_logstreams_client) - # Test with completely invalid trace_id - should raise exception - with pytest.raises(SplunkAOLoggerException, match="Invalid trace_id"): - client.get("/test", headers={TRACE_ID_HEADER: "not-a-valid-uuid"}) +def test_get_request_logger_constructs_plain_request_scoped_logger(monkeypatch: pytest.MonkeyPatch) -> None: + constructor = Mock(return_value=Mock(spec=SplunkAOLogger)) + monkeypatch.setattr("splunk_ao.middleware.tracing.SplunkAOLogger", constructor) - # Test with valid trace_id but invalid parent_id - should raise exception - valid_trace_id = "12345678-1234-4678-9abc-123456789abc" - with pytest.raises(SplunkAOLoggerException, match="Invalid span_id"): - client.get("/test", headers={TRACE_ID_HEADER: valid_trace_id, PARENT_ID_HEADER: "invalid-parent"}) + assert get_request_logger() is constructor.return_value + constructor.assert_called_once_with() diff --git a/tests/test_o11y_config.py b/tests/test_o11y_config.py index f933458b..3a49cdee 100644 --- a/tests/test_o11y_config.py +++ b/tests/test_o11y_config.py @@ -7,13 +7,14 @@ from uuid import uuid4 import pytest +from opentelemetry.instrumentation.utils import is_http_instrumentation_enabled from pydantic import SecretStr from galileo_core.constants.request_method import RequestMethod from galileo_core.constants.routes import Routes from galileo_core.helpers.api_client import ApiClient from galileo_core.schemas.base_config import GalileoConfig -from splunk_ao.config import O11yApiClient, SplunkAOConfig +from splunk_ao.config import O11yApiClient, SplunkAOConfig, _ControlPlaneApiClient from splunk_ao.shared.exceptions import AmbiguousConfigurationError, MissingConfigurationError _CONFIG_ENV_VARS = ( @@ -85,7 +86,12 @@ def test_o11y_api_client_sync_request_preserves_prefix_and_header(monkeypatch: p async def fake_make_request( request_method: RequestMethod, base_url: str, endpoint: str, headers: dict[str, str], **kwargs: object ) -> dict: - captured.update(method=request_method, url=urljoin(base_url, endpoint), headers=headers) + captured.update( + method=request_method, + url=urljoin(base_url, endpoint), + headers=headers, + instrumentation_enabled=is_http_instrumentation_enabled(), + ) return {} monkeypatch.setattr(ApiClient, "make_request", staticmethod(fake_make_request)) @@ -95,7 +101,9 @@ async def fake_make_request( "method": RequestMethod.GET, "url": "https://app.lab0.observability.splunkcloud.com/ao/api/projects", "headers": {"accept": "application/json", "Content-Type": "application/json", "X-SF-Token": "tok"}, + "instrumentation_enabled": False, } + assert is_http_instrumentation_enabled() @pytest.mark.asyncio @@ -116,12 +124,13 @@ async def fake_make_request( def test_o11y_api_client_prefixes_streaming_requests(monkeypatch: pytest.MonkeyPatch) -> None: - captured: dict[str, str] = {} + captured: dict[str, object] = {} def fake_stream_request( self: ApiClient, method: RequestMethod, path: str, *args: object, **kwargs: object ) -> contextlib.AbstractContextManager[str]: captured["path"] = path + captured["instrumentation_enabled"] = is_http_instrumentation_enabled() return nullcontext(path) monkeypatch.setattr(ApiClient, "stream_request", fake_stream_request) @@ -129,6 +138,8 @@ def fake_stream_request( pass assert captured["path"] == "/ao/api/projects" + assert captured["instrumentation_enabled"] is False + assert is_http_instrumentation_enabled() @pytest.mark.parametrize("token_var", ["SPLUNK_AO_O11Y_TOKEN", "SPLUNK_AO_O11Y_API_TOKEN"]) @@ -253,17 +264,21 @@ def test_ambiguous_environment_fails_before_config_construction(monkeypatch: pyt def test_standalone_constructor_flow_still_uses_parent_validation(monkeypatch: pytest.MonkeyPatch) -> None: calls: list[object] = [] + instrumentation_states: list[bool] = [] async def fake_make_request(request_method: RequestMethod, base_url: str, endpoint: str, **kwargs: object) -> dict: calls.append(endpoint) + instrumentation_states.append(is_http_instrumentation_enabled()) return {"status": "ok"} def fake_get_jwt_token(*args: object, **kwargs: object) -> tuple[SecretStr, None]: calls.append("jwt") + instrumentation_states.append(is_http_instrumentation_enabled()) return SecretStr("jwt-token"), None def fake_request(self: ApiClient, method: RequestMethod, path: str, **kwargs: object) -> dict[str, str]: calls.append(path) + instrumentation_states.append(is_http_instrumentation_enabled()) return {"id": str(uuid4()), "email": "user@example.com", "role": "user"} monkeypatch.setattr(SplunkAOConfig, "_instance", None) @@ -275,6 +290,10 @@ def fake_request(self: ApiClient, method: RequestMethod, path: str, **kwargs: ob cfg = SplunkAOConfig.get(console_url="https://app.galileo.ai", api_key="key") assert not isinstance(cfg.validated_api_client, O11yApiClient) + assert isinstance(cfg.validated_api_client, _ControlPlaneApiClient) assert Routes.healthcheck in calls assert "jwt" in calls assert Routes.current_user in calls + assert instrumentation_states + assert not any(instrumentation_states) + assert is_http_instrumentation_enabled() diff --git a/tests/test_openai_agents.py b/tests/test_openai_agents.py index 16ed813e..4e00e485 100644 --- a/tests/test_openai_agents.py +++ b/tests/test_openai_agents.py @@ -16,10 +16,12 @@ set_trace_processors, ) from agents.tracing import ResponseSpanData +from opentelemetry.sdk.trace import ReadableSpan from pydantic import BaseModel from pytest import MonkeyPatch, mark from galileo_core.schemas.logging.span import LlmSpan, ToolSpan +from splunk_ao import get_tracing_headers from splunk_ao.handlers.openai_agents import SplunkAOTracingProcessor from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao.schema.handlers import Node @@ -27,6 +29,81 @@ from tests.testutils.setup import setup_mock_logstreams_client, setup_mock_projects_client, setup_mock_traces_client +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + self.force_flush_calls = 0 + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + self.force_flush_calls += 1 + return True + + def shutdown(self) -> None: + return None + + +def test_openai_agents_children_enqueue_before_trace_end() -> None: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + processor = SplunkAOTracingProcessor(splunk_ao_logger=logger, flush_on_trace_end=False) + trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) + root = Node( + node_type="workflow", + run_id="root-span-id", + parent_run_id="trace-id", + span_params={ + "input": "question", + "name": "agent root", + "start_time_iso": "2025-01-01T00:00:00+00:00", + "status_code": 200, + }, + ) + child = Node( + node_type="llm", + run_id="child-span-id", + parent_run_id="root-span-id", + span_params={ + "input": "prompt", + "name": "model call", + "start_time_iso": "2025-01-01T00:00:01+00:00", + "model": "model", + "status_code": 200, + }, + ) + try: + processor.on_trace_start(trace) + processor._nodes[str(root.run_id)] = root + processor._nodes[trace.trace_id].children.append(str(root.run_id)) + processor._start_owned_root(root) + processor._start_incremental_span(root) + processor._nodes[str(child.run_id)] = child + root.children.append(str(child.run_id)) + processor._start_incremental_span(child) + + child.span_params.update(output="answer", duration_ns=10) + processor._finish_incremental_span(child) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == ["chat"] + assert str(root.run_id) in processor._active_steps + assert sink.force_flush_calls == 0 + + root.span_params.update(output="answer", duration_ns=20) + processor._finish_incremental_span(root) + processor.on_trace_end(trace) + + assert [(span.attributes or {}).get("gen_ai.operation.name") for span in sink.spans] == [ + "chat", + "invoke_workflow", + ] + assert sink.spans[0].parent == sink.spans[1].context + assert sink.force_flush_calls == 0 + finally: + logger.terminate() + + class HomeworkOutput(BaseModel): is_homework: bool reasoning: str @@ -182,18 +259,13 @@ def test_commit_failure_concludes_handler_owned_trace( processor = SplunkAOTracingProcessor(splunk_ao_logger=logger) trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) processor.on_trace_start(trace) - owned_traces = [] - - def fail_after_starting_trace(*args, **kwargs): - processor._owned_trace = logger.add_trace(input="input", name="Trace") - owned_traces.append(processor._owned_trace) - raise RuntimeError("conversion failed") + owned_trace = processor._owned_trace - with patch.object(processor, "_log_node_tree", side_effect=fail_after_starting_trace): + with patch.object(processor, "_commit_trace", side_effect=RuntimeError("conversion failed")): processor.on_trace_end(trace) assert logger.current_parent() is None - assert owned_traces[0].status_code == 500 + assert owned_trace.status_code == 500 assert processor._nodes == {} assert processor._owned_trace is None @@ -222,6 +294,47 @@ def test_commit_failure_preserves_caller_owned_trace( logger.conclude(output="done") +@patch("splunk_ao.logger.logger.AgentStreams") +@patch("splunk_ao.logger.logger.Projects") +@patch("splunk_ao.logger.logger.Traces") +def test_openai_agents_live_root_is_used_for_outbound_context_and_commit( + mock_traces_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock +) -> None: + setup_mock_traces_client(mock_traces_client) + setup_mock_projects_client(mock_projects_client) + setup_mock_logstreams_client(mock_logstreams_client) + logger = SplunkAOLogger(project="test", agent_stream="test", ingestion_hook=lambda _: None) + processor = SplunkAOTracingProcessor(splunk_ao_logger=logger) + trace = MagicMock(trace_id="trace-id", name="Agent trace", metadata={}) + processor.on_trace_start(trace) + node = Node( + node_type="agent", + run_id="root-span-id", + parent_run_id="trace-id", + span_params={ + "input": "question", + "output": "answer", + "name": "Agent root", + "start_time_iso": "2025-01-01T00:00:00+00:00", + "duration_ns": 10, + "status_code": 200, + }, + ) + processor._nodes[str(node.run_id)] = node + processor._nodes[trace.trace_id].children.append(str(node.run_id)) + processor._start_owned_root(node) + live_root = processor._owned_root + live_context = logger._otel_ids[live_root.id].span_context + + headers = get_tracing_headers() + processor.on_trace_end(trace) + + assert headers["traceparent"].split("-")[2] == format(live_context.span_id, "016x") + assert logger.traces[0].spans == [live_root] + assert logger.current_parent() is None + logger.terminate() + + def _create_mock_response_with_tools(tool_calls: list[dict]) -> dict: """Create a mock OpenAI API response with embedded tool calls.""" return { diff --git a/tests/test_otel_native_paths.py b/tests/test_otel_native_paths.py index 4896ee69..bb493f03 100644 --- a/tests/test_otel_native_paths.py +++ b/tests/test_otel_native_paths.py @@ -3,14 +3,17 @@ from unittest.mock import MagicMock, patch import pytest -from opentelemetry import trace +from opentelemetry import baggage, context, trace from opentelemetry.sdk.resources import Resource from opentelemetry.sdk.trace import Event, ReadableSpan -from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult +from opentelemetry.sdk.trace import TracerProvider as SDKTracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter, SpanExportResult from opentelemetry.sdk.util.instrumentation import InstrumentationScope from opentelemetry.trace import Link, SpanContext, SpanKind, TraceFlags from opentelemetry.trace.status import Status, StatusCode +from galileo_core.schemas.logging.span import WorkflowSpan +from splunk_ao import get_tracing_headers from splunk_ao.decorator import ( _agent_stream_context, _dataset_input_context, @@ -21,11 +24,13 @@ _session_id_context, ) from splunk_ao.deployment import DeploymentMode, O11yConfig, StandaloneConfig +from splunk_ao.logger import SplunkAOLogger from splunk_ao.otel import ( _TRACE_PROVIDER_CONTEXT_VAR, SplunkAOOTLPExporter, SplunkAOSpanProcessor, add_splunk_ao_span_processor, + start_splunk_ao_span, ) from splunk_ao.shared.exceptions import MissingConfigurationError @@ -360,6 +365,20 @@ def test_processor_does_not_put_routing_on_span_attributes() -> None: processor.shutdown() +def test_processor_reads_inbound_standard_conversation_baggage() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(SpanProcessor=RecordingSpanProcessor, _exporter=exporter) + parent_context = baggage.set_baggage("gen_ai.conversation.id", "inbound-conversation", context.Context()) + span = MagicMock() + + processor.on_start(span, parent_context=parent_context) + + calls = {args[0]: args[1] for args, _ in span.set_attribute.call_args_list} + assert calls["gen_ai.conversation.id"] == "inbound-conversation" + assert "splunk_ao.session.id" not in calls + processor.shutdown() + + def test_processor_forwards_complete_routing_to_immutable_exporter() -> None: exporter = RecordingExporter() exporter_factory = MagicMock() @@ -451,3 +470,140 @@ def test_processor_construction_does_not_replace_global_provider() -> None: assert trace.get_tracer_provider() is global_before processor.shutdown() + + +def test_sdk_native_otel_span_injects_w3c_context() -> None: + provider = SDKTracerProvider() + _TRACE_PROVIDER_CONTEXT_VAR.set(provider) + try: + with start_splunk_ao_span(WorkflowSpan(input="request", name="operation")) as span: + headers = get_tracing_headers() + + parts = headers["traceparent"].split("-") + assert int(parts[1], 16) == span.get_span_context().trace_id + assert int(parts[2], 16) == span.get_span_context().span_id + finally: + provider.shutdown() + + +def test_caller_owned_otel_span_injects_w3c_context() -> None: + provider = SDKTracerProvider() + tracer = provider.get_tracer("caller-owned") + try: + with tracer.start_as_current_span("operation") as span: + headers = get_tracing_headers() + + parts = headers["traceparent"].split("-") + assert int(parts[1], 16) == span.get_span_context().trace_id + assert int(parts[2], 16) == span.get_span_context().span_id + finally: + provider.shutdown() + + +def test_caller_owned_otel_and_path1_logger_interoperate_bidirectionally() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(_exporter=exporter) + provider = SDKTracerProvider() + provider.add_span_processor(processor) + tracer = provider.get_tracer("caller-owned") + sink = MagicMock() + sink.force_flush.return_value = True + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + try: + with tracer.start_as_current_span("otel-upstream") as upstream: + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="agent work", name="path1-operation") + operation_context = logger._otel_ids[operation.id].span_context + + with tracer.start_as_current_span("otel-downstream") as downstream: + headers = get_tracing_headers() + + logger.conclude(output="done") + logger.conclude(output="complete") + + [path1_span] = [call.args[0] for call in sink.emit.call_args_list] + assert path1_span.parent == upstream.get_span_context() + assert path1_span.context == operation_context + assert downstream.parent == operation_context + assert headers["traceparent"].split("-")[2] == format(downstream.get_span_context().span_id, "016x") + finally: + logger.terminate() + provider.shutdown() + + by_name = {span.name: span for span in exporter.spans} + assert by_name["otel-downstream"].parent == operation_context + + +def test_explicit_conversation_session_is_applied_across_paths_1_2_and_3() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(SpanProcessor=RecordingSpanProcessor, _exporter=exporter) + provider = SDKTracerProvider() + provider.add_span_processor(processor) + provider_token = _TRACE_PROVIDER_CONTEXT_VAR.set(provider) + tracer = provider.get_tracer("caller-owned") + sink = MagicMock() + sink.force_flush.return_value = True + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + try: + logger.set_session("conversation-all-paths") + logger.start_trace(input="request") + path1_operation = logger.add_workflow_span(input="path 1", name="path-1") + + with start_splunk_ao_span(WorkflowSpan(input="path 2", name="path-2")): + pass + with tracer.start_as_current_span("path-3"): + pass + + logger.conclude(output="done") + logger.conclude(output="complete") + + [path1_span] = [call.args[0] for call in sink.emit.call_args_list] + assert path1_span.name == "invoke_workflow path-1" + assert path1_span.attributes["gen_ai.conversation.id"] == "conversation-all-paths" + assert "splunk_ao.session.id" not in path1_span.attributes + + recording_processor = processor.processor + assert isinstance(recording_processor, RecordingSpanProcessor) + started = {span.name: span for span, _ in recording_processor.started} + assert started["path-2"].attributes["gen_ai.conversation.id"] == "conversation-all-paths" + assert started["path-3"].attributes["gen_ai.conversation.id"] == "conversation-all-paths" + assert "splunk_ao.session.id" not in started["path-2"].attributes + assert "splunk_ao.session.id" not in started["path-3"].attributes + assert path1_operation.id not in logger._otel_ids + finally: + logger.clear_session() + logger.terminate() + _TRACE_PROVIDER_CONTEXT_VAR.reset(provider_token) + provider.shutdown() + + +def test_sdk_native_and_caller_owned_spans_keep_topology_and_queue_without_flush() -> None: + exporter = RecordingExporter() + processor = SplunkAOSpanProcessor(_exporter=exporter) + provider = SDKTracerProvider() + provider.add_span_processor(processor) + provider_token = _TRACE_PROVIDER_CONTEXT_VAR.set(provider) + tracer = provider.get_tracer("caller-owned") + try: + assert isinstance(processor.processor, BatchSpanProcessor) + with ( + patch.object(processor, "on_end", wraps=processor.on_end) as on_end, + patch.object(processor.processor, "force_flush", wraps=processor.processor.force_flush) as force_flush, + ): + operation = WorkflowSpan(input="request", name="sdk-operation") + with start_splunk_ao_span(operation) as sdk_span: + with tracer.start_as_current_span("caller-child") as caller_span: + pass + + assert on_end.call_count == 2 + force_flush.assert_not_called() + assert caller_span.parent is not None + assert caller_span.parent.span_id == sdk_span.get_span_context().span_id + assert caller_span.get_span_context().trace_id == sdk_span.get_span_context().trace_id + finally: + _TRACE_PROVIDER_CONTEXT_VAR.reset(provider_token) + provider.shutdown() + + by_name = {span.name: span for span in exporter.spans} + assert set(by_name) == {"sdk-operation", "caller-child"} + assert by_name["caller-child"].parent == by_name["sdk-operation"].context diff --git a/tests/test_thread_pool_task_handler.py b/tests/test_thread_pool_task_handler.py deleted file mode 100644 index e5080bf4..00000000 --- a/tests/test_thread_pool_task_handler.py +++ /dev/null @@ -1,352 +0,0 @@ -from concurrent.futures import Future -from unittest.mock import Mock, patch - -import pytest -from src.splunk_ao.logger.task_handler import ThreadPoolTaskHandler - - -class TestThreadPoolTaskHandler: - """Test suite for ThreadPoolTaskHandler class.""" - - @pytest.fixture - def mock_pool(self): - """Mock EventLoopThreadPool.""" - with patch("src.splunk_ao.logger.task_handler.EventLoopThreadPool") as mock_pool_class: - mock_pool = Mock() - mock_pool_class.return_value = mock_pool - yield mock_pool - - @pytest.fixture - def handler(self, mock_pool): - """Create ThreadPoolTaskHandler instance with mocked pool.""" - return ThreadPoolTaskHandler(num_threads=2) - - @pytest.fixture - def mock_future(self): - """Create a mock Future object.""" - future = Mock(spec=Future) - future.done.return_value = False - future.result.return_value = "test_result" - future.add_done_callback = Mock() - return future - - async def dummy_async_func(self) -> str: - """Dummy async function for testing.""" - return "async_result" - - def test_init(self, mock_pool) -> None: - """Test ThreadPoolTaskHandler initialization.""" - handler = ThreadPoolTaskHandler(num_threads=4) - assert handler._tasks == {} - assert handler._retry_counts == {} - assert handler._pool is not None - - def test_submit_task_independent(self, handler, mock_pool, mock_future) -> None: - """Test submitting an independent task (dependent_on_prev=False).""" - mock_pool.submit.return_value = mock_future - - async_fn = self.dummy_async_func - task_id = "test_task_1" - - handler.submit_task(task_id, async_fn, dependent_on_prev=False) - - # Verify task was submitted to pool - mock_pool.submit.assert_called_once_with(async_fn, wait_for_result=False) - - # Verify done callback was added - mock_future.add_done_callback.assert_called_once() - - # Verify task was tracked - assert task_id in handler._tasks - assert handler._tasks[task_id]["future"] == mock_future - assert handler._tasks[task_id]["parent_task_id"] is None - assert handler._retry_counts[task_id] == 0 - - def test_submit_task_dependent_no_previous(self, handler, mock_pool, mock_future) -> None: - """Test submitting a dependent task when no previous task exists.""" - mock_pool.submit.return_value = mock_future - - async_fn = self.dummy_async_func - task_id = "test_task_1" - - handler.submit_task(task_id, async_fn, dependent_on_prev=True) - - # Should submit immediately since no previous task exists - mock_pool.submit.assert_called_once_with(async_fn, wait_for_result=False) - assert task_id in handler._tasks - assert handler._tasks[task_id]["future"] == mock_future - assert handler._tasks[task_id]["parent_task_id"] is None # No parent since no previous tasks - - def test_submit_task_dependent_previous_completed(self, handler, mock_pool, mock_future) -> None: - """Test submitting a dependent task when previous task is completed.""" - # Set up first task - mock_pool.submit.return_value = mock_future - first_task_id = "task_1" - handler.submit_task(first_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Mock the first task as completed - with patch.object(handler, "get_status", return_value="completed"): - second_mock_future = Mock(spec=Future) - mock_pool.submit.return_value = second_mock_future - - second_task_id = "task_2" - handler.submit_task(second_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Should submit immediately since previous task is completed - assert mock_pool.submit.call_count == 2 - assert second_task_id in handler._tasks - assert handler._tasks[second_task_id]["future"] == second_mock_future - - def test_submit_task_dependent_previous_running(self, handler, mock_pool, mock_future) -> None: - """Test submitting a dependent task when previous task is still running.""" - # Set up first task - mock_pool.submit.return_value = mock_future - first_task_id = "task_1" - handler.submit_task(first_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Mock the first task as still running - with patch.object(handler, "get_status", return_value="running"): - second_task_id = "task_2" - handler.submit_task(second_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Should not submit yet, but should be tracked as pending - assert mock_pool.submit.call_count == 1 # Only first task submitted - assert second_task_id in handler._tasks - assert handler._tasks[second_task_id]["future"] is None - assert handler._tasks[second_task_id]["parent_task_id"] == first_task_id - assert handler._tasks[second_task_id]["callback"] is not None - - def test_handle_task_completion_triggers_children(self, handler, mock_pool) -> None: - """Test that task completion triggers child tasks.""" - # Set up parent task - parent_future = Mock(spec=Future) - mock_pool.submit.return_value = parent_future - parent_task_id = "parent_task" - handler.submit_task(parent_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Set up child task (dependent) - with patch.object(handler, "get_status", return_value="running"): - child_task_id = "child_task" - handler.submit_task(child_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Mock child task submission - child_future = Mock(spec=Future) - mock_pool.submit.return_value = child_future - - # Trigger parent task completion - handler._handle_task_completion(parent_task_id) - - # Verify child task was submitted - assert mock_pool.submit.call_count == 2 # Parent + child - assert handler._tasks[child_task_id]["future"] == child_future - - def test_increment_retry(self, handler) -> None: - """Test retry count increment.""" - task_id = "test_task" - - # Initially 0 - assert handler.get_retry_count(task_id) == 0 - - # Increment retry - handler.increment_retry(task_id) - assert handler.get_retry_count(task_id) == 1 - - # Increment again - handler.increment_retry(task_id) - assert handler.get_retry_count(task_id) == 2 - - def test_get_status_not_found(self, handler) -> None: - """Test get_status for non-existent task.""" - assert handler.get_status("non_existent") == "not_found" - - def test_get_status_pending(self, handler, mock_pool, mock_future) -> None: - """Test get_status for pending task.""" - # Set up parent task - mock_pool.submit.return_value = mock_future - parent_task_id = "parent" - handler.submit_task(parent_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Set up child task as pending - with patch.object( - handler, "get_status", side_effect=lambda tid: "running" if tid == parent_task_id else "pending" - ): - child_task_id = "child" - handler.submit_task(child_task_id, self.dummy_async_func, dependent_on_prev=True) - - # Reset side_effect and test actual status - handler.get_status = handler.__class__.get_status.__get__(handler) - assert handler.get_status(child_task_id) == "pending" - - def test_get_status_running(self, handler, mock_pool, mock_future) -> None: - """Test get_status for running task.""" - mock_pool.submit.return_value = mock_future - mock_future.done.return_value = False - - task_id = "running_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - assert handler.get_status(task_id) == "running" - - def test_get_status_completed(self, handler, mock_pool, mock_future) -> None: - """Test get_status for completed task.""" - mock_pool.submit.return_value = mock_future - mock_future.done.return_value = True - mock_future.result.return_value = "success" - - task_id = "completed_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - assert handler.get_status(task_id) == "completed" - - def test_get_status_failed(self, handler, mock_pool, mock_future) -> None: - """Test get_status for failed task.""" - mock_pool.submit.return_value = mock_future - mock_future.done.return_value = True - mock_future.result.side_effect = Exception("Task failed") - - task_id = "failed_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - assert handler.get_status(task_id) == "failed" - - def test_get_result(self, handler, mock_pool, mock_future) -> None: - """Test getting task result.""" - mock_pool.submit.return_value = mock_future - expected_result = "test_result" - mock_future.result.return_value = expected_result - - task_id = "test_task" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - result = handler.get_result(task_id) - assert result == expected_result - - def test_get_result_not_found(self, handler) -> None: - """Test getting result for non-existent task.""" - with pytest.raises(ValueError, match="Task non_existent not found"): - handler.get_result("non_existent") - - def test_get_children(self, handler, mock_pool, mock_future) -> None: - """Test getting child tasks.""" - # Set up parent task - mock_pool.submit.return_value = mock_future - parent_task_id = "parent" - handler.submit_task(parent_task_id, self.dummy_async_func, dependent_on_prev=False) - - # Manually add two child tasks that depend on the parent - # (This simulates what would happen if we had a more complex dependency system) - def child1_callback(): - return handler._pool.submit(self.dummy_async_func, wait_for_result=False) - - def child2_callback(): - return handler._pool.submit(self.dummy_async_func, wait_for_result=False) - - handler._add_or_update_task( - task_id="child1", future=None, start_time=None, parent_task_id=parent_task_id, callback=child1_callback - ) - - handler._add_or_update_task( - task_id="child2", future=None, start_time=None, parent_task_id=parent_task_id, callback=child2_callback - ) - - children = handler.get_children(parent_task_id) - assert len(children) == 2 - - # Verify both children have the correct parent_task_id - child_task_ids = [ - task_id for task_id, task in handler._tasks.items() if task.get("parent_task_id") == parent_task_id - ] - assert len(child_task_ids) == 2 - assert "child1" in child_task_ids - assert "child2" in child_task_ids - - def test_all_tasks_completed_empty(self, handler) -> None: - """Test all_tasks_completed with no tasks.""" - assert handler.all_tasks_completed() is True - - def test_all_tasks_completed_true(self, handler, mock_pool, mock_future) -> None: - """Test all_tasks_completed when all tasks are done.""" - mock_pool.submit.return_value = mock_future - - with patch.object(handler, "get_status", return_value="completed"): - handler.submit_task("task1", self.dummy_async_func, dependent_on_prev=False) - handler.submit_task("task2", self.dummy_async_func, dependent_on_prev=False) - - assert handler.all_tasks_completed() is True - - def test_all_tasks_completed_false(self, handler, mock_pool, mock_future) -> None: - """Test all_tasks_completed when some tasks are still running.""" - mock_pool.submit.return_value = mock_future - - def mock_status(task_id) -> str: - return "completed" if task_id == "task1" else "running" - - with patch.object(handler, "get_status", side_effect=mock_status): - handler.submit_task("task1", self.dummy_async_func, dependent_on_prev=False) - handler.submit_task("task2", self.dummy_async_func, dependent_on_prev=False) - - assert handler.all_tasks_completed() is False - - def test_terminate(self, handler, mock_pool) -> None: - """Test handler termination.""" - handler.terminate() - mock_pool.stop.assert_called_once() - - def test_complex_dependency_chain(self, handler, mock_pool) -> None: - """Test a complex chain of dependent tasks.""" - # Create multiple futures for different tasks - futures = [Mock(spec=Future) for _ in range(4)] - future_index = 0 - - def get_next_future(*args, **kwargs): - nonlocal future_index - future = futures[future_index] - future_index += 1 - return future - - mock_pool.submit.side_effect = get_next_future - - # Submit task chain: A -> B -> C -> D - task_ids = ["task_A", "task_B", "task_C", "task_D"] - - # Submit first task (independent) - handler.submit_task(task_ids[0], self.dummy_async_func, dependent_on_prev=False) - - # Submit dependent tasks - for i in range(1, len(task_ids)): - with patch.object(handler, "get_status", return_value="running"): - handler.submit_task(task_ids[i], self.dummy_async_func, dependent_on_prev=True) - - # Verify only first task was submitted initially - assert mock_pool.submit.call_count == 1 - - # Simulate completion of each task in sequence - for i in range(len(task_ids) - 1): - # Mock the current task as completed for next submission - mock_pool.submit.side_effect = get_next_future - handler._handle_task_completion(task_ids[i]) - - # Verify next task was submitted - expected_calls = i + 2 # +1 for initial task, +1 for current completion - assert mock_pool.submit.call_count == expected_calls - - def test_done_callback_integration(self, handler, mock_pool) -> None: - """Test that the done callback properly triggers task completion handling.""" - mock_future = Mock(spec=Future) - mock_pool.submit.return_value = mock_future - - task_id = "callback_test" - handler.submit_task(task_id, self.dummy_async_func, dependent_on_prev=False) - - # Get the callback that was registered - callback_calls = mock_future.add_done_callback.call_args_list - assert len(callback_calls) == 1 - registered_callback = callback_calls[0][0][0] - - # Mock _handle_task_completion to verify it's called - with patch.object(handler, "_handle_task_completion") as mock_handle: - # Trigger the callback - registered_callback(mock_future) - - # Verify _handle_task_completion was called with correct task_id - mock_handle.assert_called_once_with(task_id) diff --git a/tests/test_tracing.py b/tests/test_tracing.py new file mode 100644 index 00000000..473ee7dc --- /dev/null +++ b/tests/test_tracing.py @@ -0,0 +1,218 @@ +from collections.abc import Generator + +import pytest +from opentelemetry import baggage, context, trace +from opentelemetry.sdk.trace import ReadableSpan +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags, TraceState + +from splunk_ao import extract_tracing_context, get_tracing_headers +from splunk_ao.exceptions import SplunkAOLoggerException +from splunk_ao.logger import SplunkAOLogger +from splunk_ao.logger.logger import _otel_context_state +from splunk_ao.session_context import GEN_AI_CONVERSATION_ID, _session_id_context + + +class RecordingSink: + def __init__(self) -> None: + self.spans: list[ReadableSpan] = [] + + def emit(self, span: ReadableSpan) -> None: + self.spans.append(span) + + def force_flush(self) -> bool: + return True + + def shutdown(self) -> None: + return None + + +@pytest.fixture(autouse=True) +def isolated_context() -> Generator[None, None, None]: + state_token = _otel_context_state.set(None) + session_token = _session_id_context.set(None) + token = context.attach(context.Context()) + try: + yield + finally: + context.detach(token) + _session_id_context.reset(session_token) + _otel_context_state.reset(state_token) + + +def make_logger() -> tuple[SplunkAOLogger, RecordingSink]: + sink = RecordingSink() + logger = SplunkAOLogger(project_id="project-id", agent_stream_id="stream-id", _sink=sink) + return logger, sink + + +def attach_span_context(span_context: SpanContext): + return context.attach(trace.set_span_in_context(NonRecordingSpan(span_context))) + + +def test_get_tracing_headers_injects_native_otel_context_into_same_carrier() -> None: + span_context = SpanContext( + trace_id=0x4BF92F3577B34DA6A3CE929D0E0E4736, + span_id=0x00F067AA0BA902B7, + is_remote=False, + trace_flags=TraceFlags.SAMPLED, + trace_state=TraceState([("vendor", "value")]), + ) + token = attach_span_context(span_context) + try: + carrier = {"existing": "header"} + result = get_tracing_headers(carrier) + finally: + context.detach(token) + + assert result is carrier + assert carrier["existing"] == "header" + assert carrier["traceparent"] == "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01" + assert carrier["tracestate"] == "vendor=value" + + +def test_get_tracing_headers_rejects_no_active_operation() -> None: + with pytest.raises(SplunkAOLoggerException, match="active exportable operation"): + get_tracing_headers() + + +def test_get_tracing_headers_rejects_internal_trace_envelope() -> None: + logger, _ = make_logger() + try: + logger.start_trace(input="request") + with pytest.raises(SplunkAOLoggerException, match="active exportable operation"): + get_tracing_headers() + logger.conclude(output="done") + finally: + logger.terminate() + + +def test_get_tracing_headers_uses_real_path1_operation_and_no_routing_headers() -> None: + logger, _ = make_logger() + try: + logger.start_trace(input="request") + operation = logger.add_workflow_span(input="work", name="operation") + operation_context = logger._otel_ids[operation.id].span_context + + headers = get_tracing_headers() + + assert headers["traceparent"].split("-")[2] == format(operation_context.span_id, "016x") + assert set(headers).isdisjoint( + {"project", "projectid", "logstream", "logstreamid", "experimentid", "X-SF-Token", "Splunk-AO-API-Key"} + ) + logger.conclude(output="work") + logger.conclude(output="done") + finally: + logger.terminate() + + +def test_get_tracing_headers_injects_only_standard_conversation_baggage() -> None: + logger, _ = make_logger() + try: + logger.set_session("conversation-123") + logger.start_trace(input="request") + logger.add_workflow_span(input="work", name="operation") + + headers = get_tracing_headers() + + extracted = extract_tracing_context(headers) + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-123" + assert "splunk_ao.session.id" not in headers.get("baggage", "") + assert all( + private_name not in headers.get("baggage", "") + for private_name in ( + "project", + "agent_stream", + "logstream", + "experiment", + "agent_name", + "workflow_name", + "model_name", + "token", + ) + ) + logger.conclude(output="work") + logger.conclude(output="done") + finally: + logger.terminate() + + +def test_extract_tracing_context_restores_conversation_baggage() -> None: + extracted = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=conversation-456,unrelated=value", + } + ) + + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "conversation-456" + assert baggage.get_baggage("unrelated", context=extracted) == "value" + + +@pytest.mark.parametrize("header", ["gen_ai.conversation.id=", "gen_ai.conversation.id"]) +def test_extract_tracing_context_ignores_malformed_or_empty_conversation_baggage(header: str) -> None: + extracted = extract_tracing_context( + {"traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", "baggage": header} + ) + + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) is None + + +def test_local_session_overrides_inbound_conversation_when_injecting() -> None: + inbound = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=inbound-session,unrelated=value", + } + ) + token = context.attach(inbound) + _session_id_context.set("local-session") + try: + headers = get_tracing_headers() + finally: + context.detach(token) + + extracted = extract_tracing_context(headers) + assert baggage.get_baggage(GEN_AI_CONVERSATION_ID, context=extracted) == "local-session" + assert baggage.get_baggage("unrelated", context=extracted) == "value" + + +def test_inbound_conversation_overrides_compatibility_logger_field_for_path1() -> None: + remote = extract_tracing_context( + { + "traceparent": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", + "baggage": "gen_ai.conversation.id=inbound-session", + } + ) + token = context.attach(remote) + logger, sink = make_logger() + logger.session_id = "stale-logger-session" + try: + logger.start_trace(input="request") + logger.add_workflow_span(input="work", name="operation") + logger.conclude(output="work") + logger.conclude(output="done") + + [operation] = sink.spans + assert operation.attributes[GEN_AI_CONVERSATION_ID] == "inbound-session" + finally: + logger.terminate() + context.detach(token) + + +def test_extract_tracing_context_is_case_insensitive_and_preserves_tracestate() -> None: + extracted = extract_tracing_context( + {"TRACEPARENT": "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01", "TRACESTATE": "vendor=value"} + ) + span_context = trace.get_current_span(extracted).get_span_context() + + assert span_context.is_valid + assert span_context.trace_id == 0x4BF92F3577B34DA6A3CE929D0E0E4736 + assert span_context.span_id == 0x00F067AA0BA902B7 + assert span_context.is_remote + assert span_context.trace_state.get("vendor") == "value" + + +@pytest.mark.parametrize("carrier", [{}, {"traceparent": "not-valid"}]) +def test_extract_tracing_context_rejects_missing_or_malformed_parent(carrier: dict[str, str]) -> None: + extracted = extract_tracing_context(carrier) + assert not trace.get_current_span(extracted).get_span_context().is_valid diff --git a/tests/testutils/setup.py b/tests/testutils/setup.py index 2866591e..4b32098f 100644 --- a/tests/testutils/setup.py +++ b/tests/testutils/setup.py @@ -1,15 +1,8 @@ -import copy import datetime -from collections.abc import Callable -from contextlib import contextmanager -from typing import Any from unittest.mock import AsyncMock, Mock from uuid import UUID -from pydantic import BaseModel - from splunk_ao.agent_streams import AgentStream -from splunk_ao.logger.logger import SplunkAOLogger from splunk_ao.projects import Project from splunk_ao.resources.models import ExperimentResponse, ProjectType from splunk_ao.resources.models.log_stream_response import LogStreamResponse @@ -17,171 +10,6 @@ from splunk_ao.resources.models.task_type import TaskType -class ThreadPoolTaskInfo(BaseModel): - task_id: str - function_name: str - request: Any - task_func: Callable - kwargs: dict - - -class ThreadPoolRequestCapture: - """Helper class to capture requests submitted to the thread pool in streaming methods.""" - - def __init__(self): - self.captured_tasks: list[ThreadPoolTaskInfo] = [] - self.mock_pool = None - - def capture_request(self, task_id: str, async_fn: Callable, **kwargs) -> None: - """Capture function that extracts requests and function names from ThreadPoolTaskHandler.submit_task calls.""" - - # Extract both the request and function from the lambda - captured_request = None - captured_function_name = None - captured_function = None - - # Get the first function name referenced in the lambda - # Since lambdas follow pattern: lambda: function_name(request) - code = async_fn.__code__ - - if code.co_names: - # The first (and typically only) name in co_names is our function - captured_function_name = code.co_names[0] - - # Extract request from closure and also look for callable functions - if async_fn.__closure__: - for cell in async_fn.__closure__: - cell_content = cell.cell_contents - - # Check for request object - if isinstance(cell_content, BaseModel): - captured_request = copy.deepcopy(cell_content) - - # Check for callable function (if we didn't get function name from co_names) - elif captured_function_name is None and callable(cell_content) and hasattr(cell_content, "__name__"): - captured_function_name = cell_content.__name__ - captured_function = cell_content - - # If we have function name but not function reference, look for it - elif ( - captured_function_name - and callable(cell_content) - and hasattr(cell_content, "__name__") - and cell_content.__name__ == captured_function_name - ): - captured_function = cell_content - - # Create a new lambda that uses the deep-copied request to avoid mutation issues - if captured_function and captured_request: - - def isolated_task_func(): - return captured_function(captured_request) - else: - isolated_task_func = async_fn - - task_info = ThreadPoolTaskInfo( - task_id=task_id, - function_name=captured_function_name, - request=captured_request, - task_func=isolated_task_func, - kwargs=kwargs, - ) - self.captured_tasks.append(task_info) - - def get_request_by_function_name(self, function_name: str) -> Any | None: - return next((task.request for task in self.captured_tasks if task.function_name == function_name), None) - - def get_latest_task(self) -> ThreadPoolTaskInfo | None: - """Get the most recent task info (request + function name).""" - return self.captured_tasks[-1] if self.captured_tasks else None - - def get_all_tasks(self) -> list[ThreadPoolTaskInfo]: - """ - Get all captured tasks. - - Returns: - list[ThreadPoolTaskInfo]: List of captured tasks. - """ - return self.captured_tasks - - def get_all_requests(self) -> list[Any]: - """Get all captured requests.""" - return [task.request for task in self.captured_tasks] - - def get_all_function_names(self) -> list[str]: - """Get all captured function names.""" - return [task.function_name for task in self.captured_tasks] - - def get_task_by_function_name(self, function_name: str) -> ThreadPoolTaskInfo | None: - """Get the first task info that matches the given function name. - - Args: - function_name (str): The name of the function to search for. - - Returns: - Optional[ThreadPoolTaskInfo]: The first task info that matches the given function name, or None if no match is found. - """ - return next((task for task in self.captured_tasks if task.function_name == function_name), None) - - def count_function_calls(self, function_name: str) -> int: - """Count how many times a particular function was called.""" - functions = self.get_all_function_names() - return sum(1 for func_name in functions if func_name == function_name) - - def clear(self) -> None: - """Clear all captured tasks.""" - self.captured_tasks.clear() - - def assert_function_called(self, expected_function_name: str) -> None: - """Assert that the latest captured function matches the expected name.""" - assert expected_function_name in self.get_all_function_names(), ( - f"Expected function '{expected_function_name}' not in '{self.get_all_function_names()}'" - ) - - def assert_functions_called(self, expected_function_names: list[str]) -> None: - """Assert that the captured functions match the expected sequence.""" - actual_functions = self.get_all_function_names() - assert actual_functions == expected_function_names, ( - f"Expected functions {expected_function_names}, but got {actual_functions}" - ) - - -def setup_thread_pool_request_capture(logger: SplunkAOLogger) -> ThreadPoolRequestCapture: - """ - Set up request capture for a logger's thread pool. - - Args: - logger: SplunkAOLogger instance to mock the thread pool for - - Returns: - ThreadPoolRequestCapture instance that can be used to inspect captured requests - """ - capture = ThreadPoolRequestCapture() - mock_pool = Mock(side_effect=capture.capture_request) - logger._task_handler.submit_task = mock_pool - logger._task_handler.submit_task_with_parent = mock_pool - capture.mock_pool = mock_pool - return capture - - -@contextmanager -def capture_streaming_requests(logger): - """ - Context manager for capturing streaming requests. - - Usage: - with capture_streaming_requests(logger) as capture: - logger._ingest_trace_streaming(trace) - request = capture.get_latest_request() - assert isinstance(request, TracesIngestRequest) - """ - capture = setup_thread_pool_request_capture(logger) - try: - yield capture - finally: - capture.clear() - - def setup_mock_projects_client(mock_projects_client: Mock): now = datetime.datetime.now() mock_instance = mock_projects_client.return_value