tracer-cloud--opensre/tests/core/agent/test_turn_loop.py

"""Shell-local turn loop bookkeeping tests."""

from __future__ import annotations

import io
from typing import Any

from rich.console import Console

from core.agent_harness.accounting.turn_accounting import DefaultTurnAccounting
from core.agent_harness.session.persistence.memory import InMemorySessionStorage
from core.agent_harness.turns.orchestrator import run_turn
from surfaces.interactive_shell.runtime.core.turn_accounting import (
    ToolCallingTurnResult,
)
from surfaces.interactive_shell.runtime.shell_turn_execution import execute_shell_turn
from surfaces.interactive_shell.session import Session
from surfaces.interactive_shell.utils.telemetry.recorder import LlmRunInfo


class _Recorder:
    def __init__(self) -> None:
        self.responses: list[tuple[str, LlmRunInfo | None]] = []
        self.flush_count = 0

    def set_response(self, response: str, run_info: LlmRunInfo | None = None) -> None:
        self.responses.append((response, run_info))

    def flush(self) -> None:
        self.flush_count += 1


def _console() -> Console:
    return Console(file=io.StringIO(), force_terminal=False, color_system=None, width=80)


def _unhandled_turn(*_args: object, **_kwargs: object) -> ToolCallingTurnResult:
    return ToolCallingTurnResult(
        planned_count=0,
        executed_count=0,
        executed_success_count=0,
        has_unhandled_clause=False,
        handled=False,
    )


def test_recorder_flushes_once_for_chat_fallback() -> None:
    recorder = _Recorder()
    run_info = LlmRunInfo(response_text="answered")

    def _answer(*_args: Any, **_kwargs: Any) -> LlmRunInfo:
        return run_info

    result = execute_shell_turn(
        "question",
        Session(),
        _console(),
        recorder=recorder,  # type: ignore[arg-type]
        execute_actions=_unhandled_turn,
        gather_evidence=lambda *_a, **_k: None,
        answer_agent=_answer,
    )

    assert result.answered is True
    assert result.assistant_response_text == "answered"
    assert recorder.responses == [("answered", run_info)]
    assert recorder.flush_count == 1


def test_recorder_flushes_once_for_silent_handled_turn() -> None:
    recorder = _Recorder()
    session = Session()

    def _handled(*_args: object, **_kwargs: object) -> ToolCallingTurnResult:
        return ToolCallingTurnResult(
            planned_count=1,
            executed_count=1,
            executed_success_count=1,
            has_unhandled_clause=False,
            handled=True,
            response_text="command output",
        )

    result = execute_shell_turn(
        "run something",
        session,
        _console(),
        recorder=recorder,  # type: ignore[arg-type]
        execute_actions=_handled,
        gather_evidence=lambda *_a, **_k: None,
        answer_agent=lambda *_a, **_k: None,
    )

    assert result.answered is False
    assert result.final_intent == "cli_agent_handled"
    assert recorder.responses == [("command output", None)]
    assert recorder.flush_count == 1
    assert session.cli_agent_messages[-2:] == [
        ("user", "run something"),
        ("assistant", "command output"),
    ]


def test_default_turn_accounting_persists_action_only_context() -> None:
    storage = InMemorySessionStorage()
    session = Session(storage=storage)
    storage.open_session(session)

    def _handled(*_args: object, **_kwargs: object) -> ToolCallingTurnResult:
        return ToolCallingTurnResult(
            planned_count=1,
            executed_count=1,
            executed_success_count=1,
            has_unhandled_clause=False,
            handled=True,
            response_text="Hawaii: +28C",
        )

    result = run_turn(
        "weather in Hawaii",
        session,
        execute_actions=_handled,
        gather=lambda *_args, **_kwargs: None,
        answer=lambda *_args, **_kwargs: None,
        accounting=DefaultTurnAccounting(session, "weather in Hawaii"),
    )

    records = storage.read(session.session_id)
    messages = [record for record in records if record.get("type") == "message"]

    assert result.final_intent == "cli_agent_handled"
    assert session.cli_agent_messages[-2:] == [
        ("user", "weather in Hawaii"),
        ("assistant", "Hawaii: +28C"),
    ]
    assert [
        (message.get("role"), message.get("content"), message.get("metadata"))
        for message in messages[-2:]
    ] == [
        ("user", "weather in Hawaii", {"kind": "chat"}),
        ("assistant", "Hawaii: +28C", {"kind": "chat"}),
    ]