From 24796ae61fe9f6d4184132f4a07495810c7cbc15 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Wed, 23 Sep 2026 22:09:58 +0100 Subject: [PATCH] Mint request ids on the server and trace refused requests build_agent no longer takes request_id from the request body: it becomes the primary LLM's usage request id, and quotas count distinct request ids, so a client could make every call count as one. Requests refused after setup started (unauthorized, over quota, resume conflict, setup error) now write their trace, marked error, through an after-request hook; streaming routes hand the trace to complete_stream instead. --- docsgpt/api/answer/routes/answer.py | 8 ++- docsgpt/api/answer/routes/base.py | 8 ++- docsgpt/api/answer/routes/stream.py | 10 ++- .../api/answer/services/stream_processor.py | 61 +++++++++++++++-- docsgpt/api/v1/routes.py | 8 ++- tests/api/answer/test_trace_wiring.py | 67 ++++++++++++++++++- 6 files changed, 147 insertions(+), 15 deletions(-) diff --git a/docsgpt/api/answer/routes/answer.py b/docsgpt/api/answer/routes/answer.py index bc2eb726..4d7f8dd2 100644 --- a/docsgpt/api/answer/routes/answer.py +++ b/docsgpt/api/answer/routes/answer.py @@ -13,7 +13,10 @@ from docsgpt.api.answer.services.continuation_service import ( ResumeInProgressError, ) from docsgpt.api.answer.services.persistence_policy import resolve_persistence -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) logger = logging.getLogger(__name__) @@ -88,6 +91,9 @@ class AnswerResource(Resource, BaseAnswerResource): return error decoded_token = getattr(request, "decoded_token", None) processor = StreamProcessor(data, decoded_token, trace_source="answer") + # ``complete_stream`` is consumed below and writes the trace itself; + # this covers requests refused before it runs. + flush_trace_after_request(processor) try: # ---- Continuation mode ---- if data.get("tool_actions"): diff --git a/docsgpt/api/answer/routes/base.py b/docsgpt/api/answer/routes/base.py index f23b82f3..943efedc 100644 --- a/docsgpt/api/answer/routes/base.py +++ b/docsgpt/api/answer/routes/base.py @@ -6,7 +6,7 @@ import logging import threading import time import uuid -from typing import Any, Dict, Generator, List, Optional +from typing import Any, Callable, Dict, Generator, List, Optional from flask import jsonify, make_response, Response from flask_restx import Namespace @@ -69,7 +69,9 @@ class StreamSuperseded(Exception): answer_ns = Namespace("answer", description="Answer related operations", path="/") -def _traced_stream(method): +def _traced_stream( + method: Callable[..., Generator[str, None, None]], +) -> Callable[..., Generator[str, None, None]]: """Record ``complete_stream`` as the request's execution trace. The stream runs in the SSE pump thread, not the request thread, so the @@ -81,7 +83,7 @@ def _traced_stream(method): signature = inspect.signature(method) @functools.wraps(method) - def wrapper(*args, **kwargs): + def wrapper(*args: Any, **kwargs: Any) -> Generator[str, None, None]: bound = signature.bind(*args, **kwargs) arguments = bound.arguments continuation = arguments.get("_continuation") diff --git a/docsgpt/api/answer/routes/stream.py b/docsgpt/api/answer/routes/stream.py index 36cb0472..21dbcbff 100644 --- a/docsgpt/api/answer/routes/stream.py +++ b/docsgpt/api/answer/routes/stream.py @@ -13,7 +13,10 @@ from docsgpt.api.answer.services.continuation_service import ( ResumeInProgressError, ) from docsgpt.api.answer.services.persistence_policy import resolve_persistence -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) from docsgpt.streaming.sse_keepalive import with_sse_keepalive logger = logging.getLogger(__name__) @@ -95,6 +98,7 @@ class StreamResource(Resource, BaseAnswerResource): return error decoded_token = getattr(request, "decoded_token", None) processor = StreamProcessor(data, decoded_token, trace_source="stream") + flush_trace_after_request(processor) try: # ---- Continuation mode ---- @@ -130,7 +134,7 @@ class StreamResource(Resource, BaseAnswerResource): agent_id=processor.agent_id, model_id=processor.model_id, model_user_id=processor.model_user_id, - trace=processor.trace, + trace=processor.handoff_trace(), _continuation={ "messages": messages, "tools_dict": tools_dict, @@ -183,7 +187,7 @@ class StreamResource(Resource, BaseAnswerResource): model_id=processor.model_id, model_user_id=processor.model_user_id, request_id=processor.request_id, - trace=processor.trace, + trace=processor.handoff_trace(), ), ), mimetype="text/event-stream", diff --git a/docsgpt/api/answer/services/stream_processor.py b/docsgpt/api/answer/services/stream_processor.py index bb7f47a3..27be1028 100644 --- a/docsgpt/api/answer/services/stream_processor.py +++ b/docsgpt/api/answer/services/stream_processor.py @@ -4,7 +4,9 @@ import json import logging import uuid from pathlib import Path -from typing import Any, Dict, Optional, Set +from typing import Any, Callable, Dict, Optional, Set, TypeVar + +from flask import after_this_request from docsgpt import tracing from docsgpt.agents.agent_creator import AgentCreator @@ -107,7 +109,10 @@ def get_prompt(prompt_id: str, prompts_collection=None) -> str: raise ValueError(f"Invalid prompt ID: {prompt_id}") from e -def _traced_setup(method): +T = TypeVar("T") + + +def _traced_setup(method: Callable[..., T]) -> Callable[..., T]: """Run a request-setup method inside the request's execution trace. Agent setup does real work worth seeing in the trace -- pre-fetch @@ -116,7 +121,7 @@ def _traced_setup(method): """ @functools.wraps(method) - def wrapper(self, *args, **kwargs): + def wrapper(self: "StreamProcessor", *args: Any, **kwargs: Any) -> T: trace = getattr(self, "trace", None) if trace is None: trace = tracing.start_trace(source=getattr(self, "trace_source", "stream")) @@ -136,6 +141,25 @@ def _traced_setup(method): return wrapper +def flush_trace_after_request(processor: "StreamProcessor") -> None: + """Write ``processor``'s setup trace when the request ends, unless it was claimed. + + Registered on the current request with ``after_this_request``; the hook + always hands the response back unchanged and never raises. + + Args: + processor: The request's processor. + """ + + @after_this_request + def _flush(response: Any) -> Any: + try: + processor.flush_unclaimed_trace() + except Exception: + logger.warning("Could not write an unclaimed request trace", exc_info=True) + return response + + class StreamProcessor: def __init__( self, @@ -201,6 +225,31 @@ class StreamProcessor: self._load_conversation_history() self._process_attachments() + def handoff_trace(self) -> Optional[tracing.Trace]: + """Hand the setup trace to a streaming ``complete_stream``. + + The stream writes the trace when it ends, after the view has + returned; marking the hand-off stops :meth:`flush_unclaimed_trace` + from writing it first. + + Returns: + The trace to pass as ``complete_stream(trace=...)``. + """ + self._trace_handed_off = True + return getattr(self, "trace", None) + + def flush_unclaimed_trace(self) -> None: + """Write the setup trace of a request that ended before streaming. + + A request refused after setup started (unauthorized, over its usage + limit, a resume conflict, a setup error) still records what ran, + marked ``error``. A trace the request already wrote, or handed to a + stream, is left alone. Routes arrange this with + :func:`flush_trace_after_request`. + """ + if not getattr(self, "_trace_handed_off", False): + tracing.flush(getattr(self, "trace", None), tracing.STATUS_ERROR) + @_traced_setup def build_agent(self, question: str): """One call to go from request data to a ready-to-run agent. @@ -208,10 +257,12 @@ class StreamProcessor: Combines initialize(), pre_fetch_docs(), pre_fetch_tools(), and create_agent() into a single convenience method. The request id is minted first so pre-fetch retrieval and its side-channel LLM calls - share it with the rest of the turn. + share it with the rest of the turn. It is always generated here, never + taken from the request body: request quotas count distinct request + ids, so a client-chosen id would let every call count as one. """ if not getattr(self, "request_id", None): - self.request_id = self.data.get("request_id") or str(uuid.uuid4()) + self.request_id = str(uuid.uuid4()) self.initialize() agent_type = self.agent_config.get("agent_type", "classic") diff --git a/docsgpt/api/v1/routes.py b/docsgpt/api/v1/routes.py index 9294c23a..c177ad6c 100644 --- a/docsgpt/api/v1/routes.py +++ b/docsgpt/api/v1/routes.py @@ -21,7 +21,10 @@ from docsgpt.api.answer.services.continuation_service import ( RESUME_IN_PROGRESS_MESSAGE, ResumeInProgressError, ) -from docsgpt.api.answer.services.stream_processor import StreamProcessor +from docsgpt.api.answer.services.stream_processor import ( + StreamProcessor, + flush_trace_after_request, +) from docsgpt.api.v1 import idempotency as v1_idempotency from docsgpt.api.v1.session_store import ( V1Session, @@ -258,6 +261,7 @@ def chat_completions(): try: processor = StreamProcessor(internal_data, decoded_token, trace_source="v1") + flush_trace_after_request(processor) # Set when this request took the resume claim, so a refusal can release it. claimed_conversation_id = None @@ -371,6 +375,8 @@ def chat_completions(): # safe way to re-emit a recorded SSE stream (and the regression / # b2b client is non-streaming), so a streaming request never # claims a key. This is a known, accepted limitation. + # The stream writes the trace once it runs, after this returns. + processor.handoff_trace() return Response( with_sse_keepalive( _stream_response( diff --git a/tests/api/answer/test_trace_wiring.py b/tests/api/answer/test_trace_wiring.py index 9153c424..7cf75f7a 100644 --- a/tests/api/answer/test_trace_wiring.py +++ b/tests/api/answer/test_trace_wiring.py @@ -238,14 +238,46 @@ class TestProcessorTraceSetup: assert [s.kind for s in trace.spans] == ["retrieval"] assert tracing.current_trace() is None - def test_client_supplied_request_id_is_kept(self): + def test_client_supplied_request_id_is_ignored(self): + """Quotas count distinct request ids; a client must not choose its own.""" from docsgpt.api.answer.services.stream_processor import StreamProcessor processor = StreamProcessor({"request_id": "client-rid"}, {"sub": "u1"}) with patch.object(processor, "initialize", side_effect=RuntimeError("stop")): with pytest.raises(RuntimeError): processor.build_agent("q") - assert processor.request_id == "client-rid" + assert processor.request_id and processor.request_id != "client-rid" + + def test_refused_request_still_writes_its_trace(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({}, {"sub": "u1"}) + + def _initialize(): + with tracing.span(tracing.KIND_RETRIEVAL, "retrieval"): + pass + + with patch.object(processor, "initialize", side_effect=_initialize), patch.object( + processor, "pre_fetch_docs", return_value=(None, None) + ), patch.object(processor, "pre_fetch_tools", return_value=None), patch.object( + processor, "create_agent", return_value=MagicMock() + ), patch.object(processor, "_exposure_partition", return_value=([], [])): + processor.build_agent("q") + with _captured_flushes() as flushed: + processor.flush_unclaimed_trace() + (trace,) = flushed + assert trace.status == "error" + assert [s.kind for s in trace.spans] == ["retrieval"] + + def test_handed_off_trace_is_left_to_the_stream(self): + from docsgpt.api.answer.services.stream_processor import StreamProcessor + + processor = StreamProcessor({}, {"sub": "u1"}) + processor.trace = tracing.start_trace(source="stream", capture_otel_context=False) + assert processor.handoff_trace() is processor.trace + with _captured_flushes() as flushed: + processor.flush_unclaimed_trace() + assert flushed == [] def test_tracing_disabled_leaves_no_trace(self, monkeypatch): from docsgpt.api.answer.services.stream_processor import StreamProcessor @@ -257,3 +289,34 @@ class TestProcessorTraceSetup: processor.build_agent("q") assert processor.trace is None assert processor.request_id + + +@pytest.mark.unit +class TestRouteFlushesRefusedRequests: + def test_unauthorized_answer_request_writes_its_trace(self, mock_mongo_db, flask_app): + """The route registers the flush, and the hook never replaces the response.""" + import json + + from flask_restx import Api + + from docsgpt.api.answer.routes.answer import answer_ns + + api = Api(flask_app) + api.add_namespace(answer_ns) + client = flask_app.test_client() + processor = MagicMock() + processor.decoded_token = None + processor.flush_unclaimed_trace.return_value = "not a response" + with patch( + "docsgpt.api.answer.routes.answer.StreamProcessor", return_value=processor + ), patch( + "docsgpt.api.answer.routes.answer.AnswerResource.validate_request", + return_value=None, + ): + resp = client.post( + "/api/answer", + data=json.dumps({"question": "q"}), + content_type="application/json", + ) + assert resp.status_code == 401 + processor.flush_unclaimed_trace.assert_called_once_with()