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