dev-puppeteer/my-deepagent/src/my_deepagent/session.py

"""Build a deepagents CompiledStateGraph from a Persona + run context.

Connects:
  - Persona (config) -> deepagents.create_deep_agent(...)
  - OpenRouter (model="openrouter:...") -> ChatOpenAI(base_url=openrouter)
  - Workspace dir -> LocalShellBackend (filesystem + shell execution)
  - Persona.permissions + DEFAULT_DENY -> deepagents.FilesystemPermission list
  - Subagents -> deepagents.SubAgent TypedDict list
  - Middleware list -> passed to create_deep_agent
"""

from __future__ import annotations

import os
from pathlib import Path
from typing import Any, Literal
from uuid import UUID

from deepagents import FilesystemPermission, SubAgent, create_deep_agent
from deepagents.backends import (
    CompositeBackend,
    FilesystemBackend,
    LocalShellBackend,
    StateBackend,
)
from langchain_openai import ChatOpenAI

from .config import Config
from .errors import MyDeepAgentError
from .persona import FilesystemPermissionSpec, Persona, PersonaSubagent

DEFAULT_DENY_PATHS: tuple[str, ...] = (
    "/.env*",
    "/**/*.env*",
    "/**/*token*",
    "/**/*secret*",
    "/**/*credential*",
    "/**/*.pem",
    "/**/*.key",
    "/.ssh/**",
    "/.aws/**",
    "/.config/gcloud/**",
    "/.kube/**",
    "/.gnupg/**",
)


# Mapping from our richer operation set (read/write/edit/ls) to the deepagents
# binary set (read/write). deepagents treats ls/grep/glob as read-side and
# write_file/edit_file as write-side internally, so this collapse is safe.
_OP_MAP: dict[str, Literal["read", "write"]] = {
    "read": "read",
    "write": "write",
    "edit": "write",
    "ls": "read",
}


def _map_operations(ops: tuple[str, ...] | list[str]) -> list[Literal["read", "write"]]:
    """Deduplicate-preserve-order mapping of our ops to deepagents ops."""
    seen: set[str] = set()
    out: list[Literal["read", "write"]] = []
    for op in ops:
        mapped = _OP_MAP[op]
        if mapped not in seen:
            seen.add(mapped)
            out.append(mapped)
    return out


def default_safety_permissions() -> list[FilesystemPermission]:
    """Default-allow paths and deny secret-bearing paths.

    Returned permissions are evaluated in order; first match wins.
    Allow comes first so reads/writes to the worktree succeed by default;
    then explicit denies block the secret patterns no matter what.
    """
    return [
        FilesystemPermission(
            operations=["read", "write"],
            paths=["/**"],
            mode="allow",
        ),
        FilesystemPermission(
            operations=["read", "write"],
            paths=list(DEFAULT_DENY_PATHS),
            mode="deny",
        ),
    ]


def _spec_to_permission(spec: FilesystemPermissionSpec) -> FilesystemPermission:
    """Convert pydantic FilesystemPermissionSpec to deepagents FilesystemPermission.

    Our schema accepts {read, write, edit, ls} for human-readable yaml. deepagents
    collapses these to {read, write} internally; we apply the same collapse here.
    """
    return FilesystemPermission(
        operations=_map_operations(spec.operations),
        paths=list(spec.paths),
        mode=spec.mode,
    )


def _subagent_to_dict(sub: PersonaSubagent) -> SubAgent:
    """Convert PersonaSubagent -> deepagents SubAgent TypedDict.

    Only includes optional keys when set; deepagents inherits defaults from the parent
    agent when a subagent omits ``tools`` / ``model`` / ``permissions`` / ``interrupt_on``.
    """
    out: dict[str, Any] = {
        "name": sub.name,
        "description": sub.description,
        "system_prompt": sub.system_prompt,
    }
    if sub.allowed_tools:
        out["tools"] = list(sub.allowed_tools)
    if sub.model is not None:
        out["model"] = sub.model
    if sub.permissions:
        out["permissions"] = [_spec_to_permission(p) for p in sub.permissions]
    if sub.interrupt_on:
        out["interrupt_on"] = sub.interrupt_on
    return out  # type: ignore[return-value]  # TypedDict construction from dict literal


def _resolve_openrouter_api_key(config: Config) -> str:
    """Pull the OpenRouter API key from config -> env -> error.

    Priority: config.openrouter_api_key -> MYDEEPAGENT_OPENROUTER_API_KEY -> OPENROUTER_API_KEY.
    """
    if config.openrouter_api_key:
        return config.openrouter_api_key
    env_key = os.environ.get("MYDEEPAGENT_OPENROUTER_API_KEY") or os.environ.get(
        "OPENROUTER_API_KEY"
    )
    if env_key:
        return env_key
    raise MyDeepAgentError.human_required(
        "backend_auth_failed",
        message="OpenRouter API key is not configured",
        recovery_hint=(
            "set MYDEEPAGENT_OPENROUTER_API_KEY in .env or run `mydeepagent login openrouter`"
        ),
    )


def resolve_model_instance(
    persona: Persona, config: Config, model_override: str | None = None
) -> Any:
    """Persona -> langchain BaseChatModel instance or 'provider:model' string.

    For ``openrouter:`` prefix, returns a ``ChatOpenAI`` with ``base_url=openrouter``.
    For other providers (``anthropic:``, ``openai:``, ``google:``), returns the string as-is
    so that deepagents' ``init_chat_model`` resolves it via the matching integration package.
    """
    model_spec = model_override or persona.model
    if model_spec.startswith("openrouter:"):
        params = persona.model_params
        return ChatOpenAI(
            model=model_spec.removeprefix("openrouter:"),
            api_key=_resolve_openrouter_api_key(config),
            base_url=config.openrouter_base_url,
            max_tokens=params.get("max_tokens", 4096),
            temperature=params.get("temperature", 0.2),
            top_p=params.get("top_p", 1.0),
        )
    return model_spec


def build_backend(persona: Persona, root_dir: Path) -> Any:
    """Persona.deepagents_backend -> concrete deepagents backend instance.

    Returns:
        LocalShellBackend for "local_shell" (filesystem + shell execute, the default).
        FilesystemBackend for "filesystem" (filesystem only, no shell).
        None for "state" (deepagents default StateBackend, in-process state only).
        CompositeBackend for "composite" (local_shell + state-backed /memories/ namespace).

    Raises:
        MyDeepAgentError(fatal, config_invalid) for unknown backend identifiers
        or "langsmith" which is reserved for a future milestone.
    """
    name = persona.deepagents_backend
    if name == "local_shell":
        return LocalShellBackend(
            root_dir=str(root_dir),
            virtual_mode=False,
            timeout=120,
            max_output_bytes=100_000,
            inherit_env=False,
        )
    if name == "filesystem":
        return FilesystemBackend(root_dir=str(root_dir), virtual_mode=False, max_file_size_mb=10)
    if name == "state":
        return None  # deepagents default StateBackend
    if name == "composite":
        return CompositeBackend(
            default=LocalShellBackend(root_dir=str(root_dir), virtual_mode=False),
            routes={"/memories/": StateBackend()},
        )
    raise MyDeepAgentError.fatal(
        "config_invalid",
        message=f"unsupported deepagents_backend: {name!r}",
        recovery_hint="use one of: local_shell, filesystem, state, composite",
    )


def build_agent(
    persona: Persona,
    config: Config,
    *,
    root_dir: Path,
    middleware: list[Any] | None = None,
    checkpointer: Any | None = None,
    run_id: UUID | None = None,
    phase_key: str | None = None,
    model_override: str | None = None,
) -> Any:
    """Construct a deepagents CompiledStateGraph for the given persona.

    Returns a CompiledStateGraph. Caller invokes via
    ``agent.invoke / ainvoke / astream / astream_events`` with ``{"messages": [...]}`` input.

    deepagents 0.6.1 limitation: FilesystemPermission is rejected when the backend
    implements SandboxBackendProtocol (e.g. LocalShellBackend). SafetyShellMiddleware
    enforces path + destructive-command safety in those cases instead.
    """
    from .middleware.safety import SafetyShellMiddleware

    model = resolve_model_instance(persona, config, model_override)
    backend = build_backend(persona, root_dir)

    # SafetyShellMiddleware is always first; caller-supplied middleware appends.
    all_middleware: list[Any] = [SafetyShellMiddleware()]
    if middleware:
        all_middleware.extend(middleware)

    subagents: list[SubAgent] = [_subagent_to_dict(s) for s in persona.subagents]

    kwargs: dict[str, Any] = {
        "model": model,
        "system_prompt": persona.system_prompt,
        "middleware": all_middleware,
    }
    if backend is not None:
        kwargs["backend"] = backend

    # deepagents 0.6.1: FilesystemPermission + SandboxBackendProtocol backend raises
    # NotImplementedError. Skip permissions kwarg for local_shell; SafetyShellMiddleware
    # handles path enforcement instead. Other backends (state, filesystem, composite)
    # still use the deepagents permissions system.
    use_permissions = persona.deepagents_backend != "local_shell"
    if use_permissions:
        permissions: list[FilesystemPermission] = [
            *(_spec_to_permission(p) for p in persona.permissions),
            *default_safety_permissions(),
        ]
        kwargs["permissions"] = permissions

    if persona.allowed_tools:
        kwargs["tools"] = list(persona.allowed_tools)
    if subagents:
        kwargs["subagents"] = subagents
    if persona.interrupt_on:
        kwargs["interrupt_on"] = persona.interrupt_on
    if checkpointer is not None:
        kwargs["checkpointer"] = checkpointer
    if persona.skills:
        kwargs["skills"] = list(persona.skills)
    if persona.memory_files:
        kwargs["memory"] = list(persona.memory_files)

    return create_deep_agent(**kwargs)