"""Base agent wrapper component.""" from abc import abstractmethod from collections.abc import AsyncGenerator from importlib import metadata from pathlib import Path from typing import Any, ClassVar, TYPE_CHECKING from pydantic import BaseModel from ..base_component import BaseComponent from ..outbound_proxy import BaseOutboundProxy from ...enumeration import ChunkEnum, ComponentEnum from ...schema import StreamChunk, TokenUsage from ...utils import global_counter_add_many if TYPE_CHECKING: from ..job.base_job import BaseJob class BaseAgentWrapper(BaseComponent): """Abstract base for agent wrapper components with swappable backends.""" component_type = ComponentEnum.AGENT_WRAPPER SDK_PACKAGE: ClassVar[str | None] = None TOKEN_COUNTER_PREFIX: ClassVar[str] = "__token_counter" def __init__( self, cwd: str | Path | None = None, project_path: str | Path | None = None, **kwargs, ) -> None: super().__init__(**kwargs) self._cwd = cwd self._project_path = project_path self.outbound_proxy = self.bind( "default", BaseOutboundProxy, optional=True, ) if self.SDK_PACKAGE: try: sdk_version = metadata.version(self.SDK_PACKAGE) except metadata.PackageNotFoundError: sdk_version = "unknown" self.logger.info(f"Agent SDK name={self.name} package={self.SDK_PACKAGE} version={sdk_version}") @property def cwd(self) -> Path: """Working directory shared by the agent's shell and file tools. Defaults to the project root. Override via the ``cwd`` init argument; a relative value resolves against the workspace root. """ if not self._cwd: return self.project_path cwd = Path(self._cwd) return cwd if cwd.is_absolute() else (self.workspace_path / cwd) def set_system_prompt(self, prompt: str) -> "BaseAgentWrapper": """Set the agent's system prompt. Returns self for chaining.""" self.kwargs["system_prompt"] = prompt return self def add_job_tools(self, job_tools: list[str]) -> "BaseAgentWrapper": """Append job names as tools to the agent. Returns self for chaining.""" self.kwargs.setdefault("job_tools", []).extend(job_tools) return self def add_skills(self, skills: list[str] | str) -> "BaseAgentWrapper": """Set agent skill names. Returns self for chaining.""" self.kwargs["skills"] = skills return self @property def project_path(self) -> Path: """Project root containing shared assets such as skills. A relative configured path resolves from the workspace so applications can keep runtime data in a subdirectory such as ``.reme`` while loading project assets from its parent. The workspace remains the default for backward compatibility. """ if not self._project_path: return self.workspace_path project_path = Path(self._project_path).expanduser() if not project_path.is_absolute(): project_path = self.workspace_path / project_path return project_path.resolve(strict=False) @property def project_skills_root(self) -> Path: """Project-level skills directory shared by agent backends.""" return self.project_path / "skills" def _resolve_project_skills(self, skills: list[str] | str | None) -> dict[str, Path]: """Resolve selected skill names to validated project directories.""" if skills is None: return {} if skills == "all": if not self.project_skills_root.is_dir(): raise FileNotFoundError(f"Project skills directory not found: {self.project_skills_root}") names = sorted( path.name for path in self.project_skills_root.iterdir() if path.is_dir() and (path / "SKILL.md").is_file() ) else: names = [skills] if isinstance(skills, str) else list(skills) names = list(dict.fromkeys(names)) sources: dict[str, Path] = {} for name in names: if not name or Path(name).name != name or name in {".", ".."}: raise ValueError(f"Invalid skill name: {name!r}") source = self.project_skills_root / name if not source.is_dir(): raise FileNotFoundError(f"Skill directory not found: {source}") if not (source / "SKILL.md").is_file(): raise FileNotFoundError(f"Skill '{name}' is missing SKILL.md: {source}") sources[name] = source return sources @property def subprocess_environment(self) -> dict[str, str]: """Configured environment variables for child agent processes.""" if self.app_context is None: return {} return self.app_context.app_config.environment @property def command_proxy_environment(self) -> dict[str, str]: """Managed proxy variables for agent command tools only.""" if not isinstance(self.outbound_proxy, BaseOutboundProxy): return {} return self.outbound_proxy.merge_environment() @property def bash_environment(self) -> dict[str, str]: """Configured command environment with the managed proxy applied last.""" if not isinstance(self.outbound_proxy, BaseOutboundProxy): return dict(self.subprocess_environment) return self.outbound_proxy.merge_environment(self.subprocess_environment) def set_output_schema(self, schema: dict | type[BaseModel]) -> "BaseAgentWrapper": """Set a JSON schema for structured output. Accepts dict or BaseModel class. Returns self for chaining.""" self.kwargs["output_schema"] = self._normalize_output_schema(schema) return self @staticmethod def _normalize_output_schema(schema: Any) -> dict | None: """Return a JSON-serializable output schema shared by every backend.""" if isinstance(schema, type) and issubclass(schema, BaseModel): return schema.model_json_schema() if schema is None or isinstance(schema, dict): return schema raise TypeError("output_schema must be a JSON schema dict or BaseModel class") @staticmethod def _resolve_injected_job_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]: """Collect server-owned kwargs merged into every job tool call. ``injected_job_kwargs`` values are enforced constraints added by the wrapper after receiving the model's tool arguments; the model can neither see nor override them. ``tool_context_id`` keeps its dedicated option but is carried through the same mechanism. """ injected = dict(kwargs.get("injected_job_kwargs") or {}) if tool_context_id := kwargs.get("tool_context_id"): injected["tool_context_id"] = tool_context_id return injected @staticmethod def _merge_injected_job_kwargs(model_kwargs: dict[str, Any], injected: dict[str, Any]) -> dict[str, Any]: """Merge server-owned kwargs over model tool arguments, rejecting conflicts. Silently letting model values win would make injected constraints bypassable, so any overlap is an explicit error. """ if conflicts := sorted(injected.keys() & model_kwargs.keys()): names = ", ".join(conflicts) raise ValueError(f"injected tool arguments cannot be provided by the model: {names}") return {**model_kwargs, **injected} @staticmethod def _strip_injected_parameters(parameters: dict | None, injected: dict[str, Any]) -> dict | None: """Hide injected keys from the tool parameter schema exposed to the model.""" if not parameters or not injected: return parameters parameters = dict(parameters) if "properties" in parameters: parameters["properties"] = { name: schema for name, schema in parameters["properties"].items() if name not in injected } if "required" in parameters: parameters["required"] = [name for name in parameters["required"] if name not in injected] return parameters def _resolve_job_tools(self, job_tools: list[str]) -> list["BaseJob"]: """Resolve job name strings to BaseJob instances via app_context.""" if not job_tools: return [] if self.app_context is None: raise RuntimeError("Cannot resolve job_tools without an app_context") resolved: list["BaseJob"] = [] for name in job_tools: if (job := self.app_context.jobs.get(name)) is None: raise KeyError(f"Job '{name}' not found in app_context.jobs") resolved.append(job) return resolved def _merged_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: """Merge component defaults with call-time kwargs; call-time values win.""" merged = {**self.kwargs, **kwargs} if "output_schema" in merged: merged["output_schema"] = self._normalize_output_schema(merged["output_schema"]) return merged def _merged_stream_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: """Merge stream options and reject unsupported structured output.""" merged = self._merged_kwargs(kwargs) if merged.get("output_schema") is not None: raise NotImplementedError("Structured output is not supported by reply_stream()") return merged @staticmethod def _chunk(chunk_type: ChunkEnum = ChunkEnum.CONTENT, **kwargs: Any) -> StreamChunk: """Create a StreamChunk with a short backend-friendly call site.""" return StreamChunk(chunk_type=chunk_type, **kwargs) def _record_token_usage(self, usage: TokenUsage) -> None: """Add one completed invocation to the application token tree.""" if self.app_context is None: return counters = getattr(self.app_context, "metadata", None) if not isinstance(counters, dict): return prefix = (self.TOKEN_COUNTER_PREFIX, self.name) global_counter_add_many( counters, { (*prefix, "input_tokens"): usage.input_tokens, (*prefix, "output_tokens"): usage.output_tokens, (*prefix, "total_tokens"): usage.total_tokens, }, ) @abstractmethod async def reply(self, inputs: Any, **kwargs) -> dict: """Send inputs to the agent and return a dict with session_id and last_message.""" async def compact_session(self, session_id: str) -> None: """Request compaction of one persisted agent session.""" raise NotImplementedError(f"{type(self).__name__} does not support session compaction") async def reply_stream(self, inputs: Any, **kwargs) -> AsyncGenerator[StreamChunk, None]: """Stream agent events as unified StreamChunk objects."""