TEST: Harden Copilot allocation cancellation observation

tests/unit/prompt_target/target/test_github_copilot_target.py

@@ -1985,15 +1985,23 @@ async def create_session_async(*, session_id: str = "sdk-generated-session-id",
         pytest.param("cancelled-error", id="first-delete-cancelled-error"),
     ],
 )
+@pytest.mark.parametrize(
+    "watchdog_expires",
+    [pytest.param(False, id="normal-observation"), pytest.param(True, id="expired-watchdog")],
+)
 async def test_normalizer_deletes_owned_session_when_creation_is_cancelled_after_allocation_async(
     *,
     client: NonCallableMagicMock,
     sqlite_instance: MemoryInterface,
     first_delete_failure: str,
+    watchdog_expires: bool,
 ) -> None:
     session = client.create_session.return_value
     sessions = {"unrelated-session-id"}
     allocated = asyncio.Event()
+    allow_allocation = asyncio.Event()
+    delete_started = asyncio.Event()
+    allow_delete = asyncio.Event()
     conversation_id = str(uuid4())
     cancellation_message = "cancel caller after session allocation"
     caller_cancellation: asyncio.CancelledError | None = None
@@ -2005,6 +2013,7 @@ async def test_normalizer_deletes_owned_session_when_creation_is_cancelled_after
 
     async def create_session_async(*, session_id: str = "sdk-generated-session-id", **_kwargs: Any) -> None:
         nonlocal caller_cancellation
+        await allow_allocation.wait()
         sessions.add(session_id)
         allocated.set()
         try:
@@ -2014,6 +2023,8 @@ async def create_session_async(*, session_id: str = "sdk-generated-session-id",
             raise
 
     async def delete_session_async(session_id: str) -> None:
+        delete_started.set()
+        await allow_delete.wait()
         if client.delete_session.await_count == 1 and cleanup_failure is not None:
             raise cleanup_failure
         sessions.remove(session_id)
@@ -2036,13 +2047,38 @@ async def delete_session_async(session_id: str) -> None:
                 conversation_id=conversation_id,
             )
         )
+        allocated_waiter = asyncio.create_task(allocated.wait())
+        delete_started_waiter = asyncio.create_task(delete_started.wait())
         try:
-            await asyncio.wait_for(allocated.wait(), timeout=2.0)
+            if watchdog_expires:
+                with pytest.raises(TimeoutError, match="test watchdog expired"):
+                    await wait_for_completion_async(future=allocated_waiter, timeout=0)
+                assert not allocated_waiter.done()
+                assert not request_task.done()
+                assert request_task.cancelling() == 0
+                assert sessions == {"unrelated-session-id"}
+                client.delete_session.assert_not_awaited()
+
+            allow_allocation.set()
+            await wait_for_completion_async(future=allocated_waiter)
             request_task.cancel(cancellation_message)
+            await wait_for_completion_async(future=delete_started_waiter)
+            if watchdog_expires:
+                with pytest.raises(TimeoutError, match="test watchdog expired"):
+                    await wait_for_completion_async(future=request_task, timeout=0)
+                assert not request_task.done()
+                assert request_task.cancelling() == 1
+                assert sessions == {"unrelated-session-id", client.create_session.await_args.kwargs["session_id"]}
+                client.delete_session.assert_awaited_once()
+                client.stop.assert_not_awaited()
+
+            allow_delete.set()
             with pytest.raises(asyncio.CancelledError) as cancellation_error:
-                await asyncio.wait_for(request_task, timeout=2.0)
+                await wait_for_completion_async(future=request_task)
         finally:
-            await _cancel_tasks_async(request_task)
+            allow_allocation.set()
+            allow_delete.set()
+            await _cancel_tasks_async(request_task, allocated_waiter, delete_started_waiter)
 
         add_conversation.assert_awaited_once()
         add_message.assert_not_awaited()

tests/unit/scenario/airt/test_psychosocial.py

@@ -151,22 +151,29 @@ async def test_crescendo_construction_keeps_initialization_responsive_async(
     def build(**kwargs: Any) -> AttackTechnique:
         loop.call_soon_threadsafe(entered.set)
         assert threading.get_ident() != backend_thread
-        if not release.wait(5):
-            raise TimeoutError("Crescendo construction was not released.")
+        release.wait()
         return original_build(**kwargs)
 
     with (
+        patch.object(DatasetAttackConfiguration, "prepare_async", new_callable=AsyncMock),
         _patch_base_seed_groups(_make_seed_groups()),
         patch.object(scenario, "_build_crescendo_technique", side_effect=build),
     ):
         initialize = asyncio.create_task(scenario.initialize_async())
         try:
-            await asyncio.wait_for(entered.wait(), 5)
+            # Readiness includes setup and thread scheduling, not just event-loop responsiveness.
+            await asyncio.wait_for(entered.wait(), 30)
+            heartbeat = asyncio.Event()
+            loop.call_soon(heartbeat.set)
+            await asyncio.wait_for(heartbeat.wait(), 5)
             assert not initialize.done()
+            assert not release.is_set()
             release.set()
-            await asyncio.wait_for(initialize, 5)
+            await asyncio.wait_for(initialize, 30)
         finally:
             release.set()
+            if not initialize.done():
+                initialize.cancel()
             await asyncio.gather(initialize, return_exceptions=True)
     assert [attack.atomic_attack_name for attack in scenario._atomic_attacks] == ["imminent_crisis_crescendo"]