Skip to content

fix(runner): preserve chat context when repo added mid-session #894

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Gkrumbach07 merged 3 commits into from
Mar 13, 2026
Merged

fix(runner): preserve chat context when repo added mid-session #894

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ec8a878
fix(runner): pass --resume flag when rebuilding adapter after repo add
Mar 12, 2026
6956429
fix(runner): remove unnecessary executable bits from Python files
Gkrumbach07 Mar 12, 2026
1e1c14f
fix(runner): persist Claude session IDs to disk for --resume across p…
Gkrumbach07 Mar 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 8 additions & 0 deletions components/runners/ambient-runner/ag_ui_claude_sdk/adapter.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -369,6 +369,7 @@ def build_options(
self,
input_data: Optional[RunAgentInput] = None,
thread_id: Optional[str] = None,
resume_from: Optional[str] = None,
) -> "ClaudeAgentOptions":
"""
Build ClaudeAgentOptions from stored options (object/dict/None) plus dynamic tools.
Expand All @@ -378,6 +379,8 @@ def build_options(
Args:
input_data: Optional RunAgentInput for extracting dynamic tools
thread_id: Optional thread_id for session resumption lookup
resume_from: Optional CLI session ID to resume (preserves chat history
across adapter rebuilds, e.g. after a repo is added mid-session)

Returns:
Configured ClaudeAgentOptions instance
Expand Down Expand Up @@ -451,6 +454,11 @@ def build_options(

# Remove api_key from options kwargs (handled via environment variable)
merged_kwargs.pop("api_key", None)

# Resume from a previous CLI session (preserves chat context)
if resume_from:
merged_kwargs["resume"] = resume_from

logger.debug(f"Merged kwargs after pop: {merged_kwargs}")

# Apply forwarded_props as per-run overrides (before adding dynamic tools)
Expand Down
10 changes: 9 additions & 1 deletion components/runners/ambient-runner/ambient_runner/bridges/claude/bridge.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,8 @@ def __init__(self) -> None:
self._allowed_tools: list[str] = []
self._system_prompt: dict = {}
self._stderr_lines: list[str] = []
# Preserved session IDs across adapter rebuilds (e.g. repo additions)
self._saved_session_ids: dict[str, str] = {}

# ------------------------------------------------------------------
# PlatformBridge interface
Expand Down Expand Up @@ -99,7 +101,10 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[BaseEvent]:
# 4. Get or create session worker for this thread
thread_id = input_data.thread_id or self._context.session_id
api_key = os.getenv("ANTHROPIC_API_KEY", "")
sdk_options = self._adapter.build_options(input_data, thread_id=thread_id)
saved_session_id = self._saved_session_ids.pop(thread_id, None)
sdk_options = self._adapter.build_options(
input_data, thread_id=thread_id, resume_from=saved_session_id
)
Copy link

@coderabbitai coderabbitai bot Mar 12, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Do not consume the saved resume token before the replacement worker is healthy.

Line 104 removes the only bridge-level session ID before the new worker has connected and emitted a fresh init message. If worker startup fails in that gap, the retry path no longer has any resume_from value and the next run restarts a brand-new conversation. SessionWorker.session_id is only populated later from the init SystemMessage in components/runners/ambient-runner/ambient_runner/bridges/claude/session.py Lines 133-140.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@components/runners/ambient-runner/ambient_runner/bridges/claude/bridge.py`
around lines 104 - 107, The bridge currently pops the saved resume token
(saved_session_id) before the replacement worker is confirmed healthy; instead,
pass the token into Adapter.build_options by peeking (e.g., using get or
indexing into _saved_session_ids) so you do not remove it, and only remove (pop)
the saved token once the new worker has emitted its init/SystemMessage and
SessionWorker.session_id has been populated (i.e., move the pop into the worker
init/handler in
components/runners/ambient-runner/ambient_runner/bridges/claude/session.py where
SessionWorker.session_id is set); keep the resume_from field populated on retry
paths until successful init.

worker = await self._session_manager.get_or_create(
thread_id, sdk_options, api_key
)
Expand Down Expand Up @@ -167,6 +172,9 @@ def mark_dirty(self) -> None:
self._first_run = True
self._adapter = None
if self._session_manager:
# Preserve session IDs so --resume works after adapter rebuild.
# Must be captured synchronously before the async shutdown task runs.
self._saved_session_ids.update(self._session_manager.get_all_session_ids())
manager = self._session_manager
self._session_manager = None
_async_safe_manager_shutdown(manager)
Expand Down
14 changes: 13 additions & 1 deletion components/runners/ambient-runner/ambient_runner/bridges/claude/session.py
100644 → 100755
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,11 @@ async def _run(self) -> None:

os.environ["ANTHROPIC_API_KEY"] = self._api_key

from ambient_runner.bridges.claude.mock_client import MOCK_API_KEY, MockClaudeSDKClient
from ambient_runner.bridges.claude.mock_client import (
MOCK_API_KEY,
MockClaudeSDKClient,
)

if self._api_key == MOCK_API_KEY:
logger.info("[SessionWorker] Using MockClaudeSDKClient (replay mode)")
client: Any = MockClaudeSDKClient(options=self._options)
Expand Down Expand Up @@ -302,6 +306,14 @@ def get_session_id(self, thread_id: str) -> Optional[str]:
return worker.session_id
return self._session_ids.get(thread_id)

def get_all_session_ids(self) -> dict[str, str]:
"""Return a snapshot of all known session IDs (live workers + cached)."""
result = dict(self._session_ids)
for tid, worker in self._workers.items():
if worker.session_id:
result[tid] = worker.session_id
return result

async def destroy(self, thread_id: str) -> None:
"""Stop and remove the worker for *thread_id*.

Expand Down
Loading