"""Agent context management.

This module provides the AgentContext class for managing session state
during agent execution. AgentContext holds a reference to an Environment
and derives file_operator/shell/resources from it via properties.

Architecture:
    AgentRuntime (recommended entry point)
      - Created by create_agent() factory function
      - Manages env -> ctx -> agent lifecycle as async context manager
      - stream_agent() handles runtime lifecycle automatically

    Environment (outer, long-lived)
      - Manages tmp_dir lifecycle
      - Creates and owns file_operator and shell
      - async with environment as env:

        AgentContext (inner, short-lived)
          - Manages session state (run_id, timing, handoff)
          - Holds env reference, derives file_operator/shell/resources from it
          - async with AgentContext(env=env) as ctx:

Example:
    Using create_agent and stream_agent (recommended)::

        from ya_agent_sdk.agents.main import create_agent, stream_agent

        # create_agent returns AgentRuntime (not a context manager)
        runtime = create_agent("openai-chat:gpt-4")

        # stream_agent manages runtime lifecycle automatically
        async with stream_agent(runtime, "Hello") as streamer:
            async for event in streamer:
                print(event)

    Using create_agent with manual agent.run::

        runtime = create_agent("openai-chat:gpt-4")
        async with runtime:  # Enter runtime to manage env/ctx/agent
            result = await runtime.agent.run("Hello", deps=runtime.ctx)
            print(result.output)

    Manual Environment and AgentContext setup (advanced)::

        from ya_agent_sdk.environment.local import LocalEnvironment
        from ya_agent_sdk.context import AgentContext

        async with LocalEnvironment() as env:
            async with AgentContext(env=env) as ctx:
                await ctx.file_operator.read_file("test.txt")

    Multiple sessions sharing environment::

        async with LocalEnvironment() as env:
            # First session
            async with AgentContext(env=env) as ctx1:
                await ctx1.file_operator.read_file("test.txt")

            # Second session (reuses same environment)
            async with AgentContext(env=env) as ctx2:
                ...
        # tmp_dir cleaned up when environment exits
"""

from __future__ import annotations

import asyncio
import json
from collections import defaultdict, deque
from collections.abc import Awaitable, Callable, Sequence
from contextlib import AbstractAsyncContextManager
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from enum import StrEnum
from pathlib import Path
from typing import Any, Self
from uuid import uuid4
from xml.dom.minidom import parseString
from xml.etree.ElementTree import Element, SubElement, tostring

from pydantic import BaseModel, Field, field_validator, model_validator
from pydantic_ai import (
    DeferredToolRequests,
    FunctionToolCallEvent,
    FunctionToolResultEvent,
    ModelSettings,
    OutputToolResultEvent,
    PartDeltaEvent,
    PartEndEvent,
    PartStartEvent,
    RunContext,
    ToolCallPartDelta,
    UserContent,
)
from pydantic_ai.capabilities import AbstractCapability, ProcessHistory
from pydantic_ai.messages import HandleResponseEvent as PydanticHandleResponseEvent
from pydantic_ai.messages import (
    ModelMessage,
    ModelMessagesTypeAdapter,
    ModelResponseStreamEvent,
    NativeToolCallPart,
    NativeToolReturnPart,
    RetryPromptPart,
    ToolCallPart,
    ToolReturnPart,
)
from pydantic_ai.models import Model
from pydantic_ai.usage import RunUsage
from pydantic_settings import BaseSettings, SettingsConfigDict
from ya_agent_environment import Environment, FileOperator, ResourceRegistry, Shell

from ya_agent_sdk.agents.lifecycle import LifecycleExtension
from ya_agent_sdk.events import AgentEvent, UsageSnapshotEvent
from ya_agent_sdk.usage import UsageAgentTotal, UsageSnapshot, UsageSnapshotEntry, coerce_run_usage
from ya_agent_sdk.utils import get_latest_request_usage

from .bus import BusMessage, MessageBus
from .note import NoteManager
from .tasks import TaskManager, TaskStatus

_PACKAGE_ROOT = Path(__file__).resolve().parents[2]

# =============================================================================
# Injected Context Tag Names
# =============================================================================

# XML tag names for per-turn injected context blocks.
# These are stripped from historical messages during compact to save tokens,
# since they are re-injected fresh on every request by filters.

RUNTIME_CONTEXT_TAG = "runtime-context"
"""Tag for runtime context (elapsed time, token usage, etc.). Generated by AgentContext."""

ENVIRONMENT_CONTEXT_TAG = "environment-context"
"""Tag for environment context (file system, shell config). Generated by Environment."""

PROJECT_GUIDANCE_TAG = "project-guidance"
"""Tag for project-level guidance (e.g., AGENTS.md). Injected by application layer."""

USER_RULES_TAG = "user-rules"
"""Tag for user-level rules (e.g., RULES.md). Injected by application layer."""


# =============================================================================
# Model Wrapper Type
# =============================================================================


ModelWrapper = Callable[[Model, str, dict[str, Any]], Model | Awaitable[Model]]
"""Type alias for model wrapper functions.

A model wrapper receives a pydantic-ai Model, an agent name, and a context dict,
and returns a wrapped Model. The wrapper can be sync or async.

Use cases include observability (Langfuse/DataDog), caching, rate limiting,
and cost tracking.

Args:
    model: The pydantic-ai Model to wrap.
    agent_name: Identifier for the agent using this model.
        Convention: 'main', 'debugger', 'video-understanding', etc.
    context: Additional context. May contain:
        - run_id (str): Current session identifier
        - (future fields will be documented as added)

Returns:
    Wrapped model (or original if no wrapping needed).
    Can return Model directly (sync) or Awaitable[Model] (async).

Example::

    # Sync wrapper
    def my_wrapper(
        model: Model,
        agent_name: str,
        context: dict[str, Any],
    ) -> Model:
        return LangfuseModel(model, name=agent_name, trace_id=context.get("run_id"))

    # Async wrapper
    async def async_wrapper(
        model: Model,
        agent_name: str,
        context: dict[str, Any],
    ) -> Model:
        await some_async_setup()
        return LangfuseModel(model, name=agent_name)

    runtime = create_agent(
        "openai-chat:gpt-4",
        model_wrapper=my_wrapper,
    )
"""


# =============================================================================
# Subagent Wrapper Type
# =============================================================================


SubagentWrapper = Callable[[str, str, dict[str, Any]], AbstractAsyncContextManager[None]]
"""Type alias for subagent execution wrapper functions.

A subagent wrapper receives agent_name, agent_id, and a metadata dict,
and returns an async context manager that wraps the entire subagent execution.
This enables creating observability spans (e.g., Langfuse agent observations)
that properly nest all subagent activity underneath.

Args:
    agent_name: Name of the subagent (e.g., 'debugger', 'searcher').
    agent_id: Unique identifier for this subagent invocation.
    metadata: Context dict from get_wrapper_metadata().

Returns:
    An async context manager wrapping the subagent execution.

Example::

    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def trace_subagent(name: str, id: str, meta: dict[str, Any]):
        with langfuse.start_as_current_observation(
            name=f"agent:{name}", as_type="agent",
        ):
            yield

    runtime = create_agent(
        "openai-chat:gpt-4",
        subagent_wrapper=trace_subagent,
    )
"""


# =============================================================================
# Model Capability
# =============================================================================


class ModelCapability(StrEnum):
    """Model capabilities that can be used to describe what a model supports."""

    vision = "vision"
    """Model can process and understand images."""

    video_understanding = "video_understanding"
    """Model can process and understand video content."""

    document_understanding = "document_understanding"
    """Model can process and understand documents (PDF, etc.)."""

    audio_understanding = "audio_understanding"
    """Model can process and understand audio content."""

    image_url = "image_url"
    """Model supports receiving images via URL (not just base64)."""

    video_url = "video_url"
    """Model supports receiving videos via URL."""

    reasoning_required = "reasoning_required"
    """Provider requires assistant messages to include reasoning content."""

    reasoning_foreign_incompatible = "reasoning_foreign_incompatible"
    """Provider rejects thinking parts tagged with a different provider."""


# =============================================================================
# Tool Settings
# =============================================================================


class ToolSettings(BaseSettings):
    """Tool-related settings from environment variables.

    API keys for various external services used by tools.
    All settings are loaded from environment variables or .env file.
    """

    model_config = SettingsConfigDict(
        env_file=(_PACKAGE_ROOT / ".env", ".env"),
        env_file_encoding="utf-8",
        extra="ignore",
        case_sensitive=False,
    )

    # Web search API keys
    google_search_api_key: str | None = None
    """Google Custom Search API key."""

    google_search_cx: str | None = None
    """Google Custom Search Engine ID."""

    tavily_api_key: str | None = None
    """Tavily API key for web search."""

    brave_search_api_key: str | None = None
    """Brave Search API key."""

    # Image search API keys
    pixabay_api_key: str | None = None
    """Pixabay API key for stock image search."""

    rapidapi_api_key: str | None = None
    """RapidAPI key for real-time image search."""

    # Web scraping API key
    firecrawl_api_key: str | None = None
    """Firecrawl API key for web scraping."""


def _get_tool_settings() -> ToolSettings:
    """Get ToolSettings instance.

    This function creates a new ToolSettings instance each time to ensure
    environment variables are read at the time of ToolConfig creation,
    not at module import time.
    """
    return ToolSettings()


# =============================================================================
# Stream Event Types
# =============================================================================

# Stream event type for the queue
# Includes pydantic-ai events + custom AgentEvent for user-defined events
AgentStreamEvent = ModelResponseStreamEvent | PydanticHandleResponseEvent | AgentEvent


def _create_stream_queue_factory() -> dict[str, asyncio.Queue[AgentStreamEvent]]:
    """Create a defaultdict factory for subagent stream queues."""
    return defaultdict(asyncio.Queue)


# =============================================================================
# Agent Info and Stream Event
# =============================================================================


@dataclass
class AgentInfo:
    """Metadata for a registered agent.

    Used to track agent identity and hierarchy in stream events.

    Attributes:
        agent_id: Unique identifier for the agent (e.g., "main" or 4-char short ID).
        agent_name: Human-readable name (e.g., "main", "search", "reasoning").
        parent_agent_id: ID of the parent agent, None for main agent.
    """

    agent_id: str
    agent_name: str
    parent_agent_id: str | None = None


@dataclass
class StreamEvent:
    """Stream event with agent identification.

    Wraps raw stream events with agent metadata for distinguishing
    events from different agents in a merged stream.

    Attributes:
        agent_id: ID of the agent that produced this event.
        agent_name: Name of the agent.
        event: The underlying stream event (ModelResponseStreamEvent, etc.).
    """

    agent_id: str
    agent_name: str
    event: AgentStreamEvent


# =============================================================================
# Tool ID Wrapper
# =============================================================================


class ToolIdWrapper:
    """Wrapper for tool call IDs to ensure stable, cross-provider compatible identifiers.

    This class normalizes tool call IDs from different LLM providers by mapping them to
    a consistent format ("ya-{uuid}"). This ensures:
    - Consistent ID format across all providers (OpenAI, Anthropic, Gemini, etc.)
    - Stable ID mapping within a session for proper tool call/result matching
    - Compatibility with all downstream systems expecting standardized IDs

    The wrapper maintains an internal mapping to ensure the same original ID always
    maps to the same normalized ID within a session.
    """

    def __init__(self) -> None:
        self._prefix = "ya-"
        self._tool_call_maps: dict[str, str] = {}

    def clear(self) -> None:
        """Clear the tool call ID mapping.

        This should be called when starting a new session to ensure
        fresh ID mappings.
        """
        self._tool_call_maps.clear()

    def upsert_tool_call_id(self, tool_call_id: str) -> str:
        """Normalize a tool call ID to the standard format.

        If the ID already has the standard prefix, return it unchanged.
        Otherwise, create and cache a new normalized ID.

        Args:
            tool_call_id: The original tool call ID from any provider.

        Returns:
            Normalized tool call ID with "ya-" prefix.
        """
        if tool_call_id.startswith(self._prefix):
            return tool_call_id

        if tool_call_id not in self._tool_call_maps:
            self._tool_call_maps[tool_call_id] = f"{self._prefix}{uuid4().hex}"
        return self._tool_call_maps[tool_call_id]

    def wrap_event(self, event: AgentStreamEvent) -> AgentStreamEvent:
        match event:
            case FunctionToolCallEvent():
                event.part.tool_call_id = self.upsert_tool_call_id(event.tool_call_id)
            case FunctionToolResultEvent() | OutputToolResultEvent():
                event.part.tool_call_id = self.upsert_tool_call_id(event.tool_call_id)
            case PartStartEvent() | PartEndEvent():
                if isinstance(
                    event.part,
                    (ToolCallPart, NativeToolCallPart, ToolReturnPart, NativeToolReturnPart, RetryPromptPart),
                ):
                    event.part.tool_call_id = self.upsert_tool_call_id(event.part.tool_call_id)
            case PartDeltaEvent():
                if isinstance(event.delta, ToolCallPartDelta) and event.delta.tool_call_id:
                    event.delta.tool_call_id = self.upsert_tool_call_id(event.delta.tool_call_id)
        return event

    def wrap_deferred_tool_requests(self, deferred_tool_requests: DeferredToolRequests) -> DeferredToolRequests:
        metadata = deferred_tool_requests.metadata or {}
        wrapped_metadata: dict[str, dict[str, Any]] = {}
        for call in deferred_tool_requests.calls or []:
            original_tool_call_id = call.tool_call_id
            call.tool_call_id = self.upsert_tool_call_id(call.tool_call_id)
            if original_tool_call_id in metadata:
                wrapped_metadata[call.tool_call_id] = metadata[original_tool_call_id]
        for approval in deferred_tool_requests.approvals or []:
            original_tool_call_id = approval.tool_call_id
            approval.tool_call_id = self.upsert_tool_call_id(approval.tool_call_id)
            if original_tool_call_id in metadata:
                wrapped_metadata[approval.tool_call_id] = metadata[original_tool_call_id]
        if wrapped_metadata:
            deferred_tool_requests.metadata = wrapped_metadata
        return deferred_tool_requests

    def wrap_messages(
        self,
        _: RunContext[AgentContext] | None,
        message_history: list[ModelMessage],
    ) -> list[ModelMessage]:
        """Normalize all tool call IDs in the message history.

        This method can be used directly as a pydantic-ai history_processor.

        Args:
            _: RunContext (unused, required by history_processor signature).
            message_history: List of messages to process.

        Returns:
            The same message history with normalized tool call IDs.
        """
        for m in message_history:
            for p in m.parts:
                if isinstance(
                    p, (ToolCallPart, NativeToolCallPart, ToolReturnPart, NativeToolReturnPart, RetryPromptPart)
                ):
                    p.tool_call_id = self.upsert_tool_call_id(p.tool_call_id)
        return message_history


# =============================================================================
# Type Aliases (defined after AgentContext is visible as a name)
# =============================================================================

# Hook function type for converting media data to URL.
# Can be sync or async. Returns URL string or None to use default behavior.
MediaToUrlHook = Callable[[RunContext["AgentContext"], bytes, str], Awaitable[str | None] | str | None]


# =============================================================================
# Tool Config
# =============================================================================


class ShellReviewAction(StrEnum):
    """Action when shell command review requires approval."""

    DEFER = "defer"
    DENY = "deny"


class ShellReviewRiskLevel(StrEnum):
    """Shell command review risk levels."""

    LOW = "low"
    MEDIUM = "medium"
    HIGH = "high"
    EXTRA_HIGH = "extra_high"


class ShellReviewConfig(BaseModel):
    """Configuration for shell command safety review."""

    model_config = {"arbitrary_types_allowed": True}

    enabled: bool = False
    model: str | None = None
    model_settings: dict[str, Any] | None = None
    on_needs_approval: ShellReviewAction = ShellReviewAction.DEFER
    risk_threshold: ShellReviewRiskLevel = ShellReviewRiskLevel.HIGH
    system_prompt: str | None = None

    @field_validator("model_settings", mode="before")
    @classmethod
    def resolve_model_settings_preset(
        cls,
        value: ModelSettings | str | dict[str, Any] | None,
    ) -> dict[str, Any] | None:
        """Resolve model settings preset names at the SDK config boundary."""
        if value is None:
            return None
        from ya_agent_sdk.presets import resolve_model_settings

        return resolve_model_settings(value)

    @model_validator(mode="after")
    def validate_model_when_enabled(self) -> Self:
        if self.enabled and (self.model is None or self.model.strip() == ""):
            msg = "shell_review.model is required when shell review is enabled."
            raise ValueError(msg)
        return self


class SecurityConfig(BaseModel):
    """Security-related runtime configuration."""

    shell_review: ShellReviewConfig | None = None
    """Optional shell command safety review configuration."""


class ToolConfig(BaseModel):
    """Tool-level configuration for fine-grained control.

    API keys can be passed directly or loaded from environment variables
    via ToolSettings. See .env.example for available environment variables.

    Note: Environment variables are read when ToolConfig is instantiated,
    not when the module is imported. This allows callers to set environment
    variables after importing the module.

    Extensibility:
        This class supports two extension patterns:

        1. Inheritance (recommended for type safety)::

            class MyToolConfig(ToolConfig):
                my_api_key: str | None = None

            class MyContext(AgentContext):
                tool_config: MyToolConfig = Field(default_factory=MyToolConfig)

        2. Extra attributes (for quick prototyping)::

            config = ToolConfig(my_custom_key="value")
            config.my_custom_key  # Accessible but not type-checked
    """

    model_config = {"arbitrary_types_allowed": True, "extra": "allow"}

    skip_url_verification: bool = True
    """Skip SSRF URL verification. Default enabled (skip). Set to False to enable verification for public-facing environments."""

    enable_load_document: bool = False
    """Enable document URL parsing in LoadTool. Default disabled due to poor model support."""

    image_understanding_model: str | None = None
    """Model to use for image understanding. Falls back to AgentSettings.image_understanding_model."""

    image_understanding_model_settings: ModelSettings | None = None
    """Model settings for image understanding agent."""

    video_understanding_model: str | None = None
    """Model to use for video understanding. Falls back to AgentSettings.video_understanding_model."""

    video_understanding_model_settings: ModelSettings | None = None
    """Model settings for video understanding agent."""

    audio_understanding_model: str | None = None
    """Model to use for audio understanding. Falls back to AgentSettings.audio_understanding_model."""

    audio_understanding_model_settings: ModelSettings | None = None
    """Model settings for audio understanding agent."""

    # Web search API keys
    google_search_api_key: str | None = Field(default_factory=lambda: _get_tool_settings().google_search_api_key)
    """Google Custom Search API key."""

    google_search_cx: str | None = Field(default_factory=lambda: _get_tool_settings().google_search_cx)
    """Google Custom Search Engine ID."""

    tavily_api_key: str | None = Field(default_factory=lambda: _get_tool_settings().tavily_api_key)
    """Tavily API key for web search."""

    brave_search_api_key: str | None = Field(default_factory=lambda: _get_tool_settings().brave_search_api_key)
    """Brave Search API key."""

    # Image search API keys
    pixabay_api_key: str | None = Field(default_factory=lambda: _get_tool_settings().pixabay_api_key)
    """Pixabay API key for stock image search."""

    rapidapi_api_key: str | None = Field(default_factory=lambda: _get_tool_settings().rapidapi_api_key)
    """RapidAPI key for real-time image search."""

    # Web scraping API key
    firecrawl_api_key: str | None = Field(default_factory=lambda: _get_tool_settings().firecrawl_api_key)
    """Firecrawl API key for web scraping."""

    view_max_text_file_size: int = 10 * 1024 * 1024
    """Maximum text file size in bytes that the view tool will inspect."""

    view_relaxed_text_patterns: Sequence[str] = Field(default_factory=tuple)
    """Ripgrep-style glob patterns or re: patterns where text view uses relaxed truncation limits.

    Glob patterns use the same semantics as filesystem glob/grep. Patterns prefixed with
    re: are matched as regular expressions against the agent-facing file path.
    """

    view_relaxed_text_dynamic_patterns: dict[str, tuple[str, ...]] = Field(default_factory=dict, exclude=True)
    """Runtime-registered relaxed text patterns keyed by source/toolset id."""

    def iter_view_relaxed_text_patterns(self) -> tuple[str, ...]:
        """Return static and runtime-registered relaxed text patterns."""
        dynamic_patterns = tuple(
            pattern for patterns in self.view_relaxed_text_dynamic_patterns.values() for pattern in patterns
        )
        return (*self.view_relaxed_text_patterns, *dynamic_patterns)

    def register_view_relaxed_text_patterns(self, source: str, patterns: Sequence[str]) -> None:
        """Register runtime relaxed text patterns for a source id."""
        normalized = tuple(pattern.strip() for pattern in patterns if pattern.strip())
        if normalized:
            self.view_relaxed_text_dynamic_patterns[source] = normalized
        else:
            self.view_relaxed_text_dynamic_patterns.pop(source, None)

    def unregister_view_relaxed_text_patterns(self, source: str) -> None:
        """Remove runtime relaxed text patterns registered by a source id."""
        self.view_relaxed_text_dynamic_patterns.pop(source, None)

    view_relaxed_text_file_size: int = 50 * 1024 * 1024
    """Maximum text file size in bytes for paths matched by view_relaxed_text_patterns."""

    view_relaxed_line_limit: int = 5000
    """Line limit for relaxed paths when the effective line_limit is the normal default."""

    view_relaxed_max_line_length: int = 20000
    """Max line length for relaxed paths when the effective max_line_length is the normal default."""

    view_max_content_chars: int = 60000
    """Maximum returned text characters for normal view output before content truncation."""

    view_relaxed_max_content_chars: int = 250000
    """Maximum returned text characters for relaxed view output before content truncation."""

    edit_max_file_size: int = 20 * 1024 * 1024
    """Maximum file size in bytes that edit/multi_edit tools will process. Default: 20 MB."""

    grep_max_file_size: int = 10 * 1024 * 1024
    """Maximum file size in bytes that grep will read per file. Default: 10 MB."""

    view_max_inline_image_bytes: int = 20 * 1024 * 1024
    """Maximum image size in bytes that the view tool will inline into context."""

    view_max_inline_video_bytes: int = 50 * 1024 * 1024
    """Maximum video size in bytes that the view tool will inline into context."""

    view_max_inline_audio_bytes: int = 50 * 1024 * 1024
    """Maximum audio size in bytes that the view tool will inline into context."""

    fetch_stream_chunk_size: int = 64 * 1024
    """Chunk size in bytes for streamed HTTP reads in fetch/download tools."""

    fetch_max_inline_binary_bytes: int = 30 * 1024 * 1024
    """Maximum binary response size in bytes that fetch will buffer in memory."""

    download_max_concurrency: int = 4
    """Maximum number of concurrent downloads for the download tool."""

    document_max_file_size: int = 200 * 1024 * 1024
    """Maximum file size in bytes for document conversion tools (PDF, Office). Default: 200 MB."""

    @field_validator("download_max_concurrency")
    @classmethod
    def _validate_download_max_concurrency(cls, v: int) -> int:
        if v < 1:
            raise ValueError("download_max_concurrency must be >= 1")
        return v

    # Media to URL conversion hooks
    image_to_url_hook: MediaToUrlHook | None = None
    """Hook to convert image data to URL.

    Args:
        ctx: RunContext with AgentContext
        image_data: Raw image bytes
        media_type: MIME type (e.g., 'image/png')

    Returns:
        Publicly accessible URL string, or None to use default BinaryContent behavior.
        Can be sync or async function.

    Note:
        The returned URL must be publicly accessible by the LLM provider.
        Empty strings are treated as None (fallback to default behavior).
    """

    video_to_url_hook: MediaToUrlHook | None = None
    """Hook to convert video data to URL.

    Args:
        ctx: RunContext with AgentContext
        video_data: Raw video bytes
        media_type: MIME type (e.g., 'video/mp4')

    Returns:
        Publicly accessible URL string, or None to use default BinaryContent behavior.
        Can be sync or async function.

    Note:
        The returned URL must be publicly accessible by the LLM provider.
        Empty strings are treated as None (fallback to default behavior).
    """

    audio_to_url_hook: MediaToUrlHook | None = None
    """Hook to convert audio data to URL.

    Args:
        ctx: RunContext with AgentContext
        audio_data: Raw audio bytes
        media_type: MIME type (e.g., 'audio/mpeg')

    Returns:
        Publicly accessible URL string, or None to use default BinaryContent behavior.
        Can be sync or async function.

    Note:
        The returned URL must be publicly accessible by the LLM provider.
        Empty strings are treated as None (fallback to default behavior).
    """


# =============================================================================
# Model Config
# =============================================================================


class ModelConfig(BaseModel):
    """Model configuration for context management.

    Extensibility:
        This class supports two extension patterns:

        1. Inheritance (recommended for type safety)::

            class MyModelConfig(ModelConfig):
                custom_threshold: float = 0.8

            class MyContext(AgentContext):
                model_cfg: MyModelConfig = Field(default_factory=MyModelConfig)

        2. Extra attributes (for quick prototyping)::

            config = ModelConfig(custom_field="value")
            config.custom_field  # Accessible but not type-checked
    """

    model_config = {"arbitrary_types_allowed": True, "extra": "allow"}

    context_window: int | None = None
    """Total context window size in tokens."""

    proactive_context_management_threshold: float | None = Field(default=0.65, ge=0.0, le=1.0)
    """Proactive context management threshold. When token usage exceeds this ratio, reminders are triggered."""

    compact_threshold: float = Field(default=0.90, ge=0.0, le=1.0)
    """Compact threshold for auto-compaction. When token usage exceeds this ratio, compact is triggered."""

    cold_start_trim_seconds: int = 3600
    """Seconds of inactivity after which KV cache is considered expired.

    When the gap between the last model response timestamp and now exceeds
    this value, all large tool return content is truncated before sending to
    the model, saving input token cost on cold re-encoding.
    Set to 0 to disable.
    """

    stream_resume_on_error: bool = False
    """Whether stream_agent retries failed stream attempts by default."""

    stream_resume_max_attempts: int = Field(default=3, ge=1)
    """Default maximum total attempts for stream_agent stream recovery."""

    stream_resume_prompt: str | None = None
    """Default resume prompt for stream_agent stream recovery.

    When unset, stream_agent uses its built-in resume prompt.
    """

    max_images: int = 20
    """Maximum number of images allowed in message history. Default is 20 (Claude's limit)."""

    max_videos: int = 1
    """Maximum number of videos allowed in message history. Default is 1."""

    support_gif: bool = True
    """Whether the model supports GIF images. If False, GIF images will be filtered out."""

    max_image_bytes: int = 5 * 1024 * 1024
    """Maximum allowed size in bytes for a single image as seen by the API.

    Most model APIs transmit images as base64-encoded payloads, which are
    roughly 4/3x larger than the raw bytes.  The compression filter
    automatically accounts for this overhead: it targets a raw-byte budget
    of ``max_image_bytes * 3 // 4`` so the encoded result stays within the
    API limit.

    When set to a positive value, binary images whose *encoded* size would
    exceed this limit are compressed (converted to JPEG with reduced quality
    and/or resized) before being sent to the model.  Set to 0 to disable
    compression.

    Default is 5 MB (5_242_880 bytes), matching the limit of most providers
    (e.g. Anthropic, GCP Vertex AI).
    """

    split_large_images: bool = True
    """Whether to split oversized binary images before sending to model."""

    image_split_max_height: int = 4096
    """Maximum height (in pixels) for each image segment when splitting large images."""

    image_split_overlap: int = 50
    """Overlap (in pixels) between adjacent image segments when splitting."""

    capabilities: set[ModelCapability] = Field(default_factory=set)
    """Set of capabilities supported by the model."""

    def has_capability(self, capability: ModelCapability) -> bool:
        """Check if the model has a specific capability."""
        return capability in self.capabilities

    @property
    def has_vision(self) -> bool:
        """Check if the model supports vision (image understanding)."""
        return ModelCapability.vision in self.capabilities

    @property
    def has_video_understanding(self) -> bool:
        """Check if the model supports video understanding."""
        return ModelCapability.video_understanding in self.capabilities

    @property
    def has_audio_understanding(self) -> bool:
        """Check if the model supports audio understanding."""
        return ModelCapability.audio_understanding in self.capabilities

    @property
    def has_document_understanding(self) -> bool:
        """Check if the model supports document understanding."""
        return ModelCapability.document_understanding in self.capabilities


# =============================================================================
# Resumable State
# =============================================================================


class ResumableState(BaseModel):
    """Resumable session state for AgentContext.

    This model captures the session state that can be serialized to JSON and
    restored later. It handles the special serialization requirements of
    ModelMessage using ModelMessagesTypeAdapter.

    The subagent_history is stored as serialized dict format (list[dict]) rather
    than ModelMessage objects, making the entire model JSON-serializable.

    Example:
        Saving state to JSON file::

            state = ctx.export_state()
            with open("session.json", "w") as f:
                f.write(state.model_dump_json(indent=2))

        Restoring state from JSON file::

            with open("session.json") as f:
                state = ResumableState.model_validate_json(f.read())
            new_ctx.restore_state(state)
    """

    subagent_history: dict[str, list[dict[str, Any]]] = Field(default_factory=dict)
    """Serialized subagent history, keyed by agent_id. Values are list[dict] from ModelMessagesTypeAdapter.dump_python()."""

    usage_snapshot_entries: dict[str, UsageSnapshotEntry] = Field(default_factory=dict)
    """Serialized usage ledger entries, keyed by agent/source ID."""

    user_prompts: str | Sequence[UserContent] | None = None
    """User prompts collected during the session."""

    steering_messages: list[str] = Field(default_factory=list)
    """Accumulated user steering messages for compact."""

    handoff_message: str | None = None
    """Rendered handoff message."""

    shell_env: dict[str, str] = Field(default_factory=dict)
    """Extra environment variables restored into AgentContext.shell_env."""

    deferred_tool_metadata: dict[str, dict[str, Any]] = Field(default_factory=dict)
    """Metadata for deferred tool calls."""

    agent_registry: dict[str, dict[str, Any]] = Field(default_factory=dict)
    """Serialized agent registry for tracking agent metadata."""

    need_user_approve_tools: list[str] = Field(default_factory=list)
    """List of tool names that require user approval before execution."""

    need_user_approve_mcps: list[str] = Field(default_factory=list)
    """List of MCP server names that require user approval for all tools."""

    security: SecurityConfig = Field(default_factory=SecurityConfig)
    """Security-related runtime configuration."""

    auto_load_files: list[str] = Field(default_factory=list)
    """Files to auto-load on next request. Set by handoff/compact, consumed by auto_load_files filter."""

    tasks: dict[str, dict[str, Any]] = Field(default_factory=dict)
    """Serialized tasks from TaskManager, keyed by task ID."""

    notes: dict[str, str] = Field(default_factory=dict)
    """Serialized note entries from NoteManager, keyed by entry key."""

    tool_search_loaded_tools: list[str] = Field(default_factory=list)
    """Tool names loaded via tool_search during the session."""

    tool_search_loaded_namespaces: list[str] = Field(default_factory=list)
    """Namespace IDs loaded via tool_search during the session."""

    def to_subagent_history(self) -> dict[str, list[ModelMessage]]:
        """Deserialize subagent_history to ModelMessage objects.

        Returns:
            Dict mapping agent_id to list of ModelMessage objects.
        """
        result: dict[str, list[ModelMessage]] = {}
        for key, messages_data in self.subagent_history.items():
            # subagent_history stores JSON-compatible ModelMessagesTypeAdapter.dump_python(..., mode="json") output.
            # validate_json applies val_json_bytes="base64" to nested BinaryContent inside ToolReturnPart.content.
            result[key] = ModelMessagesTypeAdapter.validate_json(json.dumps(messages_data).encode())
        return result

    def restore(self, ctx: AgentContext) -> None:
        """Restore state into an AgentContext.

        This method applies the saved state to the given context.
        Subclasses can override this method to restore additional fields.

        Args:
            ctx: The AgentContext to restore state into.

        Example::

            class MyState(ResumableState):
                custom_field: str = ""

                def restore(self, ctx: "MyContext") -> None:
                    super().restore(ctx)
                    ctx.custom_field = self.custom_field
        """
        ctx.subagent_history = self.to_subagent_history()
        ctx.usage_snapshot_entries = dict(self.usage_snapshot_entries)
        ctx.user_prompts = self.user_prompts
        ctx.steering_messages = list(self.steering_messages)
        ctx.handoff_message = self.handoff_message
        ctx.shell_env = {**ctx.shell_env, **self.shell_env}
        ctx.deferred_tool_metadata = dict(self.deferred_tool_metadata)
        # Restore agent_registry from serialized format
        ctx.agent_registry = {agent_id: AgentInfo(**info) for agent_id, info in self.agent_registry.items()}
        ctx.need_user_approve_tools = list(self.need_user_approve_tools)
        ctx.need_user_approve_mcps = list(self.need_user_approve_mcps)
        # Security is a runtime policy supplied by the current profile/config.
        # Keep the freshly constructed context security instead of restoring
        # stale or default policy from persisted conversation state.
        ctx.auto_load_files = list(self.auto_load_files)
        # Restore task_manager from serialized tasks (always reset to avoid stale state)
        ctx.task_manager = TaskManager.from_exported(self.tasks) if self.tasks else TaskManager()
        ctx.note_manager = NoteManager.from_exported(self.notes) if self.notes else NoteManager()
        ctx.tool_search_loaded_tools = list(self.tool_search_loaded_tools)
        ctx.tool_search_loaded_namespaces = list(self.tool_search_loaded_namespaces)


# =============================================================================
# Helper Functions
# =============================================================================


def _generate_run_id() -> str:
    return uuid4().hex


def _string_header_value(value: object) -> str | None:
    if isinstance(value, str) and value:
        return value
    return None


def _xml_to_string(element: Element) -> str:
    """Convert XML element to formatted string."""
    rough_string = tostring(element, encoding="unicode")
    dom = parseString(rough_string)  # noqa: S318
    # Get pretty-printed XML, skip the XML declaration line
    lines = dom.toprettyxml(indent="  ").split("\n")[1:]
    # Remove empty lines
    return "\n".join(line for line in lines if line.strip())


# =============================================================================
# Agent Context
# =============================================================================


class AgentContext(BaseModel):
    """Context for a single agent session.

    AgentContext manages session-specific state including:
    - Run identification (run_id, parent_run_id)
    - Timing (start_at, end_at, elapsed_time)
    - Deferred tool metadata
    - Handoff messages

    AgentContext holds a reference to an Environment and derives
    file_operator, shell, and resources from it via properties.

    Example:
        Using create_agent and stream_agent (recommended)::

            from ya_agent_sdk.agents.main import create_agent, stream_agent

            runtime = create_agent("openai-chat:gpt-4")
            # stream_agent manages runtime lifecycle automatically
            async with stream_agent(runtime, "Hello") as streamer:
                async for event in streamer:
                    print(event)

        Using create_agent with manual agent.run::

            runtime = create_agent("openai-chat:gpt-4")
            async with runtime:
                result = await runtime.agent.run("Hello", deps=runtime.ctx)

        Manual setup with Environment (advanced)::

            async with LocalEnvironment() as env:
                async with AgentContext(env=env) as ctx:
                    await ctx.file_operator.read_file("data.json")
    """

    model_config = {"arbitrary_types_allowed": True}

    lifecycle_extensions: list[LifecycleExtension] = Field(default_factory=list, exclude=True)
    """Runtime lifecycle extension objects attached to this context."""

    run_id: str = Field(default_factory=_generate_run_id)
    """Unique identifier for this session run."""

    parent_run_id: str | None = None
    """Parent run_id if this is a subagent context."""

    provider_session_id: str | None = None
    """Provider-facing session identifier for model request headers."""

    provider_thread_id: str | None = None
    """Provider-facing thread identifier for model request headers."""

    start_at: datetime | None = None
    """Timestamp when the context was entered."""

    end_at: datetime | None = None
    """Timestamp when the context was exited."""

    deferred_tool_metadata: dict[str, dict[str, Any]] = Field(default_factory=dict)
    """Metadata for deferred tool calls, keyed by tool_call_id."""

    handoff_message: str | None = None
    """Rendered handoff message to be injected into new context after handoff."""

    force_inject_instructions: bool = False
    """Force environment and runtime instruction injection on the next filter pass.

    Set to True by handoff/compact filters after context reset to ensure
    downstream filters inject instructions regardless of ToolReturnPart presence.
    Reset to False at the start of each filter pipeline run by the handoff filter.
    """

    env: Environment | None = None
    """Environment instance. file_operator/shell/resources are derived from it."""

    shell_env: dict[str, str] = Field(default_factory=dict)
    """Extra environment variables for shell command execution.

    These variables are merged into the shell environment for every command.
    Per-call environment variables (from tool invocation) take priority.
    Whether os.environ is included as the base layer depends on the Shell's
    include_os_env setting (see LocalShell).
    When empty (default), shell commands behave as usual (no extra env injected).
    """

    model_cfg: ModelConfig = Field(default_factory=ModelConfig)
    """Model configuration for context management."""

    tool_config: ToolConfig = Field(default_factory=ToolConfig)
    """Tool-level configuration for API keys and tool-specific settings."""

    usage_snapshot_entries: dict[str, UsageSnapshotEntry] = Field(default_factory=dict)
    """Unified per-run usage ledger, keyed by agent/source ID."""

    shell_review_records: deque[Any] = Field(default_factory=lambda: deque(maxlen=10), exclude=True)
    """Recent shell review records for current-run safety context."""

    user_prompts: str | Sequence[UserContent] | None = None
    """User prompts collected during the session for compact."""

    steering_messages: list[str] = Field(default_factory=list)
    """Accumulated user steering messages received via message bus.

    When messages with source="user" are injected via inject_bus_messages filter,
    they are also appended here. This ensures steering messages are preserved
    for compact even if the agent is interrupted and resumed.

    Cleared after compact to avoid duplication (steering content is summarized).
    """

    tool_id_wrapper: ToolIdWrapper = Field(default_factory=ToolIdWrapper)
    """Tool ID wrapper for normalizing tool call IDs across providers."""

    agent_stream_queues: dict[str, asyncio.Queue[AgentStreamEvent]] = Field(
        default_factory=_create_stream_queue_factory
    )
    """Stream queues for agent events, keyed by agent_id (e.g., "main", "explorer-a7b9").

    Each queue receives AgentStreamEvent instances during agent execution,
    enabling real-time streaming of agent responses via a sideband channel.
    The key matches agent_registry keys for consistent lookup in poll_subagents.
    """

    need_user_approve_tools: list[str] = Field(default_factory=list)
    """List of tool names that require user approval before execution.

    Tools in this list will trigger HITL (Human-in-the-Loop) flow,
    deferring execution until the user explicitly approves.
    """

    need_user_approve_mcps: list[str] = Field(default_factory=list)
    """List of MCP server names that require user approval for all tools.

    When a server name is in this list, all tools from that server will
    trigger the HITL approval flow before execution. The server name
    corresponds to the tool_prefix used when creating the MCP toolset.

    Example:
        ctx.need_user_approve_mcps = ["filesystem", "github"]
        # All tools from these servers will require approval
    """

    security: SecurityConfig = Field(default_factory=SecurityConfig)
    """Security-related runtime configuration."""

    subagent_history: dict[str, list[ModelMessage]] = Field(default_factory=dict)
    """Subagent history for resuming sessions."""

    agent_registry: dict[str, AgentInfo] = Field(default_factory=dict)
    """Registry of agent metadata, keyed by agent_id.

    Used by stream_agent to track agent identity and hierarchy.
    Populated by enter_subagent when subagents are created.
    """

    auto_load_files: list[str] = Field(default_factory=list)
    """Files to auto-load on next request. Set by handoff/compact tool, consumed by auto_load_files filter."""

    task_manager: TaskManager = Field(default_factory=TaskManager)
    """Task manager for tracking tasks and dependencies within the session."""

    note_manager: NoteManager = Field(default_factory=NoteManager)
    """Note manager for persistent key-value storage within the session."""

    tool_search_loaded_tools: list[str] = Field(default_factory=list)
    """Tool names loaded via tool_search during the session. Used for session restore."""

    tool_search_loaded_namespaces: list[str] = Field(default_factory=list)
    """Namespace IDs loaded via tool_search during the session. Used for session restore."""

    injected_context_tags: tuple[str, ...] = (RUNTIME_CONTEXT_TAG, ENVIRONMENT_CONTEXT_TAG)
    """XML tag names for per-turn injected context blocks.

    These tags identify context blocks that are re-injected fresh on every
    request by filters (runtime context, environment context, etc.).
    Historical copies are stripped during compact to save tokens.

    The SDK defaults include runtime-context and environment-context.
    Application layers can extend this tuple to register additional tags
    (e.g., project-guidance, user-rules) so that compact strips them too.

    Example::

        ctx.injected_context_tags = (
            *ctx.injected_context_tags,
            PROJECT_GUIDANCE_TAG,
            USER_RULES_TAG,
        )
    """

    context_manage_tool_names: list[str] = Field(default_factory=list)
    """Names of active context management tools (e.g., ['summarize']).

    Automatically populated by create_agent() from tools with
    is_context_manage_tool=True. Enables proactive context management
    reminders when context window usage exceeds threshold.
    """

    tool_tags: set[str] = Field(default_factory=set)
    """Active capability tags from available tools.

    Recomputed by Toolset.get_tools() on each call. Tags are collected from
    all available tools' `tags` attribute and set here (replacing previous value).
    Tools with `superseded_by_tags` matching any active tag are hidden.

    Cleared on subagent context creation to ensure fresh computation.

    This field is useful for external inspection to determine
    what tool capabilities are active in the current session.

    Example::

        if "shell" in ctx.tool_tags:
            # Shell tool is available, can use shell commands
            ...
    """

    message_bus: MessageBus = Field(default_factory=MessageBus)
    """Message bus for inter-agent communication.

    Enables real-time communication between user and agents, or between agents.
    The inject_bus_messages filter consumes and injects pending messages
    into the conversation.

    The message_bus is shared between parent and child contexts to allow
    bidirectional communication during subagent execution.

    Example:
        ctx.send_message(BusMessage(content="Please stop", source="user"))
    """

    model_wrapper: ModelWrapper | None = Field(default=None, exclude=True)
    """Optional wrapper applied to all models created within this context.

    When set, the SDK will call this wrapper for:
    - Main agent model
    - Subagent models
    - Video/image understanding models
    - Compact filter models

    The wrapper receives the model, agent name, and context dict from
    get_wrapper_metadata(), allowing customized instrumentation per agent type.

    Note:
        This field is excluded from serialization (not resumable).
        Wrappers must be re-configured when restoring sessions.

    Example::

        def my_wrapper(model: Model, agent_name: str, context: dict[str, Any]) -> Model:
            return traced_model(model, name=agent_name, trace_id=context.get("run_id"))

        runtime = create_agent(
            "openai-chat:gpt-4",
            model_wrapper=my_wrapper,
        )
    """

    subagent_wrapper: SubagentWrapper | None = Field(default=None, exclude=True)
    """Optional wrapper applied around subagent execution.

    When set, the SDK will call this wrapper as an async context manager
    around the entire subagent run. This enables consumers to create
    observability spans (e.g., Langfuse agent observations) that properly
    nest all subagent activity underneath.

    The wrapper receives agent_name, agent_id, and a metadata dict
    from get_wrapper_metadata().

    Note:
        This field is excluded from serialization (not resumable).
        Wrappers must be re-configured when restoring sessions.

    Example::

        from contextlib import asynccontextmanager

        @asynccontextmanager
        async def trace_subagent(name: str, id: str, meta: dict[str, Any]):
            with langfuse.start_as_current_observation(
                name=f"agent:{name}", as_type="agent",
            ):
                yield

        runtime = create_agent(
            "openai-chat:gpt-4",
            subagent_wrapper=trace_subagent,
        )
    """

    self_fork_agent: Any | None = Field(default=None, exclude=True)
    """Runtime-only agent used by delegate(subagent_name="self") forks.

    The self fork agent is created by create_agent() with the same model,
    system prompt, capabilities, and ordinary tools as the parent agent. It is
    excluded from serialization and restored per runtime construction.
    """

    wrapper_metadata: dict[str, Any] = Field(default_factory=dict, exclude=True)
    """Additional context passed to model_wrapper via get_wrapper_metadata().

    These fields are merged with built-in fields (run_id) when calling the wrapper.
    User-defined fields take precedence over built-in fields.

    Note:
        This field is excluded from serialization (not resumable).
        Context must be re-configured when restoring sessions.

    Example::

        runtime = create_agent(
            "openai-chat:gpt-4",
            model_wrapper=my_wrapper,
            extra_context_kwargs={
                "wrapper_metadata": {
                    "trace_id": "abc123",
                    "user_id": "user_456",
                    "tags": ["production"],
                },
            },
        )

        # Runtime modification
        ctx.wrapper_metadata["request_id"] = current_request.id
    """

    _agent_id: str = "main"
    _entered: bool = False
    _enter_lock: asyncio.Lock = None  # type: ignore[assignment]  # Initialized in __init__
    _subagent_state_lock: asyncio.Lock = None  # type: ignore[assignment]  # Initialized in __init__
    _stream_queue_enabled: bool = False
    _compact_depth: int = 0

    def __init__(self, **data: object) -> None:
        """Initialize AgentContext."""
        super().__init__(**data)
        object.__setattr__(self, "_enter_lock", asyncio.Lock())
        object.__setattr__(self, "_subagent_state_lock", asyncio.Lock())

    @property
    def agent_id(self) -> str:
        """Unique identifier for this agent context.

        Returns "main" for the main agent, or a generated short ID for subagents.
        """
        return self._agent_id

    @property
    def file_operator(self) -> FileOperator | None:
        """File operator for file system operations. Derived from env."""
        if self.env is not None:
            return self.env.file_operator
        return None

    @property
    def shell(self) -> Shell | None:
        """Shell executor for command execution. Derived from env."""
        if self.env is not None:
            return self.env.shell
        return None

    @property
    def resources(self) -> ResourceRegistry | None:
        """Resource registry for runtime resources. Derived from env."""
        if self.env is not None:
            return self.env.resources
        return None

    @property
    def elapsed_time(self) -> timedelta | None:
        """Return elapsed time since start, or None if not started.

        If session has ended, returns the final duration.
        If session is running, returns the current elapsed time.
        """
        if self.start_at is None:
            return None
        end = self.end_at if self.end_at else datetime.now(UTC)
        return end - self.start_at

    def get_model_extra_headers(self) -> dict[str, str]:
        """Return stable provider headers for model construction.

        OAuth-backed Codex models use both underscore and hyphen header variants,
        matching OpenAI Codex session header behavior.
        """
        session_id = _string_header_value(getattr(self, "provider_session_id", None)) or self.run_id
        thread_id = _string_header_value(getattr(self, "provider_thread_id", None)) or self.run_id
        return {
            "session_id": session_id,
            "session-id": session_id,
            "thread_id": thread_id,
            "thread-id": thread_id,
            "x-client-request-id": thread_id,
        }

    def get_wrapper_metadata(self) -> dict[str, Any]:
        """Return context dict for model wrapper.

        This method provides the context dictionary passed to model_wrapper
        as its third argument. Default implementation returns built-in fields
        merged with the wrapper_metadata field, with user-defined fields taking precedence.

        Built-in fields:
            - run_id: Current session identifier
            - agent_id: Current agent identifier ("main" or "debugger-a7b9")
            - parent_run_id: Parent session ID (None for main, set for subagents)

        Override this method for dynamic context generation based on runtime state.

        Returns:
            Context dict passed as third argument to model_wrapper.

        Example::

            # Simple: set wrapper_metadata field
            ctx.wrapper_metadata = {"trace_id": "abc", "user_id": "123"}

            # Advanced: override for dynamic context
            class MyContext(AgentContext):
                session_metadata: dict = Field(default_factory=dict)

                def get_wrapper_metadata(self) -> dict[str, Any]:
                    return {
                        **super().get_wrapper_metadata(),
                        "timestamp": datetime.now(UTC).isoformat(),
                        "session": self.session_metadata,
                    }
        """
        return {
            "run_id": self.run_id,
            "agent_id": self._agent_id,
            "parent_run_id": self.parent_run_id,
            **self.wrapper_metadata,
        }

    def get_current_time(self) -> datetime:
        """Return current time with timezone information.

        Override this method to provide custom time sources (e.g., NTP, mock for testing).
        Default implementation uses UTC with timezone information.

        Returns:
            Current datetime with timezone information (ISO 8601 compatible).

        Example:
            Subclass to customize time source::

                class MyContext(AgentContext):
                    def get_current_time(self) -> datetime:
                        return ntp_client.get_time()  # Custom NTP source

            Mock for testing::

                class MockContext(AgentContext):
                    def get_current_time(self) -> datetime:
                        return datetime(2025, 1, 15, 12, 0, 0, tzinfo=timezone.utc)
        """
        return datetime.now(UTC)

    def _build_active_tasks_element(self, parent: Element, detailed: bool) -> None:
        """Build active-tasks XML element and append to parent if tasks exist.

        Args:
            parent: Parent XML element to append to.
            detailed: If True, use detailed format with hints and active_form;
                otherwise use compact format.
        """
        active_tasks = [t for t in self.task_manager.list_all() if t.status != TaskStatus.COMPLETED]
        if not active_tasks:
            return

        tasks_elem = SubElement(parent, "active-tasks")
        if detailed:
            tasks_elem.set("hint", "Update status with task_update tool")

        for task in active_tasks:
            task_elem = SubElement(tasks_elem, "task")
            task_elem.set("id", task.id)
            task_elem.set("status", task.status.value)

            # Only show active (incomplete) blockers
            active_blockers = [
                bid
                for bid in task.blocked_by
                if (blocker := self.task_manager.get(bid)) and blocker.status != TaskStatus.COMPLETED
            ]
            if active_blockers:
                task_elem.set("blocked-by", ",".join(active_blockers))

            if detailed:
                # Detailed format: subject as sub-element, include active_form
                SubElement(task_elem, "subject").text = task.subject
                if task.active_form and task.status == TaskStatus.IN_PROGRESS:
                    SubElement(task_elem, "active-form").text = task.active_form
            else:
                # Compact format: subject as text content
                task_elem.text = task.subject

    async def get_context_instructions(
        self,
        run_context: RunContext[AgentContext] | None = None,
        *,
        is_user_prompt: bool = True,
    ) -> str:
        """Return runtime context instructions in XML format.

        Provides runtime information about the current session.

        Args:
            run_context: Optional RunContext for accessing message history and metadata.
            is_user_prompt: Whether this is a user prompt (True) or a tool response/retry (False).
                Some information (e.g., known subagents) is only included on user prompts
                to reduce noise. Subclasses can override to customize behavior. Default True.

        Returns:
            XML-formatted string with runtime context and optional system reminders.
        """
        parts: list[str] = []

        # Build runtime-context element
        root = Element(RUNTIME_CONTEXT_TAG)

        # Agent ID (useful for subagents to know their identity)
        SubElement(root, "agent-id").text = self._agent_id

        # Current time (ISO 8601 with timezone offset for cross-environment compatibility)
        SubElement(root, "current-time").text = self.get_current_time().isoformat(timespec="seconds")

        # Elapsed time
        elapsed = self.elapsed_time
        elapsed_str = f"{elapsed.total_seconds():.1f}s" if elapsed else "not started"
        SubElement(root, "elapsed-time").text = elapsed_str

        # Model configuration - only context_window
        if self.model_cfg.context_window is not None:
            config = SubElement(root, "model-config")
            SubElement(config, "context-window").text = str(self.model_cfg.context_window)

        # Token usage from runtime info
        if (
            run_context
            and (request_usage := get_latest_request_usage(run_context.messages))
            and request_usage.total_tokens is not None
        ):
            usage_elem = SubElement(root, "token-usage")
            SubElement(usage_elem, "total-tokens").text = str(request_usage.total_tokens)

        # Known subagents from agent_registry (excluding main agent)
        # Only include on user prompts, not tool responses
        if is_user_prompt:
            known_subagents = {
                agent_id: info for agent_id, info in self.agent_registry.items() if agent_id != self._agent_id
            }
            if known_subagents:
                subagents_elem = SubElement(root, "known-subagents")
                subagents_elem.set("hint", "Use subagent_info tool for more details")
                for _agent_id, info in known_subagents.items():
                    agent_elem = SubElement(subagents_elem, "agent")
                    agent_elem.set("id", info.agent_id)
                    agent_elem.set("name", info.agent_name)

        # Active tasks (pending + in_progress)
        self._build_active_tasks_element(root, detailed=is_user_prompt)

        # Note entries (only on user prompts to reduce noise)
        self._build_notes_element(root, is_user_prompt=is_user_prompt)

        parts.append(_xml_to_string(root))

        # Background process summary (only on user prompts)
        if is_user_prompt and self.shell is not None:
            bg_summary = self.shell.background_status_summary()
            if bg_summary:
                parts.append(bg_summary)

        # System reminders (sibling to runtime-context)
        self._build_system_reminders(parts, run_context)

        return "\n\n".join(parts)

    def _build_system_reminders(
        self,
        parts: list[str],
        run_context: RunContext[AgentContext] | None,
    ) -> None:
        """Build system-reminder XML and append to parts if needed."""
        if not (
            self.context_manage_tool_names
            and self.model_cfg.context_window is not None
            and self.model_cfg.proactive_context_management_threshold is not None
            and run_context
            and (request_usage := get_latest_request_usage(run_context.messages))
            and request_usage.total_tokens is not None
        ):
            return

        threshold_tokens = int(self.model_cfg.context_window * self.model_cfg.proactive_context_management_threshold)
        if request_usage.total_tokens < threshold_tokens:
            return

        usage_pct = int(request_usage.total_tokens / self.model_cfg.context_window * 100)
        compact_pct = int(self.model_cfg.compact_threshold * 100)
        tool_names = ", ".join(f"`{n}`" for n in self.context_manage_tool_names)
        note_guidance = ""
        if self.note_manager.entries:
            note_guidance = " Review note keys, read needed values with `note_get`, and delete stale or oversized notes before summarizing."
        reminder_text = (
            f"Context usage is at {usage_pct}% ({request_usage.total_tokens:,} / {self.model_cfg.context_window:,} tokens). "
            f"Auto-compaction will trigger at {compact_pct}%. "
            f"Please use the {tool_names} tool to summarize progress and continue when appropriate."
            f"{note_guidance}"
        )

        reminder_root = Element("system-reminder")
        item = SubElement(reminder_root, "item")
        item.text = reminder_text
        parts.append(_xml_to_string(reminder_root))

    def _build_notes_element(self, parent: Element, *, is_user_prompt: bool = True) -> None:
        """Build notes XML element and append to parent if entries exist.

        Args:
            parent: Parent XML element to append to.
            is_user_prompt: Only include note keys on user prompts to reduce noise.
        """
        if not is_user_prompt:
            return

        keys = self.note_manager.list_keys()
        if not keys:
            return

        notes_elem = SubElement(parent, "notes")
        notes_elem.set("hint", "Note values are omitted from runtime context. Use note_get to read values by key.")
        for key in keys:
            entry_elem = SubElement(notes_elem, "entry")
            entry_elem.set("key", key)

    def prepare_new_run(self) -> Self:
        """Create a fresh context copy for a new agent run.

        Returns a shallow copy with per-run state reset while sharing
        long-lived state (env, message_bus, task_manager, etc.) with the
        original context.

        This method is used by stream_agent() to create an isolated context
        for each run, preventing stale per-run state from leaking between
        consecutive runs. This is critical for avoiding ContextVar-related
        errors (pydantic-ai issue #674) where async generator cleanup from a
        previous cancelled run interferes with the current run.

        Per-run state (fresh for each run):
            - run_id, start_at, end_at
            - tool_id_wrapper, agent_stream_queues
            - usage_snapshot_entries, shell_review_records, deferred_tool_metadata
            - force_inject_instructions

        Shared state (same reference as original):
            - env, model_cfg, tool_config
            - message_bus, task_manager, note_manager
            - subagent_history, agent_registry
            - model_wrapper, subagent_wrapper
            - need_user_approve_tools, need_user_approve_mcps
            - All other fields not explicitly reset

        Returns:
            A new context instance ready for a fresh agent run.
        """
        update: dict[str, Any] = {
            "run_id": _generate_run_id(),
            "start_at": datetime.now(UTC),
            "end_at": None,
            "tool_id_wrapper": ToolIdWrapper(),
            "agent_stream_queues": _create_stream_queue_factory(),
            "usage_snapshot_entries": {},
            "shell_review_records": deque(maxlen=10),
            "deferred_tool_metadata": {},
            "force_inject_instructions": False,
        }
        new_ctx = self.model_copy(update=update)
        # Reset private state for fresh lifecycle
        object.__setattr__(new_ctx, "_entered", False)
        object.__setattr__(new_ctx, "_enter_lock", asyncio.Lock())
        object.__setattr__(new_ctx, "_subagent_state_lock", self._subagent_state_lock)
        object.__setattr__(new_ctx, "_stream_queue_enabled", False)
        object.__setattr__(new_ctx, "_compact_depth", 0)
        return new_ctx

    def reserve_subagent_id(self, agent_name: str, agent_id: str | None = None) -> str:
        """Reserve and register a subagent ID synchronously.

        The caller must hold ``_subagent_state_lock`` while calling this method.
        """
        effective_agent_id = agent_id or _generate_run_id()
        if effective_agent_id not in self.agent_registry:
            self.agent_registry[effective_agent_id] = AgentInfo(
                agent_id=effective_agent_id,
                agent_name=agent_name,
                parent_agent_id=self._agent_id,
            )
        return effective_agent_id

    def create_subagent_context(
        self,
        agent_name: str,
        agent_id: str | None = None,
        **override: object,
    ) -> Self:
        """Create a child context for subagent with independent timing.

        The subagent context inherits all fields but gets:
        - A new run_id (uses agent_id if provided, otherwise generates one)
        - parent_run_id set to current run_id
        - Fresh timing (start_at/end_at managed by __aenter__/__aexit__)
        - Shared file_operator and shell from parent
        - Registers agent info in parent's agent_registry

        Note:
            The returned context should be used with `async with` to properly
            manage timing and resource cleanup::

                sub_ctx = parent.create_subagent_context("search")
                async with sub_ctx:
                    # subagent work here

        Args:
            agent_name: Name of the subagent (e.g., "search", "reasoning").
            agent_id: ID for the subagent. If None, generates a unique ID.
                Can be tool_call_id for correlation with tool calls.
            **override: Additional fields to override in the subagent context.
                Subclasses can pass extra fields without overriding this method.

        Returns:
            A new context instance configured for the subagent.
        """
        effective_agent_id = self.reserve_subagent_id(agent_name, agent_id)

        update: dict[str, Any] = {
            "run_id": _generate_run_id(),
            "parent_run_id": self.run_id,
            "start_at": None,  # Will be set by __aenter__
            "end_at": None,  # Will be set by __aexit__
            "handoff_message": None,  # Subagents don't inherit handoff state
            "user_prompts": None,  # Subagent has its own initial prompt (set by caller)
            "steering_messages": [],  # Subagent has its own steering queue
            "tool_id_wrapper": ToolIdWrapper(),  # Fresh wrapper for subagent
            "tool_tags": set(),  # Fresh tags for subagent (recomputed by its own Toolset)
            "shell_review_records": deque(maxlen=10),
            "security": self.security.model_copy(deep=True),
            # env is inherited via model_copy (shares parent's env reference)
            **override,
        }
        new_ctx = self.model_copy(update=update)
        new_ctx._agent_id = effective_agent_id
        # Reset re-entry protection for subagent (independent lifecycle)
        object.__setattr__(new_ctx, "_entered", False)
        object.__setattr__(new_ctx, "_enter_lock", asyncio.Lock())
        object.__setattr__(new_ctx, "_subagent_state_lock", self._subagent_state_lock)
        object.__setattr__(new_ctx, "_compact_depth", 0)
        return new_ctx

    def _build_history_processors(self) -> list[Callable[..., Any]]:
        """Build the ordered context history processor callables."""
        # Import filters here to avoid circular imports
        from ya_agent_sdk.filters.auto_load_files import process_auto_load_files
        from ya_agent_sdk.filters.background_shell import inject_background_results
        from ya_agent_sdk.filters.bus_message import inject_bus_messages
        from ya_agent_sdk.filters.capability import filter_by_capability
        from ya_agent_sdk.filters.handoff import process_handoff_message
        from ya_agent_sdk.filters.image import (
            compress_large_images,
            drop_extra_images,
            drop_extra_videos,
            drop_gif_images,
            split_large_images,
        )
        from ya_agent_sdk.filters.reasoning_normalize import normalize_reasoning_for_model
        from ya_agent_sdk.filters.runtime_instructions import inject_runtime_instructions
        from ya_agent_sdk.filters.tool_args import fix_truncated_tool_args

        def dynamic_tool_id_wrapper(ctx: RunContext[AgentContext], messages: list[ModelMessage]) -> list[ModelMessage]:
            """Dynamically get tool_id_wrapper from current context."""
            return ctx.deps.tool_id_wrapper.wrap_messages(ctx, messages)

        return [
            normalize_reasoning_for_model,
            split_large_images,
            compress_large_images,
            drop_extra_images,
            drop_gif_images,
            drop_extra_videos,
            fix_truncated_tool_args,
            process_handoff_message,
            process_auto_load_files,
            filter_by_capability,
            inject_bus_messages,  # Before runtime_instructions, after handoff/auto_load
            inject_background_results,  # Consume completed bg results, before runtime_instructions
            inject_runtime_instructions,
            dynamic_tool_id_wrapper,
        ]

    def get_history_capabilities(self) -> list[AbstractCapability[AgentContext]]:
        """Return ProcessHistory capabilities for this context.

        Example::

            async with AgentContext(...) as ctx:
                agent = Agent(
                    'openai-chat:gpt-4',
                    deps_type=AgentContext,
                    capabilities=ctx.get_history_capabilities(),
                )
        """
        return [ProcessHistory(processor) for processor in self._build_history_processors()]

    def update_usage_snapshot_entry(
        self,
        *,
        agent_id: str,
        agent_name: str,
        model_id: str,
        usage: RunUsage,
        usage_id: str | None = None,
        source: str = "model_request",
        ledger_key: str | None = None,
    ) -> UsageSnapshot:
        """Update one cumulative usage ledger entry and return the latest run snapshot."""
        self.usage_snapshot_entries[ledger_key or agent_id] = UsageSnapshotEntry(
            agent_id=agent_id,
            agent_name=agent_name,
            model_id=model_id,
            usage=coerce_run_usage(usage),
            usage_id=usage_id,
            source=source,
        )
        return self.build_usage_snapshot()

    def build_usage_snapshot(self) -> UsageSnapshot:
        """Build a cumulative usage snapshot for this run."""
        total_usage = RunUsage()
        agent_usages: dict[str, UsageAgentTotal] = {}
        model_usages: dict[str, RunUsage] = {}
        entries = sorted(self.usage_snapshot_entries.values(), key=lambda entry: entry.agent_id)
        for entry in entries:
            total_usage += entry.usage
            current_agent_usage = agent_usages.get(entry.agent_id)
            if current_agent_usage is None:
                agent_usages[entry.agent_id] = UsageAgentTotal(
                    agent_name=entry.agent_name,
                    model_id=entry.model_id,
                    usage=entry.usage,
                    usage_id=entry.usage_id,
                    source=entry.source,
                )
            else:
                current_agent_usage.usage = current_agent_usage.usage + entry.usage
                if current_agent_usage.model_id != entry.model_id:
                    current_agent_usage.model_id = "multiple"
            current_model_usage = model_usages.get(entry.model_id, RunUsage())
            model_usages[entry.model_id] = current_model_usage + entry.usage
        return UsageSnapshot(
            run_id=self.parent_run_id or self.run_id,
            total_usage=total_usage,
            entries=list(entries),
            agent_usages=agent_usages,
            model_usages=model_usages,
        )

    async def emit_usage_snapshot_event(self, *, source: str = "usage_snapshot") -> None:
        """Emit the current cumulative usage snapshot without mutating the ledger."""
        if not self.usage_snapshot_entries:
            return

        snapshot = self.build_usage_snapshot()
        await self.emit_event(
            UsageSnapshotEvent(
                event_id=f"{snapshot.run_id}:usage_snapshot:{source}",
                snapshot=snapshot,
                source=source,
            )
        )

    async def emit_usage_snapshot(
        self,
        *,
        agent_id: str,
        agent_name: str,
        model_id: str,
        usage: RunUsage,
        source: str = "model_request",
        usage_id: str | None = None,
        ledger_key: str | None = None,
    ) -> None:
        """Update cumulative usage and emit a usage snapshot event.

        Deprecated: use update_usage_snapshot_entry() for ledger updates and
        emit_usage_snapshot_event() at stream/session boundaries.
        """
        self.update_usage_snapshot_entry(
            agent_id=agent_id,
            agent_name=agent_name,
            model_id=model_id,
            usage=usage,
            usage_id=usage_id,
            source=source,
            ledger_key=ledger_key,
        )
        await self.emit_usage_snapshot_event(source=source)

    async def emit_event(self, event: AgentStreamEvent) -> None:
        """Emit a custom event to the sideband stream queue.

        Events are placed in the agent_stream_queues under the current run_id,
        allowing consumers to receive custom notifications alongside pydantic-ai
        stream events.

        This method is a no-op if streaming is not enabled. Streaming is
        automatically enabled when using stream_agent().

        Args:
            event: Arbitrary event object (AgentEvent subclass, pydantic-ai events, or custom).

        Example::

            from ya_agent_sdk.events import CompactStartEvent

            await ctx.emit_event(CompactStartEvent(event_id="abc123", message_count=50))
        """
        if not self._stream_queue_enabled:
            return
        await self.agent_stream_queues[self._agent_id].put(event)

    def send_message(self, message: BusMessage) -> BusMessage:
        """Send a message to the message bus.

        This operation is idempotent: if a message with the same id already
        exists, the existing message is returned.

        Args:
            message: The message to send.

        Returns:
            The sent message, or existing message if id already exists.

        Example::

            from ya_agent_sdk.context import BusMessage

            # User sends message to main agent
            ctx.send_message(BusMessage(content="Please stop", source="user"))

            # Targeted message to subagent
            ctx.send_message(BusMessage(
                content="Focus on security",
                source="main",
                target=subagent_id,
            ))

            # With explicit id for idempotent send
            ctx.send_message(BusMessage(
                id="unique-id-123",
                content="Important",
                source="user",
            ))
        """
        return self.message_bus.send(message)

    def consume_messages(self) -> list[BusMessage]:
        """Consume pending messages for this agent (idempotent).

        Returns messages that:
        1. Are pending for this agent (targeted or broadcast)
        2. Have not been consumed before (tracked by message ID)

        This method is idempotent: each message is returned only once,
        even if called multiple times. This ensures messages are not
        duplicated on LLM retries.

        Returns:
            List of new messages, in FIFO order (oldest first).

        Example::

            # First call returns pending messages
            messages = ctx.consume_messages()
            # [BusMessage(id="abc", content="Hello"), ...]

            # Second call returns empty (already consumed)
            messages = ctx.consume_messages()
            # []
        """
        return self.message_bus.consume(self._agent_id)

    async def __aenter__(self):
        """Enter the context and start timing.

        Raises:
            RuntimeError: If the context has already been entered.
        """
        async with self._enter_lock:
            if self._entered:
                raise RuntimeError(
                    "AgentContext has already been entered. "
                    "Each AgentContext instance can only be entered once at a time."
                )
            # Must use object.__setattr__ because Pydantic v2 silently ignores
            # normal assignment on private attrs that were previously set via
            # object.__setattr__ (e.g., in prepare_new_run / create_subagent_context).
            object.__setattr__(self, "_entered", True)
        self.start_at = datetime.now(UTC)
        self.tool_id_wrapper.clear()
        # Subscribe to message bus for this agent
        self.message_bus.subscribe(self._agent_id)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Exit the context and record end time."""
        self.end_at = datetime.now(UTC)
        # For main agent: clear message bus on exit
        # For subagents: unsubscribe to free cursor
        if self._agent_id == "main":
            self.message_bus.clear()
        else:
            self.message_bus.unsubscribe(self._agent_id)
        async with self._enter_lock:
            object.__setattr__(self, "_entered", False)

    def export_state(
        self,
        *,
        include_subagent: bool = True,
        include_usage_ledger: bool = False,
        include_extra_usages: bool | None = None,
    ) -> ResumableState:
        """Export resumable session state.

        Creates a ResumableState containing all session data that can be
        serialized to JSON and restored later.

        Args:
            include_subagent: Whether to include subagent history and registry.
                Defaults to True. Set to False to exclude subagent data,
                which can reduce state size for main agent-only persistence.
            include_usage_ledger: Whether to include usage ledger entries in the state.
                Defaults to False. Usage ledger data is per-run billing data.
                Set to True for crash recovery scenarios where interrupted-run
                usage should be preserved.
            include_extra_usages: Backward-compatible alias for include_usage_ledger.

        Returns:
            ResumableState instance ready for serialization.

        Example::

            # Save full state including subagent history (default, no usage ledger records)
            state = ctx.export_state()

            # Save state without subagent data
            state = ctx.export_state(include_subagent=False)

            # Save state with usage ledger records for crash recovery
            state = ctx.export_state(include_usage_ledger=True)

            with open("session.json", "w") as f:
                f.write(state.model_dump_json(indent=2))
        """
        if include_extra_usages is not None:
            include_usage_ledger = include_extra_usages

        serialized_history: dict[str, list[dict[str, Any]]] = {}
        serialized_registry: dict[str, dict[str, Any]] = {}

        if include_subagent:
            # Serialize subagent_history using ModelMessagesTypeAdapter
            # Use mode='json' to ensure bytes (e.g., BinaryContent.data) are base64-encoded
            for key, messages in self.subagent_history.items():
                serialized_history[key] = ModelMessagesTypeAdapter.dump_python(messages, mode="json")

            # Serialize agent_registry to dict format
            serialized_registry = {
                agent_id: {
                    "agent_id": info.agent_id,
                    "agent_name": info.agent_name,
                    "parent_agent_id": info.parent_agent_id,
                }
                for agent_id, info in self.agent_registry.items()
            }

        return ResumableState(
            subagent_history=serialized_history,
            usage_snapshot_entries=dict(self.usage_snapshot_entries) if include_usage_ledger else {},
            user_prompts=self.user_prompts,
            steering_messages=list(self.steering_messages),
            handoff_message=self.handoff_message,
            shell_env=dict(self.shell_env),
            deferred_tool_metadata=dict(self.deferred_tool_metadata),
            agent_registry=serialized_registry,
            need_user_approve_tools=list(self.need_user_approve_tools),
            need_user_approve_mcps=list(self.need_user_approve_mcps),
            security=self.security.model_copy(deep=True),
            auto_load_files=list(self.auto_load_files),
            tasks=self.task_manager.export_tasks(),
            notes=self.note_manager.export_notes(),
            tool_search_loaded_tools=list(self.tool_search_loaded_tools),
            tool_search_loaded_namespaces=list(self.tool_search_loaded_namespaces),
        )

    def with_state(self, state: ResumableState | None) -> Self:
        """Restore session state from a ResumableState.

        Updates the context with state from a previously exported ResumableState.
        This allows resuming a session after serialization/deserialization.

        If state is None, returns self unchanged for convenient chaining.

        Args:
            state: ResumableState to restore from, or None to skip restoration.

        Returns:
            Self for method chaining.

        Example::

            # Load from JSON file and use with async context manager
            with open("session.json") as f:
                state = ResumableState.model_validate_json(f.read())
            async with AgentContext(...).with_state(state) as ctx:
                ...

            # Also works with None for conditional restoration
            async with AgentContext(...).with_state(maybe_state) as ctx:
                ...
        """
        if state is None:
            return self
        state.restore(self)
        return self


# =============================================================================
# FORWARD REFERENCE RESOLUTION
# =============================================================================

# Resolve forward references now that all classes are defined.
# This is required because MediaToUrlHook references "AgentContext" which
# is defined after ToolConfig. Without this, subclassing ToolConfig fails
# with PydanticUserError.
ToolConfig.model_rebuild()
