smart-org/apps/api/tests/engine_fakes.py

"""Fakes for the engine tests: a scripted model and a scripted sandbox.

Two rules the engine tests insist on: the model never answers by calling a
real provider, and nothing ever executes on the host. Both fakes also *record*
what they were given, because half of what the engine must guarantee is about
prompts (no secret ever reaches one) and commands (git never runs).

``ScriptedModel.complete`` returns the queued responses in order and raises
the queued exceptions, so a test can script "read a file, run the tests, report"
exactly the way a real agent would behave — including the failures.
"""

from __future__ import annotations

from dataclasses import dataclass

from app.application.ports import ChatMessage, CommandResult, ModelHealth, ModelResponse
from app.application.ports.model_grants import GrantedModel
from app.application.ports.sandbox_provider import (
    SessionInfo,
    SnapshotInfo,
    WorkspaceInfo,
)


@dataclass(slots=True, kw_only=True)
class ScriptedReply:
    """One queued model answer: text, or an exception to raise instead."""

    content: str = ""
    raises: BaseException | None = None
    input_tokens: int = 10
    output_tokens: int = 5


class ScriptedModel:
    """A ``ModelProviderPort`` that replays a script.

    ``seen`` holds every call it received (system prompt + messages), so a test
    can assert what the model was actually told — the no-secrets rule is only
    testable this way.
    """

    def __init__(self, replies: list[ScriptedReply] | None = None) -> None:
        self.replies: list[ScriptedReply] = list(replies or [])
        self.seen: list[dict] = []
        self.calls = 0

    def queue(self, content: str, **kwargs) -> None:
        self.replies.append(ScriptedReply(content=content, **kwargs))

    async def complete(
        self,
        *,
        model,
        messages: list[ChatMessage],
        system: str = "",
        temperature: float = 0.2,
        max_tokens: int = 2048,
        credential=None,
        credentials=None,
    ) -> ModelResponse:
        self.calls += 1
        self.seen.append(
            {
                "model": model,
                "system": system,
                "messages": list(messages),
                "credential": credential,
            }
        )
        if not self.replies:
            return ModelResponse(
                content="[[WORK_REPORT]]پاسخ پیش‌فرض[[/WORK_REPORT]]",
                model_internal_name=getattr(model, "internal_name", ""),
                input_tokens=1,
                output_tokens=1,
            )
        reply = self.replies.pop(0)
        if reply.raises is not None:
            raise reply.raises
        return ModelResponse(
            content=reply.content,
            model_internal_name=getattr(model, "internal_name", ""),
            input_tokens=reply.input_tokens,
            output_tokens=reply.output_tokens,
        )

    async def health(self) -> ModelHealth:
        return ModelHealth(ok=True, detail="scripted")

    @property
    def prompts(self) -> list[str]:
        """Every system prompt the model was given, in order."""
        return [call["system"] for call in self.seen]

    def all_text(self) -> str:
        """Everything the model saw, for the "no secret" assertion."""
        parts = [call["system"] for call in self.seen]
        for call in self.seen:
            parts.extend(message.content for message in call["messages"])
        return "\n".join(parts)


class FakeSandbox:
    """A ``SandboxProviderPort`` over a dict filesystem and scripted commands.

    ``commands`` maps the command *prefix* to a queue of results, so a test can
    make the tests red once and green afterwards: the engine must notice.
    """

    def __init__(
        self,
        *,
        files: dict[str, str] | None = None,
        commands: dict[str, list[CommandResult]] | None = None,
    ) -> None:
        self.files: dict[str, str] = dict(files or {})
        self.commands: dict[str, list[CommandResult]] = {
            key: list(value) for key, value in (commands or {}).items()
        }
        self.executed: list[dict] = []
        self.writes: list[dict] = []
        self.reads: list[str] = []

    # -- SandboxProviderPort ------------------------------------------------

    async def create_workspace(
        self, *, organization_id: str, project_id: str, template_id: str, image: str
    ) -> WorkspaceInfo:
        return WorkspaceInfo(external_id="ws-1", status="ready")

    async def start_session(self, *, workspace_external_id: str) -> SessionInfo:
        return SessionInfo(external_id="sess-1", status="running")

    async def run_command(
        self,
        *,
        session_external_id: str,
        command: str,
        timeout_seconds: int,
        max_output_bytes: int,
        env: dict[str, str] | None = None,
    ) -> CommandResult:
        self.executed.append(
            {
                "command": command,
                "timeout_seconds": timeout_seconds,
                "max_output_bytes": max_output_bytes,
                "session": session_external_id,
            }
        )
        for prefix, queue in self.commands.items():
            if command.startswith(prefix) and queue:
                return queue.pop(0)
        return CommandResult(exit_code=0, stdout="ok")

    async def read_file(self, *, session_external_id: str, path: str) -> str:
        self.reads.append(path)
        if path not in self.files:
            raise FileNotFoundError(path)
        return self.files[path]

    async def write_file(self, *, session_external_id: str, path: str, content: str) -> None:
        self.files[path] = content
        self.writes.append({"path": path, "content": content})

    async def snapshot(self, *, session_external_id: str, name: str) -> SnapshotInfo:
        return SnapshotInfo(external_id="snap-1", name=name)

    async def pause(self, *, session_external_id: str) -> None:
        return None

    async def resume(self, *, session_external_id: str) -> SessionInfo:
        return SessionInfo(external_id="sess-1", status="running")

    async def destroy(self, *, session_external_id: str) -> None:
        return None

    async def get_preview_url(self, *, session_external_id: str, port: int) -> str:
        return f"http://127.0.0.1:{port}"


class FakeGrants:
    """A ``ModelGrantPort`` with an explicit allow-list."""

    def __init__(self, granted: list[str] | None = None) -> None:
        self._granted = list(granted or ["model-strong", "model-fast"])
        self.checks: list[str] = []

    async def granted_models(self, *, organization_id: str) -> list[GrantedModel]:
        return [
            GrantedModel(ai_model_id=model_id, tier="balanced")
            for model_id in self._granted
        ]

    async def is_granted(self, *, organization_id: str, ai_model_id: str) -> bool:
        self.checks.append(ai_model_id)
        return ai_model_id in self._granted

    def grant(self, model_id: str) -> None:
        self._granted.append(model_id)


class RecordingMemory:
    """A ``MemoryPort`` that records its queries and returns fixed context."""

    def __init__(self, context: str = "") -> None:
        self.context = context
        self.queries: list[dict] = []

    async def build_context(
        self, *, organization_id: str, role: str, query: str, budget_chars: int = 2500
    ) -> str:
        self.queries.append(
            {"organization_id": organization_id, "role": role, "query": query}
        )
        return self.context

    async def record(
        self,
        *,
        organization_id: str,
        kind: str,
        content: str,
        source: str,
        evidence_task_id: str | None = None,
        key: str | None = None,
    ) -> str | None:
        return None

    async def on_task_finished(self, *, organization_id: str, task) -> None:
        return None


__all__ = [
    "FakeGrants",
    "FakeSandbox",
    "RecordingMemory",
    "ScriptedModel",
    "ScriptedReply",
]