Skip to content

fix: allow concurrent workspace exports #159

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
delete-cloud merged 1 commit into from
May 14, 2026
Merged

fix: allow concurrent workspace exports #159

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
37 changes: 30 additions & 7 deletions coding-agent/src/coding_agent/ui/session_manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -505,6 +505,7 @@ def __init__(
self._lock = asyncio.Lock()
self._store_io_guard = asyncio.Lock()
self._session_turn_locks: dict[str, asyncio.Lock] = {}
self._session_workspace_export_counts: dict[str, int] = {}

@coderabbitai coderabbitai Bot May 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Clean up export count dict when sessions are removed.

The _session_workspace_export_counts dictionary is never cleaned up when sessions are removed via _remove_session_async_no_lock (line 994 cleans _session_turn_locks but not this dict), remove_session (line 1344), or clear_sessions (line 1360). This will leak memory if sessions are repeatedly created and destroyed.

🧹 Proposed fix to add cleanup in all session removal paths

Add cleanup in _remove_session_async_no_lock after line 994:

         await self._release_owner_lease_for_session(session_id)
         self._approval_stores.pop(session_id, None)
         self._session_turn_locks.pop(session_id, None)
+        self._session_workspace_export_counts.pop(session_id, None)

Add cleanup in remove_session after line 1344:

         self._approval_stores.pop(session_id, None)
         self._session_turn_locks.pop(session_id, None)
+        self._session_workspace_export_counts.pop(session_id, None)

Add cleanup in clear_sessions after line 1360:

         self._session_cache.clear()
         self._approval_stores.clear()
         self._session_turn_locks.clear()
+        self._session_workspace_export_counts.clear()
🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@coding-agent/src/coding_agent/ui/session_manager.py` at line 508, The
_session_workspace_export_counts dict is not cleaned when sessions are removed,
causing memory leaks; update all session-removal paths to remove the session's
entry (use pop(session_id, None) or conditional del) for
_session_workspace_export_counts. Specifically, in _remove_session_async_no_lock
add cleanup alongside the existing _session_turn_locks cleanup (referencing
_session_turn_locks and _remove_session_async_no_lock), and likewise remove the
session key in remove_session and clear_sessions after their existing cleanup
logic (referencing remove_session and clear_sessions) so the export-count entry
is always cleared when a session is deleted.

data_dir = Path(os.environ.get("AGENT_DATA_DIR", "./data"))
self._tape_store = tape_store or self._create_tape_store(data_dir)
resolved_checkpoint_store = checkpoint_store or self._create_checkpoint_store(
Expand Down Expand Up @@ -689,10 +690,27 @@ def _turn_lock_for(self, session_id: str) -> asyncio.Lock:
self._session_turn_locks[session_id] = lock
return lock

def _workspace_export_in_progress(self, session_id: str) -> bool:
return self._session_workspace_export_counts.get(session_id, 0) > 0

def _begin_workspace_export(self, session_id: str) -> None:
self._session_workspace_export_counts[session_id] = (
self._session_workspace_export_counts.get(session_id, 0) + 1
)

def _end_workspace_export(self, session_id: str) -> None:
count = self._session_workspace_export_counts.get(session_id, 0)
if count <= 1:
self._session_workspace_export_counts.pop(session_id, None)
return
self._session_workspace_export_counts[session_id] = count - 1

async def prepare_session_turn(self, session_id: str) -> Session:
turn_lock = self._turn_lock_for(session_id)
if turn_lock.locked():
raise RuntimeError("turn already in progress")
if self._workspace_export_in_progress(session_id):
raise RuntimeError("turn already in progress")

session = await self.get_session_async(session_id)
await self._assert_owner(session_id)
Expand Down Expand Up @@ -1525,6 +1543,8 @@ async def run_agent(
raise RuntimeError("turn already in progress")

async with turn_lock:
if self._workspace_export_in_progress(session_id):
raise RuntimeError("turn already in progress")
await self._assert_owner(session_id)
session = await self.get_session_async(session_id)
session.last_activity = datetime.now()
Expand Down Expand Up @@ -2178,16 +2198,19 @@ async def export_workspace_archive(
if turn_lock.locked():
raise RuntimeError("turn already in progress")

async with turn_lock:
await self._assert_owner(session_id)
session = await self.get_session_async(session_id)
if session.turn_in_progress or (session.task and not session.task.done()):
raise RuntimeError("turn already in progress")
if not isinstance(session.execution_binding, CloudWorkspaceBinding):
raise ValueError("Workspace export requires cloud session")
await self._assert_owner(session_id)
session = await self.get_session_async(session_id)
if session.turn_in_progress or (session.task and not session.task.done()):
raise RuntimeError("turn already in progress")
if not isinstance(session.execution_binding, CloudWorkspaceBinding):
raise ValueError("Workspace export requires cloud session")
self._begin_workspace_export(session_id)
try:
result = await asyncio.to_thread(export_archive, session.execution_binding)
await self._assert_owner(session_id)
return result
finally:
self._end_workspace_export(session_id)

async def list_checkpoints(self, session_id: str) -> list[CheckpointMeta]:
session = await self.get_session_async(session_id)
Expand Down
44 changes: 31 additions & 13 deletions coding-agent/tests/ui/test_session_manager_public_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,10 @@
from coding_agent.approval.store import ApprovalStore
from coding_agent.approval import ApprovalPolicy
from agentkit.config.loader import ConfigError
from coding_agent.ui.execution_binding import CloudWorkspaceBinding, LocalExecutionBinding
from coding_agent.ui.execution_binding import (
CloudWorkspaceBinding,
LocalExecutionBinding,
)
from coding_agent.wire.protocol import (
ApprovalRequest,
CompletionStatus,
Expand Down Expand Up @@ -1571,11 +1574,13 @@ async def test_export_workspace_archive_rejects_active_turn() -> None:
manager.register_session(session)

with pytest.raises(RuntimeError, match="turn already in progress"):
await manager.export_workspace_archive(session_id, lambda binding: binding.workspace_id)
await manager.export_workspace_archive(
session_id, lambda binding: binding.workspace_id
)


@pytest.mark.asyncio
async def test_export_workspace_archive_rejects_when_turn_lock_is_held() -> None:
async def test_export_workspace_archive_allows_concurrent_idle_exports() -> None:
manager = SessionManager(store=InMemorySessionStore())
session_id = await manager.create_session(
execution_binding=CloudWorkspaceBinding(
Expand All @@ -1584,16 +1589,29 @@ async def test_export_workspace_archive_rejects_when_turn_lock_is_held() -> None
)
)

turn_lock = manager._turn_lock_for(session_id)
await turn_lock.acquire()
try:
with pytest.raises(RuntimeError, match="turn already in progress"):
await manager.export_workspace_archive(
session_id,
lambda binding: binding.workspace_id,
)
finally:
turn_lock.release()
first_export_started = asyncio.Event()
release_first_export = threading.Event()

def first_export(binding: CloudWorkspaceBinding) -> str:
first_export_started.set()
assert release_first_export.wait(timeout=10)
return binding.workspace_id

async def run_first_export() -> str:
return await manager.export_workspace_archive(session_id, first_export)

first_task = asyncio.create_task(run_first_export())
await asyncio.wait_for(first_export_started.wait(), timeout=1)

second_result = await manager.export_workspace_archive(
session_id,
lambda binding: f"{binding.workspace_id}-second",
)
release_first_export.set()
first_result = await first_task

assert first_result == "ws-export"
assert second_result == "ws-export-second"


@pytest.mark.asyncio
Expand Down
63 changes: 63 additions & 0 deletions coding-agent/tests/ui/test_session_manager_runtime.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

import asyncio
import threading
import types
from pathlib import Path
from typing import cast
Expand Down Expand Up @@ -659,6 +660,68 @@ def fake_create_agent(**kwargs):
await first


@pytest.mark.asyncio
async def test_prepare_session_turn_rejects_active_workspace_export() -> None:
store = InMemorySessionStore()
manager = SessionManager(store=store)
session_id = await manager.create_session(
execution_binding=CloudWorkspaceBinding(
workspace_url="docker://agent-ws-export/workspace",
workspace_id="ws-export",
)
)
export_started = asyncio.Event()
release_export = threading.Event()

def blocking_export(binding: CloudWorkspaceBinding) -> str:
export_started.set()
assert release_export.wait(timeout=10)
return binding.workspace_id

export_task = asyncio.create_task(
manager.export_workspace_archive(session_id, blocking_export)
)
await asyncio.wait_for(export_started.wait(), timeout=1)

try:
with pytest.raises(RuntimeError, match="turn already in progress"):
await manager.prepare_session_turn(session_id)
finally:
release_export.set()
await export_task


@pytest.mark.asyncio
async def test_run_agent_rejects_active_workspace_export() -> None:
store = InMemorySessionStore()
manager = SessionManager(store=store)
session_id = await manager.create_session(
execution_binding=CloudWorkspaceBinding(
workspace_url="docker://agent-ws-export/workspace",
workspace_id="ws-export",
)
)
export_started = asyncio.Event()
release_export = threading.Event()

def blocking_export(binding: CloudWorkspaceBinding) -> str:
export_started.set()
assert release_export.wait(timeout=10)
return binding.workspace_id

export_task = asyncio.create_task(
manager.export_workspace_archive(session_id, blocking_export)
)
await asyncio.wait_for(export_started.wait(), timeout=1)

try:
with pytest.raises(RuntimeError, match="turn already in progress"):
await manager.run_agent(session_id, "do work")
finally:
release_export.set()
await export_task


@pytest.mark.asyncio
async def test_rehydrated_session_rebuilds_runtime_from_persisted_tape() -> None:
store = InMemorySessionStore()
Expand Down