Treat an unfinished history tail as a continuation too

deferred_tool_results may arrive with a non-empty prompt, so the absence of a
prompt cannot be the only test for a resumption. Reproduced: resuming a live
question that way replaced its search result with the earlier-question notice
while the deferred result arrived alongside it.

_is_resumption accepts either signal — no prompt, or a history ending with a
request the model has not answered or a response whose tool calls have no
returns. A settled history ends with the previous answer, so a genuinely new
question is unaffected.

A new prompt on top of an unanswered tail is ambiguous and now counts as a
continuation: compacting costs the answer if it is one, while not compacting
only costs a larger request.
This commit is contained in:
Yiorgis Gozadinos 2026-08-10 14:13:19 +03:00
parent 8eef0b3734
commit fa0c4a4f50
No known key found for this signature in database
2 changed files with 101 additions and 11 deletions

View file

@ -123,6 +123,31 @@ def _compact_old_tool_returns(
return compacted
def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool:
"""Whether this run continues a question rather than asking a new one.
Two signals, either of which is enough, because getting this wrong hands the
model a notice where its own evidence should be:
- no prompt: how pydantic-ai resumes for interruptions and suspensions.
- an unfinished tail: the history ends with a request the model has not
answered, or with a response whose tool calls have no returns yet. Deferred
tool results may arrive *with* a prompt, so the prompt alone is not enough.
A settled history ends with the previous answer, so a genuinely new question
is not mistaken for a continuation. The framework's own first-new-message
index would be better than either signal, but it is not public here.
"""
if prompt is None:
return True
if not messages:
return False
last = messages[-1]
if isinstance(last, ModelRequest):
return True
return any(isinstance(part, ToolCallPart) for part in last.parts)
def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool:
"""Whether the model's most recent response called one of these tools."""
for message in reversed(messages):
@ -157,14 +182,11 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
"""Start a run's own copy, and decide what counts as an earlier question.
A run carrying no prompt is continuing a question rather than asking one:
pydantic-ai resumes that way for deferred tool results, interruptions and
suspended responses. ``len(ctx.messages)`` would then count the live
question's own messages and hand the model a notice where its search
result should be, so compaction is switched off for the whole run
(``turn_start=0``). The absence of a prompt is the signal the three
resume shapes differ in message layout, and reading the layout is what
broke this in the first place.
On a resumption (see ``_is_resumption``) compaction is switched off for the
whole run: ``len(ctx.messages)`` would count the live question's own
messages and replace its evidence with the earlier-question notice, leaving
the model to answer with the evidence taken away. Failing this way costs a
larger request; failing the other way costs the answer.
"""
outer = getattr(ctx.deps, "state", None)
outer_state = outer if isinstance(outer, dict) else None
@ -181,7 +203,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
search_count=0,
request_count=0,
grace_requests_used=0,
turn_start=0 if ctx.prompt is None else len(ctx.messages),
turn_start=(
0 if _is_resumption(ctx.prompt, ctx.messages) else len(ctx.messages)
),
)
run_capability._sync_state()
return run_capability

View file

@ -988,7 +988,39 @@ def _wire_returns(sent: list[Any]) -> list[str]:
@pytest.mark.asyncio
async def test_a_new_question_compacts_the_previous_one(temp_db_path):
"""The baseline the resume cases are contrasted against."""
"""The baseline the resume cases are contrasted against.
A settled history the previous question answered is what a genuinely new
question follows. An unfinished tail is ambiguous instead, and treated as a
continuation.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
wire: list[list[Any]] = []
async def model(messages, _info):
wire.append(list(messages))
return ModelResponse(parts=[TextPart("answer")])
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability])
settled = [*_in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
await agent.run("a different question", deps=Deps(), message_history=settled)
assert _wire_returns(wire[-1]) == [PRIOR_TURN_NOTICE]
@pytest.mark.asyncio
async def test_a_prompt_on_an_unanswered_tail_is_treated_as_a_continuation(
temp_db_path,
):
"""Ambiguous shape, resolved the safe way.
A history ending in a request the model never answered, plus a new prompt,
could be a fresh question or a continuation. Compacting would cost the answer
if it is a continuation; not compacting only costs a larger request.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
@ -1004,7 +1036,41 @@ async def test_a_new_question_compacts_the_previous_one(temp_db_path):
"a different question", deps=Deps(), message_history=_in_flight_history()
)
assert _wire_returns(wire[-1]) == [PRIOR_TURN_NOTICE]
assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
@pytest.mark.asyncio
async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path):
"""Deferred results may arrive with a prompt, and that is still a continuation.
``ctx.prompt`` is non-null here, so the absence of a prompt cannot be the only
signal: the history tail is unfinished a response whose tool call has no
return yet and the question it belongs to is still being answered.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
wire: list[list[Any]] = []
async def model(messages, _info):
wire.append(list(messages))
return ModelResponse(parts=[TextPart("answer")])
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability])
history = [
*_in_flight_history(),
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
]
await agent.run(
"carry on",
deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}),
message_history=history,
deps=Deps(),
)
assert "EVIDENCE FOR THE LIVE TURN" in _wire_returns(wire[-1])
assert PRIOR_TURN_NOTICE not in _wire_returns(wire[-1])
@pytest.mark.asyncio