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

@@ -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"):

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")

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",

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")

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(

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()