streaming cost updates

[dixitaniket]

• cost for streaming response

[Copilot]

Pull request overview

Adds a streaming-only billing side-channel to /v1/ohttp chunked responses so the relay can learn settled cost after the stream completes.

Changes:

• Introduces a private billing-frame marker (OHTTP_BILLING_FRAME_MAGIC) and _build_billing_frame() to serialize SessionCost into a plaintext frame.
• Updates the streaming response generator to parse the final SSE JSON, set inner cost context, and (optionally) inject the billing frame before emitting the final encrypted chunk.

Comments suppressed due to low confidence (1)

tee_gateway/controllers/ohttp_controller.py:297

• response_json is parsed from b"".join(plaintext_chunks) here, and then _set_inner_stream_cost_context() in the finally block joins and parses the same chunks again. For large streams this duplicates O(n) copying/parsing and retains two full concatenations in memory. Consider refactoring so the SSE parse happens once (e.g., compute response_json once and pass it into the final context setter, or skip the _set_inner_stream_cost_context call when response_json has already been derived).

            response_json = _parse_final_sse_json(b"".join(plaintext_chunks))
            _set_inner_cost_context(
                cost_context,
                response_json=response_json,
                status_code=status,
            )
            billing_frame = _build_billing_frame(response_json)
            if billing_frame:
                yield billing_frame

            # Always emit exactly one final chunk so the AAD=b"final"
            # marker is present — that's what protects clients from
            # undetected truncation.
            yield encrypter.encrypt_chunk(pending or b"", is_final=True)
        finally:
            _set_inner_stream_cost_context(
                cost_context, plaintext_chunks, status_code=status
            )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.