From 47b74e48147a24b2031770cf24392b8aa27bdab5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 10:22:12 +0000 Subject: [PATCH 1/8] Agents: a turn's final reply commits with its answer, in one flush A turn made 2r+1 commits of the key's database (r model requests): its prompt, each reply and each batch of tool results, and then the final reply and the answer one after the other, with nothing but memory between them. Each commit is two state writes, a run and a round trip. In `answer_once`'s run, a step that no tool waits for (a reply without tool calls, ADK's end-of-agent event) is now held in memory, and written by the next write in its transaction: so the final reply commits with the turn's answer. Replies with tool calls still commit before their tools run, and tool results before the model is asked again. What is held is written before a publish, before a call record and LangGraph's nodes' calls, and when the attempt fails; a crash loses it, and the retried attempt asks the model again. A turn makes one commit fewer in every framework, and two fewer with ADK: a plain turn 3 to 2 (ADK 4 to 2), a turn with one round of tool calls 5 to 4 (ADK 6 to 4). Co-Authored-By: Claude Fable 5.2 --- python/restate/agents/_calls.py | 8 +- python/restate/agents/_common.py | 21 +++ python/restate/agents/_conversation.py | 147 ++++++++++++++--- python/restate/agents/_turn.py | 6 +- python/restate/agents/langgraph/_graph.py | 4 +- python/restate/agents/langgraph/nodes.py | 12 +- tests/agents_adk.py | 48 ++++++ tests/agents_conversation.py | 189 ++++++++++++++++++++++ tests/agents_pydantic_ai.py | 34 ++++ 9 files changed, 442 insertions(+), 27 deletions(-) diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py index c10369a0..0d508aaa 100644 --- a/python/restate/agents/_calls.py +++ b/python/restate/agents/_calls.py @@ -46,6 +46,7 @@ from collections.abc import Callable, Iterable from dataclasses import dataclass, field +from restate.agents._common import flush_held from restate.sessions import SessionContext from restate.sessions.internal import Schema, ScopeKey @@ -91,7 +92,12 @@ class _Calls: async def _store(ctx: SessionContext, state: _Calls) -> None: - """Store ``state``; it is committed before the call it records is made.""" + """Store ``state``; it is committed before the call it records is made. + + The rows the conversation holds back are written first: the call outlives the attempt, and what led to it must + too. + """ + await flush_held(ctx) await ctx.db.execute( "INSERT OR REPLACE INTO restate_agent_calls (row, invocation, calls) VALUES (0, ?, ?)", [state.invocation, json.dumps(state.calls)], diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index e56b7af5..e8138b9b 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -65,6 +65,27 @@ def next_prompt_id(ctx: SessionContext) -> str: return ctx.request.id if n == 1 else f"{ctx.request.id}:{n}" +_HELD: ScopeKey[Callable[[], Awaitable[None]]] = ScopeKey("the held rows of the conversation") + + +def holding(ctx: SessionContext, flush: Callable[[], Awaitable[None]]) -> None: + """Have ``flush_held`` write, with ``flush``, the rows that the conversation of the attempt's turn holds back.""" + _HELD.set(ctx, flush) + + +async def flush_held(ctx: SessionContext) -> None: + """Write the rows that the conversation of the attempt's turn holds back, if any (see ``_conversation._Ledger``). + + Call it before what another party sees or what outlives the attempt: a + publish, a Restate call or send. A held row is lost in a crash, and the + retried attempt makes it again; what was seen or called must not + depend on it. + """ + flush = _HELD.get(ctx) + if flush is not None: + await flush() + + MODEL_ENV = "RESTATE_AGENT_MODEL" """The environment variable that names the model agents use when none is given.""" diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index d0cf3d3b..d1b93193 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -57,9 +57,11 @@ Each row has a ``seq``, which increases with every row written, so that a client can ask what changed since a cursor. Rows are only ever appended, in -one transaction per step, but for a reply at the end that is dropped (to ask -the model again), which is one range delete on the primary key; and a -turn's answer is one update of its row, by its unique prompt id. +one transaction per step (a step that no tool waits for, such as the final +reply, is held for the next one: see ``_Ledger``), but for a reply at the end +that is dropped (to ask the model again), which is one range delete on the +primary key; and a turn's answer is one update of its row, by its unique +prompt id. ``restate_agent_conversation_start`` has one row: the first position of the conversation (``clear`` deletes the rows, and positions carry on), the last @@ -83,11 +85,12 @@ import asyncio import bisect +import contextlib import json import logging import time import typing -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Sequence from dataclasses import dataclass, field, replace from typing import Any, Literal, Protocol @@ -103,6 +106,7 @@ Window, check_window, cut_off_result, + holding, is_cancellation, next_prompt_id, refusal_of, @@ -110,7 +114,7 @@ ) from restate.agents._turn import AttemptEnded, answering, settings_of, stop_signal_of from restate.exceptions import TerminalError -from restate.sessions import Database, DatabaseReader, DatabaseWriter, SessionContext, Transaction +from restate.sessions import Database, DatabaseReader, DatabaseWriter, SessionContext, SqlParam, Transaction from restate.sessions.internal import Schema, definition_of _logger = logging.getLogger(__name__) @@ -470,6 +474,16 @@ class _Ledger: Writes take the lock one after another (tool calls run at once append their results one after another), so each gets positions of its own. + + In ``answer_once``'s run (``Conversation.hold``), a step that no tool waits + for (a reply without tool calls, say) is held rather than written + (``held``): every write writes the held rows first, in its own + transaction, so that a turn's final reply commits with its answer, in one + flush rather than two. A held row is in memory only: lost in a crash, it + is made again by the retried attempt, which runs from what the database + holds. So the run writes them (``flush``) before anything another party + sees or that outlives the attempt: a publish, a Restate call of the + library, the attempt failing. """ def __init__(self, db: Database, items: Items, start: Start, transcript: list[Entry], next: int, seq: int) -> None: @@ -495,6 +509,37 @@ def __init__(self, db: Database, items: Items, start: Start, transcript: list[En """The latest snapshot, if the conversation was loaded with snapshots and has one.""" self.compaction = False """Whether the run asked for the history to be compacted after the turn (``Conversation.compact_after``).""" + self.holds = False + """Whether steps that no tool waits for are held (``Conversation.hold``), rather than written at once.""" + self.held: list[list[SqlParam]] = [] + """The rows of the trajectories given positions but not written yet, in order of position (``_INSERT_STEP``'s + parameters): the next write writes them first.""" + + @contextlib.asynccontextmanager + async def transaction(self) -> AsyncIterator[Transaction]: + """Open a transaction that writes the held rows first; hold the lock. + + They are no longer held once it commits; if it rolls back, they still are. + """ + async with self.db.transaction() as tx: + for row in self.held: + await tx.execute(_INSERT_STEP, row) + yield tx + self.held = [] + + async def write(self, *statements: tuple[str, Sequence[SqlParam]]) -> None: + """Run ``statements`` after the held rows, in one transaction; hold the lock. + + A single statement (one held row, say, or an answer when no row is held) commits on its own. + """ + if len(self.held) + len(statements) == 1: + sql, params = statements[0] if statements else (_INSERT_STEP, self.held[0]) + await self.db.execute(sql, params) + self.held = [] + elif self.held or statements: + async with self.transaction() as tx: + for sql, params in statements: + await tx.execute(sql, params) def called(self, call_id: str) -> None: """Note a tool call's id, and log one that an earlier call has (once).""" @@ -514,11 +559,12 @@ def check_open(self) -> None: async def store_snapshot(self, held: Sequence[Item], turn: int | None) -> Snapshot: """Store ``held`` as the snapshot that covers every position given so far, in turn ``turn``; hold the lock. - It takes the place of the earlier snapshots: only the latest is read. + It takes the place of the earlier snapshots: only the latest is read. The rows held so far, which it covers, + are written with it. """ covers = self.next - 1 stored = json.dumps([self.items.dump(item) for item in held]) - async with self.db.transaction() as tx: + async with self.transaction() as tx: await tx.execute(f"DELETE FROM {SNAPSHOTS}") await tx.execute(f"INSERT INTO {SNAPSHOTS} (covers, turn, items) VALUES (?, ?, ?)", [covers, turn, stored]) self.snapshot = _snapshot_of(self.items, covers, turn, held) @@ -597,9 +643,15 @@ def items(self) -> Sequence[Entry]: """The turn's items, in order: the model's replies, the tool calls' results, ...""" return self._entries[1:] - def _add(self, entries: list[Entry]) -> None: + def _read(self, entries: list[Entry]) -> None: + """Read what each of ``entries`` says, and the tool calls and results it holds, once.""" for entry in entries: entry.said, entry.calls, entry.results = self._ledger.items.read(entry.item) + + def _add(self, entries: list[Entry], *, read: bool = True) -> None: + if read: + self._read(entries) + for entry in entries: entry.turn = self.n self._index(entry) self._entries.extend(entries) @@ -748,7 +800,7 @@ async def _steer( Entry(n + i, steered_item(items, message["text"], message["id"]), steer=message["id"]) for i, message in enumerate(new) ] - async with ledger.db.transaction() as tx: + async with ledger.transaction() as tx: for i, entry in enumerate(added, 1): await tx.execute( _INSERT_STEP, [entry.n, self.n, ledger.seq + i, items.dump(entry.item), None, None] @@ -811,7 +863,9 @@ async def drop(self, reply: int) -> None: async with ledger.lock: self._check_writable() start = replace(ledger.start, seq=ledger.seq) - async with ledger.db.transaction() as tx: + # The held rows it drops are never written; those before it are, with the deletes. + ledger.held = [row for row in ledger.held if typing.cast(int, row[0]) < reply] + async with ledger.transaction() as tx: await tx.execute(f"DELETE FROM {STEPS} WHERE n >= ?", [reply]) await tx.execute(f"DELETE FROM {STEERED} WHERE n >= ?", [reply]) await restart(tx, start) @@ -834,6 +888,8 @@ async def answered(self, answer: str, status: Status = "answered") -> None: ends with, whose id the ``messages`` view's summary gives it. A turn that ends otherwise gets ``answer`` as its reply first, from the agent the turn is with; an empty answer gets none, and has no message. + + The final reply, which is held (see ``_Ledger``), commits with the update: in one transaction. """ at = self._answer_message() if at is None and answer.strip(): @@ -845,7 +901,7 @@ async def answered(self, answer: str, status: Status = "answered") -> None: agent = self._agent() carried = self._carried.get("carried", ledger.carried()[1]) seq = ledger.seq + 1 - await ledger.db.execute(_ANSWER, [status, answer, agent, carried, at, seq, self.n]) + await ledger.write((_ANSWER, [status, answer, agent, carried, at, seq, self.n])) ledger.seq = seq prompt = self.prompt prompt.status, prompt.answer, prompt.agent, prompt.carried = status, answer, agent, carried @@ -903,7 +959,7 @@ async def close(self, ctx: SessionContext, status: Status, note: str) -> None: answered_seq = seq + len(added) + 1 agent = self._agent() carried = self._carried.get("carried", ledger.carried()[1]) - async with ledger.db.transaction() as tx: + async with ledger.transaction() as tx: for i, entry in enumerate(added, 1): await tx.execute(_INSERT_STEP, [entry.n, self.n, seq + i, items.dump(entry.item), None, None]) await tx.execute(_ANSWER, [status, note, agent, carried, at, answered_seq, self.n]) @@ -925,27 +981,30 @@ async def _insert(self, made: Callable[[int], list[Entry]]) -> int: The position is taken once the appends before have stored theirs, so appends that run at once each get places of their own. + + Entries that hold no tool call nor result (a reply that ends the run, + say) are held, not written, in a conversation that holds them (see + ``_Ledger``): no tool waits for them, and the next write writes them. + Tool calls are written before their tools run, and results before the + model is asked again, with the rows held before them. """ ledger = self._ledger async with ledger.lock: self._check_writable() n = ledger.next added = made(n) - rows = [ + self._read(added) + rows: list[list[SqlParam]] = [ [e.n, self.n, ledger.seq + i, ledger.items.dump(e.item), e.reply, e.agent] for i, e in enumerate(added, 1) ] - if len(rows) == 1: - # One statement commits on its own. - await ledger.db.execute(_INSERT_STEP, rows[0]) + if ledger.holds and not any(e.calls or e.results for e in added): + ledger.held.extend(rows) else: - async with ledger.db.transaction() as tx: - # The same SQL for each row: the database prepares it once. - for row in rows: - await tx.execute(_INSERT_STEP, row) + await ledger.write(*((_INSERT_STEP, row) for row in rows)) ledger.seq += len(rows) ledger.next = n + len(rows) - self._add(added) + self._add(added, read=False) return n @@ -1308,6 +1367,26 @@ def compaction_asked(self) -> bool: """Whether the run asked for the history to be compacted after the turn (``compact_after``).""" return self._ledger.compaction + def hold(self) -> None: + """Hold the steps that no tool waits for, rather than write each at once: the next write writes them. + + For a run that writes them (``flush``) before each publish and Restate + call of the library, and before its attempt fails: ``answer_once``'s. + A conversation that is written otherwise (step by step, from handlers + of its own) writes each step as it comes. + """ + self._ledger.holds = True + + async def flush(self) -> None: + """Write the rows the turn's run holds (see ``_Ledger``), if any. + + Before whatever another party sees, or outlives the attempt: a publish + tells the clients to read them, and a Restate call is made once, while + a held row would be lost in a crash. + """ + async with self._ledger.lock: + await self._ledger.write() + async def seal(self) -> None: """Close the conversation to writes, once those under way are stored. @@ -1498,9 +1577,22 @@ async def answer_once( answer, status = found return stored_answer(answer, status) conversation = await Conversation.load(ctx.db, items, window, budget, snapshots=settings.snapshots) + # The turn's final reply commits with its answer; the library's publishes and Restate calls write what the run + # holds back first. + conversation.hold() + holding(ctx, conversation.flush) async def work() -> str: - return await _answer(ctx, conversation, prompt_id, text, run, check) + try: + return await _answer(ctx, conversation, prompt_id, text, run, check) + except AttemptEnded: + # Its session ended with it: nothing more is stored. + raise + except Exception: + # The attempt fails, and its retry carries on from what is stored: so what the run made so far is, + # as when each step was written as it came. A turn that was closed holds nothing back. + await _flushed(conversation) + raise def stored(call_id: str, failed: str | None) -> bool: # The Restate calls of this invocation are its turn's. @@ -1570,6 +1662,17 @@ async def answering() -> str: return stored_answer(answer) +async def _flushed(conversation: Conversation) -> None: + """Write what the run holds back, as its attempt fails: never in place of the attempt's own error. + + The attempt's session may have ended with it (the error is its sync's, say), and then nothing more is stored. + """ + try: + await conversation.flush() + except Exception as error: # pylint: disable=broad-exception-caught + _logger.debug("The rows the run held were not stored as its attempt failed: %r", error) + + async def settle(ctx: SessionContext, items: Items, invocation_id: str) -> bool: """Close the latest turn (``ended``) if invocation ``invocation_id`` made it and left it without an answer. diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index 597085c2..0e66947e 100644 --- a/python/restate/agents/_turn.py +++ b/python/restate/agents/_turn.py @@ -28,7 +28,7 @@ from restate.agents import _signals from restate.agents._api import Answering, Published, Steer from restate.agents._calls import settle -from restate.agents._common import BUDGET, WINDOW, CompactHistory, Summarize, Window, refusal_of +from restate.agents._common import BUDGET, WINDOW, CompactHistory, Summarize, Window, flush_held, refusal_of from restate.exceptions import TerminalError from restate.sessions import SessionContext from restate.sessions._link import Link @@ -224,8 +224,10 @@ async def publish(ctx: SessionContext) -> None: """Publish a change to the conversation of ``ctx``: the revision increases, and the clients read it again. While a framework is answering (inside ``answering``), the calls whose - results are stored are first forgotten. + results are stored are first forgotten. The rows the conversation holds + back are written first: the clients read them once told to. """ + await flush_held(ctx) stored = _ANSWERING.get(ctx) if stored is not None: await _settled(ctx, stored) diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index 461820e8..88a50ff3 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -118,7 +118,7 @@ def forget(state: MessagesState) -> dict[str, Any]: from pydantic_core import to_jsonable_python from restate.agents import _signals -from restate.agents._common import from_tool +from restate.agents._common import flush_held, from_tool from restate.agents._conversation import Conversation, Run, Turn from restate.agents._turn import stepped from restate.agents.langgraph._checkpointer import SqlCheckpointSaver @@ -313,6 +313,8 @@ async def _resumed(ctx: SessionContext, snapshot: StateSnapshot) -> dict[str, An pending = _interrupts(snapshot) if not pending: return None + # What the turn holds back is written before it waits: the waits outlive the attempt. + await flush_held(ctx) waits: dict[str, Any] = {} for interrupted in pending: wait_id = f"{ctx.request.id}:interrupt:{interrupted.id}" diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index ac81d768..184015bd 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -60,6 +60,7 @@ async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: from pydantic_core import to_jsonable_python from restate.agents import _signals +from restate.agents._common import flush_held from restate.agents._tools import handler_of, input_type_of, inputs_of, keyed, run_workflow, untyped from restate.agents.langgraph._context import handler_context from restate.exceptions import TerminalError @@ -86,7 +87,11 @@ class _Step(typing.NamedTuple): def _step(runtime: Runtime[Any], kind: str) -> _Step: - """Return the run of a node of ``kind``, from LangGraph's runtime and config.""" + """Return the run of a node of ``kind``, from LangGraph's runtime and config. + + Each node then writes the rows the turn holds back (``flush_held``) before its call, send or wait, which + outlives the attempt. + """ ctx = handler_context(runtime.context) if ctx is None: raise TypeError( @@ -200,6 +205,7 @@ def wait_for_signal( async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "wait_for_signal") + await flush_held(step.ctx) # The wait's id starts with the invocation's (the same in each attempt), as the waits of tools do, so that a # turn that ends without it (killed, say) has it closed: the companion closes the waits whose ids start # with the turn's. @@ -252,6 +258,7 @@ def workflow( async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "workflow") + await flush_held(step.ctx) workflow_id = key(state) if key is not None else step.idempotency_key value = called.validated(input(state)) if input is not None else None result = await run_workflow(step.ctx, target, value, workflow_id) @@ -295,6 +302,7 @@ def call( async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "call") + await flush_held(step.ctx) value = called.validated(input(state)) if input is not None else None result = await untyped(step.ctx).call( target, @@ -341,6 +349,7 @@ def send( async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "send") + await flush_held(step.ctx) value = called.validated(input(state)) if input is not None else None untyped(step.ctx).send( target, @@ -385,6 +394,7 @@ def subagent( async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "subagent") + await flush_held(step.ctx) agent_key = key(state) if key is not None else f"{step.ctx.key}.{step.node}" answer = await step.ctx.generic_call( agent_name, handler, message(state), key=agent_key, idempotency_key=step.idempotency_key diff --git a/tests/agents_adk.py b/tests/agents_adk.py index 9aad2fd3..b50701fa 100644 --- a/tests/agents_adk.py +++ b/tests/agents_adk.py @@ -615,6 +615,54 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: assert replies == ["asked"] +async def test_a_plain_turn_commits_its_final_response_end_of_agent_and_answer_at_once(): + from google.adk.events import Event # pylint: disable=import-outside-toplevel + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + from sessions_object import Host, request # pylint: disable=import-error,import-outside-toplevel + + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + class Syncing(Host): + syncs: int = 0 + + async def sync(self) -> None: + await super().sync() + self.syncs += 1 + + host: list[Syncing] = [] + asked: list[int] = [] + + class Answering(BaseLlm): + model: str = "answering" + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + asked.append(host[0].syncs) + yield LlmResponse(content=types.Content(role="model", parts=[types.Part(text="hello")])) + + asker = sessions.SessionObject("asker") + + @asker.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await adk.prompt(ctx, text, model=Answering()) + + @asker.handler + async def stored(ctx: sessions.SessionContext) -> list[bool]: + rows = await ctx.db.values("SELECT item FROM restate_agent_items ORDER BY n") + return [end_of_agent(Event.model_validate_json(str(item))) for [item] in rows] + + backend = Backend() + for text in ("first", "second"): + host[:] = [Syncing(backend, service="asker", answers={"publish": b"null"})] + assert await invoke(asker, "ask", host[0], request(), text) == "hello" + # The second turn (the conversation's tables exist): from the model's response on, one commit, which has the + # final response, the end of the agent and the answer. + assert host[0].syncs - asked[-1] == 1 + ends = await invoke(asker, "stored", Syncing(backend, service="asker"), request()) + assert ends.count(True) == 2 and len(ends) == 4 + + async def test_an_app_of_your_own_keeps_its_plugins_and_gets_this_packages_tools(): from google.adk.apps import App # pylint: disable=import-outside-toplevel from google.adk.plugins import BasePlugin # pylint: disable=import-outside-toplevel diff --git a/tests/agents_conversation.py b/tests/agents_conversation.py index f8c5f7a6..962e622d 100644 --- a/tests/agents_conversation.py +++ b/tests/agents_conversation.py @@ -957,3 +957,192 @@ async def cold(handler: str, key: str) -> tuple[Any, int]: # they come from (both tables, and the seq and answered_seq indexes). A SELECT of several maxima scanned. assert scanned >= 20 assert read <= 5 + + +# --- A step that no tool waits for commits with the next write -------------------------------------- + + +class Syncing(Host): + """A host that counts its syncs (each a commit of the database), and crashes at a publish once armed.""" + + syncs: int = 0 + crash_at_publish: bool = False + + async def sync(self) -> None: + await super().sync() + self.syncs += 1 + + def generic_call(self, service, handler, arg, key=None, idempotency_key=None, headers=None): + if handler == "publish" and self.crash_at_publish: + # The process crashes once the clients were told to read: every later sync fails. + self.crash_after(0) + return super().generic_call(service, handler, arg, key, idempotency_key, headers) + + +def _turn_rows() -> sessions.SessionObject: + """An object whose ``ask`` answers with the run that ``runs`` holds, and whose ``rows`` reads each turn back.""" + obj = sessions.SessionObject("asker") + + @obj.handler + async def rows(ctx: sessions.SessionContext) -> list[dict[str, Any]]: + conversation = await Conversation.load(ctx.db, ITEMS, "all") + return [ + {"answer": turn.answer, "items": [[entry.n, entry.item] for entry in turn.items]} + for turn in conversation.turns + ] + + return obj + + +def _syncing(backend: Backend) -> Syncing: + return Syncing(backend, service="asker", answers=PUBLISHED) + + +async def test_a_turn_that_ends_with_a_reply_commits_the_reply_and_its_answer_at_once(): + host: list[Syncing] = [] + before: list[int] = [] + + async def run(conversation: Conversation) -> str: + before.append(host[0].syncs) + await conversation.turn.append_reply({"role": "assistant", "text": "hello"}) + return "hello" + + obj = _turn_rows() + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await answer_once(ctx, text, ITEMS, run) + + backend = Backend() + for text in ("first", "second"): + host[:] = [_syncing(backend)] + assert await invoke(obj, "ask", host[0], request(), text) == "hello" + # The second turn (the conversation's tables exist): from the reply on, one commit, with the answer. + assert host[0].syncs - before[-1] == 1 + turns = await invoke(obj, "rows", _syncing(backend), request()) + assert [turn["answer"] for turn in turns] == ["hello", "hello"] + + +async def test_a_crash_at_the_final_commit_stores_neither_the_reply_nor_the_answer_and_the_retry_asks_again(): + runs: list[str] = [] + + async def run(conversation: Conversation) -> str: + runs.append("ran") + await conversation.turn.append_reply({"role": "assistant", "text": f"hello {len(runs)}"}) + return f"hello {len(runs)}" + + obj = _turn_rows() + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await answer_once(ctx, text, ITEMS, run) + + crashed = 0 + for syncs in range(20): + backend, invocation = Backend(), request() + runs.clear() + crashing = _syncing(backend) + crashing.crash_after(syncs) + try: + await invoke(obj, "ask", crashing, invocation, "hi") + break + except RuntimeError as error: + assert str(error) == "the sync failed" + stored = await invoke(obj, "rows", _syncing(backend), request()) + # Never a reply without its answer: the two commit together, or not at all. + answered = [{"answer": "hello 1", "items": [[1, {"role": "assistant", "text": "hello 1"}]]}] + assert stored in ([], [{"answer": None, "items": []}], answered), stored + # A crash after the model answered, at the final commit: neither was stored. + crashed += len(runs) == 1 and stored != answered + answer = await invoke(obj, "ask", _syncing(backend), invocation, "hi") + # The retried attempt asked the model again, or returned the stored answer; one reply, one answer. + assert answer == f"hello {len(runs)}" + assert len(runs) == (2 if stored == [{"answer": None, "items": []}] else 1) + turns = await invoke(obj, "rows", _syncing(backend), request()) + assert turns == [{"answer": answer, "items": [[1, {"role": "assistant", "text": answer}]]}] + assert crashed == 1 + + +async def test_a_reply_held_back_is_stored_before_a_publish_that_follows_it(): + from restate.agents._turn import stepped # pylint: disable=import-outside-toplevel + + obj = _turn_rows() + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + async def run(conversation: Conversation) -> str: + await conversation.turn.append_reply({"role": "assistant", "text": "on it"}) + # The clients are told to read the reply: it is stored by then, even if the process crashes after. + await stepped(ctx) + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + return await answer_once(ctx, text, ITEMS, run) + + backend = Backend() + host = _syncing(backend) + host.crash_at_publish = True + with pytest.raises(RuntimeError, match="the sync failed"): + await invoke(obj, "ask", host, request(), "hi") + turns = await invoke(obj, "rows", _syncing(backend), request()) + assert turns == [{"answer": None, "items": [[1, {"role": "assistant", "text": "on it"}]]}] + + +async def test_a_run_that_fails_its_attempt_after_a_reply_leaves_the_reply_stored(): + async def run(conversation: Conversation) -> str: + await conversation.turn.append_reply({"role": "assistant", "text": "half way"}) + raise RuntimeError("the run failed") + + obj = _turn_rows() + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await answer_once(ctx, text, ITEMS, run) + + backend = Backend() + with pytest.raises(RuntimeError, match="the run failed"): + await invoke(obj, "ask", _syncing(backend), request(), "hi") + # The retry carries on from it, as from any step stored before the attempt failed. + turns = await invoke(obj, "rows", _syncing(backend), request()) + assert turns == [{"answer": None, "items": [[1, {"role": "assistant", "text": "half way"}]]}] + + +async def test_appends_at_once_with_steps_held_back_get_positions_of_their_own_in_order(): + seen: list[list[tuple[int, Item]]] = [] + + async def run(conversation: Conversation) -> str: + turn = conversation.turn + await turn.append_reply(_calling("c1", "c2")) + rng = random.Random(7) + + async def appended(item: Item) -> None: + for _ in range(rng.randrange(3)): + await asyncio.sleep(0) + await turn.append(item) + + notes = [{"role": "assistant", "text": f"note {i}"} for i in range(4)] + results = [{"result": call_id, "output": "done"} for call_id in ("c1", "c2")] + await asyncio.gather(*(appended(item) for item in [notes[0], results[0], notes[1], notes[2], results[1]])) + await turn.append_reply(notes[3]) + seen.append([(entry.n, entry.item) for entry in turn.items]) + return "note 3" + + obj = _turn_rows() + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await answer_once(ctx, text, ITEMS, run) + + @obj.handler + async def seqs(ctx: sessions.SessionContext) -> list[Any]: + return [row[0] for row in await ctx.db.values(f"SELECT seq FROM {STEPS} ORDER BY n")] + + backend = Backend() + assert await invoke(obj, "ask", _syncing(backend), request(), "hi") == "note 3" + [turn] = await invoke(obj, "rows", _syncing(backend), request()) + stored = [(n, item) for n, item in turn["items"]] + positions = [n for n, _ in stored] + # Each item has a position of its own, the next one, and is stored where the run saw it, its seq in order. + assert positions == list(range(1, 8)) + assert [tuple(entry) for entry in stored] == seen[0] + assert await invoke(obj, "seqs", _syncing(backend), request()) == list(range(2, 9)) diff --git a/tests/agents_pydantic_ai.py b/tests/agents_pydantic_ai.py index 737c4bb1..def17862 100644 --- a/tests/agents_pydantic_ai.py +++ b/tests/agents_pydantic_ai.py @@ -960,6 +960,40 @@ def test_the_openai_sdks_error_for_a_missing_key_is_a_refusal(monkeypatch: pytes assert refused is not None and refused.status_code == 424 +async def test_a_turn_that_ends_with_a_reply_commits_the_reply_and_its_answer_at_once(): + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + from sessions_object import Host, request # pylint: disable=import-error,import-outside-toplevel + + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + class Syncing(Host): + syncs: int = 0 + + async def sync(self) -> None: + await super().sync() + self.syncs += 1 + + host: list[Syncing] = [] + asked: list[int] = [] + + def answer(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + asked.append(host[0].syncs) + return ModelResponse(parts=[TextPart("hello")]) + + asker = sessions.SessionObject("asker") + + @asker.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await pydantic_ai.prompt(ctx, text, model=FunctionModel(answer)) + + backend = Backend() + for text in ("first", "second"): + host[:] = [Syncing(backend, service="asker")] + assert await invoke(asker, "ask", host[0], request(), text) == "hello" + # The second turn (the conversation's tables exist): from the model's reply on, one commit, with the answer. + assert host[0].syncs - asked[-1] == 1 + + async def test_an_answer_stored_before_a_crash_is_not_asked_for_again(monkeypatch: pytest.MonkeyPatch): from restate.agents.pydantic_ai import _agent # pylint: disable=import-outside-toplevel from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel From 79884ef92754287bd00cd6b20b97324839a29de1 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 10:39:53 +0000 Subject: [PATCH 2/8] Agents: a turn publishes after each step it commits, not at each tool call's start and end A followed turn told its companion to publish at its start and at every tool call's start and end: 3 + 2k calls a turn (k tool calls), each an invocation of the companion. A publish carries nothing: it wakes the followers, who read the conversation again. Pydantic AI, OpenAI and ADK store a reply's results together, after the tools' end and before the next model request, and no publish followed that commit: a follower saw the results only when the next tool call started or the turn ended, a model request late. Of the 8 revisions of a turn with one reply of 3 tool calls, 6 showed a follower nothing new. Now the turn publishes once after each write of its steps that commits (a reply with tool calls, the tools' results), and waits for it: a follower sees the calls as their tools start and the results before the model is asked again, and the messages steered to the turn are those of that publish. The companion's `start` makes the turn's first publish and returns what `publish` returns (`Started.publishes`); a companion that returns nothing from `start` gets a publish from the turn, as before. `tool_call_started` only forgets the stored calls, and `tool_call_ended` does nothing. The companion no longer clears an empty `watching` at each publish. Companion calls of a turn with one reply of 3 tool calls go from 9 to 4, and of a plain turn from 3 to 2; that turn's journal, on Restate 1.8.1 with Pydantic AI, from 46-48 entries to 31-33. Co-Authored-By: Claude Fable 5.2 --- docs/agents/access-patterns.md | 18 +- docs/agents/conversation.md | 7 +- docs/agents/learnings.md | 2 +- docs/agents/porting.md | 9 + docs/agents/turns.md | 63 ++-- python/restate/agents/_conversation.py | 50 ++- python/restate/agents/_signals.py | 44 ++- python/restate/agents/_steering.py | 5 +- python/restate/agents/_tools.py | 13 +- python/restate/agents/_turn.py | 57 ++-- python/restate/agents/adk/_agent.py | 9 +- python/restate/agents/langgraph/_agent.py | 3 +- python/restate/agents/langgraph/_graph.py | 5 +- python/restate/agents/openai/_agent.py | 7 +- python/restate/agents/pydantic_ai/_agent.py | 9 +- tests/agents_publishes.py | 327 ++++++++++++++++++++ tests/agents_signals.py | 38 ++- 17 files changed, 554 insertions(+), 112 deletions(-) create mode 100644 tests/agents_publishes.py diff --git a/docs/agents/access-patterns.md b/docs/agents/access-patterns.md index 38340e2f..84366262 100644 --- a/docs/agents/access-patterns.md +++ b/docs/agents/access-patterns.md @@ -291,9 +291,10 @@ In order of value for effort: - **The companion and the call records**, not exercised by this harness: their writes per turn, and the `restate_agent_calls` read at every turn. - **Steering** adds no request to a turn that nobody steers: the companion - answers the `publish` a turn makes at each tool call's start and end with - the messages waiting for it (one state read more, of its inbox, which is - empty then), and a turn takes them at a step with no call of its own; + answers the publishes a turn makes anyway (its `start`, and one after each + write of steps that commits: a reply with tool calls, the tools' results) + with the messages waiting for it (one state read more, of its inbox, which + is empty then), and a turn takes them at a step with no call of its own; the `taken` that tells the companion is sent from the storing transaction's outbox. A steered message costs the `steer` view, one call to the companion (an inbox entry, a publish, and the placeholder sent to @@ -303,11 +304,12 @@ In order of value for effort: transcript it loads as one range of that table per page; the summary view, one range. Measured on a real server with `scripts/steering_traffic.py` (a turn of one tool call): a turn nobody - steers is what it is before steering, 47 journal entries and 5 companion - calls (`start`, 3 `publish`, `close`); a turn that takes a steered - message, 57 entries and 6 calls (the outbox's `taken`); the message - itself, the `steer` view (5 entries), the companion's `steer` (13), and - the placeholder that finds it taken (3). + steers, 40 journal entries and 4 companion calls (`start`, which + publishes; a `publish` after the reply with the call, and one after its + result; `close`); a turn that takes a steered message, 50 entries and 5 + calls (the outbox's `taken`); the message itself, the `steer` view (5 + entries), the companion's `steer` (12), and the placeholder that finds it + taken (3). - **Snapshots and `compact`:** a snapshot replaces the previous one (one `DELETE` and one `INSERT` of the whole compacted history), and `compact` reads the transcript from the summary's start. Both are occasional, but diff --git a/docs/agents/conversation.md b/docs/agents/conversation.md index 3249c7fc..a735e423 100644 --- a/docs/agents/conversation.md +++ b/docs/agents/conversation.md @@ -524,9 +524,10 @@ invocation (the next invocation's first call replaces it): again under the same id; a new call with the same input (the model asking again after the first's answer) waits anew, and never gets an earlier call's approval or denial. -- **Forgetting.** At each publish and before each tool call, a record whose - Restate call returned is forgotten once the transcript has its result - under the latest id it was made with: a success, or the very error the call +- **Forgetting.** At each publish (after each write of steps that commits) + and before each tool call, a record whose Restate call returned is + forgotten once the transcript has its result under the latest id it was + made with: a success, or the very error the call returned. Another error in the transcript (one the framework made up for a cut-off call) doesn't count: the call may be made again under a new id, and must find its record. diff --git a/docs/agents/learnings.md b/docs/agents/learnings.md index f390edb5..aacc3522 100644 --- a/docs/agents/learnings.md +++ b/docs/agents/learnings.md @@ -79,7 +79,7 @@ Each framework's details are in [The frameworks](frameworks.md). | — | **The summary's form stays per framework.** OpenAI's is one user message; ADK and LangGraph keep a user message then "Understood.". | OpenAI's was checked against both its APIs; the others run on scripted models in the tests, which can't tell what Gemini or other chat models accept. | | D6 | **Steering is a view of its own, `steer`; `chat` keeps one message per turn.** | A client that sends to `chat` gets the answer to its message as `chat`'s response, which joining another turn would take from it. A UI that wants one way in sends everything to `steer`, which starts a turn when none is answered. | | D7 | **A steered message is queued as a turn of its own the moment it is sent** (the agent's `steered`), and the turn that takes it makes that turn do nothing. | Its place in the key's queue is kept with the messages sent after it, a turn that dies leaves no message behind, and the companion forwards nothing at a turn's close. The cost is one short invocation per steered message. | -| D8 | **The turn learns of steered messages from `publish`**, which it makes at each tool call's start and end, and takes them with an outbox send. | A turn that nobody steers makes no call more: the check is free. A call of its own at each step would cost one round trip per tool round of every turn. | +| D8 | **The turn learns of steered messages from `publish`**, which it makes anyway (its `start`, and after each write of steps that commits), and takes them with an outbox send. | A turn that nobody steers makes no call more: the check is free. A call of its own at each step would cost one round trip per tool round of every turn. | | D9 | **Each framework takes a steered message in place**, in the request about to go out (OpenAI's input filter, Pydantic AI's `enqueue`, ADK's session events, LangGraph's model node); none restarts its run. | Restarting a run from the stored steps (the crash path) reran input guardrails, reset `max_turns`, usage limits and recursion limits, and left the old run able to write. | | D10 | **A message taken late is the next turn**: one that comes after the turn's last step is answered next, not folded into the turn. | The turn's answer stays the answer to what it was asked; folding it in would hide that answer from the summary view. | | — | **Steps retention is future work** ([its design](steps-retention.md)). | Nothing reads a turn's steps once it leaves the window, but deleting them changes what the full view and conversation search can show; it can come later, as an option. | diff --git a/docs/agents/porting.md b/docs/agents/porting.md index 02d455dc..372b9667 100644 --- a/docs/agents/porting.md +++ b/docs/agents/porting.md @@ -57,6 +57,15 @@ Each step stands on its own, with its tests, and the earlier ones are small. terms (`checkpointer.ts` builds such `OR`s); and, if the TypeScript LangGraph has `DeltaChannel`, keep what it is rebuilt from when pruning (`_kept` in `_checkpointer.py`). + - **The companion's `start` publishes.** A turn that says so + (`Started.publishes`, a followed turn) is published as `start` starts + it, and `start` returns what `publish` returns (`{revision, + steering?}`): the turn makes no publish of its own to start. A turn that + doesn't say so gets nothing back, and is not published: so a TypeScript + turn that this SDK's companion holds publishes its start as before, and + this SDK's turn publishes its own start when the TypeScript companion + returns nothing. `publish` (and `start`) clear `watching` only when a + watch waits. 2. **The companion's `settle`.** A turn's start says whether its agent settles (`Started.settles`); when `awaitTurn` finds the invocation gone, the companion closes the turn and sends such an agent `settle(turn)`, a diff --git a/docs/agents/turns.md b/docs/agents/turns.md index e1919343..5499abc7 100644 --- a/docs/agents/turns.md +++ b/docs/agents/turns.md @@ -112,16 +112,16 @@ sequenceDiagram participant DB as Key's database participant S as Companion (signals) C->>A: chat("Plan my week") - A->>S: start(turn, settles) + A->>S: start(turn, settles, publishes) S-)S: awaitTurn(turn) - Note over S: status.answering = turn - A->>S: publish + Note over S: status.answering = turn, publish + S-->>A: the revision, and the messages steered to the turn, if any A->>DB: the prompt's answer? (none yet) A->>DB: load the transcript and the window A->>DB: ask: the turn's row loop each step of the framework - A->>DB: append the model's reply, the tool calls' results - A->>S: publish (a tool call starts, ends) + A->>DB: append the model's reply with tool calls, or the tools' results + A->>S: publish (after the commit) S-->>A: the revision, and the messages steered to the turn, if any opt messages steered to the turn wait, at the step A->>DB: take them: steps of the turn, and taken in the outbox @@ -143,7 +143,11 @@ sequenceDiagram It also says the agent has `settle` (see [Dead turns](#dead-turns)). A retried attempt calls `start` again with the same invocation, which changes nothing (it keeps when the turn started). `start` sends - `awaitTurn(turn)`, the turn's watch. + `awaitTurn(turn)`, the turn's watch. It also publishes, as the turn's + first revision, and returns what `publish` returns: the turn makes no + publish of its own to start. A companion that returns nothing from + `start` (an older deployment, or the TypeScript SDK's) gets a `publish` + from the turn instead. 2. **The prompt.** `answer_once` reads the prompt's answer by its id, one indexed read: a retried attempt of an answered turn returns it, and reads nothing more; one of a turn that was closed (cancelled, failed, refused, or @@ -152,12 +156,15 @@ sequenceDiagram ([Which turns a turn loads](conversation.md#which-turns-a-turn-loads)), closes a turn left without an answer, and stores the prompt as the turn's row. 3. **Steps.** The framework answers: model calls and tool calls. Each step - is stored as it comes (a reply in one transaction, each result as it - comes back), and each tool call publishes as it starts and as it ends, - so a follower sees the call start and its result. The companion answers - each `publish` with the messages steered to the turn that wait: at its - next step, the turn takes them, and the model's next request has them - ([Steering a turn](#steering-a-turn)). + is stored as it comes: a reply with tool calls in one transaction, before + its tools run; the tools' results before the model is asked again. A + final reply, which no tool waits for, is held and commits with the + answer. After each write that commits, the turn publishes once, and + waits for it: so a follower sees the calls as their tools start, and + their results before the model's next request returns. The companion + answers each `publish` with the messages steered to the turn that wait: + at its next step, the turn takes them, and the model's next request has + them ([Steering a turn](#steering-a-turn)). 4. **Answer.** The answer is stored on the turn's row, with what the turn carries. If the turns before the next turn's window outgrow the agent's budget, the agent sends itself `compact` @@ -168,8 +175,8 @@ sequenceDiagram turn's error. A **background** agent (`background=True`, a subagent say) publishes its -turn's end only: no step, nor its start, wakes a follower. Its turns take no -steered messages. +turn's end only: no step, nor its start, wakes a follower (its `start` says +it does not publish). Its turns take no steered messages. ## How a turn ends @@ -450,9 +457,9 @@ sequenceDiagram V->>S: steer(id, text) S-)P: steered(id, text): queued now, behind the turn Note over S: inbox: the message, for this turn - A->>S: publish (the tool call ends) - S-->>A: revision, steering: [the message] A->>DB: the tool's result + A->>S: publish (after the result's commit) + S-->>A: revision, steering: [the message] A->>DB: the message, a step of the turn, and taken(id) in the transaction's outbox DB--)S: taken(id), once the transaction commits A->>A: the model's next request has it @@ -520,10 +527,12 @@ sequenceDiagram queue as it is sent; and, if the turn being answered takes steered messages (its `start` says so, `Started.steers`), it holds the message for that turn (`inbox`), 16 at most (`STEERS`). -- **How the turn learns of it.** From `publish`, which the turn makes at - each tool call's start and end anyway: the companion answers it with the - messages waiting for the turn. A turn that nobody steers makes no call - more. +- **How the turn learns of it.** From `publish`, which the turn makes + anyway as it starts (its `start`) and after each write of steps that + commits (a reply with tool calls, the tools' results): the companion + answers it with the messages waiting for the turn. The turn waits for the + publish after the results, and its next step then takes what it returned. + A turn that nobody steers makes no call more. - **Taken once.** At a step, the framework's integration calls `take_steers`: the turn stores each waiting message by its id, once, as a message of the user that is a step of the turn, and the integration puts it @@ -657,8 +666,8 @@ own answer, and the companion's turn is the invocation, all of them. The revision increases at: -- a turn's start, and each of its tool calls' start and end (not in the - background); +- a turn's start (its `start`), and each write of its steps that commits: + a reply with tool calls, the tools' results (not in the background); - a turn's close: its own, its watch's, or the next turn's `start` closing it; - `settle`, once it closed a dead turn in the conversation; @@ -668,7 +677,9 @@ The revision increases at: - `clear`. So a follower that reads each new revision sees every turn begin, its tool -calls, its end, and a killed turn end. +calls as their tools start, their results before the model's next request +returns, its end, and a killed turn end. A final reply, held until the +answer, shows with the turn's end. ## Agents of the TypeScript SDK @@ -684,7 +695,11 @@ does not show), and only such a turn is sent it: A turn says it takes steered messages the same way (`Started.steers`): the companion holds messages only for such a turn, and every other gets them as -turns of their own. A TypeScript companion, which has no `steer`, answers +turns of their own. And it says whether `start` publishes +(`Started.publishes`): this SDK's companion then publishes and returns what +`publish` returns; a TypeScript companion returns nothing from `start`, and +this SDK's turn then publishes its start itself. A TypeScript agent's turn +does not say so, and this SDK's companion publishes nothing at its `start`. A TypeScript companion, which has no `steer`, answers `publish` with no `steering`, so this SDK's turns take none from it; and the `steer` view fails there, as the companion has no such handler. diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index d1b93193..a4349141 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -112,7 +112,7 @@ refusal_of, stored_answer, ) -from restate.agents._turn import AttemptEnded, answering, settings_of, stop_signal_of +from restate.agents._turn import AttemptEnded, answering, settings_of, stepped, stop_signal_of from restate.exceptions import TerminalError from restate.sessions import Database, DatabaseReader, DatabaseWriter, SessionContext, SqlParam, Transaction from restate.sessions.internal import Schema, definition_of @@ -484,6 +484,9 @@ class _Ledger: holds. So the run writes them (``flush``) before anything another party sees or that outlives the attempt: a publish, a Restate call of the library, the attempt failing. + + The run publishes after each write of steps that commits (``published``), + not after the held rows: they are no step that anyone waits for. """ def __init__(self, db: Database, items: Items, start: Start, transcript: list[Entry], next: int, seq: int) -> None: @@ -514,6 +517,12 @@ def __init__(self, db: Database, items: Items, start: Start, transcript: list[En self.held: list[list[SqlParam]] = [] """The rows of the trajectories given positions but not written yet, in order of position (``_INSERT_STEP``'s parameters): the next write writes them first.""" + self.published: Callable[[], Awaitable[None]] | None = None + """What the run makes after each write of steps that commits (a publish, ``Conversation.hold``), if anything.""" + + def holds_back(self, entries: Sequence[Entry]) -> bool: + """Return whether ``entries``, which are read, are held rather than written: no tool call nor result.""" + return self.holds and not any(e.calls or e.results for e in entries) @contextlib.asynccontextmanager async def transaction(self) -> AsyncIterator[Transaction]: @@ -740,11 +749,11 @@ async def append(self, *items: Item) -> None: Appends may run at once, such as those of tool calls that run at once: each is stored after the ones before it. """ - await self._insert(lambda n: [Entry(n + i, item) for i, item in enumerate(items)]) + await self._stored(lambda n: [Entry(n + i, item) for i, item in enumerate(items)]) async def append_reply(self, *items: Item, agent: str | None = None) -> int: """Append the model's reply, of ``items`` (given by ``agent``), to the turn, and return its id.""" - return await self._insert(lambda n: [Entry(n + i, item, reply=n, agent=agent) for i, item in enumerate(items)]) + return await self._stored(lambda n: [Entry(n + i, item, reply=n, agent=agent) for i, item in enumerate(items)]) async def ask_again(self, text: str) -> None: """Ask the model again within the turn, with ``text``: a message of the user that is a step of the turn. @@ -976,6 +985,27 @@ def _check_writable(self) -> None: if self._ledger.latest is not self or self.status is not None: raise RuntimeError(f"turn {self.id!r} is not the turn being answered: only that turn is written") + async def _stored(self, made: Callable[[int], list[Entry]]) -> int: + """Store the entries that ``made`` makes (``_insert``), and return their position; then, if they committed + (they are not held), make what the run makes after a commit (``_Ledger.published``). + + That is a publish, outside the lock, awaited before the caller goes on: + the companion returns the messages steered to the turn from it, which + the caller's next step takes. Writes that run at once each make their + own. + """ + added: list[Entry] = [] + + def making(n: int) -> list[Entry]: + added[:] = made(n) + return added + + n = await self._insert(making) + published = self._ledger.published + if published is not None and not self._ledger.holds_back(added): + await published() + return n + async def _insert(self, made: Callable[[int], list[Entry]]) -> int: """Store the entries that ``made`` makes from the next position, and return that position. @@ -998,7 +1028,7 @@ async def _insert(self, made: Callable[[int], list[Entry]]) -> int: [e.n, self.n, ledger.seq + i, ledger.items.dump(e.item), e.reply, e.agent] for i, e in enumerate(added, 1) ] - if ledger.holds and not any(e.calls or e.results for e in added): + if ledger.holds_back(added): ledger.held.extend(rows) else: await ledger.write(*((_INSERT_STEP, row) for row in rows)) @@ -1367,15 +1397,21 @@ def compaction_asked(self) -> bool: """Whether the run asked for the history to be compacted after the turn (``compact_after``).""" return self._ledger.compaction - def hold(self) -> None: + def hold(self, published: Callable[[], Awaitable[None]] | None = None) -> None: """Hold the steps that no tool waits for, rather than write each at once: the next write writes them. For a run that writes them (``flush``) before each publish and Restate call of the library, and before its attempt fails: ``answer_once``'s. A conversation that is written otherwise (step by step, from handlers of its own) writes each step as it comes. + + ``published``, if given, is awaited after each write of steps that + commits (a reply with tool calls, the tools' results, with the rows + held before them): ``answer_once`` publishes there. Held rows make + none, nor does ``flush``. """ self._ledger.holds = True + self._ledger.published = published async def flush(self) -> None: """Write the rows the turn's run holds (see ``_Ledger``), if any. @@ -1578,8 +1614,8 @@ async def answer_once( return stored_answer(answer, status) conversation = await Conversation.load(ctx.db, items, window, budget, snapshots=settings.snapshots) # The turn's final reply commits with its answer; the library's publishes and Restate calls write what the run - # holds back first. - conversation.hold() + # holds back first. Each write of steps that commits publishes, unless the agent runs in the background. + conversation.hold(lambda: stepped(ctx)) holding(ctx, conversation.flush) async def work() -> str: diff --git a/python/restate/agents/_signals.py b/python/restate/agents/_signals.py index 951a077e..d39f8c6c 100644 --- a/python/restate/agents/_signals.py +++ b/python/restate/agents/_signals.py @@ -117,6 +117,10 @@ class Started(Answering, total=False): steers: bool """Whether the turn takes steered messages at its steps (``steer``). A turn that does not say so (of another SDK, or of an agent whose framework cannot take them) gets none: each is answered as a turn of its own.""" + publishes: bool + """Whether ``start`` publishes the turn's start, and returns what ``publish`` returns: a followed turn of this + SDK's agents, which then makes no publish of its own. A turn that does not say so (of another SDK, or of an + agent that runs in the background) is not published by ``start``, which returns nothing.""" class _Held(TypedDict): @@ -178,9 +182,11 @@ async def _publish(ctx: ObjectContext) -> Notifications: revision = await _get(ctx, "revision", 0) + 1 ctx.set("revision", revision) notified: Notifications = {"revision": revision} - for watching in await _get(ctx, "watching", typing.cast(list[_Watching], [])): - ctx.resolve_awakeable(watching["awakeable"], notified) - ctx.clear("watching") + watching = await _get(ctx, "watching", typing.cast(list[_Watching], [])) + for one in watching: + ctx.resolve_awakeable(one["awakeable"], notified) + if watching: + ctx.clear("watching") return notified @@ -339,21 +345,27 @@ async def listen(ctx: ObjectContext, waiting: _WaitingOn) -> None: @signals.handler(ingress_private=True) -async def start(ctx: ObjectContext, turn: Started) -> None: - """Record that the agent's ``chat`` started answering ``turn``. +async def start(ctx: ObjectContext, turn: Started) -> Published | None: + """Record that the agent's ``chat`` started answering ``turn``, and publish it if the turn ``publishes``. A retry of the same invocation keeps the original start time. A previous turn that never closed itself (because its invocation was killed) is closed now. ``awaitTurn`` then watches the new turn, so that it is closed even if no other turn follows it. + + A turn that ``publishes`` gets what ``publish`` returns (the turn's first + publish, made here rather than by a call of its own); any other gets + ``None``, and publishes itself if it is followed. """ before = await _get(ctx, "turn", typing.cast(Started | None, None)) - if before is not None and before["invocationId"] == turn["invocationId"]: - return - if before is not None: - await _closed(ctx, before["invocationId"], gone=True) - ctx.set("turn", turn) - ctx.object_send(await_turn, ctx.key(), turn["invocationId"]) + if before is None or before["invocationId"] != turn["invocationId"]: + if before is not None: + await _closed(ctx, before["invocationId"], gone=True) + ctx.set("turn", turn) + ctx.object_send(await_turn, ctx.key(), turn["invocationId"]) + if not turn.get("publishes", False): + return None + return await _published(ctx) @signals.handler( @@ -461,9 +473,15 @@ async def publish(ctx: ObjectContext) -> Published: """Increase the conversation's revision after a change, and wake up every waiting ``watch``. It returns the new revision, and the messages steered to the turn being - answered that it has not taken yet, if any: a turn publishes at each tool - call's start and end, so it learns of them at no cost of its own. + answered that it has not taken yet, if any: a turn publishes as it starts + and after each step it stores that a tool waits for (a reply with tool + calls, the tools' results), so it learns of them at no cost of its own. """ + return await _published(ctx) + + +async def _published(ctx: ObjectContext) -> Published: + """Publish (``_publish``), and return the new revision and the messages the turn being answered has not taken.""" published: Published = {**await _publish(ctx)} inbox = await _inbox(ctx) if any("text" in one for one in inbox): diff --git a/python/restate/agents/_steering.py b/python/restate/agents/_steering.py index cf02d04f..25e550b8 100644 --- a/python/restate/agents/_steering.py +++ b/python/restate/agents/_steering.py @@ -15,8 +15,9 @@ order with the messages sent after it, and, if the turn being answered takes steered messages, holds it for that turn too (see ``_signals.steer``). -The turn learns of it from its next ``publish``, which it makes at each tool -call's start and end anyway. At its next step, once the model's reply has a +The turn learns of it from its next ``publish``, which it makes anyway as it +starts and after each write of steps that commits (a reply with tool calls, +the tools' results), and waits for. At its next step, once the model's reply has a result for each of its tool calls and before the next model request, the framework integration calls ``take_steers``: the turn stores each message once, as a message of the user that is a step of the turn, and the diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 7cd100c5..c7a5c08d 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -41,7 +41,7 @@ from restate.agents import _signals from restate.agents._calls import call_of from restate.agents._common import PerSession, ToolError, configured -from restate.agents._turn import stop_signal_of, tool_call_ended, tool_call_started +from restate.agents._turn import stop_signal_of, tool_call_started from restate.exceptions import TerminalError from restate.handler import Handler, handler_from_callable from restate.sessions import SessionContext @@ -184,9 +184,9 @@ async def run(self, args: object, call: ToolCall) -> str: async def run_tool(tool: Tool, args: object, call: ToolCall) -> str: """Run ``tool`` for a tool call, the way a framework integration should. - Unlike ``Tool.run``, it publishes the tool call's start and end to the - clients that follow the conversation, and first settles the earlier calls - whose results the framework has stored. + Unlike ``Tool.run``, it first settles the earlier calls whose results the + framework has stored (``tool_call_started``). It publishes nothing: the + turn publishes as it stores the call, and its result. Args: tool: The tool to run. @@ -197,10 +197,7 @@ async def run_tool(tool: Tool, args: object, call: ToolCall) -> str: The tool's result as text, as ``Tool.run`` returns it. """ await tool_call_started(call.ctx) - try: - return await tool.run(args, call) - finally: - await tool_call_ended(call.ctx) + return await tool.run(args, call) def _text(value: object) -> str: diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index 0e66947e..b978eaf1 100644 --- a/python/restate/agents/_turn.py +++ b/python/restate/agents/_turn.py @@ -171,10 +171,11 @@ def _unfollowed(ctx: SessionContext) -> bool: class Inbox: """The messages steered to a ``chat`` turn that it has not taken yet, as the companion last published them. - The companion returns them from every ``publish``, which a turn makes at - each tool call's start and end: so the turn knows of them without a call - of its own, and finds them at its next step (``take_steers``). The - newest revision wins, as tool calls that run at once publish at once. + The companion returns them from every ``publish``, which a turn makes as + it starts and after each step it stores that a tool waits for (a reply + with tool calls, the tools' results): so the turn knows of them without a + call of its own, and finds them at its next step (``take_steers``). The + newest revision wins, as writes that run at once publish at once. """ revision: int = -1 @@ -202,7 +203,7 @@ def inbox_of(ctx: SessionContext) -> Inbox | None: """What the attempt's turn has stored, while it answers.""" _SETTLING: ScopeKey[asyncio.Future[None]] = ScopeKey("settling") -"""The attempt's latest settling of calls; they run one after another, as parallel tool calls publish at once.""" +"""The attempt's latest settling of calls; they run one after another, as parallel tool calls settle at once.""" def _settled(ctx: SessionContext, stored: Stored) -> Awaitable[None]: @@ -238,37 +239,40 @@ async def publish(ctx: SessionContext) -> None: async def stepped(ctx: SessionContext) -> None: - """Publish a step of the turn, unless the agent runs in the background.""" + """Publish a step of the turn, unless the agent runs in the background. + + ``answer_once`` has it made after each write of the turn's steps that + commits (see ``Conversation.hold``): a reply with tool calls, the tools' + results. + """ if not _unfollowed(ctx): await publish(ctx) async def tool_call_started(ctx: SessionContext) -> None: - """Publish that a tool call of the turn is starting. + """Forget the earlier calls whose results are stored, before a tool call of the turn makes calls of its own. - Publishing also forgets the earlier calls whose results are stored, before - the tool makes calls of its own. A background agent only forgets them, - without publishing. + It publishes nothing: the turn publishes once the tool call is stored + (with its reply), and once its result is. """ - if not _unfollowed(ctx): - await publish(ctx) - return stored = _ANSWERING.get(ctx) if stored is not None: await _settled(ctx, stored) async def tool_call_ended(ctx: SessionContext) -> None: - """Publish that a tool call of the turn has ended, unless the agent runs in the background.""" - await stepped(ctx) + """Do nothing: kept for integrations of an earlier version, which called it as a tool call ended. + + The turn publishes once the call's result is stored, not as the call ends. + """ async def answering(ctx: SessionContext, stored: Stored, work: Callable[[], Awaitable[T]]) -> T: """Run ``work``, a framework's turn, forgetting the calls whose results it stores as it goes. - Every publish while ``work`` runs (such as a tool call's start and end), - and the end of ``work``, forget the Restate calls of the tool calls whose - results ``stored`` says are stored (see ``_calls.py``). + Every publish while ``work`` runs (after a step is stored), each tool + call's start, and the end of ``work``, forget the Restate calls of the + tool calls whose results ``stored`` says are stored (see ``_calls.py``). Args: ctx: The handler's context. @@ -310,7 +314,8 @@ async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionCon and cancels the calls its tools made; the handler then fails with the cancellation. If the attempt ends first, the turn stops with ``AttemptEnded`` and the next attempt carries on. The turn publishes when - it starts and when it ends. + it starts (the companion's ``start`` does) and when it ends (its + ``close``). An error ends the turn for good if it is a ``TerminalError``, or a model provider's refusal (``refusal_of``), as a 424. Any other error fails the @@ -331,9 +336,17 @@ async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionCon _INBOX.set(ctx, Inbox()) # From here until the turn closes, the companion knows about it: `status` and `stop` ask it. # The agent has ``settle`` (``chat()`` adds it), so the companion may send it this turn once it is gone. - started: _signals.Started = {**turn, "settles": True, "steers": steers} - await ctx.call(_signals.start, started, key=companion) - await stepped(ctx) + followed = not _unfollowed(ctx) + started: _signals.Started = {**turn, "settles": True, "steers": steers, "publishes": followed} + published = await ctx.call(_signals.start, started, key=companion) + if isinstance(published, dict): + # The companion published the turn's start as it started it. + inbox = _INBOX.get(ctx) + if inbox is not None: + inbox.saw(published) + else: + # A companion that does not publish as it starts a turn (of an earlier deployment): publish here. + await stepped(ctx) async def end() -> None: # Have the companion close the turn, which publishes its last change diff --git a/python/restate/agents/adk/_agent.py b/python/restate/agents/adk/_agent.py index d515e3f9..8aad3f69 100644 --- a/python/restate/agents/adk/_agent.py +++ b/python/restate/agents/adk/_agent.py @@ -113,7 +113,7 @@ from restate.agents._output import answer_with from restate.agents._steering import take_steers from restate.agents._tools import Tool, current_call, tools_in -from restate.agents._turn import tool_call_ended, tool_call_started +from restate.agents._turn import tool_call_started from restate.exceptions import TerminalError from restate.sessions import SessionContext @@ -266,7 +266,7 @@ async def add_stored(self, session: Session, event: Event) -> None: class _Restate(BasePlugin): - """Publish each tool call as it starts and ends, mark a tool's error as the tool's, and give the model the + """Settle the earlier calls as each tool call starts, mark a tool's error as the tool's, and give the model the messages steered to the turn.""" def __init__(self, ctx: SessionContext, sessions: _Sessions) -> None: @@ -294,11 +294,6 @@ async def before_tool_callback( ) -> None: await tool_call_started(self.ctx) - async def after_tool_callback( - self, *, tool: BaseTool, tool_args: dict[str, Any], tool_context: ToolContext, result: dict[str, Any] - ) -> None: - await tool_call_ended(self.ctx) - async def on_tool_error_callback( self, *, tool: BaseTool, tool_args: dict[str, Any], tool_context: ToolContext, error: Exception ) -> dict[str, Any] | None: diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 2d3fe2a5..8076f31b 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -106,7 +106,7 @@ from restate.agents._output import answer_with, asking from restate.agents._steering import take_steers from restate.agents._tools import Tool, ToolCall -from restate.agents._turn import settings_of, tool_call_ended, tool_call_started +from restate.agents._turn import settings_of, tool_call_started from restate.agents.langgraph._context import handler_context from restate.agents.langgraph._errors import terminal_of from restate.agents.langgraph._graph import AnyGraph, check, run_of @@ -364,7 +364,6 @@ async def run_tools(state: State, config: RunnableConfig) -> Any: from_tool(ctx, error) raise await conversation.turn.append(*_tool_messages(ran)) - await tool_call_ended(ctx) return ran def start(state: State) -> str: diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index 88a50ff3..3c2cec75 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -120,7 +120,6 @@ def forget(state: MessagesState) -> dict[str, Any]: from restate.agents import _signals from restate.agents._common import flush_held, from_tool from restate.agents._conversation import Conversation, Run, Turn -from restate.agents._turn import stepped from restate.agents.langgraph._checkpointer import SqlCheckpointSaver from restate.agents.langgraph._context import running as running_graph from restate.agents.langgraph._errors import terminal_of @@ -217,7 +216,7 @@ async def _copied( may not have loaded. If the graph trimmed the prompt away, they are the messages whose ids are not ``known``. ``known`` is the ids the conversation has, which this adds to; the messages the conversation gave - the graph are not copied back. Then publish. + the graph are not copied back. Each copy that commits publishes (``Conversation.hold``). """ at = next((i for i, m in enumerate(messages) if m.id == prompt), None) new = [ @@ -234,8 +233,6 @@ async def _copied( else: await turn.append(message) known.add(typing.cast(str, message.id)) - if new: - await stepped(ctx) def _prompt_id(turn: Turn) -> str: diff --git a/python/restate/agents/openai/_agent.py b/python/restate/agents/openai/_agent.py index 731b4fe6..c1e65f8e 100644 --- a/python/restate/agents/openai/_agent.py +++ b/python/restate/agents/openai/_agent.py @@ -158,7 +158,7 @@ from restate.agents._output import answer_with from restate.agents._steering import take_steers from restate.agents._tools import Tool, current_call, tools_in -from restate.agents._turn import tool_call_ended, tool_call_started +from restate.agents._turn import tool_call_started from restate.exceptions import TerminalError from restate.sessions import SessionContext @@ -296,7 +296,8 @@ def _current(conversation: Conversation, root: Agent[Any], everyone: dict[str, A class _Hooks(RunHooks[Any]): - """Store each step of a run as it comes, and publish each tool call as it starts and ends.""" + """Store each step of a run as it comes (each write that commits publishes), and settle the earlier calls as + each tool call starts.""" def __init__(self, ctx: SessionContext, conversation: Conversation, config: RunConfig) -> None: self.ctx = ctx @@ -365,7 +366,6 @@ async def on_handoff(self, context: RunContextWrapper[Any], from_agent: Agent[An async def on_tool_end(self, context: RunContextWrapper[Any], agent: Agent[Any], tool: SdkTool, result: Any) -> None: self.running -= 1 - await tool_call_ended(self.ctx) class _Steers: @@ -501,7 +501,6 @@ async def _resume( # Mark the error as the tool's, so that it is never mistaken for the model provider refusing. from_tool(ctx, error) raise - await tool_call_ended(ctx) outputs.extend(ITEMS.results([(call, output)])) await turn.append(*outputs) return agent diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index 0f871fb0..dfbc93c1 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -172,7 +172,7 @@ from restate.agents._output import answer_with from restate.agents._steering import take_steers from restate.agents._tools import Tool, current_call, tools_in, wait_for -from restate.agents._turn import stop_signal_of, tool_call_ended, tool_call_started +from restate.agents._turn import stop_signal_of, tool_call_started from restate.agents.pydantic_ai._stores import TurnsHistorySource, calling, conversation_id_of from restate.exceptions import TerminalError from restate.sessions import SessionContext @@ -357,8 +357,9 @@ class _Restate(AbstractCapability[Any]): """Make a Pydantic AI run durable in a Restate handler. It stores the requests before each model request and each response as it - arrives, and publishes each tool call as it starts and ends for clients - that follow the conversation. A tool's ``restate.TerminalError`` becomes its + arrives (each write that commits publishes, for clients that follow the + conversation), and settles the earlier calls as each tool call starts. A + tool's ``restate.TerminalError`` becomes its answer to the model; any other error fails the attempt. It records the errors of model requests, the only ones that can be a refusal. """ @@ -426,7 +427,6 @@ async def after_tool_execute( self, ctx: RunContext[Any], *, call: ToolCallPart, tool_def: ToolDefinition, args: Any, result: Any ) -> Any: calling(None) - await tool_call_ended(self.ctx) return result async def handle_deferred_tool_calls( @@ -464,7 +464,6 @@ async def on_tool_execute_error( if stop is not None and stop.reason is not None: raise stop.reason from error if isinstance(error, TerminalError) and not is_cancellation(error): - await tool_call_ended(self.ctx) return f"Error: {error.message}" # Mark the error as the tool's, so that it is never mistaken for the model provider refusing. from_tool(self.ctx, error) diff --git a/tests/agents_publishes.py b/tests/agents_publishes.py new file mode 100644 index 00000000..6d639d3b --- /dev/null +++ b/tests/agents_publishes.py @@ -0,0 +1,327 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""When a followed turn publishes, and what a follower then sees: in this process, for each framework. + +A turn publishes as it starts (the companion's ``start`` does), after each +write of steps that commits (a reply with tool calls, the tools' results), and +as it closes. A follower reads the ``messages`` view once told: at each +publish, the companion here reads it as a follower would, at once. +""" + +# pylint: disable=import-error,import-outside-toplevel + +from __future__ import annotations + +import json +import typing +from collections.abc import AsyncGenerator, Callable +from dataclasses import dataclass, field +from typing import Any + +import pytest + +from restate.agents import tool +from restate.agents._object import AgentObject + +pytestmark = pytest.mark.anyio + + +@pytest.fixture +def anyio_backend(): + return "asyncio" + + +CALLS = 3 +"""How many tool calls the model's first reply makes, at once.""" + + +@dataclass +class Seen: + """What a follower saw: at each companion call, and as each model request went out.""" + + events: list[tuple[str, int]] = field(default_factory=list[tuple[str, int]]) + """Each companion call that answers with a revision, or a model request, and the results the view showed then.""" + steering: list[dict[str, str]] = field(default_factory=list[dict[str, str]]) + """The messages steered to the turn, which each publish returns.""" + + def companion_calls(self) -> list[str]: + return [event for event, _ in self.events if event != "model"] + + +class Companion: + """The agent's companion in this process: each publish (and a ``start`` that publishes) reads the ``messages`` + view, as a follower told of it would, and returns a new revision with the messages steered to the turn.""" + + def __init__(self, obj: AgentObject, backend: Any, seen: Seen, *, start_publishes: bool = True) -> None: + self.obj = obj + self.backend = backend + self.seen = seen + self.start_publishes = start_publishes + self.revision = 0 + + async def results(self) -> int: + """Return how many tool results the ``messages`` view shows.""" + from sessions_object import Host, request + + from restate.sessions.internal import invoke + + view = Host(self.backend, service=self.obj.name).view() + page = await invoke(self.obj, "messages", view, request(), {"limit": 1000}) + return sum("output" in part for message in page["messages"] for part in message["parts"]) + + async def answer(self, handler: str, arg: Any) -> bytes: + if handler == "start" and not (self.start_publishes and json.loads(arg).get("publishes")): + self.seen.events.append(("start", await self.results())) + return b"null" + if handler in ("start", "publish"): + self.revision += 1 + self.seen.events.append((handler, await self.results())) + return json.dumps({"revision": self.revision, "steering": self.seen.steering}).encode() + assert handler == "close", f"the companion has no answer for {handler}" + self.seen.events.append(("close", await self.results())) + return b"null" + + +def follow(obj: AgentObject, seen: Seen, *, start_publishes: bool = True) -> Callable[[str], Any]: + """Return a ``chat`` of ``obj`` in this process, each call a turn, whose follower records what it sees in ``seen``.""" + from sessions_database import Backend + from sessions_object import Answer, Host, request + + from restate.sessions.internal import invoke + + backend = Backend() + companion = Companion(obj, backend, seen, start_publishes=start_publishes) + + class Deferred(Answer): + def __init__(self, handler: str, arg: Any) -> None: + super().__init__(b"null") + self.handler, self.arg = handler, arg + + def __await__(self): + return companion.answer(self.handler, self.arg).__await__() + + class Followed(Host): + def generic_call(self, service, handler, arg, key=None, idempotency_key=None, headers=None): + self.calls.append((service, handler, arg, key, idempotency_key)) + return Deferred(handler, arg) + + async def chat(text: str) -> Any: + return await invoke(obj, "chat", Followed(backend, service=obj.name), request(), text) + + return chat + + +def lookup_of(seen: Seen, steer: str | None = None): + """The tool the model calls: it steers ``steer`` to the turn while it runs, if given.""" + + @tool(name="lookup") + async def lookup(about: str) -> str: + """Look something up.""" + if steer is not None and not seen.steering: + seen.steering.append({"id": "steer-1", "text": steer}) + return f"found {about}" + + return lookup + + +# --- The frameworks' agents, each with a model that calls ``lookup`` CALLS times at once, then answers ---------- + + +def pydantic_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: + from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart, UserPromptPart + from pydantic_ai.models.function import AgentInfo, FunctionModel + + from restate.agents import pydantic_ai + + async def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + seen.events.append(("model", 0)) + if any(isinstance(m, ModelResponse) for m in messages) or not calls: + told = [str(p.content) for m in messages for p in getattr(m, "parts", []) if isinstance(p, UserPromptPart)] + return ModelResponse(parts=[TextPart(f"answered after {', '.join(told)}")]) + return ModelResponse( + parts=[ToolCallPart("lookup", {"about": f"t{i}"}, tool_call_id=f"c{i}") for i in range(calls)] + ) + + return pydantic_ai.agent("followed", model=FunctionModel(model), tools=[lookup_of(seen, steer)], **options) + + +def openai_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: + from agents import Model, ModelResponse, Usage + from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText + + from restate.agents import openai + + class Scripted(Model): + async def get_response(self, system_instructions, input, *args, **kwargs): # noqa: A002 + seen.events.append(("model", 0)) + items = [{"role": "user", "content": input}] if isinstance(input, str) else list(input) + if any(item.get("type") == "function_call_output" for item in items) or not calls: + told = [str(item["content"]) for item in items if item.get("role") == "user"] + text = f"answered after {', '.join(told)}" + out: list[Any] = [ + ResponseOutputMessage( + id="msg", + type="message", + role="assistant", + status="completed", + content=[ResponseOutputText(type="output_text", text=text, annotations=[])], + ) + ] + else: + out = [ + ResponseFunctionToolCall( + type="function_call", call_id=f"c{i}", name="lookup", arguments=f'{{"about": "t{i}"}}' + ) + for i in range(calls) + ] + return ModelResponse(output=out, usage=Usage(), response_id=None) + + def stream_response(self, *args, **kwargs): + raise NotImplementedError + + return openai.agent("followed", model=Scripted(), tools=[lookup_of(seen, steer)], **options) + + +def adk_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: + from google.adk.models import BaseLlm, LlmRequest, LlmResponse + from google.genai import types + + from restate.agents import adk + + class Scripted(BaseLlm): + model: str = "scripted" + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + seen.events.append(("model", 0)) + contents = list(llm_request.contents) + responded = any(part.function_response for content in contents for part in content.parts or []) + if responded or not calls: + told = [ + str(part.text) + for content in contents + if content.role == "user" + for part in content.parts or [] + if part.text + ] + parts = [types.Part(text=f"answered after {', '.join(told)}")] + else: + parts = [ + types.Part(function_call=types.FunctionCall(id=f"c{i}", name="lookup", args={"about": f"t{i}"})) + for i in range(calls) + ] + yield LlmResponse(content=types.Content(role="model", parts=parts)) + + return adk.agent("followed", model=Scripted(), tools=[lookup_of(seen, steer)], **options) + + +AGENTS = {"pydantic_ai": pydantic_agent, "openai": openai_agent, "adk": adk_agent} + + +def make(framework: str, seen: Seen, steer: str | None = None, **options: Any) -> AgentObject: + return typing.cast(Callable[..., AgentObject], AGENTS[framework])(seen, steer, **options) + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_follower_sees_the_tools_results_before_the_model_is_asked_again(framework: str): + seen = Seen() + chat = follow(make(framework, seen), seen) + await chat("hello") + # The second model request: the publish before it showed the results, which were committed just before it. + second = [i for i, (event, _) in enumerate(seen.events) if event == "model"][1] + before = [(event, results) for event, results in seen.events[:second] if event != "model"] + assert before[-1] == ("publish", CALLS) + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_turn_of_tool_calls_publishes_as_it_starts_after_each_commit_and_as_it_closes(framework: str): + seen = Seen() + chat = follow(make(framework, seen), seen) + await chat("hello") + # The start; the reply with the calls (no results yet); the results; the close. + assert [e for e in seen.events if e[0] != "model"] == [ + ("start", 0), + ("publish", 0), + ("publish", CALLS), + ("close", CALLS), + ] + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_plain_turn_publishes_as_it_starts_and_as_it_closes_only(framework: str): + seen = Seen() + chat = follow(make(framework, seen, calls=0), seen) + await chat("hello") + assert seen.companion_calls() == ["start", "close"] + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_message_steered_while_the_tools_run_is_taken_at_the_next_step(framework: str): + seen = Seen() + chat = follow(make(framework, seen, "use Celsius"), seen) + # The tool steered it: the publish after the results returned it, and the next request had it. + assert await chat("hello") == "answered after hello, use Celsius" + assert seen.companion_calls() == ["start", "publish", "publish", "close"] + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_companion_whose_start_returns_nothing_still_gets_the_turns_first_publish(framework: str): + seen = Seen() + chat = follow(make(framework, seen, calls=0), seen, start_publishes=False) + await chat("hello") + # An older companion: the turn publishes its start itself. + assert seen.companion_calls() == ["start", "publish", "close"] + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_background_agent_publishes_only_its_end(framework: str): + seen = Seen() + chat = follow(make(framework, seen, background=True), seen) + await chat("hello") + # Its start publishes nothing, nor do its commits: only its close does. + assert seen.companion_calls() == ["start", "close"] + + +class _State: + """A companion's context in this process, for its state only: what it writes, in order.""" + + def __init__(self, **state: Any) -> None: + self.state = dict(state) + self.writes: list[tuple[str, str]] = [] + self.resolved: list[str] = [] + + async def get(self, key: str, type_hint: Any = None) -> Any: + return self.state.get(key) + + def set(self, key: str, value: Any) -> None: + self.writes.append(("set", key)) + self.state[key] = value + + def clear(self, key: str) -> None: + self.writes.append(("clear", key)) + self.state.pop(key, None) + + def resolve_awakeable(self, awakeable: str, value: Any) -> None: + self.resolved.append(awakeable) + + +async def test_a_publish_clears_the_watches_only_if_there_are_any(): + from restate.agents import _signals + + quiet = _State(revision=4) + assert await _signals._publish(typing.cast(Any, quiet)) == {"revision": 5} # pyright: ignore[reportPrivateUsage] + # Nobody watched: one write, of the revision. + assert quiet.writes == [("set", "revision")] + + watched = _State(revision=4, watching=[{"revision": 4, "awakeable": "a1"}]) + await _signals._publish(typing.cast(Any, watched)) # pyright: ignore[reportPrivateUsage] + assert watched.resolved == ["a1"] and watched.writes == [("set", "revision"), ("clear", "watching")] diff --git a/tests/agents_signals.py b/tests/agents_signals.py index 99980d83..98a5aff7 100644 --- a/tests/agents_signals.py +++ b/tests/agents_signals.py @@ -40,13 +40,15 @@ def anyio_backend(): @driver.handler() -async def start_turn(ctx: restate.Context, request: dict[str, str]) -> None: +async def start_turn(ctx: restate.Context, request: dict[str, str]) -> Published | None: turn: Started = {"invocationId": request["invocation"], "since": "2026-11-04T00:00:00.000Z"} if request.get("settles"): turn["settles"] = True if request.get("steers"): turn["steers"] = True - await ctx.object_call(_signals.start, key=request["key"], arg=turn) + if request.get("publishes"): + turn["publishes"] = True + return await ctx.object_call(_signals.start, key=request["key"], arg=turn) @driver.handler() @@ -326,3 +328,35 @@ async def test_a_turn_takes_at_most_its_share_of_steered_messages(harness: Harne assert "steering" not in await published(harness, k) assert await steer(harness, k, "later", "and more") == {"id": "later", "turn": None} assert "steering" not in await published(harness, k) + + +async def test_a_turn_that_publishes_is_published_as_it_starts_and_gets_what_publish_returns( + harness: HarnessEnvironment, +): + k = key() + + async def revision() -> int: + return (await harness.client.object_call(_signals.notifications, key=k, arg=None))["revision"] + + # A turn that does not say it publishes (of another SDK, or a background agent's): nothing is published. + before = await revision() + assert await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_quiet"}) is None + assert await revision() == before + await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_quiet"}) + + # A followed turn: its start is its first publish, with the new revision. + before = await revision() + started = await harness.client.service_call( + start_turn, arg={"key": k, "invocation": "inv_loud", "steers": "yes", "publishes": "yes"} + ) + assert started == {"revision": before + 1} + # A retried attempt starts it again: it publishes again, with the messages steered to it meanwhile. + await steer(harness, k, "s1", "use Celsius") + again = await harness.client.service_call( + start_turn, arg={"key": k, "invocation": "inv_loud", "steers": "yes", "publishes": "yes"} + ) + assert again is not None and again["revision"] == await revision() + assert again.get("steering") == [{"id": "s1", "text": "use Celsius"}] + status = typing.cast(Status, await harness.client.object_call(_signals.status, key=k, arg=None)) + assert status["answering"] is not None and status["answering"]["invocationId"] == "inv_loud" + await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_loud"}) From 2d6c44ad0664870c2ad886fcc4fae755324a1298 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 10:28:26 +0000 Subject: [PATCH 3/8] Agents: a result of a call that reuses an earlier turn's id is the turn's own A provider may give a tool call the id of an earlier turn's call. The integrations asked whether the window knew the result by call id, found the earlier turn's, and never stored the new one: later turns told the model the call was interrupted (Pydantic AI), or sent a call without output that a strict provider rejects (OpenAI). `Conversation.new_results` now takes a result for a call the turn being answered has open as the turn's own once it comes after the latest call of that id in the run's items; one before it is an earlier turn's, handed back, and is not stored again. Pydantic AI's `store` finds results through it. Co-Authored-By: Claude Fable 5.2 --- python/restate/agents/_conversation.py | 22 ++++++++++- python/restate/agents/pydantic_ai/_agent.py | 14 ++----- tests/agents_openai.py | 44 +++++++++++++++++++++ tests/agents_pydantic_ai.py | 37 +++++++++++++++++ 4 files changed, 105 insertions(+), 12 deletions(-) diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index a4349141..0d5f59a8 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -1259,8 +1259,28 @@ def new_results(self, items: Sequence[Item]) -> list[Item]: Results are found by their calls' ids, not by their place: what a framework hands back may be trimmed, or be the whole conversation again. + + But a provider may give a call of the turn being answered the id of an + earlier turn's call, whose result the window has: a result for a call + the turn has open is the turn's own once it comes after the latest + call of that id in ``items``. One before it is the earlier turn's, + handed back. """ - return [item for item in items if any(not self.knows_result(r.id) for r in self.items.read(item).results)] + running = self.running() + open_ids = {call.id for call in running.open_calls()} if running is not None else set() + parsed = [self.items.read(item) for item in items] + latest = {call.id: i for i, read in enumerate(parsed) for call in read.calls} + + def new(i: int, call_id: str) -> bool: + if call_id in open_ids and i > latest.get(call_id, -1): + return True + return not self.knows_result(call_id) + + return [ + item + for i, (item, read) in enumerate(zip(items, parsed, strict=True)) + if any(new(i, r.id) for r in read.results) + ] def window(self) -> int: """Return the position of the window's first turn: the turns before it are in the transcript only.""" diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index dfbc93c1..be5bc857 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -346,13 +346,6 @@ def _feedback_of(message: ModelMessage) -> list[str]: return [to_json(p, bytes_mode="base64").decode() for p in message.parts if isinstance(p, RetryFeedbackPart)] -def _result_ids(message: ModelMessage) -> set[str]: - """Return the ids of the tool calls whose results ``message`` holds.""" - if not isinstance(message, ModelRequest): - return set() - return {p.tool_call_id for p in message.parts if isinstance(p, ToolReturnPart) and p.tool_call_id} - - class _Restate(AbstractCapability[Any]): """Make a Pydantic AI run durable in a Restate handler. @@ -372,9 +365,7 @@ def __init__(self, ctx: SessionContext, conversation: Conversation, history: Seq self.feedback = {feedback for m in history for feedback in _feedback_of(m)} """Pydantic AI's feedback that the conversation has (each ``RetryFeedbackPart``), as JSON.""" - def _new(self, message: ModelMessage) -> bool: - if any(not self.conversation.knows_result(i) for i in _result_ids(message)): - return True + def _new_feedback(self, message: ModelMessage) -> bool: # Found by its parts, not the whole request: the run's request may be a stored one, merged or resumed. return any(feedback not in self.feedback for feedback in _feedback_of(message)) @@ -391,7 +382,8 @@ async def store(self, messages: Sequence[ModelMessage]) -> None: a request of its own. What it makes of them is the model's to see, not the conversation's. """ - new = [m for m in messages if self._new(m)] + results = {id(m) for m in self.conversation.new_results(messages)} + new = [m for m in messages if id(m) in results or self._new_feedback(m)] if new: await self.conversation.turn.append(*new) self.feedback.update(feedback for m in new for feedback in _feedback_of(m)) diff --git a/tests/agents_openai.py b/tests/agents_openai.py index a70bd336..83a0289b 100644 --- a/tests/agents_openai.py +++ b/tests/agents_openai.py @@ -1573,3 +1573,47 @@ def test_steered_messages_keep_their_order_and_go_after_their_item_or_before_wha # The item they came after taken out: before the first item that came after them, in order; else at the end. assert spliced([hi, later], taken, seen) == [hi, first, second, later] assert spliced([hi], taken, seen) == [hi, first, second] + + +async def test_a_tool_call_whose_id_an_earlier_turns_call_had_still_gets_its_output_stored(): + """A provider may give a call the id of an earlier turn's call: its output is the turn's own, and stored.""" + + def answer(system: str | None, items: list[Item]) -> list[Output]: + calls_made = [i["call_id"] for i in items if i.get("type") == "function_call"] + outputs = [i["call_id"] for i in items if i.get("type") == "function_call_output"] + # As a strict provider checks a request: each call has its output. + assert len(calls_made) == len(outputs), f"a call has no output: {items}" + last = items[-1] + if last.get("role") == "user": + return [calls("call-1", "lookup", json.dumps({"about": last["content"]}))] + return [said(f"answer: {last['output']}")] + + @tool + async def lookup(about: str) -> str: + """Look something up.""" + return f"found {about}" + + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + from sessions_object import Host, request # pylint: disable=import-error,import-outside-toplevel + + from restate.agents._conversation import STEPS # pylint: disable=import-outside-toplevel + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + obj = asker(model=Scripted(answer), tools=[lookup]) + + @obj.handler + async def outputs(ctx: sessions.SessionContext, _: Any = None) -> list[str]: + stored = [json.loads(str(item)) for [item] in await ctx.db.values(f"SELECT item FROM {STEPS} ORDER BY n")] + return [item["output"] for item in stored if item.get("type") == "function_call_output"] + + backend = Backend() + answers = [ + await invoke(obj, "ask", Host(backend, service="asker", answers=PUBLISHED), request(), w) + for w in ("one", "two", "three") + ] + assert answers == ["answer: found one", "answer: found two", "answer: found three"] + assert await invoke(obj, "outputs", Host(backend, service="asker"), request(), None) == [ + "found one", + "found two", + "found three", + ] diff --git a/tests/agents_pydantic_ai.py b/tests/agents_pydantic_ai.py index def17862..145f9072 100644 --- a/tests/agents_pydantic_ai.py +++ b/tests/agents_pydantic_ai.py @@ -1301,3 +1301,40 @@ async def answer(ctx: sessions.SessionContext, text: str) -> str: "and be brief", "and then?", ] + + +async def test_a_tool_call_whose_id_an_earlier_turns_call_had_still_gets_its_result_stored(): + """A provider may give a call the id of an earlier turn's call: its result is the turn's own, and stored.""" + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + from sessions_object import Host, request # pylint: disable=import-error,import-outside-toplevel + + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + requests: list[list[ModelMessage]] = [] + + def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + requests.append(list(messages)) + last = messages[-1] + if isinstance(last, ModelRequest) and any(isinstance(p, UserPromptPart) for p in last.parts): + word = _prompts(messages)[-1] + return ModelResponse(parts=[ToolCallPart("lookup", {"about": word}, tool_call_id="call-1")]) + return ModelResponse(parts=[TextPart("answered")]) + + def lookup(about: str) -> str: + """Look something up.""" + return f"found {about}" + + obj = pydantic_ai.agent("pai_reused_ids", model=FunctionModel(model), tools=[lookup]) + backend = Backend() + for word in ("one", "two", "three"): + host = Host(backend, service=obj.name, answers={"publish": b"null", "start": b"null", "close": b"null"}) + assert await invoke(obj, "chat", host, request(), word) == "answered" + # The last request has each turn's result, its own, and none taken for interrupted. + returned = [ + str(part.content) + for m in requests[-1] + if isinstance(m, ModelRequest) + for part in m.parts + if isinstance(part, ToolReturnPart) + ] + assert returned == ["found one", "found two", "found three"] From 6b3a4093054c82169d581a39931aa58cbeab0b8c Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 10:40:29 +0000 Subject: [PATCH 4/8] Agents: a ledger write cancelled during its commit keeps its bookkeeping `Stop.guard` cancels the run's task. A write of the ledger cancelled while it waited for its commit left its rows committed but the ledger's next position, `seq` and held rows as they were, so the turn's close wrote over them: the cancellation surfaced as a unique constraint error and the turn stayed open. As the steer store already was, every other ledger write (`Turn._insert`, `answered`, `snapshot`, `drop`, `close`, `Conversation.ask`, `snapshot_between`, `flush`) is now shielded (`_shielded`): it finishes with its bookkeeping while the caller gets the cancellation. `close` takes the lock before it reads the turn's open calls and status, so that it closes after a write under way, with the calls it made, or not at all once it stored the answer. Co-Authored-By: Claude Fable 5.2 --- python/restate/agents/_conversation.py | 54 +++++++++++++++++++++----- tests/agents_conversation.py | 54 ++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 10 deletions(-) diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index 0d5f59a8..04a2d22e 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -86,13 +86,14 @@ import asyncio import bisect import contextlib +import functools import json import logging import time import typing -from collections.abc import AsyncIterator, Awaitable, Callable, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Sequence from dataclasses import dataclass, field, replace -from typing import Any, Literal, Protocol +from typing import Any, Literal, ParamSpec, Protocol, TypeVar from restate.agents._api import Steer from restate.agents._calls import outcomes @@ -585,6 +586,29 @@ def carried(self) -> tuple[str | None, str | None]: return (latest.agent, latest.carried) if latest is not None else (None, None) +_P = ParamSpec("_P") +_T = TypeVar("_T") + + +def _shielded(write: Callable[_P, Coroutine[Any, Any, _T]]) -> Callable[_P, Coroutine[Any, Any, _T]]: + """Make ``write``, a write of the ledger, finish once it started, even if its caller is cancelled meanwhile. + + A stop cancels the run's task (``Stop.guard``). Cancelled while it waits + for its commit, a write would leave its rows committed but the ledger's + bookkeeping (the next position, the last ``seq``, the held rows) as it + was, and the next write (the turn's close) would store over them. So a + write and its bookkeeping happen together, or not at all: the caller gets + the cancellation at once, and the write finishes meanwhile, holding the + ledger's lock, so that whatever writes next (the close) writes after it. + """ + + @functools.wraps(write) + async def shielded(*args: _P.args, **kwargs: _P.kwargs) -> _T: + return await asyncio.shield(write(*args, **kwargs)) + + return shielded + + _INSERT_STEP = f"INSERT INTO {STEPS} (n, turn, seq, item, reply, agent) VALUES (?, ?, ?, ?, ?, ?)" """Store a step of a turn: its position, its turn's, its ``seq``, its item, and the reply and agent it is of.""" @@ -837,6 +861,7 @@ def carry(self, *, agent: str | None = None, state: str | None = None) -> None: if state is not None: self._carried["carried"] = state + @_shielded async def snapshot(self, items: Sequence[Item]) -> None: """Store ``items`` as a snapshot of the conversation, as the framework's own compaction left it. @@ -851,6 +876,7 @@ async def snapshot(self, items: Sequence[Item]) -> None: self._check_writable() await ledger.store_snapshot(list(items), self.n) + @_shielded async def drop(self, reply: int) -> None: """Drop the turn's reply ``reply``, and what came after it, to ask the model again. @@ -890,6 +916,7 @@ async def drop(self, reply: int) -> None: for entry in kept: self._index(entry, note=False) + @_shielded async def answered(self, answer: str, status: Status = "answered") -> None: """Store the turn's ``answer``, how it ended (``status``) and what it carries, in one update of its row. @@ -936,6 +963,7 @@ def _agent(self) -> str | None: replied = next((e.agent for e in reversed(self._entries) if e.reply is not None), None) return self._carried.get("agent") or replied or self._ledger.carried()[0] + @_shielded async def close(self, ctx: SessionContext, status: Status, note: str) -> None: """Close the turn (stopped, refused, failed, or ended once its invocation is gone) with ``note``. @@ -950,18 +978,20 @@ async def close(self, ctx: SessionContext, status: Status, note: str) -> None: crash between them would leave the turn without its answer, and the next attempt would close it again, with a second note. """ - if self.status is not None: - return ledger = self._ledger items = ledger.items - calls = self.open_calls() - returned = await outcomes(ctx, [call.id for call in calls], invocation_of(self.id)) - closing = [ - *items.results([(call, cut_off_result(returned, call.id, note)) for call in calls]), - items.reply(note), - ] + # The lock first: a write under way when the run was cancelled finishes (see ``_shielded``), and the turn + # closes after it, with the calls it made, or its answer. async with ledger.lock: + if self.status is not None: + return self._check_writable() + calls = self.open_calls() + returned = await outcomes(ctx, [call.id for call in calls], invocation_of(self.id)) + closing = [ + *items.results([(call, cut_off_result(returned, call.id, note)) for call in calls]), + items.reply(note), + ] n, seq = ledger.next, ledger.seq added = [Entry(n + i, item) for i, item in enumerate(closing)] at = added[-1].n # the note: the turn's last message, and its answer @@ -1006,6 +1036,7 @@ def making(n: int) -> list[Entry]: await published() return n + @_shielded async def _insert(self, made: Callable[[int], list[Entry]]) -> int: """Store the entries that ``made`` makes from the next position, and return that position. @@ -1368,6 +1399,7 @@ def history(self) -> list[Item]: # --- Writing --- + @_shielded async def ask(self, prompt_id: str, text: str) -> Turn: """Start the turn of prompt ``prompt_id``, with ``text``, and return it: the latest turn now.""" ledger = self._ledger @@ -1389,6 +1421,7 @@ async def ask(self, prompt_id: str, text: str) -> Turn: ledger.latest = turn return turn + @_shielded async def snapshot_between(self, items: Sequence[Item]) -> None: """Store ``items`` as a snapshot of the whole conversation, between turns (in ``compact``, say). @@ -1433,6 +1466,7 @@ def hold(self, published: Callable[[], Awaitable[None]] | None = None) -> None: self._ledger.holds = True self._ledger.published = published + @_shielded async def flush(self) -> None: """Write the rows the turn's run holds (see ``_Ledger``), if any. diff --git a/tests/agents_conversation.py b/tests/agents_conversation.py index 962e622d..3bb19d0b 100644 --- a/tests/agents_conversation.py +++ b/tests/agents_conversation.py @@ -1146,3 +1146,57 @@ async def seqs(ctx: sessions.SessionContext) -> list[Any]: assert positions == list(range(1, 8)) assert [tuple(entry) for entry in stored] == seen[0] assert await invoke(obj, "seqs", _syncing(backend), request()) == list(range(2, 9)) + + +@pytest.mark.parametrize("write", ["append", "answered", "snapshot", "drop"]) +async def test_a_write_cancelled_while_it_commits_keeps_its_bookkeeping_and_the_turn_closes_after_it(write: str): + """A stop cancels the run (``Stop.guard``), here while a write waits for its commit: the write lands with the + ledger's positions, ``seq`` and held rows, and the turn's close stores after it rather than over it.""" + gate: dict[str, Any] = {} + + class Gated(Host): + async def sync(self) -> None: + if gate.pop("armed", False): + gate["blocked"].set() + await gate["release"].wait() + await super().sync() + + async def run(conversation: Conversation) -> str: + turn = conversation.turn + await turn.append_reply(_calling("c1")) + await turn.append({"result": "c1", "output": "done"}) + # Held: written with the next write. + await turn.append_reply({"role": "assistant", "text": "noted"}) + writes = { + "append": lambda: turn.append_reply(_calling("c2")), + "answered": lambda: turn.answered("noted"), + "snapshot": lambda: turn.snapshot([{"role": "user", "text": "summary"}]), + "drop": lambda: turn.drop(typing.cast(int, turn.last_reply())), + } + gate.update(armed=True, blocked=asyncio.Event(), release=asyncio.Event()) + task = asyncio.ensure_future(writes[write]()) + await gate["blocked"].wait() + task.cancel() + gate["release"].set() + with pytest.raises(asyncio.CancelledError): + await task + raise restate.TerminalError("cancelled", status_code=409) + + obj = _turn_rows() + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await answer_once(ctx, text, ITEMS, run) + + backend = Backend() + with pytest.raises(restate.TerminalError, match="cancelled"): + await invoke(obj, "ask", Gated(backend, service="asker", answers=PUBLISHED), request(), "hi") + [turn] = await invoke(obj, "rows", _syncing(backend), request()) + texts = [item.get("text", item.get("output")) for _, item in turn["items"]] + expected = { + "append": (CANCELLED, ["", "done", "noted", "", CANCELLED, CANCELLED]), + "answered": ("noted", ["", "done", "noted"]), + "snapshot": (CANCELLED, ["", "done", "noted", CANCELLED]), + "drop": (CANCELLED, ["", "done", CANCELLED]), + } + assert (turn["answer"], texts) == expected[write] From 20f0fb1dc72da0611e95269935425a96599e60b6 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 10:42:11 +0000 Subject: [PATCH 5/8] Agents: a plan read in a tool call commits nothing on its own The plan store journaled each read in a tool call in a commit of its own, though a read's row matters only if a write follows it in the call. The reads are now kept on the call and written in its next write's transaction: a call that only reads commits nothing, and a call that runs again replays its reads as before, since one that crashed before its next write had not changed the plan since them. Co-Authored-By: Claude Fable 5.2 --- docs/agents/frameworks.md | 4 +- python/restate/agents/pydantic_ai/_stores.py | 31 ++++++++++---- tests/agents_pydantic_ai_stores.py | 44 ++++++++++++++++++++ 3 files changed, 70 insertions(+), 9 deletions(-) diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index a0cfc039..0057273f 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -265,7 +265,9 @@ handlers have no `ctx.run`; it disables harness features), `StepPersistence` order in the call. Replaying only the writes isn't enough: the call's earlier writes are committed, and reading them changes what it writes next (which steps a status change unblocks), so its writes no longer line up. - This is what the harness's own durable stores assume ("the store's + A read's journal row is kept in memory and written with the call's next + write: only a write makes it matter, so a call that only reads commits + nothing. This is what the harness's own durable stores assume ("the store's recorded result keeps the original ids"). - **A snapshot is the history the model will get,** so: strip system prompt parts (or `ReinjectSystemPrompt` stops reinjecting it); drop the input diff --git a/python/restate/agents/pydantic_ai/_stores.py b/python/restate/agents/pydantic_ai/_stores.py index b72be899..0d2edc01 100644 --- a/python/restate/agents/pydantic_ai/_stores.py +++ b/python/restate/agents/pydantic_ai/_stores.py @@ -22,7 +22,9 @@ progress, and ``clear`` empties them with the rest of the conversation. A tool call runs again after a crash, so the plan's reads and writes in a tool call are journaled: the call gets what each returned before, and so makes the -writes it made before, each done once (see ``_journaled``). A spill is named +writes it made before, each done once (see ``_journaled``). A read is +journaled with the write that follows it: a call that only reads commits +nothing. A spill is named by the tool call, as the harness's keys are made of the run's id, which every attempt has its own of. """ @@ -77,6 +79,9 @@ class _Call: id: str steps: int = 0 """How many reads and writes of the plan it made (``_journaled``).""" + reads: list[tuple[int, str]] = dataclasses.field(default_factory=list[tuple[int, str]]) + """The reads of the plan it made since its last write, by their order in it, with what each returned: journaled + with its next write, if it makes one (``_journaled``).""" writes: int = 0 """How many writes of memory it made (``_memory.py``).""" @@ -123,9 +128,13 @@ async def _journaled( writes it made before, each done once. Its reads are journaled too: its earlier writes are committed, and reading them would change what it writes next (which steps a status change unblocks, say), so that its writes would - not line up with those it made before. Outside a tool call, ``op`` is done - as it comes. The journals of earlier invocations, whose calls no longer - run, are dropped. + not line up with those it made before. A read's journal row matters only + if a write follows it, so it is kept on the call and written in the next + write's transaction: a call that only reads commits nothing, and one that + crashes before its next write did not change the plan since those reads, + and so reads it as it did. Outside a tool call, ``op`` is done as it + comes. The journals of earlier invocations, whose calls no longer run, + are dropped. """ ctx = running() await SCHEMA.ensure(ctx.db) @@ -142,13 +151,19 @@ async def _journaled( ) if done is not None: return load(str(done["result"])) + if not writes: + result = await op(ctx.db) + call.reads.append((n, dump(result))) + return result async with ctx.db.transaction() as tx: result = await op(tx) await tx.execute(f"DELETE FROM {JOURNAL} WHERE invocation != ?", [invocation]) - await tx.execute( - f"INSERT INTO {JOURNAL} (invocation, call, n, result) VALUES (?, ?, ?, ?)", - [invocation, call.id, n, dump(result)], - ) + for step, read in [*call.reads, (n, dump(result))]: + await tx.execute( + f"INSERT INTO {JOURNAL} (invocation, call, n, result) VALUES (?, ?, ?, ?)", + [invocation, call.id, step, read], + ) + call.reads = [] return result diff --git a/tests/agents_pydantic_ai_stores.py b/tests/agents_pydantic_ai_stores.py index 9c76a272..77362ad1 100644 --- a/tests/agents_pydantic_ai_stores.py +++ b/tests/agents_pydantic_ai_stores.py @@ -296,3 +296,47 @@ def lookup(about: str) -> str: # Each run's messages are its turn's, as stored: the search's own run has its prompt first. history = await _history(chat, runs[-1].run_id) assert _prompt(history[:1]) == "search fact" + + +async def test_a_tool_call_that_only_reads_the_plan_commits_nothing(monkeypatch: pytest.MonkeyPatch): + """A read's journal row is needed only if a write follows in the same tool call: it is written with that write.""" + from restate.agents.pydantic_ai import _stores # pylint: disable=import-outside-toplevel + + model = calls( + { + "plan": ("write_plan", {"items": [{"content": "Pack"}, {"content": "Travel"}]}), + "read": ("read_plan", {}), + } + ) + mine: Any = Agent(model, capabilities=[Planning(store=pydantic_ai.plan_store())]) + chat = Chat(pydantic_ai.agent("pai_plans_read", agent=mine)) + await chat("plan the trip") + + @chat.obj.handler(name="journal") + async def journal(ctx: Any, _: Any = None) -> int: + [[count]] = await ctx.db.values(f"SELECT COUNT(*) FROM {_stores.JOURNAL}") + return int(count) + + journaled = await invoke(chat.obj, "journal", chat.host(), request(), None) + syncs: list[bool] = [] + reading: list[bool] = [] + sync, get_items = Host.sync, _stores.KeyPlanStore.get_items + + async def counted(self: Any) -> None: + syncs.extend(reading) + await sync(self) + + async def read(self: Any) -> Any: + reading.append(True) + try: + return await get_items(self) + finally: + reading.clear() + + monkeypatch.setattr(Host, "sync", counted) + monkeypatch.setattr(_stores.KeyPlanStore, "get_items", read) + answer = await chat("read it") + assert "Pack" in answer and "Travel" in answer + # No commit while the plan was read, and no journal row for the read. + assert syncs == [] + assert await invoke(chat.obj, "journal", chat.host(), request(), None) == journaled From 0ddea6da19017e33e92690f7cffe767685502bb7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 10:43:49 +0000 Subject: [PATCH 6/8] Agents: a key's first turn applies both its schemas in one commit A turn reads the tool call records as it ends (`settle`, from `answering`), whether or not a tool ran, so a key's first turn applied the calls table's schema in a commit of its own after the conversation's. `restate.sessions.internal.ensure_schemas` applies several schemas, those missing in one transaction (the sessions layer already did so for a list), and `answer_once` ensures the conversation's and the calls' with it. The call records still ensure their own schema, for callers without a conversation: within an attempt that already applied it, that costs nothing. Co-Authored-By: Claude Fable 5.2 --- docs/agents/conversation.md | 3 +++ python/restate/agents/_conversation.py | 7 +++++-- python/restate/sessions/_schema.py | 18 ++++++++++++++++-- python/restate/sessions/internal.py | 6 ++++-- tests/agents_conversation.py | 23 +++++++++++++++++++++++ 5 files changed, 51 insertions(+), 6 deletions(-) diff --git a/docs/agents/conversation.md b/docs/agents/conversation.md index a735e423..c6205659 100644 --- a/docs/agents/conversation.md +++ b/docs/agents/conversation.md @@ -551,6 +551,9 @@ invocation (the next invocation's first call replaces it): discards, so a key never written reads as empty tables. The names are `restate_agents.turns`, `restate-agents.calls`, `restate_agents.langgraph`, `restate_agents.pydantic_ai` and `restate_agents.pydantic_ai.memory`. + A turn (`answer_once`) applies the conversation's and the tool calls' + together (`ensure_schemas`), as it reads the call records as it ends, + whether or not a tool ran: a key's first turn pays one schema commit. - **Migrations.** `restate.sessions`' rule is: never change a migration that has run; append one. This release is the exception (D1 in [What we learned](learnings.md#decisions)): its schemas are written for new diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index 04a2d22e..a80d99bd 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -96,6 +96,7 @@ from typing import Any, Literal, ParamSpec, Protocol, TypeVar from restate.agents._api import Steer +from restate.agents._calls import SCHEMA as CALLS_SCHEMA from restate.agents._calls import outcomes from restate.agents._common import ( BUDGET, @@ -116,7 +117,7 @@ from restate.agents._turn import AttemptEnded, answering, settings_of, stepped, stop_signal_of from restate.exceptions import TerminalError from restate.sessions import Database, DatabaseReader, DatabaseWriter, SessionContext, SqlParam, Transaction -from restate.sessions.internal import Schema, definition_of +from restate.sessions.internal import Schema, definition_of, ensure_schemas _logger = logging.getLogger(__name__) @@ -1661,7 +1662,9 @@ async def answer_once( window = check_window(settings.window if window is None else window) budget = settings.budget if budget is None else budget prompt_id = next_prompt_id(ctx) - await SCHEMA.ensure(ctx.db) + # With the tool calls' (``_calls.py``), which the turn reads as it ends, whether or not a tool ran: a key's + # first turn applies both in one commit. + await ensure_schemas(ctx.db, [SCHEMA, CALLS_SCHEMA]) if (found := await answer_in(ctx.db, prompt_id)) is not None: # Answered (or closed) already: a retry answers, or fails, as the turn ended. answer, status = found diff --git a/python/restate/sessions/_schema.py b/python/restate/sessions/_schema.py index 6286fd68..6b917f65 100644 --- a/python/restate/sessions/_schema.py +++ b/python/restate/sessions/_schema.py @@ -97,8 +97,22 @@ async def ensure(self, db: Any) -> None: Args: db: A session object's ``ctx.db``, or a database from ``database``. """ - ensure = ensure_of(db) - await (ensure([self]) if ensure is not None else apply(db, [self])) + await ensure_schemas(db, [self]) + + +async def ensure_schemas(db: Any, schemas: Sequence[Schema]) -> None: + """Apply the migrations of ``schemas`` that ``db`` has not applied yet, as ``Schema.ensure`` does each. + + Those missing are applied in one transaction: a database that lacks + several schemas (a key's first handler, say) gets them in one commit + rather than one each. + + Args: + db: A session object's ``ctx.db``, or a database from ``database``. + schemas: The schemas. + """ + ensure = ensure_of(db) + await (ensure(schemas) if ensure is not None else apply(db, schemas)) Ensure = Callable[[Sequence[Schema]], Awaitable[None]] diff --git a/python/restate/sessions/internal.py b/python/restate/sessions/internal.py index c6b18dd2..dad60e4e 100644 --- a/python/restate/sessions/internal.py +++ b/python/restate/sessions/internal.py @@ -14,7 +14,8 @@ it. It provides: - ``Schema``: tables that a library owns and creates where it uses them - (``await schema.ensure(ctx.db)``). + (``await schema.ensure(ctx.db)``), and ``ensure_schemas``, which applies + several in one commit. - ``database``: another named database under the same key. - ``reset``: empty the key's database inside a transaction. - ``cancelled`` and ``cancel_calls``: explicit cancellation, for handlers @@ -29,7 +30,7 @@ """ from restate.sessions._object import attach_invocation, cancel_calls, cancelled, definition_of, invoke, start -from restate.sessions._schema import Schema +from restate.sessions._schema import Schema, ensure_schemas from restate.sessions._scope import ScopeKey from restate.sessions._session import ObjectConfig, SqlHost, database, reset @@ -43,6 +44,7 @@ "cancelled", "database", "definition_of", + "ensure_schemas", "invoke", "reset", "start", diff --git a/tests/agents_conversation.py b/tests/agents_conversation.py index 3bb19d0b..7eb6c462 100644 --- a/tests/agents_conversation.py +++ b/tests/agents_conversation.py @@ -1200,3 +1200,26 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: "drop": (CANCELLED, ["", "done", CANCELLED]), } assert (turn["answer"], texts) == expected[write] + + +async def test_a_keys_first_turn_applies_its_schemas_in_one_commit(): + """The conversation's tables and the tool calls' (``_calls.py``), which the turn reads as it ends, are made in + one commit, though no tool runs: a key's first turn pays one commit for its schemas, not two.""" + + async def run(conversation: Conversation) -> str: + await conversation.turn.append_reply({"role": "assistant", "text": "hello"}) + return "hello" + + obj = _turn_rows() + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await answer_once(ctx, text, ITEMS, run) + + backend = Backend() + syncs: list[int] = [] + for text in ("first", "second"): + host = _syncing(backend) + assert await invoke(obj, "ask", host, request(), text) == "hello" + syncs.append(host.syncs) + assert syncs[0] - syncs[1] == 1 From 5de21bf43f55442e39c199582dc496fe52d76428 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 10:44:29 +0000 Subject: [PATCH 7/8] Docs: Pydantic AI's snapshots are _Snapshots', and _Steering is innermost The design of the Pydantic AI integration said `_Restate` runs innermost and stores the snapshots. `_Snapshots` stores them, innermost but for `_Steering`, which it wraps; `_Restate` comes first of the run's capabilities and stores the results, the enqueued messages and the responses. The text and the sequence diagram now say so. Co-Authored-By: Claude Fable 5.2 --- docs/agents/frameworks.md | 35 +++++++++++++++++++++++------------ 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index 0057273f..a59d0266 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -176,33 +176,44 @@ and it leans on the harness wherever a capability does what our code would ### The design **History from snapshots.** The harness's compaction replaces our window, -budget and summaries. When compaction rewrote the history, `_Restate`, which -runs innermost, stores the rewritten list as a snapshot of the turn before the -request goes out; the next turn starts from it, and its input only grows until -the next compaction. An `Agent` of your own with no compaction capability gets -the same default (Q3). +budget and summaries. `_Restate` comes first of the run's capabilities (it +has no position of its own): before each request it stores the run's new tool +results, and after it the response. When compaction rewrote the +history, `_Snapshots`, which runs after every other capability's +`before_model_request` but `_Steering`'s, stores the rewritten list as a +snapshot of the turn before the request goes out; the next turn starts from +it, and its input only grows until the next compaction. `_Steering`, innermost +inside it, takes the messages steered to the turn last, so that a snapshot +this request stores comes before them. An `Agent` of your own with no +compaction capability gets the same default (Q3). ```mermaid sequenceDiagram autonumber participant R as Run (turn N) + participant G as _Restate (first) participant T as TieredCompaction - participant G as _Restate (innermost) + participant S as _Snapshots + participant St as _Steering (innermost) participant DB as Key's database participant M as Model R->>DB: load: the latest snapshot, then every row after it, whole loop each model request - R->>T: before_model_request + R->>G: before_model_request + G->>DB: store new tool results, by call id + G->>T: before_model_request opt over the target T->>T: clamp, clear old tool results, then summarize if needed end - T->>G: before_model_request (the messages as they will be sent) + T->>S: before_model_request (the messages as they will be sent) opt the history was rewritten, not extended - G->>DB: store a snapshot of the turn: the compacted history, minus system prompt parts + S->>DB: store a snapshot of the turn: the compacted history, minus system prompt parts end - G->>DB: store new tool results, by call id - G->>M: the request - M-->>G: the response + S->>St: before_model_request + St->>DB: store the messages steered to the turn (take_steers), then enqueue them + R->>M: the request, with the enqueued messages + M-->>R: the response + R->>G: after_model_request G->>DB: store the response end Note over DB: turn N+1 loads from this snapshot: its input only grows until the next compaction From 7c37cfba3bbbf099dfea11a9d9b789cd52cc8dc7 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 10 Nov 2026 11:17:11 +0000 Subject: [PATCH 8/8] Agents: a graph's step of tool results is one write; typing fixes for CI A graph of your own copied each ToolMessage of a step in a write of its own: since each write that commits now publishes, a step of N tools made N commits and N publishes, where the built-in agent makes one of each. Consecutive tool messages are now copied in one write. Also: an annotation pyright needed in `new_results`, and ADK's request and response types imported from their own modules in the publishes test. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01RePHkSS8KRKmJ93hcd9asV --- python/restate/agents/_conversation.py | 2 +- python/restate/agents/langgraph/_graph.py | 22 ++++++++++++---- tests/agents_langgraph_graph.py | 31 ++++++++++++++++++++++- tests/agents_publishes.py | 4 ++- 4 files changed, 51 insertions(+), 8 deletions(-) diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index a80d99bd..62b9ca62 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -1299,7 +1299,7 @@ def new_results(self, items: Sequence[Item]) -> list[Item]: handed back. """ running = self.running() - open_ids = {call.id for call in running.open_calls()} if running is not None else set() + open_ids: set[str] = {call.id for call in running.open_calls()} if running is not None else set() parsed = [self.items.read(item) for item in items] latest = {call.id: i for i, read in enumerate(parsed) for call in read.calls} diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index 3c2cec75..f2b21120 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -216,7 +216,8 @@ async def _copied( may not have loaded. If the graph trimmed the prompt away, they are the messages whose ids are not ``known``. ``known`` is the ids the conversation has, which this adds to; the messages the conversation gave - the graph are not copied back. Each copy that commits publishes (``Conversation.hold``). + the graph are not copied back. Tool messages that follow one another (a step's results) are copied in one + write. Each write that commits publishes (``Conversation.hold``). """ at = next((i for i, m in enumerate(messages) if m.id == prompt), None) new = [ @@ -227,12 +228,23 @@ async def _copied( and m.id not in known and not m.id.startswith(_CARRIED) ] + results: list[ToolMessage] = [] + + async def stored() -> None: + # The results of a step in one write: one commit and one publish, however many tools ran. + if results: + await turn.append(*results) + known.update(typing.cast(str, m.id) for m in results) + results.clear() + for message in new: - if isinstance(message, AIMessage): - await turn.append_reply(message) - else: - await turn.append(message) + if isinstance(message, ToolMessage): + results.append(message) + continue + await stored() + await turn.append_reply(message) known.add(typing.cast(str, message.id)) + await stored() def _prompt_id(turn: Turn) -> str: diff --git a/tests/agents_langgraph_graph.py b/tests/agents_langgraph_graph.py index 1a6a9b22..9ac2d9a7 100644 --- a/tests/agents_langgraph_graph.py +++ b/tests/agents_langgraph_graph.py @@ -48,7 +48,7 @@ from restate.agents.internal import CANCELLED, ENDED # noqa: E402 from restate.agents.langgraph import nodes # noqa: E402 from restate.agents.langgraph._errors import terminal_of # noqa: E402 -from restate.agents.langgraph._graph import at_breakpoint # noqa: E402 +from restate.agents.langgraph._graph import _copied, at_breakpoint # noqa: E402 # pyright: ignore[reportPrivateUsage] from restate.exceptions import TerminalError # noqa: E402 @@ -871,6 +871,35 @@ async def test_output_validates_the_last_ai_message(harness: HarnessEnvironment) assert (await chat(harness, "graph_planner", key, f"plan {key}")).json() == {"title": key} +async def test_a_steps_tool_messages_are_copied_in_one_write(): + """The results of tools that ran at once are one write of the conversation: one commit, and one publish.""" + writes: list[list[str]] = [] + + class Recording: + async def append(self, *items: BaseMessage) -> None: + writes.append([f"result {m.id}" for m in items]) + + async def append_reply(self, *items: BaseMessage) -> int: + writes.append([f"reply {m.id}" for m in items]) + return 0 + + calls = [{"name": "lookup", "args": {}, "id": f"c{i}"} for i in range(3)] + messages: list[BaseMessage] = [ + HumanMessage("hi", id="prompt:1"), + AIMessage("", id="a1", tool_calls=calls), + *(ToolMessage("done", id=f"t{i}", tool_call_id=f"c{i}") for i in range(3)), + AIMessage("all three are done", id="a2"), + ] + known: set[str] = set() + await _copied(typing.cast(Any, None), typing.cast(Any, Recording()), known, messages, "prompt:1") + assert writes == [["reply a1"], ["result t0", "result t1", "result t2"], ["reply a2"]] + assert known == {"a1", "t0", "t1", "t2", "a2"} + # Results that end the state (the step's, before the model is asked again) are written too. + writes.clear() + await _copied(typing.cast(Any, None), typing.cast(Any, Recording()), set(), messages[:-1], "prompt:1") + assert writes == [["reply a1"], ["result t0", "result t1", "result t2"]] + + async def test_langgraph_streams_a_step_only_once_its_checkpoint_is_stored(): """What copying messages relies on: with durability="sync", a streamed state is stored, even by a slow saver.""" from langgraph.checkpoint.memory import InMemorySaver # pylint: disable=import-outside-toplevel diff --git a/tests/agents_publishes.py b/tests/agents_publishes.py index 6d639d3b..53d213af 100644 --- a/tests/agents_publishes.py +++ b/tests/agents_publishes.py @@ -191,7 +191,9 @@ def stream_response(self, *args, **kwargs): def adk_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: - from google.adk.models import BaseLlm, LlmRequest, LlmResponse + from google.adk.models import BaseLlm + from google.adk.models.llm_request import LlmRequest + from google.adk.models.llm_response import LlmResponse from google.genai import types from restate.agents import adk