From 47797d19ee90ae2080e422d719d3398af8b3b524 Mon Sep 17 00:00:00 2001 From: kina0630 Date: Thu, 14 May 2026 19:53:11 +0800 Subject: [PATCH] fix: allow concurrent workspace exports --- .../src/coding_agent/ui/session_manager.py | 37 ++++++++--- .../ui/test_session_manager_public_api.py | 44 +++++++++---- .../tests/ui/test_session_manager_runtime.py | 63 +++++++++++++++++++ 3 files changed, 124 insertions(+), 20 deletions(-) diff --git a/coding-agent/src/coding_agent/ui/session_manager.py b/coding-agent/src/coding_agent/ui/session_manager.py index 58db41f71..31ac6ebfa 100644 --- a/coding-agent/src/coding_agent/ui/session_manager.py +++ b/coding-agent/src/coding_agent/ui/session_manager.py @@ -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] = {} 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( @@ -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) @@ -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() @@ -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) diff --git a/coding-agent/tests/ui/test_session_manager_public_api.py b/coding-agent/tests/ui/test_session_manager_public_api.py index c1b068e5a..9162cbab5 100644 --- a/coding-agent/tests/ui/test_session_manager_public_api.py +++ b/coding-agent/tests/ui/test_session_manager_public_api.py @@ -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, @@ -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( @@ -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 diff --git a/coding-agent/tests/ui/test_session_manager_runtime.py b/coding-agent/tests/ui/test_session_manager_runtime.py index 79aab9c33..8ed005734 100644 --- a/coding-agent/tests/ui/test_session_manager_runtime.py +++ b/coding-agent/tests/ui/test_session_manager_runtime.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import threading import types from pathlib import Path from typing import cast @@ -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()