fix(chat): persist reasoning for Responses stream and WS final snapshot

- Parse reasoning_summary_text (and related) SSE events in OpenAI Responses
  streaming; forward deltas to on_token and return reasoning_content on
  LLMResponse so thinking-mode rows get event_payload in the store.
- When emitting chat/final from turn_runner, include event_type/event_payload
  and merge reasoning_content from earlier assistant rows in the same
  turn_uuid so the collapsed reasoning fold matches what streamed.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
oliver 2026-05-12 20:18:15 +08:00
parent 46ebfeb54f
commit a3365b77fa
3 changed files with 164 additions and 11 deletions

View file

@ -556,7 +556,18 @@ async def run_agent_turn_via_bridge(
pass
elif run_status != "failed":
try:
persisted = store.get_messages(session_id=session_id, limit=64)
def _event_payload_as_dict(raw: Any) -> dict[str, Any]:
if raw is None or not str(raw).strip():
return {}
if isinstance(raw, dict):
return dict(raw)
try:
obj = json.loads(str(raw))
except Exception:
return {}
return dict(obj) if isinstance(obj, dict) else {}
persisted = store.get_messages(session_id=session_id, limit=200)
for m in reversed(list(persisted or [])):
if str(getattr(m, "role", "") or "").lower() != "assistant":
continue
@ -565,6 +576,8 @@ async def run_agent_turn_via_bridge(
if not content.strip() and not _persisted_chat_attachments_nonempty(atraw):
continue
final_text = content
et = str(getattr(m, "event_type", "") or "").strip()
ep = getattr(m, "event_payload", None)
final_msg = {
"id": int(getattr(m, "id", 0) or 0),
"role": "assistant",
@ -573,6 +586,31 @@ async def run_agent_turn_via_bridge(
"tool_calls": getattr(m, "tool_calls", None),
"attachments": atraw,
}
# Admin / WS clients build the reasoning fold from event_payload.reasoning_content (thinking mode)
# or event_type=reasoning rows after reload; including these on ``final`` avoids "stream had
# reasoning but terminal bubble only shows tools" when the UI hydrates from this snapshot.
if et:
final_msg["event_type"] = et
if ep is not None and str(ep).strip():
final_msg["event_payload"] = ep
# Thinking-mode: reasoning_content is often stored on an earlier assistant row in the same turn
# (e.g. first tool_call snapshot) while the newest row is post-tool prose only — merge it in.
ep_d = _event_payload_as_dict(final_msg.get("event_payload"))
rc0 = str(ep_d.get("reasoning_content") or "").strip()
turn_key = str(getattr(m, "turn_uuid", "") or "").strip()
if not rc0 and turn_key:
for m2 in reversed(list(persisted or [])):
if str(getattr(m2, "role", "") or "").lower() != "assistant":
continue
if str(getattr(m2, "turn_uuid", "") or "").strip() != turn_key:
continue
ep2 = _event_payload_as_dict(getattr(m2, "event_payload", None))
rc2 = str(ep2.get("reasoning_content") or "").strip()
if rc2:
merged = dict(ep_d)
merged["reasoning_content"] = rc2
final_msg["event_payload"] = json.dumps(merged, ensure_ascii=False)
break
break
except Exception:
pass

View file

@ -185,20 +185,88 @@ def _collect_tool_calls_from_response_dict(resp: dict[str, Any]) -> list[LLMTool
return out
def _stream_reasoning_delta_text(ev: dict[str, Any]) -> str:
"""Extract a textual chunk from a Responses streaming reasoning-related event."""
delta = ev.get("delta")
if isinstance(delta, dict):
for k in ("text", "delta", "summary", "content"):
v = delta.get(k)
if isinstance(v, str) and v:
return v
return ""
if delta is not None:
s = str(delta)
return s
for k in ("text", "summary", "content"):
v = ev.get(k)
if isinstance(v, str) and v:
return v
return ""
def _reasoning_from_completed_response(resp: dict[str, Any]) -> str:
"""Best-effort reasoning string from a terminal Responses ``response`` object."""
for key in ("reasoning", "reasoning_summary", "reasoning_text", "reasoning_content"):
v = resp.get(key)
if isinstance(v, str) and v.strip():
return v.strip()
items = resp.get("output")
if not isinstance(items, list):
return ""
chunks: list[str] = []
for it in items:
d = _as_dict(it) or {}
typ = str(d.get("type") or "").lower()
if "reasoning" not in typ:
continue
for k in ("summary", "text", "content"):
x = d.get(k)
if isinstance(x, str) and x.strip():
chunks.append(x.strip())
break
return "\n".join(chunks).strip()
def parse_openai_responses_stream_events(
events: Iterable[Any],
*,
on_token: Optional[Callable[[str], None]] = None,
) -> tuple[str, list[LLMToolCall], dict[str, Any] | None]:
) -> tuple[str, list[LLMToolCall], dict[str, Any] | None, str]:
"""Pure stream parser (offline-testable).
Returns: (assembled_text, tool_calls, final_response_dict?)
Returns: (assembled_output_text, tool_calls, final_response_dict?, reasoning_text)
Reasoning models emit ``response.reasoning_summary_text.delta`` (and similar) events; those are
forwarded to ``on_token`` for live UI, accumulated into ``reasoning_text`` for persistence, and
are **not** merged into the returned ``assembled_output_text`` (assistant body).
"""
parts: list[str] = []
reasoning_parts: list[str] = []
final_resp: dict[str, Any] | None = None
def _emit_reasoning_chunk(s: str) -> None:
s = str(s or "")
if not s:
return
reasoning_parts.append(s)
if on_token:
on_token(s)
for ev in events:
d = _as_dict(ev) or {}
typ = str(d.get("type") or "")
# Reasoning / thinking summary stream (do not append to output_text body).
if typ in (
"response.reasoning_summary_text.delta",
"response.reasoning_summary_text",
"response.reasoning_text.delta",
"response.reasoning_text",
"response.reasoning_summary.delta",
) or (typ.startswith("response.reasoning") and typ.endswith(".delta")):
rs = _stream_reasoning_delta_text(d)
if rs:
_emit_reasoning_chunk(rs)
continue
# Text deltas
if typ in ("response.output_text.delta", "response.output_text"):
delta = d.get("delta")
@ -219,6 +287,10 @@ def parse_openai_responses_stream_events(
parts.append(txt)
if on_token:
on_token(txt)
for rkey in ("reasoning_summary_text", "reasoning_text", "reasoning"):
rv = delta.get(rkey)
if isinstance(rv, str) and rv:
_emit_reasoning_chunk(rv)
continue
# Terminal response object
if typ in ("response.completed", "response.complete", "response.done"):
@ -233,7 +305,10 @@ def parse_openai_responses_stream_events(
text = "".join(parts)
tool_calls = _collect_tool_calls_from_response_dict(final_resp) if final_resp else []
return text, tool_calls, final_resp
reasoning_text = "".join(reasoning_parts).strip()
if not reasoning_text and final_resp:
reasoning_text = _reasoning_from_completed_response(final_resp)
return text, tool_calls, final_resp, reasoning_text
class OpenAIResponsesModel(ChatModel):
@ -606,7 +681,9 @@ class OpenAIResponsesModel(ChatModel):
_log_openai_responses_wire_kwargs(payload, variant=str(cand_label), phase="stream")
try:
stream = self._client.responses.create(**payload)
text, tool_calls, final_resp = parse_openai_responses_stream_events(stream, on_token=on_token)
text, tool_calls, final_resp, reasoning_text = parse_openai_responses_stream_events(
stream, on_token=on_token
)
if (not text.strip()) and final_resp:
ot = final_resp.get("output_text")
if isinstance(ot, str) and ot.strip():
@ -615,7 +692,7 @@ class OpenAIResponsesModel(ChatModel):
on_token(text)
if cand_label != input_candidates[0][0]:
logger.info("openai_responses: succeeded with input variant %s", cand_label)
return LLMResponse(content=text, tool_calls=tool_calls)
return LLMResponse(content=text, tool_calls=tool_calls, reasoning_content=reasoning_text)
except Exception as exc:
emsg = str(exc)
# Fallback: provider thinking-mode replay contract.
@ -633,7 +710,9 @@ class OpenAIResponsesModel(ChatModel):
phase="stream_retry_thinking",
)
stream = self._client.responses.create(**forced)
text, tool_calls, final_resp = parse_openai_responses_stream_events(stream, on_token=on_token)
text, tool_calls, final_resp, reasoning_text = parse_openai_responses_stream_events(
stream, on_token=on_token
)
if (not text.strip()) and final_resp:
ot = final_resp.get("output_text")
if isinstance(ot, str) and ot.strip():
@ -642,7 +721,7 @@ class OpenAIResponsesModel(ChatModel):
on_token(text)
if cand_label != input_candidates[0][0]:
logger.info("openai_responses: succeeded with input variant %s", cand_label)
return LLMResponse(content=text, tool_calls=tool_calls)
return LLMResponse(content=text, tool_calls=tool_calls, reasoning_content=reasoning_text)
except Exception:
pass
if _env_truthy("AIA_RESPONSES_LOG_API_ERROR_DETAIL"):
@ -675,11 +754,12 @@ class OpenAIResponsesModel(ChatModel):
d = _as_dict(resp) or {}
text = str(d.get("output_text") or "")
tool_calls = _collect_tool_calls_from_response_dict(d)
reasoning_text = _reasoning_from_completed_response(d)
if on_token and text:
on_token(text)
if cand_label != input_candidates[0][0]:
logger.info("openai_responses non-stream: succeeded with input variant %s", cand_label)
return LLMResponse(content=text, tool_calls=tool_calls)
return LLMResponse(content=text, tool_calls=tool_calls, reasoning_content=reasoning_text)
except Exception as e2:
emsg2 = str(e2)
if "reasoning_content" in emsg2 and "thinking mode" in emsg2 and "must be passed back" in emsg2:
@ -699,6 +779,7 @@ class OpenAIResponsesModel(ChatModel):
d = _as_dict(resp) or {}
text = str(d.get("output_text") or "")
tool_calls = _collect_tool_calls_from_response_dict(d)
reasoning_text = _reasoning_from_completed_response(d)
if on_token and text:
on_token(text)
if cand_label != input_candidates[0][0]:
@ -706,7 +787,7 @@ class OpenAIResponsesModel(ChatModel):
"openai_responses non-stream: succeeded with input variant %s",
cand_label,
)
return LLMResponse(content=text, tool_calls=tool_calls)
return LLMResponse(content=text, tool_calls=tool_calls, reasoning_content=reasoning_text)
except Exception:
pass
if _env_truthy("AIA_RESPONSES_LOG_API_ERROR_DETAIL"):

View file

@ -25,12 +25,46 @@ def test_openai_responses_stream_parses_text_deltas_and_final_response_tool_call
},
]
buf: list[str] = []
text, tool_calls, final_resp = parse_openai_responses_stream_events(events, on_token=buf.append)
text, tool_calls, final_resp, reasoning_text = parse_openai_responses_stream_events(events, on_token=buf.append)
assert text == "Hello"
assert "".join(buf) == "Hello"
assert reasoning_text == ""
assert final_resp and final_resp.get("id") == "resp_1"
assert len(tool_calls) == 1
assert tool_calls[0].id == "call_1"
assert tool_calls[0].name == "query_route"
assert tool_calls[0].arguments.get("destination") == "1.1.1.1"
def test_openai_responses_stream_emits_reasoning_summary_deltas() -> None:
events = [
{"type": "response.reasoning_summary_text.delta", "delta": "step "},
{"type": "response.reasoning_summary_text.delta", "delta": "one"},
{"type": "response.output_text.delta", "delta": "Hi"},
{"type": "response.completed", "response": {"id": "resp_2", "output_text": "Hi"}},
]
buf: list[str] = []
text, tool_calls, final_resp, reasoning_text = parse_openai_responses_stream_events(events, on_token=buf.append)
assert text == "Hi"
assert reasoning_text == "step one"
assert "".join(buf) == "step oneHi"
assert final_resp and final_resp.get("id") == "resp_2"
assert tool_calls == []
def test_openai_responses_stream_reasoning_fallback_from_completed_output() -> None:
events = [
{"type": "response.output_text.delta", "delta": "Z"},
{
"type": "response.completed",
"response": {
"id": "resp_3",
"output_text": "Z",
"output": [{"type": "reasoning", "summary": "silent plan"}],
},
},
]
text, _, _, reasoning_text = parse_openai_responses_stream_events(events)
assert text == "Z"
assert reasoning_text == "silent plan"