diff --git a/test/async_client_test.py b/test/async_client_test.py index 218e9da..c7036fe 100644 --- a/test/async_client_test.py +++ b/test/async_client_test.py @@ -510,11 +510,25 @@ async def test_async_sample(self, uid): await _complete_sandbox_task(cape.id, 'cape') await _complete_sandbox_task(triage.id, 'triage') + # Poll until BOTH sandbox deps read COMPLETED *and* the LLM report was + # auto-triggered — the exact postconditions asserted below, not a proxy. + # Keying off llm_report alone raced: the report auto-triggers on *any + # one* completed dep (the scan or a single sandbox), so a response can + # show it triggered while sandbox_cape still projects NOT_TRIGGERED (the + # per-task projection doesn't update atomically). This test drives both + # sandboxes to SUCCEEDED, so both projections reach COMPLETED; waiting on + # them directly closes the window. See sync test_sample. _PRE_TRIGGER = {None, 'NOT_TRIGGERED', 'WAITING_FOR_OTHER_TASKS'} result = await api.sample(sha) for _ in range(90): result = await api.sample(sha) - if result.tasks.get('llm_report', {}).get('requested_status') not in _PRE_TRIGGER: + tasks = result.tasks or {} + sandboxes_completed = all( + tasks.get(f'sandbox_{s}', {}).get('requested_status') == 'COMPLETED' + for s in ('cape', 'triage') + ) + llm_triggered = tasks.get('llm_report', {}).get('requested_status') not in _PRE_TRIGGER + if sandboxes_completed and llm_triggered: break await asyncio.sleep(1) assert isinstance(result.artifact_instance, dict) diff --git a/test/client_scan_test.py b/test/client_scan_test.py index a099797..44d7dc0 100644 --- a/test/client_scan_test.py +++ b/test/client_scan_test.py @@ -880,18 +880,35 @@ def test_sample(self): _complete_sandbox_task(cape.id, 'cape') _complete_sandbox_task(triage.id, 'triage') - # Poll the sample until the LLM report has been auto-triggered. The view - # triggers it once a sandbox dep is COMPLETED, so the report's status - # moves NOT_TRIGGERED/WAITING_FOR_OTHER_TASKS -> PENDING (then FAILED here, - # since e2e has no OPENAI_API_KEY). We key off requested_status: the - # requested_id stays null until a report actually renders, which can't - # happen without an LLM. Reaching a triggered status also confirms the - # sandbox deps completed. + # Poll the sample until BOTH sandbox deps read COMPLETED *and* the LLM + # report has been auto-triggered — i.e. wait on the exact postconditions + # asserted below, not a proxy for them. The report's requested_status moves + # NOT_TRIGGERED/WAITING_FOR_OTHER_TASKS -> PENDING (then FAILED here, since + # e2e has no OPENAI_API_KEY); requested_id stays null until a report + # actually renders (impossible without an LLM), so requested_status is the + # trigger signal. + # + # Keying the loop off llm_report alone raced: the report auto-triggers as + # soon as *any one* dependency completes (the scan or a single sandbox), so + # a response can show it triggered while sandbox_cape still projects + # NOT_TRIGGERED — the delayed COLLECTING_DATA->SUCCEEDED transition (see + # _complete_sandbox_task) not yet folded into the per-task projection, which + # doesn't update atomically. That mismatch was the flake. This test drives + # BOTH sandboxes to SUCCEEDED before polling, so both projections do reach + # COMPLETED; waiting on them directly (not on the report as a proxy) closes + # the window. A 90x1s timeout here therefore means a dependency never + # settled — the scan's bounty window or a sandbox — not this loop. _PRE_TRIGGER = {None, 'NOT_TRIGGERED', 'WAITING_FOR_OTHER_TASKS'} result = api.sample(sha) for _ in range(90): result = api.sample(sha) - if result.tasks.get('llm_report', {}).get('requested_status') not in _PRE_TRIGGER: + tasks = result.tasks or {} + sandboxes_completed = all( + tasks.get(f'sandbox_{s}', {}).get('requested_status') == 'COMPLETED' + for s in ('cape', 'triage') + ) + llm_triggered = tasks.get('llm_report', {}).get('requested_status') not in _PRE_TRIGGER + if sandboxes_completed and llm_triggered: break time.sleep(1) assert isinstance(result.artifact_instance, dict)