Repository: ErikBjare/gptme
File: gptme/hooks/agents_md_inject.py (Python)

Code before:
loaded yet. If found, their contents are injected as system messages.

This extends the tree-walking AGENTS.md loading from prompt_workspace() (which runs
at startup) to also work mid-session when the CWD changes.

See: https://github.com/gptme/gptme/issues/1513
"""

import logging
import os
from collections.abc import Generator
from contextvars import ContextVar
from pathlib import Path
from typing import Any

from ..hooks import HookType, StopPropagation, register_hook
from ..logmanager import Log
from ..message import Message
from ..prompts import ALWAYS_LOAD_FILES

logger = logging.getLogger(__name__)

# Track which agent files have already been loaded (by resolved path).
# Seeded at session start from the initial workspace tree walk,
# then updated as new files are discovered on CWD changes.
_loaded_agent_files: set[str] = set()

# Track CWD before tool execution (same pattern as cwd_tracking.py)
_cwd_before_var: ContextVar[str | None] = ContextVar(
    "agents_md_cwd_before", default=None
)


def _find_agent_files_in_tree(directory: Path) -> list[Path]:
    """Find AGENTS.md/CLAUDE.md/GEMINI.md files from home down to the given directory.

    Walks from home → directory (most general first), checking each directory
    for agent instruction files. Only returns files not already in _loaded_agent_files.
    """
    new_files: list[Path] = []
    home_dir = Path.home().resolve()
    target = directory.resolve()

    # Build path from home to target
    dirs_to_check: list[Path] = []
    current = target
    while current != current.parent:
        dirs_to_check.append(current)
        if current == home_dir:
            break
        current = current.parent

Review comment (suggestion):
Should at least be a ContextVar, probably also integrate properly with https://github.com/gptme/gptme/blob/a8ac588abcc03abb9c417880b9fa22737f582f6c/gptme/prompts.py#L523-L528

(maybe by initializing a _loaded_agent_files ContextVar there in prompts.md that this hook can reuse, with comments that clearly connect the two)

Also, I notice that this hook itself manages/detects cwd_changed, but it might make sense to turn the "cwd changed" into a proper hook type itself (leave a comment in the code and/or make an issue).

Code after:
loaded yet. If found, their contents are injected as system messages.

This extends the tree-walking AGENTS.md loading from prompt_workspace() (which runs
at startup) to also work mid-session when the CWD changes.

The set of already-loaded files is shared with prompt_workspace() via the
_loaded_agent_files_var ContextVar defined in prompts.py, which seeds it at startup.

See: https://github.com/gptme/gptme/issues/1513
"""

import logging
import os
from collections.abc import Generator
from contextvars import ContextVar
from pathlib import Path
from typing import Any

from ..hooks import HookType, StopPropagation, register_hook
from ..logmanager import Log
from ..message import Message
from ..prompts import _loaded_agent_files_var, find_agent_files_in_tree

logger = logging.getLogger(__name__)

# Track CWD before tool execution (same pattern as cwd_tracking.py)
# TODO: Consider making "cwd changed" a proper hook type (e.g. HookType.CWD_CHANGED)
#       so that hooks can subscribe to directory changes directly instead of each hook
#       independently comparing pre/post CWD. See: https://github.com/gptme/gptme/issues/1521
_cwd_before_var: ContextVar[str | None] = ContextVar(
    "agents_md_cwd_before", default=None
)


def _get_loaded_files() -> set[str]:
    """Get (or lazily initialize) the loaded agent files set for this context.

    Normally populated by prompt_workspace() at session start. If called before
    that (e.g., in tests), initializes to an empty set.
    """
    files = _loaded_agent_files_var.get()
    if files is None:
        files = set()
        _loaded_agent_files_var.set(files)
    return files


def pre_execute(
    log: Log, workspace: Path | None, tool_use: Any
) -> Generator[Message | StopPropagation, None, None]:
    """Store CWD before tool execution."""