cowork-local/tests/fakes/fake_provider.py

"""FakeProvider - a scripted, offline stand-in for a real LLM provider (R01-T02).

The real providers (``providers/openai_compat.py``, ``providers/anthropic.py``)
open HTTP connections, need API keys and stream at the mercy of the network, so
nothing above them could be tested deterministically. This double implements the
same :class:`providers.base.Provider` contract from a list of scripted turns:

    provider = FakeProvider([
        ScriptedTurn(tool_calls=[("save_file", {"filename": "a.md", "content": "hi"})]),
        ScriptedTurn(text="Saved it."),
    ])

Turn 1 asks the agent loop to call a tool, turn 2 ends the loop with plain text -
exactly the two-step shape ``run_cowork`` exercises, with zero I/O.

It records every call it received (:attr:`FakeProvider.calls`) so a test can
assert on what the layer above actually sent (message list, tool catalogue),
which is how the characterization and contract suites pin current behaviour.
"""
from __future__ import annotations

import itertools
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

from cowork_local.providers.base import (
    CancelFn,
    Provider,
    ProviderError,
    TextCallback,
    ToolSpec,
)

# One scripted tool call: (name, arguments). Ids are generated by the provider so
# a test never has to invent them, mirroring what a real gateway does.
ToolCallScript = Tuple[str, Dict[str, Any]]


@dataclass(frozen=True)
class ScriptedTurn:
    """What :class:`FakeProvider` should do for ONE ``chat()`` call.

    ``text`` is streamed through ``on_text`` and returned as the assistant
    message content. ``reasoning`` goes to ``on_reasoning`` only - it must never
    leak into the answer, and asserting that is one of this double's jobs.

    ``tool_calls`` makes the agent loop run tools and come back for another turn;
    an empty tuple ends the loop.

    ``error``, when set, raises :class:`ProviderError` instead of answering, so
    error/recovery paths are testable without simulating a network fault.

    ``chunk_size`` > 0 splits ``text`` into fixed-size pieces to exercise
    chunk-boundary handling in stream consumers (the ``<think>`` splitter and the
    UI's incremental markdown renderer both have boundary logic worth covering).
    """

    text: str = ""
    reasoning: str = ""
    tool_calls: Sequence[ToolCallScript] = ()
    error: Optional[str] = None
    chunk_size: int = 0


@dataclass
class RecordedCall:
    """A snapshot of one ``chat()`` invocation, for assertions after the fact."""

    messages: List[Dict[str, Any]]
    tool_names: List[str]
    cancelled: bool = False


class FakeProvider(Provider):
    """A ``Provider`` that replays :class:`ScriptedTurn` objects.

    Args:
        turns: the scripted turns, consumed in order.
        model: the model id reported through ``describe()`` / usage records.
        models: what :meth:`list_models` returns (Settings' "Load models").
        strict: when True (default) running past the end of the script raises
            ``AssertionError``. That is intentional noise: a silent extra turn
            usually means the code under test looped more than the test author
            expected, and hiding it behind an empty answer would turn a real
            behaviour change into a passing test.
    """

    name = "fake"
    # The double can accept image content blocks, so vision code paths are
    # reachable in tests without a real vision-capable gateway.
    supports_vision = True

    def __init__(
        self,
        turns: Optional[Sequence[ScriptedTurn]] = None,
        *,
        model: str = "fake-model",
        models: Optional[Sequence[str]] = None,
        strict: bool = True,
        conf: Optional[Dict[str, Any]] = None,
    ) -> None:
        super().__init__(dict(conf or {}, model=model))
        self._turns: List[ScriptedTurn] = list(turns or [])
        self._models = list(models or [model])
        self._strict = strict
        self._ids = itertools.count(1)  # deterministic tool-call ids: call_1, call_2, ...
        self.calls: List[RecordedCall] = []

    # -- introspection helpers used by tests ---------------------------- #
    @property
    def call_count(self) -> int:
        """How many times the layer above asked this provider to run a turn."""
        return len(self.calls)

    @property
    def remaining_turns(self) -> int:
        """Scripted turns not consumed yet - assert 0 to prove the script was
        fully used (an unused turn means the code stopped earlier than intended)."""
        return len(self._turns)

    def last_messages(self) -> List[Dict[str, Any]]:
        """The message list sent on the most recent call (empty if never called)."""
        return self.calls[-1].messages if self.calls else []

    # -- Provider contract ---------------------------------------------- #
    def chat(
        self,
        messages: List[Dict[str, Any]],
        tools: Optional[List[ToolSpec]] = None,
        on_text: Optional[TextCallback] = None,
        cancel: Optional[CancelFn] = None,
        on_reasoning: Optional[TextCallback] = None,
    ) -> Dict[str, Any]:
        """Replay the next scripted turn, honouring cancel and both callbacks.

        The message list is deep-ish copied into the recording because the agent
        loop keeps appending to the SAME list object; without the copy every
        recorded call would show the final state and assertions on "what was
        sent at step 1" would be meaningless.
        """
        record = RecordedCall(
            messages=[dict(m) for m in messages],
            tool_names=[t.name for t in (tools or [])],
        )
        self.calls.append(record)

        turn = self._next_turn()

        # Checked before streaming anything: a provider that already knows the
        # caller gave up must not spend callbacks on text nobody will render.
        if self._is_cancelled(cancel):
            record.cancelled = True
            return {"role": "assistant", "content": "", "tool_calls": []}

        if turn.error:
            raise ProviderError(turn.error)

        if turn.reasoning and on_reasoning:
            on_reasoning(turn.reasoning)

        for piece in self._stream_pieces(turn):
            # Re-checked between chunks so a mid-stream Stop truncates the answer
            # the same way a real streamed response does.
            if self._is_cancelled(cancel):
                record.cancelled = True
                break
            if on_text:
                on_text(piece)

        return {
            "role": "assistant",
            "content": turn.text,
            "tool_calls": [
                {"id": f"call_{next(self._ids)}", "name": name, "arguments": dict(args)}
                for name, args in turn.tool_calls
            ],
        }

    def list_models(self) -> List[str]:
        """Configured model ids. Clears ``last_error`` so ``test_connection()``
        reports success, matching how a healthy real provider behaves."""
        self.last_error = ""
        return list(self._models)

    # -- internals ------------------------------------------------------- #
    def _next_turn(self) -> ScriptedTurn:
        """Pop the next scripted turn, or fail loudly when the script ran out."""
        if self._turns:
            return self._turns.pop(0)
        if self._strict:
            raise AssertionError(
                f"FakeProvider script exhausted: chat() was called {len(self.calls)} "
                "time(s) but fewer turns were scripted. Add a ScriptedTurn, or pass "
                "strict=False if the extra call is genuinely expected."
            )
        return ScriptedTurn()

    @staticmethod
    def _stream_pieces(turn: ScriptedTurn) -> List[str]:
        """Split a turn's answer into the fragments to stream.

        ``chunk_size == 0`` streams the whole answer in one piece (the common
        case); a positive size slices it so tests can drive chunk-boundary logic.
        """
        if not turn.text:
            return []
        if turn.chunk_size <= 0:
            return [turn.text]
        size = turn.chunk_size
        return [turn.text[i:i + size] for i in range(0, len(turn.text), size)]


__all__ = ["FakeProvider", "ScriptedTurn", "RecordedCall"]