diff --git a/.github/workflows/run-tests.yaml b/.github/workflows/run-tests.yaml index 84d1af731..f48631911 100644 --- a/.github/workflows/run-tests.yaml +++ b/.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 \ diff --git a/docs/FACTORY_DISTRIBUTED_INFERENCE.md b/docs/FACTORY_DISTRIBUTED_INFERENCE.md index d248decd8..2fb0dc2b5 100644 --- a/docs/FACTORY_DISTRIBUTED_INFERENCE.md +++ b/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. diff --git a/src/drift/api/server.py b/src/drift/api/server.py index 00a3d91e5..c14de9608 100644 --- a/src/drift/api/server.py +++ b/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 diff --git a/src/drift/api/text_response.py b/src/drift/api/text_response.py index 08a30f0c5..cbffaf122 100644 --- a/src/drift/api/text_response.py +++ b/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: diff --git a/tests/test_factory_gateway.py b/tests/test_factory_gateway.py index 803de11a8..5b857c3a1 100644 --- a/tests/test_factory_gateway.py +++ b/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): diff --git a/tests/test_factory_sdk_asgi_hold.py b/tests/test_factory_sdk_asgi_hold.py new file mode 100644 index 000000000..cd6e51d7c --- /dev/null +++ b/tests/test_factory_sdk_asgi_hold.py @@ -0,0 +1,256 @@ +"""Offline SDK-wire-to-ASGI fixture; no original issuer or physical dispatch.""" + +import hashlib +import json +import unittest +from dataclasses import replace + +import httpx + +from drift.api.server import create_factory_app +from drift.factory_admission import FactoryAdmission, RequestAdmissionDenied +from drift.node.model_manager import ModelDescriptor, ModelManager, ModelRuntime + +MANIFEST = "sha256:" + "1" * 64 +# Copied from FACTORY's accepted owner-wire SDK run 438df4a3: actualWire.bodyUtf8 +# equals its independently precommitted wire. Keep these bytes and digest fixed; +# constructing a fresh JSON object here would lose the SDK-final wire evidence. +SDK_BODY_UTF8 = ( + b'{"model":"sha256:1111111111111111111111111111111111111111111111111111111111111111",' + b'"messages":[{"role":"user","content":"offline bridge fixture"}],' + b'"temperature":1,"stream":true,"max_tokens":8,' + b'"stream_options":{"include_usage":true}}' +) +SDK_BODY_SHA256 = "44c04671db48d10c164afcdf6558032c4ca68cb8f0a1c5e9ebdaeb7f449843b5" +# Independent R4 expectation for PR35's receiver-normalized, null-excluded body. +# In particular, PR35 inserts n=1 and renders temperature as a Python float. +NORMALIZED_UTF8 = ( + b'{"max_tokens":8,"messages":[{"content":"offline bridge fixture","role":"user"}],' + b'"model":"sha256:1111111111111111111111111111111111111111111111111111111111111111",' + b'"n":1,"stream":true,"stream_options":{"include_usage":true},"temperature":1.0}' +) +NORMALIZED_SHA256 = "21a0d9be9d903deba02504bb66506208338bc4423c8669a71db1904d21aa4890" +# Copied byte-for-byte from FACTORY's accepted real Flow R2 case-report +# actualWire.bodyUtf8. The Flow request adds a system message and a cache key; +# the key is only an opaque bound field at this receiver. +FLOW_SDK_BODY_UTF8 = ( + b'{"model":"sha256:1111111111111111111111111111111111111111111111111111111111111111",' + b'"messages":[{"role":"user","content":"offline bridge fixture"},' + b'{"role":"system","content":""}],"temperature":1,"stream":true,"max_tokens":8,' + b'"stream_options":{"include_usage":true},"prompt_cache_key":"flujo-c1c5k0kx"}' +) +FLOW_SDK_BODY_SHA256 = "b86e125b6a031a226acc03147b5d865402ad149a621626fcdbf7629f5541fd03" +FLOW_NORMALIZED_UTF8 = ( + b'{"max_tokens":8,"messages":[{"content":"offline bridge fixture","role":"user"},' + b'{"content":"","role":"system"}],' + b'"model":"sha256:1111111111111111111111111111111111111111111111111111111111111111",' + b'"n":1,"prompt_cache_key":"flujo-c1c5k0kx","stream":true,' + b'"stream_options":{"include_usage":true},"temperature":1.0}' +) +FLOW_NORMALIZED_SHA256 = "ab288a143d34de92156f01d0d4c4f5d05bbdc51eda4a335e16c8816205cd08b7" +# FACTORY's offline-request-1 is not a valid CommunityAI request identity. +SYNTHETIC_ASGI_REQUEST_ID = "d" * 32 + + +def canonical_bytes(body): + return json.dumps(body, ensure_ascii=False, allow_nan=False, sort_keys=True, separators=(",", ":")).encode("utf-8") + + +class HeldPeer: + supports_request_context = True + supports_single_attempt = True + + def __init__(self): + self.entries = 0 + self.after_claim = 0 + self.request_ids = [] + self.bodies = [] + + async def stream(self, body, *, chat, context): + self.entries += 1 + self.request_ids.append(context.request_id) + self.bodies.append(body) + context.begin_dispatch() # The fixture claim below always refuses. + self.after_claim += 1 + yield {"type": "done", "finish_reason": "stop", "usage": {}} + + +class SDKWireASGIHoldTests(unittest.IsolatedAsyncioTestCase): + async def asyncSetUp(self): + self.manager = ModelManager() + self.assertEqual(len(SDK_BODY_UTF8), 232) + self.assertEqual(hashlib.sha256(SDK_BODY_UTF8).hexdigest(), SDK_BODY_SHA256) + self.assertEqual(len(NORMALIZED_UTF8), 240) + self.assertEqual(hashlib.sha256(NORMALIZED_UTF8).hexdigest(), NORMALIZED_SHA256) + self.assertNotEqual(SDK_BODY_UTF8, NORMALIZED_UTF8) + self.assertEqual(len(FLOW_SDK_BODY_UTF8), 299) + self.assertEqual(hashlib.sha256(FLOW_SDK_BODY_UTF8).hexdigest(), FLOW_SDK_BODY_SHA256) + self.assertEqual(len(FLOW_NORMALIZED_UTF8), 307) + self.assertEqual(hashlib.sha256(FLOW_NORMALIZED_UTF8).hexdigest(), FLOW_NORMALIZED_SHA256) + self.assertNotEqual(FLOW_SDK_BODY_UTF8, FLOW_NORMALIZED_UTF8) + + self.peer = HeldPeer() + self.loads = [] + self.admissions = [] + self.auth_checks = [] + self.claim_attempts = 0 + + def load(): + self.loads.append(True) + return ModelRuntime(None, None, text_client=self.peer) + + self.manager.register(ModelDescriptor("Fixture", aliases=(MANIFEST,), manifest_digest=MANIFEST), load) + + def claim_hold(): + self.claim_attempts += 1 + raise RequestAdmissionDenied() + + self.synthetic_admission = FactoryAdmission( + request_id=SYNTHETIC_ASGI_REQUEST_ID, + body_digest=NORMALIZED_SHA256, + manifest_digest=MANIFEST, + num_blocks=1, + confidentiality_class="ordinary", + class_version=1, + authority=lambda: None, + route_snapshot=lambda: { + "manifest_digest": MANIFEST, + "status": "complete", + "covered_blocks": 1, + "total_blocks": 1, + "missing_blocks": [], + "chat_ready": True, + "peer_count": 1, + "last_updated_age": 0.0, + "replica_counts": [1], + }, + claim_dispatch=claim_hold, + ) + + def admission(body): + self.admissions.append(body) + return self.synthetic_admission + + def authenticate(key): + self.auth_checks.append(key) + return key == "fixture-key" + + self.app = create_factory_app( + model_manager=self.manager, + api_key_verifier=authenticate, + factory_admission=admission, + ) + + async def asyncTearDown(self): + self.manager.shutdown() + + async def post(self, raw): + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=self.app), base_url="http://test") as client: + return await client.post( + "/v1/chat/completions", + content=raw, + headers={"Content-Type": "application/json", "Authorization": "Bearer fixture-key"}, + ) + + async def test_exact_sdk_bytes_normalize_and_stop_before_peer_dispatch(self): + response = await self.post(SDK_BODY_UTF8) + + self.assertEqual(len(self.admissions), 1) + normalized = canonical_bytes(self.admissions[0]) + self.assertEqual(normalized, NORMALIZED_UTF8) + self.assertEqual(hashlib.sha256(normalized).hexdigest(), NORMALIZED_SHA256) + self.assertEqual(self.auth_checks, ["fixture-key"]) + self.assertEqual(self.loads, [True]) + self.assertEqual(response.status_code, 200) # SSE headers precede the held claim. + frames = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")] + self.assertEqual(frames[-1], "[DONE]") + error = next(json.loads(frame)["error"] for frame in frames[:-1] if '"error"' in frame) + self.assertEqual(error["code"], "request_not_dispatched") + self.assertEqual(error["request_id"], SYNTHETIC_ASGI_REQUEST_ID) + self.assertIs(error["retryable"], False) + self.assertEqual(self.claim_attempts, 1) + self.assertEqual(self.peer.entries, 1) + self.assertEqual(self.peer.request_ids, [SYNTHETIC_ASGI_REQUEST_ID]) + self.assertEqual(self.peer.after_claim, 0) + self.assertNotIn('"choices": []', response.text) + self.assertNotIn('"finish_reason": "stop"', response.text) + + async def test_changed_usage_bit_cannot_reuse_fixed_admission(self): + changed = SDK_BODY_UTF8.replace(b'"include_usage":true', b'"include_usage":false') + self.assertNotEqual(changed, SDK_BODY_UTF8) + response = await self.post(changed) + + self.assertEqual(response.status_code, 403) + self.assertEqual(len(self.admissions), 1) + self.assertNotEqual(canonical_bytes(self.admissions[0]), NORMALIZED_UTF8) + self.assertEqual(self.auth_checks, []) + self.assertEqual(self.loads, []) + self.assertEqual(self.claim_attempts, 0) + self.assertEqual(self.peer.entries, 0) + + async def test_real_flow_sdk_body_normalizes_and_stops_at_held_claim(self): + self.synthetic_admission = replace(self.synthetic_admission, body_digest=FLOW_NORMALIZED_SHA256) + response = await self.post(FLOW_SDK_BODY_UTF8) + + self.assertEqual(len(self.admissions), 1) + normalized = canonical_bytes(self.admissions[0]) + self.assertEqual(normalized, FLOW_NORMALIZED_UTF8) + self.assertEqual(hashlib.sha256(normalized).hexdigest(), FLOW_NORMALIZED_SHA256) + self.assertEqual(self.auth_checks, ["fixture-key"]) + self.assertEqual(self.loads, [True]) + self.assertEqual(response.status_code, 200) # SSE headers precede the held claim. + frames = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")] + self.assertEqual(frames[-1], "[DONE]") + error = next(json.loads(frame)["error"] for frame in frames[:-1] if '"error"' in frame) + self.assertEqual(error["code"], "request_not_dispatched") + self.assertEqual(error["request_id"], SYNTHETIC_ASGI_REQUEST_ID) + self.assertIs(error["retryable"], False) + self.assertEqual(self.claim_attempts, 1) + self.assertEqual(self.peer.entries, 1) + self.assertEqual(self.peer.request_ids, [SYNTHETIC_ASGI_REQUEST_ID]) + self.assertEqual(self.peer.bodies, [self.admissions[0]]) + self.assertEqual(self.peer.bodies[0]["prompt_cache_key"], "flujo-c1c5k0kx") + self.assertEqual(self.peer.after_claim, 0) + self.assertNotIn('"choices": []', response.text) + self.assertNotIn('"finish_reason": "stop"', response.text) + + async def test_changed_flow_cache_key_cannot_reuse_fixed_admission(self): + self.synthetic_admission = replace(self.synthetic_admission, body_digest=FLOW_NORMALIZED_SHA256) + changed = FLOW_SDK_BODY_UTF8.replace(b'"flujo-c1c5k0kx"', b'"different"') + response = await self.post(changed) + + self.assertEqual(response.status_code, 403) + self.assertEqual(len(self.admissions), 1) + self.assertNotEqual(canonical_bytes(self.admissions[0]), FLOW_NORMALIZED_UTF8) + self.assertEqual(self.auth_checks, []) + self.assertEqual(self.loads, []) + self.assertEqual(self.claim_attempts, 0) + self.assertEqual(self.peer.entries, 0) + + async def test_invalid_flow_cache_key_refuses_before_admission(self): + for invalid in (0, {}, [], "", "x" * 257): + with self.subTest(invalid=invalid): + changed = json.loads(FLOW_SDK_BODY_UTF8) + changed["prompt_cache_key"] = invalid + response = await self.post(json.dumps(changed).encode("utf-8")) + self.assertEqual(response.status_code, 422) + self.assertEqual(self.admissions, []) + self.assertEqual(self.auth_checks, []) + self.assertEqual(self.loads, []) + self.assertEqual(self.claim_attempts, 0) + self.assertEqual(self.peer.entries, 0) + + async def test_unknown_nested_option_refuses_before_admission(self): + changed = SDK_BODY_UTF8.replace( + b'"stream_options":{"include_usage":true}', + b'"stream_options":{"include_usage":true,"unexpected":true}', + ) + self.assertNotEqual(changed, SDK_BODY_UTF8) + response = await self.post(changed) + + self.assertEqual(response.status_code, 422) + self.assertEqual(self.admissions, []) + self.assertEqual(self.auth_checks, []) + self.assertEqual(self.loads, []) + self.assertEqual(self.claim_attempts, 0) + self.assertEqual(self.peer.entries, 0) diff --git a/tests/test_official_openai_client.py b/tests/test_official_openai_client.py index 051f90b4a..3a77ef04e 100644 --- a/tests/test_official_openai_client.py +++ b/tests/test_official_openai_client.py @@ -94,6 +94,26 @@ def test_official_python_client_models_completions_chat_and_streaming(): assert streamed.replace(" ", "") == "tok101tok102tok103" +@pytest.mark.parametrize("chat", [False, True]) +def test_official_python_client_consumes_requested_terminal_usage(chat): + app = create_app(_Model(), _Tokenizer(), model_name="local-test", api_keys=["local-key"]) + with _live_server(app) as base_url: + client = openai.OpenAI(api_key="local-key", base_url=base_url, max_retries=0, timeout=5) + arguments = dict(model="local-test", max_tokens=3, stream=True, stream_options={"include_usage": True}) + if chat: + stream = client.chat.completions.create(messages=[{"role": "user", "content": "hello"}], **arguments) + else: + stream = client.completions.create(prompt="hello", **arguments) + chunks = list(stream) + assert all(chunk.usage is None for chunk in chunks[:-1]) + assert chunks[-2].choices[0].finish_reason == "length" + assert chunks[-1].choices == [] + assert chunks[-1].usage.prompt_tokens == 2 + assert chunks[-1].usage.completion_tokens == 3 + assert chunks[-1].usage.total_tokens == 5 + assert {chunk.id for chunk in chunks} == {chunks[0].id} + + def test_official_python_client_observes_bearer_authentication(): app = create_app(_Model(), _Tokenizer(), model_name="local-test", api_keys=["local-key"]) with _live_server(app) as base_url: diff --git a/tests/test_openai_api.py b/tests/test_openai_api.py index 6ba9c9470..aa395828e 100644 --- a/tests/test_openai_api.py +++ b/tests/test_openai_api.py @@ -336,6 +336,117 @@ def test_chat_completion_stream(api): assert chunks[-1]["choices"][0]["finish_reason"] == "stop" +@pytest.mark.parametrize("include_usage", [False, True]) +@pytest.mark.parametrize( + "path,payload", + [ + ("/v1/chat/completions", {"messages": [{"role": "user", "content": "hi"}]}), + ("/v1/completions", {"prompt": "hi"}), + ], +) +def test_stream_options_emit_terminal_usage_only_when_requested(api, include_usage, path, payload): + response = api.client.post( + path, + json={ + **payload, + "stream": True, + "stream_options": {"include_usage": include_usage}, + }, + ) + assert response.status_code == 200 + lines = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")] + assert lines[-1] == "[DONE]" + chunks = [json.loads(line) for line in lines[:-1]] + if include_usage: + assert len([chunk for chunk in chunks if chunk["choices"] == []]) == 1 + assert all(chunk["usage"] is None for chunk in chunks[:-1]) + assert chunks[-2]["choices"][0]["finish_reason"] == "stop" + assert chunks[-1]["choices"] == [] + prompt_tokens = 3 if path == "/v1/chat/completions" else 1 + assert chunks[-1]["usage"] == { + "prompt_tokens": prompt_tokens, + "completion_tokens": 3, + "total_tokens": prompt_tokens + 3, + } + assert {chunk["id"] for chunk in chunks} == {chunks[0]["id"]} + else: + assert all("usage" not in chunk for chunk in chunks) + assert chunks[-1]["choices"][0]["finish_reason"] == "stop" + + +@pytest.mark.parametrize( + "options,stream,expected_status", + [ + ({}, True, 422), + ({"include_usage": "true"}, True, 422), + ({"include_usage": 1}, True, 422), + ({"include_usage": True, "unsupported": True}, True, 422), + ([], True, 422), + ({"include_usage": True}, False, 400), + ], +) +@pytest.mark.parametrize( + "path,payload", + [ + ("/v1/chat/completions", {"messages": [{"role": "user", "content": "hi"}]}), + ("/v1/completions", {"prompt": "hi"}), + ], +) +def test_invalid_stream_options_fail_before_model_load(options, stream, expected_status, path, payload): + loads = [] + manager = ModelManager() + + def load(): + loads.append(True) + raise AssertionError("invalid request must not load a model") + + manager.register(ModelDescriptor("lazy"), load) + client = TestClient(create_app(model_manager=manager)) + response = client.post( + path, + json={ + "model": "lazy", + **payload, + "stream": stream, + "stream_options": options, + }, + ) + assert response.status_code == expected_status + assert loads == [] + + +@pytest.mark.parametrize( + "path,payload", + [ + ("/v1/chat/completions", {"messages": [{"role": "user", "content": "hi"}]}), + ("/v1/completions", {"prompt": "hi"}), + ], +) +def test_failed_stream_does_not_claim_usage(path, payload): + class FailingAfterEnd(FakeModel): + def generate(self, input_ids, *, streamer, **kwargs): + streamer.put(input_ids) + streamer.end() + raise RuntimeError("offline generation failure") + + client = TestClient(create_app(FailingAfterEnd(), FakeTokenizer(), model_name="failing")) + response = client.post( + path, + json={ + **payload, + "stream": True, + "stream_options": {"include_usage": True}, + }, + ) + assert response.status_code == 200 + lines = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")] + chunks = [json.loads(line) for line in lines[:-1]] + assert lines[-1] == "[DONE]" + assert any("error" in chunk for chunk in chunks) + assert all(chunk.get("usage") is None for chunk in chunks) + assert not any(chunk.get("choices") == [] for chunk in chunks) + + def test_completions_endpoint(api): response = api.client.post("/v1/completions", json={"prompt": "one two three", "temperature": 0}) body = response.json() diff --git a/tests/test_text_mesh.py b/tests/test_text_mesh.py index 22fb284ef..62f482925 100644 --- a/tests/test_text_mesh.py +++ b/tests/test_text_mesh.py @@ -1,6 +1,7 @@ """Consumer acceptance: no local model, full text responses, real peer transport.""" import asyncio +import json import tempfile import unittest from pathlib import Path @@ -107,6 +108,73 @@ async def test_streaming_and_plain_completions_keep_openai_shape(self): self.assertEqual(result.json()["choices"][0]["text"], "Hello from peers.") self.assertEqual(self.peer.closed_requests, 2) + async def test_stream_options_format_validated_peer_usage_without_forwarding(self): + for endpoint, payload in ( + ("/v1/chat/completions", {"messages": [{"role": "user", "content": "Hi"}]}), + ("/v1/completions", {"prompt": "Hi"}), + ): + for include_usage in (False, True): + response = await self.client.post( + endpoint, + json={ + "model": "auto", + **payload, + "stream": True, + "stream_options": {"include_usage": include_usage}, + }, + ) + self.assertEqual(response.status_code, 200) + self.assertNotIn("stream_options", self.peer.requests[-1][0]) + lines = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")] + self.assertEqual(lines[-1], "[DONE]") + chunks = [json.loads(line) for line in lines[:-1]] + if include_usage: + self.assertTrue(all(chunk["usage"] is None for chunk in chunks[:-1])) + self.assertEqual(chunks[-2]["choices"][0]["finish_reason"], "stop") + self.assertEqual(chunks[-1]["choices"], []) + self.assertEqual( + chunks[-1]["usage"], {"prompt_tokens": 3, "completion_tokens": 3, "total_tokens": 6} + ) + self.assertEqual({chunk["id"] for chunk in chunks}, {chunks[0]["id"]}) + else: + self.assertEqual(chunks[-1]["choices"][0]["finish_reason"], "stop") + self.assertEqual(chunks[-1]["usage"]["total_tokens"], 6) + self.assertFalse(any(chunk["choices"] == [] for chunk in chunks)) + + async def test_prompt_cache_key_is_not_forwarded_to_ordinary_peers(self): + response = await self.client.post( + "/v1/chat/completions", + json={ + "model": "auto", + "messages": [{"role": "user", "content": "Hi"}], + "prompt_cache_key": "opaque-fixture-key", + }, + ) + self.assertEqual(response.status_code, 200) + self.assertNotIn("prompt_cache_key", self.peer.requests[-1][0]) + + async def test_failed_peer_stream_does_not_claim_usage(self): + async def fail_after_delta(body, *, chat): + yield {"type": "delta", "text": "partial"} + raise TextPeerUnavailable("peer stopped") + + self.peer.stream = fail_after_delta + response = await self.client.post( + "/v1/chat/completions", + json={ + "model": "auto", + "messages": [{"role": "user", "content": "Hi"}], + "stream": True, + "stream_options": {"include_usage": True}, + }, + ) + self.assertEqual(response.status_code, 200) + lines = [line[6:] for line in response.text.splitlines() if line.startswith("data: ")] + chunks = [json.loads(line) for line in lines[:-1]] + self.assertTrue(any("error" in chunk for chunk in chunks)) + self.assertFalse(any(chunk.get("choices") == [] for chunk in chunks)) + self.assertTrue(all(chunk.get("usage") is None for chunk in chunks)) + async def test_fallback_only_when_community_path_unavailable_then_returns(self): self.assertEqual(self.manager.resolve("auto").model_id, "Community") self.health.update(chat_ready=False)