From 78fa25982ef3732f6cb7f374c0e3c8ec55520363 Mon Sep 17 00:00:00 2001 From: Yash Pratap Solanky <101447028+ysolanky@users.noreply.github.com> Date: Fri, 9 Oct 2026 01:56:38 -0400 Subject: [PATCH] fix: persist ClaudeAgent SDK session id across restarts (#10903) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary `ClaudeAgent` kept the Agno `session_id` -> Claude SDK session id mapping only in memory (`self._sdk_session_ids`). After a process restart, or on a different AgentOS worker, the same Agno session silently started a fresh Claude conversation. The AgentOS UI still showed the full history from the DB, so it looked as if the session had continued. - **Persist the mapping:** the SDK session id is now stored in `session.session_data["claude_sdk_session_id"]`, which the base class already upserts after each run. The in-memory map stays as the fallback when no `db` is configured. This matches what `CodexAgent` and `AntigravityAgent` already do. - **Fall back when a stored id can't be resumed:** the SDK keeps its transcript on local disk under the agent's `cwd`, so a stored id may point at a transcript that isn't there (another host, changed `cwd`, deleted transcript). In that case the CLI exits 1 before yielding any message, and without a fallback the session would fail on every run from then on. If a resume fails before the first message, the adapter forgets the id and starts a fresh SDK session, with the Agno history added to the start of the prompt. Errors after messages have started streaming are still raised, so output is never duplicated. - **Share the prompt builder:** `_build_prompt` moved from `CodexAgent` to `BaseExternalAgent` so both adapters use it. Codex behavior is unchanged. - **Simpler adapter hooks:** both hooks now go through a single `_aquery` helper, which removes the duplicated session-id capture code in the streaming and non-streaming paths. ## Type of change - [x] Bug fix - [ ] 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) - [ ] Examples and guides: Relevant cookbook examples have been included or updated (if applicable) - [ ] 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 **Tests:** new file `libs/agno/tests/unit/agents/test_claude.py`, using a fake `claude_agent_sdk`. It covers: - persistence and resume across a fresh agent instance (streaming and non-streaming) - the stale-id fallback, including the history added to the prompt - the in-memory mapping when no db is configured - no retry when the run wasn't resuming The 3 persistence and fallback tests fail on `main`; the other 3 are regression guards. All 58 tests in `tests/unit/agents/` pass. **End-to-end check with the real SDK (0.2.95):** each turn ran in a separate Python process against a shared SQLite db. 1. Turn 1 stored the SDK session id. 2. Turn 2, in the same `cwd`, resumed that session and recalled a fact from turn 1. 3. Turn 3, in a different `cwd`, logged the resume warning, fell back to a fresh SDK session, still recalled the fact from the added history, and stored the new id. **validate.sh:** ruff passes. mypy reports 0 errors in `agno/agents`. The 36 mypy errors it does report are in modules that wrap third-party SDKs (`tools/mcp`, `utils/models/_genai_compat`, `models/google`, `models/anthropic`). They come from SDK-version drift in my local venv and are not touched by this PR. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: kausmeows --- libs/agno/agno/agents/base.py | 18 ++ libs/agno/agno/agents/claude/agent.py | 111 +++++++--- libs/agno/agno/agents/codex/agent.py | 38 ++-- libs/agno/tests/unit/agents/test_claude.py | 226 +++++++++++++++++++++ libs/agno/tests/unit/agents/test_codex.py | 21 ++ 5 files changed, 361 insertions(+), 53 deletions(-) create mode 100644 libs/agno/tests/unit/agents/test_claude.py diff --git a/libs/agno/agno/agents/base.py b/libs/agno/agno/agents/base.py index 4137f888bc..7b7694802f 100644 --- a/libs/agno/agno/agents/base.py +++ b/libs/agno/agno/agents/base.py @@ -618,6 +618,24 @@ class BaseExternalAgent: history.append({"role": msg.role, "content": str(msg.content)}) return history + @staticmethod + def _build_prompt(input: Any, history: Optional[List[Dict[str, Any]]], resumed: bool) -> str: + """Plain prompt when the framework's own session carries the context; otherwise + prepend the persisted chat history so a fresh session does not lose it.""" + text = str(input) + if resumed or not history: + return text + lines = ["Previous conversation (for context, do not repeat it):"] + for message in history: + role = message.get("role") + content = message.get("content") + if role in ("user", "assistant") and content: + lines.append(f"{role}: {content}") + if len(lines) == 1: + return text + lines.extend(["", "Current message:", text]) + return "\n".join(lines) + # --------------------------------------------------------------------------- # Internal: non-streaming # --------------------------------------------------------------------------- diff --git a/libs/agno/agno/agents/claude/agent.py b/libs/agno/agno/agents/claude/agent.py index 04905fbf45..a65400a01e 100644 --- a/libs/agno/agno/agents/claude/agent.py +++ b/libs/agno/agno/agents/claude/agent.py @@ -10,6 +10,7 @@ from agno.run.agent import ( ToolCallCompletedEvent, ToolCallStartedEvent, ) +from agno.utils.log import log_warning def _sdk() -> Any: @@ -76,10 +77,13 @@ class ClaudeAgent(BaseExternalAgent): options_kwargs: Dict[str, Any] = field(default_factory=dict) framework: str = "claude-agent-sdk" - # Maps Agno session_id -> SDK session id. Keyed per session to avoid cross-session bleed. + # Key under which the SDK session id is stored in the Agno session's session_data. + _SESSION_KEY = "claude_sdk_session_id" + + # 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) - def _build_options(self, *, streaming: bool = False, **kwargs: Any) -> Any: + def _build_options(self, *, streaming: bool = False, resume: Optional[str] = None) -> Any: """Build ClaudeAgentOptions from agent config.""" sdk = _sdk() @@ -108,16 +112,79 @@ class ClaudeAgent(BaseExternalAgent): if streaming: opts["include_partial_messages"] = True - # Resume only the SDK session tied to this Agno session_id. - session_id = kwargs.get("session_id") - if session_id: - sdk_session_id = self._sdk_session_ids.get(session_id) - if sdk_session_id: - opts["resume"] = sdk_session_id + if resume: + opts["resume"] = resume opts.update(self.options_kwargs) return sdk.ClaudeAgentOptions(**opts) + # --------------------------------------------------------------------------- + # Agno session <-> SDK session mapping + # --------------------------------------------------------------------------- + + def _get_sdk_session_id(self, session: Any, session_id: Optional[str]) -> Optional[str]: + if session is not None and session.session_data: + sdk_session_id = session.session_data.get(self._SESSION_KEY) + if sdk_session_id: + return str(sdk_session_id) + if session_id: + return self._sdk_session_ids.get(session_id) + return None + + def _remember_sdk_session(self, session: Any, session_id: Optional[str], sdk_session_id: Optional[str]) -> None: + """Stash the SDK session id on the session (persisted by the base class) and in memory.""" + if not sdk_session_id: + return + if session is not None: + if session.session_data is None: + session.session_data = {} + session.session_data[self._SESSION_KEY] = sdk_session_id + if session_id: + self._sdk_session_ids[session_id] = sdk_session_id + + def _forget_sdk_session(self, session: Any, session_id: Optional[str]) -> None: + if session is not None and session.session_data: + session.session_data.pop(self._SESSION_KEY, None) + if session_id: + self._sdk_session_ids.pop(session_id, None) + + async def _aquery( + self, input: Any, history: Optional[List[Dict[str, Any]]], *, streaming: bool, **kwargs: Any + ) -> AsyncIterator[Any]: + """Run sdk.query() against the SDK session tied to this Agno session, recording its id. + + The SDK transcript lives on local disk under the agent's cwd, so a stored id may not + be resumable (another host, changed cwd, deleted transcript). If the resume fails + before any message arrives, start a fresh SDK session seeded with the Agno history. + """ + sdk = _sdk() + session = kwargs.get("session") + session_id = kwargs.get("session_id") + resume = self._get_sdk_session_id(session, session_id) + + while True: + options = self._build_options(streaming=streaming, resume=resume) + prompt = self._build_prompt(input, history, resumed=resume is not None) + received = False + try: + async for message in sdk.query(prompt=prompt, options=options): + received = True + 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")) + elif isinstance(message, sdk.ResultMessage): + self._remember_sdk_session(session, session_id, getattr(message, "session_id", None)) + yield message + return + except Exception as e: + if resume is None or received: + raise + log_warning( + f"Claude SDK: could not resume session {resume} for session {session_id}: {e}. Starting a new one." + ) + self._forget_sdk_session(session, session_id) + resume = None + @staticmethod def _check_result_message(sdk: Any, message: Any) -> None: """Raise if the SDK reported an error result so the base class can surface it.""" @@ -139,20 +206,11 @@ class ClaudeAgent(BaseExternalAgent): """Non-streaming: collect all messages and return final content.""" sdk = _sdk() - options = self._build_options(**kwargs) - agno_session_id = kwargs.get("session_id") assistant_text = "" final_result = "" - async for message in sdk.query(prompt=str(input), options=options): - if isinstance(message, sdk.SystemMessage): - if hasattr(message, "subtype") and message.subtype == "init": - data = getattr(message, "data", {}) or {} - sdk_session_id = data.get("session_id") - if sdk_session_id and agno_session_id: - self._sdk_session_ids[agno_session_id] = sdk_session_id - - elif isinstance(message, sdk.AssistantMessage): + async for message in self._aquery(input, history, streaming=False, **kwargs): + if isinstance(message, sdk.AssistantMessage): # Accumulate every text block; multiple blocks per message are valid for block in message.content: if isinstance(block, sdk.TextBlock): @@ -160,8 +218,6 @@ class ClaudeAgent(BaseExternalAgent): elif isinstance(message, sdk.ResultMessage): self._check_result_message(sdk, message) - if hasattr(message, "session_id") and message.session_id and agno_session_id: - self._sdk_session_ids[agno_session_id] = message.session_id if hasattr(message, "result") and message.result: final_result = str(message.result) @@ -182,8 +238,6 @@ class ClaudeAgent(BaseExternalAgent): sdk = _sdk() run_id = kwargs.get("run_id", str(uuid4())) - agno_session_id = kwargs.get("session_id") - options = self._build_options(streaming=True, **kwargs) # Track whether we got any StreamEvents (token-level streaming) got_stream_events = False @@ -192,7 +246,7 @@ class ClaudeAgent(BaseExternalAgent): # Map tool_use_id -> (tool_name, tool_args) for carrying forward to ToolCallCompleted tool_info_map: Dict[str, Dict[str, Any]] = {} - async for message in sdk.query(prompt=str(input), options=options): + async for message in self._aquery(input, history, streaming=True, **kwargs): if isinstance(message, sdk.StreamEvent): got_stream_events = True event = message.event @@ -213,13 +267,6 @@ class ClaudeAgent(BaseExternalAgent): content=text, ) - elif isinstance(message, sdk.SystemMessage): - if hasattr(message, "subtype") and message.subtype == "init": - data = getattr(message, "data", {}) or {} - sdk_session_id = data.get("session_id") - if sdk_session_id and agno_session_id: - self._sdk_session_ids[agno_session_id] = sdk_session_id - elif isinstance(message, sdk.AssistantMessage): # Always extract tool calls from complete AssistantMessage # (has full name + args). For text, only use if no StreamEvents. @@ -279,5 +326,3 @@ class ClaudeAgent(BaseExternalAgent): elif isinstance(message, sdk.ResultMessage): self._check_result_message(sdk, message) - if hasattr(message, "session_id") and message.session_id and agno_session_id: - self._sdk_session_ids[agno_session_id] = message.session_id diff --git a/libs/agno/agno/agents/codex/agent.py b/libs/agno/agno/agents/codex/agent.py index caf770218d..6db8a9736c 100644 --- a/libs/agno/agno/agents/codex/agent.py +++ b/libs/agno/agno/agents/codex/agent.py @@ -274,11 +274,25 @@ class CodexAgent(BaseExternalAgent): if session_id: self._thread_ids.pop(session_id, None) + @staticmethod + def _is_missing_thread(error: Exception) -> bool: + """True when the app-server reports the resume target as gone, not a failed call. + + The app-server answers a resume of an unknown thread with a JSON-RPC invalid + request error whose message reads "no rollout found for thread id ..."; a + malformed id reads "invalid session id ...". + """ + text = str(getattr(error, "message", None) or error) + return "no rollout found" in text or "invalid session id" in text + async def _open_thread(self, codex: Any, sdk: Any, session: Any, session_id: Optional[str]) -> Tuple[Any, bool]: """Resume the thread tied to this session, or start a new one. - Returns (thread, resumed). A failed resume (thread files removed, - ephemeral thread) falls back to a fresh thread. + Returns (thread, resumed). Only a resume the app-server rejects because the + thread no longer exists (rollout files removed, ephemeral thread, bad id) + falls back to a fresh thread. Any other failure (busy server, closed + transport, bad cwd) is raised as-is and keeps the stored thread id, so a + transient problem cannot unlink the session from its conversation. """ thread_id = self._get_thread_id(session, session_id) if thread_id: @@ -287,6 +301,8 @@ class CodexAgent(BaseExternalAgent): log_debug(f"Codex: resumed thread {thread_id} for session {session_id}") return thread, True except Exception as e: + if not self._is_missing_thread(e): + raise log_warning( f"Codex: could not resume thread {thread_id} for session {session_id}: {e}. Starting a new one." ) @@ -297,24 +313,6 @@ class CodexAgent(BaseExternalAgent): self._remember_thread(session, session_id, getattr(thread, "id", None)) return thread, False - @staticmethod - def _build_prompt(input: Any, history: Optional[List[Dict[str, Any]]], resumed: bool) -> str: - """Plain prompt when the Codex thread carries its own context; otherwise - prepend the persisted chat history so a fresh thread does not lose it.""" - text = str(input) - if resumed or not history: - return text - lines = ["Previous conversation (for context, do not repeat it):"] - for message in history: - role = message.get("role") - content = message.get("content") - if role in ("user", "assistant") and content: - lines.append(f"{role}: {content}") - if len(lines) == 1: - return text - lines.extend(["", "Current message:", text]) - return "\n".join(lines) - # --------------------------------------------------------------------------- # Adapter hooks # --------------------------------------------------------------------------- diff --git a/libs/agno/tests/unit/agents/test_claude.py b/libs/agno/tests/unit/agents/test_claude.py new file mode 100644 index 0000000000..3b5a4b2974 --- /dev/null +++ b/libs/agno/tests/unit/agents/test_claude.py @@ -0,0 +1,226 @@ +"""Unit tests for ClaudeAgent. + +These exercise SDK session persistence and resume fallback against a fake +claude_agent_sdk module. They do not launch Claude Code. +""" + +import asyncio +import os +import tempfile +from dataclasses import dataclass, field +from types import ModuleType +from typing import Any, Dict, List, Optional + +import pytest + +from agno.agents.claude import ClaudeAgent +from agno.agents.claude import agent as claude_module +from agno.db.sqlite import SqliteDb +from agno.run.agent import RunCompletedEvent, RunContentEvent, RunErrorEvent + +# --------------------------------------------------------------------------- +# Fake claude_agent_sdk +# --------------------------------------------------------------------------- + + +@dataclass +class ClaudeAgentOptions: + resume: Optional[str] = None + include_partial_messages: bool = False + extra: Dict[str, Any] = field(default_factory=dict) + + def __init__(self, **kwargs: Any) -> None: + self.resume = kwargs.pop("resume", None) + self.include_partial_messages = kwargs.pop("include_partial_messages", False) + self.extra = kwargs + + +@dataclass +class SystemMessage: + subtype: str + data: Dict[str, Any] + + +@dataclass +class TextBlock: + text: str + + +@dataclass +class AssistantMessage: + content: List[Any] + + +@dataclass +class ResultMessage: + session_id: str + result: str + subtype: str = "success" + is_error: bool = False + + +class StreamEvent: + pass + + +class UserMessage: + pass + + +class ToolUseBlock: + pass + + +class ToolResultBlock: + pass + + +class ProcessError(Exception): + pass + + +class FakeState: + def __init__(self) -> None: + self.known_sessions: set = set() + self.session_counter = 0 + self.calls: List[Dict[str, Any]] = [] + self.reply = "ok" + + +@pytest.fixture +def fake_sdk(monkeypatch) -> FakeState: + state = FakeState() + + async def query(prompt: str, options: ClaudeAgentOptions): + state.calls.append({"prompt": prompt, "resume": options.resume}) + if options.resume is not None and options.resume not in state.known_sessions: + # Mirrors the CLI: "No conversation found" on stderr, exit 1, no messages yielded + raise ProcessError("Command failed with exit code 1") + if options.resume is not None: + sdk_session_id = options.resume + else: + state.session_counter += 1 + sdk_session_id = f"sdk-{state.session_counter}" + state.known_sessions.add(sdk_session_id) + yield SystemMessage(subtype="init", data={"session_id": sdk_session_id}) + yield AssistantMessage(content=[TextBlock(text=state.reply)]) + yield ResultMessage(session_id=sdk_session_id, result=state.reply) + + module = ModuleType("claude_agent_sdk") + for cls in ( + ClaudeAgentOptions, + SystemMessage, + TextBlock, + AssistantMessage, + ResultMessage, + StreamEvent, + UserMessage, + ToolUseBlock, + ToolResultBlock, + ): + setattr(module, cls.__name__, cls) + module.query = query # type: ignore[attr-defined] + monkeypatch.setattr(claude_module, "_sdk", lambda: module) + return state + + +@pytest.fixture +def tmp_db(): + fd, path = tempfile.mkstemp(suffix=".db") + os.close(fd) + try: + yield SqliteDb(db_file=path) + finally: + if os.path.exists(path): + os.unlink(path) + + +def _collect(agent: ClaudeAgent, text: str, **kwargs: Any) -> List[Any]: + async def consume(): + return [event async for event in agent._arun_stream(text, **kwargs)] + + return asyncio.run(consume()) + + +# --------------------------------------------------------------------------- +# Session persistence +# --------------------------------------------------------------------------- + + +def _run(agent: ClaudeAgent, text: str, stream: bool, **kwargs: Any) -> None: + if stream: + list(agent.run(text, stream=True, **kwargs)) + else: + agent.run(text, **kwargs) + + +@pytest.mark.parametrize("stream", [True, False]) +def test_sdk_session_persisted_and_resumed_after_restart(fake_sdk, tmp_db, stream): + agent = ClaudeAgent(name="Claude", id="claude", db=tmp_db) + _run(agent, "first", stream, session_id="s1", user_id="u1") + + session = agent.read_or_create_session("s1") + assert session.session_data["claude_sdk_session_id"] == "sdk-1" + + # A fresh agent object (simulating a restart) resumes the same SDK session from the DB + agent2 = ClaudeAgent(name="Claude", id="claude", db=tmp_db) + _run(agent2, "second", stream, session_id="s1", user_id="u1") + + assert fake_sdk.calls[-1] == {"prompt": "second", "resume": "sdk-1"} + + +def test_unresumable_session_starts_fresh_with_history(fake_sdk, tmp_db): + agent = ClaudeAgent(name="Claude", id="claude", db=tmp_db) + _collect(agent, "first", session_id="s1", user_id="u1") + + # Transcript gone (other host, changed cwd): the stored id can no longer be resumed + fake_sdk.known_sessions.clear() + events = _collect(agent, "second", session_id="s1", user_id="u1") + + assert isinstance(events[-1], RunCompletedEvent) + assert not any(isinstance(e, RunErrorEvent) for e in events) + assert [c["resume"] for c in fake_sdk.calls] == [None, "sdk-1", None] + prompt = fake_sdk.calls[-1]["prompt"] + assert "user: first" in prompt + assert "assistant: ok" in prompt + assert prompt.endswith("Current message:\nsecond") + + session = agent.read_or_create_session("s1") + assert session.session_data["claude_sdk_session_id"] == "sdk-2" + + +def test_in_memory_mapping_without_db(fake_sdk): + agent = ClaudeAgent(name="Claude", id="claude") + + _collect(agent, "first", session_id="s1") + _collect(agent, "second", session_id="s1") + _collect(agent, "other", session_id="s2") + + assert [c["resume"] for c in fake_sdk.calls] == [None, "sdk-1", None] + assert fake_sdk.calls[1]["prompt"] == "second" + + +def test_streamed_text_and_non_stream_content(fake_sdk): + agent = ClaudeAgent(name="Claude", id="claude") + fake_sdk.reply = "hello" + + events = _collect(agent, "hi", session_id="s1") + assert "".join(e.content for e in events if isinstance(e, RunContentEvent)) == "hello" + assert agent.run("hi", session_id="s2").content == "hello" + + +def test_failure_without_resume_is_not_retried(fake_sdk, monkeypatch): + agent = ClaudeAgent(name="Claude", id="claude") + calls: List[Any] = [] + + async def failing_query(prompt: str, options: ClaudeAgentOptions): + calls.append(options.resume) + raise ProcessError("boom") + yield # pragma: no cover + + module = claude_module._sdk() + monkeypatch.setattr(module, "query", failing_query) + + events = _collect(agent, "hi", session_id="s1") + assert isinstance(events[-1], RunErrorEvent) + assert calls == [None] diff --git a/libs/agno/tests/unit/agents/test_codex.py b/libs/agno/tests/unit/agents/test_codex.py index a717eec627..3c7a3f9abc 100644 --- a/libs/agno/tests/unit/agents/test_codex.py +++ b/libs/agno/tests/unit/agents/test_codex.py @@ -25,6 +25,7 @@ from agno.run.agent import ( ToolCallCompletedEvent, ToolCallStartedEvent, ) +from agno.run.base import RunStatus # --------------------------------------------------------------------------- # Fake openai_codex SDK @@ -474,6 +475,26 @@ def test_unresumable_thread_starts_fresh_with_history(fake_sdk, tmp_db): assert session.session_data["codex_thread_id"] == "thread-2" +def test_transient_resume_failure_keeps_thread(fake_sdk, tmp_db, monkeypatch): + """A resume that fails for any reason other than a missing thread must keep the stored id.""" + fake_sdk.notifications = [_delta("m1", "ok"), _turn_completed()] + agent = CodexAgent(name="Codex", id="codex", db=tmp_db) + agent.run("first", session_id="s1", user_id="u1") + assert agent.read_or_create_session("s1").session_data["codex_thread_id"] == "thread-1" + + async def busy_resume(self, thread_id: str, **kwargs: Any) -> FakeThread: + fake_sdk.calls.append({"op": "thread_resume", "thread_id": thread_id, "kwargs": kwargs}) + raise RuntimeError("server busy: retry limit exceeded") + + monkeypatch.setattr(FakeAsyncCodex, "thread_resume", busy_resume) + run_output = agent.run("second", session_id="s1", user_id="u1") + + assert run_output.status == RunStatus.error + ops = [c["op"] for c in fake_sdk.calls if c["op"] in ("thread_start", "thread_resume")] + assert ops == ["thread_start", "thread_resume"], "an unrelated failure must not start a new thread" + assert agent.read_or_create_session("s1").session_data["codex_thread_id"] == "thread-1" + + def test_in_memory_thread_mapping_without_db(fake_sdk): fake_sdk.notifications = [_delta("m1", "ok"), _turn_completed()] agent = CodexAgent(name="Codex", id="codex")