fix(runtime): isolate nested invocation ownership

src/agent_loop.py

@@ -84,6 +84,7 @@ from src.tool_types import ToolBlock
 from src.turn_contract import selected_tools_for_request, with_turn_contract
 from src.agent_runtime.journal import propose_action, execute_action
 from src.agent_runtime.completion import with_completion_gate
+from src.teacher_escalation import with_teacher_takeover, request_teacher_takeover
 from src.tool_utils import _truncate, get_mcp_manager
 from src.agent_tools import (
     parse_tool_blocks,
@@ -20326,6 +20327,7 @@ def _blocks_before_inference(turn_contract) -> bool:
 
 
 @with_turn_contract
+@with_teacher_takeover
 @with_completion_gate
 async def stream_agent_loop(
     endpoint_url: str,
@@ -20368,6 +20370,7 @@ async def stream_agent_loop(
     thinking_mode: Optional[str] = None,
     suppress_skills: bool = False,
     reasoning_effort: Optional[str] = None,
+    _parent_run_id: Optional[str] = None,
 ) -> AsyncGenerator[str, None]:
     """Streaming agent loop generator.
 
@@ -37239,29 +37242,25 @@ async def stream_agent_loop(
     )
     yield f"data: {json.dumps({'type': 'metrics', 'data': metrics})}\n\n"
 
-    # Teacher-escalation: inline takeover visible in the chat stream.
-    # The student just finished; if Tier 1 flags failure, the teacher
-    # gets a turn (with its own tool calls forwarded to the user) and
-    # a skill is saved ONLY if the teacher actually succeeds. Skipped
-    # when we ARE the teacher to avoid recursion.
+    # Queue the existing teacher hook. The outer adapter executes it only
+    # after this invocation's completion gate and action context have closed.
     if not _is_teacher_run and not guide_only and not _awaiting_user:
-        try:
-            from src.teacher_escalation import run_teacher_inline
-            async for evt in run_teacher_inline(
-                student_endpoint_url=endpoint_url,
-                student_messages=messages,
-                student_tool_events=tool_events,
-                student_reply=full_response,
-                owner=owner,
-                session_id=session_id,
-                workspace=workspace,
-                disabled_tools=disabled_tools,
-                tool_policy=tool_policy,
-                active_document=active_document,
-                active_email=active_email,
-            ):
-                yield evt
-        except Exception as _esc_err:
-            logger.warning(f"teacher escalation hook failed: {_esc_err}", exc_info=True)
+        request_teacher_takeover(
+            student_endpoint_url=endpoint_url,
+            student_messages=messages,
+            student_tool_events=tool_events,
+            student_reply=full_response,
+            owner=owner,
+            session_id=session_id,
+            workspace=workspace,
+            disabled_tools=disabled_tools,
+            tool_policy=tool_policy,
+            active_document=active_document,
+            active_email=active_email,
+            turn_contract=turn_contract,
+            external_untrusted_context_seen=run_security.external_untrusted_context_seen,
+            client_runtime_context=client_runtime_context,
+            plan_mode=plan_mode,
+        )
 
     yield "data: [DONE]\n\n"

src/agent_runtime/completion.py

@@ -181,8 +181,9 @@ def with_completion_gate(func):
         trusted_workspace = vet_workspace(bound.get('workspace')) if bound.get('workspace') else ''
         requirements = replace(requirements, workspace_root=trusted_workspace or '')
         parent = current_journal()
-        journal = parent if parent is not None and parent.workspace == requirements.workspace_root else ActionJournal(
-            workspace=requirements.workspace_root, observed_artifacts=requirements.required_artifacts)
+        journal = ActionJournal(
+            workspace=requirements.workspace_root, observed_artifacts=requirements.required_artifacts,
+            parent_run_id=bound.get('_parent_run_id') or (parent.run_id if parent is not None else None))
         answer_events: list[dict] = []
         metrics_events: list[dict] = []
         answer = ''
@@ -308,7 +309,8 @@ def with_completion_gate(func):
             for event in metrics_events:
                 metadata = event.setdefault('data', {})
                 metadata.update(completion_decision=decision.to_dict(), evidence_events=ledger.to_list(),
-                                action_receipts=journal.to_list(), completion_requirements=requirements.to_dict())
+                                action_receipts=journal.to_list(), completion_requirements=requirements.to_dict(),
+                                run_id=journal.run_id, parent_run_id=journal.parent_run_id)
                 metadata['completion_gate'] = {
                     'buffer_seconds': released_at - first_answer_at if first_answer_at is not None else 0,
                     'first_visible_answer_seconds': released_at - started,

src/agent_runtime/journal.py

@@ -66,6 +66,7 @@ class ActionJournal:
     actions: list[ActionReceipt] = field(default_factory=list)
     workspace: str = ''
     observed_artifacts: tuple[str, ...] = ()
+    parent_run_id: str | None = None
 
     def capture_versions(self, action: ActionReceipt) -> None:
         if self.workspace:
@@ -111,9 +112,11 @@ _ACTION: ContextVar[ActionReceipt | None] = ContextVar('runtime_current_action',
 @contextmanager
 def bind_journal(journal: ActionJournal):
     token = _JOURNAL.set(journal)
+    action_token = _ACTION.set(None)
     try:
         yield journal
     finally:
+        _ACTION.reset(action_token)
         _JOURNAL.reset(token)
 
 

src/teacher_escalation.py

@@ -24,6 +24,10 @@ itself wasn't confident about.
 from __future__ import annotations
 
 import asyncio
+from contextlib import aclosing
+from contextvars import ContextVar
+from copy import deepcopy
+from functools import wraps
 import logging
 import re
 from typing import Any, Dict, List, Optional, Tuple
@@ -32,6 +36,58 @@ from urllib.parse import urlparse
 logger = logging.getLogger(__name__)
 
 
+_TAKEOVER: ContextVar[dict | None] = ContextVar('teacher_takeover_request', default=None)
+
+
+def request_teacher_takeover(**parameters):
+    """Record a finished student's handoff; execution waits for its gate to close."""
+    from src.agent_runtime.journal import current_journal
+    request = _TAKEOVER.get()
+    if request is not None:
+        journal = current_journal()
+        request.update(parameters, parent_run_id=journal.run_id if journal is not None else None)
+
+
+def with_teacher_takeover(func):
+    """Orchestrate gated invocations and own the single outer stream terminator."""
+    @wraps(func)
+    async def wrapped(*args, **kwargs):
+        request = {}
+        token = _TAKEOVER.set(request)
+        done = False
+        failed = False
+        try:
+            async with aclosing(func(*args, **kwargs)) as stream:
+                async for chunk in stream:
+                    if chunk.strip() == 'data: [DONE]':
+                        done = True
+                        continue
+                    failed |= chunk.startswith('event: error')
+                    yield chunk
+            # The parent generator and gate have both unwound. Child control
+            # events now belong only to the child, never to the parent gate.
+            if request and not failed:
+                try:
+                    async with aclosing(run_teacher_inline(**request)) as stream:
+                        async for chunk in stream:
+                            if chunk.strip() == 'data: [DONE]':
+                                continue
+                            failed |= chunk.startswith('event: error')
+                            yield chunk
+                except Exception as exc:
+                    logger.warning('teacher escalation hook failed: %s', exc, exc_info=True)
+                    if not failed:
+                        import json
+                        yield 'data: ' + json.dumps({'type': 'escalation_failed', 'reason': str(exc),
+                                                    'teacher': True}) + '\n\n'
+            if done and not failed:
+                yield 'data: [DONE]\n\n'
+        finally:
+            _TAKEOVER.reset(token)
+
+    return wrapped
+
+
 # Hosts considered SOTA / paid APIs — if the student's endpoint URL
 # hits one of these, the loop is OFF (the user is already paying for
 # a top-tier model; no need to escalate).
@@ -524,6 +580,11 @@ async def run_teacher_inline(
     tool_policy: Any = None,
     active_document: Any = None,
     active_email: Optional[Dict[str, str]] = None,
+    turn_contract=None,
+    parent_run_id: Optional[str] = None,
+    external_untrusted_context_seen: bool = False,
+    client_runtime_context: Optional[Dict[str, Any]] = None,
+    plan_mode: bool = False,
 ):
     """Async generator. Yields SSE event strings.
 
@@ -606,7 +667,7 @@ async def run_teacher_inline(
     # user/assistant/tool history so the teacher sees what the student
     # tried. The appended note leads with the user request text so RAG
     # tool selection picks the right tools for the teacher's turn.
-    history = [m for m in student_messages if m.get("role") != "system"]
+    history = deepcopy([m for m in student_messages if m.get("role") != "system"])
     note_content = (
         f"{user_request or '(no user request captured)'}\n\n"
         "[teacher-takeover] The previous attempt by the student model "
@@ -623,8 +684,9 @@ async def run_teacher_inline(
     captured_tool_events: List[Dict[str, Any]] = []
     captured_text_parts: List[str] = []
     captured_metrics: Dict[str, Any] = {}
+    captured_decision: Dict[str, Any] = {}
 
-    async for evt_str in stream_agent_loop(
+    async with aclosing(stream_agent_loop(
         endpoint_url=teacher_url,
         model=teacher_model,
         messages=teacher_messages,
@@ -632,51 +694,63 @@ async def run_teacher_inline(
         owner=owner,
         session_id=session_id,
         workspace=workspace,
-        disabled_tools=disabled_tools,
+        disabled_tools=set(disabled_tools) if disabled_tools is not None else None,
         tool_policy=tool_policy,
         active_document=active_document,
         active_email=active_email,
+        turn_contract=turn_contract,
+        _parent_run_id=parent_run_id,
+        external_untrusted_context_seen=external_untrusted_context_seen,
+        client_runtime_context=deepcopy(client_runtime_context),
+        plan_mode=plan_mode,
         _is_teacher_run=True,
-    ):
-        # Swallow teacher's own [DONE] — outer loop emits the real one
-        if "[DONE]" in evt_str:
-            continue
-        if evt_str.startswith("data: "):
-            try:
-                payload = json.loads(evt_str[6:].strip())
-            except Exception:
+    )) as stream:
+        async for evt_str in stream:
+            # Swallow teacher's own [DONE] — outer loop emits the real one
+            if evt_str.strip() == 'data: [DONE]':
+                continue
+            if evt_str.startswith('event: error'):
                 yield evt_str
-                continue
-            if isinstance(payload, dict):
-                payload["teacher"] = True
-                typ = payload.get("type")
-                if typ == "metrics" and isinstance(payload.get("data"), dict):
-                    # The outer chat route persists only the last metrics
-                    # payload. Keep a copy so any approval produced after the
-                    # recursive teacher run's metrics remains reloadable.
-                    captured_metrics = dict(payload["data"])
-                if typ == "tool_output":
-                    captured_tool_event = {
-                        "tool": payload.get("tool"),
-                        "command": payload.get("command"),
-                        "output": payload.get("output"),
-                        "exit_code": payload.get("exit_code"),
-                    }
-                    if isinstance(payload.get("ask_user"), dict):
-                        captured_tool_event["ask_user"] = payload["ask_user"]
-                    captured_tool_events.append(captured_tool_event)
-                if "delta" in payload and isinstance(payload["delta"], str):
-                    if payload.get("thinking"):
-                        continue
-                    captured_text_parts.append(payload["delta"])
-                yield 'data: ' + json.dumps(payload) + '\n\n'
-                continue
-        yield evt_str
+                return
+            if evt_str.startswith("data: "):
+                try:
+                    payload = json.loads(evt_str[6:].strip())
+                except Exception:
+                    yield evt_str
+                    continue
+                if isinstance(payload, dict):
+                    payload["teacher"] = True
+                    typ = payload.get("type")
+                    if typ == 'completion_decision' and isinstance(payload.get('data'), dict):
+                        captured_decision = payload['data']
+                    if typ == "metrics" and isinstance(payload.get("data"), dict):
+                        # The outer chat route persists only the last metrics
+                        # payload. Keep a copy so any approval produced after the
+                        # recursive teacher run's metrics remains reloadable.
+                        captured_metrics = dict(payload["data"])
+                    if typ == "tool_output":
+                        captured_tool_event = {
+                            "tool": payload.get("tool"),
+                            "command": payload.get("command"),
+                            "output": payload.get("output"),
+                            "exit_code": payload.get("exit_code"),
+                        }
+                        if isinstance(payload.get("ask_user"), dict):
+                            captured_tool_event["ask_user"] = payload["ask_user"]
+                        captured_tool_events.append(captured_tool_event)
+                    if "delta" in payload and isinstance(payload["delta"], str):
+                        if payload.get("thinking"):
+                            continue
+                        captured_text_parts.append(payload["delta"])
+                    yield 'data: ' + json.dumps(payload) + '\n\n'
+                    continue
+            yield evt_str
 
     # A takeover that paused for a question or exact action has not completed
     # yet. Its server-owned approval card is already in the live/persisted tool
     # events; do not evaluate the partial trace or distill it into a skill.
-    if any(event.get("ask_user") for event in captured_tool_events):
+    if (any(event.get("ask_user") for event in captured_tool_events)
+            or (captured_decision and not captured_decision.get('can_complete', False))):
         return
 
     teacher_text = "".join(captured_text_parts).strip()