wuyang6
466 lines
16 KiB
Python
466 lines
16 KiB
Python
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()
|
|
if scratchpad_directory:
|
|
scratchpad_path = Path(scratchpad_directory)
|
|
plan_runtime = PlanRuntime.from_storage_path(scratchpad_path / 'plan_runtime.json')
|
|
if plan_runtime.steps:
|
|
context['planRuntime'] = plan_runtime.render_summary()
|
|
task_runtime = TaskRuntime.from_storage_path(scratchpad_path / 'task_runtime.json')
|
|
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()
|