Source code for loom.ai.fastapi.response

"""Single-pass JSON response for the agent endpoints.

:class:`~loom.rest.fastapi.response.MsgspecJSONResponse` raises on any type
``msgspec`` cannot encode natively, and an agent's answer is decoded from the
model's output into whatever shape the artifact declared — including backend
types that arrive through a capability. A bodyless 500 at the last moment of a
paid run is the failure mode this module removes, exactly as
``loom.rest.fastapi.sql`` had to for query results.
"""

from __future__ import annotations

import base64
from collections.abc import Mapping
from ipaddress import IPv4Address, IPv6Address
from typing import Any

import msgspec
from starlette.responses import Response

from loom.ai.abc import AgentResult, FinalEvent
from loom.rest.fastapi.response import MsgspecJSONResponse


def encode_exotic(obj: Any) -> str:
    """Encode values ``msgspec`` does not handle natively.

    ``msgspec`` already covers datetime, date, UUID and Decimal; this hook adds
    IPv4/IPv6 addresses, bytes as base64, and a documented ``str()`` fallback so
    an unexpected type degrades into a readable value instead of a bodyless 500.

    Args:
        obj: Value the encoder could not serialise on its own.

    Returns:
        The textual form written to the response body.
    """
    if isinstance(obj, (IPv4Address, IPv6Address)):
        return str(obj)
    if isinstance(obj, bytes):
        return base64.b64encode(obj).decode("ascii")
    return str(obj)


ENCODER = msgspec.json.Encoder(enc_hook=encode_exotic)
"""Module-level encoder shared by the JSON and the SSE surfaces: built once."""


def result_payload(result: AgentResult | FinalEvent) -> Mapping[str, object]:
    """Project a completed run onto the published ``/run`` body or ``final`` frame.

    Both surfaces publish the same four keys, always present (``null`` when
    absent: a fixed shape). The result is projected rather than encoded
    wholesale so a field added to :class:`~loom.ai.abc.AgentResult` or
    :class:`~loom.ai.abc.FinalEvent` for the runtime's own use — the run's new
    ``messages`` — can never reach the wire by accident.

    Args:
        result: The completed run, or its terminal stream event.

    Returns:
        ``{"output", "usage", "interaction_id", "hook_result"}``.
    """
    return {
        "output": result.output,
        "usage": result.usage,
        "interaction_id": result.interaction_id,
        "hook_result": result.hook_result,
    }




class AgentJSONResponse(MsgspecJSONResponse):
    """Agent response encoded once by the module-level agent encoder.

    Example::

        return AgentJSONResponse(content=result_payload(result))
    """



    def render(self, content: object) -> bytes:
        """Encode *content* to JSON bytes in a single pass.

        Args:
            content: Value to serialise, typically the mapping
                :func:`result_payload` or :func:`error_response` built.

        Returns:
            The UTF-8 encoded JSON body.
        """
        return ENCODER.encode(content)




def error_response(
    status_code: int, code: str, message: str, *, interaction_id: str | None = None
) -> Response:
    """Build the flat ``{"code", "message", "interaction_id"}`` body agent surfaces refuse with.

    Both the HTTP surface and the A2A one answer a pre-first-byte failure with
    the same fields, so the shape is defined once here rather than once per
    transport. ``interaction_id`` is always on the wire — ``null`` when the
    failure happened before a run was admitted — because a fixed shape beats a
    conditional one.

    Args:
        status_code: HTTP status of the response.
        code: Stable machine-readable code.
        message: Human-readable description, safe to return to the caller.
        interaction_id: Identifier of the admitted run the failure belongs to,
            the only handle correlating it with the server-side log line.

    Returns:
        The encoded error response.

    Example::

        return error_response(404, "AGENT_NOT_FOUND", "no agent named 'x' is exposed")
    """
    return AgentJSONResponse(
        content={"code": code, "message": message, "interaction_id": interaction_id},
        status_code=status_code,
    )