fix(openai): align first-chunk fields and cover response combinations

Author: arekay-nv

AGENTS.md

@@ -351,6 +351,7 @@ See [Development Guide](docs/DEVELOPMENT.md) for full setup and workflow details
 ### Data Types & Serialization
 
 - **Core types** (`Query`, `QueryResult`, `StreamChunk`): `msgspec.Struct` with `frozen=True`, `array_like=True`, `gc=False`, `omit_defaults=True`
+- **Streaming response boundary**: Entry zero of each streamed `TextModelOutput` field represents the same first meaningful delta. Fields that start later need an empty first entry (`""` for text, `()` for tool-call batches). TPOT drops entry zero independently in every field; full-output consumers combine all entries. Text tails may be aggregated; tool-call batches retain their fragment indices for merging.
 - **Config types**: `pydantic.BaseModel` for validation
 - **Enums**: `str, Enum` pattern for serializable enums (e.g., `LoadPatternType`, `APIType`)
 - **Serialization**: `msgspec.json` for hot-path (ZMQ transport), `pydantic` for config

src/inference_endpoint/core/types.py

@@ -122,7 +122,11 @@ class TextModelOutput(
     """Structured output from a text model.
 
     Supports main output, optional reasoning (e.g. chain-of-thought), and tool calls.
-    Each field may be a string (non-streaming) or tuple of strings (streaming chunks).
+    Text fields are strings for non-streaming responses or tuples for streaming.
+    Entry zero in every streamed field belongs to the same first meaningful delta.
+    Fields that start later use an empty first entry ("" for text, () for tools).
+    Remaining entries hold the tail; text entries may be aggregated. Tool calls
+    are a flat tuple for non-streaming or a tuple of delta batches for streaming.
 
     AT-RISK (gc=False): Has mutable container field `tool_calls`. Any change that
     mutates `tool_calls` after construction or stores cyclic references in it
@@ -182,38 +186,13 @@ def text_after_first_chunk(self) -> str:
         For non-streaming (str fields), there is no "first chunk" concept so
         this returns an empty string.
 
-        Streamed tool-call chunks are merged after dropping the first tool-call
-        chunk when the response starts with a pure tool-call delta.
+        All streamed fields exclude entry zero, including tool-call batches.
         """
-        parts: list[str] = []
-        if self.reasoning:
-            if isinstance(self.reasoning, tuple) and len(self.reasoning) > 1:
-                parts.extend(self.reasoning[1:])
-            # str reasoning: single chunk, skip entirely (it IS the first chunk)
-        if self.output:
-            if isinstance(self.output, str):
-                # Non-streaming: if reasoning was present and was the first chunk,
-                # include the full output. Otherwise no first chunk to skip.
-                if parts or (self.reasoning and isinstance(self.reasoning, tuple)):
-                    parts.append(self.output)
-            elif isinstance(self.output, tuple):
-                if parts or self.reasoning:
-                    # First chunk was in reasoning; include all output chunks.
-                    parts.extend(self.output)
-                elif len(self.output) > 1:
-                    # No reasoning; first chunk is output[0], skip it.
-                    parts.extend(self.output[1:])
-        tool_calls = self.tool_calls
-        if (
-            tool_calls
-            and isinstance(tool_calls[0], list | tuple)
-            and not (self.output or self.reasoning)
-        ):
-            tool_calls = tool_calls[1:]
-        merged_tool_calls = merge_tool_calls(tool_calls)
-        if merged_tool_calls:
-            parts.append(msgspec.json.encode(list(merged_tool_calls)).decode())
-        return "".join(parts)
+        content, reasoning, tool_calls = self.as_message_parts_after_first_chunk()
+        text = (reasoning or "") + content
+        if tool_calls:
+            text += msgspec.json.encode(list(tool_calls)).decode()
+        return text
 
     def as_message_parts(
         self,
@@ -236,38 +215,17 @@ def as_message_parts(
     def as_message_parts_after_first_chunk(
         self,
     ) -> tuple[str, str | None, TOOL_CALLS_TYPE | None]:
-        """Return message parts emitted after the first stream chunk."""
-        reasoning_after: str | None = None
-        if isinstance(self.reasoning, tuple) and len(self.reasoning) > 1:
-            reasoning_after = "".join(self.reasoning[1:])
-
-        # has_reasoning_tail: reasoning[1:] is non-empty (mirrors `parts` logic in text_after_first_chunk)
-        has_reasoning_tail = reasoning_after is not None
-
-        content_after = ""
-        if self.output:
-            if isinstance(self.output, str):
-                # Include if reasoning is any non-empty tuple (it was the first chunk)
-                if has_reasoning_tail or (
-                    self.reasoning and isinstance(self.reasoning, tuple)
-                ):
-                    content_after = self.output
-            elif isinstance(self.output, tuple):
-                if has_reasoning_tail or self.reasoning:
-                    # First chunk was in reasoning; include all output chunks.
-                    content_after = "".join(self.output)
-                elif len(self.output) > 1:
-                    # No reasoning; first chunk is output[0], skip it.
-                    content_after = "".join(self.output[1:])
-
-        tool_calls = self.tool_calls
-        if (
-            tool_calls
-            and isinstance(tool_calls[0], list | tuple)
-            and not (self.output or self.reasoning)
-        ):
-            tool_calls = tool_calls[1:]
-        return content_after, reasoning_after, merge_tool_calls(tool_calls)
+        """Drop the aligned first entry in each streamed field for TPOT."""
+        content = "".join(self.output[1:]) if isinstance(self.output, tuple) else ""
+        reasoning = (
+            "".join(self.reasoning[1:]) or None
+            if isinstance(self.reasoning, tuple)
+            else None
+        )
+        tool_calls = None
+        if self.tool_calls and isinstance(self.tool_calls[0], list | tuple):
+            tool_calls = merge_tool_calls(self.tool_calls[1:])
+        return content, reasoning, tool_calls
 
 
 OUTPUT_TYPE = TextModelOutput

src/inference_endpoint/metrics/steady_state_diagnostics.py

@@ -586,9 +586,8 @@ def build_natl_result(
 def text_after_first_chunk(data: object) -> str:
     """Return output text after the first streamed chunk.
 
-    This is the text used as the TPOT numerator. ``data`` is the COMPLETE event
-    payload: ``[tag, output, reasoning?, tool_calls?]`` with trailing defaults
-    omitted by msgspec ``array_like`` and ``omit_defaults``.
+    This text is tokenized for the TPOT denominator. ``data`` is the COMPLETE
+    event payload: ``[tag, output, reasoning?, tool_calls?]``.
 
     ``output`` and ``reasoning`` are each either a string for non-streaming
     output or a list of streamed chunks. The logic mirrors
@@ -599,21 +598,10 @@ def text_after_first_chunk(data: object) -> str:
     output = data[1] if len(data) > 1 else ""
     reasoning = data[2] if len(data) > 2 else None
     parts: list[str] = []
-    if reasoning:
-        if isinstance(reasoning, list) and len(reasoning) > 1:
-            parts.extend(reasoning[1:])
-        # str reasoning is a single (first) chunk -> skip entirely
-    if output:
-        if isinstance(output, str):
-            # Non-streaming output: keep it only if a first chunk already lived in a
-            # (streaming) reasoning trace; otherwise the str output IS the first chunk.
-            if parts or (reasoning and isinstance(reasoning, list)):
-                parts.append(output)
-        elif isinstance(output, list):
-            if parts or reasoning:
-                parts.extend(output)
-            elif len(output) > 1:
-                parts.extend(output[1:])
+    if isinstance(reasoning, list):
+        parts.extend(reasoning[1:])
+    if isinstance(output, list):
+        parts.extend(output[1:])
     # Tool-call reconstruction is intentionally omitted: tool-call samples use a
     # chat-template tokenization path this diagnostic does not replicate.
     return "".join(parts)
@@ -1353,7 +1341,7 @@ def cov_pass_row(
     # CoV is undefined with fewer than two points. Report inconclusive, not
     # PASS, so a short or empty window cannot masquerade as steady.
     if len(values) < 2:
-        return {b: None for b in bounds}
+        return dict.fromkeys(bounds)
     c = cov(values)
     return {b: c <= b for b in bounds}
 

src/inference_endpoint/openai/accumulator.py

@@ -21,7 +21,7 @@
 from inference_endpoint.endpoint_client.accumulator_protocol import (
     SSEAccumulatorProtocol,
 )
-from inference_endpoint.openai.types import SSEChoice
+from inference_endpoint.openai.types import SSEChoice, SSEDelta
 
 
 class OpenAISSEAccumulator(SSEAccumulatorProtocol):
@@ -33,6 +33,7 @@ def __init__(self, query_id: str, stream_all_chunks: bool):
         self.tool_call_chunks: list[tuple[dict[str, Any], ...]] = []
         self._finish_reason: str | None = None
 
+        self._first_delta: SSEDelta | None = None
         self.first_chunk_sent = False
         self.query_id = query_id
         self.stream_all_chunks = stream_all_chunks
@@ -57,6 +58,8 @@ def add_chunk(self, choice: SSEChoice | None) -> StreamChunk | None:
         if delta.content:
             self.output_chunks.append(delta.content)
         content = (rc or "") + (delta.content or "")
+        if not self.first_chunk_sent and (content or delta.tool_calls):
+            self._first_delta = delta
         if not content and delta.tool_calls and not self.first_chunk_sent:
             # Pure tool-call delta with no text: emit a zero-length sentinel so
             # RECV_FIRST / TTFT fires for agentic responses that have no content.
@@ -86,32 +89,22 @@ def add_chunk(self, choice: SSEChoice | None) -> StreamChunk | None:
             return None
 
     def get_final_output(self) -> QueryResult:
-        tool_calls = tuple(self.tool_call_chunks) if self.tool_call_chunks else None
-
-        if self.reasoning_chunks:
-            resp_reasoning: list[str] = [self.reasoning_chunks[0]]
-            if len(self.reasoning_chunks) > 1:
-                resp_reasoning.append("".join(self.reasoning_chunks[1:]))
-            text_output = TextModelOutput(
-                output="".join(self.output_chunks),
-                reasoning=resp_reasoning,
-                tool_calls=tool_calls,
-            )
-        elif self.output_chunks:
-            resp_output: list[str] = [self.output_chunks[0]]
-            if len(self.output_chunks) > 1:
-                resp_output.append("".join(self.output_chunks[1:]))
-            text_output = TextModelOutput(
-                output=resp_output,
-                reasoning=None,
-                tool_calls=tool_calls,
-            )
-        else:
-            text_output = TextModelOutput(
-                output=[],
-                reasoning=None,
-                tool_calls=tool_calls,
-            )
+        first = self._first_delta
+        first_content = (first.content or "") if first else ""
+        first_reasoning = (
+            (first.reasoning_content or first.reasoning or "") if first else ""
+        )
+        tool_calls = None
+        if self.tool_call_chunks:
+            tool_calls = tuple(self.tool_call_chunks)
+            if first is not None and not first.tool_calls:
+                tool_calls = ((), *tool_calls)
+
+        text_output = TextModelOutput(
+            output=_with_first_chunk(self.output_chunks, first_content),
+            reasoning=_with_first_chunk(self.reasoning_chunks, first_reasoning) or None,
+            tool_calls=tool_calls,
+        )
 
         metadata: dict[str, Any] = {
             "first_chunk": not self.first_chunk_sent,
@@ -128,3 +121,13 @@ def get_final_output(self) -> QueryResult:
             response_output=text_output,
             metadata=metadata,
         )
+
+
+def _with_first_chunk(chunks: list[str], first: str) -> tuple[str, ...]:
+    """Keep the first delta's contribution separate from the accumulated tail."""
+    if not chunks:
+        return ()
+    tail_start = 1 if first else 0
+    if len(chunks) == tail_start:
+        return (first,)
+    return first, "".join(chunks[tail_start:])

tests/unit/async_utils/services/metrics_aggregator/test_aggregator.py

@@ -1389,7 +1389,9 @@ async def test_tpot_osl_for_tool_call_complete(self, tmp_path):
                             SampleEventType.COMPLETE,
                             "s1",
                             ts=5000,
-                            data=TextModelOutput(output="ok", tool_calls=(tool_call,)),
+                            data=TextModelOutput(
+                                output=("ok",), tool_calls=((), (tool_call,))
+                            ),
                         ),
                     ]
                 )

tests/unit/async_utils/services/metrics_aggregator/test_metrics_table.py

@@ -402,7 +402,7 @@ async def count_batch_async(self, inputs, _loop, live=False):
 class TestTpotTriggerToolCalls:
     """TpotTrigger routes to message path when tool_calls are present."""
 
-    async def test_tpot_tool_calls_only_response(self):
+    async def test_tpot_tools_after_first_content_delta(self):
         """TpotTrigger includes tool_calls in TPOT denominator for agentic responses."""
         from inference_endpoint.async_utils.services.metrics_aggregator.metrics_table import (
             SampleField,
@@ -428,7 +428,7 @@ async def test_tpot_tool_calls_only_response(self):
                 "function": {"name": "f", "arguments": "{}"},
             },
         )
-        tmo = TextModelOutput(output=[], tool_calls=tool_calls)
+        tmo = TextModelOutput(output=("checking",), tool_calls=((), tool_calls))
         ev = EventRecord(
             event_type=SampleEventType.COMPLETE,
             timestamp_ns=2000,