Skip to content

Commit 2f0b13d

Browse files
Gkrumbach07Ambient Code Botclaude
authored
fix(runner): preserve chat context when repo added mid-session (#894)
## Summary - Pass saved CLI session ID to adapter rebuild via `resume_from` parameter - Preserves session IDs in the bridge across `mark_dirty()` adapter teardowns - Adds `get_all_session_ids()` public method to `SessionManager` to avoid private attribute access ## Test plan - [x] Unit tests pass (452 passed) - [x] Lints pass (ruff check + ruff format) - [ ] Manual test: add repo mid-session, verify chat history preserved **Jira:** [RHOAIENG-52263](https://issues.redhat.com/browse/RHOAIENG-52263) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Ambient Code Bot <bot@ambient-code.local> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 7c84b25 commit 2f0b13d

3 files changed

Lines changed: 84 additions & 4 deletions

File tree

components/runners/ambient-runner/ag_ui_claude_sdk/adapter.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -369,6 +369,7 @@ def build_options(
369369
self,
370370
input_data: Optional[RunAgentInput] = None,
371371
thread_id: Optional[str] = None,
372+
resume_from: Optional[str] = None,
372373
) -> "ClaudeAgentOptions":
373374
"""
374375
Build ClaudeAgentOptions from stored options (object/dict/None) plus dynamic tools.
@@ -378,6 +379,8 @@ def build_options(
378379
Args:
379380
input_data: Optional RunAgentInput for extracting dynamic tools
380381
thread_id: Optional thread_id for session resumption lookup
382+
resume_from: Optional CLI session ID to resume (preserves chat history
383+
across adapter rebuilds, e.g. after a repo is added mid-session)
381384
382385
Returns:
383386
Configured ClaudeAgentOptions instance
@@ -451,6 +454,11 @@ def build_options(
451454

452455
# Remove api_key from options kwargs (handled via environment variable)
453456
merged_kwargs.pop("api_key", None)
457+
458+
# Resume from a previous CLI session (preserves chat context)
459+
if resume_from:
460+
merged_kwargs["resume"] = resume_from
461+
454462
logger.debug(f"Merged kwargs after pop: {merged_kwargs}")
455463

456464
# Apply forwarded_props as per-run overrides (before adding dynamic tools)

components/runners/ambient-runner/ambient_runner/bridges/claude/bridge.py

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,8 @@ def __init__(self) -> None:
5656
self._allowed_tools: list[str] = []
5757
self._system_prompt: dict = {}
5858
self._stderr_lines: list[str] = []
59+
# Preserved session IDs across adapter rebuilds (e.g. repo additions)
60+
self._saved_session_ids: dict[str, str] = {}
5961

6062
# ------------------------------------------------------------------
6163
# PlatformBridge interface
@@ -99,7 +101,13 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[BaseEvent]:
99101
# 4. Get or create session worker for this thread
100102
thread_id = input_data.thread_id or self._context.session_id
101103
api_key = os.getenv("ANTHROPIC_API_KEY", "")
102-
sdk_options = self._adapter.build_options(input_data, thread_id=thread_id)
104+
saved_session_id = (
105+
self._saved_session_ids.pop(thread_id, None)
106+
or self._session_manager.get_session_id(thread_id)
107+
)
108+
sdk_options = self._adapter.build_options(
109+
input_data, thread_id=thread_id, resume_from=saved_session_id
110+
)
103111
worker = await self._session_manager.get_or_create(
104112
thread_id, sdk_options, api_key
105113
)
@@ -121,6 +129,11 @@ async def run(self, input_data: RunAgentInput) -> AsyncIterator[BaseEvent]:
121129
async for event in wrapped_stream:
122130
yield event
123131

132+
# Persist session ID after turn completes (for --resume on pod restart)
133+
if worker.session_id:
134+
self._session_manager._session_ids[thread_id] = worker.session_id
135+
self._session_manager._persist_session_ids()
136+
124137
self._first_run = False
125138

126139
async def interrupt(self, thread_id: Optional[str] = None) -> None:
@@ -167,6 +180,9 @@ def mark_dirty(self) -> None:
167180
self._first_run = True
168181
self._adapter = None
169182
if self._session_manager:
183+
# Preserve session IDs so --resume works after adapter rebuild.
184+
# Must be captured synchronously before the async shutdown task runs.
185+
self._saved_session_ids.update(self._session_manager.get_all_session_ids())
170186
manager = self._session_manager
171187
self._session_manager = None
172188
_async_safe_manager_shutdown(manager)
@@ -279,7 +295,11 @@ async def _setup_platform(self) -> None:
279295
"""Full platform setup: auth, workspace, MCP, observability."""
280296
# Session manager
281297
if self._session_manager is None:
282-
self._session_manager = SessionManager()
298+
state_dir = os.path.join(
299+
os.getenv("WORKSPACE_PATH", "/workspace"),
300+
os.getenv("RUNNER_STATE_DIR", ".claude"),
301+
)
302+
self._session_manager = SessionManager(state_dir=state_dir)
283303

284304
# Claude-specific auth
285305
from ambient_runner.bridges.claude.auth import setup_sdk_authentication

components/runners/ambient-runner/ambient_runner/bridges/claude/session.py

Lines changed: 54 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,11 @@
2626
"""
2727

2828
import asyncio
29+
import json
2930
import logging
3031
import os
3132
from contextlib import suppress
33+
from pathlib import Path
3234
from typing import Any, AsyncIterator, Optional
3335

3436
logger = logging.getLogger(__name__)
@@ -252,13 +254,18 @@ class SessionManager:
252254
mix messages on the single underlying SDK client).
253255
254256
Tracks session IDs returned by the CLI so that workers can be recreated
255-
with ``--resume`` after a pod restart.
257+
with ``--resume`` after a pod restart. Session IDs are persisted to disk
258+
so they survive pod restarts.
256259
"""
257260

258-
def __init__(self) -> None:
261+
_SESSION_IDS_FILE = "claude_session_ids.json"
262+
263+
def __init__(self, state_dir: str = "") -> None:
259264
self._workers: dict[str, SessionWorker] = {}
260265
self._locks: dict[str, asyncio.Lock] = {}
261266
self._session_ids: dict[str, str] = {} # thread_id -> CLI session_id
267+
self._state_dir = state_dir
268+
self._restore_session_ids()
262269

263270
async def get_or_create(
264271
self,
@@ -306,6 +313,14 @@ def get_session_id(self, thread_id: str) -> Optional[str]:
306313
return worker.session_id
307314
return self._session_ids.get(thread_id)
308315

316+
def get_all_session_ids(self) -> dict[str, str]:
317+
"""Return a snapshot of all known session IDs (live workers + cached)."""
318+
result = dict(self._session_ids)
319+
for tid, worker in self._workers.items():
320+
if worker.session_id:
321+
result[tid] = worker.session_id
322+
return result
323+
309324
async def destroy(self, thread_id: str) -> None:
310325
"""Stop and remove the worker for *thread_id*.
311326
@@ -316,6 +331,7 @@ async def destroy(self, thread_id: str) -> None:
316331
if worker is not None:
317332
if worker.session_id:
318333
self._session_ids[thread_id] = worker.session_id
334+
self._persist_session_ids()
319335
await worker.stop()
320336
self._locks.pop(thread_id, None)
321337
logger.debug(f"[SessionManager] Destroyed worker for thread={thread_id}")
@@ -326,3 +342,39 @@ async def shutdown(self) -> None:
326342
for tid in thread_ids:
327343
await self.destroy(tid)
328344
logger.info("[SessionManager] All workers shut down")
345+
346+
# ── session ID persistence ──
347+
348+
def _session_ids_path(self) -> Path | None:
349+
if not self._state_dir:
350+
return None
351+
return Path(self._state_dir) / self._SESSION_IDS_FILE
352+
353+
def _persist_session_ids(self) -> None:
354+
"""Save session IDs to disk for --resume across pod restarts."""
355+
path = self._session_ids_path()
356+
if not path or not self._session_ids:
357+
return
358+
try:
359+
path.parent.mkdir(parents=True, exist_ok=True)
360+
with open(path, "w") as f:
361+
json.dump(self._session_ids, f)
362+
logger.info("Persisted %d session ID(s) to %s", len(self._session_ids), path)
363+
except OSError:
364+
logger.debug("Could not persist session IDs to %s", path, exc_info=True)
365+
366+
def _restore_session_ids(self) -> None:
367+
"""Restore session IDs from disk (written by a previous pod)."""
368+
path = self._session_ids_path()
369+
if not path or not path.exists():
370+
return
371+
try:
372+
with open(path) as f:
373+
restored = json.load(f)
374+
if isinstance(restored, dict):
375+
self._session_ids.update(restored)
376+
logger.info(
377+
"Restored %d Claude session ID(s) from %s", len(restored), path
378+
)
379+
except (OSError, json.JSONDecodeError):
380+
logger.debug("Could not restore session IDs from %s", path, exc_info=True)

0 commit comments

Comments
 (0)