LCORE-2981: Add OpenTelemetry instrumentation for POST /v1/streaming_query endpoint - #2415
LCORE-2981: Add OpenTelemetry instrumentation for POST /v1/streaming_query endpoint#2415anik120 wants to merge 1 commit into
Conversation
WalkthroughThe streaming query endpoint now creates an OpenTelemetry root span, records anonymized request data, and propagates the span through standard, compaction, and agent response streaming. Completion, error, cancellation, and validation paths now close or annotate spans. ChangesStreaming OpenTelemetry instrumentation
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant streaming_query
participant generate_response_with_compaction
participant generate_agent_response
participant OpenTelemetryExporter
Client->>streaming_query: Submit streaming query
streaming_query->>OpenTelemetryExporter: Create root span and record request attributes
streaming_query->>generate_response_with_compaction: Pass root span
generate_response_with_compaction->>generate_agent_response: Pass root span
generate_agent_response->>OpenTelemetryExporter: Record completion attributes and events
generate_agent_response->>OpenTelemetryExporter: End root span
generate_response_with_compaction-->>Client: Stream response or SSE error
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 7✅ Passed checks (7 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/utils/agents/streaming.py (1)
251-257: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winThe span leaks when the consumer closes the generator.
root_span.end()runs only on three explicit code paths. If the client disconnects mid-stream, Starlette callsaclose()on this async generator.GeneratorExitis raised at the currentyieldinside thetry, so onlyderegister_streamruns. Execution never reaches lines 254-256, 285-286, or 338, and the span is never ended. An unended span is never exported, so traces are lost for exactly the aborted requests you want to inspect.Guard the whole body with a single
try/finally. That also removes the three duplicatedend()calls and makes the double-end on the compaction path insrc/app/endpoints/streaming_query.py(lines 503-505) harmless to reason about.🛡️ Sketch: single owner for span termination
media_type = context.query_request.media_type or MEDIA_TYPE_JSON + span_ended = False + try: ... - if not stream_completed: - if root_span is not None: - root_span.end() - return + if not stream_completed: + return + ... + finally: + if root_span is not None and not span_ended: + root_span.end() + span_ended = TrueAn
async withhelper orcontextlib.AsyncExitStackalso works and keeps the indentation flat.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/utils/agents/streaming.py` around lines 251 - 257, Wrap the entire async generator body in a single try/finally that owns root_span termination, including all yield and early-return paths. Move deregister_stream and root_span.end into that finalizer, remove the duplicated end calls from the explicit completion, error, and compaction paths, and preserve existing stream behavior while ensuring GeneratorExit from aclose() ends the span.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/app/endpoints/streaming_query.py`:
- Around line 492-505: Update generate_agent_response so it no longer ends
root_span on normal completion or handled error paths, leaving span termination
solely to generate_response_with_compaction’s existing finally block. Preserve
the delegated generator’s span attribute recording and ensure all paths still
allow the wrapper to end the span exactly once.
In `@tests/unit/app/endpoints/test_streaming_query.py`:
- Around line 877-906: Add a non-vacuous assertion in
test_child_spans_nested_under_root that at least one child span is emitted
before validating each child’s parent, or configure one of the
_setup_common_mocks collaborators to create a child span. Preserve the existing
requirement that every emitted child span is parented to the
streaming_query.handle_request root span.
In `@tests/unit/utils/agents/test_streaming.py`:
- Around line 823-834: Move the duplicated otel_fixture definition into a shared
conftest.py. Remove the class-local otel_fixture from
tests/unit/utils/agents/test_streaming.py lines 823-834 and
tests/unit/app/endpoints/test_streaming_query.py lines 589-600, ensuring both
test suites consume the shared otel fixture without changing their test
behavior.
---
Outside diff comments:
In `@src/utils/agents/streaming.py`:
- Around line 251-257: Wrap the entire async generator body in a single
try/finally that owns root_span termination, including all yield and
early-return paths. Move deregister_stream and root_span.end into that
finalizer, remove the duplicated end calls from the explicit completion, error,
and compaction paths, and preserve existing stream behavior while ensuring
GeneratorExit from aclose() ends the span.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: d340d83c-79bb-448b-89bc-5e7a4e4f09b6
📒 Files selected for processing (4)
src/app/endpoints/streaming_query.pysrc/utils/agents/streaming.pytests/unit/app/endpoints/test_streaming_query.pytests/unit/utils/agents/test_streaming.py
📜 Review details
⏰ Context from checks skipped due to timeout. (14)
- GitHub Check: build-pr
- GitHub Check: unit_tests (3.12)
- GitHub Check: unit_tests (3.13)
- GitHub Check: integration_tests (3.12)
- GitHub Check: integration_tests (3.13)
- GitHub Check: Pylinter
- GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-0-8-on-pull-request
- GitHub Check: E2E: server mode / ci / group 2
- GitHub Check: E2E: library mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 2
- GitHub Check: E2E Tests for Lightspeed Evaluation job
🧰 Additional context used
📓 Path-based instructions (3)
**/*
📄 CodeRabbit inference engine (Custom checks)
**/*: Flag meaningful O(n^2)+ algorithms on non-trivial inputs, including handlers and Kubernetes list operations.
Flag N+1 patterns that list items and then query once per item, including Kubernetes API and database access.
Flag expensive work inside loops, including API calls, JSON parsing, and regex compilation.
Flag unbounded growth in caches, watchers, or buffers when eviction or limits are missing.
Flag missing pagination or limits on list operations and API endpoints.
Flag secrets or tokens logged in plaintext or hardcoded in source.
Flag API endpoints missing authentication or authorization.
Flag injection vulnerabilities, including SQL injection, command injection, and path traversal.
Flag sensitive data leaked in API responses, WebSocket messages, or logs.
Flag Kubernetes Secrets and Red Hat secrets missing OwnerReferences.
Files:
tests/unit/utils/agents/test_streaming.pysrc/utils/agents/streaming.pytests/unit/app/endpoints/test_streaming_query.pysrc/app/endpoints/streaming_query.py
tests/unit/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use pytest for unit tests, shared fixtures in
conftest.py,pytest-mockfor mocks,pytest.mark.asynciofor async tests, and maintain at least 60% unit-test coverage.
Files:
tests/unit/utils/agents/test_streaming.pytests/unit/app/endpoints/test_streaming_query.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Use absolute imports for internal modules and follow the prescribed FastAPI and Llama Stack import conventions.
All modules must begin with descriptive docstrings; uselogger = get_logger(__name__)fromlog.pyfor module logging; package__init__.pyfiles must contain brief package descriptions.
Define shared constants in the centralconstants.pymodule, add descriptive comments, and annotate constants withFinal[type].
Use complete type annotations for function parameters, return types, class attributes, and type aliases; prefer specific types overAny, use modern union syntax, and usetyping_extensions.Selffor model validators.
All functions and classes require descriptive Google-style docstrings, including appropriateParameters,Returns,Raises, andAttributessections.
Use descriptive snake_case, action-oriented function names such asget_,validate_, andcheck_; use PascalCase class names with standard suffixes such asConfiguration,Error/Exception,Resolver, andInterface.
Avoid modifying input parameters in place; return a newly constructed data structure instead.
Useasync deffor I/O operations and external API calls; API endpoints should raise FastAPIHTTPExceptionwith appropriate status codes and handle Llama StackAPIConnectionError.
Usefrom log import get_loggerand standard logger levels:debugfor diagnostics,infofor general execution,warningfor unexpected conditions or potential problems, anderrorfor serious failures.
Configuration models must extendConfigurationBase, setextra="forbid"to reject unknown fields, use Pydantic validators for custom validation, and use types such asOptional[FilePath],PositiveInt, andSecretStrwhere appropriate.
Abstract interfaces must useABCand@abstractmethoddecorators.
Never commit secrets or keys; use environment variables for sensitive data.
Files:
src/utils/agents/streaming.pysrc/app/endpoints/streaming_query.py
🧠 Learnings (4)
📚 Learning: 2026-06-24T13:45:37.249Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 1971
File: src/utils/markdown_repair.py:31-36
Timestamp: 2026-06-24T13:45:37.249Z
Learning: In the lightspeed-stack repository, docstrings must use the section header name "Parameters:" (not "Args:") for function arguments, even if the project references Google Python docstring conventions. Ensure docstrings follow the project’s established "Parameters:" header format for any documented function parameters.
Applied to files:
tests/unit/utils/agents/test_streaming.pysrc/utils/agents/streaming.pytests/unit/app/endpoints/test_streaming_query.pysrc/app/endpoints/streaming_query.py
📚 Learning: 2026-07-06T15:26:18.398Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 2071
File: src/models/config.py:2416-2422
Timestamp: 2026-07-06T15:26:18.398Z
Learning: In this repo’s Python code under src/**, don’t treat differences in string concatenation style as a style inconsistency when Black has effectively forced (or made clearer) use of explicit `+` string concatenation in multi-line logger/string expressions. If adjacent-literal implicit concatenation is avoided/changed specifically to accommodate Black’s formatting in these call sites, accept the `+` usage and don’t recommend converting it solely for consistency with nearby blocks that use implicit concatenation.
Applied to files:
src/utils/agents/streaming.pysrc/app/endpoints/streaming_query.py
📚 Learning: 2026-07-17T19:25:05.325Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 2166
File: src/utils/saved_prompts.py:129-157
Timestamp: 2026-07-17T19:25:05.325Z
Learning: For any endpoint that handles saved prompts and calls `src/utils/saved_prompts.py::create_saved_prompt`, treat the endpoint as the validation boundary. Before calling `create_saved_prompt`, validate the incoming saved-prompt name and content, specifically using `validate_saved_prompt_name` and then persist (store) the normalized value it returns. Do not call `create_saved_prompt` with unvalidated/raw name/content.
Applied to files:
src/utils/agents/streaming.pysrc/app/endpoints/streaming_query.py
📚 Learning: 2026-04-06T20:18:07.852Z
Learnt from: major
Repo: lightspeed-core/lightspeed-stack PR: 1463
File: src/app/endpoints/rlsapi_v1.py:266-271
Timestamp: 2026-04-06T20:18:07.852Z
Learning: In the lightspeed-stack codebase, within `src/app/endpoints/` inference/MCP endpoints, treat `tools: Optional[list[Any]]` in MCP tool definitions as an intentional, consistent typing pattern (used across `query`, `responses`, `streaming_query`, `rlsapi_v1`). Do not raise or suggest this as a typing issue during code review; changing it in isolation could break endpoint typing consistency across the codebase.
Applied to files:
src/app/endpoints/streaming_query.py
🔇 Additional comments (4)
src/app/endpoints/streaming_query.py (1)
170-180: LGTM!Also applies to: 211-220, 240-242, 353-353, 379-379
tests/unit/app/endpoints/test_streaming_query.py (1)
580-584: LGTM!Also applies to: 602-684, 686-875
src/utils/agents/streaming.py (1)
164-192: LGTM!Also applies to: 320-338
tests/unit/utils/agents/test_streaming.py (1)
8-18: LGTM!Also applies to: 70-70, 836-1086
| async for event in generate_agent_response( | ||
| generator, | ||
| context, | ||
| responses_params, | ||
| turn_summary, | ||
| background_topic_summary_tasks=_background_topic_summary_tasks, | ||
| emit_start=False, | ||
| original_input=compacted_original_input, | ||
| root_span=root_span, | ||
| ): | ||
| yield event | ||
| finally: | ||
| if root_span is not None: | ||
| root_span.end() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
opentelemetry-python Span.end() called twice behavior warning
💡 Result:
In the OpenTelemetry Python SDK, the official specification for the Span.end method states that only the first call to end should modify the span [1][2][3]. Implementations are free to either ignore subsequent calls or raise an error when end is called more than once [2][4][5]. If your application triggers a second call to span.end, you may encounter warnings or logs such as "Calling end on an ended span" depending on the specific library implementation being used [6]. Because the Span object is intended to represent a single, discrete operation with a fixed duration, calling end multiple times is considered an anti-pattern and often indicates a logic error in the instrumentation code, such as improper handling of conditional paths (e.g., calling end on both success and failure paths, or within a context manager block that already handles closing the span) [6][2][7]. To resolve this: 1. Review your instrumentation logic to ensure that span.end is called exactly once per span lifecycle [6]. 2. If you are using a context manager (the with statement), note that it automatically calls end upon exiting the block, so you should not call end manually within that block [2][8][7]. 3. If you encounter this warning in third-party libraries, it often points to a known issue in that library's instrumentation code, which may need to be addressed via an upstream fix or configuration adjustment [6].
Citations:
- 1: https://opentelemetry-python.readthedocs.io/en/stable/_modules/opentelemetry/trace/span.html
- 2: https://opentelemetry-python.readthedocs.io/en/stable/api/trace.html
- 3: https://github.com/open-telemetry/opentelemetry-python/blob/eed100c3/opentelemetry-api/src/opentelemetry/trace/span.py
- 4: https://opentelemetry-python.readthedocs.io/en/latest/sdk/trace.html
- 5: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-api/src/opentelemetry/trace/span.py
- 6: GRPC traces called twice when result evaluates to false open-telemetry/opentelemetry-python-contrib#3380
- 7: https://opentelemetry.io/docs/languages/python/instrumentation
- 8: https://opentelemetry-python.readthedocs.io/en/latest/_modules/opentelemetry/trace.html
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- candidate files ---'
fd -t f 'streaming_query\.py|streaming\.py' src
printf '%s\n' '--- relevant definitions and calls ---'
rg -n -C 5 'def generate_agent_response|generate_agent_response\(|root_span|end_span' src/app/endpoints/streaming_query.py src/utils/agents/streaming.py
printf '%s\n' '--- focused source ranges ---'
sed -n '220,355p' src/utils/agents/streaming.py
sed -n '450,525p' src/app/endpoints/streaming_query.pyRepository: lightspeed-core/lightspeed-stack
Length of output: 23484
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- OpenTelemetry version pins ---'
rg -n -i -C 2 'opentelemetry|otel' pyproject.toml poetry.lock uv.lock requirements*.txt setup.cfg setup.py 2>/dev/null || true
printf '%s\n' '--- full shared generator control flow ---'
sed -n '164,345p' src/utils/agents/streaming.py
printf '%s\n' '--- compaction wrapper control flow ---'
sed -n '403,510p' src/app/endpoints/streaming_query.py
printf '%s\n' '--- available installed OpenTelemetry modules ---'
python3 - <<'PY'
try:
import inspect
import opentelemetry.sdk.trace
from opentelemetry.sdk.trace import ReadableSpan
print("sdk_trace:", inspect.getfile(opentelemetry.sdk.trace))
print("ReadableSpan.end source:")
print(inspect.getsource(ReadableSpan.end))
except Exception as exc:
print(type(exc).__name__ + ":", exc)
PYRepository: lightspeed-core/lightspeed-stack
Length of output: 30922
🌐 Web query:
OpenTelemetry Python SDK 1.44.0 ReadableSpan end called twice warning "Calling end on ended span"
💡 Result:
The warning "Calling end on an ended span" in the OpenTelemetry Python SDK indicates that the end method is being invoked more than once on the same Span object [1][2]. In the OpenTelemetry Python SDK, the Span.end method includes a thread-safe check to verify if the span has already been ended [1][2]. If a second call to end is made, the SDK logs this warning to alert you that the operation has already been completed [1][2]. This behavior aligns with the OpenTelemetry specification, which states that implementations should ignore subsequent calls to end [3][4]. Common causes for this warning include: 1. Double Instrumentation or Logic Errors: Often caused by custom code or instrumentation libraries that inadvertently call end multiple times on the same span [5]. For example, if an instrumentation wrapper ends a span upon encountering a specific result, but the main application logic or another decorator also calls end, the second call will trigger this warning [5]. 2. Context Manager Issues: If you are using a span as a context manager (e.g., with tracer.start_as_current_span(...) as span:), the span is automatically ended when the block exits [6]. Manually calling span.end inside such a block is redundant and will cause this warning. 3. Buggy Instrumentation: In some cases, community-contributed instrumentation libraries (such as those in the opentelemetry-python-contrib repository) may have bugs where they incorrectly trigger an end call under specific execution paths (e.g., certain gRPC response conditions) [5]. To debug this, check your call stacks to identify where the duplicate calls are originating. If the warning occurs in third-party instrumentation, ensure your dependencies are updated to the latest versions, as these issues are often tracked and resolved in upstream patch releases [5].
Citations:
- 1: https://github.com/open-telemetry/opentelemetry-python/blob/eed100c3/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
- 2: https://github.com/open-telemetry/opentelemetry-python/blob/main/opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py
- 3: https://github.com/open-telemetry/opentelemetry-specification/blob/v1.44.0/specification/trace/api.md
- 4: https://opentelemetry-python.readthedocs.io/en/stable/_modules/opentelemetry/trace/span.html
- 5: GRPC traces called twice when result evaluates to false open-telemetry/opentelemetry-python-contrib#3380
- 6: https://opentelemetry-python.readthedocs.io/en/latest/api/trace.html
Give root_span a single owner. generate_agent_response ends the span on normal completion and handled error paths, while generate_response_with_compaction ends it again in finally. OpenTelemetry SDK 1.44.0 logs Calling end on an ended span for the second call. Keep the wrapper’s finally and make the delegated generator record attributes without ending the span.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/endpoints/streaming_query.py` around lines 492 - 505, Update
generate_agent_response so it no longer ends root_span on normal completion or
handled error paths, leaving span termination solely to
generate_response_with_compaction’s existing finally block. Preserve the
delegated generator’s span attribute recording and ensure all paths still allow
the wrapper to end the span exactly once.
| @pytest.fixture(name="otel") | ||
| def otel_fixture( | ||
| self, | ||
| ) -> Generator[tuple[Any, InMemorySpanExporter], None, None]: | ||
| """Provide an isolated tracer and exporter for OTEL tests.""" | ||
| exporter = InMemorySpanExporter() | ||
| provider = TracerProvider() | ||
| provider.add_span_processor(SimpleSpanProcessor(exporter)) | ||
| tracer = provider.get_tracer("unit-test-tracer") | ||
| yield tracer, exporter | ||
| exporter.clear() | ||
| provider.shutdown() |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Duplicated otel_fixture. The same tracer/exporter fixture is defined twice, byte for byte. Move it to a shared conftest.py so both suites use one definition.
tests/unit/utils/agents/test_streaming.py#L823-L834: remove the class-localotel_fixtureand consume the sharedotelfixture.tests/unit/app/endpoints/test_streaming_query.py#L589-L600: remove the class-localotel_fixtureand consume the sharedotelfixture.
As per coding guidelines: "Use pytest for unit tests, shared fixtures in conftest.py".
📍 Affects 2 files
tests/unit/utils/agents/test_streaming.py#L823-L834(this comment)tests/unit/app/endpoints/test_streaming_query.py#L589-L600
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/utils/agents/test_streaming.py` around lines 823 - 834, Move the
duplicated otel_fixture definition into a shared conftest.py. Remove the
class-local otel_fixture from tests/unit/utils/agents/test_streaming.py lines
823-834 and tests/unit/app/endpoints/test_streaming_query.py lines 589-600,
ensuring both test suites consume the shared otel fixture without changing their
test behavior.
Source: Coding guidelines
03f83b5 to
c33d499
Compare
…query endpoint
Adds OpenTelemetry (OTEL) tracing instrumentation for the POST /v1/streaming_query endpoint to enable distributed tracing and observability.
**Instrumented Components**
1. Streaming Query Endpoint Handler (`src/app/endpoints/streaming_query.py`)
2. Agent Streaming Response (`src/utils/agents/streaming.py`)
**Span Hierarchy**
streaming_query.handle_request (root span)
├── quota.check
├── shield.moderate
├── rag.retrieve
└── llm.inference
└── tool.execution (attributes only)
**Design Note:** Uses manual span management (`tracer.start_span()` + `trace.use_span()` with `end_on_exit=False`) instead of `tracer.start_as_current_span()` because `StreamingResponse` generators run after the handler returns
— a context manager would close the span prematurely. The span is ended explicitly in `generate_agent_response` at all exit paths (success, stream error, cancellation, topic summary failure), and in
`generate_response_with_compaction` via try/finally. `span.end()` is idempotent in the OTel Python SDK.
c33d499 to
46e385b
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/app/endpoints/streaming_query.py`:
- Around line 208-217: Wrap the anonymized attribute construction in the
streaming query handler around set_span_attributes in a try/except for
ValueError, so missing OTEL anonymization configuration cannot fail the request.
On failure, log a warning and continue without setting the affected span
attributes, preserving normal request validation and endpoint execution.
In `@tests/unit/app/endpoints/test_streaming_query.py`:
- Around line 798-830: Add a test alongside
test_passes_root_span_to_generate_agent_response that forces
needs_compaction_path to return True, exercises streaming_query_endpoint_handler
through the compaction flow, drains the response, and uses the in-memory
exporter to assert exactly one finished span named
streaming_query.handle_request.
In `@tests/unit/utils/agents/test_streaming.py`:
- Around line 982-1027: Strengthen test_no_spans_finished_when_root_span_is_none
by collecting and asserting the yielded events from generate_agent_response,
rather than only checking exporter.get_finished_spans(). Verify the expected
completion/event output so the test confirms the root_span=None success path
still yields its events while retaining the no-finished-spans assertion.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro Plus
Run ID: 81e27dd5-327e-4b81-a77b-f3bdc1801e97
📒 Files selected for processing (4)
src/app/endpoints/streaming_query.pytests/unit/app/endpoints/test_streaming_query.pytests/unit/conftest.pytests/unit/utils/agents/test_streaming.py
📜 Review details
⏰ Context from checks skipped due to timeout. (7)
- GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-0-8-on-pull-request
- GitHub Check: E2E: server mode / ci / group 2
- GitHub Check: E2E: library mode / ci / group 2
- GitHub Check: E2E: library mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 3
- GitHub Check: E2E: library mode / ci / group 1
- GitHub Check: E2E Tests for Lightspeed Evaluation job
🧰 Additional context used
📓 Path-based instructions (3)
**/*
📄 CodeRabbit inference engine (Custom checks)
**/*: Flag meaningful O(n^2)+ algorithms on non-trivial inputs, including handlers and Kubernetes list operations.
Flag N+1 patterns that list items and then query once per item, including Kubernetes API and database access.
Flag expensive work inside loops, including API calls, JSON parsing, and regex compilation.
Flag unbounded growth in caches, watchers, or buffers when eviction or limits are missing.
Flag missing pagination or limits on list operations and API endpoints.
Flag secrets or tokens logged in plaintext or hardcoded in source.
Flag API endpoints missing authentication or authorization.
Flag injection vulnerabilities, including SQL injection, command injection, and path traversal.
Flag sensitive data leaked in API responses, WebSocket messages, or logs.
Flag Kubernetes Secrets and Red Hat secrets missing OwnerReferences.
Files:
tests/unit/conftest.pysrc/app/endpoints/streaming_query.pytests/unit/utils/agents/test_streaming.pytests/unit/app/endpoints/test_streaming_query.py
tests/unit/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Use pytest for unit tests, shared fixtures in
conftest.py,pytest-mockfor mocks,pytest.mark.asynciofor async tests, and maintain at least 60% unit-test coverage.
Files:
tests/unit/conftest.pytests/unit/utils/agents/test_streaming.pytests/unit/app/endpoints/test_streaming_query.py
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Use absolute imports for internal modules and follow the prescribed FastAPI and Llama Stack import conventions.
All modules must begin with descriptive docstrings; uselogger = get_logger(__name__)fromlog.pyfor module logging; package__init__.pyfiles must contain brief package descriptions.
Define shared constants in the centralconstants.pymodule, add descriptive comments, and annotate constants withFinal[type].
Use complete type annotations for function parameters, return types, class attributes, and type aliases; prefer specific types overAny, use modern union syntax, and usetyping_extensions.Selffor model validators.
All functions and classes require descriptive Google-style docstrings, including appropriateParameters,Returns,Raises, andAttributessections.
Use descriptive snake_case, action-oriented function names such asget_,validate_, andcheck_; use PascalCase class names with standard suffixes such asConfiguration,Error/Exception,Resolver, andInterface.
Avoid modifying input parameters in place; return a newly constructed data structure instead.
Useasync deffor I/O operations and external API calls; API endpoints should raise FastAPIHTTPExceptionwith appropriate status codes and handle Llama StackAPIConnectionError.
Usefrom log import get_loggerand standard logger levels:debugfor diagnostics,infofor general execution,warningfor unexpected conditions or potential problems, anderrorfor serious failures.
Configuration models must extendConfigurationBase, setextra="forbid"to reject unknown fields, use Pydantic validators for custom validation, and use types such asOptional[FilePath],PositiveInt, andSecretStrwhere appropriate.
Abstract interfaces must useABCand@abstractmethoddecorators.
Never commit secrets or keys; use environment variables for sensitive data.
Files:
src/app/endpoints/streaming_query.py
🧠 Learnings (4)
📚 Learning: 2026-06-24T13:45:37.249Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 1971
File: src/utils/markdown_repair.py:31-36
Timestamp: 2026-06-24T13:45:37.249Z
Learning: In the lightspeed-stack repository, docstrings must use the section header name "Parameters:" (not "Args:") for function arguments, even if the project references Google Python docstring conventions. Ensure docstrings follow the project’s established "Parameters:" header format for any documented function parameters.
Applied to files:
tests/unit/conftest.pysrc/app/endpoints/streaming_query.pytests/unit/utils/agents/test_streaming.pytests/unit/app/endpoints/test_streaming_query.py
📚 Learning: 2026-04-06T20:18:07.852Z
Learnt from: major
Repo: lightspeed-core/lightspeed-stack PR: 1463
File: src/app/endpoints/rlsapi_v1.py:266-271
Timestamp: 2026-04-06T20:18:07.852Z
Learning: In the lightspeed-stack codebase, within `src/app/endpoints/` inference/MCP endpoints, treat `tools: Optional[list[Any]]` in MCP tool definitions as an intentional, consistent typing pattern (used across `query`, `responses`, `streaming_query`, `rlsapi_v1`). Do not raise or suggest this as a typing issue during code review; changing it in isolation could break endpoint typing consistency across the codebase.
Applied to files:
src/app/endpoints/streaming_query.py
📚 Learning: 2026-07-06T15:26:18.398Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 2071
File: src/models/config.py:2416-2422
Timestamp: 2026-07-06T15:26:18.398Z
Learning: In this repo’s Python code under src/**, don’t treat differences in string concatenation style as a style inconsistency when Black has effectively forced (or made clearer) use of explicit `+` string concatenation in multi-line logger/string expressions. If adjacent-literal implicit concatenation is avoided/changed specifically to accommodate Black’s formatting in these call sites, accept the `+` usage and don’t recommend converting it solely for consistency with nearby blocks that use implicit concatenation.
Applied to files:
src/app/endpoints/streaming_query.py
📚 Learning: 2026-07-17T19:25:05.325Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 2166
File: src/utils/saved_prompts.py:129-157
Timestamp: 2026-07-17T19:25:05.325Z
Learning: For any endpoint that handles saved prompts and calls `src/utils/saved_prompts.py::create_saved_prompt`, treat the endpoint as the validation boundary. Before calling `create_saved_prompt`, validate the incoming saved-prompt name and content, specifically using `validate_saved_prompt_name` and then persist (store) the normalized value it returns. Do not call `create_saved_prompt` with unvalidated/raw name/content.
Applied to files:
src/app/endpoints/streaming_query.py
🔇 Additional comments (5)
src/app/endpoints/streaming_query.py (2)
170-177: 📐 Maintainability & Code Quality | ⚡ Quick winSpan ownership is still split.
generate_agent_responseendsroot_spanon the early-return and topic-summary-failure paths, andgenerate_response_with_compactionends it again in itsfinally. The OpenTelemetry SDK ignores the second call and logsCalling end on an ended span. Give the span one owner: let the wrapper'sfinallyend it, and let the delegated generator only record attributes and events.
405-405: LGTM!Also applies to: 421-425, 426-499
tests/unit/app/endpoints/test_streaming_query.py (1)
579-583: LGTM!Also applies to: 585-670, 672-796, 832-861, 863-909
tests/unit/conftest.py (1)
9-18: LGTM!Also applies to: 57-66
tests/unit/utils/agents/test_streaming.py (1)
14-16: LGTM!Also applies to: 68-68, 818-885, 888-926, 928-980, 1029-1071
| set_span_attributes( | ||
| root_span, | ||
| { | ||
| SpanAttributes.USER_ID: anonymize_value(user_id), | ||
| SpanAttributes.INPUT: anonymize_value(query_request.query), | ||
| SpanAttributes.REQUEST_ATTACHMENTS_COUNT: ( | ||
| len(query_request.attachments) if query_request.attachments else 0 | ||
| ), | ||
| }, | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Anonymization failure now fails the request. anonymize_value raises ValueError when OTEL_ANONYMIZATION_SECRET is unset and OTEL_SDK_DISABLED is not true/1. This call runs before any request validation, so a deployment that enables the OTEL SDK without the secret returns 500 for every POST /v1/streaming_query. Tracing should not be able to break the endpoint.
Fail soft here and log a warning instead.
🛡️ Proposed guard
# Set initial span attributes
- set_span_attributes(
- root_span,
- {
- SpanAttributes.USER_ID: anonymize_value(user_id),
- SpanAttributes.INPUT: anonymize_value(query_request.query),
- SpanAttributes.REQUEST_ATTACHMENTS_COUNT: (
- len(query_request.attachments) if query_request.attachments else 0
- ),
- },
- )
+ try:
+ set_span_attributes(
+ root_span,
+ {
+ SpanAttributes.USER_ID: anonymize_value(user_id),
+ SpanAttributes.INPUT: anonymize_value(query_request.query),
+ SpanAttributes.REQUEST_ATTACHMENTS_COUNT: (
+ len(query_request.attachments) if query_request.attachments else 0
+ ),
+ },
+ )
+ except ValueError as exc:
+ logger.warning("Skipping OTEL request attributes: %s", exc)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| set_span_attributes( | |
| root_span, | |
| { | |
| SpanAttributes.USER_ID: anonymize_value(user_id), | |
| SpanAttributes.INPUT: anonymize_value(query_request.query), | |
| SpanAttributes.REQUEST_ATTACHMENTS_COUNT: ( | |
| len(query_request.attachments) if query_request.attachments else 0 | |
| ), | |
| }, | |
| ) | |
| # Set initial span attributes | |
| try: | |
| set_span_attributes( | |
| root_span, | |
| { | |
| SpanAttributes.USER_ID: anonymize_value(user_id), | |
| SpanAttributes.INPUT: anonymize_value(query_request.query), | |
| SpanAttributes.REQUEST_ATTACHMENTS_COUNT: ( | |
| len(query_request.attachments) if query_request.attachments else 0 | |
| ), | |
| }, | |
| ) | |
| except ValueError as exc: | |
| logger.warning("Skipping OTEL request attributes: %s", exc) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/app/endpoints/streaming_query.py` around lines 208 - 217, Wrap the
anonymized attribute construction in the streaming query handler around
set_span_attributes in a try/except for ValueError, so missing OTEL
anonymization configuration cannot fail the request. On failure, log a warning
and continue without setting the affected span attributes, preserving normal
request validation and endpoint execution.
| @pytest.mark.asyncio | ||
| async def test_passes_root_span_to_generate_agent_response( | ||
| self, | ||
| dummy_request: Request, # pylint: disable=redefined-outer-name | ||
| setup_configuration: AppConfig, | ||
| mocker: MockerFixture, | ||
| otel: tuple[Any, InMemorySpanExporter], | ||
| ) -> None: | ||
| """Test that root_span is forwarded to generate_agent_response.""" | ||
| tracer, _exporter = otel | ||
| self._setup_common_mocks(mocker, setup_configuration, tracer) | ||
|
|
||
| mock_gen = mocker.patch( | ||
| "app.endpoints.streaming_query.generate_agent_response", | ||
| ) | ||
|
|
||
| async def gen_side_effect(*_a: Any, **_kw: Any) -> AsyncIterator[str]: | ||
| yield "data: test\n\n" | ||
|
|
||
| mock_gen.side_effect = gen_side_effect | ||
|
|
||
| response = await streaming_query_endpoint_handler( | ||
| request=dummy_request, | ||
| query_request=QueryRequest( | ||
| query="test" | ||
| ), # pyright: ignore[reportCallIssue] | ||
| auth=MOCK_AUTH_STREAMING, | ||
| mcp_headers={}, | ||
| ) | ||
| await _drain_response(response) | ||
|
|
||
| mock_gen.assert_called_once() | ||
| assert mock_gen.call_args.kwargs["root_span"] is not None |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
No coverage for the compaction path. generate_response_with_compaction receives root_span and ends it in a finally block, but no test in this class exercises that branch. That branch is where the double-end risk lives. Add a test that forces needs_compaction_path to return True and asserts exactly one finished streaming_query.handle_request span.
Do you want me to generate that test?
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/app/endpoints/test_streaming_query.py` around lines 798 - 830, Add
a test alongside test_passes_root_span_to_generate_agent_response that forces
needs_compaction_path to return True, exercises streaming_query_endpoint_handler
through the compaction flow, drains the response, and uses the in-memory
exporter to assert exactly one finished span named
streaming_query.handle_request.
| async def test_no_spans_finished_when_root_span_is_none( | ||
| self, | ||
| mocker: MockerFixture, | ||
| make_generator_context: Callable[..., ResponseGeneratorContext], | ||
| responses_params: ResponsesApiParams, | ||
| otel: tuple[Any, InMemorySpanExporter], | ||
| ) -> None: | ||
| """Test that no spans are finished when root_span is None.""" | ||
| _tracer, exporter = otel | ||
| context = make_generator_context() | ||
| turn_summary = TurnSummary() | ||
| turn_summary.token_usage = TokenCounter(input_tokens=3, output_tokens=7) | ||
|
|
||
| async def inner() -> AsyncIterator[str]: | ||
| yield serialize_event( | ||
| TokenStreamPayload.create(chunk_id=0, token="Hi"), | ||
| MEDIA_TYPE_JSON, | ||
| ) | ||
|
|
||
| mocker.patch("utils.agents.streaming.consume_query_tokens") | ||
| mocker.patch( | ||
| "utils.agents.streaming.get_available_quotas", | ||
| return_value={"daily": 100}, | ||
| ) | ||
| mocker.patch( | ||
| "utils.agents.streaming.maybe_get_topic_summary", | ||
| new=mocker.AsyncMock(return_value=None), | ||
| ) | ||
| mocker.patch("utils.agents.streaming.store_query_results") | ||
| mock_config = mocker.Mock() | ||
| mock_config.quota_limiters = [] | ||
| mocker.patch("utils.agents.streaming.configuration", mock_config) | ||
|
|
||
| [ | ||
| event | ||
| async for event in generate_agent_response( | ||
| inner(), | ||
| context, | ||
| responses_params, | ||
| turn_summary, | ||
| [], | ||
| root_span=None, | ||
| ) | ||
| ] | ||
|
|
||
| assert len(exporter.get_finished_spans()) == 0 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value
The assertion is weaker than the test name implies. No span is ever started in this test, and the otel exporter is per-test. get_finished_spans() == 0 therefore holds regardless of what generate_agent_response does with root_span. What the test actually proves is that the success path does not raise on root_span=None. Assert the yielded events too, so a regression that skips the completion path is caught.
♻️ Suggested tightening
- [
- event
- async for event in generate_agent_response(
+ events = [
+ event
+ async for event in generate_agent_response(
inner(),
context,
responses_params,
turn_summary,
[],
root_span=None,
)
]
+ assert _sse_event_types(events) == ["start", "token", "end"]
assert len(exporter.get_finished_spans()) == 0📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def test_no_spans_finished_when_root_span_is_none( | |
| self, | |
| mocker: MockerFixture, | |
| make_generator_context: Callable[..., ResponseGeneratorContext], | |
| responses_params: ResponsesApiParams, | |
| otel: tuple[Any, InMemorySpanExporter], | |
| ) -> None: | |
| """Test that no spans are finished when root_span is None.""" | |
| _tracer, exporter = otel | |
| context = make_generator_context() | |
| turn_summary = TurnSummary() | |
| turn_summary.token_usage = TokenCounter(input_tokens=3, output_tokens=7) | |
| async def inner() -> AsyncIterator[str]: | |
| yield serialize_event( | |
| TokenStreamPayload.create(chunk_id=0, token="Hi"), | |
| MEDIA_TYPE_JSON, | |
| ) | |
| mocker.patch("utils.agents.streaming.consume_query_tokens") | |
| mocker.patch( | |
| "utils.agents.streaming.get_available_quotas", | |
| return_value={"daily": 100}, | |
| ) | |
| mocker.patch( | |
| "utils.agents.streaming.maybe_get_topic_summary", | |
| new=mocker.AsyncMock(return_value=None), | |
| ) | |
| mocker.patch("utils.agents.streaming.store_query_results") | |
| mock_config = mocker.Mock() | |
| mock_config.quota_limiters = [] | |
| mocker.patch("utils.agents.streaming.configuration", mock_config) | |
| [ | |
| event | |
| async for event in generate_agent_response( | |
| inner(), | |
| context, | |
| responses_params, | |
| turn_summary, | |
| [], | |
| root_span=None, | |
| ) | |
| ] | |
| assert len(exporter.get_finished_spans()) == 0 | |
| async def test_no_spans_finished_when_root_span_is_none( | |
| self, | |
| mocker: MockerFixture, | |
| make_generator_context: Callable[..., ResponseGeneratorContext], | |
| responses_params: ResponsesApiParams, | |
| otel: tuple[Any, InMemorySpanExporter], | |
| ) -> None: | |
| """Test that no spans are finished when root_span is None.""" | |
| _tracer, exporter = otel | |
| context = make_generator_context() | |
| turn_summary = TurnSummary() | |
| turn_summary.token_usage = TokenCounter(input_tokens=3, output_tokens=7) | |
| async def inner() -> AsyncIterator[str]: | |
| yield serialize_event( | |
| TokenStreamPayload.create(chunk_id=0, token="Hi"), | |
| MEDIA_TYPE_JSON, | |
| ) | |
| mocker.patch("utils.agents.streaming.consume_query_tokens") | |
| mocker.patch( | |
| "utils.agents.streaming.get_available_quotas", | |
| return_value={"daily": 100}, | |
| ) | |
| mocker.patch( | |
| "utils.agents.streaming.maybe_get_topic_summary", | |
| new=mocker.AsyncMock(return_value=None), | |
| ) | |
| mocker.patch("utils.agents.streaming.store_query_results") | |
| mock_config = mocker.Mock() | |
| mock_config.quota_limiters = [] | |
| mocker.patch("utils.agents.streaming.configuration", mock_config) | |
| events = [ | |
| event | |
| async for event in generate_agent_response( | |
| inner(), | |
| context, | |
| responses_params, | |
| turn_summary, | |
| [], | |
| root_span=None, | |
| ) | |
| ] | |
| assert _sse_event_types(events) == ["start", "token", "end"] | |
| assert len(exporter.get_finished_spans()) == 0 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/utils/agents/test_streaming.py` around lines 982 - 1027,
Strengthen test_no_spans_finished_when_root_span_is_none by collecting and
asserting the yielded events from generate_agent_response, rather than only
checking exporter.get_finished_spans(). Verify the expected completion/event
output so the test confirms the root_span=None success path still yields its
events while retaining the no-finished-spans assertion.
|
Screenshots from Jaeger: 
|
1 participant
Description
Adds OpenTelemetry (OTEL) tracing instrumentation for the POST /v1/streaming_query endpoint to enable distributed tracing and observability.
Instrumented Components
src/app/endpoints/streaming_query.py)src/utils/agents/streaming.py)Span Hierarchy
streaming_query.handle_request (root span)
├── quota.check
├── shield.moderate
├── rag.retrieve
└── llm.inference
└── tool.execution (attributes only)
Design Note: Uses manual span management (
tracer.start_span()+trace.use_span()withend_on_exit=False) instead oftracer.start_as_current_span()becauseStreamingResponsegenerators run after the handler returns — a context manager would close the span prematurely. The span is ended explicitly ingenerate_agent_responseat all exit paths (success, stream error, cancellation, topic summary failure), and ingenerate_response_with_compactionvia try/finally.span.end()is idempotent in the OTel Python SDK.Type of change
pyproject.toml+uv.lock]requirements.*.txtfor Konflux]Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
Observability
Tests