"""Agent — the core building block of Logos ADK."""

from __future__ import annotations

import asyncio
from collections import OrderedDict
from dataclasses import asdict, is_dataclass
from datetime import datetime, UTC
from fnmatch import fnmatch
import hashlib
import json
import time
from typing import Any, AsyncIterator
from uuid import uuid4

from logos_adk.config import AgentConfig, ReliabilityConfig, SecurityConfig
from logos_adk.config_loader import load_agent_config, load_agent_from_config, resolve_agent_class
from logos_adk.context import RunContext, get_run_context, reset_run_context, set_run_context
from logos_adk.cost import DEFAULT_RATES
from logos_adk.errors import (
    ConfigError,
    ProviderError,
    ProviderTimeoutError,
)
from logos_adk.memory import Memory
from logos_adk.observe import Exporter
from logos_adk.providers import Provider, provider_registry
from logos_adk.reliability import CircuitBreaker, ProviderHealth, RetryPolicy, ensure_circuit_allows
from logos_adk.runnable import Runnable
from logos_adk.spend_ledger import (
    SpendLedger,
)
from logos_adk.security import (
    AuditLogger,
    AuditRetentionPolicy,
    InputValidator,
    PIIDetector,
    PromptShield,
)
from logos_adk.state import AgentState
from logos_adk.state_store import StateStore, build_state_store
from logos_adk.tools import Tool
from logos_adk.tools.toolkit import Toolkit
from logos_adk.types import (
    ContentPart,
    MediaChunk,
    Message,
    Response,
    StructuredResponse,
    StreamDone,
    StreamChunk,
    TextChunk,
    ToolCall,
    ToolEnd,
    ToolResult,
    ToolStart,
    Usage,
    content_text,
    normalize_content_parts,
)
from logos_adk.context_window import trim_context as _trim_context
from logos_adk.token_counter import count_messages as _count_messages


from logos_adk.agent._budgeting import BudgetMixin
from logos_adk.agent._caching import CachingMixin
from logos_adk.agent._collaboration import CollaborationMixin
from logos_adk.agent._constants import _UNSET
from logos_adk.agent._messages import MessagesMixin
from logos_adk.agent._observability import ObservabilityMixin
from logos_adk.agent._routing import RoutingMixin
from logos_adk.agent._security import SecurityMixin
from logos_adk.agent._self_qa import SelfQAMixin
from logos_adk.agent._state import StateMixin


class Agent(
    CollaborationMixin,
    StateMixin,
    MessagesMixin,
    CachingMixin,
    BudgetMixin,
    RoutingMixin,
    SelfQAMixin,
    SecurityMixin,
    ObservabilityMixin,
    Runnable,
):
    """An AI agent that can use tools and communicate with LLMs."""

    state_class: type[AgentState] = AgentState

    def __init__(self, name: str, *, model: Any | None = None, system: str | None = None) -> None:
        self._config = AgentConfig(name=name, model=model, system_prompt=system)
        self._runtime: Any | None = None
        self._provider: Provider | None = None
        self._sessions: OrderedDict[str | None, AgentState] = OrderedDict()
        self._state_store: StateStore | None = None
        self._max_cached_sessions = 128
        self._team: Any | None = None
        self._pending_topic_subscriptions: list[dict[str, Any]] = []
        self._mcp_connected = False
        self._config_path: str | None = None
        self._input_validator = InputValidator()
        self._prompt_shield = PromptShield()
        self._pii_detector = PIIDetector()
        self._audit_logger = AuditLogger()
        self._retry_policy = RetryPolicy(
            retry_on=(ProviderError, TimeoutError, ProviderTimeoutError)
        )
        self._provider_timeout = 60.0
        self._fallback_models: list[Any] = []
        self._model_routing_learning: Any | None = None
        self._tool_workflow_learning: Any | None = None
        self._provider_circuits: dict[str, CircuitBreaker] = {}
        self._spend_ledger: SpendLedger | None = None
        self._default_task_budget_usd: float | None = None
        self._default_task_budget_tokens: int | None = None
        self._default_budget_threshold_ratio: float = 0.8
        self._input_validation_enabled = True
        self._prompt_shield_enabled = True
        self._pii_detection_enabled = True
        self._audit_logging_enabled = True
        self._circuit_breaker_failure_threshold = 5
        self._circuit_breaker_recovery_timeout = 30.0
        self._output_max_length = 20_000
        self._input_pii_action = "selective"
        self._output_pii_action = "selective"
        self._cached_tool_schemas: list[dict] | None = None
        self._run_cache_depth: int = 0
        self._cached_system_prompt: str | None | object = _UNSET
        self._semantic_cache: Any | None = None  # P2: SemanticCache instance
        self._apply_policy_config()
        self._restore_persistent_tools()

    @property
    def name(self) -> str:
        return self._config.name

    def model(self, model: Any) -> Agent:
        self._config.model = model
        return self

    def system(self, prompt: str) -> Agent:
        self._config.system_prompt = prompt
        return self

    def role(self, role: str, *, goal: str | None = None, backstory: str | None = None) -> Agent:
        self._config.role = role
        if goal is not None:
            self._config.goal = goal
        if backstory is not None:
            self._config.backstory = backstory
        return self

    def goal(self, goal: str) -> Agent:
        self._config.goal = goal
        return self

    def backstory(self, backstory: str) -> Agent:
        self._config.backstory = backstory
        return self

    @property
    def crew_profile(self) -> dict[str, str]:
        profile: dict[str, str] = {}
        if self._config.role:
            profile["role"] = self._config.role
        if self._config.goal:
            profile["goal"] = self._config.goal
        if self._config.backstory:
            profile["backstory"] = self._config.backstory
        return profile

    def tools(self, tools: list[Tool] | Toolkit) -> Agent:
        if isinstance(tools, Toolkit):
            tool_list = tools.get_tools()
        elif isinstance(tools, list):
            tool_list = tools
        else:
            return self
        known_names = {tool.name for tool in self._config.tools}
        for item in tool_list:
            if item.name in known_names:
                continue
            self._config.tools.append(item)
            known_names.add(item.name)
        return self

    def artifact_sandbox(self, *, root: str | None = None) -> Agent:
        """Attach the shared artifact sandbox toolkit."""
        from logos_adk.tools.artifacts import ArtifactToolkit

        return self.tools(ArtifactToolkit(root=root))

    def _restore_persistent_tools(self) -> None:
        try:
            from logos_adk.tool_synthesis import GeneratedToolStore

            GeneratedToolStore().attach_tools_for_agent(self)
        except Exception:
            return

    def mcp(self, source: Any) -> Agent:
        """Connect to an MCP server and import its tools.

        Args:
            source: MCPClient instance, or a string (auto-detects stdio vs SSE).
                - URL string: connects via SSE (e.g., "http://localhost:3000/sse")
                - Command string: launches via stdio (e.g., "npx @server-github")
                - MCPClient: uses the client directly

        The connection is established lazily on first run().
        """
        from logos_adk.mcp import MCPClient

        if isinstance(source, str):
            source = MCPClient.from_string(source)
        self._config.mcp_servers.append(source)
        return self

    def memory(self, *memories: Memory | list[Memory] | tuple[Memory, ...]) -> Agent:
        """Attach one or more memory strategies to the agent."""
        for memory in memories:
            if isinstance(memory, (list, tuple)):
                self._config.memories.extend(memory)
            else:
                self._config.memories.append(memory)
        return self

    def observe(self, *exporters: Exporter | list[Exporter] | tuple[Exporter, ...]) -> Agent:
        """Attach one or more observability exporters to the agent."""
        for exporter in exporters:
            if isinstance(exporter, (list, tuple)):
                self._config.observers.extend(exporter)
            else:
                self._config.observers.append(exporter)
        return self

    def state_backend(self, store: StateStore) -> Agent:
        """Attach a persistent state backend to the agent."""
        self._state_store = store
        return self

    def session_cache(self, max_cached_sessions: int) -> Agent:
        """Bound the number of in-memory session states kept hot."""
        self._max_cached_sessions = max(1, max_cached_sessions)
        self._evict_cold_sessions()
        return self

    def security(
        self,
        *,
        input_validator: InputValidator | None = None,
        prompt_shield: PromptShield | None = None,
        pii_detector: PIIDetector | None = None,
        audit_logger: AuditLogger | None = None,
    ) -> Agent:
        """Attach security primitives to the agent."""
        if input_validator is not None:
            self._input_validator = input_validator
        if prompt_shield is not None:
            self._prompt_shield = prompt_shield
        if pii_detector is not None:
            self._pii_detector = pii_detector
        if audit_logger is not None:
            self._audit_logger = audit_logger
        return self

    def reliability(
        self,
        *,
        retry_policy: RetryPolicy | None = None,
        provider_timeout: float | None = None,
    ) -> Agent:
        """Attach reliability controls to the agent."""
        if retry_policy is not None:
            self._retry_policy = retry_policy
        if provider_timeout is not None:
            self._provider_timeout = max(0.1, provider_timeout)
        return self

    def self_qa(
        self,
        *,
        enabled: bool = True,
        min_assertions: int | None = None,
        max_assertions: int | None = None,
        max_retries: int | None = None,
        fail_closed: bool | None = None,
    ) -> Agent:
        """Enable autonomous self-QA before returning the final answer."""
        self._config.self_qa.enabled = enabled
        if min_assertions is not None:
            self._config.self_qa.min_assertions = max(1, min_assertions)
        if max_assertions is not None:
            self._config.self_qa.max_assertions = max(
                self._config.self_qa.min_assertions,
                max_assertions,
            )
        if max_retries is not None:
            self._config.self_qa.max_retries = max(0, max_retries)
        if fail_closed is not None:
            self._config.self_qa.fail_closed = fail_closed
        return self

    def context_window(
        self,
        *,
        max_tokens: int = 0,
        strategy: str = "chunked_summary",
        reserved_output_tokens: int = 4096,
        chunk_size: int = 10,
        summary_max_tokens: int = 300,
        max_summary_chain: int = 5,
        keep_recent: int = 6,
        summary_model: str | None = None,
        summary_system_prompt: str | None = None,
        smart_tool_selection: bool = False,
        semantic_cache_max_entries: int = 0,
        auto_detect_model_limit: bool = False,
    ) -> Agent:
        """Configure context window management for token optimization."""
        from logos_adk.config import ContextWindowConfig

        self._config.context_window = ContextWindowConfig(
            max_context_tokens=max(0, max_tokens),
            reserved_output_tokens=max(0, reserved_output_tokens),
            strategy=strategy,
            chunk_size=max(2, chunk_size),
            summary_max_tokens=max(50, summary_max_tokens),
            max_summary_chain=max(1, max_summary_chain),
            keep_recent_messages=max(1, keep_recent),
            summary_model=summary_model,
            summary_system_prompt=summary_system_prompt,
            smart_tool_selection=smart_tool_selection,
            semantic_cache_max_entries=semantic_cache_max_entries,
            auto_detect_model_limit=auto_detect_model_limit,
        )
        return self

    def fallback(self, *models: Any) -> Agent:
        """Register fallback providers/models."""
        self._fallback_models.extend(models)
        return self

    def cost_aware_routing(
        self,
        *,
        enabled: bool = True,
        budget_preference: str | None = None,
        workspace_soft_limit_usd: float | None = None,
        session_soft_limit_usd: float | None = None,
        threshold_ratio: float | None = None,
        fail_on_budget_exhausted: bool | None = None,
    ) -> Agent:
        """Enable autonomous model routing based on task complexity and budget preference."""
        self._config.routing.enabled = enabled
        if budget_preference is not None:
            self._config.routing.budget_preference = budget_preference
        if workspace_soft_limit_usd is not None:
            self._config.routing.workspace_soft_limit_usd = workspace_soft_limit_usd
        if session_soft_limit_usd is not None:
            self._config.routing.session_soft_limit_usd = session_soft_limit_usd
        if threshold_ratio is not None:
            self._config.routing.threshold_ratio = threshold_ratio
        if fail_on_budget_exhausted is not None:
            self._config.routing.fail_on_budget_exhausted = bool(fail_on_budget_exhausted)
        return self

    def budget(
        self,
        *,
        task_budget_usd: float | None = None,
        task_budget_tokens: int | None = None,
        threshold_ratio: float | None = None,
        fail_on_budget_exhausted: bool | None = None,
    ) -> Agent:
        """Attach a default per-run task budget."""
        self._default_task_budget_usd = (
            None if task_budget_usd is None else max(0.0, float(task_budget_usd))
        )
        self._default_task_budget_tokens = (
            None if task_budget_tokens is None else max(0, int(task_budget_tokens))
        )
        if threshold_ratio is not None:
            self._default_budget_threshold_ratio = max(0.0, float(threshold_ratio))
            self._config.routing.threshold_ratio = self._default_budget_threshold_ratio
        if task_budget_usd is not None:
            self._config.routing.task_budget_usd = self._default_task_budget_usd
        if task_budget_tokens is not None:
            self._config.routing.task_budget_tokens = self._default_task_budget_tokens
        if fail_on_budget_exhausted is not None:
            self._config.routing.fail_on_budget_exhausted = bool(fail_on_budget_exhausted)
        return self

    def spend_ledger(self, ledger: SpendLedger) -> Agent:
        """Attach a spend ledger used for runtime budget pressure decisions."""
        self._spend_ledger = ledger
        return self

    def model_routing_learning(self, engine: Any) -> Agent:
        """Attach model routing learning for adaptive provider selection."""
        self._model_routing_learning = engine
        return self

    def tool_workflow_learning(self, engine: Any) -> Agent:
        """Attach tool/workflow learning for runtime ordering decisions."""
        self._tool_workflow_learning = engine
        return self

    def temporal_autonomy(self) -> Agent:
        """Attach the temporal scheduling toolkit."""
        from logos_adk.tools.temporal import TemporalToolkit

        return self.tools(TemporalToolkit(self))

    @classmethod
    def from_definition(cls, config: AgentConfig) -> Agent:
        """Build an agent from a typed AgentConfig object."""
        resolved_cls = resolve_agent_class(config, default_cls=cls)
        if resolved_cls is not cls:
            return resolved_cls.from_definition(config)
        agent = cls(config.name)
        agent._config = config
        agent._apply_policy_config()
        return agent

    @classmethod
    def from_config(cls, path: str) -> Agent:
        agent = load_agent_from_config(path, agent_cls=cls)
        agent._config_path = path
        return agent

    def reload(self, path: str | None = None) -> None:
        """Reload agent config from YAML. Preserves session state.

        Args:
            path: Path to YAML config file. If None, uses the original path
                  from from_config().

        Raises:
            ConfigError: If no config path is available.
        """
        reload_path = path or self._config_path
        if reload_path is None:
            raise ConfigError(
                "Cannot reload: agent has no config path. Use Agent.from_config() or pass a path."
            )
        new_config = load_agent_config(reload_path)
        self._config = new_config
        self._config_path = reload_path
        self._provider = None
        self._mcp_connected = False
        self._provider_circuits.clear()
        self._apply_policy_config()

    def on_event(self, event: str, handler: Any) -> Agent:
        self._config.event_handlers.setdefault(event, []).append(handler)
        return self

    def serve(self, **kwargs: Any) -> None:
        """Start a web server for this agent.

        Shortcut for ``serve(agents=[self], ...)``.
        """
        from logos_adk.serve import serve as _serve

        _serve(agents=[self], **kwargs)

    async def _connect_mcp_servers(self) -> None:
        """Lazily connect to configured MCP servers and import their tools."""
        if self._mcp_connected or not self._config.mcp_servers:
            return
        for client in self._config.mcp_servers:
            if not client.connected:
                await client.connect()
            self._config.tools.extend(client.get_tools())
        self._mcp_connected = True

    def _ensure_spawned(self) -> None:
        """Auto-spawn into the default runtime if needed."""
        if self._runtime is None:
            from logos_adk.runtime import _default_runtime_holder

            runtime = _default_runtime_holder.get_or_create()
            runtime.spawn(self)

    def _resolve_model_name(self) -> str:
        """Get a string model name for span attributes and cost tracking."""
        model = self._config.model
        if isinstance(model, str):
            return model
        if model is None:
            return "unknown"
        return self._provider_model_name(model)

    def _create_provider_from_config(self) -> Provider:
        """Create a provider from AgentConfig.provider."""
        from logos_adk.config import ProviderConfig
        from logos_adk.providers.gemini import GeminiProvider
        from logos_adk.providers.anthropic import AnthropicProvider
        from logos_adk.providers.openai import OpenAICompatibleProvider, OpenAIProvider
        from logos_adk.providers.openai_responses import OpenAIResponsesProvider

        provider_config = self._config.provider
        if provider_config is None:
            raise ConfigError("No provider configuration found")

        pc: ProviderConfig = provider_config

        if pc.type == "openai":
            return OpenAIProvider(
                model=pc.model or "gpt-4o",
                api_key=pc.api_key,
                base_url=pc.base_url,
                proxy=pc.proxy,
            )
        elif pc.type == "openai-compatible":
            if not pc.base_url:
                raise ConfigError("openai-compatible provider requires 'base_url'")
            return OpenAICompatibleProvider(
                model=pc.model or "gpt-4o",
                base_url=pc.base_url,
                api_key=pc.api_key,
                api_key_env=pc.api_key_env,
                proxy=pc.proxy,
            )
        elif pc.type == "anthropic":
            return AnthropicProvider(
                model=pc.model or "claude-sonnet-4-20250514",
                api_key=pc.api_key,
                proxy=pc.proxy,
            )
        elif pc.type == "gemini":
            return GeminiProvider(
                model=pc.model or "gemini-2.0-flash",
                api_key=pc.api_key,
                base_url=pc.base_url,
                proxy=pc.proxy,
            )
        elif pc.type == "openai-responses":
            return OpenAIResponsesProvider(
                model=pc.model or "gpt-4.1",
                api_key=pc.api_key,
                base_url=pc.base_url,
                proxy=pc.proxy,
            )
        elif pc.type == "ollama":
            return OpenAICompatibleProvider(
                model=pc.model or "llama3",
                base_url=pc.base_url or "http://localhost:11434/v1",
                api_key=pc.api_key or "not-needed",
                proxy=pc.proxy,
            )
        else:
            raise ConfigError(f"Unknown provider type: {pc.type}")

    def _resolve_provider(self) -> Provider:
        """Resolve the configured model to a concrete provider."""
        if self._provider is not None:
            return self._provider

        if self._config.provider is not None:
            self._provider = self._create_provider_from_config()
            return self._provider

        model = self._config.model
        if model is None:
            raise ConfigError("Agent has no model configured. Use .model() to set one.")
        if isinstance(model, Provider):
            self._provider = model
        elif isinstance(model, str):
            self._provider = provider_registry.resolve(model)
        else:
            raise ConfigError(f"Invalid model type: {type(model)}")

        return self._provider

    def _get_state(self, session_id: str | None) -> AgentState:
        """Get or create per-session agent state."""
        if session_id not in self._sessions:
            state = self._resolve_state_store().load(self.name, session_id)
            if state is None:
                state = self.state_class(session_id=session_id)
            self._sessions[session_id] = state
        self._sessions.move_to_end(session_id)
        self._evict_cold_sessions()
        return self._sessions[session_id]

    def _resolve_state_store(self) -> StateStore:
        if self._state_store is None:
            default_persistent = self._config_path is not None
            self._state_store = build_state_store(
                self._config.state,
                default_persistent=default_persistent,
            )
        return self._state_store

    def _provider_model_name(self, provider: Any) -> str:
        return (
            getattr(provider, "model_name", None)
            or getattr(provider, "_model", None)
            or getattr(provider, "model", None)
            or type(provider).__name__
        )

    def _provider_model_unit_cost(self, provider: Provider) -> float | None:
        model_name = self._provider_model_name(provider)
        for pattern, rates in DEFAULT_RATES.items():
            if pattern == model_name or ("*" in pattern and fnmatch(model_name, pattern)):
                return float(rates[0] + rates[1])
        return None

    def as_tool(self, description: str | None = None) -> Tool:
        """Convert this agent into a Tool that can be used by other agents."""
        from logos_adk.tools import Tool

        async def call_agent(task: str) -> str:
            response = await self.run(task)
            return response.content

        return Tool(
            fn=call_agent,
            name=f"ask_{self.name.lower()}",
            description=description
            or f"Consult with {self.name}. Task should be a clear instruction.",
        )

    def _apply_policy_config(self) -> None:
        security: SecurityConfig = self._config.security
        reliability: ReliabilityConfig = self._config.reliability

        self._input_validation_enabled = security.input_validation_enabled
        self._prompt_shield_enabled = security.prompt_shield_enabled
        self._pii_detection_enabled = security.pii_detection_enabled
        self._audit_logging_enabled = security.audit_logging_enabled
        self._output_max_length = security.output_max_length
        self._input_pii_action = security.input_pii_action
        self._output_pii_action = security.output_pii_action

        self._input_validator = InputValidator(
            min_length=security.input_min_length,
            max_length=security.input_max_length,
            reject_null_bytes=security.reject_null_bytes,
        )
        self._prompt_shield = PromptShield(
            patterns=security.prompt_injection_patterns or None,
            threshold=security.prompt_injection_threshold,
        )
        self._pii_detector = PIIDetector(blocked_types=set(security.pii_blocked_types))
        self._audit_logger = AuditLogger(
            redact_pii=security.audit_redact_pii,
            retention=AuditRetentionPolicy(ttl_days=security.audit_retention_days),
        )

        self._provider_timeout = reliability.timeout_seconds
        self._retry_policy = RetryPolicy(
            max_attempts=reliability.retry_max_attempts,
            base_delay=reliability.retry_base_delay,
            backoff=reliability.retry_backoff,
            retry_on=(ProviderError, TimeoutError, ProviderTimeoutError),
        )
        self._fallback_models = list(reliability.fallback_models)
        self._circuit_breaker_failure_threshold = reliability.circuit_breaker_failure_threshold
        self._circuit_breaker_recovery_timeout = reliability.circuit_breaker_recovery_timeout

    def health(self) -> dict[str, Any]:
        """Return provider health snapshot."""
        providers = []
        try:
            candidates = self._provider_candidates()
        except Exception as exc:
            return {"agent": self.name, "providers": [], "error": str(exc)}
        for provider in candidates:
            circuit = self._provider_circuit(provider)
            providers.append(
                ProviderHealth(
                    provider=type(provider).__name__,
                    circuit_state=circuit.state,
                    failures=circuit.failures,
                    healthy=circuit.state != "open",
                ).__dict__
            )
        return {"agent": self.name, "providers": providers}

    async def run_batch(
        self,
        prompts: list[str],
        *,
        session_ids: list[str | None] | None = None,
        max_concurrency: int = 4,
        **kwargs: Any,
    ) -> list[Response]:
        """Run multiple prompts with bounded concurrency."""
        if not prompts:
            return []
        if session_ids is not None and len(session_ids) != len(prompts):
            raise ValueError("session_ids length must match prompts length")

        indexed_results: list[Response | None] = [None] * len(prompts)
        semaphore = asyncio.Semaphore(max(1, max_concurrency))

        async def run_one(index: int, prompt: str) -> None:
            session_id = session_ids[index] if session_ids is not None else None
            async with semaphore:
                indexed_results[index] = await self.run(prompt, session_id=session_id, **kwargs)

        await asyncio.gather(*(run_one(index, prompt) for index, prompt in enumerate(prompts)))
        return [result for result in indexed_results if result is not None]

    def _build_tool_schemas(self) -> list[dict] | None:
        """Build tool schemas for the provider (cached per run)."""
        if not self._config.tools:
            return None
        if self._cached_tool_schemas is not None:
            return self._cached_tool_schemas
        schemas = [tool.schema for tool in self._ordered_tools()]
        from logos_adk.tool_dedup import deduplicate_tool_schemas

        schemas = deduplicate_tool_schemas(schemas) or []
        self._cached_tool_schemas = schemas
        return schemas

    def _invalidate_tool_schema_cache(self) -> None:
        """Clear the cached tool schemas."""
        self._cached_tool_schemas = None

    def _get_semantic_cache(self):
        """Lazily initialize semantic cache if enabled."""
        max_entries = self._config.context_window.semantic_cache_max_entries
        if max_entries <= 0:
            return None
        if self._semantic_cache is None:
            from logos_adk.semantic_cache import SemanticCache

            self._semantic_cache = SemanticCache(max_entries=max_entries)
        return self._semantic_cache

    def _begin_run_caches(self) -> None:
        """Begin run-scoped caching. Supports nesting (stream->run re-entrancy)."""
        self._run_cache_depth += 1
        if self._run_cache_depth == 1:
            self._cached_system_prompt = _UNSET
            self._cached_tool_schemas = None
            # P2: Auto-configure context window on first run
            cw = self._config.context_window
            if cw.auto_detect_model_limit and cw.max_context_tokens == 0:
                from logos_adk.auto_context import auto_configure_context_window

                model = self._config.model or ""
                self._config.context_window = auto_configure_context_window(cw, model=model)

    def _end_run_caches(self) -> None:
        """End run-scoped caching. Only clears on outermost exit."""
        self._run_cache_depth = max(0, self._run_cache_depth - 1)
        if self._run_cache_depth == 0:
            self._cached_system_prompt = _UNSET
            self._cached_tool_schemas = None

    def _ordered_tools(self) -> list[Tool]:
        # Filter out disabled tools
        disabled = set(self._config.disabled_tools)
        tools = [t for t in self._config.tools if t.name not in disabled]

        if not tools or self._tool_workflow_learning is None:
            return tools
        run_context = get_run_context()
        ordered_names = self._tool_workflow_learning.recommended_tool_names(
            [tool.name for tool in tools],
            task_type=run_context.task_type if run_context else None,
            workspace_id=run_context.workspace_id if run_context else None,
        )
        by_name = {tool.name: tool for tool in tools}
        return [by_name[name] for name in ordered_names if name in by_name]

    def _find_tool(self, name: str) -> Tool | None:
        """Find a configured tool by name."""
        for tool in self._config.tools:
            if tool.name == name:
                return tool
        return None

    async def _execute_tool_calls(self, tool_calls: list[ToolCall]) -> list[Message]:
        """Execute tool calls and return tool messages for the next provider turn."""
        tool_messages: list[Message] = []
        run_context = get_run_context()

        for tool_call in tool_calls:
            tool = self._find_tool(tool_call.name)
            if tool is None:
                tool_result = ToolResult(
                    call_id=tool_call.id,
                    content=None,
                    error=f"Tool '{tool_call.name}' not found",
                )
            else:
                tool_result = await tool.execute(tool_call)
            if self._tool_workflow_learning is not None:
                self._tool_workflow_learning.record_tool_call(
                    tool_name=tool_call.name,
                    success=not bool(tool_result.error),
                    duration_ms=float(tool_result.duration_ms or 0.0),
                    task_type=run_context.task_type if run_context else None,
                    workspace_id=run_context.workspace_id if run_context else None,
                    trace_id=run_context.trace_id if run_context else None,
                    error=tool_result.error,
                )

            tool_messages.append(
                Message(
                    role="tool",
                    content=(
                        ""
                        if tool_result.error
                        else content_text(
                            str(tool_result.content) if tool_result.content is not None else "",
                            tool_result.content_parts,
                        )
                    ),
                    content_parts=normalize_content_parts(tool_result.content_parts),
                    tool_result=tool_result,
                )
            )

        return tool_messages

    def _build_provider_messages(
        self,
        conversation: list[Message],
        memory_messages: list[Message],
        *,
        system_prompt_override: Any = _UNSET,
        tool_schemas: list[dict] | None = None,
        state: AgentState | None = None,
    ) -> list[Message]:
        """Build the provider-facing conversation payload."""
        provider_messages: list[Message] = []
        system_prompt = self._resolved_system_prompt(system_prompt_override)
        if system_prompt:
            provider_messages.append(Message(role="system", content=system_prompt))

        provider_messages.extend(memory_messages)

        # Add conversation messages with sender context
        for msg in conversation:
            content = msg.content
            content_parts = normalize_content_parts(msg.content_parts)
            # If the message is from another agent, prefix it to clarify who said it
            if msg.role == "assistant" and msg.sender and msg.sender != self.name:
                prefix = f"[{msg.sender}]: "
                content = f"{prefix}{content}" if content else prefix.strip()
                if content_parts and content_parts[0].type == "text":
                    first = content_parts[0]
                    content_parts = [
                        ContentPart(
                            type="text",
                            text=f"{prefix}{first.text or ''}",
                            uri=first.uri,
                            data=first.data,
                            mime_type=first.mime_type,
                            name=first.name,
                            metadata=dict(first.metadata),
                        )
                    ] + content_parts[1:]
                elif content_parts:
                    content_parts = [ContentPart(type="text", text=prefix.strip())] + content_parts

            provider_messages.append(
                Message(
                    role=msg.role,
                    content=content,
                    content_parts=content_parts,
                    tool_calls=msg.tool_calls,
                    tool_result=msg.tool_result,
                )
            )

        # --- Context window trimming ---
        cw = self._config.context_window
        if cw.max_context_tokens > 0 and state is not None:
            provider_messages = _trim_context(
                provider_messages,
                config=cw,
                state=state,
                model=self._provider_model_name(self._resolve_provider())
                if self._config.model
                else "",
                tool_schemas=tool_schemas,
            )

        return provider_messages

    def _serialize_config_value(self, value: Any) -> Any:
        if value is None or isinstance(value, (str, int, float, bool)):
            return value
        if is_dataclass(value) and not isinstance(value, type):
            return {key: self._serialize_config_value(item) for key, item in asdict(value).items()}
        if isinstance(value, dict):
            return {
                str(key): self._serialize_config_value(item)
                for key, item in sorted(value.items(), key=lambda item: str(item[0]))
            }
        if isinstance(value, (list, tuple, set)):
            return [self._serialize_config_value(item) for item in value]
        if callable(value):
            return getattr(value, "__name__", type(value).__name__)
        return repr(value)

    def _config_fingerprint(self) -> str:
        payload = {
            "name": self._config.name,
            "model": self._resolve_model_name(),
            "system_prompt": self._config.system_prompt,
            "role": self._config.role,
            "goal": self._config.goal,
            "backstory": self._config.backstory,
            "tools": [tool.name for tool in self._config.tools],
            "memories": [type(memory).__name__ for memory in self._config.memories],
            "observers": [type(observer).__name__ for observer in self._config.observers],
            "state": self._serialize_config_value(self._config.state),
            "security": self._serialize_config_value(self._config.security),
            "reliability": self._serialize_config_value(self._config.reliability),
        }
        serialized = json.dumps(payload, sort_keys=True, separators=(",", ":"))
        return hashlib.sha256(serialized.encode()).hexdigest()[:16]

    def output_schema(self, schema: type) -> Agent:
        """Set a Pydantic model as the structured output schema.

        Args:
            schema: A Pydantic BaseModel subclass defining the expected output.

        Usage::

            from pydantic import BaseModel

            class Analysis(BaseModel):
                sentiment: str
                confidence: float

            agent = Agent("analyst", model="...").output_schema(Analysis)
            response = await agent.run("Analyze: great product!")
            # response.parsed is an Analysis instance
        """
        self._config.output_schema = schema
        return self

    def cache(
        self,
        backend: Any,
        ttl: int | None = 3600,
        semantic: bool = False,
        embed_fn: Any = None,
        similarity_threshold: float = 0.95,
    ) -> Agent:
        """Enable response caching.

        Args:
            backend: A CacheBackend instance.
            ttl: Time-to-live in seconds (None = no expiry).
            semantic: Enable semantic similarity matching.
            embed_fn: Async function to embed text into vectors (required if semantic=True).
            similarity_threshold: Cosine similarity threshold for semantic matching.
        """
        from logos_adk.cache import CacheConfig
        from logos_adk.errors import ConfigError

        if semantic and embed_fn is None:
            raise ConfigError("semantic caching requires embed_fn")
        self._config.cache_backend = backend
        self._config.cache_config = CacheConfig(
            ttl=ttl,
            semantic=semantic,
            embed_fn=embed_fn,
            similarity_threshold=similarity_threshold,
        )
        return self

    def knowledge(self, kb: Any, *, mode: str = "memory", limit: int = 5) -> Agent:
        """Attach a KnowledgeBase to the agent.

        Args:
            kb: A KnowledgeBase instance.
            mode: Integration mode — "memory" (auto-inject context),
                  "tool" (agent-driven search), or "both".
            limit: Maximum number of chunks to retrieve.
        """
        from logos_adk.rag.memory import RAGMemory
        from logos_adk.rag.tool import rag_tool

        self._config.knowledge_base = kb
        if mode in ("memory", "both"):
            self._config.memories.append(RAGMemory(kb, limit=limit))
        if mode in ("tool", "both"):
            self._config.tools.append(rag_tool(kb, limit=limit))
        return self

    def guardrail(self, guard: Any) -> Agent:
        """Add a guardrail that checks both input and output."""
        self._config.input_guardrails.append(guard)
        self._config.output_guardrails.append(guard)
        return self

    def input_guardrail(self, guard: Any) -> Agent:
        """Add a guardrail that checks only input (user prompts)."""
        self._config.input_guardrails.append(guard)
        return self

    def output_guardrail(self, guard: Any) -> Agent:
        """Add a guardrail that checks only output (LLM responses)."""
        self._config.output_guardrails.append(guard)
        return self

    def _parse_structured_response(self, response: Response, schema: type) -> StructuredResponse:
        """Parse a provider response into a StructuredResponse using the schema."""
        import json as json_mod
        from pydantic import BaseModel, ValidationError

        parsed = None
        content = content_text(response.content, response.content_parts).strip()

        # Try to extract JSON from the response content
        try:
            # Handle markdown code blocks
            if content.startswith("```"):
                # Extract content between ```json\n...\n``` or ```\n...\n```
                lines = content.split("\n")
                start = 1
                end = len(lines) - 1
                if lines[end].strip() == "```":
                    content = "\n".join(lines[start:end])

            data = json_mod.loads(content)
            if issubclass(schema, BaseModel):
                parsed = schema.model_validate(data)  # type: ignore[attr-defined]
        except (json_mod.JSONDecodeError, ValidationError):
            parsed = None

        return StructuredResponse(
            content=response.content,
            content_parts=response.content_parts,
            tool_calls=response.tool_calls,
            usage=response.usage,
            raw=response.raw,
            parsed=parsed,
        )

    async def run(
        self,
        prompt: str | list[ContentPart],
        *,
        session_id: str | None = None,
        **kwargs: Any,
    ) -> Response:
        """Run the agent with a user prompt."""
        import logging

        agent_logger = logging.getLogger(f"logos_adk.agent.{self.name}")
        prompt_text, prompt_parts = self._normalize_prompt_input(prompt)
        agent_logger.info(f"Начинаю выполнение задачи. Промпт: {prompt_text[:100]}...")

        trace_id = kwargs.pop("_trace_id", uuid4().hex)
        request_id = kwargs.pop("_request_id", None)
        root_started_at = datetime.now(UTC)
        root_span_id = kwargs.pop("_root_span_id", uuid4().hex)
        on_text = kwargs.pop("on_text", None)
        on_tool_start = kwargs.pop("on_tool_start", None)
        on_tool_end = kwargs.pop("on_tool_end", None)
        on_done = kwargs.pop("on_done", None)
        output_schema = kwargs.pop("output_schema", None) or self._config.output_schema
        system_prompt_override = kwargs.pop("system_prompt_override", _UNSET)
        workspace_id = kwargs.pop("_workspace_id", None)
        task_type = kwargs.pop("_task_type", None)
        user_segment = kwargs.pop("_user_segment", None)
        prompt_name = kwargs.pop("_prompt_name", None)
        prompt_version = kwargs.pop("_prompt_version", None)
        prompt_variant = kwargs.pop("_prompt_variant", None)
        experiment_id = kwargs.pop("_experiment_id", None)
        routing_hint = kwargs.pop("_routing_hint", None)
        criticality = kwargs.pop("_criticality", None)
        kwargs.pop("task_run_id", None)
        budget_preference = kwargs.pop("_budget_preference", None)
        fail_on_budget_exhausted = kwargs.pop("fail_on_budget_exhausted", None)
        task_budget_usd = kwargs.pop("task_budget_usd", None)
        task_budget_tokens = kwargs.pop("task_budget_tokens", None)
        team_budget_usd = kwargs.pop("team_budget_usd", None)
        team_budget_tokens = kwargs.pop("team_budget_tokens", None)
        budget_threshold_ratio = kwargs.pop("budget_threshold_ratio", None)
        task_budget = kwargs.pop("_task_budget", None)
        team_budget = kwargs.pop("_team_budget", None)
        task_budget, team_budget = self._resolve_runtime_budget_accounts(
            task_budget=task_budget,
            team_budget=team_budget,
            task_budget_usd=task_budget_usd,
            task_budget_tokens=task_budget_tokens,
            team_budget_usd=team_budget_usd,
            team_budget_tokens=team_budget_tokens,
            threshold_ratio=budget_threshold_ratio,
        )

        run_context = RunContext(
            trace_id=trace_id,
            request_id=request_id,
            session_id=session_id,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            prompt_text=prompt_text,
            prompt_name=prompt_name or self.name,
            prompt_version=prompt_version,
            prompt_variant=prompt_variant,
            experiment_id=experiment_id,
            routing_hint=routing_hint,
            criticality=criticality,
            budget_preference=budget_preference,
            fail_on_budget_exhausted=bool(
                self._config.routing.fail_on_budget_exhausted
                if fail_on_budget_exhausted is None
                else fail_on_budget_exhausted
            ),
            task_budget=task_budget,
            team_budget=team_budget,
        )
        context_token = set_run_context(run_context)
        try:
            await self._connect_mcp_servers()
            callbacks_requested = any(
                callback is not None for callback in (on_text, on_tool_start, on_tool_end, on_done)
            )
            self_qa_enabled = self._self_qa_enabled()
            if callbacks_requested and not self_qa_enabled:
                async for chunk in self.stream(
                    prompt,
                    session_id=session_id,
                    _trace_id=trace_id,
                    **kwargs,
                ):
                    if isinstance(chunk, TextChunk):
                        await self._invoke_callback(on_text, chunk)
                    elif isinstance(chunk, ToolStart):
                        await self._invoke_callback(on_tool_start, chunk)
                    elif isinstance(chunk, ToolEnd):
                        await self._invoke_callback(on_tool_end, chunk)
                    elif isinstance(chunk, StreamDone):
                        await self._invoke_callback(on_done, chunk.response)
                        await self._emit_span(
                            name="Agent.run",
                            trace_id=trace_id,
                            span_id=root_span_id,
                            parent_span_id=None,
                            started_at=root_started_at,
                            session_id=session_id,
                            attributes={"mode": "callbacks"},
                        )
                        await self._flush_observers()
                        result = chunk.response
                        if output_schema:
                            result = self._parse_structured_response(result, output_schema)
                        return result

                raise RuntimeError("Agent.stream() completed without emitting StreamDone")

            prompt_text = await self._secure_prompt(
                prompt_text,
                session_id=session_id,
                trace_id=trace_id,
                request_id=request_id,
            )
            if self._config.input_guardrails:
                prompt_text = await self._check_guardrails(
                    prompt_text,
                    self._config.input_guardrails,
                    role="user",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                )

            self._ensure_spawned()
            provider = self._resolve_provider()
            state = self._get_state(session_id)
            self._begin_run_caches()
            tool_schemas = self._build_tool_schemas()

            conversation: list[Message] = list(state.messages)
            conversation.append(
                Message(role="user", content=prompt_text, content_parts=prompt_parts)
            )

            remember_started_at = datetime.now(UTC)
            remembered_entries = await self._remember(conversation, state, session_id)
            await self._emit_span(
                name="Memory.remember",
                trace_id=trace_id,
                parent_span_id=root_span_id,
                started_at=remember_started_at,
                session_id=session_id,
                attributes={"entries_found": len(remembered_entries)},
            )
            memory_messages = self._memory_messages(remembered_entries)

            provider_messages_for_cache = self._build_provider_messages(
                conversation,
                memory_messages,
                system_prompt_override=system_prompt_override,
                tool_schemas=tool_schemas,
                state=state,
            )
            cached = None
            if not self_qa_enabled:
                cached = await self._cache_lookup(
                    provider_messages_for_cache,
                    trace_id,
                    root_span_id,
                    session_id,
                )
            if cached is not None:
                conversation.append(
                    Message(
                        role="assistant",
                        content=cached.content,
                        content_parts=cached.content_parts,
                        sender=self.name,
                    )
                )
                state.messages = conversation
                state = self._persist_state(state)
                await self._emit_span(
                    name="Agent.run",
                    trace_id=trace_id,
                    span_id=root_span_id,
                    parent_span_id=None,
                    started_at=root_started_at,
                    session_id=session_id,
                    attributes={"mode": "cache_hit"},
                )
                await self._flush_observers()
                if output_schema:
                    return self._parse_structured_response(
                        self._attach_budget_metadata(cached, run_context=run_context),
                        output_schema,
                    )
                return self._attach_budget_metadata(cached, run_context=run_context)

            budget_failure = self._strict_budget_failure_response(
                run_context=run_context,
                stage="preflight",
            )
            if budget_failure is not None:
                conversation.append(
                    Message(
                        role="assistant",
                        content=budget_failure.content,
                        content_parts=budget_failure.content_parts,
                        sender=self.name,
                    )
                )
                state.messages = conversation
                self._persist_state(state)
                return budget_failure

            max_tool_rounds = self._effective_max_tool_rounds(run_context)
            for _ in range(max_tool_rounds + 1):
                provider_messages = self._build_provider_messages(
                    conversation,
                    memory_messages,
                    system_prompt_override=system_prompt_override,
                    tool_schemas=tool_schemas,
                    state=state,
                )

                await self._emit_metric(
                    name="adk_context_tokens",
                    value=_count_messages(
                        provider_messages, model=self._provider_model_name(self._resolve_provider())
                    ),
                    unit="tokens",
                    attributes={"agent_name": self.name, "stage": "after_trim"},
                )

                llm_started_at = datetime.now(UTC)
                provider, response = await self._provider_generate_with_reliability(
                    provider_messages,
                    tools=tool_schemas,
                    output_schema=output_schema,
                    trace_id=trace_id,
                    session_id=session_id,
                    request_id=request_id,
                    **kwargs,
                )
                await self._emit_span(
                    name="LLM.generate",
                    trace_id=trace_id,
                    parent_span_id=root_span_id,
                    started_at=llm_started_at,
                    session_id=session_id,
                    attributes={
                        "provider": type(provider).__name__,
                        "model": self._provider_model_name(provider),
                        "tool_calls": len(response.tool_calls),
                        "input_tokens": response.usage.input_tokens,
                        "output_tokens": response.usage.output_tokens,
                    },
                )
                await self._emit_metric(
                    name="adk_llm_tokens_input",
                    value=response.usage.input_tokens,
                    unit="tokens",
                    attributes={
                        "agent_name": self.name,
                        "model": self._provider_model_name(provider),
                    },
                )
                await self._emit_metric(
                    name="adk_llm_tokens_output",
                    value=response.usage.output_tokens,
                    unit="tokens",
                    attributes={
                        "agent_name": self.name,
                        "model": self._provider_model_name(provider),
                    },
                )

                budget_failure = self._strict_budget_failure_response(
                    run_context=run_context,
                    stage="llm_generate",
                    usage=response.usage,
                    raw=response.raw,
                )
                if budget_failure is not None:
                    conversation.append(
                        Message(
                            role="assistant",
                            content=budget_failure.content,
                            content_parts=budget_failure.content_parts,
                            sender=self.name,
                        )
                    )
                    state.messages = conversation
                    self._persist_state(state)
                    await self._flush_observers()
                    return budget_failure

                if not response.tool_calls:
                    if self._config.output_guardrails:
                        checked = await self._check_guardrails(
                            content_text(response.content, response.content_parts),
                            self._config.output_guardrails,
                            role="assistant",
                            session_id=session_id,
                            trace_id=trace_id,
                            request_id=request_id,
                        )
                        if checked != response.content:
                            response = Response(
                                content=checked,
                                content_parts=response.content_parts,
                                tool_calls=response.tool_calls,
                                usage=response.usage,
                                raw=response.raw,
                            )
                    response = await self._secure_output(
                        response,
                        session_id=session_id,
                        trace_id=trace_id,
                        request_id=request_id,
                    )
                    response = await self._run_self_qa(
                        prompt_text=prompt_text,
                        response=response,
                        trace_id=trace_id,
                        session_id=session_id,
                        request_id=request_id,
                    )
                    response = self._attach_budget_metadata(
                        response,
                        run_context=run_context,
                        self_qa_max_retries=self._effective_self_qa_max_retries(run_context),
                        max_tool_rounds=max_tool_rounds,
                    )
                    conversation.append(
                        Message(
                            role="assistant",
                            content=response.content,
                            content_parts=response.content_parts,
                            sender=self.name,
                        )
                    )
                    state.messages = conversation
                    if response.status == "completed":
                        commit_started_at = datetime.now(UTC)
                        await self._commit_memories(conversation, response, state, session_id)
                        await self._emit_span(
                            name="Memory.commit",
                            trace_id=trace_id,
                            parent_span_id=root_span_id,
                            started_at=commit_started_at,
                            session_id=session_id,
                            attributes={"strategies": len(self._config.memories)},
                        )
                    await self._emit_span(
                        name="Agent.run",
                        trace_id=trace_id,
                        span_id=root_span_id,
                        parent_span_id=None,
                        started_at=root_started_at,
                        session_id=session_id,
                        status="ok" if response.status == "completed" else "error",
                        attributes={
                            "mode": "generate",
                            "final_status": response.status,
                            "self_qa_enabled": self_qa_enabled,
                        },
                    )
                    await self._emit_metric(
                        name="adk_session_memory_messages",
                        value=len(state.messages),
                        unit="count",
                        attributes={"agent_name": self.name},
                    )
                    state = self._persist_state(state)
                    await self._flush_observers()
                    if response.status == "completed":
                        await self._cache_store(
                            provider_messages,
                            response,
                            trace_id,
                            root_span_id,
                            session_id,
                        )
                    if callbacks_requested:
                        final_text = content_text(response.content, response.content_parts)
                        if final_text:
                            await self._invoke_callback(on_text, TextChunk(text=final_text))
                        await self._invoke_callback(on_done, response)
                    if output_schema and response.status == "completed":
                        return self._parse_structured_response(response, output_schema)
                    return response

                conversation.append(
                    Message(
                        role="assistant",
                        content=response.content,
                        content_parts=response.content_parts,
                        tool_calls=response.tool_calls,
                    )
                )
                for tool_call in response.tool_calls:
                    tool_started_at = datetime.now(UTC)
                    agent_logger.info(
                        f"Вызываю инструмент '{tool_call.name}' с аргументами: {tool_call.arguments}"
                    )
                    tool_message = (await self._execute_tool_calls([tool_call]))[0]
                    agent_logger.info(
                        f"Инструмент '{tool_call.name}' вернул: {str(tool_message.content)[:200]}..."
                    )
                    conversation.append(tool_message)

                    # Handle dynamic pause from tool
                    if tool_message.tool_result and tool_message.tool_result.status == "paused":
                        state.messages = conversation
                        self._persist_state(state)
                        await self._emit_span(
                            name=f"Tool.{tool_call.name}",
                            trace_id=trace_id,
                            parent_span_id=root_span_id,
                            started_at=tool_started_at,
                            session_id=session_id,
                            attributes={"status": "paused"},
                        )
                        await self._flush_observers()
                        return Response(
                            content=tool_message.content or "Execution paused by agent.",
                            content_parts=tool_message.content_parts,
                            tool_calls=response.tool_calls,
                            usage=response.usage,
                            raw=response.raw,
                            status="paused",
                        )

                    await self._emit_span(
                        name=f"Tool.{tool_call.name}",
                        trace_id=trace_id,
                        parent_span_id=root_span_id,
                        started_at=tool_started_at,
                        session_id=session_id,
                        status="error"
                        if tool_message.tool_result and tool_message.tool_result.error
                        else "ok",
                        attributes={
                            "call_id": tool_call.id,
                            "error": tool_message.tool_result.error
                            if tool_message.tool_result
                            else None,
                        },
                    )
                    await self._emit_metric(
                        name="adk_tool_calls_total",
                        value=1.0,
                        unit="count",
                        attributes={
                            "agent_name": self.name,
                            "tool_name": tool_call.name,
                            "status": "error"
                            if tool_message.tool_result and tool_message.tool_result.error
                            else "ok",
                        },
                    )

            state.messages = conversation
            await self._emit_span(
                name="Agent.run",
                trace_id=trace_id,
                span_id=root_span_id,
                parent_span_id=None,
                started_at=root_started_at,
                session_id=session_id,
                status="error",
                attributes={"reason": "max_tool_rounds_exceeded"},
            )
            await self._flush_observers()
            return response
        except Exception as exc:
            await self._emit_span(
                name="Agent.run",
                trace_id=trace_id,
                span_id=root_span_id,
                parent_span_id=None,
                started_at=root_started_at,
                session_id=session_id,
                status="error",
                attributes={"exception_type": type(exc).__name__},
            )
            await self._flush_observers()
            raise
        finally:
            self._end_run_caches()
            reset_run_context(context_token)

    async def stream(
        self,
        prompt: str | list[ContentPart],
        *,
        session_id: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[StreamChunk]:
        """Stream the agent response as chunks."""
        import logging

        agent_logger = logging.getLogger(f"logos_adk.agent.{self.name}")
        trace_id = kwargs.pop("_trace_id", uuid4().hex)
        request_id = kwargs.pop("_request_id", None)
        root_span_id = kwargs.pop("_root_span_id", uuid4().hex)
        root_started_at = datetime.now(UTC)
        system_prompt_override = kwargs.pop("system_prompt_override", _UNSET)
        workspace_id = kwargs.pop("_workspace_id", None)
        task_type = kwargs.pop("_task_type", None)
        user_segment = kwargs.pop("_user_segment", None)
        prompt_name = kwargs.pop("_prompt_name", None)
        prompt_version = kwargs.pop("_prompt_version", None)
        prompt_variant = kwargs.pop("_prompt_variant", None)
        experiment_id = kwargs.pop("_experiment_id", None)
        routing_hint = kwargs.pop("_routing_hint", None)
        criticality = kwargs.pop("_criticality", None)
        kwargs.pop("task_run_id", None)
        budget_preference = kwargs.pop("_budget_preference", None)
        fail_on_budget_exhausted = kwargs.pop("fail_on_budget_exhausted", None)
        task_budget_usd = kwargs.pop("task_budget_usd", None)
        task_budget_tokens = kwargs.pop("task_budget_tokens", None)
        team_budget_usd = kwargs.pop("team_budget_usd", None)
        team_budget_tokens = kwargs.pop("team_budget_tokens", None)
        budget_threshold_ratio = kwargs.pop("budget_threshold_ratio", None)
        task_budget = kwargs.pop("_task_budget", None)
        team_budget = kwargs.pop("_team_budget", None)
        effective_fail_on_budget_exhausted = bool(
            self._config.routing.fail_on_budget_exhausted
            if fail_on_budget_exhausted is None
            else fail_on_budget_exhausted
        )
        if self._self_qa_enabled() or effective_fail_on_budget_exhausted:
            response = await self.run(
                prompt,
                session_id=session_id,
                _trace_id=trace_id,
                _request_id=request_id,
                _root_span_id=root_span_id,
                system_prompt_override=system_prompt_override,
                _workspace_id=workspace_id,
                _task_type=task_type,
                _user_segment=user_segment,
                _prompt_name=prompt_name,
                _prompt_version=prompt_version,
                _prompt_variant=prompt_variant,
                _experiment_id=experiment_id,
                _routing_hint=routing_hint,
                _criticality=criticality,
                _budget_preference=budget_preference,
                fail_on_budget_exhausted=effective_fail_on_budget_exhausted,
                task_budget_usd=task_budget_usd,
                task_budget_tokens=task_budget_tokens,
                team_budget_usd=team_budget_usd,
                team_budget_tokens=team_budget_tokens,
                budget_threshold_ratio=budget_threshold_ratio,
                _task_budget=task_budget,
                _team_budget=team_budget,
                **kwargs,
            )
            final_text = content_text(response.content, response.content_parts)
            if final_text:
                yield TextChunk(text=final_text)
            for part in response.content_parts:
                if part.type != "text":
                    yield MediaChunk(part=part)
            yield StreamDone(response=response)
            return
        prompt_text, prompt_parts = self._normalize_prompt_input(prompt)
        task_budget, team_budget = self._resolve_runtime_budget_accounts(
            task_budget=task_budget,
            team_budget=team_budget,
            task_budget_usd=task_budget_usd,
            task_budget_tokens=task_budget_tokens,
            team_budget_usd=team_budget_usd,
            team_budget_tokens=team_budget_tokens,
            threshold_ratio=budget_threshold_ratio,
        )
        run_context = RunContext(
            trace_id=trace_id,
            request_id=request_id,
            session_id=session_id,
            workspace_id=workspace_id,
            task_type=task_type,
            user_segment=user_segment,
            prompt_text=prompt_text,
            prompt_name=prompt_name or self.name,
            prompt_version=prompt_version,
            prompt_variant=prompt_variant,
            experiment_id=experiment_id,
            routing_hint=routing_hint,
            criticality=criticality,
            budget_preference=budget_preference,
            fail_on_budget_exhausted=effective_fail_on_budget_exhausted,
            task_budget=task_budget,
            team_budget=team_budget,
        )
        context_token = set_run_context(run_context)
        try:
            await self._connect_mcp_servers()
            prompt_text = await self._secure_prompt(
                prompt_text,
                session_id=session_id,
                trace_id=trace_id,
                request_id=request_id,
            )
            if self._config.input_guardrails:
                prompt_text = await self._check_guardrails(
                    prompt_text,
                    self._config.input_guardrails,
                    role="user",
                    session_id=session_id,
                    trace_id=trace_id,
                    request_id=request_id,
                )
            self._ensure_spawned()
            state = self._get_state(session_id)
            self._begin_run_caches()
            tool_schemas = self._build_tool_schemas()

            conversation: list[Message] = list(state.messages)
            conversation.append(
                Message(role="user", content=prompt_text, content_parts=prompt_parts)
            )
            remember_started_at = datetime.now(UTC)
            remembered_entries = await self._remember(conversation, state, session_id)
            await self._emit_span(
                name="Memory.remember",
                trace_id=trace_id,
                parent_span_id=root_span_id,
                started_at=remember_started_at,
                session_id=session_id,
                attributes={"entries_found": len(remembered_entries)},
            )
            memory_messages = self._memory_messages(remembered_entries)

            provider_messages_for_cache = self._build_provider_messages(
                conversation,
                memory_messages,
                system_prompt_override=system_prompt_override,
                tool_schemas=tool_schemas,
                state=state,
            )
            cached = await self._cache_lookup(
                provider_messages_for_cache,
                trace_id,
                root_span_id,
                session_id,
            )
            if cached is not None:
                conversation.append(
                    Message(
                        role="assistant",
                        content=cached.content,
                        content_parts=cached.content_parts,
                        sender=self.name,
                    )
                )
                state.messages = conversation
                state = self._persist_state(state)
                if cached.content:
                    yield TextChunk(text=cached.content)
                for part in cached.content_parts:
                    if part.type != "text":
                        yield MediaChunk(part=part)
                await self._emit_span(
                    name="Agent.stream",
                    trace_id=trace_id,
                    span_id=root_span_id,
                    parent_span_id=None,
                    started_at=root_started_at,
                    session_id=session_id,
                    attributes={"mode": "cache_hit"},
                )
                await self._flush_observers()
                yield StreamDone(
                    response=self._attach_budget_metadata(cached, run_context=run_context)
                )
                return

            budget_failure = self._strict_budget_failure_response(
                run_context=run_context,
                stage="preflight",
            )
            if budget_failure is not None:
                conversation.append(
                    Message(
                        role="assistant",
                        content=budget_failure.content,
                        content_parts=budget_failure.content_parts,
                        sender=self.name,
                    )
                )
                state.messages = conversation
                self._persist_state(state)
                yield StreamDone(response=budget_failure)
                return

            max_tool_rounds = self._effective_max_tool_rounds(run_context)
            for _ in range(max_tool_rounds + 1):
                provider_messages = self._build_provider_messages(
                    conversation,
                    memory_messages,
                    system_prompt_override=system_prompt_override,
                    tool_schemas=tool_schemas,
                    state=state,
                )

                streamed_text: list[str] = []
                streamed_parts: list[ContentPart] = []
                streamed_response: Response | None = None
                provider: Provider | None = None
                stream_started = False

                llm_started_at = datetime.now(UTC)
                last_error: Exception | None = None
                primary_provider = self._resolve_provider()
                run_context = get_run_context()
                candidates = self._provider_candidates()
                routing_decision = self._routing_decision_for_candidates(candidates)
                ordered_candidates = self._ordered_candidates_from_decision(
                    candidates, routing_decision
                )
                await self._emit_model_routing_decision_metrics(
                    routing_decision=routing_decision,
                    run_context=run_context,
                )
                for candidate in ordered_candidates:
                    circuit = self._provider_circuit(candidate)
                    candidate_started = time.monotonic()
                    candidate_model_name = self._provider_model_name(candidate)
                    try:
                        ensure_circuit_allows(circuit, type(candidate).__name__)
                        provider = candidate
                        async for chunk in candidate.stream(
                            provider_messages, tools=tool_schemas, **kwargs
                        ):
                            stream_started = True
                            if isinstance(chunk, TextChunk):
                                streamed_text.append(chunk.text)
                                yield chunk
                                continue

                            if isinstance(chunk, MediaChunk):
                                streamed_parts.append(chunk.part)
                                yield chunk
                                continue

                            if isinstance(chunk, str):
                                text_chunk = TextChunk(text=chunk)
                                streamed_text.append(chunk)
                                yield text_chunk
                                continue

                            if isinstance(chunk, Response):
                                streamed_response = chunk
                                continue

                            yield chunk
                        circuit.record_success()
                        if self._tool_workflow_learning is not None:
                            self._tool_workflow_learning.record_fallback(
                                provider_key=self._provider_key(candidate),
                                success=True,
                                duration_ms=(time.monotonic() - candidate_started) * 1000,
                                task_type=run_context.task_type if run_context else None,
                                workspace_id=run_context.workspace_id if run_context else None,
                                trace_id=trace_id,
                                was_primary=candidate is primary_provider,
                            )
                        break
                    except asyncio.TimeoutError:
                        circuit.record_failure()
                        last_error = ProviderTimeoutError(
                            f"Provider timeout after {self._provider_timeout}s"
                        )
                        if self._model_routing_learning is not None:
                            self._model_routing_learning.record_route(
                                model_key=self._provider_key(candidate),
                                success=False,
                                latency_ms=(time.monotonic() - candidate_started) * 1000,
                                cost_usd=0.0,
                                confidence_tier=(
                                    routing_decision.confidence_tier
                                    if routing_decision is not None
                                    else "medium"
                                ),
                                workspace_id=run_context.workspace_id if run_context else None,
                                task_type=run_context.task_type if run_context else None,
                                user_segment=run_context.user_segment if run_context else None,
                                trace_id=trace_id,
                                was_canary=(
                                    routing_decision.used_canary
                                    and routing_decision.selected_model_key
                                    == self._provider_key(candidate)
                                    if routing_decision is not None
                                    else False
                                ),
                            )
                            await self._emit_model_routing_result_metric(
                                model_key=self._provider_key(candidate),
                                success=False,
                                confidence_tier=(
                                    routing_decision.confidence_tier
                                    if routing_decision is not None
                                    else "medium"
                                ),
                                run_context=run_context,
                                was_canary=(
                                    routing_decision.used_canary
                                    and routing_decision.selected_model_key
                                    == self._provider_key(candidate)
                                    if routing_decision is not None
                                    else False
                                ),
                            )
                        if self._tool_workflow_learning is not None:
                            self._tool_workflow_learning.record_fallback(
                                provider_key=self._provider_key(candidate),
                                success=False,
                                duration_ms=(time.monotonic() - candidate_started) * 1000,
                                task_type=run_context.task_type if run_context else None,
                                workspace_id=run_context.workspace_id if run_context else None,
                                trace_id=trace_id,
                                was_primary=candidate is primary_provider,
                                error=str(last_error),
                            )
                    except Exception as exc:
                        circuit.record_failure()
                        last_error = exc
                        if self._model_routing_learning is not None:
                            self._model_routing_learning.record_route(
                                model_key=self._provider_key(candidate),
                                success=False,
                                latency_ms=(time.monotonic() - candidate_started) * 1000,
                                cost_usd=0.0,
                                confidence_tier=(
                                    routing_decision.confidence_tier
                                    if routing_decision is not None
                                    else "medium"
                                ),
                                workspace_id=run_context.workspace_id if run_context else None,
                                task_type=run_context.task_type if run_context else None,
                                user_segment=run_context.user_segment if run_context else None,
                                trace_id=trace_id,
                                was_canary=(
                                    routing_decision.used_canary
                                    and routing_decision.selected_model_key
                                    == self._provider_key(candidate)
                                    if routing_decision is not None
                                    else False
                                ),
                            )
                            await self._emit_model_routing_result_metric(
                                model_key=self._provider_key(candidate),
                                success=False,
                                confidence_tier=(
                                    routing_decision.confidence_tier
                                    if routing_decision is not None
                                    else "medium"
                                ),
                                run_context=run_context,
                                was_canary=(
                                    routing_decision.used_canary
                                    and routing_decision.selected_model_key
                                    == self._provider_key(candidate)
                                    if routing_decision is not None
                                    else False
                                ),
                            )
                        if self._tool_workflow_learning is not None:
                            self._tool_workflow_learning.record_fallback(
                                provider_key=self._provider_key(candidate),
                                success=False,
                                duration_ms=(time.monotonic() - candidate_started) * 1000,
                                task_type=run_context.task_type if run_context else None,
                                workspace_id=run_context.workspace_id if run_context else None,
                                trace_id=trace_id,
                                was_primary=candidate is primary_provider,
                                error=str(exc),
                            )
                        if stream_started:
                            raise
                        continue
                if provider is None:
                    if last_error is not None:
                        raise last_error
                    raise RuntimeError("No provider candidates available")

                response = streamed_response or Response(
                    content="".join(streamed_text),
                    content_parts=streamed_parts,
                    tool_calls=[],
                    usage=Usage(input_tokens=0, output_tokens=0),
                    raw={},
                )
                if streamed_text and response.content != "".join(streamed_text):
                    response = Response(
                        content="".join(streamed_text),
                        content_parts=response.content_parts,
                        tool_calls=response.tool_calls,
                        usage=response.usage,
                        raw=response.raw,
                    )
                response = self._attach_routing_metadata(
                    response,
                    provider=provider,
                    routing_decision=routing_decision,
                )
                self._record_spend(
                    provider=provider,
                    usage=response.usage,
                    run_context=run_context,
                    trace_id=trace_id,
                )
                if self._model_routing_learning is not None and provider is not None:
                    self._model_routing_learning.record_route(
                        model_key=self._provider_key(provider),
                        success=True,
                        latency_ms=(time.monotonic() - candidate_started) * 1000,
                        cost_usd=self._estimate_response_cost_usd(
                            candidate_model_name, response.usage
                        ),
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        workspace_id=run_context.workspace_id if run_context else None,
                        task_type=run_context.task_type if run_context else None,
                        user_segment=run_context.user_segment if run_context else None,
                        trace_id=trace_id,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                    await self._emit_model_routing_result_metric(
                        model_key=self._provider_key(provider),
                        success=True,
                        confidence_tier=(
                            routing_decision.confidence_tier
                            if routing_decision is not None
                            else "medium"
                        ),
                        run_context=run_context,
                        was_canary=(
                            routing_decision.used_canary
                            and routing_decision.selected_model_key == self._provider_key(provider)
                            if routing_decision is not None
                            else False
                        ),
                    )
                await self._emit_span(
                    name="LLM.stream",
                    trace_id=trace_id,
                    parent_span_id=root_span_id,
                    started_at=llm_started_at,
                    session_id=session_id,
                    attributes={
                        "provider": type(provider).__name__,
                        "model": self._provider_model_name(provider),
                        "tool_calls": len(response.tool_calls),
                        "input_tokens": response.usage.input_tokens,
                        "output_tokens": response.usage.output_tokens,
                    },
                )
                await self._emit_metric(
                    name="adk_llm_tokens_input",
                    value=response.usage.input_tokens,
                    unit="tokens",
                    attributes={
                        "agent_name": self.name,
                        "model": self._provider_model_name(provider),
                    },
                )
                await self._emit_metric(
                    name="adk_llm_tokens_output",
                    value=response.usage.output_tokens,
                    unit="tokens",
                    attributes={
                        "agent_name": self.name,
                        "model": self._provider_model_name(provider),
                    },
                )

                budget_failure = self._strict_budget_failure_response(
                    run_context=run_context,
                    stage="llm_stream",
                    usage=response.usage,
                    raw=response.raw,
                )
                if budget_failure is not None:
                    conversation.append(
                        Message(
                            role="assistant",
                            content=budget_failure.content,
                            content_parts=budget_failure.content_parts,
                            sender=self.name,
                        )
                    )
                    state.messages = conversation
                    self._persist_state(state)
                    await self._flush_observers()
                    yield StreamDone(response=budget_failure)
                    return

                if not response.tool_calls:
                    conversation.append(
                        Message(
                            role="assistant",
                            content=response.content,
                            content_parts=response.content_parts,
                            sender=self.name,
                        )
                    )
                    state.messages = conversation
                    commit_started_at = datetime.now(UTC)
                    await self._commit_memories(conversation, response, state, session_id)
                    await self._emit_span(
                        name="Memory.commit",
                        trace_id=trace_id,
                        parent_span_id=root_span_id,
                        started_at=commit_started_at,
                        session_id=session_id,
                        attributes={"strategies": len(self._config.memories)},
                    )
                    await self._emit_span(
                        name="Agent.stream",
                        trace_id=trace_id,
                        span_id=root_span_id,
                        parent_span_id=None,
                        started_at=root_started_at,
                        session_id=session_id,
                    )
                    await self._emit_metric(
                        name="adk_session_memory_messages",
                        value=len(state.messages),
                        unit="count",
                        attributes={"agent_name": self.name},
                    )
                    state = self._persist_state(state)
                    await self._flush_observers()
                    await self._cache_store(
                        provider_messages,
                        response,
                        trace_id,
                        root_span_id,
                        session_id,
                    )
                    if self._config.output_guardrails:
                        checked = await self._check_guardrails(
                            content_text(response.content, response.content_parts),
                            self._config.output_guardrails,
                            role="assistant",
                            session_id=session_id,
                            trace_id=trace_id,
                            request_id=request_id,
                        )
                        if checked != response.content:
                            response = Response(
                                content=checked,
                                content_parts=response.content_parts,
                                tool_calls=response.tool_calls,
                                usage=response.usage,
                                raw=response.raw,
                            )
                    response = await self._secure_output(
                        response,
                        session_id=session_id,
                        trace_id=trace_id,
                        request_id=request_id,
                    )
                    response = self._attach_budget_metadata(
                        response,
                        run_context=run_context,
                        max_tool_rounds=max_tool_rounds,
                    )
                    for part in response.content_parts:
                        if part.type != "text":
                            yield MediaChunk(part=part)
                    yield StreamDone(response=response)
                    return

                conversation.append(
                    Message(
                        role="assistant",
                        content=response.content,
                        content_parts=response.content_parts,
                        tool_calls=response.tool_calls,
                    )
                )
                for tool_call in response.tool_calls:
                    yield ToolStart(
                        call_id=tool_call.id,
                        name=tool_call.name,
                        arguments=tool_call.arguments,
                    )
                    tool_started_at = datetime.now(UTC)
                    agent_logger.info(
                        f"Вызываю инструмент '{tool_call.name}' с аргументами: {tool_call.arguments}"
                    )
                    tool_message = (await self._execute_tool_calls([tool_call]))[0]
                    agent_logger.info(
                        f"Инструмент '{tool_call.name}' вернул: {str(tool_message.content)[:200]}..."
                    )
                    conversation.append(tool_message)
                    await self._emit_span(
                        name=f"Tool.{tool_call.name}",
                        trace_id=trace_id,
                        parent_span_id=root_span_id,
                        started_at=tool_started_at,
                        session_id=session_id,
                        status="error"
                        if tool_message.tool_result and tool_message.tool_result.error
                        else "ok",
                        attributes={
                            "call_id": tool_call.id,
                            "error": tool_message.tool_result.error
                            if tool_message.tool_result
                            else None,
                        },
                    )
                    await self._emit_metric(
                        name="adk_tool_calls_total",
                        value=1.0,
                        unit="count",
                        attributes={
                            "agent_name": self.name,
                            "tool_name": tool_call.name,
                            "status": "error"
                            if tool_message.tool_result and tool_message.tool_result.error
                            else "ok",
                        },
                    )
                    yield ToolEnd(
                        call_id=tool_call.id,
                        result=tool_message.tool_result
                        or ToolResult(call_id=tool_call.id, content=None),
                    )

            state.messages = conversation
            final_response = Response(
                content="",
                tool_calls=[],
                usage=Usage(input_tokens=0, output_tokens=0),
                raw={"reason": "max_tool_rounds_exceeded"},
            )
            await self._emit_span(
                name="Agent.stream",
                trace_id=trace_id,
                span_id=root_span_id,
                parent_span_id=None,
                started_at=root_started_at,
                session_id=session_id,
                status="error",
                attributes={"reason": "max_tool_rounds_exceeded"},
            )
            await self._flush_observers()
            yield StreamDone(response=final_response)
        except Exception as exc:
            await self._emit_span(
                name="Agent.stream",
                trace_id=trace_id,
                span_id=root_span_id,
                parent_span_id=None,
                started_at=root_started_at,
                session_id=session_id,
                status="error",
                attributes={"exception_type": type(exc).__name__},
            )
            await self._flush_observers()
            raise
        finally:
            self._end_run_caches()
            reset_run_context(context_token)
