mirror of
https://github.com/agno-agi/agno.git
synced 2026-10-11 22:09:58 +00:00
feat: continue ClaudeAgent runs from a step (#10917)
## Summary Adds step replay to `ClaudeAgent` using the same API as native `Agent.continue_run`. Stacked on #10905 (`feat/harness-agents-phase2-astra`); it needs #10904's database-backed Claude SDK transcript store. - **`ClaudeAgent.continue_run` / `acontinue_run`**: accept `continue_from` (`"end"`, `"last_user"` or a message index), `fork` and `input`, and support `stream=True` and `background=True`. Continuing forks the stored SDK transcript at the chosen message (`fork_session_via_store(..., up_to_message_id=...)`) and resumes the fork, so the model only remembers what happened up to that point. - `fork=True` writes a sibling run with `forked_from_run_id` / `forked_from_message_index`; `fork=False` replaces the source run in place and keeps its fork lineage. - `continue_from="last_user"` without `input` replays the whole turn: it forks just before the prompt (or starts a fresh SDK session for the session's first turn) and re-sends the original input. - A mid-turn continue without `input` sends "Continue from where you left off." - Later turns in the session continue from the replayed branch, because the session's SDK pointer moves to the fork. - **Transcript positions on runs**: the adapter enables the SDK's `replay-user-messages` so the prompt arrives with its transcript uuid. It records the SDK session id and uuid for each top-level prompt, tool call, tool result and final reply in `Message.provider_data["claude_sdk"]`. Tool results are marked as checkpoints, so `GET /agents/{id}/runs/{run_id}/checkpoints` lists each step. - **AgentOS**: no router changes. `/continue` and `/checkpoints` already dispatch on `acontinue_run` / `aget_run_output`, so `ClaudeAgent` moves from 501 to supported. Unsupported options (`regenerate`, `replace_original`, `additional_instructions`, HITL `tools`) return 400. - **`BaseExternalAgent`**: generic plumbing only. `ExternalContinuation` carries the kept transcript into `_arun_non_stream` / `_arun_stream`, a per-run `run_state` dict reaches the adapter, and an `_annotate_run_output` hook (no-op by default) lets an adapter attach per-message data. Other adapters are unaffected: none forwards run kwargs to its framework. - **Resume fallback unchanged for normal runs**: the echoed prompt is not counted as a received message, so #10903's "stored session cannot be resumed, start fresh with history" fallback still works. Continuations never fall back; a failed fork resume errors instead of silently replaying the full history. ## Type of change - [ ] Bug fix - [x] New feature - [ ] Breaking change - [ ] Improvement - [ ] Model update - [ ] Other: --- ## Checklist - [x] Code complies with style guidelines - [x] Ran format/validation scripts (`./scripts/format.sh` and `./scripts/validate.sh`) - [x] Self-review completed - [x] Documentation updated (comments, docstrings) - [x] Examples and guides: Relevant cookbook examples have been included or updated (if applicable) - [x] Tested in clean environment - [x] Tests added/updated (if applicable) ### Duplicate and AI-Generated PR Check - [x] I have searched existing [open pull requests](https://github.com/agno-agi/agno/pulls) and confirmed that no other PR already addresses this issue - [ ] If a similar PR exists, I have explained below why this PR is a better approach - [x] Check if this PR was entirely AI-generated (by Copilot, Claude Code, Cursor, etc.) --- ## Additional Notes Base is `feat/harness-agents-phase2-astra` (#10905), not `main`. Merge after #10903, #10904 and #10905. Limits: - Only runs recorded with transcript storage (PostgresDb, AsyncPostgresDb, SqliteDb, AsyncSqliteDb) can be continued; older runs and other databases raise a clear error. - Files and other side effects after the checkpoint are not rewound. The SDK's file checkpointing does not carry over to forked sessions. - Steps inside subagents (`parent_tool_use_id`) are not checkpoints; they live in subagent transcripts. - `regenerate` / `replace_original` and Codex are follow-ups. Verification: - Confirmed against the real SDK (0.2.95) that streamed message uuids equal stored transcript uuids, and that forks cut at a tool result or before the prompt only remember up to that point. - Live end-to-end on SQLite: fork from a tool result, replay of a first and a later turn, in-place streamed continue, follow-up turn on the branch, and AgentOS `/checkpoints` + `/continue` (non-stream and SSE). `cookbook/frameworks/claude-agent-sdk/continue_from.py` passes with `claude-sonnet-4-6`. - New unit tests against the fake SDK: transcript positions and checkpoints (stream and non-stream), fork from a tool result, in-place continue lineage (sync API), turn replay with and without a parent entry, no fallback on a failed continuation, rejected options, and prompt echo not blocking the resume fallback. Each guard was mutation-checked: breaking lineage, allowing fallback, or counting the echo as received makes a test fail. - `validate.sh`: no ruff or mypy findings in touched files; the remaining mypy errors are pre-existing model/MCP SDK typing issues from local dependency versions. Postgres integration tests were not run locally. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Ashpreet <[email protected]> Co-authored-by: Ashpreet <[email protected]> Co-authored-by: kausmeows <[email protected]> Co-authored-by: Claude Fable 5.1 <[email protected]>
This commit is contained in:
Yash Pratap Solanky
co-authored by
Ashpreet
Ashpreet
kausmeows
Claude Fable 5.1
Ashpreet
Ashpreet
kausmeows
Claude Fable 5.1
parent
c15c55ea7d
commit
3be277189d
@@ -10,16 +10,16 @@ Run: .venvs/demo/bin/python cookbook/05_agent_os/27_filesystem/basic.py
|
||||
Try: Connect Agno OS to http://localhost:7777 and open File System
|
||||
"""
|
||||
|
||||
from os import getenv
|
||||
|
||||
from agno.agent import Agent
|
||||
from agno.db.postgres import PostgresDb
|
||||
from agno.models.openai import OpenAIResponses
|
||||
from agno.os import AgentOS
|
||||
|
||||
from os import getenv
|
||||
db = PostgresDb(
|
||||
id="agentos-knowledge-postgres",
|
||||
db_url=getenv("DATABASE_URL", "postgresql+psycopg://ai:ai@localhost:5532/ai"),
|
||||
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -19,6 +19,18 @@ Use PostgreSQL in production. Transcript storage supports PostgresDb, AsyncPostg
|
||||
|
||||
The `agno_transcripts` table is created on first use. A development database that ran an earlier build of this feature has an older table shape and raises a schema mismatch; drop the table and it is recreated.
|
||||
|
||||
## Continue from a step
|
||||
|
||||
The `continue_from/` folder has one cookbook per scenario: continue a finished run with a new instruction, fork from a tool checkpoint, replay or rewrite the prompt, continue in the background, and the same flow over the AgentOS API. See `continue_from/README.md`.
|
||||
|
||||
```bash
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/01_continue_finished_run.py
|
||||
```
|
||||
|
||||
`ClaudeAgent.continue_run` / `acontinue_run` take the same `continue_from` and `fork` arguments as `Agent.continue_run`, and AgentOS serves them through `/agents/{id}/runs/{run_id}/continue` and `/checkpoints`. Each tool result is a checkpoint. Continuing forks the Claude SDK transcript at that message, so the model only remembers what happened up to it; `continue_from="last_user"` without `input` re-sends the original prompt. A finished run is never rewritten in place: like native agents, its continuation is a new sibling run with `forked_from_run_id`, whatever `fork` is set to. Later turns in the session continue from the newest branch.
|
||||
|
||||
Requires a database with transcript storage, and only runs recorded with it can be continued. Files the agent changed after the checkpoint are not rewound. `regenerate` and HITL `requirements` are not supported.
|
||||
|
||||
## Background runs and cancellation
|
||||
|
||||
`background_cancel.py` serves the agent through AgentOS. Submit runs with `background=true`, poll the run endpoint, or use `stream=true` for indexed SSE. Runs continue after disconnects; the resume endpoint reads the configured event stream. Cancel through the run cancellation endpoint.
|
||||
|
||||
@@ -2,6 +2,27 @@
|
||||
|
||||
## 2026-10-09
|
||||
|
||||
### continue_from.py (finished runs always fork)
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Reran after `continue_run(fork=False)` on a completed run was changed to behave like native agents: the continuation becomes a new sibling run with fork lineage instead of rewriting the source run. Also exercised live with `background=True` on a completed run (ALPHA -> BETA -> GAMMA, three runs in the session, source run untouched).
|
||||
|
||||
**Result:** Checkpoints listed at 3, 5 and 6; branch from step 3 kept only `echo alpha`; replayed turn returned DONE. The in-place-with-background ValueError is gone.
|
||||
|
||||
---
|
||||
|
||||
### continue_from.py (review fixes)
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Reran after the review fixes: checkpoints are exposed per tool batch (one per sequential step, one per parallel batch, anchored at the batch's last transcript entry), `continue_from=0` starts a branch before the prompt, `continue_from="last_user"` replays the selected user turn rather than only the first, cancelled runs are rejected, and in-place continuation refuses `background=True`. The cookbook now asserts on the tool results the branch carries instead of on the model's wording, so a run where Claude batches both commands still passes.
|
||||
|
||||
**Result:** Sequential run: checkpoints at messages 3, 5 and the end; the branch from step 3 carried only the first result and answered with the first command. Replay of the whole turn completed. A separate parallel run exposed a single checkpoint at the batch end.
|
||||
|
||||
---
|
||||
## 2026-10-09
|
||||
|
||||
### transcript_store.py
|
||||
|
||||
**Status:** PASS
|
||||
@@ -116,8 +137,6 @@ checks; the earlier live provider cancellation runs were not repeated.
|
||||
|
||||
---
|
||||
|
||||
---
|
||||
|
||||
### session_store.py (transcript schema revision)
|
||||
|
||||
**Status:** PASS
|
||||
@@ -125,3 +144,13 @@ checks; the earlier live provider cancellation runs were not repeated.
|
||||
**Description:** Reran the two-process verification after scoping transcript rows by framework, project, session and subpath, numbering positions per transcript and recording the owning Agno session. Used claude-agent-sdk 0.2.95 from the demo environment, the `ANTHROPIC_API_KEY` from `.envrc` and empty config directories.
|
||||
|
||||
**Result:** Process A replied `OK`; process B, with its Agno run deleted, replied `cobalt orchard 742`. The PostgreSQL contract test (PostgresDb and AsyncPostgresDb) passed against PostgreSQL 14.
|
||||
|
||||
---
|
||||
|
||||
### continue_from.py (2026-10-08)
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Live run with claude-agent-sdk 0.2.95 and `claude-sonnet-4-6` on SQLite transcript storage. A two-step Bash turn listed checkpoints at both tool results and the end. Continuing from the first tool result with a question produced a branch that only knew `echo alpha`; `continue_from="last_user"` replayed the turn as a forked sibling. Separately, with `claude-haiku-4-5`, exercised AgentOS `/checkpoints` and `/continue` (non-stream and SSE), an in-place streamed continue, and a replay of a non-first turn that kept earlier context.
|
||||
|
||||
**Result:** Continue and checkpoints work end to end against the real SDK. Files touched after a checkpoint are not rewound.
|
||||
|
||||
@@ -0,0 +1,105 @@
|
||||
"""
|
||||
Continue a finished ClaudeAgent run
|
||||
===================================
|
||||
A completed run can be continued with a follow-up instruction. The
|
||||
continuation is a new run in the same session that carries the full Claude
|
||||
conversation, so the model remembers the code it wrote and the decisions it
|
||||
made in the source run.
|
||||
|
||||
Scenario: the first run writes a small Python module and runs it. The
|
||||
continuation asks for a test file for "the function you just wrote" without
|
||||
naming it, which only works if the model still has the source run's context.
|
||||
|
||||
Like native Agno agents, a finished run is never rewritten in place: the
|
||||
continuation gets its own run_id and records where it came from in
|
||||
forked_from_run_id, whatever the fork argument is set to.
|
||||
|
||||
Requirements:
|
||||
pip install claude-agent-sdk
|
||||
A database with transcript storage (SQLite or PostgreSQL).
|
||||
|
||||
Usage:
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/01_continue_finished_run.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
SESSION_ID = "continue-finished-run"
|
||||
|
||||
|
||||
async def main():
|
||||
with tempfile.TemporaryDirectory(prefix="agno-continue-") as workdir:
|
||||
workspace = Path(workdir)
|
||||
agent = ClaudeAgent(
|
||||
id="continue-demo",
|
||||
model="claude-sonnet-4-6",
|
||||
db=SqliteDb(db_file=str(workspace / "runs.db")),
|
||||
cwd=workdir,
|
||||
# The SDK's default system prompt does not name the working directory, so tell
|
||||
# the model where to create files; Bash already runs there.
|
||||
system_prompt=f"Your working directory is {workdir}. Create and run files there.",
|
||||
allowed_tools=["Write", "Read", "Bash"],
|
||||
permission_mode="bypassPermissions",
|
||||
max_turns=8,
|
||||
max_budget_usd=0.5,
|
||||
)
|
||||
|
||||
# 1. The source run: write a module and run it.
|
||||
source = await agent.arun(
|
||||
"Create slugify.py with a function slugify(text) that lowercases the text, replaces "
|
||||
"every run of non-alphanumeric characters with a single hyphen, and strips leading and "
|
||||
"trailing hyphens. Then run it with python3 to print slugify('Hello, World!') and reply "
|
||||
"with just the printed output.",
|
||||
session_id=SESSION_ID,
|
||||
)
|
||||
assert source.status == RunStatus.completed, source.content
|
||||
assert (workspace / "slugify.py").exists(), (
|
||||
"the source run should have written slugify.py"
|
||||
)
|
||||
print(f"Source run {source.run_id[:8]}: {source.content}")
|
||||
|
||||
# 2. Continue it. continue_from defaults to "end": the follow-up is appended after
|
||||
# the last message of the source run. The prompt never names the function or the
|
||||
# file, so the model has to remember them from the source run.
|
||||
continued = await agent.acontinue_run(
|
||||
run_id=source.run_id,
|
||||
session_id=SESSION_ID,
|
||||
input="Write test_slugify.py with three plain assert statements covering the function "
|
||||
"you just wrote, including the example you ran, and a final print('ok'). Run it with "
|
||||
"python3 and reply with just what it printed.",
|
||||
)
|
||||
assert continued.status == RunStatus.completed, continued.content
|
||||
assert (workspace / "test_slugify.py").exists(), (
|
||||
"the continuation should have written the test file"
|
||||
)
|
||||
print(f"Continued run {continued.run_id[:8]}: {continued.content}")
|
||||
print(
|
||||
f"Workspace: {sorted(p.name for p in workspace.iterdir() if p.suffix == '.py')}"
|
||||
)
|
||||
|
||||
# The continuation is a new run with lineage back to the source.
|
||||
assert continued.run_id != source.run_id
|
||||
assert continued.forked_from_run_id == source.run_id
|
||||
assert continued.forked_from_message_index == len(source.messages or [])
|
||||
|
||||
# The source run is untouched and the session holds both runs.
|
||||
stored = await agent.aget_run_output(source.run_id, SESSION_ID)
|
||||
assert stored is not None and stored.content == source.content
|
||||
session = await agent.aget_session(SESSION_ID)
|
||||
print(f"Runs in session: {len(session.runs or [])}")
|
||||
for run in session.runs or []:
|
||||
lineage = (run.forked_from_run_id or "")[:8] or "-"
|
||||
tools = sum(1 for m in run.messages or [] if m.role == "tool")
|
||||
print(
|
||||
f" {run.run_id[:8]} forked_from={lineage} tool_calls={tools} content={run.content!r}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
Fork a ClaudeAgent run from a checkpoint
|
||||
========================================
|
||||
Every tool result in a stored run is a checkpoint. Forking from one creates
|
||||
a sibling run whose Claude conversation stops at that point: the model
|
||||
remembers what happened up to the checkpoint and nothing after it.
|
||||
|
||||
Scenario: the source run analyses a CSV file in three Bash steps (count the
|
||||
rows, sum a column, write a summary file). The fork branches after the first
|
||||
step and asks the model, without tools, what it knows about the file. The
|
||||
branch knows the row count and does not know the total, because that result
|
||||
is not part of its conversation.
|
||||
|
||||
Files are not rewound: summary.txt written in step three stays on disk.
|
||||
|
||||
Requirements:
|
||||
pip install claude-agent-sdk
|
||||
A database with transcript storage (SQLite or PostgreSQL).
|
||||
|
||||
Usage:
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/02_fork_from_checkpoint.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.os.checkpoints import list_run_checkpoints
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
SESSION_ID = "fork-from-checkpoint"
|
||||
|
||||
SALES_CSV = """region,amount
|
||||
north,120
|
||||
south,80
|
||||
east,200
|
||||
west,45
|
||||
north,55
|
||||
"""
|
||||
|
||||
|
||||
async def main():
|
||||
with tempfile.TemporaryDirectory(prefix="agno-continue-") as workdir:
|
||||
workspace = Path(workdir)
|
||||
(workspace / "sales.csv").write_text(SALES_CSV)
|
||||
agent = ClaudeAgent(
|
||||
id="continue-demo",
|
||||
model="claude-sonnet-4-6",
|
||||
db=SqliteDb(db_file=str(workspace / "runs.db")),
|
||||
cwd=workdir,
|
||||
system_prompt=f"Your working directory is {workdir}. Use Bash for every step.",
|
||||
allowed_tools=["Bash"],
|
||||
permission_mode="bypassPermissions",
|
||||
max_turns=8,
|
||||
max_budget_usd=0.5,
|
||||
)
|
||||
|
||||
# 1. A three-step tool run. One Bash call per step gives three checkpoints.
|
||||
run = await agent.arun(
|
||||
"sales.csv has a header row and an amount column. Do these steps in order, one Bash "
|
||||
"call each, waiting for each result before the next: "
|
||||
"(1) count the data rows with `tail -n +2 sales.csv | wc -l`; "
|
||||
"(2) sum the amount column with awk; "
|
||||
"(3) write both numbers to summary.txt. "
|
||||
"Then reply with just: rows=<n> total=<sum>.",
|
||||
session_id=SESSION_ID,
|
||||
)
|
||||
assert run.status == RunStatus.completed, run.content
|
||||
assert (workspace / "summary.txt").exists(), (
|
||||
"step three should have written summary.txt"
|
||||
)
|
||||
print(f"Source run: {run.content}")
|
||||
|
||||
# 2. List the checkpoints: one per tool step and one at the end.
|
||||
# message_index is the value to pass as continue_from.
|
||||
checkpoints = list_run_checkpoints(run)
|
||||
print("Checkpoints:")
|
||||
for checkpoint in checkpoints:
|
||||
print(
|
||||
f" {checkpoint['message_index']}: {checkpoint['reason']} "
|
||||
f"({checkpoint['message_role']}) {checkpoint['message_preview']}"
|
||||
)
|
||||
|
||||
# 3. Fork from the first checkpoint. With one call per step that is the row count.
|
||||
# If Claude batched calls in parallel, the batch is one checkpoint and the branch
|
||||
# keeps every result in it.
|
||||
first_step = checkpoints[0]["message_index"]
|
||||
kept_results = [
|
||||
m for m in (run.messages or [])[:first_step] if m.role == "tool"
|
||||
]
|
||||
branch = await agent.acontinue_run(
|
||||
run_id=run.run_id,
|
||||
session_id=SESSION_ID,
|
||||
continue_from=first_step,
|
||||
fork=True,
|
||||
input="Do not run any tools. From this conversation only: how many data rows does "
|
||||
"sales.csv have, and what is the total of the amount column? If you have not "
|
||||
"computed one of them yet, say exactly 'not computed yet' for it.",
|
||||
)
|
||||
assert branch.status == RunStatus.completed, branch.content
|
||||
print(f"Branch from step {first_step}: {branch.content}")
|
||||
|
||||
# The branch carries exactly the tool results up to the checkpoint and no new ones.
|
||||
branch_results = [m for m in (branch.messages or []) if m.role == "tool"]
|
||||
assert [m.tool_call_id for m in branch_results] == [
|
||||
m.tool_call_id for m in kept_results
|
||||
]
|
||||
assert branch.forked_from_run_id == run.run_id
|
||||
assert branch.forked_from_message_index == first_step
|
||||
|
||||
# The source run keeps all of its steps, and the file from step three is still there.
|
||||
stored = await agent.aget_run_output(run.run_id, SESSION_ID)
|
||||
assert stored is not None and len(stored.messages or []) == len(
|
||||
run.messages or []
|
||||
)
|
||||
print(f"Source run still has {len(stored.messages or [])} messages")
|
||||
print(
|
||||
f"summary.txt on disk: {(workspace / 'summary.txt').read_text().strip()!r}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,106 @@
|
||||
"""
|
||||
Replay or rewrite a ClaudeAgent turn
|
||||
====================================
|
||||
Two ways to start over from a run without losing it:
|
||||
|
||||
- continue_from="last_user" with no input re-sends the original prompt. The
|
||||
model sees the same conversation as the first time and runs its tools again
|
||||
against the workspace as it is now.
|
||||
- continue_from=0 with input starts a branch before the prompt, so the same
|
||||
run lineage gets a different prompt.
|
||||
|
||||
Scenario: the source run lists the markdown files in a folder. A second file
|
||||
is added, then the turn is replayed: the replay's tool call sees both files.
|
||||
The rewrite asks a different question about the same folder.
|
||||
|
||||
Both create sibling runs with forked_from_run_id set; the source run is kept.
|
||||
|
||||
Requirements:
|
||||
pip install claude-agent-sdk
|
||||
A database with transcript storage (SQLite or PostgreSQL).
|
||||
|
||||
Usage:
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/03_replay_user_turn.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
SESSION_ID = "replay-user-turn"
|
||||
|
||||
|
||||
def tool_outputs(run) -> str:
|
||||
return "\n".join(str(m.content) for m in run.messages or [] if m.role == "tool")
|
||||
|
||||
|
||||
async def main():
|
||||
with tempfile.TemporaryDirectory(prefix="agno-continue-") as workdir:
|
||||
notes = Path(workdir) / "notes"
|
||||
notes.mkdir()
|
||||
(notes / "roadmap.md").write_text(
|
||||
"# Roadmap\n\n- ship continue-from\n- ship sandboxes\n"
|
||||
)
|
||||
agent = ClaudeAgent(
|
||||
id="continue-demo",
|
||||
model="claude-sonnet-4-6",
|
||||
db=SqliteDb(db_file=str(Path(workdir) / "runs.db")),
|
||||
cwd=workdir,
|
||||
allowed_tools=["Bash"],
|
||||
permission_mode="bypassPermissions",
|
||||
max_turns=4,
|
||||
max_budget_usd=0.5,
|
||||
)
|
||||
|
||||
run = await agent.arun(
|
||||
"List the markdown files in notes/ with `ls notes` and reply with just the file names, "
|
||||
"comma separated.",
|
||||
session_id=SESSION_ID,
|
||||
)
|
||||
assert run.status == RunStatus.completed, run.content
|
||||
assert "roadmap.md" in tool_outputs(run) and "retro.md" not in tool_outputs(run)
|
||||
print(f"Original answer: {run.content}")
|
||||
|
||||
# The workspace changes after the first run.
|
||||
(notes / "retro.md").write_text("# Retro\n\n- what went well\n")
|
||||
|
||||
# 1. Replay the same prompt. No input is given, so the stored user message is
|
||||
# re-sent and the ls runs again, now seeing the new file.
|
||||
replay = await agent.acontinue_run(
|
||||
run_id=run.run_id,
|
||||
session_id=SESSION_ID,
|
||||
continue_from="last_user",
|
||||
fork=True,
|
||||
)
|
||||
assert replay.status == RunStatus.completed, replay.content
|
||||
assert replay.forked_from_run_id == run.run_id
|
||||
assert "retro.md" in tool_outputs(replay), (
|
||||
"the replayed tool call should see the new file"
|
||||
)
|
||||
print(f"Replayed answer: {replay.content}")
|
||||
|
||||
# 2. Rewrite the prompt. continue_from=0 branches before the first message, so the
|
||||
# new input replaces the original prompt in this branch.
|
||||
rewrite = await agent.acontinue_run(
|
||||
run_id=run.run_id,
|
||||
session_id=SESSION_ID,
|
||||
continue_from=0,
|
||||
fork=True,
|
||||
input="Count the bullet points across all files in notes/ with "
|
||||
"`grep -c '^- ' notes/*.md` and reply with the per-file counts.",
|
||||
)
|
||||
assert rewrite.status == RunStatus.completed, rewrite.content
|
||||
assert rewrite.forked_from_run_id == run.run_id
|
||||
assert rewrite.forked_from_message_index == 0
|
||||
print(f"Rewritten answer: {rewrite.content}")
|
||||
|
||||
session = await agent.aget_session(SESSION_ID)
|
||||
print(f"Runs in session: {len(session.runs or [])}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,132 @@
|
||||
"""
|
||||
Continue a ClaudeAgent run in the background, streamed
|
||||
======================================================
|
||||
A continuation can run in the background like any other run. With
|
||||
background=True and stream=True the call returns an indexed event stream
|
||||
that survives client disconnects: the run keeps going on the server, and a
|
||||
client can reconnect to the same stream later through the AgentOS resume
|
||||
endpoint. This is what the AgentOS UI uses for long-running agent turns.
|
||||
|
||||
Scenario: the source run writes a script but does not run it. The
|
||||
continuation is submitted in the background and streamed: tool events and
|
||||
content arrive as they happen, the final RunOutput is yielded last, and the
|
||||
same result is then read back from the database.
|
||||
|
||||
A finished run is always continued as a new run, so the background job gets
|
||||
a fresh run_id with forked_from_run_id pointing at the source.
|
||||
|
||||
Requirements:
|
||||
pip install claude-agent-sdk
|
||||
A database with transcript storage (SQLite or PostgreSQL).
|
||||
|
||||
Usage:
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/04_background_continue.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.run.agent import RunOutput
|
||||
from agno.run.background import await_background_runs
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
SESSION_ID = "background-continue"
|
||||
|
||||
|
||||
def sse_events(chunk: str):
|
||||
"""Background streams yield SSE text; each data line is one JSON event."""
|
||||
for line in chunk.splitlines():
|
||||
if line.startswith("data: "):
|
||||
yield json.loads(line[6:])
|
||||
|
||||
|
||||
async def main():
|
||||
with tempfile.TemporaryDirectory(prefix="agno-continue-") as workdir:
|
||||
workspace = Path(workdir)
|
||||
agent = ClaudeAgent(
|
||||
id="continue-demo",
|
||||
model="claude-sonnet-4-6",
|
||||
db=SqliteDb(db_file=str(workspace / "runs.db")),
|
||||
cwd=workdir,
|
||||
# The SDK's default system prompt does not name the working directory, so tell
|
||||
# the model where to create files; Bash already runs there.
|
||||
system_prompt=f"Your working directory is {workdir}. Create and run files there.",
|
||||
allowed_tools=["Write", "Bash"],
|
||||
permission_mode="bypassPermissions",
|
||||
max_turns=8,
|
||||
max_budget_usd=0.5,
|
||||
)
|
||||
|
||||
source = await agent.arun(
|
||||
"Write fizzbuzz.py that prints FizzBuzz for 1 through 15, one value per line. "
|
||||
"Do not run it. Reply with just the file name.",
|
||||
session_id=SESSION_ID,
|
||||
)
|
||||
assert source.status == RunStatus.completed, source.content
|
||||
assert (workspace / "fizzbuzz.py").exists()
|
||||
print(f"Source run {source.run_id[:8]}: {source.content}")
|
||||
|
||||
# Submit the continuation in the background and follow its event stream.
|
||||
# yield_run_output=True makes the final RunOutput the last item.
|
||||
stream = agent.acontinue_run(
|
||||
run_id=source.run_id,
|
||||
session_id=SESSION_ID,
|
||||
input="Now run the script you wrote with python3 and reply with just its last line.",
|
||||
background=True,
|
||||
stream=True,
|
||||
yield_run_output=True,
|
||||
)
|
||||
print("Streaming:")
|
||||
result = None
|
||||
seen = []
|
||||
async for item in stream:
|
||||
if isinstance(item, RunOutput):
|
||||
result = item
|
||||
continue
|
||||
for event in sse_events(item):
|
||||
kind = event.get("event")
|
||||
seen.append(kind)
|
||||
if kind == "RunStarted":
|
||||
print(f" {kind} run_id={event.get('run_id', '')[:8]}")
|
||||
elif kind == "ToolCallStarted":
|
||||
tool = event.get("tool") or {}
|
||||
print(
|
||||
f" {kind} {tool.get('tool_name')} {json.dumps(tool.get('tool_args'))[:80]}"
|
||||
)
|
||||
elif kind == "ToolCallCompleted":
|
||||
tool = event.get("tool") or {}
|
||||
print(f" {kind} {str(tool.get('result'))[:60]!r}")
|
||||
elif kind == "RunContent" and event.get("content"):
|
||||
print(f" {kind} {event['content']!r}")
|
||||
elif kind in ("RunCompleted", "RunError", "RunCancelled"):
|
||||
print(f" {kind}")
|
||||
|
||||
assert result is not None, "the stream should end with the RunOutput"
|
||||
assert result.status == RunStatus.completed, result.content
|
||||
assert result.run_id != source.run_id
|
||||
assert result.forked_from_run_id == source.run_id
|
||||
assert "RunCompleted" in seen, seen
|
||||
print(f"Finished run {result.run_id[:8]}: {result.content!r}")
|
||||
|
||||
# The same result is in the database for any client that reconnects later.
|
||||
stored = await agent.aget_run_output(result.run_id, SESSION_ID)
|
||||
assert stored is not None and stored.status == RunStatus.completed
|
||||
outputs = [str(m.content) for m in stored.messages or [] if m.role == "tool"]
|
||||
assert any("FizzBuzz" in output for output in outputs), (
|
||||
"the script should have been run"
|
||||
)
|
||||
print(
|
||||
f"Stored run {stored.run_id[:8]}: {stored.status} with {len(outputs)} tool results"
|
||||
)
|
||||
|
||||
# The stream ends before the run marks its event stream complete. Let that finish
|
||||
# before the loop closes; a long-lived server does not need this.
|
||||
await await_background_runs()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
Continue and checkpoints over the AgentOS API
|
||||
=============================================
|
||||
Serve a ClaudeAgent through AgentOS and use the same continue endpoints as
|
||||
native agents:
|
||||
|
||||
GET /agents/{agent_id}/runs/{run_id}/checkpoints?session_id=...
|
||||
POST /agents/{agent_id}/runs/{run_id}/continue
|
||||
|
||||
Run normally to serve AgentOS; pass --verify for a live end-to-end check
|
||||
through the HTTP API: a three-step analysis of a CSV file, its checkpoint
|
||||
list, a fork after the first step, and a streamed follow-up on the finished
|
||||
run.
|
||||
|
||||
Requirements:
|
||||
pip install claude-agent-sdk
|
||||
|
||||
Usage:
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/05_agentos_api.py
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/05_agentos_api.py --verify
|
||||
|
||||
Then call the API (the server runs in the current directory, so give it a CSV to look at):
|
||||
# Start a run
|
||||
curl -X POST http://localhost:7777/agents/continue-demo/runs \
|
||||
-F "message=Count the data rows in sales.csv with Bash, then sum its amount column with awk in a second call, then reply rows=<n> total=<sum>." \
|
||||
-F "session_id=demo" -F "stream=false"
|
||||
|
||||
# List its checkpoints
|
||||
curl "http://localhost:7777/agents/continue-demo/runs/<run_id>/checkpoints?session_id=demo"
|
||||
|
||||
# Fork from the first checkpoint
|
||||
curl -X POST http://localhost:7777/agents/continue-demo/runs/<run_id>/continue \
|
||||
-F "session_id=demo" -F "continue_from=3" -F "fork=true" -F "stream=false" \
|
||||
-F "input=Without tools: what do you know about sales.csv so far?"
|
||||
|
||||
# Continue the finished run with a follow-up (streamed)
|
||||
curl -N -X POST http://localhost:7777/agents/continue-demo/runs/<run_id>/continue \
|
||||
-F "session_id=demo" -F "continue_from=end" -F "input=Which region had the largest single amount?"
|
||||
"""
|
||||
|
||||
import json
|
||||
import sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.os import AgentOS
|
||||
|
||||
AGENT_ID = "continue-demo"
|
||||
|
||||
SALES_CSV = """region,amount
|
||||
north,120
|
||||
south,80
|
||||
east,200
|
||||
west,45
|
||||
north,55
|
||||
"""
|
||||
|
||||
db = SqliteDb(db_file="tmp/claude-agent-sdk-continue.db")
|
||||
agent = ClaudeAgent(
|
||||
id=AGENT_ID,
|
||||
name="Continue Demo",
|
||||
model="claude-sonnet-4-6",
|
||||
db=db,
|
||||
allowed_tools=["Bash"],
|
||||
permission_mode="bypassPermissions",
|
||||
max_turns=8,
|
||||
max_budget_usd=0.5,
|
||||
)
|
||||
agent_os = AgentOS(agents=[agent], db=db)
|
||||
app = agent_os.get_app()
|
||||
|
||||
|
||||
def verify():
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix="agno-continue-") as workdir:
|
||||
(Path(workdir) / "sales.csv").write_text(SALES_CSV)
|
||||
agent.db = SqliteDb(db_file=str(Path(workdir) / "runs.db"))
|
||||
agent.cwd = workdir
|
||||
agent.system_prompt = (
|
||||
f"Your working directory is {workdir}. Use Bash for every step."
|
||||
)
|
||||
session_id = "agentos-api"
|
||||
with TestClient(AgentOS(agents=[agent], db=agent.db).get_app()) as client:
|
||||
# 1. Start a three-step tool run.
|
||||
response = client.post(
|
||||
f"/agents/{AGENT_ID}/runs",
|
||||
data={
|
||||
"message": "sales.csv has a header row and an amount column. Do these steps in "
|
||||
"order, one Bash call each, waiting for each result before the next: "
|
||||
"(1) count the data rows with `tail -n +2 sales.csv | wc -l`; "
|
||||
"(2) sum the amount column with awk; "
|
||||
"(3) print the largest amount and its region with sort. "
|
||||
"Then reply with just: rows=<n> total=<sum> max=<region>.",
|
||||
"session_id": session_id,
|
||||
"stream": "false",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
run = response.json()
|
||||
print(f"Run {run['run_id'][:8]}: {run['status']} {run['content']!r}")
|
||||
|
||||
# 2. List checkpoints.
|
||||
response = client.get(
|
||||
f"/agents/{AGENT_ID}/runs/{run['run_id']}/checkpoints",
|
||||
params={"session_id": session_id},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
checkpoints = response.json()["checkpoints"]
|
||||
print("Checkpoints:")
|
||||
for checkpoint in checkpoints:
|
||||
print(
|
||||
f" {checkpoint['message_index']}: {checkpoint['reason']} {checkpoint['message_preview']!r}"
|
||||
)
|
||||
|
||||
# 3. Fork from the first checkpoint and ask what the branch knows.
|
||||
first_step = checkpoints[0]["message_index"]
|
||||
response = client.post(
|
||||
f"/agents/{AGENT_ID}/runs/{run['run_id']}/continue",
|
||||
data={
|
||||
"session_id": session_id,
|
||||
"continue_from": str(first_step),
|
||||
"fork": "true",
|
||||
"stream": "false",
|
||||
"input": "Do not run any tools. From this conversation only: how many data rows does "
|
||||
"sales.csv have, and what is the total of the amount column? If you have not "
|
||||
"computed one of them yet, say exactly 'not computed yet' for it.",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
branch = response.json()
|
||||
assert branch["forked_from_run_id"] == run["run_id"]
|
||||
branch_tools = [
|
||||
m for m in branch.get("messages", []) if m.get("role") == "tool"
|
||||
]
|
||||
kept_tools = [
|
||||
m for m in run["messages"][:first_step] if m.get("role") == "tool"
|
||||
]
|
||||
assert len(branch_tools) == len(kept_tools), (
|
||||
"the branch should carry only the kept tool results"
|
||||
)
|
||||
print(f"Branch from step {first_step}: {branch['content']!r}")
|
||||
|
||||
# 4. Continue the finished run from its end, streamed.
|
||||
response = client.post(
|
||||
f"/agents/{AGENT_ID}/runs/{run['run_id']}/continue",
|
||||
data={
|
||||
"session_id": session_id,
|
||||
"continue_from": "end",
|
||||
"stream": "true",
|
||||
"input": "Without running tools: which region had the largest single amount? "
|
||||
"Reply with just the region.",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
events = [
|
||||
json.loads(line[6:])
|
||||
for line in response.text.splitlines()
|
||||
if line.startswith("data: ")
|
||||
]
|
||||
kinds = [event.get("event") for event in events]
|
||||
final = next(
|
||||
(e for e in reversed(events) if e.get("event") == "RunCompleted"), None
|
||||
)
|
||||
assert final is not None, kinds
|
||||
print(
|
||||
f"Streamed follow-up: {len(events)} events, answer={final.get('content')!r}"
|
||||
)
|
||||
|
||||
# 5. The session holds the source run and two sibling runs.
|
||||
response = client.get(
|
||||
f"/sessions/{session_id}/runs", params={"type": "agent"}
|
||||
)
|
||||
runs = response.json()
|
||||
print("Session runs:")
|
||||
for item in runs:
|
||||
lineage = (item.get("forked_from_run_id") or "")[:8] or "-"
|
||||
print(f" {item['run_id'][:8]} {item['status']} forked_from={lineage}")
|
||||
assert len(runs) == 3
|
||||
print("PASS: checkpoints, fork and continue work over the AgentOS API")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if "--verify" in sys.argv:
|
||||
verify()
|
||||
else:
|
||||
agent_os.serve(app="05_agentos_api:app", reload=False)
|
||||
@@ -0,0 +1,107 @@
|
||||
"""
|
||||
Recover after a tool succeeds but the Claude run fails
|
||||
=====================================================
|
||||
Record an invoice in a local audit log, then deliberately exhaust max_turns.
|
||||
Continue from the failed run's end using its stored tool result. The invoice
|
||||
must remain recorded exactly once: recovery must not repeat the write.
|
||||
|
||||
Use --stream to exercise the streaming failure path as well.
|
||||
|
||||
Requirements:
|
||||
pip install claude-agent-sdk
|
||||
Claude credentials configured for the SDK.
|
||||
|
||||
Usage:
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/06_recover_failed_tool_run.py
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/06_recover_failed_tool_run.py --stream
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.os.checkpoints import list_run_checkpoints
|
||||
from agno.run.agent import RunOutput
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
SESSION_ID = "recover-failed-tool-run"
|
||||
|
||||
|
||||
async def main(stream: bool = False):
|
||||
with tempfile.TemporaryDirectory(prefix="agno-recovery-") as workdir:
|
||||
workspace = Path(workdir)
|
||||
(workspace / "record_invoice.py").write_text(
|
||||
"from pathlib import Path\n"
|
||||
"audit = Path('audit.log')\n"
|
||||
"with audit.open('a') as handle:\n"
|
||||
" handle.write('invoice-1042 recorded\\n')\n"
|
||||
"print('Recorded invoice-1042; audit entries:', len(audit.read_text().splitlines()))\n"
|
||||
)
|
||||
agent = ClaudeAgent(
|
||||
id="failed-run-demo",
|
||||
model="claude-sonnet-4-6",
|
||||
db=SqliteDb(db_file=str(workspace / "runs.db")),
|
||||
cwd=workdir,
|
||||
system_prompt=f"Your working directory is {workdir}. Follow the requested command exactly.",
|
||||
allowed_tools=["Bash"],
|
||||
options_kwargs={"tools": ["Bash"]},
|
||||
permission_mode="bypassPermissions",
|
||||
max_turns=1,
|
||||
max_budget_usd=0.5,
|
||||
)
|
||||
prompt = "Run `python3 record_invoice.py` with Bash exactly once, then summarize its output."
|
||||
if stream:
|
||||
failed = None
|
||||
async for event in agent.arun(
|
||||
prompt, session_id=SESSION_ID, stream=True, yield_run_output=True
|
||||
):
|
||||
if isinstance(event, RunOutput):
|
||||
failed = event
|
||||
else:
|
||||
print(f"Event: {event.event}")
|
||||
assert failed is not None
|
||||
else:
|
||||
failed = await agent.arun(prompt, session_id=SESSION_ID)
|
||||
|
||||
assert failed.status == RunStatus.error, failed.content
|
||||
assert "max_turns" in str(failed.content), failed.content
|
||||
audit = workspace / "audit.log"
|
||||
assert audit.read_text().splitlines() == ["invoice-1042 recorded"]
|
||||
assert failed.tools and len(failed.tools) == 1
|
||||
assert "audit entries: 1" in str(failed.tools[0].result)
|
||||
assert failed.messages[-1].role == "tool"
|
||||
print(f"Failed run: {failed.run_id}")
|
||||
print(f"Error: {failed.content}")
|
||||
print(f"Preserved tool result: {failed.tools[0].result}")
|
||||
print("Recovery boundaries:", list_run_checkpoints(failed))
|
||||
|
||||
# Increase the turn limit and resume after the completed tool result.
|
||||
# fork=True preserves the failed source run for inspection.
|
||||
agent.max_turns = 3
|
||||
recovered = await agent.acontinue_run(
|
||||
run_id=failed.run_id,
|
||||
session_id=SESSION_ID,
|
||||
continue_from="end",
|
||||
fork=True,
|
||||
input="Do not run any tools or repeat the invoice write. From the tool result already "
|
||||
"in this conversation, report which invoice was recorded and the audit entry count.",
|
||||
)
|
||||
assert recovered.status == RunStatus.completed, recovered.content
|
||||
assert recovered.forked_from_run_id == failed.run_id
|
||||
assert [tool.tool_call_id for tool in recovered.tools or []] == [
|
||||
tool.tool_call_id for tool in failed.tools
|
||||
], "recovery must retain the completed call without executing another tool"
|
||||
assert audit.read_text().splitlines() == ["invoice-1042 recorded"]
|
||||
print(f"Recovered answer: {recovered.content}")
|
||||
print(
|
||||
"PASS: recovered from the stored result; the invoice was written only once"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--stream", action="store_true")
|
||||
asyncio.run(main(stream=parser.parse_args().stream))
|
||||
@@ -0,0 +1,121 @@
|
||||
"""
|
||||
Inspect checkpoints when Claude delegates to a subagent
|
||||
=======================================================
|
||||
A sales analyst subagent reads a CSV using Bash. Its internal tool results
|
||||
remain visible in the run, but only the top-level Agent result is a replay
|
||||
checkpoint. Fork from that result without re-running the analysis.
|
||||
|
||||
Requirements:
|
||||
pip install claude-agent-sdk
|
||||
Claude credentials configured for the SDK.
|
||||
|
||||
Usage:
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/07_subagent_checkpoints.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.os.checkpoints import list_run_checkpoints
|
||||
from agno.run.agent import RunOutput
|
||||
from agno.run.base import RunStatus
|
||||
from claude_agent_sdk import AgentDefinition
|
||||
|
||||
SESSION_ID = "subagent-checkpoints"
|
||||
|
||||
|
||||
async def main():
|
||||
with tempfile.TemporaryDirectory(prefix="agno-subagent-") as workdir:
|
||||
workspace = Path(workdir)
|
||||
(workspace / "sales.csv").write_text(
|
||||
"region,amount\nnorth,120\nsouth,80\neast,200\n"
|
||||
)
|
||||
agent = ClaudeAgent(
|
||||
id="subagent-checkpoint-demo",
|
||||
model="claude-sonnet-4-6",
|
||||
db=SqliteDb(db_file=str(workspace / "runs.db")),
|
||||
cwd=workdir,
|
||||
system_prompt=f"Your working directory is {workdir}. Delegate CSV analysis to sales-analyst.",
|
||||
allowed_tools=["Agent", "Bash"],
|
||||
permission_mode="bypassPermissions",
|
||||
max_turns=6,
|
||||
max_budget_usd=1,
|
||||
options_kwargs={
|
||||
"tools": ["Agent", "Bash"],
|
||||
"agents": {
|
||||
"sales-analyst": AgentDefinition(
|
||||
description="Computes sales totals from CSV files.",
|
||||
prompt=f"Your working directory is {workdir}. Use Bash to read sales.csv "
|
||||
"and calculate the row count and sum of the amount column. Return both numbers.",
|
||||
tools=["Bash"],
|
||||
model="sonnet",
|
||||
)
|
||||
},
|
||||
},
|
||||
)
|
||||
run = None
|
||||
async for event in agent.arun(
|
||||
"Use the sales-analyst subagent to analyze sales.csv. Do not run Bash yourself. "
|
||||
"Wait for the subagent and summarize its row count and total.",
|
||||
session_id=SESSION_ID,
|
||||
stream=True,
|
||||
yield_run_output=True,
|
||||
):
|
||||
if isinstance(event, RunOutput):
|
||||
run = event
|
||||
assert run is not None and run.status == RunStatus.completed, run
|
||||
print(f"Source answer: {run.content}")
|
||||
|
||||
checkpoints = list_run_checkpoints(run)
|
||||
checkpoint_indexes = {checkpoint["message_index"] for checkpoint in checkpoints}
|
||||
nested_results = []
|
||||
top_level_results = []
|
||||
for index, message in enumerate(run.messages or [], start=1):
|
||||
if message.role != "tool":
|
||||
continue
|
||||
recorded = bool((message.provider_data or {}).get("claude_sdk"))
|
||||
print(
|
||||
f"Tool result at message {index}: top_level={recorded}, checkpoint={index in checkpoint_indexes}"
|
||||
)
|
||||
if recorded:
|
||||
top_level_results.append(index)
|
||||
else:
|
||||
nested_results.append(index)
|
||||
assert index not in checkpoint_indexes, (
|
||||
"subagent steps cannot be replayed independently"
|
||||
)
|
||||
assert message.checkpoint_status is None
|
||||
assert nested_results, (
|
||||
"the SDK must emit subagent tool messages to demonstrate filtering"
|
||||
)
|
||||
assert top_level_results, "the parent Agent tool must complete"
|
||||
checkpoint = next(
|
||||
item for item in checkpoints if item["message_index"] in top_level_results
|
||||
)
|
||||
print("Selected top-level checkpoint:", checkpoint)
|
||||
branch = await agent.acontinue_run(
|
||||
run_id=run.run_id,
|
||||
session_id=SESSION_ID,
|
||||
continue_from=checkpoint["message_index"],
|
||||
fork=True,
|
||||
input="Without using tools or delegating again, summarize the sales analyst's result already in context.",
|
||||
)
|
||||
assert branch.status == RunStatus.completed, branch.content
|
||||
assert branch.forked_from_run_id == run.run_id
|
||||
kept_ids = {
|
||||
message.tool_call_id
|
||||
for message in run.messages[: checkpoint["message_index"]]
|
||||
if message.role == "tool"
|
||||
}
|
||||
assert {tool.tool_call_id for tool in branch.tools or []} == kept_ids
|
||||
print(f"Branch answer: {branch.content}")
|
||||
print(
|
||||
"PASS: subagent steps were excluded; the parent checkpoint continued successfully"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
@@ -0,0 +1,56 @@
|
||||
# Continue a ClaudeAgent run from a step
|
||||
|
||||
`ClaudeAgent.continue_run` / `acontinue_run` take the same `continue_from` and `fork` arguments as `Agent.continue_run`, and AgentOS serves them through the same `/continue` and `/checkpoints` endpoints.
|
||||
|
||||
## How it works
|
||||
|
||||
Every run stores its messages in Agno, and with transcript storage the Claude conversation itself is mirrored into the `agno_transcripts` table. Each fully recorded top-level tool batch is a checkpoint. Continuing from a checkpoint forks the Claude conversation at that message, so the model only remembers what happened up to it, and the continuation runs as a new turn on that branch.
|
||||
|
||||
| Call | Effect |
|
||||
|------|--------|
|
||||
| `continue_from="end"` (default) with `input` | Append a new instruction after the last message. |
|
||||
| `continue_from=<message_index>` with `input` | Branch at that checkpoint; the model forgets everything after it. |
|
||||
| `continue_from="last_user"` without `input` | Re-send the original prompt and produce a fresh answer. |
|
||||
| `continue_from=0` with `input` | Branch before the prompt, so `input` replaces it. |
|
||||
|
||||
A finished run is never rewritten in place. Whatever `fork` is set to, the continuation is a new run in the same session with `forked_from_run_id` and `forked_from_message_index` set, and the source run is kept as is. Later turns in the session continue from the newest branch.
|
||||
|
||||
If a later turn cannot resume its SDK session, history recovery follows the active branch's ancestry. Discarded sibling replies are excluded, and the retained source messages appear only once. AgentOS streams and buffers continuation events under the new run's ID, including when a completed run forks automatically; use that ID to resume the branch's stream.
|
||||
|
||||
## Cookbooks
|
||||
|
||||
Each cookbook seeds a temporary workspace with a small realistic task and asserts on files, tool results and run lineage rather than on the model's wording.
|
||||
|
||||
| File | Scenario | What it shows |
|
||||
|------|----------|---------------|
|
||||
| `01_continue_finished_run.py` | Write a Python module and run it, then ask for tests for "the function you just wrote". | A continuation remembers the source run; it is a new run with `forked_from_run_id`. |
|
||||
| `02_fork_from_checkpoint.py` | Three Bash steps over a CSV file: count rows, sum a column, write a summary. Fork after step one. | The branch knows the row count and not the total; checkpoint listing; files are not rewound. |
|
||||
| `03_replay_user_turn.py` | List markdown files, add a file, replay the prompt; then rewrite the prompt from message 0. | Replay re-runs tools against the current workspace; `continue_from=0` replaces the prompt. |
|
||||
| `04_background_continue.py` | Write a script, then run it in a background continuation with `stream=True`. | Live tool and content events from a background run, the final `RunOutput`, and the stored result. |
|
||||
| `05_agentos_api.py` | The CSV scenario over HTTP. `--verify` runs it end to end. | `/checkpoints`, `/continue` with `fork=true`, a streamed follow-up, and the session's run lineage. |
|
||||
| `06_recover_failed_tool_run.py` | Record an invoice once, deliberately exhaust the turn limit, then recover from the stored result. | Completed tool results survive a failed run; continuation does not repeat the write. Add `--stream` to test streaming recovery. |
|
||||
| `07_subagent_checkpoints.py` | Delegate CSV analysis to a sales analyst subagent, inspect checkpoints, and continue from its parent result. | Internal subagent steps are excluded from checkpoints; the top-level result can be continued. |
|
||||
|
||||
A script that consumes a background stream and exits at once can log "Failed to complete event stream": the final output reaches the client before the run finishes its event-stream bookkeeping, and closing the loop cancels that step. `await_background_runs()` from `agno.run.background` waits for it; a server keeps its loop alive and does not need it.
|
||||
|
||||
The SDK's default system prompt does not tell the model its working directory, so the cookbooks that create files set `system_prompt` to name it. Bash already runs in `cwd`.
|
||||
|
||||
```bash
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/01_continue_finished_run.py
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/02_fork_from_checkpoint.py
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/03_replay_user_turn.py
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/04_background_continue.py
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/05_agentos_api.py --verify
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/06_recover_failed_tool_run.py
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/06_recover_failed_tool_run.py --stream
|
||||
.venvs/demo/bin/python cookbook/frameworks/claude-agent-sdk/continue_from/07_subagent_checkpoints.py
|
||||
```
|
||||
|
||||
## Limits
|
||||
|
||||
- Requires a database with transcript storage (SqliteDb, AsyncSqliteDb, PostgresDb, AsyncPostgresDb). Only runs recorded with it can be continued.
|
||||
- Files the agent changed after a checkpoint are not rewound; only the conversation is.
|
||||
- Cancelled runs cannot be continued.
|
||||
- Failed runs retain completed tool results. The error stays in the run's content; it is not treated as an assistant transcript entry when continuing.
|
||||
- Only fully recorded top-level tool batches are checkpoints. Subagent tool steps are not checkpoints. When Claude issues several tool calls in parallel, the batch is a single checkpoint and a branch keeps all of its results; an incomplete batch cannot be used as a continuation boundary.
|
||||
- `regenerate` and HITL `requirements` are not supported for ClaudeAgent.
|
||||
@@ -0,0 +1,87 @@
|
||||
# Test log
|
||||
|
||||
## 2026-10-09
|
||||
|
||||
Environment: demo venv, claude-agent-sdk, claude-sonnet-4-6, SQLite transcript storage in a temporary workspace per run. A finished run is always continued as a new sibling run with fork lineage. Assertions are on files, tool results and lineage, not on the model's wording.
|
||||
|
||||
Finding while writing these: the SDK's default system prompt does not tell the model its working directory, so the Write tool wrote to `/slugify.py` and failed with a read-only file system error. Cookbooks that create files now set `system_prompt` naming the workspace. Bash was never affected because the subprocess runs in `cwd`.
|
||||
|
||||
### 01_continue_finished_run.py
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** The source run writes `slugify.py` and runs it. The continuation asks for `test_slugify.py` covering "the function you just wrote" without naming it, then runs the tests. Asserts both files exist, the continuation has a new run_id with `forked_from_run_id` and `forked_from_message_index` set, the source run is unchanged, and the session holds two runs.
|
||||
|
||||
**Result:** Source replied `hello-world`; continuation wrote and ran the tests and replied `ok`. Two runs in the session, the second forked from the first.
|
||||
|
||||
---
|
||||
|
||||
### 02_fork_from_checkpoint.py
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Three Bash steps over a seeded `sales.csv` (count rows, sum amount, write summary.txt). Lists checkpoints, forks after the first step, and asks the branch without tools what it knows. Asserts the branch carries exactly the kept tool results, lineage, and that the source run and summary.txt are intact.
|
||||
|
||||
**Result:** Checkpoints at 3 (`5`), 5 (`500`), 7 (file written), 8 (end). The branch answered rows 5 and "not computed yet" for the total. summary.txt still on disk.
|
||||
|
||||
---
|
||||
|
||||
### 03_replay_user_turn.py
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Lists `notes/*.md`, adds `retro.md`, replays the prompt with `continue_from="last_user"`, then rewrites the prompt with `continue_from=0`. Asserts the replayed tool output contains the new file and that both branches carry lineage.
|
||||
|
||||
**Result:** Original answer `roadmap.md`; replay answered `retro.md, roadmap.md`; rewrite counted bullets per file. Three runs in the session.
|
||||
|
||||
---
|
||||
|
||||
### 04_background_continue.py
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Writes `fizzbuzz.py`, then continues with `background=True, stream=True, yield_run_output=True` to run it. Prints RunStarted, ToolCallStarted, ToolCallCompleted, RunContent and RunCompleted as they arrive, checks the final RunOutput and lineage, then reads the stored run. Ends with `await_background_runs()`.
|
||||
|
||||
**Result:** Events streamed live; final content `FizzBuzz`; stored run COMPLETED with two tool results. Before `await_background_runs()` was added, the script logged "Failed to complete event stream" on exit because the loop closed while the run was marking its event stream complete; the run itself was already COMPLETED. The helper was added to `agno.run.background` and the warning no longer appears.
|
||||
|
||||
---
|
||||
|
||||
### 05_agentos_api.py --verify
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Through the AgentOS test client: POST a three-step CSV run, GET `/checkpoints`, POST `/continue` with `fork=true` from the first checkpoint asking what the branch knows, POST `/continue` from `end` streamed with a question about the data, then list the session runs.
|
||||
|
||||
**Result:** Run replied `rows=5 total=500 max=east`; checkpoints 3, 5, 7, 8; the branch knew the row count and said the total was not computed yet; the streamed follow-up answered `east`; the session listed the source plus two forked runs.
|
||||
|
||||
---
|
||||
|
||||
### 06_recover_failed_tool_run.py
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Runs a script that appends one invoice entry to an audit log, deliberately hits `max_turns=1`, then forks from the failed run's end. Checks the completed tool result survives and that recovery executes no additional tools or invoice writes.
|
||||
|
||||
**Result:** Live Claude SDK 0.2.110 with claude-sonnet-4-6 and SQLite: the initial run returned `error_max_turns` after recording invoice-1042. Its stored result reported one audit entry, and continuation answered from that result. The audit log still contained exactly one entry.
|
||||
|
||||
---
|
||||
|
||||
### 07_subagent_checkpoints.py
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Delegates CSV analysis to a sales analyst subagent, checks that nested tool results have no checkpoint markers, then forks from the top-level Agent result without re-running the analysis.
|
||||
|
||||
**Result:** Live Claude SDK 0.2.110 with claude-sonnet-4-6 and SQLite: the parent result at message 3 was a checkpoint; nested results at messages 5 and 7 were excluded. The branch recalled 3 rows and a total of 400 without new tool calls.
|
||||
|
||||
---
|
||||
|
||||
### 06_recover_failed_tool_run.py --stream
|
||||
|
||||
**Status:** PASS
|
||||
|
||||
**Description:** Repeats the invoice-recovery scenario with streamed events and a final `RunOutput`. Checks that the terminal RunError retains the completed tool checkpoint for continuation.
|
||||
|
||||
**Result:** Live Claude SDK 0.2.110: streamed run ended with `error_max_turns`, preserved the result at message 3, and recovered successfully. No additional tools ran and the audit log contained exactly one invoice entry.
|
||||
|
||||
---
|
||||
@@ -27,13 +27,12 @@ from pathlib import Path
|
||||
from typing import Any, Dict, List, Sequence
|
||||
from uuid import uuid4
|
||||
|
||||
from sqlalchemy import text
|
||||
|
||||
from agno.agents.claude import ClaudeAgent
|
||||
from agno.db.base import BaseDb
|
||||
from agno.db.postgres import PostgresDb
|
||||
from agno.db.sqlite import SqliteDb
|
||||
from agno.run.base import RunStatus
|
||||
from sqlalchemy import text
|
||||
|
||||
AGENT_ID = "transcript-store-demo"
|
||||
SECRET = "cobalt-orchard-742"
|
||||
|
||||
@@ -42,8 +42,8 @@ from agno.utils.agent import (
|
||||
execute_system_message,
|
||||
)
|
||||
from agno.utils.common import is_typed_dict
|
||||
from agno.utils.log import log_debug, log_info, log_warning
|
||||
from agno.utils.knowledge import get_run_response_kwarg, get_user_id_kwarg
|
||||
from agno.utils.log import log_debug, log_info, log_warning
|
||||
from agno.utils.message import copy_history_message, filter_tool_calls, get_text_from_message, render_instructions
|
||||
from agno.utils.prompts import get_json_output_prompt, get_response_model_format_prompt
|
||||
from agno.utils.timer import Timer
|
||||
|
||||
@@ -45,7 +45,6 @@ from agno.utils.log import log_debug, log_error, log_warning
|
||||
from agno.utils.merge_dict import merge_dictionaries
|
||||
from agno.utils.string import generate_id_from_name
|
||||
|
||||
|
||||
# MemoryManager.__init__ (agno/memory/manager.py) auto-generates
|
||||
# ``memory_manager_<8 hex>`` when no id is passed. Such an id is minted fresh
|
||||
# every process, so a config carrying it can never resolve against a registry
|
||||
|
||||
@@ -33,8 +33,8 @@ from agno.agent import (
|
||||
_tools,
|
||||
_utils,
|
||||
)
|
||||
from agno.compaction.compaction import Compaction
|
||||
from agno.agent.followup import FollowupConfig, resolve_followup_settings
|
||||
from agno.compaction.compaction import Compaction
|
||||
from agno.compression.manager import CompressionManager
|
||||
from agno.db.base import AsyncBaseDb, BaseDb, ComponentType, UserMemory
|
||||
from agno.eval.base import BaseEval
|
||||
|
||||
@@ -76,6 +76,19 @@ class ExternalRunResult:
|
||||
warnings: Optional[List[Dict[str, Any]]] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ExternalContinuation:
|
||||
"""Replay of a stored run from one of its message boundaries."""
|
||||
|
||||
messages: List[Message]
|
||||
tools: Optional[List[ToolExecution]]
|
||||
source_input: Optional[RunInput]
|
||||
record_input: Any
|
||||
anchor: Dict[str, Any]
|
||||
forked_from_run_id: Optional[str] = None
|
||||
forked_from_message_index: Optional[int] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaseExternalAgent:
|
||||
"""Base class for external framework adapters.
|
||||
@@ -863,6 +876,21 @@ class BaseExternalAgent:
|
||||
created_at=now,
|
||||
)
|
||||
|
||||
def _finish_run_output(
|
||||
self, run: RunOutput, run_state: Dict[str, Any], continuation: Optional[ExternalContinuation]
|
||||
) -> None:
|
||||
"""Prepend a continuation's kept transcript, then let the adapter annotate messages."""
|
||||
if continuation is not None:
|
||||
run.messages = continuation.messages + (run.messages or [])
|
||||
run.tools = (continuation.tools or []) + (run.tools or []) or None
|
||||
run.input = continuation.source_input
|
||||
run.forked_from_run_id = continuation.forked_from_run_id
|
||||
run.forked_from_message_index = continuation.forked_from_message_index
|
||||
self._annotate_run_output(run, run_state)
|
||||
|
||||
def _annotate_run_output(self, run: RunOutput, run_state: Dict[str, Any]) -> None:
|
||||
"""Attach adapter-specific per-message data recorded in run_state during the run."""
|
||||
|
||||
def _get_history_from_session(
|
||||
self, session: AgentSession, exclude_run_id: Optional[str] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
@@ -880,10 +908,35 @@ class BaseExternalAgent:
|
||||
history: List[Dict[str, Any]] = []
|
||||
if not session.runs:
|
||||
return history
|
||||
# A fork already contains the retained prefix of its source run. Its
|
||||
# predecessor is the source's predecessor, not the most recent sibling.
|
||||
# Keep predecessor links so forks of older branches also discard any
|
||||
# intervening turns, without copying the entire history for every run.
|
||||
runs: List[RunOutput] = []
|
||||
predecessors: List[Optional[int]] = []
|
||||
indexes: Dict[str, int] = {}
|
||||
head: Optional[int] = None
|
||||
for run in session.runs:
|
||||
if not isinstance(run, RunOutput) or not run.messages or run.run_id == exclude_run_id:
|
||||
continue
|
||||
for msg in run.messages:
|
||||
predecessor = head
|
||||
if run.forked_from_run_id:
|
||||
source_index = indexes.get(run.forked_from_run_id)
|
||||
# If the source was deleted, its earlier ancestry is unknown;
|
||||
# only the prefix retained in the fork is safe to replay.
|
||||
predecessor = predecessors[source_index] if source_index is not None else None
|
||||
head = len(runs)
|
||||
runs.append(run)
|
||||
predecessors.append(predecessor)
|
||||
if run.run_id:
|
||||
indexes[run.run_id] = head
|
||||
|
||||
branch: List[RunOutput] = []
|
||||
while head is not None:
|
||||
branch.append(runs[head])
|
||||
head = predecessors[head]
|
||||
for run in reversed(branch):
|
||||
for msg in run.messages or []:
|
||||
if msg.role == "assistant" and msg.tool_calls:
|
||||
# Assistant message with tool calls (no text content)
|
||||
history.append(
|
||||
@@ -976,16 +1029,23 @@ class BaseExternalAgent:
|
||||
session_id = kwargs.get("session_id") or str(uuid4())
|
||||
kwargs["session_id"] = session_id
|
||||
user_id = kwargs.get("user_id")
|
||||
continuation: Optional[ExternalContinuation] = kwargs.get("continuation")
|
||||
record_input = input if continuation is None else continuation.record_input
|
||||
run_state: Dict[str, Any] = {}
|
||||
session = await self.aread_or_create_session(session_id, user_id) if self.db else None
|
||||
history = self._get_history_from_session(session, exclude_run_id=run_id) if session else None
|
||||
history = (
|
||||
self._get_history_from_session(session, exclude_run_id=run_id) if session and not continuation else None
|
||||
)
|
||||
try:
|
||||
async with self._run_cancellation(run_id):
|
||||
content = await self._arun_adapter(input, history=history, run_id=run_id, session=session, **kwargs)
|
||||
content = await self._arun_adapter(
|
||||
input, history=history, run_id=run_id, session=session, run_state=run_state, **kwargs
|
||||
)
|
||||
run_output = self._build_run_output(
|
||||
run_id,
|
||||
session_id,
|
||||
user_id,
|
||||
input,
|
||||
record_input,
|
||||
content.content if isinstance(content, ExternalRunResult) else content,
|
||||
RunStatus.completed,
|
||||
tools=content.tools if isinstance(content, ExternalRunResult) else None,
|
||||
@@ -994,11 +1054,26 @@ class BaseExternalAgent:
|
||||
run_output.metadata = {"warnings": content.warnings}
|
||||
except RunCancelledException:
|
||||
run_output = self._build_run_output(
|
||||
run_id, session_id, user_id, input, "Run cancelled", RunStatus.cancelled
|
||||
run_id,
|
||||
session_id,
|
||||
user_id,
|
||||
record_input,
|
||||
"Run cancelled",
|
||||
RunStatus.cancelled,
|
||||
tools=list(run_state.get("tools", {}).values()) or None,
|
||||
)
|
||||
except Exception as error:
|
||||
log_exception(f"Error in {self.framework} agent '{self.id}': {error}")
|
||||
run_output = self._build_run_output(run_id, session_id, user_id, input, str(error), RunStatus.error)
|
||||
run_output = self._build_run_output(
|
||||
run_id,
|
||||
session_id,
|
||||
user_id,
|
||||
record_input,
|
||||
str(error),
|
||||
RunStatus.error,
|
||||
tools=list(run_state.get("tools", {}).values()) or None,
|
||||
)
|
||||
self._finish_run_output(run_output, run_state, continuation)
|
||||
if session is not None:
|
||||
await self._apersist_run_in_session(session, run_output)
|
||||
return run_output
|
||||
@@ -1009,8 +1084,12 @@ class BaseExternalAgent:
|
||||
kwargs["session_id"] = session_id
|
||||
user_id = kwargs.get("user_id")
|
||||
yield_run_output = kwargs.pop("yield_run_output", False)
|
||||
continuation: Optional[ExternalContinuation] = kwargs.get("continuation")
|
||||
run_state: Dict[str, Any] = {}
|
||||
session = await self.aread_or_create_session(session_id, user_id) if self.db else None
|
||||
history = self._get_history_from_session(session, exclude_run_id=run_id) if session else None
|
||||
history = (
|
||||
self._get_history_from_session(session, exclude_run_id=run_id) if session and not continuation else None
|
||||
)
|
||||
yield RunStartedEvent(run_id=run_id, agent_id=self.get_id(), agent_name=self.name or "", session_id=session_id)
|
||||
accumulated_content = ""
|
||||
warnings: List[Dict[str, Any]] = []
|
||||
@@ -1020,7 +1099,7 @@ class BaseExternalAgent:
|
||||
try:
|
||||
async with self._run_cancellation(run_id):
|
||||
async for event in self._arun_adapter_stream(
|
||||
input, history=history, run_id=run_id, session=session, **kwargs
|
||||
input, history=history, run_id=run_id, session=session, run_state=run_state, **kwargs
|
||||
):
|
||||
if isinstance(event, ExternalRunWarningEvent) and event.warning is not None:
|
||||
warnings.append(event.warning)
|
||||
@@ -1044,13 +1123,14 @@ class BaseExternalAgent:
|
||||
run_id,
|
||||
session_id,
|
||||
user_id,
|
||||
input,
|
||||
input if continuation is None else continuation.record_input,
|
||||
str(run_error) if run_error else accumulated_content,
|
||||
status,
|
||||
list(tools.values()) or None,
|
||||
)
|
||||
if warnings:
|
||||
run.metadata = {"warnings": warnings}
|
||||
self._finish_run_output(run, run_state, continuation)
|
||||
if session is not None:
|
||||
await self._apersist_run_in_session(session, run)
|
||||
fields: Dict[str, Any] = dict(
|
||||
|
||||
@@ -1,16 +1,21 @@
|
||||
import copy
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, AsyncIterator, Dict, List, Optional
|
||||
from typing import Any, AsyncIterator, Awaitable, Dict, Iterator, List, Literal, Optional, Tuple, Union, cast
|
||||
from uuid import uuid4
|
||||
|
||||
from agno.agents.base import BaseExternalAgent, ExternalRunResult, ExternalRunWarningEvent
|
||||
from agno.agents.base import BaseExternalAgent, ExternalContinuation, ExternalRunResult, ExternalRunWarningEvent
|
||||
from agno.db.base import AsyncBaseDb, BaseDb
|
||||
from agno.exceptions import RunNotContinuableError
|
||||
from agno.models.message import Message
|
||||
from agno.models.response import ToolExecution
|
||||
from agno.run.agent import (
|
||||
RunContentEvent,
|
||||
RunOutput,
|
||||
RunOutputEvent,
|
||||
ToolCallCompletedEvent,
|
||||
ToolCallStartedEvent,
|
||||
)
|
||||
from agno.run.base import RunStatus
|
||||
from agno.run.cancel import araise_if_cancelled
|
||||
from agno.utils.log import log_debug, log_warning
|
||||
|
||||
@@ -85,6 +90,9 @@ class ClaudeAgent(BaseExternalAgent):
|
||||
|
||||
# Key under which the SDK session id is stored in the Agno session's session_data.
|
||||
_SESSION_KEY = "claude_sdk_session_id"
|
||||
# Message.provider_data key holding the SDK session id and transcript uuid behind a message.
|
||||
_MESSAGE_REF_KEY = "claude_sdk"
|
||||
_CONTINUE_PROMPT = "Continue from where you left off."
|
||||
|
||||
# Fallback Agno session_id -> SDK session id map, used when no db is configured.
|
||||
_sdk_session_ids: Dict[str, str] = field(default_factory=dict, init=False, repr=False)
|
||||
@@ -183,6 +191,8 @@ class ClaudeAgent(BaseExternalAgent):
|
||||
opts["resume"] = resume
|
||||
|
||||
opts.update(self.options_kwargs)
|
||||
# Echo the prompt with its transcript uuid so a run can later be continued from it.
|
||||
opts["extra_args"] = {"replay-user-messages": None, **(self.options_kwargs.get("extra_args") or {})}
|
||||
store = session_store if session_store is not None else self._transcript_store(agno_session_id)
|
||||
if store is not None:
|
||||
opts["session_store"] = store
|
||||
@@ -231,10 +241,17 @@ class ClaudeAgent(BaseExternalAgent):
|
||||
sdk = _sdk()
|
||||
session = kwargs.get("session")
|
||||
session_id = kwargs.get("session_id")
|
||||
resume = self._get_sdk_session_id(session, session_id)
|
||||
run_state = kwargs.get("run_state")
|
||||
continuation: Optional[ExternalContinuation] = kwargs.get("continuation")
|
||||
store = self._transcript_store(session.session_id) if session is not None else None
|
||||
if continuation is not None:
|
||||
resume = await self._afork_sdk_session(continuation.anchor, store)
|
||||
else:
|
||||
resume = self._get_sdk_session_id(session, session_id)
|
||||
|
||||
while True:
|
||||
if run_state is not None:
|
||||
run_state.clear()
|
||||
options = self._build_options(streaming=streaming, resume=resume, session_store=store)
|
||||
prompt = self._build_prompt(input, history, resumed=resume is not None)
|
||||
received = False
|
||||
@@ -246,20 +263,28 @@ class ClaudeAgent(BaseExternalAgent):
|
||||
await client.query(prompt)
|
||||
self._set_run_handle(run_id, client)
|
||||
async for message in client.receive_response():
|
||||
self._record_message_ref(sdk, message, run_state)
|
||||
if isinstance(message, sdk.UserMessage) and not self._tool_result_ids(sdk, message):
|
||||
# The echoed prompt is bookkeeping only; adapters never see it.
|
||||
continue
|
||||
if isinstance(message, sdk.ResultMessage):
|
||||
self._check_result_message(sdk, message)
|
||||
if not isinstance(message, sdk.SystemMessage):
|
||||
received = True
|
||||
sdk_session_id = None
|
||||
if isinstance(message, sdk.SystemMessage) and getattr(message, "subtype", None) == "init":
|
||||
data = getattr(message, "data", {}) or {}
|
||||
self._remember_sdk_session(session, session_id, data.get("session_id"))
|
||||
sdk_session_id = (getattr(message, "data", {}) or {}).get("session_id")
|
||||
elif isinstance(message, sdk.ResultMessage):
|
||||
self._remember_sdk_session(session, session_id, getattr(message, "session_id", None))
|
||||
sdk_session_id = getattr(message, "session_id", None)
|
||||
if sdk_session_id:
|
||||
self._remember_sdk_session(session, session_id, sdk_session_id)
|
||||
if run_state is not None:
|
||||
run_state["session_id"] = sdk_session_id
|
||||
yield message
|
||||
return
|
||||
except Exception as e:
|
||||
await araise_if_cancelled(run_id)
|
||||
if resume is None or received or not self._is_missing_session(sdk, e):
|
||||
if resume is None or received or continuation is not None or not self._is_missing_session(sdk, e):
|
||||
raise
|
||||
log_warning(
|
||||
f"Claude SDK: could not resume session {resume} for session {session_id}: {e}. Starting a new one."
|
||||
@@ -273,6 +298,397 @@ class ClaudeAgent(BaseExternalAgent):
|
||||
async def _ainterrupt_run(self, handle: Any) -> None:
|
||||
await handle.interrupt()
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Continue from a message boundary
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _tool_result_ids(sdk: Any, message: Any) -> List[str]:
|
||||
content = getattr(message, "content", None)
|
||||
if not isinstance(content, list):
|
||||
return []
|
||||
return [block.tool_use_id for block in content if isinstance(block, sdk.ToolResultBlock)]
|
||||
|
||||
def _record_message_ref(self, sdk: Any, message: Any, run_state: Optional[Dict[str, Any]]) -> None:
|
||||
"""Remember which transcript uuid produced each top-level prompt, tool call, tool result and reply."""
|
||||
uuid = getattr(message, "uuid", None)
|
||||
if run_state is None or not uuid or getattr(message, "parent_tool_use_id", None):
|
||||
return
|
||||
# Transcript order: parallel tool results can arrive in a different order than their calls.
|
||||
order = run_state.setdefault("order", {})
|
||||
if uuid not in order:
|
||||
order[uuid] = len(order)
|
||||
if isinstance(message, sdk.AssistantMessage):
|
||||
run_state["final"] = uuid
|
||||
run_state["last"] = uuid
|
||||
run_state["final_text"] = (
|
||||
"".join(block.text for block in message.content if isinstance(block, sdk.TextBlock))
|
||||
if not any(isinstance(block, sdk.ToolUseBlock) for block in message.content)
|
||||
else None
|
||||
)
|
||||
for block in message.content:
|
||||
if isinstance(block, sdk.ToolUseBlock):
|
||||
run_state[f"call:{block.id}"] = uuid
|
||||
elif isinstance(message, sdk.UserMessage):
|
||||
run_state["last"] = uuid
|
||||
result_ids = self._tool_result_ids(sdk, message)
|
||||
for tool_use_id in result_ids:
|
||||
run_state[f"result:{tool_use_id}"] = uuid
|
||||
if not result_ids:
|
||||
run_state.setdefault("prompt", uuid)
|
||||
|
||||
def _annotate_run_output(self, run: RunOutput, run_state: Dict[str, Any]) -> None:
|
||||
"""Store each message's transcript position and mark the end of each tool batch as a checkpoint.
|
||||
|
||||
A fork keeps every transcript entry up to a position. Tool results from one batch of
|
||||
parallel calls can sit in the transcript in a different order than Agno lists them, so a
|
||||
checkpoint is only exposed where the kept messages and the kept transcript agree: at a
|
||||
tool result that is later in the transcript than every tool result before it and earlier
|
||||
than every tool result after it.
|
||||
"""
|
||||
sdk_session_id = run_state.get("session_id")
|
||||
if not sdk_session_id or not run.messages:
|
||||
return
|
||||
if run.status in (RunStatus.error, RunStatus.cancelled):
|
||||
# The base class puts the error text in a synthetic assistant
|
||||
# message. It is not a transcript entry; retain the error in
|
||||
# run.content and end the messages at the work actually received.
|
||||
run.messages.pop()
|
||||
if run_state.get("final_text") and run_state.get("last") == run_state.get("final"):
|
||||
run.messages.append(Message(role="assistant", content=run_state["final_text"]))
|
||||
order: Dict[str, int] = run_state.get("order") or {}
|
||||
last = len(run.messages) - 1
|
||||
for index, message in enumerate(run.messages):
|
||||
if (message.provider_data or {}).get(self._MESSAGE_REF_KEY):
|
||||
continue
|
||||
if message.role == "user":
|
||||
key = "prompt"
|
||||
elif message.role == "tool":
|
||||
key = f"result:{message.tool_call_id}"
|
||||
elif message.tool_calls:
|
||||
key = f"call:{message.tool_calls[0].get('id')}"
|
||||
elif index == last:
|
||||
key = "final"
|
||||
else:
|
||||
continue
|
||||
uuid = run_state.get(key)
|
||||
if uuid is None:
|
||||
continue
|
||||
ref: Dict[str, Any] = {"session_id": sdk_session_id, "uuid": uuid}
|
||||
if uuid in order:
|
||||
ref["position"] = order[uuid]
|
||||
message.provider_data = {**(message.provider_data or {}), self._MESSAGE_REF_KEY: ref}
|
||||
for message in run.messages:
|
||||
if message.role == "tool":
|
||||
message.checkpoint_status = None
|
||||
message.checkpoint_created_at = None
|
||||
for index in self._checkpoint_indexes(run.messages):
|
||||
message = run.messages[index]
|
||||
message.checkpoint_status = RunStatus.running.value
|
||||
message.checkpoint_created_at = message.created_at
|
||||
|
||||
def _ref(self, message: Message) -> Optional[Dict[str, Any]]:
|
||||
return (message.provider_data or {}).get(self._MESSAGE_REF_KEY)
|
||||
|
||||
@staticmethod
|
||||
def _tool_batches(messages: List[Message]) -> List[List[int]]:
|
||||
"""Zero-based indexes of tool results, grouped into batches issued before the next reply.
|
||||
|
||||
Parallel calls produce one batch; a fork can only keep a batch whole, because its results
|
||||
may sit in the transcript in a different order than Agno lists them.
|
||||
"""
|
||||
batches: List[List[int]] = []
|
||||
current: List[int] = []
|
||||
last_result_position: Optional[int] = None
|
||||
for index, message in enumerate(messages):
|
||||
if message.role == "tool":
|
||||
current.append(index)
|
||||
position = ((message.provider_data or {}).get("claude_sdk") or {}).get("position")
|
||||
if position is not None:
|
||||
last_result_position = max(last_result_position or -1, position)
|
||||
elif message.role == "assistant" and message.tool_calls:
|
||||
# A call issued before the previous result arrived belongs to the same parallel batch;
|
||||
# a call issued after it starts a new, sequential step.
|
||||
position = ((message.provider_data or {}).get("claude_sdk") or {}).get("position")
|
||||
parallel = (
|
||||
current
|
||||
and position is not None
|
||||
and last_result_position is not None
|
||||
and position < last_result_position
|
||||
)
|
||||
if current and not parallel:
|
||||
batches.append(current)
|
||||
current = []
|
||||
last_result_position = None
|
||||
else:
|
||||
if current:
|
||||
batches.append(current)
|
||||
current = []
|
||||
last_result_position = None
|
||||
if current:
|
||||
batches.append(current)
|
||||
return batches
|
||||
|
||||
def _checkpoint_indexes(self, messages: List[Message]) -> List[int]:
|
||||
"""Zero-based indexes of the last tool result of each batch: the only places a fork is exact."""
|
||||
return [
|
||||
batch[-1] for batch in self._tool_batches(messages) if all(self._ref(messages[index]) for index in batch)
|
||||
]
|
||||
|
||||
def _batch_end(self, messages: List[Message], index: int) -> int:
|
||||
"""Move a boundary inside a batch of tool results to the end of that batch."""
|
||||
if index <= 0 or index > len(messages) or messages[index - 1].role != "tool":
|
||||
return index
|
||||
for batch in self._tool_batches(messages):
|
||||
if index - 1 in batch:
|
||||
return batch[-1] + 1
|
||||
return index
|
||||
|
||||
def _batch_anchor(self, messages: List[Message], index: int) -> Optional[Dict[str, Any]]:
|
||||
"""The transcript entry a fork must keep up to so the whole batch ending at index is kept."""
|
||||
for batch in self._tool_batches(messages):
|
||||
if index - 1 in batch:
|
||||
refs = [self._ref(messages[i]) or {} for i in batch]
|
||||
if not all(refs):
|
||||
return None
|
||||
return max(refs, key=lambda r: r.get("position", -1)) or None
|
||||
return self._ref(messages[index - 1])
|
||||
|
||||
async def _afork_sdk_session(self, anchor: Dict[str, Any], store: Any) -> Optional[str]:
|
||||
"""Fork the stored SDK transcript at the anchor; None starts a fresh SDK session."""
|
||||
if store is None:
|
||||
raise ValueError("Continuing a ClaudeAgent run requires a database with transcript storage")
|
||||
up_to: Optional[str] = anchor["uuid"]
|
||||
if anchor.get("before"):
|
||||
key = {"project_key": store.project_key, "session_id": anchor["session_id"]}
|
||||
entries = await store.load(key) or []
|
||||
entry = next((entry for entry in entries if entry.get("uuid") == up_to), None)
|
||||
if entry is None:
|
||||
raise ValueError(f"Claude SDK transcript entry {up_to} was not found")
|
||||
up_to = entry.get("parentUuid")
|
||||
if up_to is None:
|
||||
return None
|
||||
forked = await _sdk().fork_session_via_store(
|
||||
store, anchor["session_id"], directory=self.cwd, up_to_message_id=up_to
|
||||
)
|
||||
return forked.session_id
|
||||
|
||||
def _build_continuation(
|
||||
self,
|
||||
source: RunOutput,
|
||||
*,
|
||||
continue_from: Union[int, Literal["end", "last_user"]],
|
||||
fork: bool,
|
||||
input: Optional[str],
|
||||
) -> Tuple[Any, str, ExternalContinuation]:
|
||||
"""Resolve a boundary in a stored run into the prompt, run id and kept transcript for the replay."""
|
||||
from agno.agent._run import _resolve_continue_from
|
||||
from agno.utils.message import safe_truncation_index
|
||||
|
||||
status = getattr(source.status, "value", source.status)
|
||||
if status == RunStatus.cancelled.value:
|
||||
raise RunNotContinuableError(f"Cannot continue run {source.run_id}: run is cancelled")
|
||||
if not fork and status == RunStatus.completed.value:
|
||||
# Like native agents: a finished run is never rewritten in place. Its continuation is a
|
||||
# new sibling run with fork lineage, which also keeps the background path's PENDING row new.
|
||||
fork = True
|
||||
messages = source.messages or []
|
||||
if not messages:
|
||||
raise ValueError("The run has no messages to continue from")
|
||||
requested = _resolve_continue_from(source, continue_from=continue_from)
|
||||
index = 0 if requested == 0 else safe_truncation_index(messages, requested)
|
||||
if not 0 <= index <= len(messages):
|
||||
raise ValueError(f"continue_from must resolve to a message boundary between 0 and {len(messages)}")
|
||||
index = self._batch_end(messages, index)
|
||||
# The boundary message decides the mode: a user message without new input is replayed from
|
||||
# just before it; index 0 drops everything and needs new input; anything else continues after.
|
||||
boundary = messages[index - 1] if index > 0 else messages[0]
|
||||
ref = self._batch_anchor(messages, index) if boundary.role == "tool" else self._ref(boundary)
|
||||
if not ref:
|
||||
raise ValueError(
|
||||
"This run has no Claude SDK transcript position at that boundary; "
|
||||
"only runs recorded with transcript storage can be continued"
|
||||
)
|
||||
replay = index > 0 and input is None and boundary.role == "user"
|
||||
if index == 0:
|
||||
if input is None:
|
||||
raise ValueError("continue_from=0 drops the whole transcript; provide input to start the branch")
|
||||
prompt: Any = input
|
||||
kept: List[Any] = []
|
||||
tools = None
|
||||
record_input: Any = input
|
||||
before = True
|
||||
elif replay:
|
||||
prompt = boundary.content if boundary.content else None
|
||||
if prompt is None and index == 1 and source.input is not None:
|
||||
prompt = source.input.input_content
|
||||
if prompt is None:
|
||||
raise ValueError("The selected user message has no content to replay")
|
||||
kept = copy.deepcopy(messages[: index - 1])
|
||||
call_ids = {message.tool_call_id for message in kept if message.tool_call_id}
|
||||
tools = [copy.deepcopy(tool) for tool in source.tools or [] if tool.tool_call_id in call_ids] or None
|
||||
record_input = prompt
|
||||
before = True
|
||||
else:
|
||||
prompt = input or self._CONTINUE_PROMPT
|
||||
kept = copy.deepcopy(messages[:index])
|
||||
call_ids = {message.tool_call_id for message in kept if message.tool_call_id}
|
||||
tools = [copy.deepcopy(tool) for tool in source.tools or [] if tool.tool_call_id in call_ids] or None
|
||||
record_input = input
|
||||
before = False
|
||||
continuation = ExternalContinuation(
|
||||
messages=kept,
|
||||
tools=tools,
|
||||
source_input=copy.deepcopy(source.input),
|
||||
record_input=record_input,
|
||||
anchor={"session_id": ref["session_id"], "uuid": ref["uuid"], "before": before},
|
||||
forked_from_run_id=source.run_id if fork else source.forked_from_run_id,
|
||||
forked_from_message_index=index if fork else source.forked_from_message_index,
|
||||
)
|
||||
return prompt, str(uuid4()) if fork else str(source.run_id), continuation
|
||||
|
||||
@staticmethod
|
||||
def _check_continue_args(
|
||||
requirements: Any,
|
||||
updated_tools: Any,
|
||||
regenerate: bool,
|
||||
replace_original: Optional[bool],
|
||||
additional_instructions: Optional[str],
|
||||
) -> None:
|
||||
if requirements or updated_tools:
|
||||
raise ValueError("ClaudeAgent runs have no human-in-the-loop requirements to resolve")
|
||||
if regenerate or replace_original is not None or additional_instructions is not None:
|
||||
raise ValueError("ClaudeAgent does not support regenerate; use continue_from='last_user' with fork=True")
|
||||
|
||||
def continue_run(
|
||||
self,
|
||||
run_response: Optional[RunOutput] = None,
|
||||
*,
|
||||
run_id: Optional[str] = None,
|
||||
requirements: Optional[List[Any]] = None,
|
||||
updated_tools: Optional[List[ToolExecution]] = None,
|
||||
input: Optional[str] = None,
|
||||
continue_from: Union[int, Literal["end", "last_user"]] = "end",
|
||||
fork: bool = False,
|
||||
regenerate: bool = False,
|
||||
replace_original: Optional[bool] = None,
|
||||
additional_instructions: Optional[str] = None,
|
||||
stream: bool = False,
|
||||
session_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> Union[RunOutput, Iterator[RunOutputEvent]]:
|
||||
"""Continue a stored run from a message boundary by forking its SDK transcript there.
|
||||
|
||||
continue_from accepts "end", "last_user" or a message index. With fork=True the replay is a new
|
||||
sibling run. A finished run is always continued as a new sibling run with fork lineage, as native
|
||||
agents do, whatever fork is set to. Files the agent changed are not rewound.
|
||||
"""
|
||||
self._check_continue_args(requirements, updated_tools, regenerate, replace_original, additional_instructions)
|
||||
source = run_response or self.get_run_output(run_id or "", session_id, user_id)
|
||||
if source is None:
|
||||
raise ValueError(f"Run {run_id} not found in session {session_id}")
|
||||
prompt, new_run_id, continuation = self._build_continuation(
|
||||
source, continue_from=continue_from, fork=fork, input=input
|
||||
)
|
||||
return self.run(
|
||||
prompt,
|
||||
stream=stream,
|
||||
session_id=source.session_id,
|
||||
user_id=user_id or source.user_id,
|
||||
run_id=new_run_id,
|
||||
continuation=continuation,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def acontinue_run(
|
||||
self,
|
||||
run_response: Optional[RunOutput] = None,
|
||||
*,
|
||||
run_id: Optional[str] = None,
|
||||
requirements: Optional[List[Any]] = None,
|
||||
updated_tools: Optional[List[ToolExecution]] = None,
|
||||
input: Optional[str] = None,
|
||||
continue_from: Union[int, Literal["end", "last_user"]] = "end",
|
||||
fork: bool = False,
|
||||
regenerate: bool = False,
|
||||
replace_original: Optional[bool] = None,
|
||||
additional_instructions: Optional[str] = None,
|
||||
stream: bool = False,
|
||||
session_id: Optional[str] = None,
|
||||
user_id: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
"""Async continue_run: returns a coroutine, or an async iterator when stream=True."""
|
||||
self._check_continue_args(requirements, updated_tools, regenerate, replace_original, additional_instructions)
|
||||
kwargs.pop("background_tasks", None)
|
||||
call = dict(
|
||||
run_response=run_response,
|
||||
run_id=run_id,
|
||||
input=input,
|
||||
continue_from=continue_from,
|
||||
fork=fork,
|
||||
session_id=session_id,
|
||||
user_id=user_id,
|
||||
**kwargs,
|
||||
)
|
||||
if stream:
|
||||
return self._acontinue_stream(**call)
|
||||
return self._acontinue(**call)
|
||||
|
||||
async def _aresolve_continuation(
|
||||
self,
|
||||
run_response: Optional[RunOutput],
|
||||
run_id: Optional[str],
|
||||
session_id: Optional[str],
|
||||
user_id: Optional[str],
|
||||
continue_from: Union[int, Literal["end", "last_user"]],
|
||||
fork: bool,
|
||||
input: Optional[str],
|
||||
) -> Tuple[RunOutput, Any, str, ExternalContinuation]:
|
||||
source = run_response or await self.aget_run_output(run_id or "", session_id, user_id)
|
||||
if source is None:
|
||||
raise ValueError(f"Run {run_id} not found in session {session_id}")
|
||||
prompt, new_run_id, continuation = self._build_continuation(
|
||||
source, continue_from=continue_from, fork=fork, input=input
|
||||
)
|
||||
return source, prompt, new_run_id, continuation
|
||||
|
||||
async def _acontinue(
|
||||
self, *, run_response, run_id, input, continue_from, fork, session_id, user_id, **kwargs: Any
|
||||
) -> RunOutput:
|
||||
source, prompt, new_run_id, continuation = await self._aresolve_continuation(
|
||||
run_response, run_id, session_id, user_id, continue_from, fork, input
|
||||
)
|
||||
run = self.arun(
|
||||
prompt,
|
||||
stream=False,
|
||||
session_id=source.session_id,
|
||||
user_id=user_id or source.user_id,
|
||||
run_id=new_run_id,
|
||||
continuation=continuation,
|
||||
**kwargs,
|
||||
)
|
||||
return await cast(Awaitable[RunOutput], run)
|
||||
|
||||
async def _acontinue_stream(
|
||||
self, *, run_response, run_id, input, continue_from, fork, session_id, user_id, **kwargs: Any
|
||||
) -> AsyncIterator[Any]:
|
||||
source, prompt, new_run_id, continuation = await self._aresolve_continuation(
|
||||
run_response, run_id, session_id, user_id, continue_from, fork, input
|
||||
)
|
||||
stream = self.arun(
|
||||
prompt,
|
||||
stream=True,
|
||||
session_id=source.session_id,
|
||||
user_id=user_id or source.user_id,
|
||||
run_id=new_run_id,
|
||||
continuation=continuation,
|
||||
**kwargs,
|
||||
)
|
||||
async for item in cast(AsyncIterator[Any], stream):
|
||||
yield item
|
||||
|
||||
@staticmethod
|
||||
def _is_missing_session(sdk: Any, error: Exception) -> bool:
|
||||
"""True when the resume target has no transcript, as opposed to a failed run such as an API error."""
|
||||
@@ -327,6 +743,10 @@ class ClaudeAgent(BaseExternalAgent):
|
||||
tools: Dict[str, ToolExecution] = {}
|
||||
|
||||
async for message in self._aquery(input, history, streaming=False, **kwargs):
|
||||
# Share the live accumulator so the base class can persist work
|
||||
# already completed if the SDK raises before returning a result.
|
||||
if kwargs.get("run_state") is not None:
|
||||
kwargs["run_state"]["tools"] = tools
|
||||
warning = self._mirror_warning(message)
|
||||
if warning is not None:
|
||||
warnings.append(warning)
|
||||
|
||||
@@ -1,10 +1,10 @@
|
||||
"""LocalFileSystem: the disk-based backend for FileSystem."""
|
||||
|
||||
import os
|
||||
from urllib.parse import unquote
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
from urllib.parse import unquote
|
||||
|
||||
from agno.exceptions import PathSecurityError
|
||||
from agno.fs._paths import build_chunk, path_in_directory
|
||||
@@ -13,7 +13,6 @@ from agno.fs.errors import InvalidPathError, QuotaExceededError, UnsupportedOper
|
||||
from agno.fs.types import FileData, FileMeta
|
||||
from agno.utils.path_safety import safe_join_relative_path
|
||||
|
||||
|
||||
# Characters a user id keeps literally in its directory name; see _encode_user_id.
|
||||
_USER_ID_SAFE = frozenset("abcdefghijklmnopqrstuvwxyz0123456789-_")
|
||||
|
||||
|
||||
@@ -51,8 +51,8 @@ from agno.knowledge.page.types import (
|
||||
Page,
|
||||
PageChanged,
|
||||
PageError,
|
||||
PageMoved,
|
||||
PageList,
|
||||
PageMoved,
|
||||
PageNotFound,
|
||||
PageRead,
|
||||
PageSearchConfig,
|
||||
|
||||
@@ -23,7 +23,6 @@ from agno.db.base import AsyncBaseDb, BaseDb
|
||||
from agno.job_queue import QueueConfig
|
||||
from agno.knowledge.knowledge import Knowledge
|
||||
from agno.media.storage.base import AsyncMediaStorage, MediaStorage
|
||||
from agno.os.middleware.cors import OriginPolicy
|
||||
from agno.os.config import (
|
||||
AgentOSConfig,
|
||||
AuthorizationConfig,
|
||||
@@ -50,6 +49,7 @@ from agno.os.config import (
|
||||
from agno.os.event_streams import BaseEventStream, set_event_stream
|
||||
from agno.os.interfaces.base import BaseInterface
|
||||
from agno.os.job_queue import apply_queue_config, queue_lifespan
|
||||
from agno.os.middleware.cors import OriginPolicy
|
||||
from agno.os.router import get_base_router, get_info_router, get_websocket_router
|
||||
from agno.os.routers.agents import get_agent_router
|
||||
from agno.os.routers.approvals import get_approval_router
|
||||
|
||||
@@ -306,20 +306,27 @@ async def agent_continue_response_streamer(
|
||||
# continue-streamer parity): this response is otherwise their only
|
||||
# copy - after an inline continue of a formerly-queued/streamed run,
|
||||
# /resume would replay just the pre-pause prefix and the stream
|
||||
# status would stay PAUSED forever. Skipped for remote agents (the
|
||||
# remote OS owns that run's stream) and for fork/regenerate (they
|
||||
# mint a NEW run_id; publishing under the original would corrupt it).
|
||||
_sync_stream = not isinstance(agent, RemoteAgent) and not fork and not regenerate
|
||||
if _sync_stream:
|
||||
await amark_continue_stream_running(run_id, component=agent, session_id=session_id, user_id=user_id)
|
||||
# status would stay PAUSED forever. Resolve the stream identity from
|
||||
# the emitted run: completed runs can fork even when fork=False.
|
||||
# Remote agents own their event streams on the remote OS.
|
||||
_sync_stream = not isinstance(agent, RemoteAgent)
|
||||
_stream_run_id: Optional[str] = None
|
||||
try:
|
||||
async for run_response_chunk in continue_response:
|
||||
if _sync_stream and not isinstance(run_response_chunk, RunOutput):
|
||||
if _sync_stream and _stream_run_id is None:
|
||||
_stream_run_id = getattr(run_response_chunk, "run_id", None)
|
||||
if not _stream_run_id and not fork and not regenerate:
|
||||
_stream_run_id = run_id
|
||||
if _stream_run_id:
|
||||
await amark_continue_stream_running(
|
||||
_stream_run_id, component=agent, session_id=session_id, user_id=user_id
|
||||
)
|
||||
if _stream_run_id and not isinstance(run_response_chunk, RunOutput):
|
||||
with contextlib.suppress(Exception):
|
||||
await get_event_stream().add_event(run_id, run_response_chunk)
|
||||
await get_event_stream().add_event(_stream_run_id, run_response_chunk)
|
||||
yield format_sse_event(run_response_chunk) # type: ignore
|
||||
finally:
|
||||
if _sync_stream:
|
||||
if _stream_run_id:
|
||||
# Stream close + paused-ticket settle as one cancellation-
|
||||
# proof unit: a client disconnect cancels this generator,
|
||||
# and an interrupted finalizer abandoned the stream view as
|
||||
@@ -335,9 +342,9 @@ async def agent_continue_response_streamer(
|
||||
_cancelled = _exc is not None and issubclass(_exc, (asyncio.CancelledError, GeneratorExit))
|
||||
await afinalize_continue_stream(
|
||||
agent,
|
||||
run_id,
|
||||
_stream_run_id,
|
||||
session_id,
|
||||
queue_worker=queue_worker,
|
||||
queue_worker=queue_worker if _stream_run_id == run_id else None,
|
||||
final_status=RunStatus.cancelled if _cancelled else None,
|
||||
)
|
||||
except (InputCheckError, OutputCheckError) as e:
|
||||
@@ -1737,7 +1744,7 @@ def get_agent_router(
|
||||
# continue settles - only_if_tracked leaves never-streamed
|
||||
# runs alone. Skipped for remote agents and fork/regenerate
|
||||
# (they mint a NEW run_id).
|
||||
if not isinstance(agent, RemoteAgent) and not fork and not regenerate:
|
||||
if not isinstance(agent, RemoteAgent) and run_response_obj.run_id == run_id:
|
||||
# Stream close + paused-ticket settle as one
|
||||
# cancellation-proof unit (see the streaming twin)
|
||||
await afinalize_continue_stream(
|
||||
|
||||
@@ -25,11 +25,11 @@ from agno.os.routers.filesystem.schema import (
|
||||
)
|
||||
from agno.os.routers.filesystem.utils import _filesystem_backend_key
|
||||
from agno.os.schema import (
|
||||
SortOrder,
|
||||
BadRequestResponse,
|
||||
InternalServerErrorResponse,
|
||||
NotFoundResponse,
|
||||
PaginationInfo,
|
||||
SortOrder,
|
||||
UnauthenticatedResponse,
|
||||
ValidationErrorResponse,
|
||||
)
|
||||
|
||||
@@ -19,6 +19,20 @@ def _spawn_background(coro: Awaitable[None]) -> None:
|
||||
task.add_done_callback(_background_tasks.discard)
|
||||
|
||||
|
||||
async def await_background_runs(timeout: Optional[float] = None) -> None:
|
||||
"""Wait for every background run started on this event loop to finish its bookkeeping.
|
||||
|
||||
A background run hands its final output to the attached client before it marks the event
|
||||
stream complete, so a script that exits right after consuming the stream cancels that last
|
||||
step and logs a warning. Call this before closing the loop. A server keeps its loop alive
|
||||
and does not need it.
|
||||
"""
|
||||
pending = [task for task in _background_tasks if not task.done()]
|
||||
if not pending:
|
||||
return
|
||||
await asyncio.wait(pending, timeout=timeout)
|
||||
|
||||
|
||||
async def _execute_background(
|
||||
run: Any,
|
||||
execute: Callable[[], Awaitable[None]],
|
||||
|
||||
@@ -53,6 +53,11 @@ from agno.run.agent import (
|
||||
RunOutput,
|
||||
RunOutputEvent,
|
||||
)
|
||||
|
||||
# Strong references to background tasks so they aren't garbage-collected mid-execution.
|
||||
# See: https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
|
||||
# One registry for every detached run, shared with agents, so AgentOS shutdown drains one set.
|
||||
from agno.run.background import _background_tasks # noqa: E402
|
||||
from agno.run.cancel import (
|
||||
acancel_run as acancel_run_global,
|
||||
)
|
||||
@@ -128,11 +133,6 @@ from agno.utils.log import (
|
||||
log_warning,
|
||||
)
|
||||
|
||||
# Strong references to background tasks so they aren't garbage-collected mid-execution.
|
||||
# See: https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
|
||||
# One registry for every detached run, shared with agents, so AgentOS shutdown drains one set.
|
||||
from agno.run.background import _background_tasks # noqa: E402
|
||||
|
||||
# Cancel raises immediately on every event. Only terminal events bypass so the
|
||||
# member's own cancel handler can yield them to the stream. RunError is excluded —
|
||||
# in a cancel race the team's cancel wins.
|
||||
|
||||
@@ -7,10 +7,10 @@ import pytest
|
||||
|
||||
from agno.agent.agent import Agent
|
||||
from agno.models.openai import OpenAIChat
|
||||
from agno.run import RunContext
|
||||
from agno.run.agent import RunContentEvent
|
||||
from agno.run.base import RunStatus
|
||||
from agno.run.cancel import cancel_run
|
||||
from agno.run import RunContext
|
||||
from agno.run.workflow import StepProgressEvent, WorkflowCancelledEvent
|
||||
from agno.workflow import HumanReview, Step, Workflow
|
||||
from agno.workflow.condition import Condition
|
||||
|
||||
@@ -482,3 +482,593 @@ def test_user_session_store_and_file_checkpointing_disable_injection(fake_sdk, t
|
||||
checkpointing._build_options(agno_session_id="agno-session")
|
||||
assert len(warnings) == 1
|
||||
assert "enable_file_checkpointing" in warnings[0]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Continue from a message boundary
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _msg(message: Any, uuid: str) -> Any:
|
||||
message.uuid = uuid
|
||||
return message
|
||||
|
||||
|
||||
def _user(uuid: str, content: Any) -> Any:
|
||||
user = UserMessage()
|
||||
user.content = content
|
||||
return _msg(user, uuid)
|
||||
|
||||
|
||||
def _tool_turn(prompt: str, sdk_session_id: str = "sdk-1") -> List[Any]:
|
||||
tool = ToolUseBlock()
|
||||
tool.id, tool.name, tool.input = "call-1", "Bash", {"command": "echo alpha"}
|
||||
result = ToolResultBlock()
|
||||
result.tool_use_id, result.content = "call-1", "alpha"
|
||||
return [
|
||||
SystemMessage("init", {"session_id": sdk_session_id}),
|
||||
_user("u-prompt", prompt),
|
||||
_msg(AssistantMessage([tool]), "u-call"),
|
||||
_user("u-result", [result]),
|
||||
_msg(AssistantMessage([TextBlock("done")]), "u-final"),
|
||||
ResultMessage(sdk_session_id, "done"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def scripted(fake_sdk, monkeypatch):
|
||||
"""Serve one scripted SDK turn per query and record prompts, resumes and forks."""
|
||||
turns: List[List[Any]] = []
|
||||
calls: List[Dict[str, Any]] = []
|
||||
forks: List[Dict[str, Any]] = []
|
||||
|
||||
async def query(prompt, options):
|
||||
calls.append({"prompt": prompt, "resume": options.resume, "extra_args": options.extra.get("extra_args")})
|
||||
for message in turns.pop(0):
|
||||
yield message
|
||||
|
||||
async def fork_session_via_store(store, session_id, directory=None, up_to_message_id=None):
|
||||
forks.append({"session_id": session_id, "up_to": up_to_message_id})
|
||||
return type("Fork", (), {"session_id": f"fork-{len(forks)}"})()
|
||||
|
||||
module = claude_module._sdk()
|
||||
monkeypatch.setattr(module, "query", query)
|
||||
module.fork_session_via_store = fork_session_via_store
|
||||
return turns, calls, forks
|
||||
|
||||
|
||||
def _ref(message: Any) -> Optional[Dict[str, Any]]:
|
||||
return (message.provider_data or {}).get("claude_sdk")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
async def test_runs_record_transcript_positions_and_tool_checkpoints(scripted, tmp_db, stream):
|
||||
turns, calls, _ = scripted
|
||||
turns.append(_tool_turn("go"))
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
if stream:
|
||||
events = [event async for event in agent.arun("go", session_id="s", stream=True)]
|
||||
assert sum(type(event).__name__ == "ToolCallCompletedEvent" for event in events) == 1
|
||||
run = (await agent.aget_session("s")).runs[0]
|
||||
else:
|
||||
run = await agent.arun("go", session_id="s")
|
||||
assert calls[0]["extra_args"] == {"replay-user-messages": None}
|
||||
assert [(m.role, (_ref(m) or {}).get("uuid")) for m in run.messages] == [
|
||||
("user", "u-prompt"),
|
||||
("assistant", "u-call"),
|
||||
("tool", "u-result"),
|
||||
("assistant", "u-final"),
|
||||
]
|
||||
assert {_ref(m)["session_id"] for m in run.messages} == {"sdk-1"}
|
||||
assert [m.checkpoint_status for m in run.messages] == [None, None, "RUNNING", None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fork_from_tool_result_branches_the_sdk_transcript(scripted, tmp_db):
|
||||
turns, calls, forks = scripted
|
||||
turns.append(_tool_turn("go"))
|
||||
turns.append(
|
||||
[
|
||||
SystemMessage("init", {"session_id": "fork-1"}),
|
||||
_user("v-prompt", "next"),
|
||||
_msg(AssistantMessage([TextBlock("branched")]), "v-final"),
|
||||
ResultMessage("fork-1", "branched"),
|
||||
]
|
||||
)
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
source = await agent.arun("go", session_id="s")
|
||||
branch = await agent.acontinue_run(run_id=source.run_id, session_id="s", continue_from=3, fork=True, input="next")
|
||||
assert forks == [{"session_id": "sdk-1", "up_to": "u-result"}]
|
||||
assert calls[1]["resume"] == "fork-1" and calls[1]["prompt"] == "next"
|
||||
assert branch.run_id != source.run_id
|
||||
assert (branch.forked_from_run_id, branch.forked_from_message_index) == (source.run_id, 3)
|
||||
assert [m.role for m in branch.messages] == ["user", "assistant", "tool", "user", "assistant"]
|
||||
assert _ref(branch.messages[3]) == {"session_id": "fork-1", "uuid": "v-prompt", "position": 0}
|
||||
assert [t.tool_call_id for t in branch.tools] == ["call-1"]
|
||||
assert len((await agent.aget_run_output(source.run_id, "s")).messages) == 4
|
||||
assert len((await agent.aget_session("s")).runs) == 2
|
||||
assert (await agent.aget_session("s")).session_data["claude_sdk_session_id"] == "fork-1"
|
||||
|
||||
|
||||
def test_continuing_a_finished_run_without_fork_still_forks(scripted, tmp_db):
|
||||
"""Like native agents: a completed run is never rewritten in place, whatever fork is set to."""
|
||||
turns, calls, forks = scripted
|
||||
turns.append(_tool_turn("go"))
|
||||
turns.append(
|
||||
[
|
||||
SystemMessage("init", {"session_id": "fork-1"}),
|
||||
_user("v-prompt", "next"),
|
||||
_msg(AssistantMessage([TextBlock("branched")]), "v-final"),
|
||||
ResultMessage("fork-1", "branched"),
|
||||
]
|
||||
)
|
||||
turns.append([SystemMessage("init", {"session_id": "fork-2"}), ResultMessage("fork-2", "more")])
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
source = agent.run("go", session_id="s")
|
||||
branch = agent.continue_run(run_id=source.run_id, session_id="s", continue_from=3, fork=True, input="next")
|
||||
continued = agent.continue_run(run_id=branch.run_id, session_id="s", fork=False)
|
||||
assert calls[2]["prompt"] == ClaudeAgent._CONTINUE_PROMPT
|
||||
assert forks[-1] == {"session_id": "fork-1", "up_to": "v-final"}
|
||||
assert continued.run_id != branch.run_id
|
||||
assert continued.forked_from_run_id == branch.run_id
|
||||
assert continued.forked_from_message_index == len(branch.messages)
|
||||
assert continued.content == "more"
|
||||
assert len(agent.get_session("s").runs) == 3
|
||||
assert agent.get_run_output(branch.run_id, "s").content == "branched", "the source run is untouched"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("parent", [None, "u-previous"])
|
||||
async def test_replaying_a_turn_resends_its_prompt_from_before_it(scripted, tmp_db, parent):
|
||||
turns, calls, forks = scripted
|
||||
turns.append(_tool_turn("go"))
|
||||
turns.append(_tool_turn("go", sdk_session_id="sdk-2"))
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
source = await agent.arun("go", session_id="s")
|
||||
tmp_db.append_transcript_entries(
|
||||
framework="claude-agent-sdk",
|
||||
project_key=agent.project_key,
|
||||
session_id="sdk-1",
|
||||
entries=[{"type": "user", "uuid": "u-prompt", "parentUuid": parent}],
|
||||
agno_session_id="s",
|
||||
)
|
||||
replay = await agent.acontinue_run(run_id=source.run_id, session_id="s", continue_from="last_user", fork=True)
|
||||
assert calls[1]["prompt"] == "go"
|
||||
if parent is None:
|
||||
assert forks == [] and calls[1]["resume"] is None
|
||||
else:
|
||||
assert forks == [{"session_id": "sdk-1", "up_to": "u-previous"}] and calls[1]["resume"] == "fork-1"
|
||||
assert [m.role for m in replay.messages] == ["user", "assistant", "tool", "assistant"]
|
||||
assert replay.messages[0].content == "go"
|
||||
assert replay.forked_from_message_index == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_continue_rejects_unrecorded_runs_and_does_not_fall_back(fake_sdk, scripted, tmp_db):
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
turns, calls, forks = scripted
|
||||
turns.append([SystemMessage("init", {"session_id": "sdk-1"}), ResultMessage("sdk-1", "untracked")])
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
untracked = await agent.arun("go", session_id="s")
|
||||
with pytest.raises(ValueError, match="no Claude SDK transcript position"):
|
||||
await agent.acontinue_run(run_id=untracked.run_id, session_id="s", input="next")
|
||||
with pytest.raises(ValueError, match="regenerate"):
|
||||
agent.acontinue_run(run_id=untracked.run_id, session_id="s", regenerate=True)
|
||||
|
||||
turns.append(_tool_turn("go"))
|
||||
source = await agent.arun("go", session_id="s")
|
||||
|
||||
async def failing_query(prompt, options):
|
||||
calls.append({"prompt": prompt, "resume": options.resume})
|
||||
raise ProcessError("Command failed with exit code 1")
|
||||
yield # pragma: no cover
|
||||
|
||||
monkeypatch_target = claude_module._sdk()
|
||||
original = monkeypatch_target.query
|
||||
monkeypatch_target.query = failing_query
|
||||
try:
|
||||
failed = await agent.acontinue_run(run_id=source.run_id, session_id="s", continue_from=3, input="next")
|
||||
finally:
|
||||
monkeypatch_target.query = original
|
||||
assert failed.status == RunStatus.error
|
||||
assert [call["resume"] for call in calls[-1:]] == ["fork-1"]
|
||||
assert len([call for call in calls if call.get("prompt") == "next"]) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_prompt_echo_does_not_block_resume_fallback(fake_sdk, tmp_path, monkeypatch):
|
||||
agent = ClaudeAgent(db=SqliteDb(db_file=str(tmp_path / "db")))
|
||||
await agent.arun("Remember the code bluejay", session_id="session")
|
||||
calls = []
|
||||
|
||||
async def query(prompt, options):
|
||||
calls.append(options.resume)
|
||||
if options.resume:
|
||||
yield _user("echo", prompt)
|
||||
yield ResultMessage(options.resume, "No conversation found", is_error=True)
|
||||
else:
|
||||
yield ResultMessage("new", "bluejay")
|
||||
|
||||
monkeypatch.setattr(claude_module._sdk(), "query", query)
|
||||
run = await agent.arun("What code?", session_id="session")
|
||||
assert run.content == "bluejay"
|
||||
assert calls == ["sdk-1", None]
|
||||
|
||||
|
||||
def _recorded(role: str, uuid: str, position: int, content: Any = None, **kwargs: Any) -> Any:
|
||||
from agno.models.message import Message
|
||||
|
||||
return Message(
|
||||
role=role,
|
||||
content=content,
|
||||
provider_data={"claude_sdk": {"session_id": "sdk-1", "uuid": uuid, "position": position}},
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _recorded_run(messages: List[Any], status: Any = None) -> Any:
|
||||
from agno.models.response import ToolExecution
|
||||
from agno.run.agent import RunInput, RunOutput
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
return RunOutput(
|
||||
run_id="source",
|
||||
session_id="s",
|
||||
status=status or RunStatus.completed,
|
||||
input=RunInput(input_content="go"),
|
||||
messages=messages,
|
||||
tools=[ToolExecution(tool_call_id=m.tool_call_id, tool_name="Bash") for m in messages if m.role == "tool"],
|
||||
)
|
||||
|
||||
|
||||
def test_cancelled_runs_cannot_be_continued(fake_sdk, tmp_db):
|
||||
from agno.exceptions import RunNotContinuableError
|
||||
from agno.run.base import RunStatus
|
||||
|
||||
run = _recorded_run([_recorded("user", "u-prompt", 0, "go")], status=RunStatus.cancelled)
|
||||
with pytest.raises(RunNotContinuableError, match="cancelled"):
|
||||
ClaudeAgent(db=tmp_db)._build_continuation(run, continue_from="end", fork=True, input="next")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_continue_from_zero_starts_a_branch_before_the_prompt(scripted, tmp_db):
|
||||
turns, calls, forks = scripted
|
||||
turns.append(_tool_turn("go"))
|
||||
turns.append(
|
||||
[SystemMessage("init", {"session_id": "sdk-2"}), _user("w-prompt", "fresh"), ResultMessage("sdk-2", "ok")]
|
||||
)
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
source = await agent.arun("go", session_id="s")
|
||||
tmp_db.append_transcript_entries(
|
||||
framework="claude-agent-sdk",
|
||||
project_key=agent.project_key,
|
||||
session_id="sdk-1",
|
||||
entries=[{"type": "user", "uuid": "u-prompt", "parentUuid": None}],
|
||||
agno_session_id="s",
|
||||
)
|
||||
with pytest.raises(ValueError, match="continue_from=0"):
|
||||
await agent.acontinue_run(run_id=source.run_id, session_id="s", continue_from=0, fork=True)
|
||||
branch = await agent.acontinue_run(run_id=source.run_id, session_id="s", continue_from=0, fork=True, input="fresh")
|
||||
assert forks == [] and calls[1]["resume"] is None and calls[1]["prompt"] == "fresh"
|
||||
assert [m.role for m in branch.messages] == ["user", "assistant"] and branch.messages[0].content == "fresh"
|
||||
assert branch.forked_from_message_index == 0
|
||||
|
||||
|
||||
def test_last_user_replays_the_selected_user_turn_not_only_the_first(fake_sdk, tmp_db):
|
||||
messages = [
|
||||
_recorded("user", "u1", 0, "first"),
|
||||
_recorded("assistant", "a1", 1, "reply one"),
|
||||
_recorded("user", "u2", 2, "second"),
|
||||
_recorded("assistant", "a2", 3, "reply two"),
|
||||
]
|
||||
prompt, _, continuation = ClaudeAgent(db=tmp_db)._build_continuation(
|
||||
_recorded_run(messages), continue_from="last_user", fork=True, input=None
|
||||
)
|
||||
assert prompt == "second"
|
||||
assert [m.content for m in continuation.messages] == ["first", "reply one"]
|
||||
assert continuation.anchor == {"session_id": "sdk-1", "uuid": "u2", "before": True}
|
||||
assert continuation.record_input == "second"
|
||||
|
||||
|
||||
def test_parallel_tool_batch_has_one_checkpoint_and_forks_whole(fake_sdk, tmp_db):
|
||||
# Transcript order: prompt 0, call1 1, call2 2, result2 3, result1 4, final 5; Agno lists result1 first.
|
||||
messages = [
|
||||
_recorded("user", "u", 0, "go"),
|
||||
_recorded(
|
||||
"assistant",
|
||||
"c1",
|
||||
1,
|
||||
tool_calls=[{"id": "call-1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
|
||||
),
|
||||
_recorded("tool", "r1", 4, "alpha", tool_call_id="call-1"),
|
||||
_recorded(
|
||||
"assistant",
|
||||
"c2",
|
||||
2,
|
||||
tool_calls=[{"id": "call-2", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
|
||||
),
|
||||
_recorded("tool", "r2", 3, "beta", tool_call_id="call-2"),
|
||||
_recorded("assistant", "f", 5, "done"),
|
||||
]
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
assert agent._checkpoint_indexes(messages) == [4], "only the end of the batch is a checkpoint"
|
||||
_, _, continuation = agent._build_continuation(_recorded_run(messages), continue_from=3, fork=True, input="next")
|
||||
assert continuation.forked_from_message_index == 5, "a boundary inside the batch moves to its end"
|
||||
assert continuation.anchor["uuid"] == "r1", "the fork keeps up to the transcript-latest result of the batch"
|
||||
assert [m.role for m in continuation.messages] == ["user", "assistant", "tool", "assistant", "tool"]
|
||||
assert sorted(t.tool_call_id for t in continuation.tools) == ["call-1", "call-2"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finished_run_continues_in_the_background_as_a_fork(scripted, tmp_db):
|
||||
import asyncio
|
||||
|
||||
turns, _, _ = scripted
|
||||
turns.append(_tool_turn("go"))
|
||||
turns.append([SystemMessage("init", {"session_id": "fork-1"}), ResultMessage("fork-1", "more")])
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
source = await agent.arun("go", session_id="s")
|
||||
accepted = await agent.acontinue_run(run_response=source, fork=False, background=True, input="more")
|
||||
assert accepted.run_id != source.run_id, "background continuation of a finished run must be a new run"
|
||||
for _ in range(100):
|
||||
stored = await agent.aget_run_output(accepted.run_id, "s")
|
||||
if stored is not None and getattr(stored.status, "value", stored.status) == "COMPLETED":
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
assert stored is not None and stored.content == "more"
|
||||
assert stored.forked_from_run_id == source.run_id
|
||||
assert (await agent.aget_run_output(source.run_id, "s")).content == "done", "the source run is untouched"
|
||||
|
||||
|
||||
def test_sequential_tool_calls_keep_a_checkpoint_per_step(fake_sdk, tmp_db):
|
||||
# Transcript order follows Agno order: each call is issued after the previous result arrived.
|
||||
messages = [
|
||||
_recorded("user", "u", 0, "go"),
|
||||
_recorded(
|
||||
"assistant",
|
||||
"c1",
|
||||
1,
|
||||
tool_calls=[{"id": "call-1", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
|
||||
),
|
||||
_recorded("tool", "r1", 2, "alpha", tool_call_id="call-1"),
|
||||
_recorded(
|
||||
"assistant",
|
||||
"c2",
|
||||
3,
|
||||
tool_calls=[{"id": "call-2", "type": "function", "function": {"name": "Bash", "arguments": "{}"}}],
|
||||
),
|
||||
_recorded("tool", "r2", 4, "beta", tool_call_id="call-2"),
|
||||
_recorded("assistant", "f", 5, "done"),
|
||||
]
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
assert agent._checkpoint_indexes(messages) == [2, 4]
|
||||
_, _, continuation = agent._build_continuation(_recorded_run(messages), continue_from=3, fork=True, input="next")
|
||||
assert continuation.forked_from_message_index == 3 and continuation.anchor["uuid"] == "r1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
async def test_resume_fallback_after_fork_only_replays_active_branch(scripted, tmp_db, monkeypatch, stream):
|
||||
turns, _, _ = scripted
|
||||
source_turn = _tool_turn("go")
|
||||
source_turn[-2] = _msg(AssistantMessage([TextBlock("discarded reply")]), "u-final")
|
||||
source_turn[-1] = ResultMessage("sdk-1", "discarded reply")
|
||||
turns.append(source_turn)
|
||||
turns.append(
|
||||
[
|
||||
SystemMessage("init", {"session_id": "fork-1"}),
|
||||
_user("v-prompt", "next"),
|
||||
_msg(AssistantMessage([TextBlock("branch reply")]), "v-final"),
|
||||
ResultMessage("fork-1", "branch reply"),
|
||||
]
|
||||
)
|
||||
agent = ClaudeAgent(id="branch-history", db=tmp_db)
|
||||
source = await agent.arun("go", session_id="s")
|
||||
await agent.acontinue_run(run_response=source, continue_from=3, input="next")
|
||||
prompts = []
|
||||
|
||||
async def missing_session(prompt, options):
|
||||
prompts.append(prompt)
|
||||
if options.resume:
|
||||
raise ProcessError("No conversation found")
|
||||
yield SystemMessage("init", {"session_id": "fresh"})
|
||||
yield AssistantMessage([TextBlock("recovered")])
|
||||
yield ResultMessage("fresh", "recovered")
|
||||
|
||||
monkeypatch.setattr(claude_module._sdk(), "query", missing_session)
|
||||
# A fresh adapter verifies that branch selection survives a process restart.
|
||||
restarted = ClaudeAgent(id="branch-history", db=tmp_db)
|
||||
if stream:
|
||||
events = [event async for event in restarted.arun("remember?", session_id="s", stream=True)]
|
||||
assert isinstance(events[-1], RunCompletedEvent)
|
||||
else:
|
||||
result = await restarted.arun("remember?", session_id="s")
|
||||
assert result.content == "recovered"
|
||||
assert len(prompts) == 2
|
||||
assert "discarded reply" not in prompts[-1]
|
||||
assert "branch reply" in prompts[-1]
|
||||
assert prompts[-1].count("user: go") == 1
|
||||
assert prompts[-1].count("assistant called Bash") == 1
|
||||
assert prompts[-1].count("tool result: alpha") == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
async def test_agentos_automatic_fork_does_not_change_source_stream(scripted, tmp_db, monkeypatch, stream):
|
||||
from httpx import ASGITransport, AsyncClient
|
||||
|
||||
import agno.os.event_streams as event_streams
|
||||
from agno.os import AgentOS
|
||||
from agno.os.event_streams import InMemoryEventStream
|
||||
from agno.os.managers import EventsBuffer, SSESubscriberManager
|
||||
|
||||
event_stream = InMemoryEventStream(events_buffer=EventsBuffer(), subscriber_manager=SSESubscriberManager())
|
||||
monkeypatch.setattr(event_streams, "_event_stream", event_stream)
|
||||
turns, _, _ = scripted
|
||||
turns.append(_tool_turn("go"))
|
||||
turns.append(
|
||||
[
|
||||
SystemMessage("init", {"session_id": "fork-1"}),
|
||||
_user("v-prompt", "next"),
|
||||
ResultMessage("fork-1", "branch failed", is_error=True),
|
||||
]
|
||||
)
|
||||
agent = ClaudeAgent(id="stream-identity", db=tmp_db)
|
||||
source = await agent.arun("go", session_id="s")
|
||||
await event_stream.register_run(source.run_id)
|
||||
await event_stream.add_event(source.run_id, RunCompletedEvent(run_id=source.run_id))
|
||||
await event_stream.complete_run(source.run_id, RunStatus.completed)
|
||||
original_events = await event_stream.replay(source.run_id)
|
||||
app = AgentOS(agents=[agent], db=tmp_db).get_app()
|
||||
async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
|
||||
response = await client.post(
|
||||
f"/agents/{agent.id}/runs/{source.run_id}/continue",
|
||||
data={"session_id": "s", "input": "next", "stream": str(stream).lower()},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
branch = (await agent.aget_session("s")).runs[-1]
|
||||
assert branch.run_id != source.run_id and branch.status == RunStatus.error
|
||||
assert await event_stream.get_run_status(source.run_id) == RunStatus.completed
|
||||
assert await event_stream.replay(source.run_id) == original_events
|
||||
if stream:
|
||||
branch_events = await event_stream.replay(branch.run_id)
|
||||
assert branch_events and all(event.run_id == branch.run_id for _, event in branch_events)
|
||||
assert await event_stream.get_run_status(branch.run_id) == RunStatus.error
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
@pytest.mark.parametrize("failure", ["result", "exception"])
|
||||
async def test_failed_run_retains_completed_tool_checkpoint(scripted, tmp_db, monkeypatch, stream, failure):
|
||||
from agno.os.checkpoints import list_run_checkpoints
|
||||
|
||||
turns, _, forks = scripted
|
||||
failed_turn = _tool_turn("go")[:4]
|
||||
if failure == "result":
|
||||
turns.append(failed_turn + [ResultMessage("sdk-1", "provider failed", is_error=True)])
|
||||
else:
|
||||
original_query = claude_module._sdk().query
|
||||
|
||||
async def query(prompt, options):
|
||||
if options.resume is None:
|
||||
for message in failed_turn:
|
||||
yield message
|
||||
raise ProcessError("provider failed")
|
||||
async for message in original_query(prompt, options):
|
||||
yield message
|
||||
|
||||
monkeypatch.setattr(claude_module._sdk(), "query", query)
|
||||
turns.append(
|
||||
[
|
||||
SystemMessage("init", {"session_id": "fork-1"}),
|
||||
_user("v-prompt", ClaudeAgent._CONTINUE_PROMPT),
|
||||
_msg(AssistantMessage([TextBlock("recovered")]), "v-final"),
|
||||
ResultMessage("fork-1", "recovered"),
|
||||
]
|
||||
)
|
||||
agent = ClaudeAgent(id="failed-checkpoint", db=tmp_db)
|
||||
if stream:
|
||||
events = [event async for event in agent.arun("go", session_id="s", stream=True)]
|
||||
assert isinstance(events[-1], RunErrorEvent)
|
||||
failed = (await agent.aget_session("s")).runs[0]
|
||||
else:
|
||||
failed = await agent.arun("go", session_id="s")
|
||||
assert failed.status == RunStatus.error and "provider failed" in failed.content
|
||||
assert [message.role for message in failed.messages] == ["user", "assistant", "tool"]
|
||||
assert failed.tools[0].result == "alpha"
|
||||
assert _ref(failed.messages[-1])["uuid"] == "u-result"
|
||||
assert failed.messages[-1].checkpoint_status == RunStatus.running.value
|
||||
assert [checkpoint["message_index"] for checkpoint in list_run_checkpoints(failed)] == [3]
|
||||
# Load from the database, so the assertion also checks serialized recovery.
|
||||
continued = await agent.acontinue_run(run_id=failed.run_id, session_id="s", fork=True)
|
||||
assert forks[-1] == {"session_id": "sdk-1", "up_to": "u-result"}
|
||||
assert continued.status == RunStatus.completed
|
||||
assert continued.tools[0].result == "alpha"
|
||||
assert [tool.tool_call_id for tool in continued.tools] == ["call-1"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
async def test_failed_run_retains_real_trailing_assistant_text(scripted, tmp_db, stream):
|
||||
turns, _, _ = scripted
|
||||
turns.append(_tool_turn("go")[:-1] + [ResultMessage("sdk-1", "provider failed", is_error=True)])
|
||||
agent = ClaudeAgent(id="failed-text", db=tmp_db)
|
||||
if stream:
|
||||
_ = [event async for event in agent.arun("go", session_id="s", stream=True)]
|
||||
failed = (await agent.aget_session("s")).runs[0]
|
||||
else:
|
||||
failed = await agent.arun("go", session_id="s")
|
||||
assert failed.messages[-1].content == "done"
|
||||
assert "provider failed" in failed.content
|
||||
_, _, continuation = agent._build_continuation(failed, continue_from="end", fork=True, input="next")
|
||||
assert continuation.anchor["uuid"] == "u-final"
|
||||
assert continuation.tools[0].result == "alpha"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("stream", [False, True])
|
||||
async def test_subagent_tools_are_not_exposed_as_checkpoints(scripted, tmp_db, stream):
|
||||
from agno.os.checkpoints import list_run_checkpoints
|
||||
|
||||
turns, _, _ = scripted
|
||||
parent = _tool_turn("go")
|
||||
parent[2].content[0].id = "parent-call"
|
||||
parent[2].content[0].name = "Agent"
|
||||
parent[3].content[0].tool_use_id = "parent-call"
|
||||
child = _tool_turn("child")
|
||||
for message in child[2:4]:
|
||||
message.parent_tool_use_id = "parent-call"
|
||||
turns.append([parent[0], parent[1], parent[2], child[2], child[3], *parent[3:]])
|
||||
agent = ClaudeAgent(id="subagent-checkpoints", db=tmp_db)
|
||||
if stream:
|
||||
_ = [event async for event in agent.arun("go", session_id="s", stream=True)]
|
||||
run = (await agent.aget_session("s")).runs[0]
|
||||
else:
|
||||
run = await agent.arun("go", session_id="s")
|
||||
child_result = next(message for message in run.messages if message.tool_call_id == "call-1")
|
||||
assert child_result.checkpoint_status is None
|
||||
assert _ref(child_result) is None
|
||||
checkpoints = list_run_checkpoints(run)
|
||||
assert [checkpoint["message_index"] for checkpoint in checkpoints] == [3, 6]
|
||||
for checkpoint in checkpoints:
|
||||
_, _, continuation = agent._build_continuation(
|
||||
run, continue_from=checkpoint["message_index"], fork=True, input="next"
|
||||
)
|
||||
assert continuation.anchor["uuid"] in ("u-result", "u-final")
|
||||
|
||||
|
||||
def test_incomplete_parallel_batch_is_not_a_checkpoint(fake_sdk, tmp_db):
|
||||
from agno.models.message import Message
|
||||
|
||||
messages = [
|
||||
_recorded("user", "u", 0, "go"),
|
||||
_recorded("assistant", "c1", 1, tool_calls=[{"id": "call-1"}]),
|
||||
_recorded("tool", "r1", 3, "completed", tool_call_id="call-1"),
|
||||
_recorded("assistant", "c2", 2, tool_calls=[{"id": "call-2"}]),
|
||||
Message(role="tool", tool_call_id="call-2", content=""),
|
||||
]
|
||||
agent = ClaudeAgent(db=tmp_db)
|
||||
assert agent._checkpoint_indexes(messages) == []
|
||||
with pytest.raises(ValueError, match="no Claude SDK transcript position"):
|
||||
agent._build_continuation(_recorded_run(messages), continue_from=3, fork=True, input="next")
|
||||
|
||||
|
||||
def test_sync_continue_failed_run_preserves_tool_results(scripted, tmp_db):
|
||||
turns, _, forks = scripted
|
||||
turns.append(_tool_turn("go")[:4] + [ResultMessage("sdk-1", "provider failed", is_error=True)])
|
||||
turns.append(
|
||||
[
|
||||
SystemMessage("init", {"session_id": "fork-1"}),
|
||||
_msg(AssistantMessage([TextBlock("recovered")]), "v-final"),
|
||||
ResultMessage("fork-1", "recovered"),
|
||||
]
|
||||
)
|
||||
agent = ClaudeAgent(id="sync-recovery", db=tmp_db)
|
||||
failed = agent.run("go", session_id="s")
|
||||
recovered = agent.continue_run(run_id=failed.run_id, session_id="s")
|
||||
assert recovered.run_id == failed.run_id
|
||||
assert recovered.status == RunStatus.completed
|
||||
assert recovered.tools[0].result == "alpha"
|
||||
assert forks == [{"session_id": "sdk-1", "up_to": "u-result"}]
|
||||
|
||||
@@ -87,3 +87,36 @@ def test_run_indexes_are_sequential(agent):
|
||||
asyncio.run(agent._arun_non_stream(text, session_id="s1", user_id="u1"))
|
||||
rows, _ = agent.db.get_runs(session_id="s1", deserialize=False)
|
||||
assert sorted(r["run_index"] for r in rows) == [0, 1, 2]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("source", ["source", "branch", "missing"])
|
||||
def test_history_follows_fork_ancestry_and_subsequent_turns(agent, source):
|
||||
from agno.models.message import Message
|
||||
from agno.run.agent import RunOutput
|
||||
from agno.session.agent import AgentSession
|
||||
|
||||
def run(run_id, content, forked_from=None):
|
||||
return RunOutput(
|
||||
run_id=run_id,
|
||||
messages=[Message(role="user", content=content)],
|
||||
forked_from_run_id=forked_from,
|
||||
)
|
||||
|
||||
session = AgentSession(
|
||||
session_id="s",
|
||||
runs=[
|
||||
run("earlier", "earlier turn"),
|
||||
run("source", "original turn"),
|
||||
run("discarded", "discarded later turn"),
|
||||
run("branch", "first branch", "source"),
|
||||
run("later", "first branch follow-up"),
|
||||
run("second-branch", "second branch", source),
|
||||
run("follow-up", "active follow-up"),
|
||||
run("current", "exclude current input"),
|
||||
],
|
||||
)
|
||||
history = agent._get_history_from_session(session, exclude_run_id="current")
|
||||
expected = ["second branch", "active follow-up"]
|
||||
if source != "missing":
|
||||
expected.insert(0, "earlier turn")
|
||||
assert [message["content"] for message in history] == expected
|
||||
|
||||
@@ -2,11 +2,10 @@
|
||||
|
||||
import sqlite3
|
||||
|
||||
import pytest
|
||||
from sqlalchemy import inspect as sa_inspect
|
||||
|
||||
from agno.db.sqlite import SqliteDb
|
||||
import pytest
|
||||
|
||||
from agno.fs import FileSystem, InvalidPathError, SchemaOutdatedError, UnsupportedOperationError
|
||||
from agno.fs.db import DbFileSystem
|
||||
|
||||
|
||||
@@ -351,3 +351,39 @@ class TestInlineContinueSettlesTicket:
|
||||
worker, store = await self._paused_ticket_worker("other")
|
||||
await asettle_paused_ticket(worker, "never-queued", RunStatus.completed) # no ticket: no-op
|
||||
assert (await store.get_job("other"))["status"] == "paused"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("fork,regenerate", [(False, False), (True, False), (False, True)])
|
||||
async def test_agent_fork_stream_uses_emitted_run_id(stream, fork, regenerate):
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from agno.os.routers.agents.router import agent_continue_response_streamer
|
||||
from agno.run.agent import RunCompletedEvent, RunStartedEvent
|
||||
|
||||
await _park_paused(stream, "source")
|
||||
await stream.complete_run("source", RunStatus.completed)
|
||||
original_events = await stream.replay("source")
|
||||
final_run = SimpleNamespace(run_id="branch", status=RunStatus.completed)
|
||||
agent = FakeAgent([RunStartedEvent(run_id="branch"), RunCompletedEvent(run_id="branch")], final_run)
|
||||
queue_worker = SimpleNamespace(store=AsyncMock())
|
||||
frames = [
|
||||
frame
|
||||
async for frame in agent_continue_response_streamer(
|
||||
agent,
|
||||
run_id="source",
|
||||
session_id="s",
|
||||
fork=fork,
|
||||
regenerate=regenerate,
|
||||
queue_worker=queue_worker,
|
||||
)
|
||||
]
|
||||
assert len(frames) == 2
|
||||
assert await stream.replay("source") == original_events
|
||||
assert await stream.get_run_status("source") == RunStatus.completed
|
||||
assert await stream.get_run_status("branch") == RunStatus.completed
|
||||
branch_events = await stream.replay("branch")
|
||||
assert [event.event for _, event in branch_events] == ["RunStarted", "RunCompleted"]
|
||||
assert all(event.run_id == "branch" for _, event in branch_events)
|
||||
assert not queue_worker.store.mock_calls, "a fork must not settle its source's queue ticket"
|
||||
|
||||
@@ -154,3 +154,34 @@ async def test_clean_completion_still_finalizes(stream_harness):
|
||||
|
||||
assert await stream_harness.get_run_status(RUN_ID) == RunStatus.completed
|
||||
assert (await store.get_job(RUN_ID))["status"] == "completed"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_disconnect_from_automatic_fork_only_closes_branch_stream(stream_harness):
|
||||
from agno.os.routers.agents.router import agent_continue_response_streamer
|
||||
from agno.run.agent import RunStartedEvent
|
||||
|
||||
class ForkAgent(ContinueFakeAgent):
|
||||
def acontinue_run(self, **kwargs):
|
||||
async def gen():
|
||||
yield RunStartedEvent(run_id="branch")
|
||||
await self._never.wait()
|
||||
|
||||
return gen()
|
||||
|
||||
await stream_harness.register_run(RUN_ID)
|
||||
await stream_harness.complete_run(RUN_ID, RunStatus.completed)
|
||||
first_frame = asyncio.Event()
|
||||
|
||||
async def consume():
|
||||
async for _ in agent_continue_response_streamer(ForkAgent(), RUN_ID, session_id=SESSION_ID):
|
||||
first_frame.set()
|
||||
|
||||
task = asyncio.create_task(consume())
|
||||
await asyncio.wait_for(first_frame.wait(), timeout=1)
|
||||
assert await stream_harness.get_run_status("branch") == RunStatus.running
|
||||
task.cancel()
|
||||
await asyncio.gather(task, return_exceptions=True)
|
||||
assert await stream_harness.get_run_status("branch") == RunStatus.cancelled
|
||||
assert await stream_harness.get_run_status(RUN_ID) == RunStatus.completed
|
||||
assert await stream_harness.replay(RUN_ID) == []
|
||||
|
||||
@@ -0,0 +1,43 @@
|
||||
"""Unit tests for the shared background run helpers."""
|
||||
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
from agno.run import background
|
||||
from agno.run.background import _spawn_background, await_background_runs
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_await_background_runs_waits_for_spawned_tasks():
|
||||
finished = []
|
||||
|
||||
async def work():
|
||||
await asyncio.sleep(0.05)
|
||||
finished.append("done")
|
||||
|
||||
_spawn_background(work())
|
||||
_spawn_background(work())
|
||||
assert finished == []
|
||||
await await_background_runs()
|
||||
assert finished == ["done", "done"]
|
||||
assert not [task for task in background._background_tasks if not task.done()]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_await_background_runs_returns_when_nothing_is_pending():
|
||||
await await_background_runs()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_await_background_runs_honours_timeout():
|
||||
release = asyncio.Event()
|
||||
|
||||
async def work():
|
||||
await release.wait()
|
||||
|
||||
_spawn_background(work())
|
||||
await await_background_runs(timeout=0.05)
|
||||
assert [task for task in background._background_tasks if not task.done()], "task should still be running"
|
||||
release.set()
|
||||
await await_background_runs()
|
||||
Reference in New Issue
Block a user