Accept Flow usage and cache key at Factory ingress

.github/workflows/run-tests.yaml

@@ -51,7 +51,7 @@ jobs:
         run: |
           source .venv/bin/activate
           python -m pytest -p pytest_asyncio.plugin -q --disable-warnings \
-            tests/test_factory_admission.py tests/test_factory_gateway.py \
+            tests/test_factory_admission.py tests/test_factory_gateway.py tests/test_factory_sdk_asgi_hold.py \
             tests/test_text_peer_request_guards.py tests/test_text_request_deadline.py \
             tests/test_openai_api.py tests/test_official_openai_client.py tests/test_text_mesh.py \
             -k 'not RealTransportTests'
@@ -93,6 +93,7 @@ jobs:
             tests/test_external_qualification.py \
             tests/test_factory_admission.py \
             tests/test_factory_gateway.py \
+            tests/test_factory_sdk_asgi_hold.py \
             tests/test_device_portability.py \
             tests/test_gate16_catalog_channel.py \
             tests/test_gate16_catalog_drill.py \

docs/FACTORY_DISTRIBUTED_INFERENCE.md

@@ -68,4 +68,8 @@ All required-confidential requests remain denied in this profile. Future qualifi
 
 ## Local validation scope
 
+Both completion endpoints recognize `stream_options: {"include_usage": true}` when `stream` is true. The supported options object requires an explicit boolean `include_usage` and rejects unknown keys. When requested, content and finish chunks carry `usage: null`, followed by one `choices: []` chunk containing actual generated token counts or validated peer usage, before `[DONE]`. Failed or interrupted generation does not emit a total-usage chunk. This follows the [OpenAI stream-usage format](https://developers.openai.com/api/reference/resources/chat#chatcompletionstreamoptions). With the option omitted or false, tensor streams retain their original format and text-peer streams retain their legacy usage on the finish chunk.
+
+Factory admission includes the parsed option in its normalized body digest, and the logical peer request retains that same field. Changing the option cannot reuse the original admission record. Updated Factory peers must recognize this schema; the ordinary mixed-version path strips this HTTP response-formatting option before sending to older peers. The option is excluded from generation sampling arguments. Token usage is inference metadata, not settled provider billing or evidence of confidential execution. These schema and SSE checks do not establish compatibility with FACTORY's original fixed journal or qualify the physical sender.
+
 Focused tests use HTTP/peer fixtures and synthetic original authority/claim readers. They exercise strict field refusal, missing/unknown/protected class, body/model mismatch, complete fresh block coverage, revocation after waits, repeated original POST refusal including unknown outcome, original JSON/SSE identity, no peer failover and retained cleanup. They do not execute paid jobs, contact a model, establish genuine Factory authority or prove hardware/channel protection. The real host adapter, physical send qualification, chosen manifest, independent participant formation and lifecycle evidence remain adoption blockers.

src/drift/api/server.py

@@ -27,7 +27,7 @@
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import StreamingResponse
 from hivemind.utils.logging import get_logger
-from pydantic import BaseModel, ConfigDict, StrictBool, StrictStr
+from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictStr
 from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 
 from drift.factory_admission import FactoryAdmission, RequestAdmissionDenied, synchronous_result
@@ -85,6 +85,12 @@ class ChatMessage(BaseModel):
     content: Union[StrictStr, List[TextContentPart]] = ""
 
 
+class StreamOptions(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    include_usage: StrictBool
+
+
 class ChatCompletionRequest(BaseModel):
     model_config = ConfigDict(extra="forbid")
     model: Optional[str] = None
@@ -95,6 +101,10 @@ class ChatCompletionRequest(BaseModel):
     top_p: Optional[float] = None
     stop: Optional[Union[str, List[str]]] = None
     stream: bool = False
+    stream_options: Optional[StreamOptions] = None
+    # An opaque SDK hint, bound into Factory's normalized body. This receiver
+    # does not implement prompt-cache behavior.
+    prompt_cache_key: Optional[StrictStr] = Field(default=None, min_length=1, max_length=256)
     n: int = 1
     # Optional extension for verified templates that expose a reasoning switch.
     # Omission preserves the model's own template default.
@@ -110,6 +120,7 @@ class CompletionRequest(BaseModel):
     top_p: Optional[float] = None
     stop: Optional[Union[str, List[str]]] = None
     stream: bool = False
+    stream_options: Optional[StreamOptions] = None
     n: int = 1
 
 
@@ -419,7 +430,14 @@ def usage(input_ids: torch.Tensor, output_ids: torch.Tensor) -> Dict[str, int]:
     def finish_reason(completion_tokens: int, gen_kwargs: Dict[str, Any]) -> str:
         return "length" if completion_tokens >= gen_kwargs["max_new_tokens"] else "stop"
 
-    async def sse_stream(loaded: LoadedModel, input_ids: torch.Tensor, gen_kwargs: Dict[str, Any], *, chat: bool):
+    async def sse_stream(
+        loaded: LoadedModel,
+        input_ids: torch.Tensor,
+        gen_kwargs: Dict[str, Any],
+        *,
+        chat: bool,
+        include_usage: bool = False,
+    ):
         """Yield OpenAI-format SSE chunks while generate() runs in a worker thread."""
         request_id = f"{'chatcmpl' if chat else 'cmpl'}-{uuid.uuid4().hex[:24]}"
         created = int(time.time())
@@ -432,6 +450,19 @@ def chunk(payload: Dict[str, Any], reason: Optional[str] = None) -> str:
             choice.update(payload)
             body = {"id": request_id, "object": object_name, "created": created, "model": selected_model}
             body["choices"] = [choice]
+            if include_usage:
+                body["usage"] = None
+            return f"data: {json.dumps(body)}\n\n"
+
+        def usage_chunk(output_ids: torch.Tensor) -> str:
+            body = {
+                "id": request_id,
+                "object": object_name,
+                "created": created,
+                "model": selected_model,
+                "choices": [],
+                "usage": usage(input_ids, output_ids),
+            }
             return f"data: {json.dumps(body)}\n\n"
 
         future = None
@@ -460,11 +491,15 @@ def chunk(payload: Dict[str, Any], reason: Optional[str] = None) -> str:
                 try:
                     output_ids = await asyncio.shield(future)
                     reason = finish_reason(output_ids.shape[1] - input_ids.shape[1], gen_kwargs)
+                    completed = True
                 except Exception as exc:
                     logger.exception("Generation failed mid-stream")
                     yield f"data: {json.dumps({'error': {'message': str(exc), 'type': 'server_error'}})}\n\n"
                     reason = "stop"
+                    completed = False
                 yield chunk({"delta": {}} if chat else {"text": ""}, reason)
+                if include_usage and completed:
+                    yield usage_chunk(output_ids)
                 yield "data: [DONE]\n\n"
         finally:
             cancelled.event.set()
@@ -509,6 +544,8 @@ async def list_models(request: Request):
     async def chat_completions(body: ChatCompletionRequest, request: Request):
         if body.n != 1:
             raise HTTPException(status_code=400, detail="n > 1 is not supported")
+        if body.stream_options is not None and not body.stream:
+            raise HTTPException(status_code=400, detail="stream_options requires stream=true")
         context = request_context(body, request)
         loaded = await load_with_budget(body.model, context)
         if context.single_attempt and (
@@ -547,7 +584,14 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
 
         if body.stream:
             return StreamingResponse(
-                sse_stream(loaded, input_ids, gen_kwargs, chat=True), media_type="text/event-stream"
+                sse_stream(
+                    loaded,
+                    input_ids,
+                    gen_kwargs,
+                    chat=True,
+                    include_usage=body.stream_options is not None and body.stream_options.include_usage,
+                ),
+                media_type="text/event-stream",
             )
 
         future = None
@@ -586,6 +630,8 @@ async def chat_completions(body: ChatCompletionRequest, request: Request):
     async def completions(body: CompletionRequest, request: Request):
         if body.n != 1:
             raise HTTPException(status_code=400, detail="n > 1 is not supported")
+        if body.stream_options is not None and not body.stream:
+            raise HTTPException(status_code=400, detail="stream_options requires stream=true")
         context = request_context(body, request)
         loaded = await load_with_budget(body.model, context)
         if context.single_attempt and (
@@ -624,7 +670,14 @@ async def completions(body: CompletionRequest, request: Request):
 
         if body.stream:
             return StreamingResponse(
-                sse_stream(loaded, input_ids, gen_kwargs, chat=False), media_type="text/event-stream"
+                sse_stream(
+                    loaded,
+                    input_ids,
+                    gen_kwargs,
+                    chat=False,
+                    include_usage=body.stream_options is not None and body.stream_options.include_usage,
+                ),
+                media_type="text/event-stream",
             )
 
         future = None

src/drift/api/text_response.py

@@ -24,7 +24,14 @@ async def text_peer_response(loaded, body, *, chat, semaphore, context=None):
     created = int(time.time())
     model_id = loaded.descriptor.manifest_digest if context.single_attempt else loaded.descriptor.model_id
     request = body.model_dump(exclude_none=True)
+    # Ordinary mixed-version peers retain their old request surface. Factory
+    # sends retain admitted fields in the worker request, but the cache key is
+    # only an opaque bound hint here; it does not enable caching.
+    if not context.single_attempt:
+        request.pop("stream_options", None)
+        request.pop("prompt_cache_key", None)
     request["model"] = loaded.descriptor.manifest_digest
+    include_usage = body.stream_options is not None and body.stream_options.include_usage
     events_started = False
 
     def error_detail(exc):
@@ -53,10 +60,28 @@ def chunk(text=None, *, done=None, role=False):
             "model": model_id,
             "choices": [choice],
         }
-        if done is not None:
+        if include_usage:
+            result["usage"] = None
+        elif done is not None:
             result["usage"] = done["usage"]
         return "data: " + json.dumps(result) + "\n\n"
 
+    def usage_chunk(done):
+        return (
+            "data: "
+            + json.dumps(
+                {
+                    "id": request_id,
+                    "object": "chat.completion.chunk" if chat else "text_completion",
+                    "created": created,
+                    "model": model_id,
+                    "choices": [],
+                    "usage": done["usage"],
+                }
+            )
+            + "\n\n"
+        )
+
     async def events():
         nonlocal events_started
         events_started = True
@@ -128,6 +153,8 @@ async def sse():
                 elif frame["type"] == "done":
                     done = True
                     yield chunk(done=frame)
+                    if include_usage:
+                        yield usage_chunk(frame)
                 else:
                     yield ": waiting for community\n\n"
             if not done:

tests/test_factory_gateway.py

@@ -18,11 +18,13 @@ class Client:
 
     def __init__(self):
         self.contexts = []
+        self.bodies = []
         self.fail = False
 
     async def stream(self, body, *, chat, context):
         context.begin_dispatch()
         self.contexts.append(context)
+        self.bodies.append(body)
         if self.fail:
             raise TextPeerOutcomeUnknown(context.request_id)
         yield {"type": "delta", "text": "answer"}
@@ -191,6 +193,54 @@ async def test_json_and_sse_use_original_identity_and_terminal_usage(self):
         self.assertTrue(all(context.single_attempt for context in self.peer.contexts))
         self.assertTrue(all(item.active_requests == 0 for item in self.manager.snapshots()))
 
+    async def test_stream_options_preserve_admitted_peer_body_and_original_identity(self):
+        for endpoint, schema, payload in (
+            ("/v1/chat/completions", ChatCompletionRequest, {"messages": [{"role": "user", "content": "hello"}]}),
+            ("/v1/completions", CompletionRequest, {"prompt": "hello"}),
+        ):
+            for include_usage in (False, True):
+                body = {
+                    "model": MANIFEST,
+                    **payload,
+                    "stream": True,
+                    "stream_options": {"include_usage": include_usage},
+                }
+                normalized = schema(**body).model_dump(exclude_none=True)
+                request_id = f"{len(self.claimed) + 1:032x}"
+                self.admission = self.record(normalized, request_id)
+                response = await self.post(body, path=endpoint)
+                self.assertEqual(response.status_code, 200)
+                self.assertEqual(self.peer.bodies[-1], normalized)
+                lines = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")]
+                self.assertEqual(lines[-1], "[DONE]")
+                frames = [json.loads(line) for line in lines[:-1]]
+                prefix = "chatcmpl-" if endpoint.endswith("chat/completions") else "cmpl-"
+                self.assertTrue(all(frame["id"] == prefix + request_id for frame in frames))
+                if include_usage:
+                    self.assertTrue(all(frame["usage"] is None for frame in frames[:-1]))
+                    self.assertEqual(frames[-2]["choices"][0]["finish_reason"], "stop")
+                    self.assertEqual(frames[-1]["choices"], [])
+                else:
+                    self.assertEqual(frames[-1]["choices"][0]["finish_reason"], "stop")
+                self.assertEqual(frames[-1]["usage"]["total_tokens"], 2)
+                self.assertEqual(self.peer.contexts[-1].request_id, request_id)
+        self.assertEqual(len(self.claimed), 4)
+        self.assertEqual(len(self.peer.contexts), 4)
+
+    async def test_changed_stream_options_cannot_reuse_original_body(self):
+        for endpoint, schema, payload in (
+            ("/v1/chat/completions", ChatCompletionRequest, {"messages": [{"role": "user", "content": "hello"}]}),
+            ("/v1/completions", CompletionRequest, {"prompt": "hello"}),
+        ):
+            admitted = {"model": MANIFEST, **payload, "stream": True, "stream_options": {"include_usage": True}}
+            self.admission = self.record(schema(**admitted).model_dump(exclude_none=True))
+            response = await self.post({**admitted, "stream_options": {"include_usage": False}}, path=endpoint)
+            self.assertEqual(response.status_code, 403)
+            self.assertEqual(self.auth_checks, [])
+            self.assertEqual(self.loads, [])
+            self.assertEqual(self.claimed, set())
+            self.assertEqual(self.peer.bodies, [])
+
     async def test_unknown_json_and_sse_do_not_claim_completion_or_retry(self):
         self.peer.fail = True
         response = await self.post()
@@ -200,10 +250,12 @@ async def test_unknown_json_and_sse_do_not_claim_completion_or_retry(self):
         self.assertEqual(detail["request_id"], REQUEST)
         self.assertFalse(detail["retryable"])
         self.body["stream"] = True
+        self.body["stream_options"] = {"include_usage": True}
         self.admission = self.record(ChatCompletionRequest(**self.body).model_dump(exclude_none=True), "c" * 32)
         response = await self.post()
         self.assertIn('"code": "inference_outcome_unknown"', response.text)
         self.assertNotIn('"finish_reason": "stop"', response.text)
+        self.assertNotIn('"choices": []', response.text)
         self.assertEqual(len(self.peer.contexts), 2)  # one entry per explicit fixture request
 
     async def test_plain_completions_share_original_gate_and_identity(self):