fix(oci): record finish state so [DONE] fallback keeps tool_calls reason

This commit is contained in:
Cursor Agent 2026-09-03 08:23:46 +00:00
parent 1cfc402763
commit 06db628d2d
No known key found for this signature in database
2 changed files with 76 additions and 1 deletions

View file

@ -748,6 +748,24 @@ class OCIStreamWrapper(CustomStreamWrapper):
def _with_stream_identity(self, parsed: ModelResponseStream) -> ModelResponseStream:
return self.model_response_creator(chunk={"choices": parsed.choices})
def _record_terminal_state(self, parsed: ModelResponseStream) -> None:
"""Mirror the finish state the parent wrapper's ``chunk_creator`` would
normally record, so its ``StopIteration`` fallback emits a trailing
chunk with the real finish reason instead of defaulting to ``"stop"``.
Without this, skipping ``[DONE]`` lets the source iterator exhaust
cleanly and the parent's ``finish_reason_handler()`` runs with
``received_finish_reason`` and ``tool_call`` both unset, producing a
trailing ``"stop"`` chunk that overwrites an earlier ``"tool_calls"``
finish reason for consumers (and ``stream_chunk_builder``) that keep
the last non-null value.
"""
for choice in parsed.choices:
if getattr(choice, "finish_reason", None):
self.received_finish_reason = choice.finish_reason
if getattr(choice.delta, "tool_calls", None):
self.tool_call = True
def chunk_creator(self, chunk: Any) -> ModelResponseStream | None:
if not isinstance(chunk, str):
raise ValueError(f"Chunk is not a string: {chunk}")
@ -780,8 +798,11 @@ class OCIStreamWrapper(CustomStreamWrapper):
if getattr(choice.delta, "content", None):
self._cohere_text_emitted = True
break
self._record_terminal_state(result)
return self._with_stream_identity(result)
return self._with_stream_identity(handle_generic_stream_chunk(dict_chunk))
generic_result: Final = handle_generic_stream_chunk(dict_chunk)
self._record_terminal_state(generic_result)
return self._with_stream_identity(generic_result)
__all__ = [

View file

@ -2041,3 +2041,57 @@ class TestOCIStreamWrapperDoneSentinel:
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
with pytest.raises(OCIError, match="Chunk cannot be parsed as JSON"):
wrapper.chunk_creator("data: [DONE] trailing garbage")
class TestOCIStreamWrapperRecordsFinishState:
"""After skipping ``[DONE]`` the source iterator ends normally, so
``CustomStreamWrapper.__next__`` falls back to ``finish_reason_handler()``
to emit a trailing chunk. That helper reads ``received_finish_reason`` and
``tool_call``, so ``chunk_creator`` must mirror both onto ``self`` or the
trailing chunk defaults to ``"stop"`` and overwrites an earlier
``"tool_calls"`` finish reason for consumers that keep the last one."""
def test_generic_tool_calls_terminal_chunk_records_finish_state(self):
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
tool_call_event = (
'data: {"index":0,"message":{"role":"ASSISTANT","content":null,'
'"toolCalls":[{"id":"call_1","type":"FUNCTION","name":"get_weather","arguments":"{}"}]},'
'"finishReason":"TOOL_CALLS"}'
)
wrapper.chunk_creator(tool_call_event)
assert wrapper.received_finish_reason == "tool_calls"
assert wrapper.tool_call is True
trailing = wrapper.finish_reason_handler()
assert trailing.choices[0].finish_reason == "tool_calls"
def test_generic_stop_terminal_chunk_records_finish_state(self):
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
wrapper.chunk_creator(_GENERIC_TERMINAL_EVENT)
assert wrapper.received_finish_reason == "stop"
assert wrapper.tool_call is False
trailing = wrapper.finish_reason_handler()
assert trailing.choices[0].finish_reason == "stop"
def test_cohere_terminal_chunk_records_finish_state(self):
wrapper = _make_stream_wrapper(_STREAM_COHERE_MODEL)
wrapper.chunk_creator(_COHERE_TERMINAL_EVENT)
assert wrapper.received_finish_reason == "stop"
trailing = wrapper.finish_reason_handler()
assert trailing.choices[0].finish_reason == "stop"
def test_non_terminal_generic_chunk_leaves_finish_state_untouched(self):
wrapper = _make_stream_wrapper(_STREAM_GENERIC_MODEL)
wrapper.chunk_creator(_GENERIC_TEXT_EVENT.format(text="hi"))
assert wrapper.received_finish_reason is None
assert wrapper.tool_call is False