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(