[object Object]

← back to Exo

Ciaran/gpt oss tool call finish reason (#1673)

d01636100a8d5211f917feaf5a06aa326446eeec · 2026-03-06 21:09:20 +0000 · ciaranbor

## Motivation

When gpt-oss truncates a response mid-tool-call (e.g. hitting max
tokens), the finish_reason was swallowed because the parser was in
tool-accumulation mode and never yielded it back to the caller. This
caused the API to hang or fail to signal completion.

## Changes

In parse_gpt_oss, when accumulating tool call arguments and a
finish_reason arrives, yield the response (with empty text) so the
finish reason propagates

## Why It Works

The parser now checks for finish_reason on every chunk while inside a
tool call. When the model stops generating (truncation), the finish
signal is forwarded immediately rather than being silently consumed.

## Test Plan

### Automated Testing

Added tests covering truncated tool calls (finish_reason="length") and
plain text truncation

Files touched

Diff

commit d01636100a8d5211f917feaf5a06aa326446eeec
Author: ciaranbor <81697641+ciaranbor@users.noreply.github.com>
Date:   Fri Mar 6 21:09:20 2026 +0000

    Ciaran/gpt oss tool call finish reason (#1673)
    
    ## Motivation
    
    When gpt-oss truncates a response mid-tool-call (e.g. hitting max
    tokens), the finish_reason was swallowed because the parser was in
    tool-accumulation mode and never yielded it back to the caller. This
    caused the API to hang or fail to signal completion.
    
    ## Changes
    
    In parse_gpt_oss, when accumulating tool call arguments and a
    finish_reason arrives, yield the response (with empty text) so the
    finish reason propagates
    
    ## Why It Works
    
    The parser now checks for finish_reason on every chunk while inside a
    tool call. When the model stops generating (truncation), the finish
    signal is forwarded immediately rather than being silently consumed.
    
    ## Test Plan
    
    ### Automated Testing
    
    Added tests covering truncated tool calls (finish_reason="length") and
    plain text truncation
---
 .../runner/llm_inference/model_output_parsers.py   |  3 +
 .../unittests/test_runner/test_parse_gpt_oss.py    | 84 +++++++++++++++++++++-
 2 files changed, 85 insertions(+), 2 deletions(-)

diff --git a/src/exo/worker/runner/llm_inference/model_output_parsers.py b/src/exo/worker/runner/llm_inference/model_output_parsers.py
index 99e4b6f2..bfdd6e4a 100644
--- a/src/exo/worker/runner/llm_inference/model_output_parsers.py
+++ b/src/exo/worker/runner/llm_inference/model_output_parsers.py
@@ -116,6 +116,9 @@ def parse_gpt_oss(
         if current_tool_name is not None:
             if delta:
                 tool_arg_parts.append(delta)
+            if response.finish_reason is not None:
+                yield response.model_copy(update={"text": "".join(tool_arg_parts)})
+                tool_arg_parts = []
             continue
 
         if ch == "analysis" and not thinking:
diff --git a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
index 90914ea0..f83dcd36 100644
--- a/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
+++ b/src/exo/worker/tests/unittests/test_runner/test_parse_gpt_oss.py
@@ -1,5 +1,6 @@
 from collections.abc import Generator
 
+from exo.shared.types.api import FinishReason
 from exo.shared.types.worker.runner_response import (
     GenerationResponse,
     ToolCallResponse,
@@ -83,6 +84,7 @@ THINKING_THEN_TOOL_TOKENS: list[tuple[int, str]] = [
 
 def _make_gen_responses(
     tokens: list[tuple[int, str]],
+    last_finish_reason: FinishReason = "stop",
 ) -> list[GenerationResponse]:
     """Build GenerationResponse list from (token_id, text) pairs."""
     responses: list[GenerationResponse] = []
@@ -92,7 +94,7 @@ def _make_gen_responses(
             GenerationResponse(
                 text=text,
                 token=tid,
-                finish_reason="stop" if is_last else None,
+                finish_reason=last_finish_reason if is_last else None,
                 usage=None,
             )
         )
@@ -101,11 +103,12 @@ def _make_gen_responses(
 
 def _collect(
     tokens: list[tuple[int, str]],
+    last_finish_reason: FinishReason = "stop",
 ) -> list[GenerationResponse | ToolCallResponse]:
     """Feed tokens through parse_gpt_oss and collect all yielded responses."""
 
     def _gen() -> Generator[GenerationResponse, None, None]:
-        yield from _make_gen_responses(tokens)
+        yield from _make_gen_responses(tokens, last_finish_reason)
 
     return list(x for x in parse_gpt_oss(_gen()) if x is not None)
 
@@ -171,3 +174,80 @@ class TestParseGptOssThinkingThenToolCall:
         tc = _get_tool_call(results)
         assert tc.tool_calls[0].name == "get_current_weather"
         assert "Tokyo" in tc.tool_calls[0].arguments
+
+
+# fmt: off
+# Truncated tool call: recipient + channel + message + partial args, no <|call|>
+TRUNCATED_TOOL_CALL_TOKENS: list[tuple[int, str]] = [
+    (316,    " to"),
+    (28,     "="),
+    (44580,  "functions"),
+    (775,    ".get"),
+    (23981,  "_current"),
+    (170154, "_weather"),
+    (_CHANNEL, "<|channel|>"),
+    (12606,  "comment"),
+    (815,    "ary"),
+    (5701,   " json"),
+    (_MESSAGE, "<|message|>"),
+    (10848,  '{"'),
+    (7693,   "location"),
+    (1243,   '":'),
+    (392,    ' "'),
+    (173844, "Tokyo"),
+    # No <|call|> — generation truncated here
+]
+
+# Plain text tokens (no tool call)
+PLAIN_TEXT_TOKENS: list[tuple[int, str]] = [
+    (_CHANNEL, "<|channel|>"),
+    (35644,  "analysis"),
+    (_MESSAGE, "<|message|>"),
+    (12845,  "Let"),
+    (668,    " me"),
+    (2411,   " think"),
+    (1078,   " about"),
+    (495,    " this"),
+    (13,     "."),
+    (_END,   "<|end|>"),
+    (_START, "<|start|>"),
+    (_ASSISTANT, "assistant"),
+    (_CHANNEL, "<|channel|>"),
+    (12606,  "comment"),
+    (815,    "ary"),
+    (_MESSAGE, "<|message|>"),
+    (9906,   "Hello"),
+    (14,     ","),
+    (2989,   " world"),
+]
+# fmt: on
+
+
+class TestParseGptOssMaxTokensTruncation:
+    """Truncated tool calls must still yield finish_reason."""
+
+    def test_truncated_tool_call_yields_finish_reason(self):
+        results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length")
+        gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
+        finish_reasons = [
+            r.finish_reason for r in gen_responses if r.finish_reason is not None
+        ]
+        assert "length" in finish_reasons
+
+    def test_truncated_tool_call_emits_partial_args(self):
+        results = _collect(TRUNCATED_TOOL_CALL_TOKENS, last_finish_reason="length")
+        gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
+        last = [r for r in gen_responses if r.finish_reason is not None][-1]
+        assert len(last.text) > 0
+
+    def test_truncated_plain_text_still_works(self):
+        results = _collect(PLAIN_TEXT_TOKENS, last_finish_reason="length")
+        gen_responses = [r for r in results if isinstance(r, GenerationResponse)]
+        finish_reasons = [
+            r.finish_reason for r in gen_responses if r.finish_reason is not None
+        ]
+        assert "length" in finish_reasons
+        # Verify non-empty text was yielded (delta text differs from raw token text
+        # due to Harmony encoding, so we just check something was emitted)
+        all_text = "".join(r.text for r in gen_responses)
+        assert len(all_text) > 0

← 79c8dbea Ciaran/minor download bugs (#1674)  ·  back to Exo  ·  Implement continuous batching (#1642) 131ad0ff →