Agent Runtime & Turn Budgets¶

litestar_mcp.agent provides a guarded execution runtime (AgentRuntime), declarative agent specifications (Agent), type-driven tools (@tool and ToolContext), strict execution budgets (TurnBudget), and streaming model adapters (ModelClient, GoogleGenAIClient, and MockModelClient).

Defining an Agent¶

An Agent (litestar_mcp.agent.Agent) encapsulates instructions, direct tools, reusable SkillController classes or instances, and an optional model binding:

During __post_init__, Agent:

1. Instantiates any SkillController classes passed in skills once (agent.skill_instances).
2. Combines direct tools with every tool returned by skill.get_tools() into agent.tool_set.
3. Joins agent.instructions and non-empty skill.get_instructions() blocks with blank lines via agent.combined_instructions.
4. Rejects duplicate tool names or any user tool named "transfer_to_agent" (reserved for AgentGroup handoffs) with ImproperlyConfiguredException.

Authoring Tools with @tool and ToolContext¶

Decorate functions or SkillController methods with @tool (litestar_mcp.tool / litestar_mcp.core.tools.tool) to compile a callable into a Tool with a JSON Schema 2020-12 parameters contract and a msgspec argument validation struct:

class BaristaSkillController(SkillController):
    """Barista skill providing beverage menu lookup and order preparation tools."""

    name = "barista"
    description = "Specialized skills for querying menu items and preparing coffee orders"
    instructions = (
        "Barista Operational Grounding:\n"
        "- Suggest seasonal single-origin roasts when asked for recommendations.\n"
        "- Confirm milk alternatives (oat, almond, soy) before finalizing espresso drinks.\n"
        "- Never accept orders for items outside the published menu.\n"
    )

    def __init__(self, owner: Any = None) -> None:
        """Initialize the barista skill with catalog menu and empty order list."""
        super().__init__(owner)
        self.menu: dict[str, dict[str, Any]] = {
            "espresso": {"price": 3.50, "roast": "Ethiopian Yirgacheffe", "decaf_available": True},
            "flat_white": {"price": 4.75, "roast": "Guatemalan Antigua", "decaf_available": True},
            "cold_brew": {"price": 5.00, "roast": "Colombian Huila", "decaf_available": False},
        }
        self.orders: list[dict[str, Any]] = []

    @tool
    def get_menu(self) -> dict[str, dict[str, Any]]:
        """Retrieve available coffee drinks and pricing."""
        return self.menu

    @tool
    def place_order(self, beverage: str, milk: str = "whole", shots: int = 2) -> dict[str, Any]:
        """Place a customized beverage order."""
        bev = beverage.lower().replace(" ", "_")
        if bev not in self.menu:
            return {"status": "error", "message": f"'{beverage}' is not on the menu."}

        order_record = {
            "order_id": len(self.orders) + 1,
            "beverage": bev,
            "milk": milk,
            "shots": shots,
            "price": self.menu[bev]["price"],
            "status": "in_queue",
        }
        self.orders.append(order_record)
        return {"status": "success", "order": order_record}

Type-Based ToolContext and DI Injection¶

When @tool builds the tool's public JSON Schema (Tool.parameters), it inspects parameter type annotations and strips internal context and dependency injection parameters so the LLM only sees domain arguments:

• ToolContext (litestar_mcp.core.context.ToolContext): Any parameter annotated as ToolContext receives the active turn's ToolContext instance. You can also access it anywhere inside an executing tool via current_tool_context().
• ASGIConnection / Request: Any parameter annotated as Request or ASGIConnection receives ctx.request.
• Universal DI (Provide, FromDishka[T], Inject(), Depends(fn)): Parameters matched by controller/tool dependencies={"svc": Provide(...)}, Dishka FromDishka[T], or Annotated[T, Inject()] / Depends(fn) are stripped from Tool.parameters (tool.stripped_parameters) and resolved inside an AsyncExitStack on each await tool.execute(arguments, context=ctx) call.
• Sync and Async Callables: Async functions are awaited directly; synchronous functions are offloaded via anyio.to_thread.run_sync.

ToolContext carries the caller's security and session metadata across HTTP requests, A2A tasks, and queued worker turns:

AgentRuntime & Litestar Plugin Integration¶

AgentRuntime (litestar_mcp.agent.AgentRuntime) implements Litestar's InitPluginProtocol and manages the multi-turn model invocation loop, parallel tool execution (execute_tool_calls), session locking/persistence, history compaction, and OpenTelemetry span emission.

Registering AgentRuntime in Litestar(plugins=[...]) automatically registers the "runtime" dependency provider (Provide(self._provide_runtime, sync_to_thread=False)) consumed by AgentChatController and custom route handlers:

app = Litestar(
    route_handlers=[AgentChatController, CoffeePromptController, BaristaSkillController],
    plugins=[
        AgentRuntime(setup_coffee_agent_group(), default_model=create_model_client()),
        LitestarMCP(),
    ],
)

Constructor Parameters¶

Execution Methods (TurnRequest & TurnResponse)¶

You can invoke AgentRuntime directly from Python code, background tasks, or custom Litestar handlers using TurnRequest:

• await runtime.run_turn(TurnRequest(...)) -> TurnResponse: Runs the turn loop to completion and returns a TurnResponse(session_id=..., turn_id=..., output=..., agent_name=..., messages=..., token_usage=...).
• runtime.stream(TurnRequest(...)) -> AsyncIterator[AgentStreamFrame]: Streams typed AgentStreamFrame objects (session, thought, delta, tool_call, tool_result, agent_transfer, complete, error) as the turn progresses.
• async with runtime.open_stream(TurnRequest(...)) as receive_stream:: Opens a background anyio task group running runtime.produce(turn, send) and yields a MemoryObjectReceiveStream[AgentStreamFrame].
• await runtime.get_history(context, session_id) -> list[AgentMessage] | None: Loads conversation history scoped to context.owner_key() (or ANONYMOUS_OWNER).

TurnBudget Guardrails & BudgetExceededError¶

Every turn executed by AgentRuntime is bounded by a TurnBudget (litestar_mcp.agent.TurnBudget) to prevent runaway tool loops, unbounded token consumption, or hung upstream model calls:

All three fields must be strictly positive (> 0); passing 0 or a negative value raises ValueError at TurnBudget construction time.

When BudgetExceededError (litestar_mcp.agent.BudgetExceededError, a subclass of LitestarMCPError) is raised:

1. AgentRuntime.produce records the error on the active OpenTelemetry span, emits a terminal AgentStreamFrame(event_type="error", payload={"error": ...}), and re-raises BudgetExceededError.
2. AgentChatController.create_turn (POST /agent/turns) catches BudgetExceededError and translates it into an HTTP 504 Gateway Timeout response.

ModelClient Protocol & MockModelClient¶

AgentRuntime communicates with LLMs through the runtime-checkable ModelClient protocol (litestar_mcp.agent.ModelClient):

from collections.abc import AsyncIterator, Sequence
from typing import Protocol
from litestar_mcp.agent import AgentMessage, ModelDelta
from litestar_mcp.core.tools import Tool


class ModelClient(Protocol):
    def stream_turn(
        self,
        *,
        system_instruction: str | None,
        messages: Sequence[AgentMessage],
        tools: Sequence[Tool],
    ) -> AsyncIterator[ModelDelta]: ...

Each yielded ModelDelta (litestar_mcp.agent.ModelDelta) carries an event_type ("delta", "thought", "tool_call", or "usage"), incremental text / thought, tool call metadata (tool_name, call_id, arguments, provider_metadata), or token counts (prompt_tokens, completion_tokens, thought_tokens).

Deterministic Testing with MockModelClient¶

For fast offline unit and integration tests without network calls or API keys, pass a MockModelClient (litestar_mcp.agent.MockModelClient) pre-loaded with scripted ModelDelta sequences per model iteration:

from litestar_mcp.agent import Agent, AgentRuntime, MockModelClient, ModelDelta, TurnRequest
from litestar_mcp.core.tools import tool


@tool
def get_weather(city: str) -> dict[str, str]:
    """Return current weather for a city."""
    return {"city": city, "forecast": "Sunny, 22C"}


mock_model = MockModelClient(
    responses=[
        [
            ModelDelta(
                event_type="tool_call",
                tool_name="get_weather",
                call_id="call_1",
                arguments={"city": "Seattle"},
            )
        ],
        [
            ModelDelta(event_type="thought", text="Formatting weather report for Seattle."),
            ModelDelta(event_type="delta", text="It is Sunny and 22C in Seattle."),
            ModelDelta(event_type="usage", prompt_tokens=24, completion_tokens=10, thought_tokens=6),
        ],
    ]
)

agent = Agent(name="weather_bot", tools=[get_weather], model=mock_model)
runtime = AgentRuntime(agent)

mock_model.call_history records every {"system_instruction": ..., "messages": [...], "tools": [...]} invocation so tests can assert exact prompt assembly and tool schemas.

GoogleGenAIClient (gemini-3.8-flash)¶

GoogleGenAIClient (litestar_mcp.agent.GoogleGenAIClient, enabled via litestar-mcp[google-genai]) wraps the official google-genai SDK with native support for Gemini thinking models, Vertex AI, and shared httpx2.AsyncClient connection pooling:

def create_model_client() -> GoogleGenAIClient:
    """Instantiate GoogleGenAIClient with a thinking budget and low temperature."""
    return GoogleGenAIClient(thinking_budget=2048, temperature=0.2)


def setup_coffee_agent_group() -> AgentGroup:
    """Build and configure the coffee shop agent group with specialists."""
    barista_agent = Agent(
        name="barista",
        description="Prepares coffee drinks, looks up menu pricing, and manages customer orders.",
        instructions="You are the lead barista. Handle menu questions and beverage customizations.",
        skills=[BaristaSkillController],
    )

    triage_coordinator = Agent(
        name="concierge",
        description="Greets customers and routes coffee questions to the barista specialist.",
        instructions="Greet customers warmly. Delegate ordering and coffee questions to the barista.",
    )

    return AgentGroup(
        coordinator=triage_coordinator,
        specialists=[barista_agent],
    )

Configuration & Capabilities¶

Thought-Signature Preservation Across Function Calls¶

Gemini thinking models emit an opaque thought_signature on Part(function_call=...) chunks and require that exact signature to be echoed back in conversation history on subsequent function-response turns. GoogleGenAIClient handles this round-trip automatically:

1. When parsing a streaming candidate part containing a function_call, _parse_part_deltas captures part.thought_signature and function_call.id inside ModelDelta.provider_metadata["google"].
2. AgentRuntime stores provider_metadata on the ToolCall inside the assistant AgentMessage persisted in SessionStore.
3. On subsequent iterations and follow-up turns, _build_contents reconstructs types.Part(function_call=types.FunctionCall(id=..., name=..., args=...), thought_signature=google.get("thought_signature")) alongside the correlated types.FunctionResponse(id=..., name=..., response=...), preserving reasoning continuity across multi-step tool chains.