agent_context.py

from __future__ import annotations

import os
import platform
import subprocess
from dataclasses import dataclass
from datetime import date
from functools import lru_cache
from pathlib import Path

from .agent_plugin_cache import load_plugin_cache_summary
from .account_runtime import AccountRuntime
from .ask_user_runtime import AskUserRuntime
from .config_runtime import ConfigRuntime
from .hook_policy import HookPolicyRuntime
from .lsp_runtime import LSPRuntime
from .mcp_runtime import MCPRuntime
from .plan_runtime import PlanRuntime
from .plugin_runtime import PluginRuntime
from .remote_runtime import RemoteRuntime
from .remote_trigger_runtime import RemoteTriggerRuntime
from .search_runtime import SearchRuntime
from .task_runtime import TaskRuntime
from .team_runtime import TeamRuntime
from .workflow_runtime import WorkflowRuntime
from .worktree_runtime import WorktreeRuntime
from .agent_types import AgentRuntimeConfig

MAX_STATUS_CHARS = 2000
MAX_MEMORY_CHARACTER_COUNT = 40000
MEMORY_INSTRUCTION_PROMPT = (
    'Codebase and user instructions are shown below. Be sure to adhere to '
    'these instructions. IMPORTANT: These instructions override default '
    'behavior when they directly apply to the task.'
)

_SYSTEM_PROMPT_INJECTION: str | None = None


@dataclass(frozen=True)
class AgentContextSnapshot:
    cwd: Path
    shell: str
    platform_name: str
    os_version: str
    current_date: str
    is_git_repo: bool
    is_git_worktree: bool
    scratchpad_directory: str | None
    additional_working_directories: tuple[str, ...]
    user_context: dict[str, str]
    system_context: dict[str, str]


def clear_context_caches() -> None:
    _get_git_status_cached.cache_clear()
    _get_system_context_cached.cache_clear()
    _get_user_context_cached.cache_clear()


def get_system_prompt_injection() -> str | None:
    return _SYSTEM_PROMPT_INJECTION


def set_system_prompt_injection(value: str | None) -> None:
    global _SYSTEM_PROMPT_INJECTION
    _SYSTEM_PROMPT_INJECTION = value
    clear_context_caches()


def build_context_snapshot(
    runtime_config: AgentRuntimeConfig,
    *,
    scratchpad_directory: Path | None = None,
) -> AgentContextSnapshot:
    cwd = runtime_config.cwd.resolve()
    additional_dirs = tuple(
        str(path.resolve()) for path in runtime_config.additional_working_directories
    )
    return AgentContextSnapshot(
        cwd=cwd,
        shell=os.environ.get('SHELL', 'unknown'),
        platform_name=platform.system().lower() or os.name,
        os_version=_get_os_version(),
        current_date=date.today().isoformat(),
        is_git_repo=_is_git_repo(cwd),
        is_git_worktree=_is_git_worktree(cwd),
        scratchpad_directory=(
            str(scratchpad_directory.resolve()) if scratchpad_directory is not None else None
        ),
        additional_working_directories=additional_dirs,
        user_context=get_user_context(
            cwd,
            additional_dirs,
            runtime_config.disable_claude_md_discovery,
            scratchpad_directory=scratchpad_directory,
        ),
        system_context=get_system_context(cwd, scratchpad_directory=scratchpad_directory),
    )


def get_git_status(cwd: Path) -> str | None:
    return _get_git_status_cached(str(cwd.resolve()))


def get_system_context(
    cwd: Path,
    *,
    scratchpad_directory: Path | None = None,
) -> dict[str, str]:
    scratchpad = str(scratchpad_directory.resolve()) if scratchpad_directory is not None else ''
    return dict(_get_system_context_cached(str(cwd.resolve()), scratchpad))


def get_user_context(
    cwd: Path,
    additional_working_directories: tuple[str, ...] = (),
    disable_claude_md_discovery: bool = False,
    scratchpad_directory: Path | None = None,
) -> dict[str, str]:
    normalized_dirs = tuple(
        str(Path(path).resolve()) for path in additional_working_directories
    )
    return dict(
        _get_user_context_cached(
            str(cwd.resolve()),
            normalized_dirs,
            disable_claude_md_discovery,
            str(scratchpad_directory.resolve()) if scratchpad_directory is not None else '',
        )
    )


def render_context_report(snapshot: AgentContextSnapshot, model: str) -> str:
    lines = [
        '# Context',
        '',
        '## Environment',
        f'- Primary working directory: {snapshot.cwd}',
        f'- Model: {model}',
        f'- Shell: {Path(snapshot.shell).name or snapshot.shell}',
        f'- Platform: {snapshot.platform_name}',
        f'- OS Version: {snapshot.os_version}',
        f'- Is a git repository: {snapshot.is_git_repo}',
        f'- Is a git worktree: {snapshot.is_git_worktree}',
        f'- Current date: {snapshot.current_date}',
    ]
    if snapshot.scratchpad_directory:
        lines.append(f'- Scratchpad directory: {snapshot.scratchpad_directory}')
    if snapshot.additional_working_directories:
        lines.extend(
            [
                '',
                '## Additional Working Directories',
                *[f'- {path}' for path in snapshot.additional_working_directories],
            ]
        )
    if snapshot.user_context:
        lines.extend(['', '## User Context'])
        for key, value in snapshot.user_context.items():
            lines.extend([f'### {key}', value, ''])
        while lines and not lines[-1]:
            lines.pop()
    if snapshot.system_context:
        lines.extend(['', '## System Context'])
        for key, value in snapshot.system_context.items():
            lines.extend([f'### {key}', value, ''])
        while lines and not lines[-1]:
            lines.pop()
    return '\n'.join(lines)


@lru_cache(maxsize=32)
def _get_system_context_cached(cwd: str, scratchpad_directory: str) -> dict[str, str]:
    context: dict[str, str] = {}
    git_status = _get_git_status_cached(cwd)
    if git_status is not None:
        context['gitStatus'] = git_status
    injection = get_system_prompt_injection()
    if injection:
        context['cacheBreaker'] = f'[CACHE_BREAKER: {injection}]'
    if scratchpad_directory:
        context['scratchpadDirectory'] = scratchpad_directory
    return context


@lru_cache(maxsize=32)
def _get_user_context_cached(
    cwd: str,
    additional_working_directories: tuple[str, ...],
    disable_claude_md_discovery: bool,
    scratchpad_directory: str,
) -> dict[str, str]:
    context: dict[str, str] = {
        'currentDate': f"Today's date is {date.today().isoformat()}.",
    }
    if scratchpad_directory:
        context['scratchpad'] = (
            'Use this session-specific scratchpad directory for temporary files instead '
            f'of /tmp when you need throwaway workspace: {scratchpad_directory}'
        )
    if disable_claude_md_discovery:
        return context

    memory_bundle = _load_memory_bundle(Path(cwd), additional_working_directories)
    if memory_bundle:
        context['claudeMd'] = memory_bundle
    plugin_cache = load_plugin_cache_summary(Path(cwd), additional_working_directories)
    if plugin_cache:
        context['pluginCache'] = plugin_cache
    plugin_runtime = PluginRuntime.from_workspace(Path(cwd), additional_working_directories)
    if plugin_runtime.manifests:
        context['pluginRuntime'] = plugin_runtime.render_summary()
    hook_policy_runtime = HookPolicyRuntime.from_workspace(Path(cwd), additional_working_directories)
    if hook_policy_runtime.manifests:
        context['hookPolicy'] = hook_policy_runtime.render_summary()
        managed_settings = hook_policy_runtime.managed_settings()
        if managed_settings:
            context['managedSettings'] = '\n'.join(
                f'{key}={value}'
                for key, value in sorted(managed_settings.items())
            )
        safe_env = hook_policy_runtime.safe_env()
        if safe_env:
            context['safeEnv'] = '\n'.join(
                f'{key}={value}'
                for key, value in sorted(safe_env.items())
            )
        context['trustMode'] = (
            'Workspace trust mode: trusted'
            if hook_policy_runtime.is_trusted()
            else 'Workspace trust mode: untrusted'
        )
    mcp_runtime = MCPRuntime.from_workspace(Path(cwd), additional_working_directories)
    if mcp_runtime.resources:
        context['mcpRuntime'] = mcp_runtime.render_summary()
    remote_runtime = RemoteRuntime.from_workspace(Path(cwd), additional_working_directories)
    if remote_runtime.has_remote_config():
        context['remoteRuntime'] = remote_runtime.render_summary()
    remote_trigger_runtime = RemoteTriggerRuntime.from_workspace(
        Path(cwd),
        additional_working_directories,
    )
    if remote_trigger_runtime.has_state():
        context['remoteTriggerRuntime'] = remote_trigger_runtime.render_summary()
    search_runtime = SearchRuntime.from_workspace(Path(cwd), additional_working_directories)
    if search_runtime.has_search_runtime():
        context['searchRuntime'] = search_runtime.render_summary()
    account_runtime = AccountRuntime.from_workspace(Path(cwd), additional_working_directories)
    if account_runtime.has_account_state():
        context['accountRuntime'] = account_runtime.render_summary()
    ask_user_runtime = AskUserRuntime.from_workspace(Path(cwd), additional_working_directories)
    if ask_user_runtime.has_state():
        context['askUserRuntime'] = ask_user_runtime.render_summary()
    config_runtime = ConfigRuntime.from_workspace(Path(cwd))
    if config_runtime.has_config():
        context['configRuntime'] = config_runtime.render_summary()
    lsp_runtime = LSPRuntime.from_workspace(Path(cwd), additional_working_directories)
    if lsp_runtime.has_lsp_support():
        context['lspRuntime'] = lsp_runtime.render_summary()
    plan_runtime = PlanRuntime.from_workspace(Path(cwd))
    if plan_runtime.steps:
        context['planRuntime'] = plan_runtime.render_summary()
    task_runtime = TaskRuntime.from_workspace(Path(cwd))
    if task_runtime.tasks:
        context['taskRuntime'] = task_runtime.render_summary()
    team_runtime = TeamRuntime.from_workspace(Path(cwd), additional_working_directories)
    if team_runtime.has_team_state():
        context['teamRuntime'] = team_runtime.render_summary()
    workflow_runtime = WorkflowRuntime.from_workspace(Path(cwd), additional_working_directories)
    if workflow_runtime.has_workflows():
        context['workflowRuntime'] = workflow_runtime.render_summary()
    worktree_runtime = WorktreeRuntime.from_workspace(Path(cwd))
    if worktree_runtime.repo_root is not None or worktree_runtime.has_state():
        context['worktreeRuntime'] = worktree_runtime.render_summary()
    return context


@lru_cache(maxsize=32)
def _get_git_status_cached(cwd: str) -> str | None:
    root = Path(cwd)
    if not _is_git_repo(root):
        return None

    branch = _run_command(['git', 'branch', '--show-current'], root)
    main_branch = _detect_default_branch(root)
    status = _run_command(['git', '--no-optional-locks', 'status', '--short'], root) or ''
    log = _run_command(['git', '--no-optional-locks', 'log', '--oneline', '-n', '5'], root) or '(none)'
    user_name = _run_command(['git', 'config', 'user.name'], root)

    if len(status) > MAX_STATUS_CHARS:
        status = (
            status[:MAX_STATUS_CHARS]
            + '\n... (truncated because it exceeds 2k characters. Use bash for full git status.)'
        )

    parts = [
        'This is the git status at the start of the conversation. It is a snapshot and does not update automatically during the run.',
        f'Current branch: {branch or "(unknown)"}',
        f'Main branch: {main_branch or "(unknown)"}',
    ]
    if user_name:
        parts.append(f'Git user: {user_name}')
    parts.extend(
        [
            f'Status:\n{status or "(clean)"}',
            f'Recent commits:\n{log}',
        ]
    )
    return '\n\n'.join(parts)


def _load_memory_bundle(cwd: Path, additional_working_directories: tuple[str, ...]) -> str | None:
    discovered: list[Path] = []
    seen: set[Path] = set()

    for candidate in _discover_global_memory_files():
        _remember_path(candidate, discovered, seen)

    for directory in _walk_upwards(cwd):
        for candidate in _discover_memory_files_for_directory(directory):
            _remember_path(candidate, discovered, seen)

    for raw_path in additional_working_directories:
        for candidate in _discover_memory_files_for_directory(Path(raw_path)):
            _remember_path(candidate, discovered, seen)

    if not discovered:
        return None

    blocks = [MEMORY_INSTRUCTION_PROMPT]
    for path in discovered:
        try:
            content = path.read_text(encoding='utf-8', errors='replace').strip()
        except OSError:
            continue
        if not content:
            continue
        if len(content) > MAX_MEMORY_CHARACTER_COUNT:
            content = (
                content[:MAX_MEMORY_CHARACTER_COUNT]
                + '\n... (truncated because it exceeds the memory size limit)'
            )
        blocks.append(f'## {path}\n{content}')
    if len(blocks) == 1:
        return None
    return '\n\n'.join(blocks)


def _discover_global_memory_files() -> list[Path]:
    home_memory = Path.home() / '.claude' / 'CLAUDE.md'
    return [home_memory] if home_memory.is_file() else []


def _discover_memory_files_for_directory(directory: Path) -> list[Path]:
    files: list[Path] = []
    for candidate in (
        directory / 'CLAUDE.md',
        directory / '.claude' / 'CLAUDE.md',
        directory / 'CLAUDE.local.md',
    ):
        if candidate.is_file():
            files.append(candidate.resolve())

    rules_dir = directory / '.claude' / 'rules'
    if rules_dir.is_dir():
        files.extend(
            path.resolve()
            for path in sorted(rules_dir.glob('*.md'))
            if path.is_file()
        )
    return files


def _walk_upwards(cwd: Path) -> list[Path]:
    parents = list(cwd.resolve().parents)
    parents.reverse()
    return [*parents, cwd.resolve()]


def _remember_path(path: Path, discovered: list[Path], seen: set[Path]) -> None:
    resolved = path.resolve()
    if resolved in seen:
        return
    seen.add(resolved)
    discovered.append(resolved)


def _detect_default_branch(cwd: Path) -> str | None:
    origin_head = _run_command(
        ['git', 'symbolic-ref', '--quiet', '--short', 'refs/remotes/origin/HEAD'],
        cwd,
    )
    if origin_head and '/' in origin_head:
        return origin_head.split('/', 1)[1]

    for candidate in ('main', 'master'):
        try:
            completed = subprocess.run(
                ['git', 'show-ref', '--verify', f'refs/heads/{candidate}'],
                cwd=cwd,
                capture_output=True,
                text=True,
                timeout=2.0,
                check=False,
            )
        except OSError:
            return None
        if completed.returncode == 0:
            return candidate
    return None


@lru_cache(maxsize=32)
def _is_git_repo(cwd: Path) -> bool:
    if (cwd / '.git').exists():
        return True
    try:
        completed = subprocess.run(
            ['git', 'rev-parse', '--is-inside-work-tree'],
            cwd=cwd,
            capture_output=True,
            text=True,
            timeout=2.0,
            check=False,
        )
    except OSError:
        return False
    return completed.returncode == 0 and completed.stdout.strip() == 'true'


def _is_git_worktree(cwd: Path) -> bool:
    if not _is_git_repo(cwd):
        return False
    git_dir = _run_command(['git', 'rev-parse', '--git-dir'], cwd)
    git_common_dir = _run_command(['git', 'rev-parse', '--git-common-dir'], cwd)
    return bool(git_dir and git_common_dir and git_dir != git_common_dir)


def _run_command(command: list[str], cwd: Path) -> str | None:
    try:
        completed = subprocess.run(
            command,
            cwd=cwd,
            capture_output=True,
            text=True,
            timeout=2.0,
            check=False,
        )
    except OSError:
        return None
    if completed.returncode != 0:
        return None
    output = completed.stdout.strip()
    return output or None


def _get_os_version() -> str:
    system = platform.system()
    release = platform.release()
    if system and release:
        return f'{system} {release}'
    return platform.platform()