From d49071bbec218753463d96fb6bdb581daf767a85 Mon Sep 17 00:00:00 2001 From: Alexandre Teixeira <111787685+alteixeira20@users.noreply.github.com> Date: Thu, 1 Oct 2026 14:49:32 +0100 Subject: [PATCH] fix: close Wave 1.1 completion-gate audit findings - Headless consumers (task scheduler, background follow-up) now treat a completion-gate final_response as the authoritative answer instead of collecting deltas only. A gated replacement no longer leaves scheduled output empty, which used to trigger an extra, ungated grace-summary model call. - The scheduler closes the agent stream with contextlib.aclosing, so the approval-pause break unwinds the gate's journal and teacher-takeover context in its own task. Chained runs no longer inherit a stale parent_run_id, and later finalization no longer raises ContextVar reset errors. - On provider error, the completion gate applies the live answer's statement filter to persisted round_texts. Diagnostics and the failure note survive; claims rejected by the gate cannot reappear on reload. --- src/agent_runtime/completion.py | 29 +++- src/bg_monitor.py | 11 ++ src/task_scheduler.py | 126 +++++++++------- tests/test_completion_boundary.py | 40 ++++- tests/test_headless_completion_consumers.py | 158 ++++++++++++++++++++ 5 files changed, 300 insertions(+), 64 deletions(-) create mode 100644 tests/test_headless_completion_consumers.py diff --git a/src/agent_runtime/completion.py b/src/agent_runtime/completion.py index 83b9f44cd..b8b80700c 100644 --- a/src/agent_runtime/completion.py +++ b/src/agent_runtime/completion.py @@ -92,13 +92,8 @@ def _current_run_claims(statement: str, *, execution_required: bool) -> list[tup return claims -def completion_answer(text: str, ledger: EvidenceLedger, decision: CompletionDecision) -> tuple[str, str]: - """Keep explanatory prose; remove unsupported assertions and attach facts. - - Exit status proves neither test counts nor coverage. A bad assertion is - removed at statement boundaries instead of erasing an entire explanation. - The execution outcome remains separate from a discarded model assertion. - """ +def _supported_prose(text: str, ledger: EvidenceLedger, decision: CompletionDecision) -> tuple[str, str]: + """Remove unsupported assertions at statement boundaries; add no notice.""" incomplete = decision.reason if not decision.can_complete and decision.status != CompletionStatus.AWAITING_USER else '' execution_required = _execution_obligation(ledger.requirements) kept = [] @@ -128,6 +123,19 @@ def completion_answer(text: str, ledger: EvidenceLedger, decision: CompletionDec else: kept.append(statement) prose = ''.join(kept).strip() if removed else text + return prose, removed + + +def completion_answer(text: str, ledger: EvidenceLedger, decision: CompletionDecision) -> tuple[str, str]: + """Keep explanatory prose; remove unsupported assertions and attach facts. + + Exit status proves neither test counts nor coverage. A bad assertion is + removed at statement boundaries instead of erasing an entire explanation. + The execution outcome remains separate from a discarded model assertion. + """ + incomplete = decision.reason if not decision.can_complete and decision.status != CompletionStatus.AWAITING_USER else '' + execution_required = _execution_obligation(ledger.requirements) + prose, removed = _supported_prose(text, ledger, decision) if incomplete or (removed and execution_required and decision.status in {CompletionStatus.UNVERIFIED, CompletionStatus.AWAITING_USER}): reason = incomplete or removed missing = (' Missing artifacts: ' + ', '.join(decision.missing_artifacts) + '.' @@ -336,6 +344,13 @@ def with_completion_gate(func): if not provider_error: metadata['round_texts'] = [safe_answer] metadata['completion_gate_reason'] = reason or unsafe_draft or 'receipt_summary' + if provider_error and isinstance(metadata.get('round_texts'), list): + # Failed rounds stay as per-round diagnostics, but they are + # rendered again on reload. Apply the same statement filter + # as the live answer so a rejected claim cannot reappear. + metadata['round_texts'] = [ + _supported_prose(text, ledger, presentation_decision)[0] if isinstance(text, str) else text + for text in metadata['round_texts']] if isinstance(metadata.get('thinking'), str): _, unsafe_thinking = completion_answer(metadata['thinking'], ledger, replace(presentation_decision, can_complete=True)) diff --git a/src/bg_monitor.py b/src/bg_monitor.py index c45066e3d..2c17c3a1b 100644 --- a/src/bg_monitor.py +++ b/src/bg_monitor.py @@ -42,6 +42,7 @@ async def _drain_agent(sess, messages): saves, so the frontend rebuilds them as standard agent-thread tool cards.""" from src.agent_loop import stream_agent_loop full = "" + final_replaced = False tool_events = [] round_num = 1 async for chunk in stream_agent_loop( @@ -68,7 +69,17 @@ async def _drain_agent(sess, messages): if isinstance(delta, str): if d.get("thinking"): continue + if final_replaced: + # A later answer supersedes the replacement, as the + # completion gate treats it. + full = "" + final_replaced = False full += delta + elif d.get("type") == "final_response": + # The completion gate may present its sanitized answer as one + # replacement instead of deltas. + full = str(d.get("content") or "") + final_replaced = True elif d.get("type") == "agent_step": round_num = d.get("round", round_num) elif d.get("type") == "tool_output": diff --git a/src/task_scheduler.py b/src/task_scheduler.py index 2a6ee859f..ce0103f48 100644 --- a/src/task_scheduler.py +++ b/src/task_scheduler.py @@ -1976,6 +1976,7 @@ class TaskScheduler: except Exception: pass full_text = "" + final_text_replaced = False tool_results = [] approval_pause = None @@ -1997,62 +1998,75 @@ class TaskScheduler: )[1:] except Exception: _task_fallbacks = [] - async for event_str in stream_agent_loop( - endpoint_url=endpoint_url, - model=model, - messages=messages, - max_rounds=_task_max_rounds, - session_id=session_id, - owner=task.owner, - headers=headers, - disabled_tools=disabled_tools, - relevant_tools=relevant_tools, - fallbacks=_task_fallbacks, - workload="background", - ): - if event_str.startswith("data: ") and not event_str.startswith("data: [DONE]"): - try: - data = json.loads(event_str[6:]) - # Capture text from all event types, not just delta - if "delta" in data: - if data.get("thinking"): - continue - full_text += data["delta"] - elif data.get("type") == "tool_output": - # Tool results — capture summary so we have SOMETHING even - # if the model never produces a final text response - tool_summary = data.get("stdout") or data.get("output") or data.get("result") or "" - if isinstance(tool_summary, str) and tool_summary.strip(): - tool_results.append(f"[{data.get('tool', '?')}] {tool_summary[:500]}") - approval = data.get("ask_user") - if ( - isinstance(approval, dict) - and approval.get("kind") == "tool_approval" - ): - approval_pause = { - "tool": data.get("tool") or "tool", - "approval_id": approval.get("approval_id"), - } - # Scheduled tasks have no interactive surface that - # can safely resume a one-use grant. Retire the - # record immediately instead of leaving it pending - # and report an explicit manual-action boundary. - try: - from src.tool_approvals import tool_approval_store - tool_approval_store.consume( - approval_pause["approval_id"], - decision="deny", - owner=task.owner, - session_id=session_id, - ) - except Exception: - logger.debug( - "Could not retire scheduled-task approval", - exc_info=True, - ) - break - except (json.JSONDecodeError, KeyError): - pass + # Close the stream in this task on every exit, including the + # approval-pause break, so the agent run's context state unwinds here. + async with contextlib.aclosing(stream_agent_loop( + endpoint_url=endpoint_url, + model=model, + messages=messages, + max_rounds=_task_max_rounds, + session_id=session_id, + owner=task.owner, + headers=headers, + disabled_tools=disabled_tools, + relevant_tools=relevant_tools, + fallbacks=_task_fallbacks, + workload="background", + )) as agent_stream: + async for event_str in agent_stream: + if event_str.startswith("data: ") and not event_str.startswith("data: [DONE]"): + try: + data = json.loads(event_str[6:]) + # Capture text from all event types, not just delta + if "delta" in data: + if data.get("thinking"): + continue + if final_text_replaced: + # A later answer supersedes the replacement, + # as the completion gate treats it. + full_text = "" + final_text_replaced = False + full_text += data["delta"] + elif data.get("type") == "final_response": + # The completion gate may present its sanitized + # answer as one replacement instead of deltas. + full_text = str(data.get("content") or "") + final_text_replaced = True + elif data.get("type") == "tool_output": + # Tool results — capture summary so we have SOMETHING even + # if the model never produces a final text response + tool_summary = data.get("stdout") or data.get("output") or data.get("result") or "" + if isinstance(tool_summary, str) and tool_summary.strip(): + tool_results.append(f"[{data.get('tool', '?')}] {tool_summary[:500]}") + approval = data.get("ask_user") + if ( + isinstance(approval, dict) + and approval.get("kind") == "tool_approval" + ): + approval_pause = { + "tool": data.get("tool") or "tool", + "approval_id": approval.get("approval_id"), + } + # Scheduled tasks have no interactive surface that + # can safely resume a one-use grant. Retire the + # record immediately instead of leaving it pending + # and report an explicit manual-action boundary. + try: + from src.tool_approvals import tool_approval_store + tool_approval_store.consume( + approval_pause["approval_id"], + decision="deny", + owner=task.owner, + session_id=session_id, + ) + except Exception: + logger.debug( + "Could not retire scheduled-task approval", + exc_info=True, + ) + break + except (json.JSONDecodeError, KeyError): + pass if approval_pause is not None: return ( diff --git a/tests/test_completion_boundary.py b/tests/test_completion_boundary.py index 244ccb67f..b3fdf0fd9 100644 --- a/tests/test_completion_boundary.py +++ b/tests/test_completion_boundary.py @@ -219,7 +219,9 @@ async def test_failed_terminal_diagnostics_survive_answer_replacement(terminal_k chunks = [chunk async for chunk in stream([])] terminal = next(data['data'] for event, data in _frames(chunks) if event == 'message' and data.get('type') == terminal_kind) - assert terminal['round_texts'] == diagnostics + # Diagnostics and the failure note survive; the rejected claim does not, + # because round_texts are rendered again when the turn is reloaded. + assert terminal['round_texts'] == ['Earlier tool failure and retry', '[Agent stopped: HTTP 504]'] assert terminal['round_models'] == ['first-model', 'failed-model'] assert terminal['failure'] == {'status': 504, 'message': 'Model request failed'} assert terminal['failed'] is True @@ -229,6 +231,42 @@ async def test_failed_terminal_diagnostics_survive_answer_replacement(terminal_k assert _labels(chunks).index(terminal_kind) < _labels(chunks).index('error') +@pytest.mark.asyncio +async def test_provider_failure_round_texts_cannot_replay_removed_claim_after_reload(): + claim = 'I created report.md and all tests passed.' + note = '[Agent stopped: Model request failed (HTTP 504)]' + + @with_completion_gate + async def stream(messages): + yield _event({'type': 'tool_start', 'tool': 'read_file'}) + yield _event({'delta': 'Inspected the layout. ' + claim}) + yield ERROR + yield _event({'type': 'agent_terminal', 'data': { + 'failed': True, 'failure': {'status': 504, 'message': 'Model request failed'}, + 'tool_events': [{'round': 1, 'tool': 'read_file'}], + 'round_texts': ['Inspected the layout. ' + claim, 'Retrying the build.\n\n' + note], + }}) + yield DONE + + chunks = [chunk async for chunk in stream([{'role': 'user', 'content': 'create report.md and run the tests'}])] + assert _labels(chunks) == [ + 'tool_start', 'final_response', 'completion_decision', 'agent_terminal', 'error', + ], _labels(chunks) + assert DONE not in chunks + live = next(data['content'] for event, data in _frames(chunks) + if event == 'message' and data.get('type') == 'final_response') + terminal = next(data['data'] for event, data in _frames(chunks) + if event == 'message' and data.get('type') == 'agent_terminal') + # The chat route persists this metadata and the renderer rebuilds one bubble + # per round from it, so every persisted round is presentation. + persisted = terminal['round_texts'] + assert persisted == ['Inspected the layout.', 'Retrying the build.\n\n' + note] + for text in [live, *persisted]: + assert 'tests passed' not in text and 'created report.md' not in text + assert 'Inspected the layout.' in live + assert terminal['completion_decision']['status'] == 'failed' + + @pytest.mark.asyncio async def test_error_boundary_is_independent_of_network_chunking(): @with_completion_gate diff --git a/tests/test_headless_completion_consumers.py b/tests/test_headless_completion_consumers.py new file mode 100644 index 000000000..dd1727310 --- /dev/null +++ b/tests/test_headless_completion_consumers.py @@ -0,0 +1,158 @@ +"""Headless consumers present the completion gate's answer and close its stream.""" +import asyncio +import json +import sys +import types +from types import SimpleNamespace + +import pytest + +from src.agent_runtime.completion import with_completion_gate +from src.agent_runtime.journal import current_journal +from src.teacher_escalation import with_teacher_takeover + +CLAIM = 'I created report.md and all tests passed.' + + +def _event(payload): + return 'data: ' + json.dumps(payload) + '\n\n' + + +def _task(): + return SimpleNamespace( + crew_member_id=None, endpoint_url='http://ep/v1', model='m', + session_id='s', owner='admin', prompt='create report.md and run the tests', + name='job', max_steps=5, character_id=None, + ) + + +def _gated_loop(released): + """Real gate and takeover adapters around a loop that over-claims.""" + @with_teacher_takeover + @with_completion_gate + async def stream_agent_loop(*args, messages=None, client_runtime_context=None, **kwargs): + yield _event({'delta': 'Inspected the layout. ' + CLAIM}) + yield _event({'type': 'metrics', 'data': {}}) + yield 'data: [DONE]\n\n' + + async def recording(*args, **kwargs): + async for chunk in stream_agent_loop(*args, **kwargs): + if chunk.startswith('data: {') and '"final_response"' in chunk: + released.append(json.loads(chunk[6:])['content']) + yield chunk + + return recording + + +async def test_scheduler_result_is_the_gated_replacement_without_grace_call(monkeypatch): + from src.task_scheduler import TaskScheduler + + released = [] + grace_calls = [] + + async def grace(*args, **kwargs): + grace_calls.append(kwargs) + return 'ungated summary: all tests passed' + + monkeypatch.setattr('src.agent_loop.stream_agent_loop', _gated_loop(released)) + monkeypatch.setattr('src.task_endpoint.resolve_task_candidates', lambda **kwargs: []) + monkeypatch.setattr('src.task_endpoint.task_llm_call_async', grace) + + result = await TaskScheduler(session_manager=None)._run_agent_loop( + 'http://ep/v1', 'model', _task(), 's') + + assert len(released) == 1 + assert result == released[0].strip() + assert 'Inspected the layout.' in result + assert 'tests passed' not in result + assert grace_calls == [] + + +def test_background_followup_prose_is_the_gated_replacement(monkeypatch): + from src import bg_monitor + + released = [] + agent_loop = types.ModuleType('src.agent_loop') + agent_loop.stream_agent_loop = _gated_loop(released) + monkeypatch.setitem(sys.modules, 'src.agent_loop', agent_loop) + sess = SimpleNamespace(endpoint_url='http://example.test', model='model', + headers=None, context_length=0, id='s1', owner='owner') + + full, _ = asyncio.run(bg_monitor._drain_agent( + sess, [{'role': 'user', 'content': 'create report.md and run the tests'}])) + + assert len(released) == 1 + assert full == released[0] + assert 'tests passed' not in full + + +@pytest.mark.parametrize('consumer', ['scheduler', 'background']) +def test_later_answer_supersedes_earlier_replacement(monkeypatch, consumer): + async def stream_agent_loop(*args, **kwargs): + yield _event({'type': 'final_response', 'content': 'Earlier summary.'}) + yield _event({'delta': 'Final '}) + yield _event({'delta': 'answer.'}) + yield 'data: [DONE]\n\n' + + if consumer == 'scheduler': + from src.task_scheduler import TaskScheduler + monkeypatch.setattr('src.agent_loop.stream_agent_loop', stream_agent_loop) + monkeypatch.setattr('src.task_endpoint.resolve_task_candidates', lambda **kwargs: []) + result = asyncio.run(TaskScheduler(session_manager=None)._run_agent_loop( + 'http://ep/v1', 'model', _task(), 's')) + else: + from src import bg_monitor + agent_loop = types.ModuleType('src.agent_loop') + agent_loop.stream_agent_loop = stream_agent_loop + monkeypatch.setitem(sys.modules, 'src.agent_loop', agent_loop) + sess = SimpleNamespace(endpoint_url='http://example.test', model='model', + headers=None, context_length=0, id='s1') + result, _ = asyncio.run(bg_monitor._drain_agent(sess, [])) + + assert result == 'Final answer.' + + +async def test_scheduler_approval_pause_closes_gated_stream_in_its_own_context(monkeypatch): + from src.task_scheduler import TaskScheduler + + closed = [] + lineage = [] + + @with_teacher_takeover + @with_completion_gate + async def paused_loop(*args, messages=None, client_runtime_context=None, **kwargs): + try: + yield _event({'type': 'tool_output', 'tool': 'bash', 'output': 'Waiting for an exact user approval.', + 'ask_user': {'kind': 'tool_approval', 'approval_id': 'missing'}}) + yield _event({'delta': 'not reached'}) + finally: + closed.append(current_journal() is not None) + + @with_teacher_takeover + @with_completion_gate + async def later_loop(*args, messages=None, client_runtime_context=None, **kwargs): + yield _event({'delta': 'Later run.'}) + yield _event({'type': 'metrics', 'data': {}}) + yield 'data: [DONE]\n\n' + + async def later_run(): + async for chunk in later_loop(messages=[{'role': 'user', 'content': 'x'}]): + if chunk.startswith('data: {') and '"metrics"' in chunk: + lineage.append(json.loads(chunk[6:])['data']['parent_run_id']) + + monkeypatch.setattr('src.agent_loop.stream_agent_loop', paused_loop) + monkeypatch.setattr('src.task_endpoint.resolve_task_candidates', lambda **kwargs: []) + + result = await TaskScheduler(session_manager=None)._run_agent_loop( + 'http://ep/v1', 'model', _task(), 's') + + assert 'paused safely' in result + # Closed during the pause, while its own journal was still bound. + assert closed == [True] + assert current_journal() is None + # Neither a chained task (which copies this context) nor a later run in + # this task inherits the paused run's journal as its parent. + chained = asyncio.create_task(later_run()) + await chained + await later_run() + assert lineage == [None, None]