(null);
// --- Getters ---
const isBoardMode = computed(
@@ -1806,6 +1837,7 @@ export function useChatStream(deps: ChatStreamDeps) {
boardState,
debateState,
collaborationState,
+ autonomyPaused,
conversations: deps.conversations,
currentConversationId: deps.currentConversationId,
appendMessage: deps.appendMessage,
@@ -1830,6 +1862,7 @@ export function useChatStream(deps: ChatStreamDeps) {
boardState,
debateState,
collaborationState,
+ autonomyPaused,
// Getters
isBoardMode,
isPlanExec,
diff --git a/src/agentkit/server/frontend/src/views/ChatView.vue b/src/agentkit/server/frontend/src/views/ChatView.vue
index 05e50aa..71b0283 100644
--- a/src/agentkit/server/frontend/src/views/ChatView.vue
+++ b/src/agentkit/server/frontend/src/views/ChatView.vue
@@ -73,6 +73,16 @@
+
+
@@ -108,6 +118,7 @@ import ChatMessage from '@/components/chat/ChatMessage.vue'
import ChatInput from '@/components/chat/ChatInput.vue'
import StickyModeHeader from '@/components/chat/StickyModeHeader.vue'
import PhaseIndicator from '@/components/chat/PhaseIndicator.vue'
+import AutonomyPausedCard from '@/components/chat/messages/AutonomyPausedCard.vue'
const ATypographyText = ATypography.Text
@@ -158,6 +169,18 @@ watch(
}
)
+// IQ-Boost/U1: scroll the autonomy pause card into view when it appears so
+// the user sees the resume button without manual scrolling.
+watch(
+ () => chatStore.autonomyPaused,
+ async (paused) => {
+ if (paused) {
+ await nextTick()
+ scrollToBottom()
+ }
+ }
+)
+
function scrollToBottom(): void {
if (messagesContainer.value) {
messagesContainer.value.scrollTop = messagesContainer.value.scrollHeight
@@ -171,6 +194,11 @@ function handleSend(message: string, model?: string): void {
chatStore.sendMessage(message)
}
}
+
+/** IQ-Boost/U1: resume autonomy-paused execution. */
+function handleResumeAutonomy(resumeToken: string): void {
+ chatStore.resumeAutonomy(resumeToken)
+}
diff --git a/src/agentkit/server/frontend/tests/unit/components/AutonomyPausedCard.test.ts b/src/agentkit/server/frontend/tests/unit/components/AutonomyPausedCard.test.ts
new file mode 100644
index 0000000..149ac5a
--- /dev/null
+++ b/src/agentkit/server/frontend/tests/unit/components/AutonomyPausedCard.test.ts
@@ -0,0 +1,150 @@
+/**
+ * IQ-Boost/U1: unit tests for AutonomyPausedCard.
+ *
+ * Mount strategy: native Vue createApp + h (no @vue/test-utils dependency),
+ * consistent with BoardBannerCard.test.ts. Verifies:
+ * - reason → 中文标签映射 (timeout / consecutive_failures / manual)
+ * - progress 摘要渲染 (步骤 N / M · 工具 X)
+ * - resume 按钮点击 emit resume 事件携带 resume_token
+ */
+
+import { afterEach, describe, expect, it } from 'vitest'
+import { createApp, h, type App } from 'vue'
+import AutonomyPausedCard from '@/components/chat/messages/AutonomyPausedCard.vue'
+import type { IAutonomyPausedData } from '@/api/types'
+
+interface Mounted {
+ container: HTMLElement
+ root: HTMLElement
+ app: App
+ unmount: () => void
+}
+
+function mountCard(
+ props: { data: IAutonomyPausedData },
+ onResume?: (token: string) => void,
+): Mounted {
+ const container = document.createElement('div')
+ document.body.appendChild(container)
+ const app = createApp({
+ setup() {
+ return () =>
+ h(AutonomyPausedCard as never, {
+ data: props.data as never,
+ onResume: (token: string) => onResume?.(token),
+ } as never)
+ },
+ })
+ app.mount(container)
+ const root = container.querySelector('.autonomy-card') as HTMLElement
+ return {
+ container,
+ root,
+ app,
+ unmount: () => {
+ app.unmount()
+ container.remove()
+ },
+ }
+}
+
+function makeData(
+ overrides: Partial = {},
+): IAutonomyPausedData {
+ return {
+ reason: 'timeout',
+ progress: { step: 3, tool_name: 'shell', total_steps: 10 },
+ resume_token: 'autonomy_pause:timeout:3',
+ consecutive_failures: 0,
+ elapsed_seconds: 305.4,
+ ...overrides,
+ }
+}
+
+describe('AutonomyPausedCard (IQ-Boost/U1)', () => {
+ let mounted: Mounted | null = null
+
+ afterEach(() => {
+ mounted?.unmount()
+ mounted = null
+ })
+
+ it('renders reason=timeout as "自主执行超时"', () => {
+ mounted = mountCard({ data: makeData({ reason: 'timeout' }) })
+ const reason = mounted.container.querySelector(
+ '.autonomy-card__reason',
+ ) as HTMLElement
+ expect(reason).toBeTruthy()
+ expect(reason.textContent).toBe('自主执行超时')
+ })
+
+ it('renders reason=consecutive_failures as "连续失败触发暂停"', () => {
+ mounted = mountCard({
+ data: makeData({ reason: 'consecutive_failures', consecutive_failures: 3 }),
+ })
+ const reason = mounted.container.querySelector(
+ '.autonomy-card__reason',
+ ) as HTMLElement
+ expect(reason.textContent).toBe('连续失败触发暂停')
+ // consecutive_failures count is also shown in meta.
+ const meta = mounted.container.querySelector(
+ '.autonomy-card__meta',
+ ) as HTMLElement
+ expect(meta.textContent).toContain('连续失败 3 次')
+ })
+
+ it('renders reason=manual as "用户手动暂停"', () => {
+ mounted = mountCard({ data: makeData({ reason: 'manual' }) })
+ const reason = mounted.container.querySelector(
+ '.autonomy-card__reason',
+ ) as HTMLElement
+ expect(reason.textContent).toBe('用户手动暂停')
+ })
+
+ it('renders progress summary "步骤 3 / 10 · 工具 shell"', () => {
+ mounted = mountCard({ data: makeData() })
+ const progress = mounted.container.querySelector(
+ '.autonomy-card__progress',
+ ) as HTMLElement
+ expect(progress).toBeTruthy()
+ expect(progress.textContent).toContain('步骤 3 / 10')
+ expect(progress.textContent).toContain('工具 shell')
+ })
+
+ it('renders elapsed time "已运行 5m 5s" for 305.4s', () => {
+ mounted = mountCard({ data: makeData({ elapsed_seconds: 305.4 }) })
+ const meta = mounted.container.querySelector(
+ '.autonomy-card__meta',
+ ) as HTMLElement
+ expect(meta.textContent).toContain('已运行 5m 5s')
+ })
+
+ it('emits resume with resume_token when 继续 button is clicked', async () => {
+ let emittedToken: string | null = null
+ mounted = mountCard(
+ { data: makeData({ resume_token: 'autonomy_pause:timeout:7' }) },
+ (token) => {
+ emittedToken = token
+ },
+ )
+ const button = mounted.container.querySelector(
+ '.autonomy-card__actions button',
+ ) as HTMLButtonElement
+ expect(button).toBeTruthy()
+ button.click()
+ // Vue event flush — await a microtask.
+ await Promise.resolve()
+ expect(emittedToken).toBe('autonomy_pause:timeout:7')
+ })
+
+ it('hides progress block when progress has no known fields', () => {
+ mounted = mountCard({
+ data: makeData({ progress: {} }),
+ })
+ const progress = mounted.container.querySelector(
+ '.autonomy-card__progress',
+ )
+ // v-if="progressSummary" — empty progress → no render.
+ expect(progress).toBeNull()
+ })
+})
diff --git a/src/agentkit/server/frontend/tests/unit/stores/chatStream.autonomy.test.ts b/src/agentkit/server/frontend/tests/unit/stores/chatStream.autonomy.test.ts
new file mode 100644
index 0000000..3e6fa78
--- /dev/null
+++ b/src/agentkit/server/frontend/tests/unit/stores/chatStream.autonomy.test.ts
@@ -0,0 +1,192 @@
+/**
+ * IQ-Boost/U1: unit tests for the autonomy_paused dispatch branch.
+ *
+ * dispatchWsEvent is a pure function over ChatStreamState, so we build a
+ * fixture (mirrors chatStream.test.ts) and assert:
+ * - autonomy_paused event stores the payload in state.autonomyPaused
+ * - a "自主执行已暂停" milestone step is appended for visibility
+ * - resumeAutonomy (chatStore) clears autonomyPaused and sends the resume
+ * WS message — covered here by asserting the dispatch-side clear path
+ * (the store action is a thin wrapper over socket.send + state clear).
+ */
+
+import { beforeEach, describe, expect, it, vi } from 'vitest'
+import { ref, type Ref } from 'vue'
+import {
+ dispatchWsEvent,
+ type ChatStreamState,
+ type IStreamingStep,
+} from '@/stores/chatStream'
+import type {
+ IAutonomyPausedData,
+ IChatMessage,
+ IConversation,
+ IExpertTeamState,
+ WsServerMessage,
+} from '@/api/types'
+
+// Mock ant-design-vue so any dynamic import in dispatchWsEvent is harmless.
+vi.mock('ant-design-vue', () => ({
+ message: { warning: vi.fn() },
+}))
+
+vi.mock('@/api/documents', () => ({
+ isDocumentMeta: vi.fn(() => true),
+}))
+
+interface Fixture {
+ state: ChatStreamState
+ conversations: Ref
+ currentConversationId: Ref
+ appendMessageSpy: ReturnType
+ updateMessageSpy: ReturnType
+ markConversationDoneSpy: ReturnType
+ teamStore: {
+ teamState: IExpertTeamState | null
+ setTeamState: ReturnType
+ updatePhases: ReturnType
+ updatePhaseStatus: ReturnType
+ clearTeam: ReturnType
+ }
+ calendarStore: { handleWsEvent: ReturnType }
+ documentsStore: { addDocument: ReturnType }
+}
+
+function createFixture(convId: string = 'conv-1'): Fixture {
+ const conversations = ref([
+ {
+ id: convId,
+ title: '新对话',
+ messages: [],
+ created_at: new Date().toISOString(),
+ updated_at: new Date().toISOString(),
+ },
+ ])
+ const currentConversationId = ref(convId)
+ const appendMessageSpy = vi.fn((cid: string, msg: IChatMessage) => {
+ const conv = conversations.value.find((c) => c.id === cid)
+ if (conv) conv.messages.push(msg)
+ })
+ const updateMessageSpy = vi.fn(
+ (cid: string, mid: string, updates: Partial) => {
+ const conv = conversations.value.find((c) => c.id === cid)
+ if (!conv) return
+ const msg = conv.messages.find((m) => m.id === mid)
+ if (msg) Object.assign(msg, updates)
+ },
+ )
+ const markConversationDoneSpy = vi.fn()
+
+ const teamStore = {
+ teamState: null as IExpertTeamState | null,
+ setTeamState: vi.fn(),
+ updatePhases: vi.fn(),
+ updatePhaseStatus: vi.fn(),
+ clearTeam: vi.fn(),
+ }
+ const calendarStore = { handleWsEvent: vi.fn() }
+ const documentsStore = { addDocument: vi.fn() }
+
+ const state: ChatStreamState = {
+ streamingStepsByConv: ref(new Map()),
+ currentPhase: ref(null),
+ phaseViolations: ref([]),
+ boardState: ref(null),
+ debateState: ref(null),
+ collaborationState: ref(null),
+ autonomyPaused: ref(null),
+ conversations,
+ currentConversationId,
+ appendMessage: appendMessageSpy,
+ updateMessage: updateMessageSpy,
+ markConversationDone: markConversationDoneSpy,
+ resolveIncomingConvId: () => currentConversationId.value ?? '',
+ getTeamStore: () => teamStore as unknown as ChatStreamState['getTeamStore'] extends () => infer R ? R : never,
+ getCalendarStore: () => calendarStore as unknown as ChatStreamState['getCalendarStore'] extends () => infer R ? R : never,
+ getDocumentsStore: () => documentsStore as unknown as ChatStreamState['getDocumentsStore'] extends () => infer R ? R : never,
+ }
+
+ return {
+ state,
+ conversations,
+ currentConversationId,
+ appendMessageSpy,
+ updateMessageSpy,
+ markConversationDoneSpy,
+ teamStore,
+ calendarStore,
+ documentsStore,
+ }
+}
+
+function autonomyEvent(
+ overrides: Partial = {},
+): WsServerMessage {
+ const data: IAutonomyPausedData = {
+ reason: 'timeout',
+ progress: { step: 3, tool_name: 'shell', total_steps: 10 },
+ resume_token: 'autonomy_pause:timeout:3',
+ consecutive_failures: 0,
+ elapsed_seconds: 305.4,
+ ...overrides,
+ }
+ return { type: 'autonomy_paused', data }
+}
+
+describe('dispatchWsEvent — autonomy_paused (IQ-Boost/U1)', () => {
+ let f: Fixture
+ beforeEach(() => {
+ f = createFixture()
+ })
+
+ it('stores the autonomy_paused payload in state.autonomyPaused', () => {
+ expect(f.state.autonomyPaused.value).toBeNull()
+ dispatchWsEvent(autonomyEvent(), f.state)
+ expect(f.state.autonomyPaused.value).not.toBeNull()
+ expect(f.state.autonomyPaused.value?.reason).toBe('timeout')
+ expect(f.state.autonomyPaused.value?.resume_token).toBe(
+ 'autonomy_pause:timeout:3',
+ )
+ expect(f.state.autonomyPaused.value?.progress.step).toBe(3)
+ })
+
+ it('appends a "自主执行已暂停" milestone step for visibility', () => {
+ dispatchWsEvent(autonomyEvent(), f.state)
+ const steps = f.state.streamingStepsByConv.value.get('conv-1') ?? []
+ expect(steps.length).toBe(1)
+ expect(steps[0].label).toBe('自主执行已暂停')
+ expect(steps[0].type).toBe('milestone')
+ expect(steps[0].status).toBe('running')
+ expect(steps[0].detail).toBe('timeout')
+ })
+
+ it('overwrites a previous pause payload when a second event arrives', () => {
+ dispatchWsEvent(autonomyEvent({ reason: 'timeout' }), f.state)
+ dispatchWsEvent(
+ autonomyEvent({
+ reason: 'consecutive_failures',
+ resume_token: 'autonomy_pause:consecutive_failures:5',
+ consecutive_failures: 3,
+ }),
+ f.state,
+ )
+ expect(f.state.autonomyPaused.value?.reason).toBe('consecutive_failures')
+ expect(f.state.autonomyPaused.value?.resume_token).toBe(
+ 'autonomy_pause:consecutive_failures:5',
+ )
+ expect(f.state.autonomyPaused.value?.consecutive_failures).toBe(3)
+ // Two milestone steps accumulated (one per event).
+ const steps = f.state.streamingStepsByConv.value.get('conv-1') ?? []
+ expect(steps.length).toBe(2)
+ })
+
+ it('does not throw when resolveIncomingConvId returns empty (no conv)', () => {
+ f.currentConversationId.value = null
+ // No pending conversations either → resolver returns "".
+ expect(() => dispatchWsEvent(autonomyEvent(), f.state)).not.toThrow()
+ // State is still set (dispatch sets it before the conv-id-gated step).
+ expect(f.state.autonomyPaused.value?.reason).toBe('timeout')
+ // No step appended (no conv to attach it to).
+ expect(f.state.streamingStepsByConv.value.size).toBe(0)
+ })
+})
diff --git a/src/agentkit/server/frontend/tests/unit/stores/chatStream.test.ts b/src/agentkit/server/frontend/tests/unit/stores/chatStream.test.ts
index a3060d6..976731b 100644
--- a/src/agentkit/server/frontend/tests/unit/stores/chatStream.test.ts
+++ b/src/agentkit/server/frontend/tests/unit/stores/chatStream.test.ts
@@ -98,6 +98,7 @@ function createFixture(convId: string = 'conv-1'): Fixture {
boardState: ref(null),
debateState: ref(null),
collaborationState: ref(null),
+ autonomyPaused: ref(null),
conversations,
currentConversationId,
appendMessage: appendMessageSpy,
diff --git a/src/agentkit/server/routes/chat.py b/src/agentkit/server/routes/chat.py
index 089fce2..0c907ed 100644
--- a/src/agentkit/server/routes/chat.py
+++ b/src/agentkit/server/routes/chat.py
@@ -1620,10 +1620,20 @@ async def _handle_chat_message(
)
async def _resume_handler(resume_token: str, reason: str) -> bool:
- """Block until the user sends ``resume`` for the given token."""
- loop = asyncio.get_running_loop()
- future: asyncio.Future[bool] = loop.create_future()
- _pending_autonomy_resumes[resume_token] = future
+ """Block until the user sends ``resume`` for the given token.
+
+ U6: The Future is pre-registered in the ``autonomy_paused`` event
+ handler below (before ``send_json``) so the WebSocket loop can
+ resolve it even if the resume message arrives before the engine
+ calls this handler. This closure only awaits the pre-registered
+ Future; the fallback (create on miss) covers non-WS callers/tests.
+ """
+ future = _pending_autonomy_resumes.get(resume_token)
+ if future is None:
+ # Fallback: non-WS caller or event not seen yet — create on demand.
+ loop = asyncio.get_running_loop()
+ future = loop.create_future()
+ _pending_autonomy_resumes[resume_token] = future
logger.info(f"Autonomy paused ({reason}), waiting for resume: {resume_token}")
try:
# Wait up to 30 minutes for user resume (long task availability).
@@ -1719,6 +1729,16 @@ async def _handle_chat_message(
# The _resume_handler closure is already blocking the engine
# waiting for the user's ``resume`` message. The frontend
# shows a pause card with reason + progress + resume button.
+ #
+ # U6: Pre-register the resume Future BEFORE send_json so the
+ # WebSocket loop can resolve it even if the resume message
+ # arrives while send_json is suspended (race fix). Without
+ # this, the resume message would be dropped with "token not
+ # found" and the engine would deadlock.
+ resume_token = event.data.get("resume_token")
+ if resume_token and resume_token not in _pending_autonomy_resumes:
+ loop = asyncio.get_running_loop()
+ _pending_autonomy_resumes[resume_token] = loop.create_future()
await websocket.send_json(
{
"type": "autonomy_paused",
diff --git a/tests/unit/test_autonomy_paused.py b/tests/unit/test_autonomy_paused.py
index 417befb..9a2fdb2 100644
--- a/tests/unit/test_autonomy_paused.py
+++ b/tests/unit/test_autonomy_paused.py
@@ -9,10 +9,12 @@ Verifies:
- Non-autonomy mode → no pause (gate is a no-op)
- Disabled thresholds (0) → no pause
- _track_tool_result_for_autonomy increments on error, resets on success
+- U6: resume Future pre-registration race fix
"""
from __future__ import annotations
+import asyncio
import time
from unittest.mock import AsyncMock, MagicMock
@@ -43,9 +45,7 @@ async def _detect_and_await_pause(
return True, []
reason, token, event_data = pause_info
event = ReActEvent(event_type="autonomy_paused", step=step, data=event_data)
- should_continue = await engine._await_autonomy_resume(
- token, reason, resume_handler
- )
+ should_continue = await engine._await_autonomy_resume(token, reason, resume_handler)
return should_continue, [event]
@@ -391,3 +391,191 @@ class TestAutonomyConfigParsing:
cfg = DangerousToolsConfig.from_dict({})
assert cfg.autonomy_timeout_minutes == 30
assert cfg.max_consecutive_failures == 3
+
+
+# ---------------------------------------------------------------------------
+# U6: resume Future pre-registration race fix
+# ---------------------------------------------------------------------------
+
+
+class TestResumeFutureRaceFix:
+ """U6: Future is pre-registered before send_json so the WebSocket loop
+ can resolve it even if the resume message arrives before _resume_handler
+ is called by the engine.
+
+ These tests mirror the chat.py ``autonomy_paused`` event handler +
+ ``_resume_handler`` closure interaction without spinning up a full
+ WebSocket server.
+ """
+
+ @pytest.mark.asyncio
+ async def test_future_registered_before_send(self):
+ """Future exists in pending_autonomy_resumes before send_json fires.
+
+ Simulates the autonomy_paused event handler: the pre-registration
+ step (U6) must insert the Future before the event is forwarded to
+ the frontend.
+ """
+ pending: dict[str, asyncio.Future] = {}
+ event_data = {
+ "resume_token": "autonomy_pause:timeout:1",
+ "reason": "timeout",
+ }
+ # Mirror chat.py: autonomy_paused event handler pre-registration.
+ resume_token = event_data.get("resume_token")
+ if resume_token and resume_token not in pending:
+ loop = asyncio.get_running_loop()
+ pending[resume_token] = loop.create_future()
+ # At this point send_json has NOT run yet — Future is already
+ # registered, so a concurrent resume message can resolve it.
+ assert resume_token in pending
+ assert not pending[resume_token].done()
+
+ @pytest.mark.asyncio
+ async def test_resume_message_not_lost_when_arriving_before_handler(self):
+ """Resume message arriving before _resume_handler is awaited still
+ resolves the Future (no lost message).
+
+ This is the core race the fix addresses: pre-registration lets the
+ WebSocket loop find the Future even if the engine has not yet
+ called _resume_handler.
+ """
+ pending: dict[str, asyncio.Future] = {}
+ resume_token = "autonomy_pause:timeout:1"
+
+ # Step 1: autonomy_paused event handler pre-registers Future.
+ loop = asyncio.get_running_loop()
+ pre_future: asyncio.Future[bool] = loop.create_future()
+ pending[resume_token] = pre_future
+
+ # Step 2: WebSocket loop receives resume BEFORE engine calls
+ # _resume_handler (the race window the fix closes).
+ assert resume_token in pending
+ fut = pending[resume_token]
+ if not fut.done():
+ fut.set_result(True)
+
+ # Step 3: engine finally calls _resume_handler, which awaits the
+ # already-resolved Future — no deadlock.
+ assert pre_future.done()
+ assert pre_future.result() is True
+
+ @pytest.mark.asyncio
+ async def test_resume_handler_uses_pre_registered_future(self):
+ """_resume_handler finds and awaits the pre-registered Future
+ instead of creating a new one (which would orphan the pre-registered
+ Future and lose the resume message)."""
+ pending: dict[str, asyncio.Future] = {}
+ resume_token = "autonomy_pause:consecutive_failures:3"
+
+ # Pre-register (autonomy_paused event handler).
+ loop = asyncio.get_running_loop()
+ pre_future: asyncio.Future[bool] = loop.create_future()
+ pending[resume_token] = pre_future
+
+ # Mirror _resume_handler: look up existing Future first.
+ future = pending.get(resume_token)
+ assert future is pre_future # Same object — no orphaned Future.
+
+ # Resolve and await.
+ pre_future.set_result(True)
+ result = await asyncio.wait_for(future, timeout=1.0)
+ assert result is True
+
+ @pytest.mark.asyncio
+ async def test_resume_handler_fallback_creates_future(self):
+ """When no pre-registered Future exists (non-WS caller or event not
+ seen), _resume_handler creates one as a fallback."""
+ pending: dict[str, asyncio.Future] = {}
+ resume_token = "autonomy_pause:timeout:1"
+
+ # Mirror _resume_handler fallback path.
+ future = pending.get(resume_token)
+ if future is None:
+ loop = asyncio.get_running_loop()
+ future = loop.create_future()
+ pending[resume_token] = future
+
+ assert resume_token in pending
+ assert pending[resume_token] is future
+ assert not future.done()
+
+ @pytest.mark.asyncio
+ async def test_double_resolve_guard(self):
+ """Duplicate resume messages do not crash — the second is ignored
+ because the Future is already done.
+
+ Mirrors the WebSocket loop's double-resolve guard
+ (``if not fut.done(): fut.set_result(...)``).
+ """
+ pending: dict[str, asyncio.Future] = {}
+ resume_token = "autonomy_pause:timeout:1"
+ loop = asyncio.get_running_loop()
+ future: asyncio.Future[bool] = loop.create_future()
+ pending[resume_token] = future
+
+ # First resume resolves the Future.
+ assert not future.done()
+ future.set_result(True)
+
+ # Second resume: guard prevents double-resolve (would raise
+ # InvalidStateError without the done() check).
+ if not future.done():
+ future.set_result(True) # Not reached — guard holds.
+
+ assert future.done()
+ assert future.result() is True
+
+ @pytest.mark.asyncio
+ async def test_pre_registration_does_not_overwrite_existing(self):
+ """Pre-registration skips when a Future already exists (idempotent).
+
+ Guards against the edge case where the engine emits autonomy_paused
+ twice with the same token (e.g. retry) — the second pre-registration
+ must not replace the first Future.
+ """
+ pending: dict[str, asyncio.Future] = {}
+ resume_token = "autonomy_pause:timeout:1"
+ loop = asyncio.get_running_loop()
+ original: asyncio.Future[bool] = loop.create_future()
+ pending[resume_token] = original
+
+ # Mirror chat.py: pre-registration checks ``not in``.
+ if resume_token and resume_token not in pending:
+ pending[resume_token] = loop.create_future()
+
+ assert pending[resume_token] is original # Unchanged.
+
+ @pytest.mark.asyncio
+ async def test_full_flow_pre_register_then_handler_awaits(self):
+ """End-to-end: pre-register → resolve from WS loop → handler awaits
+ and returns True.
+
+ Validates the complete U6 fix: Future is pre-registered before the
+ event is sent, the WS loop resolves it, and _resume_handler returns
+ the resolved value without creating a new Future.
+ """
+ pending: dict[str, asyncio.Future] = {}
+ resume_token = "autonomy_pause:timeout:5"
+ event_data = {"resume_token": resume_token, "reason": "timeout"}
+
+ # 1. autonomy_paused event handler: pre-register BEFORE send_json.
+ rt = event_data.get("resume_token")
+ if rt and rt not in pending:
+ loop = asyncio.get_running_loop()
+ pending[rt] = loop.create_future()
+
+ # 2. WS loop resolves the Future (resume message arrives).
+ fut = pending[rt]
+ if not fut.done():
+ fut.set_result(True)
+
+ # 3. _resume_handler: look up pre-registered Future, await it.
+ handler_future = pending.get(rt)
+ assert handler_future is not None
+ result = await asyncio.wait_for(handler_future, timeout=1.0)
+
+ assert result is True
+ # 4. Cleanup (finally block in _resume_handler).
+ pending.pop(rt, None)
+ assert rt not in pending
diff --git a/tests/unit/test_lead_reflection_retrieval.py b/tests/unit/test_lead_reflection_retrieval.py
index 1554358..f6f6b94 100644
--- a/tests/unit/test_lead_reflection_retrieval.py
+++ b/tests/unit/test_lead_reflection_retrieval.py
@@ -240,7 +240,9 @@ class TestDecomposeWithReflection:
)
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
- await orchestrator._decompose_task(team.lead_expert, "test task")
+ await orchestrator._decompose_task(
+ team.lead_expert, "test task", trigger_reason="verify_retry"
+ )
# Verify retrieve was called
reflexion.retrieve_prompt_reflection.assert_awaited_once()
@@ -278,7 +280,9 @@ class TestDecomposeWithReflection:
reflexion.retrieve_prompt_reflection = AsyncMock(return_value=None)
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
- await orchestrator._decompose_task(team.lead_expert, "test task")
+ await orchestrator._decompose_task(
+ team.lead_expert, "test task", trigger_reason="verify_retry"
+ )
call_kwargs = gw.chat.await_args.kwargs
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
@@ -296,7 +300,9 @@ class TestDecomposeWithReflection:
reflexion.retrieve_prompt_reflection = AsyncMock(side_effect=RuntimeError("search failed"))
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
- await orchestrator._decompose_task(team.lead_expert, "test task")
+ await orchestrator._decompose_task(
+ team.lead_expert, "test task", trigger_reason="verify_retry"
+ )
# Default prompt used despite retrieval failure
call_kwargs = gw.chat.await_args.kwargs
@@ -317,9 +323,126 @@ class TestDecomposeWithReflection:
)
orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
- await orchestrator._decompose_task(team.lead_expert, "test task")
+ await orchestrator._decompose_task(
+ team.lead_expert, "test task", trigger_reason="verify_retry"
+ )
call_kwargs = gw.chat.await_args.kwargs
messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
assert "Historical Reflection" not in prompt_content
+
+
+# ── U4/KTD5: retrieve_prompt_reflection trigger gating ────────────────
+
+
+class TestDecomposeTriggerGating:
+ """U4/KTD5: _decompose_task only calls retrieve_prompt_reflection when
+ trigger_reason is one of {verify_retry, schema_validation, loop_detection}.
+ Non-trigger decompositions skip retrieval entirely (default prompt)."""
+
+ def _make_orchestrator_with_reflexion(
+ self, retrieve_return: object = None, retrieve_side_effect: object | None = None
+ ) -> tuple[TeamOrchestrator, MagicMock, MagicMock]:
+ """Build orchestrator + gateway + reflexion mock.
+
+ Returns (orchestrator, gateway, reflexion). retrieve_side_effect, if
+ set, takes precedence over retrieve_return.
+ """
+ team = _make_team_with_experts()
+ gw = _make_llm_gateway_mock()
+ team._experts["lead"].agent._llm_gateway = gw
+
+ reflexion = MagicMock(spec=ReflexionEngine)
+ if retrieve_side_effect is not None:
+ reflexion.retrieve_prompt_reflection = AsyncMock(side_effect=retrieve_side_effect)
+ else:
+ reflexion.retrieve_prompt_reflection = AsyncMock(return_value=retrieve_return)
+
+ orchestrator = TeamOrchestrator(team, reflexion_engine=reflexion)
+ return orchestrator, gw, reflexion
+
+ @pytest.mark.asyncio
+ async def test_retrieves_on_verify_retry_trigger(self):
+ orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
+ retrieve_return={
+ "improved_prompt": "IMPROVED",
+ "score": 0.8,
+ "reflection": "r",
+ "version": 1,
+ "task_hash": "abc",
+ }
+ )
+ await orchestrator._decompose_task(
+ orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
+ )
+ reflexion.retrieve_prompt_reflection.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ async def test_retrieves_on_schema_validation_trigger(self):
+ orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(retrieve_return=None)
+ await orchestrator._decompose_task(
+ orchestrator._team.lead_expert, "task", trigger_reason="schema_validation"
+ )
+ reflexion.retrieve_prompt_reflection.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ async def test_retrieves_on_loop_detection_trigger(self):
+ orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(retrieve_return=None)
+ await orchestrator._decompose_task(
+ orchestrator._team.lead_expert, "task", trigger_reason="loop_detection"
+ )
+ reflexion.retrieve_prompt_reflection.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ async def test_skips_retrieval_when_trigger_is_none(self):
+ """Fresh decomposition (trigger_reason=None) → no retrieval."""
+ orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
+ retrieve_return={"improved_prompt": "would-be-used", "score": 0.9}
+ )
+ await orchestrator._decompose_task(orchestrator._team.lead_expert, "task")
+ reflexion.retrieve_prompt_reflection.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_skips_retrieval_on_invalid_trigger_reason(self):
+ """trigger_reason not in the allowed set → no retrieval."""
+ orchestrator, _, reflexion = self._make_orchestrator_with_reflexion(
+ retrieve_return={"improved_prompt": "would-be-used", "score": 0.9}
+ )
+ await orchestrator._decompose_task(
+ orchestrator._team.lead_expert, "task", trigger_reason="random"
+ )
+ reflexion.retrieve_prompt_reflection.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_triggered_retrieval_failure_degrades_to_default_prompt(self):
+ """trigger_reason valid but retrieve raises → default prompt (non-blocking)."""
+ orchestrator, gw, reflexion = self._make_orchestrator_with_reflexion(
+ retrieve_side_effect=RuntimeError("search down")
+ )
+ await orchestrator._decompose_task(
+ orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
+ )
+ # Retrieval was attempted (triggered) but failed — default prompt used.
+ reflexion.retrieve_prompt_reflection.assert_awaited_once()
+ call_kwargs = gw.chat.await_args.kwargs
+ messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
+ prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
+ assert "Historical Reflection" not in prompt_content
+
+ @pytest.mark.asyncio
+ async def test_skips_retrieval_when_no_reflexion_engine_even_with_trigger(self):
+ """No reflexion_engine → retrieval skipped regardless of trigger_reason."""
+ team = _make_team_with_experts()
+ gw = _make_llm_gateway_mock()
+ team._experts["lead"].agent._llm_gateway = gw
+
+ orchestrator = TeamOrchestrator(team, reflexion_engine=None)
+ await orchestrator._decompose_task(
+ orchestrator._team.lead_expert, "task", trigger_reason="verify_retry"
+ )
+ # Default prompt — no Historical Reflection section.
+ call_kwargs = gw.chat.await_args.kwargs
+ messages = call_kwargs.get("messages") or gw.chat.await_args.args[0]
+ prompt_content = messages[0]["content"] if isinstance(messages, list) else str(messages)
+ assert "Historical Reflection" not in prompt_content
diff --git a/tests/unit/test_react_autonomy.py b/tests/unit/test_react_autonomy.py
index bb56f56..53d622f 100644
--- a/tests/unit/test_react_autonomy.py
+++ b/tests/unit/test_react_autonomy.py
@@ -300,3 +300,210 @@ class TestAutonomyConstructorWiring:
)
assert engine._autonomy_mode is True
assert engine._dangerous_tools_config is cfg
+
+
+# ---------------------------------------------------------------------------
+# U7: _track_tool_result_for_autonomy denied parameter
+# ---------------------------------------------------------------------------
+
+
+class TestTrackToolResultDenied:
+ """U7: user denial (denied=True) does NOT increment the failure counter.
+
+ Verifies:
+ - denied=True leaves _consecutive_failures unchanged
+ - denied=False with an error result still increments (backward compat)
+ - denied=False with a permission_denied result still increments
+ - denied=True with a permission_denied result does NOT increment
+ """
+
+ def test_denied_does_not_increment(self):
+ """denied=True with an error result leaves the counter unchanged."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(),
+ autonomy_mode=True,
+ )
+ engine._consecutive_failures = 1
+ engine._track_tool_result_for_autonomy({"error": "boom", "is_error": True}, denied=True)
+ assert engine._consecutive_failures == 1 # Unchanged
+
+ def test_denied_permission_denied_does_not_increment(self):
+ """denied=True with a permission_denied result leaves the counter
+ unchanged — this is the exact scenario U7 fixes (user rejected
+ confirmation)."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(),
+ autonomy_mode=True,
+ )
+ engine._consecutive_failures = 2
+ engine._track_tool_result_for_autonomy(
+ {
+ "output": "",
+ "exit_code": 126,
+ "is_error": True,
+ "error_type": "permission_denied",
+ "message": "用户拒绝执行命令",
+ },
+ denied=True,
+ )
+ assert engine._consecutive_failures == 2 # Unchanged
+
+ def test_not_denied_error_increments(self):
+ """denied=False (default) with an error result still increments
+ (backward compat)."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(),
+ autonomy_mode=True,
+ )
+ engine._consecutive_failures = 1
+ engine._track_tool_result_for_autonomy({"error": "boom"})
+ assert engine._consecutive_failures == 2
+
+ def test_not_denied_permission_denied_increments(self):
+ """denied=False with a permission_denied result increments — without
+ the denied flag, a permission_denied result is treated as a failure.
+ This confirms the denied flag is required to skip counting."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(),
+ autonomy_mode=True,
+ )
+ engine._consecutive_failures = 0
+ engine._track_tool_result_for_autonomy(
+ {"error_type": "permission_denied", "is_error": True}
+ )
+ assert engine._consecutive_failures == 1
+
+ def test_denied_does_not_reset_counter(self):
+ """denied=True leaves the counter at whatever value it had — it
+ neither increments nor resets. This matters when denials interleave
+ with real failures."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(),
+ autonomy_mode=True,
+ )
+ # One real failure, then a denial.
+ engine._consecutive_failures = 0
+ engine._track_tool_result_for_autonomy({"error": "boom"})
+ assert engine._consecutive_failures == 1
+ engine._track_tool_result_for_autonomy({"error_type": "permission_denied"}, denied=True)
+ assert engine._consecutive_failures == 1 # Still 1, not reset to 0
+
+ def test_denied_no_effect_when_not_autonomy(self):
+ """denied=True is a no-op when autonomy mode is off (same as the
+ base behavior — tracking is skipped entirely)."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(),
+ autonomy_mode=False,
+ )
+ engine._consecutive_failures = 0
+ engine._track_tool_result_for_autonomy({"error": "boom"}, denied=True)
+ assert engine._consecutive_failures == 0
+
+
+# ---------------------------------------------------------------------------
+# U7: denial does not trigger autonomy_paused
+# ---------------------------------------------------------------------------
+
+
+class TestDenialDoesNotTriggerPause:
+ """U7: a sequence of user denials does NOT trigger autonomy_paused
+ (reason=consecutive_failures), while a sequence of real failures does.
+
+ Uses _detect_autonomy_pause to check whether a pause would fire.
+ """
+
+ @pytest.mark.asyncio
+ async def test_three_denials_no_pause(self):
+ """3 consecutive denials (denied=True) do NOT trigger
+ autonomy_paused — the counter never reaches the threshold."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(
+ autonomy_timeout_minutes=0, # Disabled
+ max_consecutive_failures=3,
+ ),
+ autonomy_mode=True,
+ )
+ engine._autonomy_started_at = 0 # Disable timeout check
+ denied_result = {
+ "error_type": "permission_denied",
+ "is_error": True,
+ "output": "",
+ }
+ # Simulate 3 denials — counter must stay at 0.
+ for _ in range(3):
+ engine._track_tool_result_for_autonomy(denied_result, denied=True)
+ assert engine._consecutive_failures == 0
+ # No pause should fire.
+ pause_info = engine._detect_autonomy_pause(step=1, progress={})
+ assert pause_info is None
+
+ @pytest.mark.asyncio
+ async def test_three_failures_trigger_pause(self):
+ """3 consecutive failures (denied=False) DO trigger autonomy_paused
+ — confirms the threshold still works for real failures."""
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(
+ autonomy_timeout_minutes=0, # Disabled
+ max_consecutive_failures=3,
+ ),
+ autonomy_mode=True,
+ )
+ engine._autonomy_started_at = 0 # Disable timeout check
+ error_result = {"error": "tool crashed"}
+ # Simulate 3 real failures.
+ for _ in range(3):
+ engine._track_tool_result_for_autonomy(error_result)
+ assert engine._consecutive_failures == 3
+ # Pause should fire with reason=consecutive_failures.
+ pause_info = engine._detect_autonomy_pause(step=1, progress={})
+ assert pause_info is not None
+ reason, _token, _data = pause_info
+ assert reason == "consecutive_failures"
+
+ @pytest.mark.asyncio
+ async def test_mixed_denials_and_failures_no_false_pause(self):
+ """Interleaved denials and failures: denials don't increment, so
+ 2 failures + 5 denials + 1 failure = 3 failures total → pause.
+
+ Confirms denials neither help nor hurt the failure count — only
+ real failures count toward the threshold.
+ """
+ engine = ReActEngine(
+ llm_gateway=MagicMock(),
+ dangerous_tools_config=DangerousToolsConfig(
+ autonomy_timeout_minutes=0,
+ max_consecutive_failures=3,
+ ),
+ autonomy_mode=True,
+ )
+ engine._autonomy_started_at = 0
+ error_result = {"error": "crash"}
+ denied_result = {"error_type": "permission_denied", "is_error": True}
+
+ # 2 real failures.
+ engine._track_tool_result_for_autonomy(error_result)
+ engine._track_tool_result_for_autonomy(error_result)
+ assert engine._consecutive_failures == 2
+
+ # 5 denials — counter must stay at 2.
+ for _ in range(5):
+ engine._track_tool_result_for_autonomy(denied_result, denied=True)
+ assert engine._consecutive_failures == 2
+ # No pause yet (below threshold).
+ assert engine._detect_autonomy_pause(step=1, progress={}) is None
+
+ # 1 more real failure → threshold reached.
+ engine._track_tool_result_for_autonomy(error_result)
+ assert engine._consecutive_failures == 3
+ pause_info = engine._detect_autonomy_pause(step=1, progress={})
+ assert pause_info is not None
+ reason, _token, _data = pause_info
+ assert reason == "consecutive_failures"
diff --git a/tests/unit/test_reflexion_persist.py b/tests/unit/test_reflexion_persist.py
index 14fe06f..05d7654 100644
--- a/tests/unit/test_reflexion_persist.py
+++ b/tests/unit/test_reflexion_persist.py
@@ -74,6 +74,9 @@ class TestStorePromptReflection:
mem._session_factory.return_value.__aenter__ = AsyncMock(return_value=mock_db)
mem._session_factory.return_value.__aexit__ = AsyncMock(return_value=None)
mem._episodic_model = MagicMock()
+ # U2: store_prompt_reflection now calls cleanup_expired lazily; provide
+ # an AsyncMock so the real code path can await it without DB side effects.
+ mem.cleanup_expired = AsyncMock(return_value=0)
# Track the ORM entry passed to db.add for assertions
mem._mock_db = mock_db
return mem
@@ -434,3 +437,78 @@ class TestMultiVersionCoexistence:
# All three ORM entries created via db.add
assert mem._mock_db.add.call_count == 3
assert mem._mock_db.commit.await_count == 3
+
+
+# ── U2: lazy cleanup_expired on store ─────────────────────────────────
+
+
+class TestLazyCleanupOnStore:
+ """U2: store_prompt_reflection triggers cleanup_expired with 10%
+ probability (lazy, fire-and-forget). Cleanup failure never blocks the
+ store; cleanup is always filtered to task_type='prompt_reflection'."""
+
+ def _make_persistable_mock(self) -> MagicMock:
+ # Reuse the persistable mock (now includes async cleanup_expired).
+ return TestStorePromptReflection()._make_persistable_mock()
+
+ @pytest.mark.asyncio
+ async def test_lazy_cleanup_triggered_below_threshold(self):
+ """random.random() < 0.1 → cleanup_expired awaited once."""
+ mem = self._make_persistable_mock()
+ with patch("agentkit.memory.episodic.random.random", return_value=0.05):
+ await EpisodicMemory.store_prompt_reflection(
+ mem, task_input="test", reflection="r", improved_prompt="p"
+ )
+ mem.cleanup_expired.assert_awaited_once_with(task_type="prompt_reflection")
+
+ @pytest.mark.asyncio
+ async def test_lazy_cleanup_not_triggered_at_or_above_threshold(self):
+ """random.random() >= 0.1 → cleanup_expired not awaited."""
+ mem = self._make_persistable_mock()
+ with patch("agentkit.memory.episodic.random.random", return_value=0.9):
+ await EpisodicMemory.store_prompt_reflection(
+ mem, task_input="test", reflection="r", improved_prompt="p"
+ )
+ mem.cleanup_expired.assert_not_awaited()
+
+ @pytest.mark.asyncio
+ async def test_lazy_cleanup_failure_does_not_block_store(self):
+ """cleanup_expired raises → store still returns a valid key."""
+ mem = self._make_persistable_mock()
+ mem.cleanup_expired = AsyncMock(side_effect=RuntimeError("cleanup boom"))
+ with patch("agentkit.memory.episodic.random.random", return_value=0.05):
+ key = await EpisodicMemory.store_prompt_reflection(
+ mem, task_input="test", reflection="r", improved_prompt="p"
+ )
+ assert key is not None # store succeeded despite cleanup failure
+ mem.cleanup_expired.assert_awaited_once()
+
+ @pytest.mark.asyncio
+ async def test_lazy_cleanup_filters_to_prompt_reflection_task_type(self):
+ """cleanup_expired is always called with task_type='prompt_reflection'."""
+ mem = self._make_persistable_mock()
+ with patch("agentkit.memory.episodic.random.random", return_value=0.05):
+ await EpisodicMemory.store_prompt_reflection(
+ mem, task_input="test", reflection="r", improved_prompt="p"
+ )
+ # Keyword arg, never None — prevents purging other task types.
+ assert mem.cleanup_expired.await_args.kwargs == {"task_type": "prompt_reflection"}
+
+ @pytest.mark.asyncio
+ async def test_lazy_cleanup_triggered_at_least_once_over_many_calls(self):
+ """10+ calls with mixed random draws → cleanup_expired called at least once."""
+ mem = self._make_persistable_mock()
+ # 9 non-triggering draws (0.9) then 1 triggering (0.05).
+ draws = [0.9] * 9 + [0.05]
+ with patch("agentkit.memory.episodic.random.random", side_effect=draws):
+ for i in range(10):
+ await EpisodicMemory.store_prompt_reflection(
+ mem,
+ task_input=f"task {i}",
+ reflection="r",
+ improved_prompt="p",
+ )
+ assert mem.cleanup_expired.await_count >= 1
+ # Every cleanup call uses the prompt_reflection filter.
+ for call in mem.cleanup_expired.await_args_list:
+ assert call.kwargs == {"task_type": "prompt_reflection"}
diff --git a/tests/unit/test_team_parallel.py b/tests/unit/test_team_parallel.py
index 8bd8f1e..cc27298 100644
--- a/tests/unit/test_team_parallel.py
+++ b/tests/unit/test_team_parallel.py
@@ -10,6 +10,7 @@ Covers:
from __future__ import annotations
+import asyncio
import json
from unittest.mock import AsyncMock, MagicMock
@@ -17,6 +18,7 @@ import pytest
from agentkit.core.handoff_transport import InProcessHandoffTransport
from agentkit.core.protocol import TaskResult, TaskStatus
+from agentkit.core.shared_workspace import SharedWorkspace
from agentkit.experts.config import ExpertConfig
from agentkit.experts.expert import Expert
from agentkit.experts.orchestrator import TeamOrchestrator
@@ -462,3 +464,106 @@ class TestExecuteParallelSubtasks:
# phase_results has all 3
phase_results: dict = result["phase_results"]
assert len(phase_results) == 3
+
+
+# ── U3: SharedWorkspace.lock_section TOCTOU protection ───────────────
+
+
+class TestLockSectionToctou:
+ """U3/KTD3: SharedWorkspace.lock_section guards parallel subtask
+ read-then-write critical sections against TOCTOU races.
+
+ lock_section is the async-context-manager form of lock/unlock: it polls
+ lock() until acquired (or timeout), holds for the block, and releases on
+ exit. Used by _finalize_phase to wrap workspace.write().
+ """
+
+ @pytest.mark.asyncio
+ async def test_lock_section_acquires_and_releases(self):
+ """Lock acquired on enter, released on exit — next caller can acquire."""
+ ws = SharedWorkspace()
+ async with ws.lock_section("resource", "agent_a", timeout=1.0):
+ # Held by agent_a — agent_b cannot acquire via raw lock()
+ assert await ws.lock("resource", agent_id="agent_b") is False
+ # Released after block — agent_b can now acquire
+ assert await ws.lock("resource", agent_id="agent_b") is True
+ await ws.unlock("resource", "agent_b")
+
+ @pytest.mark.asyncio
+ async def test_lock_section_calls_lock_and_unlock(self):
+ """lock_section delegates to lock() on enter and unlock() on exit."""
+ ws = SharedWorkspace()
+ lock_spy = AsyncMock(wraps=ws.lock)
+ unlock_spy = AsyncMock(wraps=ws.unlock)
+ ws.lock = lock_spy
+ ws.unlock = unlock_spy
+
+ async with ws.lock_section("k", "agent_a", timeout=1.0):
+ pass
+
+ lock_spy.assert_awaited_with("k", "agent_a", timeout=1.0)
+ unlock_spy.assert_awaited_with("k", "agent_a")
+
+ @pytest.mark.asyncio
+ async def test_lock_section_raises_timeout_when_held(self):
+ """Lock held by another agent → lock_section raises TimeoutError."""
+ ws = SharedWorkspace()
+ await ws.lock("shared", "holder") # hold the lock
+
+ with pytest.raises(TimeoutError):
+ async with ws.lock_section("shared", "waiter", timeout=0.2):
+ pass # pragma: no cover — never reached
+
+ # Waiter never acquired, so holder still owns it.
+ assert await ws.lock("shared", agent_id="other") is False
+
+ @pytest.mark.asyncio
+ async def test_parallel_writes_to_different_keys_do_not_conflict(self):
+ """2 subtasks writing different keys concurrently both succeed."""
+ ws = SharedWorkspace()
+ results: list[str] = []
+
+ async def writer(key: str, agent: str) -> None:
+ async with ws.lock_section(key, agent, timeout=2.0):
+ await ws.write(key, f"val-{agent}", agent)
+ results.append(agent)
+
+ await asyncio.gather(writer("k1", "a"), writer("k2", "b"))
+
+ assert set(results) == {"a", "b"}
+ assert (await ws.read("k1"))["value"] == "val-a"
+ assert (await ws.read("k2"))["value"] == "val-b"
+
+ @pytest.mark.asyncio
+ async def test_same_key_writes_serialize(self):
+ """2 subtasks writing the same key: the second waits for the first
+ to release — enter/exit events never interleave."""
+ ws = SharedWorkspace()
+ log: list[str] = []
+
+ async def writer(name: str) -> None:
+ async with ws.lock_section("shared_key", name, timeout=2.0):
+ log.append(f"{name}-enter")
+ await asyncio.sleep(0.05)
+ log.append(f"{name}-exit")
+
+ await asyncio.gather(writer("a"), writer("b"))
+
+ # One fully completes before the other begins (serialized).
+ assert log in (
+ ["a-enter", "a-exit", "b-enter", "b-exit"],
+ ["b-enter", "b-exit", "a-enter", "a-exit"],
+ )
+
+ @pytest.mark.asyncio
+ async def test_lock_section_releases_on_exception(self):
+ """If the wrapped block raises, the lock is still released."""
+ ws = SharedWorkspace()
+
+ with pytest.raises(RuntimeError, match="boom"):
+ async with ws.lock_section("k", "agent_a", timeout=1.0):
+ raise RuntimeError("boom")
+
+ # Lock released despite exception — another agent can acquire.
+ assert await ws.lock("k", agent_id="agent_b") is True
+ await ws.unlock("k", "agent_b")