harnessed 0.1.0__py3-none-any.whl

harnessed/__init__.py

@@ -0,0 +1,58 @@
+import logging
+from importlib import import_module
+from typing import Any
+
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+
+# Keep the standalone kernel independent of Pydantic and optional provider SDKs.
+_EXPORTS = {
+    "Agent": ("harnessed.agent", "Agent"),
+    "ApprovalAnswer": ("harnessed.core.messages", "ApprovalAnswer"),
+    "ApprovalRequest": ("harnessed.core.session.events", "ApprovalRequested"),
+    "Context": ("harnessed.kernel", "Context"),
+    "File": ("harnessed.core.attachments", "File"),
+    "Image": ("harnessed.core.attachments", "Image"),
+    "ImageDetail": ("harnessed.core.messages", "ImageDetail"),
+    "Inject": ("harnessed.kernel", "Inject"),
+    "Observation": ("harnessed.core.messages", "Observation"),
+    "OutputContract": ("harnessed.core.outputs", "OutputContract"),
+    "Plugin": ("harnessed.core.plugin", "Plugin"),
+    "PluginConfig": ("harnessed.core.plugin", "PluginConfig"),
+    "PluginState": ("harnessed.kernel", "PluginState"),
+    "PluginConflict": ("harnessed.core.errors", "PluginConflict"),
+    "PolicyDecision": ("harnessed.core.messages", "PolicyDecision"),
+    "Role": ("harnessed.core.messages", "Role"),
+    "RunHandle": ("harnessed.core.runtime.handles", "RunHandle"),
+    "RunResult": ("harnessed.core.runtime.result", "RunResult"),
+    "Section": ("harnessed.core.prompt.values", "Section"),
+    "Session": ("harnessed.core.runtime.session", "Session"),
+    "SessionClosed": ("harnessed.core.errors", "SessionClosed"),
+    "SessionStatus": ("harnessed.core.runtime.control", "SessionStatus"),
+    "StopReason": ("harnessed.core.model", "StopReason"),
+    "StreamOverflow": ("harnessed.core.errors", "StreamOverflow"),
+    "ToolCall": ("harnessed.core.messages", "ToolCall"),
+    "ToolContext": ("harnessed.core.tools.context", "ToolContext"),
+    "ToolError": ("harnessed.core.tools.registry", "ToolError"),
+    "ToolResult": ("harnessed.core.messages", "ToolResult"),
+    "ToolStatus": ("harnessed.core.messages", "ToolStatus"),
+    "hook": ("harnessed.kernel", "hook"),
+    "plugin": ("harnessed.kernel", "plugin"),
+    "tool": ("harnessed.core.tools.decorator", "tool"),
+}
+
+
+def __getattr__(name: str) -> Any:
+    if name not in _EXPORTS:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module, attribute = _EXPORTS[name]
+    value = getattr(import_module(module), attribute)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_EXPORTS))
+
+
+__all__ = sorted(_EXPORTS)

harnessed/__init__.pyi

@@ -0,0 +1,68 @@
+from harnessed.agent import Agent as Agent
+from harnessed.core.attachments import File as File
+from harnessed.core.attachments import Image as Image
+from harnessed.core.errors import PluginConflict as PluginConflict
+from harnessed.core.errors import SessionClosed as SessionClosed
+from harnessed.core.errors import StreamOverflow as StreamOverflow
+from harnessed.core.messages import ApprovalAnswer as ApprovalAnswer
+from harnessed.core.messages import ImageDetail as ImageDetail
+from harnessed.core.messages import Observation as Observation
+from harnessed.core.messages import PolicyDecision as PolicyDecision
+from harnessed.core.messages import Role as Role
+from harnessed.core.messages import ToolCall as ToolCall
+from harnessed.core.messages import ToolResult as ToolResult
+from harnessed.core.messages import ToolStatus as ToolStatus
+from harnessed.core.model import StopReason as StopReason
+from harnessed.core.outputs import OutputContract as OutputContract
+from harnessed.core.plugin import Plugin as Plugin
+from harnessed.core.plugin import PluginConfig as PluginConfig
+from harnessed.core.prompt.values import Section as Section
+from harnessed.core.runtime.control import SessionStatus as SessionStatus
+from harnessed.core.runtime.handles import RunHandle as RunHandle
+from harnessed.core.runtime.result import RunResult as RunResult
+from harnessed.core.runtime.session import Session as Session
+from harnessed.core.session.events import ApprovalRequested as ApprovalRequest
+from harnessed.core.tools.context import ToolContext as ToolContext
+from harnessed.core.tools.decorator import tool as tool
+from harnessed.core.tools.registry import ToolError as ToolError
+from harnessed.kernel import Context as Context
+from harnessed.kernel import Inject as Inject
+from harnessed.kernel import PluginState as PluginState
+from harnessed.kernel import hook as hook
+from harnessed.kernel import plugin as plugin
+
+from . import events as events
+
+__all__ = [
+    "Agent",
+    "ApprovalAnswer",
+    "ApprovalRequest",
+    "Context",
+    "File",
+    "Image",
+    "ImageDetail",
+    "Inject",
+    "Observation",
+    "OutputContract",
+    "Plugin",
+    "PluginConfig",
+    "PluginConflict",
+    "PolicyDecision",
+    "Role",
+    "RunHandle",
+    "RunResult",
+    "Section",
+    "Session",
+    "SessionClosed",
+    "SessionStatus",
+    "StopReason",
+    "StreamOverflow",
+    "ToolCall",
+    "ToolContext",
+    "ToolError",
+    "ToolResult",
+    "ToolStatus",
+    "hook",
+    "plugin",
+    "tool",
+]

harnessed/agent.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+
+from collections.abc import AsyncGenerator, Iterable
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, Unpack, overload
+
+from harnessed.core.approvals import Approvals
+from harnessed.core.errors import PluginConflict
+from harnessed.core.event_registry import EventRegistry
+from harnessed.core.model import Model
+from harnessed.core.outputs import OutputContract, OutputTypes
+from harnessed.core.plugin import Service, plugin_contribution, plugin_events
+from harnessed.core.prompt.instructions import Instructions
+from harnessed.core.prompt.values import InstructionValues, PromptComposition
+from harnessed.core.runtime.agent_loop import AgentLoop
+from harnessed.core.runtime.opener import open_session
+from harnessed.core.runtime.options import RunOptions
+from harnessed.core.runtime.result import RunResult
+from harnessed.core.runtime.services import RuntimeServices
+from harnessed.core.runtime.session import Session
+from harnessed.core.session.repository import SessionRepository
+from harnessed.core.session.session import SessionLog
+from harnessed.core.session.store import MemorySessions
+from harnessed.core.tools.output_limit import ToolOutputLimit
+from harnessed.core.tracing import NoopTracer
+from harnessed.kernel import Context, PluginLike
+from harnessed.kernel.plugin import as_plugin_spec
+from harnessed.models.pydantic_ai import PydanticAIModel
+
+if TYPE_CHECKING:
+    from harnessed.models.pydantic_ai import NativeModel
+
+
+class Agent:
+    """Resolve model, instructions and plugins into a reusable configuration.
+
+    Each opened session owns its runtime and plugin lifetimes. The configured
+    session-store plugin retains conversations and enforces a single writer.
+    """
+
+    def __init__(
+        self,
+        *,
+        model: Model | str | NativeModel | None = None,
+        instructions: InstructionValues | None = None,
+        plugins: Iterable[PluginLike] = (),
+    ) -> None:
+        supplied: list[PluginLike] = list(plugins)
+        if instructions is not None:
+            supplied.insert(0, Instructions(instructions))
+        if model is not None:
+            supplied.insert(0, Service(Model, model))
+        normalized: list[PluginLike] = []
+        for value in supplied:
+            if isinstance(value, Service) and value.key is Model:
+                implementation = value.value
+                if not isinstance(implementation, Model):
+                    implementation = PydanticAIModel(implementation)
+                value = Service(Model, implementation)
+            normalized.append(value)
+        declared = next((value for value in normalized if isinstance(value, Instructions)), None)
+        self.prompt = PromptComposition(
+            declared.values if declared is not None else (),
+            (plugin_contribution(value) for value in normalized),
+        )
+        defaults: list[PluginLike] = [
+            AgentLoop(),
+            MemorySessions(),
+            ToolOutputLimit(),
+            Approvals(),
+            NoopTracer(),
+            Instructions(composition=self.prompt),
+        ]
+        occupied = {as_plugin_spec(value).slot: index for index, value in enumerate(defaults)}
+        explicit: set[str] = set()
+        for value in normalized:
+            if isinstance(value, Instructions):
+                value = Instructions(value.values, composition=self.prompt)
+            spec = as_plugin_spec(value)
+            if spec.slot is not None:
+                if spec.slot in explicit:
+                    raise PluginConflict(f"Multiple plugins in slot {spec.slot!r}")
+                explicit.add(spec.slot)
+                if spec.slot in occupied:
+                    defaults[occupied[spec.slot]] = value
+                    continue
+                occupied[spec.slot] = len(defaults)
+            defaults.append(value)
+        self.plugins = (RuntimeServices(), *defaults)
+        self._events = EventRegistry(
+            kind for plugin in self.plugins for kind in plugin_events(plugin)
+        )
+        self.sessions = SessionRepository(self.plugins, registry=self._events)
+
+    @asynccontextmanager
+    async def session(
+        self,
+        session_id: str | SessionLog | None = None,
+        *,
+        restore_outputs: Iterable[type[Any] | OutputContract[Any]] = (),
+    ) -> AsyncGenerator[Session]:
+        """Open or resume a conversation; omit its ID to create a new one."""
+        async with Context() as ctx:
+            ctx.provide(EventRegistry, self._events.copy())
+            # Validators are per session, so dynamic output types never leak between them.
+            outputs = OutputTypes()
+            for output in restore_outputs:
+                outputs.register(output)
+            ctx.provide(OutputTypes, outputs)
+            await ctx.install(self.plugins, settle=True)
+            yield await open_session(ctx, session_id)
+
+    @overload
+    async def run[T](
+        self, prompt: str, *, output: type[T] | OutputContract[T], **options: Unpack[RunOptions]
+    ) -> RunResult[T]: ...
+
+    @overload
+    async def run(
+        self, prompt: str, *, output: None = None, **options: Unpack[RunOptions]
+    ) -> RunResult[None]: ...
+
+    async def run(
+        self,
+        prompt: str,
+        *,
+        output: type[Any] | OutputContract[Any] | None = None,
+        **options: Unpack[RunOptions],
+    ) -> RunResult[Any]:
+        """Run one task in a new session."""
+        async with self.session() as session:
+            return await session.run(prompt, output=output, **options)

harnessed/core/approvals.py

@@ -0,0 +1,291 @@
+import inspect
+from collections.abc import Awaitable, Callable, Iterable
+from dataclasses import dataclass, replace
+from typing import Annotated, Protocol, Self, runtime_checkable
+
+from pydantic import ConfigDict, Field, TypeAdapter
+
+from harnessed.core.hooks import ToolPolicy
+from harnessed.core.messages import ApprovalAnswer, PolicyDecision
+from harnessed.core.plugin import Plugin, PluginConfig
+from harnessed.core.session.events import (
+    ApprovalRequested,
+    ApprovalResolved,
+    ApprovalRevoked,
+    ApprovalStatus,
+    SessionEvent,
+)
+from harnessed.core.tools.references import ToolReference, ToolSelector, tool_selector
+from harnessed.core.tools.registry import ToolRegistry
+from harnessed.kernel import Context, Inject, Next, Priority, hook
+
+from .messages import ToolCall
+
+
+@runtime_checkable
+class Approver(Protocol):
+    async def ask(self, request: ApprovalRequested) -> ApprovalAnswer | None: ...
+
+
+type ApprovalCallback = Callable[
+    [ApprovalRequested], ApprovalAnswer | bool | None | Awaitable[ApprovalAnswer | bool | None]
+]
+
+
+class CallbackApprover:
+    """Adapt a request callback to the narrow Approver service."""
+
+    def __init__(self, callback: ApprovalCallback) -> None:
+        self.callback = callback
+
+    async def ask(self, request: ApprovalRequested) -> ApprovalAnswer | None:
+        value = self.callback(request)
+        value = await value if inspect.isawaitable(value) else value
+        if value is None or isinstance(value, ApprovalAnswer):
+            return value
+        if type(value) is bool:
+            return ApprovalAnswer.ALLOW if value else ApprovalAnswer.DENY
+        raise TypeError("Approvers must return ApprovalAnswer or bool")
+
+
+class ApprovalGateConfig(PluginConfig):
+    risky: Annotated[
+        tuple[Annotated[str, Field(min_length=1)] | Callable | ToolReference, ...],
+        Field(strict=False),
+    ]
+    ask: ApprovalCallback | None = None
+    scope: Annotated[str, Field(min_length=1)] = "approval-gate"
+    priority: int = Priority.GUARD
+
+
+class ApprovalGate(Plugin):
+    """Keep a scoped guard active even when its optional approver is unavailable."""
+
+    Config = ApprovalGateConfig
+
+    config: ApprovalGateConfig
+
+    def __init__(
+        self,
+        *,
+        risky: Iterable[ToolSelector],
+        ask: ApprovalCallback | None = None,
+        scope: str = "approval-gate",
+        priority: int = Priority.GUARD,
+    ) -> None:
+        super().__init__(risky=risky, ask=ask, scope=scope, priority=priority)
+        self.risky = tuple(tool_selector(item) for item in self.config.risky)
+        self.ask = self.config.ask
+
+    async def setup(self, ctx: Context, tools: Inject[ToolRegistry]) -> None:
+        self.risky = tuple(tools.reference(item, owner=ctx) for item in self.risky)
+        ctx.on(ToolPolicy, self._gate, priority=self.config.priority)
+
+    async def _gate(
+        self, event: ToolPolicy, next_: Next[PolicyDecision], tools: Inject[ToolRegistry]
+    ) -> PolicyDecision:
+        if not any(tools.matches(item, event.call.name) for item in self.risky):
+            return await next_()
+        answer = await event.session.approvals.request(
+            event.call,
+            f"{event.call.name} is marked risky",
+            scope=self.config.scope,
+            approver=CallbackApprover(self.ask) if self.ask is not None else None,
+        )
+        if answer not in (ApprovalAnswer.ALLOW, ApprovalAnswer.ALWAYS):
+            return PolicyDecision.DENY
+        return await next_()
+
+
+_APPROVAL_TIMEOUT = TypeAdapter(
+    Annotated[float, Field(gt=0)],
+    config=ConfigDict(
+        title="approval_timeout", strict=True, allow_inf_nan=False, hide_input_in_errors=True
+    ),
+)
+
+
+def validate_approval_timeout(value: float) -> None:
+    _APPROVAL_TIMEOUT.validate_python(value)
+
+
+@dataclass(frozen=True)
+class Approval:
+    request: ApprovalRequested
+    resolution: ApprovalResolved | None = None
+    revoked: bool = False
+
+    @property
+    def request_id(self) -> str:
+        return self.request.request_id
+
+    @property
+    def pending(self) -> bool:
+        return self.resolution is None
+
+    @property
+    def decision(self) -> ApprovalAnswer | None:
+        return self.resolution.decision if self.resolution is not None else None
+
+    @property
+    def is_grant(self) -> bool:
+        return (
+            not self.revoked
+            and self.decision is ApprovalAnswer.ALWAYS
+            and self.request.grant_id is None
+        )
+
+
+@dataclass(frozen=True, init=False)
+class ApprovalState:
+    records: tuple[Approval, ...]
+
+    @classmethod
+    def from_records(cls, records: Iterable[Approval]) -> Self:
+        snapshot = object.__new__(cls)
+        object.__setattr__(snapshot, "records", tuple(records))
+        return snapshot
+
+    def __init__(self, events: Iterable[SessionEvent]) -> None:
+        records: dict[str, Approval] = {}
+        for event in events:
+            record = approval_transition(records, event)
+            if record is not None:
+                records[record.request_id] = record
+        object.__setattr__(self, "records", tuple(records.values()))
+
+    @property
+    def pending(self) -> tuple[Approval, ...]:
+        return tuple(record for record in self.records if record.pending)
+
+    @property
+    def grants(self) -> tuple[Approval, ...]:
+        return tuple(record for record in self.records if record.is_grant)
+
+    def grant_for(self, tool: str, *, scope: str) -> Approval | None:
+        return next(
+            (
+                record
+                for record in reversed(self.grants)
+                if record.request.call.name == tool and record.request.scope == scope
+            ),
+            None,
+        )
+
+
+def approval_transition(records, event: SessionEvent) -> Approval | None:
+    if isinstance(event, ApprovalRequested):
+        if event.request_id in records:
+            raise ValueError("Duplicate approval request_id")
+        if event.grant_id is not None:
+            grant = records.get(event.grant_id)
+            if (
+                grant is None
+                or not grant.is_grant
+                or grant.request.scope != event.scope
+                or grant.request.call.name != event.call.name
+            ):
+                raise ValueError(
+                    "Cached approval must reference an active grant for its tool and scope"
+                )
+        return Approval(event)
+    elif isinstance(event, ApprovalResolved):
+        record = records.get(event.request_id)
+        if record is None or not record.pending:
+            raise ValueError("Approval resolution requires a pending request")
+        if record.request.grant_id is not None and event.status == ApprovalStatus.ANSWERED:
+            grant = records[record.request.grant_id]
+            if not grant.is_grant or event.decision is not ApprovalAnswer.ALWAYS:
+                raise ValueError("Cached approval grant is no longer active")
+        return replace(record, resolution=event)
+    elif isinstance(event, ApprovalRevoked):
+        record = records.get(event.request_id)
+        if record is None or not record.is_grant:
+            raise ValueError("Approval revocation requires an active ALWAYS grant")
+        return replace(record, revoked=True)
+    return None
+
+
+class ApprovalsConfig(PluginConfig):
+    timeout: Annotated[float, Field(gt=0)] = 300.0
+
+
+class Approvals(Plugin, slot="approvals"):
+    """Compose default decisions and explicit tool rules; required approval wins trust."""
+
+    Config = ApprovalsConfig
+
+    config: ApprovalsConfig
+
+    def __init__(
+        self,
+        *,
+        approver: Approver | ApprovalCallback | None = None,
+        default: PolicyDecision = PolicyDecision.ASK,
+        trust: Iterable[ToolSelector] = (),
+        require: Iterable[ToolSelector] = (),
+        deny: Iterable[ToolSelector] = (),
+        policy: Callable[[ToolCall], PolicyDecision | None] | None = None,
+        timeout: float = 300.0,
+    ) -> None:
+        super().__init__(timeout=timeout)
+        self.approver: Approver | None = (
+            approver
+            if approver is None or isinstance(approver, Approver)
+            else CallbackApprover(approver)
+        )
+        if not isinstance(default, PolicyDecision):
+            raise TypeError("default must be a PolicyDecision")
+        if any(isinstance(items, str) for items in (trust, require, deny)):
+            raise TypeError(
+                "Tool rules must be collections; use default=PolicyDecision.ALLOW to allow all"
+            )
+        self.default = default
+        self.trust = tuple(tool_selector(item) for item in trust)
+        self.require = tuple(tool_selector(item) for item in require)
+        self.deny = tuple(tool_selector(item) for item in deny)
+        self.policy = policy
+
+    async def setup(self, ctx: Context, tools: Inject[ToolRegistry]) -> None:
+        self.trust = tuple(tools.reference(item, owner=ctx) for item in self.trust)
+        self.require = tuple(tools.reference(item, owner=ctx) for item in self.require)
+        self.deny = tuple(tools.reference(item, owner=ctx) for item in self.deny)
+        if self.approver is not None:
+            # Tools and other guards asking mid-call reach the same human.
+            ctx.provide(Approver, self.approver)
+
+    @hook(ToolPolicy, priority=Priority.GUARD)
+    async def authorize(
+        self, event: ToolPolicy, next_: Next[PolicyDecision], tools: Inject[ToolRegistry]
+    ) -> PolicyDecision:
+        session = event.session
+        if any(tools.matches(item, event.call.name) for item in self.deny):
+            return PolicyDecision.DENY
+        required = any(tools.matches(item, event.call.name) for item in self.require)
+        decision = (
+            PolicyDecision.ASK
+            if required
+            else (self.policy(event.call) if self.policy is not None else None)
+        )
+        if decision is None:
+            decision = (
+                PolicyDecision.ALLOW
+                if any(tools.matches(item, event.call.name) for item in self.trust)
+                or tools.is_readonly(event.call.name)
+                else self.default
+            )
+        if decision is not None and not isinstance(decision, PolicyDecision):
+            raise TypeError("Approval policies must return PolicyDecision or None")
+        if decision is PolicyDecision.ASK:
+            answer = await session.approvals.request(
+                event.call,
+                "Tool may have side effects",
+                approver=self.approver,
+                timeout=self.config.timeout,
+            )
+            decision = (
+                PolicyDecision.DENY if answer is ApprovalAnswer.DENY else PolicyDecision.ALLOW
+            )
+        if decision is PolicyDecision.ALLOW:
+            return await next_()
+        return PolicyDecision.DENY

harnessed/core/attachments.py

@@ -0,0 +1,86 @@
+import asyncio
+import base64
+import json
+import mimetypes
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+from harnessed.core.messages import (
+    ContentPart,
+    ImageDetail,
+    ImagePart,
+    Message,
+    Role,
+    TextPart,
+    User,
+    freeze,
+    system,
+    to_dict,
+    user,
+)
+
+
+@dataclass(frozen=True)
+class File:
+    path: str | Path
+    encoding: str = "utf-8"
+
+    async def parts(self) -> tuple[ContentPart, ...]:
+        path = Path(self.path)
+        text = await asyncio.to_thread(path.read_text, encoding=self.encoding)
+        return (TextPart(f"File: {path.name}\n{text}"),)
+
+
+@dataclass(frozen=True)
+class Image:
+    path: str | Path
+    detail: ImageDetail = ImageDetail.AUTO
+
+    async def parts(self) -> tuple[ContentPart, ...]:
+        location = str(self.path)
+        if location.startswith(("https://", "http://", "data:image/")):
+            return (ImagePart(location, self.detail),)
+        path = Path(self.path)
+        mime, _ = mimetypes.guess_type(path.name)
+        if mime not in ("image/png", "image/jpeg", "image/gif", "image/webp"):
+            raise ValueError("Images must be PNG, JPEG, GIF or WebP")
+        data = await asyncio.to_thread(path.read_bytes)
+        encoded = base64.b64encode(data).decode("ascii")
+        return (ImagePart(f"data:{mime};base64,{encoded}", self.detail),)
+
+
+type Attachment = File | Image
+
+
+async def prepare_input(
+    content: str | Message,
+    *,
+    attachments: Iterable[Attachment] = (),
+    context: Mapping[str, Any] | None = None,
+) -> tuple[Message, tuple[Message, ...]]:
+    message = user(content) if isinstance(content, str) else content
+    attachments = tuple(attachments)
+    if attachments:
+        if message.role != Role.USER:
+            raise ValueError("Attachments belong to user messages")
+        parts = (
+            [TextPart(message.content)]
+            if isinstance(message.content, str)
+            else list(message.content)
+        )
+        for attachment in attachments:
+            parts.extend(await attachment.parts())
+        message = User(tuple(parts), id=message.id)
+    contributions = []
+    for name, value in (context or {}).items():
+        if not isinstance(name, str):
+            raise TypeError("Context keys must be strings")
+        text = (
+            value
+            if isinstance(value, str)
+            else json.dumps(to_dict(freeze(value)), ensure_ascii=False)
+        )
+        contributions.append(system(f"Task context [{name}]:\n{text}"))
+    return message, tuple(contributions)

harnessed/core/budget.py

@@ -0,0 +1,49 @@
+import inspect
+from typing import Annotated
+
+from pydantic import ConfigDict, Field
+from pydantic.dataclasses import dataclass
+
+from harnessed.core.errors import ContextBudgetExceeded
+from harnessed.core.model import ModelRequest
+from harnessed.core.tokens import TokenCounter
+
+
+def truncate(text: str, max_tokens: int, counter: TokenCounter) -> str:
+    low, high = 0, len(text)
+    while low < high:
+        middle = (low + high + 1) // 2
+        if counter.count_text(text[:middle]) <= max_tokens:
+            low = middle
+        else:
+            high = middle - 1
+    return text[:low]
+
+
+@dataclass(
+    frozen=True,
+    config=ConfigDict(strict=True, extra="forbid", allow_inf_nan=False, hide_input_in_errors=True),
+)
+class ContextBudget:
+    limit: Annotated[int, Field(ge=1)] | None = None
+    reserve_tokens: Annotated[int, Field(ge=0)] = 0
+
+    def window_limit(self, context_window: int | None) -> int | None:
+        if context_window is None:
+            return self.limit
+        return min(context_window, self.limit) if self.limit is not None else context_window
+
+    def output_reserve(self, request: ModelRequest) -> int:
+        return max(request.max_tokens or 0, self.reserve_tokens)
+
+    async def check(
+        self, request: ModelRequest, *, context_window: int | None, counter: TokenCounter
+    ) -> None:
+        limit = self.window_limit(context_window)
+        if limit is None:
+            return
+        counted = counter.count_request(request)
+        used = await counted if inspect.isawaitable(counted) else counted
+        reserve = self.output_reserve(request)
+        if used + reserve > limit:
+            raise ContextBudgetExceeded(used, reserve, limit)