portfolio-optimiser/spikes/_harness.py

"""Shared test seams for the Fase 1 de-risk spikes (throwaway, dev-only).

Three reusable pieces every spike leans on:

1. **Cost/stop invariant (B4 / D6).** ``Budget`` refuses to start without positive
   token and round caps (``ValueError`` on a bad construction argument), and
   ``TokenMeter`` raises ``BudgetExceeded`` the moment a cap is crossed at runtime.
   Two exception types is intentional: ``ValueError`` = bad ctor argument (you never
   even started), ``BudgetExceeded`` = a cap was breached while running.

2. **Deterministic fake model.** ``FakeChatClient`` subclasses the GA
   ``BaseChatClient`` with scripted, deterministic replies and counts "tokens" by
   word-count — no network, no endpoint. ``fake_agent`` wraps it in a real
   ``agent_framework.Agent`` so the orchestration builders get genuine participants.
   The Step 2 builder smoke (in ``tests/spikes/test_harness.py``) proves this client
   can actually drive the GA ``GroupChatBuilder`` / ``ConcurrentBuilder``.

3. **Gated live arm.** ``live_local_client_or_skip`` builds an
   ``agent_framework.openai.OpenAIChatClient`` against an OpenAI-compatible LOCAL
   endpoint **directly** (the D2 ``LocalBackend`` seam is deliberately left un-wired
   until Fase 2, so ``src/`` stays untouched), or ``pytest.skip``s when the
   ``PORTFOLIO_LOCAL_*`` env is unset. No silent egress (D6).
"""

from __future__ import annotations

import os
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any

from agent_framework import (
    Agent,
    BaseChatClient,
    ChatResponse,
    ChatResponseUpdate,
    Message,
)


class BudgetExceeded(RuntimeError):
    """Raised the moment a runtime cap (tokens or rounds) is crossed (B4).

    Carries the breached ``kind`` ("tokens" | "rounds"), the ``limit`` that was
    set, and the ``observed`` value that crossed it — a structured stop event,
    never a silent hang.
    """

    def __init__(self, kind: str, limit: int, observed: int) -> None:
        self.kind = kind
        self.limit = limit
        self.observed = observed
        super().__init__(f"budget exceeded: {kind} limit={limit} observed={observed}")


@dataclass(frozen=True)
class Budget:
    """Hard token + round/iteration caps, required at startup (A4 / D6).

    Refuses to construct without positive caps — fail-fast, never an unbounded loop.
    """

    max_tokens: int
    max_rounds: int

    def __post_init__(self) -> None:
        if self.max_tokens <= 0:
            raise ValueError(f"max_tokens must be positive, got {self.max_tokens}")
        if self.max_rounds <= 0:
            raise ValueError(f"max_rounds must be positive, got {self.max_rounds}")


class TokenMeter:
    """Accumulates token and round usage against a ``Budget``; raises the moment
    a cap is crossed."""

    def __init__(self, budget: Budget) -> None:
        self.budget = budget
        self.tokens = 0
        self.rounds = 0

    def charge(self, tokens: int) -> int:
        """Add ``tokens`` to the running total; raise ``BudgetExceeded`` if over cap."""
        self.tokens += tokens
        if self.tokens > self.budget.max_tokens:
            raise BudgetExceeded("tokens", self.budget.max_tokens, self.tokens)
        return self.tokens

    def tick_round(self) -> int:
        """Increment the round counter; raise ``BudgetExceeded`` if over cap."""
        self.rounds += 1
        if self.rounds > self.budget.max_rounds:
            raise BudgetExceeded("rounds", self.budget.max_rounds, self.rounds)
        return self.rounds


def _word_tokens(text: str) -> int:
    """Token proxy: word count. Deterministic, endpoint-free."""
    return len(text.split())


def message_texts(messages: Sequence[Message]) -> list[str]:
    """Extract the text payloads from a sequence of MAF ``Message`` objects.

    MAF content items vary (objects with ``.text``, bare strings, or ``{"text": ...}``
    dicts); this normalizes them to a flat list of strings. Used by ``FakeChatClient``
    to record exactly what an agent *received* per call — the observable Spike B uses to
    detect cross-run conversation bleed (G2/B7)."""
    out: list[str] = []
    for m in messages:
        for c in getattr(m, "contents", []) or []:
            text = getattr(c, "text", None)
            if text is None and isinstance(c, str):
                text = c
            if text is None and isinstance(c, dict):
                text = c.get("text")
            if text is not None:
                out.append(str(text))
    return out


class FakeChatClient(BaseChatClient):
    """A deterministic, network-free ``BaseChatClient`` for driving MAF agents in tests.

    Returns scripted replies in order; once the script is exhausted it falls back to
    ``default_reply``. Counts "tokens" by word-count of each reply it emits, exposing
    ``total_tokens`` and ``call_count`` for the spike measurements.
    """

    OTEL_PROVIDER_NAME = "fake"

    def __init__(self, scripted: Sequence[str] | None = None, *, default_reply: str = "ok") -> None:
        super().__init__()
        self._scripted: list[str] = list(scripted or [])
        self._idx = 0
        self._default = default_reply
        self.total_tokens = 0
        self.call_count = 0
        # One entry per call: the text payloads this client RECEIVED that call. Lets a
        # spike observe whether a reused workflow feeds run N+1 the prior runs' history.
        self.received_texts: list[list[str]] = []

    def _next_reply(self) -> str:
        reply = self._scripted[self._idx] if self._idx < len(self._scripted) else self._default
        self._idx += 1
        self.call_count += 1
        self.total_tokens += _word_tokens(reply)
        return reply

    def _inner_get_response(
        self,
        *,
        messages: Sequence[Message],
        stream: bool,
        options: Any,
        **kwargs: Any,
    ) -> Any:
        # Matches the GA BaseChatClient contract: return a ResponseStream when
        # streaming, otherwise an awaitable resolving to a ChatResponse.
        self.received_texts.append(message_texts(messages))
        reply = self._next_reply()

        if stream:

            async def _agen() -> Any:
                yield ChatResponseUpdate(
                    role="assistant", contents=[{"type": "text", "text": reply}]
                )

            return self._build_response_stream(_agen())

        async def _coro() -> ChatResponse:
            return ChatResponse(
                messages=[Message(role="assistant", contents=[reply])],
                response_id="fake",
            )

        return _coro()


def fake_agent(
    client: BaseChatClient,
    name: str,
    instructions: str = "You are a terse participant. Answer in one short line.",
) -> Agent:
    """Build a minimal real ``agent_framework.Agent`` backed by ``client`` so the
    orchestration builders get a genuine participant."""
    return Agent(client, instructions, name=name)


def live_local_client_or_skip() -> Any:
    """Build an ``OpenAIChatClient`` against the OpenAI-compatible LOCAL endpoint
    (``PORTFOLIO_LOCAL_BASE_URL`` + ``PORTFOLIO_LOCAL_MODEL``), or ``pytest.skip``
    when unset.

    The D2 ``LocalBackend`` seam is intentionally NOT used here — its live wiring is a
    Fase 2 concern; the throwaway spike builds the client directly so ``src/`` stays
    untouched. No silent egress: without the env vars the live arm simply skips (D6).
    """
    base_url = os.environ.get("PORTFOLIO_LOCAL_BASE_URL")
    model = os.environ.get("PORTFOLIO_LOCAL_MODEL")
    if not base_url or not model:
        import pytest

        pytest.skip(
            "LOCAL endpoint not configured "
            "(set PORTFOLIO_LOCAL_BASE_URL and PORTFOLIO_LOCAL_MODEL to run the live arm)"
        )

    from agent_framework.openai import OpenAIChatClient

    # Most local OpenAI-compatible servers (Ollama / LM Studio) accept any non-empty
    # key; allow an override but default to a dummy so construction never blocks.
    api_key = os.environ.get("PORTFOLIO_LOCAL_API_KEY", "local")
    return OpenAIChatClient(model=model, api_key=api_key, base_url=base_url)