sub_loop.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_chat/core/sub_loop.py
 
# OCT-LINT: disable=tests_exist reason="Covered by oc_chat/tests/unit/test_sub_loop_hygiene.py"
"""
Headless mini-ReAct engine for offloading cheap tasks out of the main loop.
 
Purpose
-------
:func:`run_sub_loop` runs a self-contained, headless ReAct loop with its
own minimal conversation context and a fixed model.  The sub-loop has full
access to the tools registry, logs every iteration to a session artifact,
and returns **only the final assistant text** — the main loop never sees
the sub-loop's internal tool calls or intermediate reasoning.
 
This lets the main loop preserve its prefix cache (the main model never
changes mid-turn) while still routing cheap tasks (reads, lookups) to a
cheaper model like ``deepseek-v4-flash``.
 
Responsibilities
----------------
- :func:`run_sub_loop` — construct a minimal conversation, run ≤N
  ReAct iterations, stream from the API, dispatch tool calls, archive
  the full sub-loop log, return the final answer.
- ``_DEFAULT_SUB_SYSTEM`` — task-focused framing combined with
  ``CANONICAL_POLICY`` + the compact tool map to compose the default
  sub-agent prompt per call (plan 135 §3.8).
 
Diagnostics
-----------
Domain: OC_CHAT.SUB_LOOP
Levels:
    L2 — lifecycle (spawn / complete / cap hit / error)
    L3 — per-iteration summary
    L4 — full tool-call dump
 
Contracts
---------
- The sub-loop never mutates the caller's session or message list.
- The task description must be self-contained — no dependency on
  conversation history, project tree, or AGENT_CONTEXT.
- Returns a :class:`SubLoopResult` (full text + merged usage + counts +
  archive path).  The caller injects ``.text`` into the main conversation
  as a (bounded) user message and folds ``.usage`` into session cost.
- If the iteration cap is reached, the last assistant message is
  returned with a ``[SUB_LOOP_CAP]`` prefix.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
- oc_chat.core.persistence (lazy; session artifact logging)
- oc_chat.core.event_bus (default_bus for synchronous M2 event publication)
- oc_chat.core.events (make_subagent_* TurnEvent factories and new_event_id)
- oc_chat.core.client_interface (lazy; ClientInterface - the abstract streaming API)
- oc_chat.core.cancellation (CancellationError - a cancelled sub-loop stream, and CancellationToken for type hints)
- oc_chat.core.policy (lazy; CANONICAL_POLICY - the base of the default sub-agent prompt)
- oc_chat.core.session (lazy; build_system_prompt - adds the tools section to the default sub-agent prompt)
- oc_chat.tools._sandbox (lazy; redact_secrets - redacts tool results before the sub-model context and the archive)
- oc_chat.tools.shell (lazy; deny_unconfirmed_tier2 - fails a gated action closed in the headless sub-loop)
 
See Also
--------
oc_chat/docs/conventions/CONTRACTS.md §5 Streaming chunk tuples — the
chunk_type / data contract that stream_chat yields.
"""
 
from __future__ import annotations
 
import json
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, TYPE_CHECKING
 
from oc_chat.core.cancellation import CancellationError
from oc_chat.core.event_bus import default_bus
from oc_chat.core.events import (
    make_subagent_completed,
    make_subagent_started,
    new_event_id,
)
 
if TYPE_CHECKING:
    from oc_chat.core.cancellation import CancellationToken
 
from oc_diagnostics import _dbg
 
# Lazy imports to avoid circular dependency at module level
 
# Plan 135 §3.8: the short task-focused framing appended to CANONICAL_POLICY when
# a caller does not supply an explicit ``system_prompt``. The full default prompt
# is composed PER CALL inside :func:`run_sub_loop` (CANONICAL_POLICY + this framing
# + the conditional tools section / compact map for the sub-loop's tool set) — a
# per-call compact map cannot live in a module constant. The env block is
# deliberately excluded (the sub-loop is headless / cache-isolated by design).
_DEFAULT_SUB_SYSTEM = (
    "You are a task-focused sub-agent. "
    "Complete the assigned task using available tools. "
    "Return ONLY the requested information — no commentary, no extra text."
)
 
MAX_SUB_ITERATIONS = 5
SUB_LOOP_CAP_MARKER = "[SUB_LOOP_CAP] "
 
 
def _build_default_sub_system(tool_schemas: list) -> str:
    """Compose the default sub-agent system prompt (plan 135 §3.8).
 
    ``CANONICAL_POLICY`` + the task-focused framing + the conditional tools
    section / compact map for *tool_schemas* (the no-tools policy when the set is
    empty). The env block is deliberately excluded — the sub-loop is headless and
    cache-isolated. Stays small (no JSON-schema dump; schemas travel via ``tools``).
    """
    from oc_chat.core.policy import CANONICAL_POLICY
    from oc_chat.core.session import build_system_prompt
 
    return build_system_prompt(f"{CANONICAL_POLICY}\n\n{_DEFAULT_SUB_SYSTEM}", tool_schemas)
 
 
@dataclass
class SubLoopResult:
    """Structured return of :func:`run_sub_loop` (plan 134 B2).
 
    The sub-loop used to return a bare ``str``; that hid its cost from the
    parent turn.  Every exit path now returns this — including the early
    cancellation / API-error / cap-hit paths — so the parent can fold the
    sub-loop's token spend into session cost + the live display, and build
    A4's truncation marker from :attr:`artifact_path`.
 
    Fields
    ------
    text:
        The final assistant reply (full, untruncated — the parent bounds it
        at the injection site).  Cap/cancel/error texts carry their markers.
    usage:
        Merged token counters summed across inner iterations (named fields
        only, types kept distinct), or ``None`` if no usage chunk arrived.
    iterations:
        Inner ReAct iterations entered.
    tool_calls:
        Total tool calls dispatched across the sub-loop.
    artifact_path:
        Absolute path of the archived full transcript, or ``None`` when
        archiving is off (``session_name is None``) — the marker then omits
        the path.
    unknown_usage_raw:
        Provider usage fields the merge did not recognise — carried raw
        (never folded into cost math).  The **parent** emits the throttled
        "new usage field" warning, so the sub-loop needs no session handle.
    """
 
    text: str
    usage: dict | None = None
    iterations: int = 0
    tool_calls: int = 0
    artifact_path: str | None = None
    unknown_usage_raw: list[dict] = field(default_factory=list)
 
 
# Token-usage fields summed across sub-loop iterations.  Types are kept
# distinct (prices differ per type); anything outside this set is carried raw
# (see :func:`_merge_usage`) so a new provider field is handled deliberately
# rather than silently folded into a cost number.
_KNOWN_USAGE_SCALARS = (
    "prompt_tokens",
    "completion_tokens",
    "total_tokens",
    "prompt_cache_hit_tokens",
    "prompt_cache_miss_tokens",
)
_KNOWN_USAGE_DETAILS = ("prompt_tokens_details", "completion_tokens_details")
_KNOWN_PROMPT_DETAILS = ("cached_tokens", "audio_tokens", "cache_write_tokens")
_KNOWN_COMPLETION_DETAILS = (
    "reasoning_tokens",
    "audio_tokens",
    "accepted_prediction_tokens",
    "rejected_prediction_tokens",
)
 
 
def _merge_usage(
    merged: dict[str, Any],
    new: Any,
    *,
    model: str,
    unknown_raw: list[dict[str, Any]],
) -> None:
    """Accumulate one provider usage dict into *merged* (in place).
 
    Named scalar token fields and the closed vocabulary of numeric detail fields
    are summed. ``None`` for a known scalar or any detail sub-field means absent
    and is skipped silently. A novel detail name, any non-null non-numeric value,
    or an unrecognised top-level field is appended to *unknown_raw* as
    ``{field, value, model}`` and is **never** folded into cost math.
    """
    if not isinstance(new, dict):
        return
    for key, val in new.items():
        if key in _KNOWN_USAGE_SCALARS:
            if val is None:
                continue
            if isinstance(val, (int, float)) and not isinstance(val, bool):
                merged[key] = merged.get(key, 0) + val
            else:
                unknown_raw.append({"field": key, "value": val, "model": model})
        elif key in _KNOWN_USAGE_DETAILS:
            if isinstance(val, dict):
                dst = merged.setdefault(key, {})
                known_details = (
                    _KNOWN_PROMPT_DETAILS
                    if key == "prompt_tokens_details"
                    else _KNOWN_COMPLETION_DETAILS
                )
                for sub, sval in val.items():
                    if sval is None:
                        continue
                    if (
                        sub in known_details
                        and isinstance(sval, (int, float))
                        and not isinstance(sval, bool)
                    ):
                        dst[sub] = dst.get(sub, 0) + sval
                    else:
                        unknown_raw.append(
                            {"field": f"{key}.{sub}", "value": sval, "model": model})
            else:
                unknown_raw.append({"field": key, "value": val, "model": model})
        else:
            unknown_raw.append({"field": key, "value": val, "model": model})
 
 
def _now_iso() -> str:
    """Millisecond-precision local ISO-8601 timestamp for event envelopes (M2.4)."""
    return datetime.now().astimezone().isoformat(timespec="milliseconds")
 
 
def run_sub_loop(
    client: "ClientInterface",  # noqa: F821
    task: str,
    model: str,
    tool_registry: dict,
    tool_schemas: list,
    *,
    system_prompt: str | None = None,
    max_iterations: int = MAX_SUB_ITERATIONS,
    session_name: str | None = None,
    max_tokens: int | None = None,
    presenter: object | None = None,
    cancellation: "CancellationToken | None" = None,
    parent_event_id: str | None = None,
    parent_session_id: str | None = None,
) -> SubLoopResult:
    """Execute a headless mini-ReAct loop and return a :class:`SubLoopResult`.
 
    Parameters
    ----------
    client:
        The same :class:`ClientInterface` the main loop uses.
    task:
        A self-contained description of what the sub-agent should do.
    model:
        Model ID to use (e.g. ``deepseek-v4-flash``).
    tool_registry:
        ``{name: callable}`` dict from the main loop.
    tool_schemas:
        OpenAI tool schemas for the API ``tools`` parameter.
    system_prompt:
        Explicit system prompt, used verbatim. When ``None`` (the default),
        run_sub_loop composes one per call: ``CANONICAL_POLICY`` + the
        task-focused framing + the conditional tools section / compact map for
        ``tool_schemas`` (plan 135 §3.8). No env block (headless by design).
    max_iterations:
        Maximum ReAct iterations (default 5).
    session_name:
        Passed through to artifact archiving.  ``None`` disables.
    max_tokens:
        Completion token limit (default: model default).
 
    Returns
    -------
    SubLoopResult
        The final answer (or a ``[SUB_LOOP_CAP]`` / ``[SUB_LOOP_ERROR]``
        marker) plus the merged token usage, iteration / tool-call counts,
        the archive path, and any unrecognised usage fields.  Returned on
        **every** exit path except ``KeyboardInterrupt``, which re-raises so
        the parent's handler records cost and the REPL backstop recovers.
    """
    from oc_chat.core.client_interface import ClientInterface
    from oc_chat.core.persistence import save_tool_artifact
 
    func = "run_sub_loop"
    _dbg("OC_CHAT.SUB_LOOP", func, f"spawn model={model} task={task[:80]}…", 2)
 
    # Plan 135 §3.8: compose the default sub-agent prompt per call so it carries
    # the compact map for THIS sub-loop's tool set. An explicit system_prompt is
    # used verbatim (not overridden).
    if system_prompt is None:
        system_prompt = _build_default_sub_system(tool_schemas)
 
    # M2.4: typed sub-agent events. Gated on both parent ids so legacy
    # callers and unit tests that omit them emit nothing (plan 102 §4.1).
    # subagent_session_id is fresh + stable across the started/completed pair;
    # parent_event_id is the parent turn's turn_id (ADR-1), reused as turn_id.
    _events_enabled = parent_event_id is not None and parent_session_id is not None
    _subagent_sid = new_event_id()
 
    def _emit(make: Any, **kw: Any) -> None:
        if not _events_enabled:
            return
        default_bus().publish(make(
            event_id=new_event_id(), parent_event_id=parent_event_id,
            session_id=parent_session_id, ts=_now_iso(), turn_id=parent_event_id,
            subagent_session_id=_subagent_sid, **kw))
 
    _emit(make_subagent_started)
 
    # ── Build minimal conversation ──────────────────────────────────
    messages: list[dict[str, Any]] = [
        {"role": "system", "content": system_prompt.strip()},
        {"role": "user", "content": task},
    ]
    log_archive: list[dict[str, Any]] = []  # full transcript for artifact
 
    # ── Loop-scoped accounting (plan 134 B2 / Opus M2) ──────────────
    # These live OUTSIDE the iteration loop so every early-exit path can
    # return the partial usage that arrived before the exit; per-iteration
    # locals reset and would be lost.
    merged_usage: dict[str, Any] = {}
    unknown_usage_raw: list[dict[str, Any]] = []
    total_tool_calls = 0
    iteration = 0
    assistant_msg: dict[str, Any] = {}
 
    # ── Mini ReAct loop ─────────────────────────────────────────────
    for iteration in range(max_iterations):
        _dbg("OC_CHAT.SUB_LOOP", func,
             f"iter {iteration}/{max_iterations} model={model}", 3)
 
        full_content = ""
        full_reasoning = ""
        finish_reason: str | None = None  # None == provider reported none
        tool_calls_flat: list[dict] = []  # provider {id,name,arguments} for dispatch
 
        # Phase H.4a-iii: bracket the inner stream with a spinner. Stops on
        # first yielded chunk to avoid overlap with streamed text; finally
        # block is a belt-and-braces safety net (spinner_stop is idempotent).
        _spinner = presenter.spinner_start("thinking (sub-loop)…") if presenter is not None else None
 
        def _partial(text: str) -> SubLoopResult:
            """Build a SubLoopResult for an early-exit path, preserving usage."""
            return SubLoopResult(
                text=text,
                usage=merged_usage or None,
                iterations=iteration + 1,
                tool_calls=total_tool_calls,
                artifact_path=None,
                unknown_usage_raw=unknown_usage_raw,
            )
 
        try:
            for chunk_type, data in client.stream_chat(
                messages,
                model,
                tools=tool_schemas,
                max_tokens=max_tokens,
                cancellation=cancellation,
            ):
                if _spinner is not None and presenter is not None:
                    presenter.spinner_stop(_spinner)
                if chunk_type == "reasoning":
                    full_reasoning += data
                elif chunk_type == "content":
                    full_content += data
                elif chunk_type == "tool_call":
                    # Providers stream tool calls ONE at a time (singular —
                    # plan 134 A1). Accumulate them like the main loop does;
                    # the old plural "tool_calls" branch never matched in
                    # production, so sub-loop tools were silently dropped.
                    tool_calls_flat.append(data)
                elif chunk_type == "usage":
                    # Merge across iterations so the parent can fold the
                    # sub-loop's token spend into session cost (B2).
                    _merge_usage(merged_usage, data,
                                 model=model, unknown_raw=unknown_usage_raw)
                elif chunk_type == "finish_reason":
                    finish_reason = data
        except CancellationError:
            # Cooperative cancel from a sister presenter or signal handler.
            _dbg("OC_CHAT.SUB_LOOP", func,
                 "sub-loop stream cancelled via CancellationToken", 2, override=True)
            _emit(make_subagent_completed, outcome="cancelled")  # M2.4 ADR-4
            return _partial(SUB_LOOP_CAP_MARKER + "Sub-loop cancelled.")
        except KeyboardInterrupt:
            # Propagate cleanly to the REPL backstop (main.py). A single path
            # cannot both return a SubLoopResult and re-raise; the parent's
            # interrupt handler records partial cost and recovers (plan 134
            # H3). Folding interrupt usage here is a deferred refactor.
            _dbg("OC_CHAT.SUB_LOOP", func,
                 "sub-loop stream cancelled by user (Ctrl+C)", 2, override=True)
            _emit(make_subagent_completed, outcome="cancelled")  # M2.4 ADR-4
            raise
        except Exception as e:
            _dbg("OC_CHAT.SUB_LOOP", func,
                 f"API error: {type(e).__name__}: {e}", 1, override=True)
            _emit(make_subagent_completed, outcome="error")  # M2.4 ADR-4
            return _partial(f"[SUB_LOOP_ERROR: {type(e).__name__}: {e}]")
        finally:
            # Phase H.4a-iii: safety stop in case no chunks ever arrived
            # (or first-chunk stop somehow missed). Idempotent.
            if _spinner is not None and presenter is not None:
                presenter.spinner_stop(_spinner)
 
        # A truncated sub-loop hands the parent a partial answer that reads
        # like a finished one — headless, so nobody sees it happen. Log it.
        if finish_reason == "length":
            _dbg("OC_CHAT.SUB_LOOP", func,
                 f"TRUNCATED at max_tokens: iter={iteration} model={model} "
                 f"max_tokens={max_tokens} content={len(full_content)}c "
                 f"reasoning={len(full_reasoning)}c "
                 f"tool_calls={len(tool_calls_flat)}", 1, override=True)
 
        # Record assistant message. tool_calls are stored in the NESTED API
        # schema ({id, type, function:{name, arguments}}) — the flat provider
        # shape would be rejected when re-sent on the next iteration (plan 134
        # A1 (ii)). The flat records (tool_calls_flat) drive dispatch below.
        api_tool_calls = [
            {
                "id": tc.get("id", tc.get("name", "?")),
                "type": "function",
                "function": {
                    "name": tc.get("name"),
                    "arguments": tc.get("arguments", ""),
                },
            }
            for tc in tool_calls_flat
        ]
        assistant_msg = {"role": "assistant", "content": full_content or None}
        if full_reasoning:
            assistant_msg["reasoning_content"] = full_reasoning
        if api_tool_calls:
            assistant_msg["tool_calls"] = api_tool_calls
        messages.append(assistant_msg)
        log_archive.append(dict(assistant_msg))
 
        # ── No tool calls → this is the final answer ────────────────
        if not tool_calls_flat:
            _dbg("OC_CHAT.SUB_LOOP", func, "final answer received", 2)
            break
 
        # ── Execute tool calls (sequential dispatch) ────────────────
        total_tool_calls += len(tool_calls_flat)
        for tc in tool_calls_flat:
            tool_fn = tool_registry.get(tc.get("name", ""))
            if not tool_fn:
                result = {"error": f"Unknown tool: {tc.get('name')}"}
            else:
                try:
                    args = json.loads(tc.get("arguments", "{}"))
                    result = tool_fn(**args)
                except Exception as e:
                    _dbg("OC_CHAT.SUB_LOOP", func,
                         f"tool {tc.get('name')} raised {type(e).__name__}: {e}",
                         1, override=True)
                    result = {"error": str(e)}
 
            # M8.0-C / F67: a sub-loop is always headless -- no human to confirm a
            # gated action (a Tier-2 command or an external-content fetch). Fail
            # closed: convert any `requires_confirmation` to a terminal error so the
            # sub-agent stops asking (it can never self-authorize). The interactive
            # confirm channels live only in the main react-loop.
            if isinstance(result, dict) and result.get("requires_confirmation"):
                from oc_chat.tools import shell
                result = shell.deny_unconfirmed_tier2(
                    result, note="A sub-agent runs headless and cannot confirm a gated action.")
                _dbg("OC_CHAT.SUB_LOOP", func,
                     "gated action denied headless (sub-loop): requires_confirmation", 2, override=True)
 
            # Convert result dict to string for the API. M8.0 WU-9 review fix
            # (F46/F51): redact secret patterns before the result enters the
            # sub-model context AND the on-disk archive -- recall returns memory
            # content unredacted (the filesystem tools redact at the tool level).
            from oc_chat.tools._sandbox import redact_secrets
            result_str = redact_secrets(json.dumps(result, ensure_ascii=False, default=str))
            result_msg = {
                "role": "tool",
                "tool_call_id": tc.get("id", tc.get("name", "?")),
                "content": result_str,
            }
            messages.append(result_msg)
            log_archive.append(dict(result_msg))
 
    else:
        # Loop completed without break → hit iteration cap
        _dbg("OC_CHAT.SUB_LOOP", func,
             f"cap hit after {max_iterations} iterations", 2)
        _emit(make_subagent_completed, outcome="cap_hit")  # M2.4: both cap returns
        if messages and messages[-1].get("role") == "assistant":
            last = messages[-1].get("content") or ""
            cap_text = SUB_LOOP_CAP_MARKER + last
        else:
            cap_text = SUB_LOOP_CAP_MARKER + "no assistant reply"
        return SubLoopResult(
            text=cap_text,
            usage=merged_usage or None,
            iterations=max_iterations,
            tool_calls=total_tool_calls,
            artifact_path=None,
            unknown_usage_raw=unknown_usage_raw,
        )
 
    # ── Extract final answer ────────────────────────────────────────
    # The full text is returned verbatim. The chain-of-thought line "filter"
    # (plan 134 A2) and the silent 10k truncation (A4) are gone — the former
    # mangled legitimate answers ("First,", "Step 1:" …); the latter is
    # replaced by the parent bounding the injected message at the model's
    # tool-result cap (react_loop / resolve_tool_result_cap). Sub-agent
    # terseness is the system prompt's job, not a regex.
    final_text = (assistant_msg.get("content") or "").strip()
 
    # M2.4: happy-path completion emitted before the archive side-effect (ADR-3)
    # so it fires whether or not session_name is set, and even if archiving fails.
    _emit(make_subagent_completed, outcome="ok")
 
    # ── Archive full transcript ─────────────────────────────────────
    artifact_path: str | None = None
    if session_name is not None:
        archive_entry = {
            "sub_loop": {
                "model": model,
                "task": task,
                "iterations": iteration + 1,
                "max_iterations": max_iterations,
                "final_length": len(final_text),
            },
            "messages": log_archive,
        }
        archive_json = json.dumps(archive_entry, ensure_ascii=False, default=str)
        artifact_path = str(save_tool_artifact(
            session_name, f"sub_loop_{model}", archive_json))
        _dbg("OC_CHAT.SUB_LOOP", func,
                 f"done model={model} iters={iteration+1} chars={len(final_text)}", 2)
    return SubLoopResult(
        text=final_text,
        usage=merged_usage or None,
        iterations=iteration + 1,
        tool_calls=total_tool_calls,
        artifact_path=artifact_path,
        unknown_usage_raw=unknown_usage_raw,
    )