fix: persist ClaudeAgent SDK session id across restarts (#10903)

## 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 <[email protected]>
2026-10-12 02:10:54 +00:00 · 2026-10-09 11:26:38 +05:30
parent 5f1fd0ce28
commit 78fa25982e
5 changed files with 361 additions and 53 deletions
@@ -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
    # ---------------------------------------------------------------------------
@@ -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
@@ -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
    # ---------------------------------------------------------------------------
@@ -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]
@@ -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")