flash-context 0.1.0__py3-none-any.whl

flash_context/mcp_server.py

@@ -0,0 +1,89 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from importlib import import_module
+from pathlib import Path
+from typing import Any
+
+from flash_context.agent import explore
+from flash_context.config import Settings, SettingsOverrides, load_settings
+from flash_context.types import ExploreRequest, ExplorerError, ExploreResult
+
+CoreFn = Callable[[ExploreRequest, Settings], ExploreResult]
+
+
+def explore_repository_handler(
+    *,
+    query: str,
+    repo_root: str | None = None,
+    max_turns: int | None = None,
+    citation: bool | None = None,
+    settings: Settings | None = None,
+    core: CoreFn = explore,
+) -> dict[str, object]:
+    root = Path(repo_root) if repo_root is not None else Path.cwd()
+    try:
+        effective_settings = settings or load_settings(
+            repo_root=root,
+            overrides=SettingsOverrides(max_turns=max_turns),
+        )
+        if max_turns is not None and settings is not None:
+            effective_settings = Settings(
+                base_url=settings.base_url,
+                model=settings.model,
+                api_key=settings.api_key,
+                max_turns=max_turns,
+                max_read_bytes=settings.max_read_bytes,
+                max_grep_results=settings.max_grep_results,
+                traj_dir=settings.traj_dir,
+                ignore=settings.ignore,
+                timeout_seconds=settings.timeout_seconds,
+                max_observation_chars=settings.max_observation_chars,
+                max_read_lines=settings.max_read_lines,
+                max_completion_tokens=settings.max_completion_tokens,
+                temperature=settings.temperature,
+                max_parallel_tools=settings.max_parallel_tools,
+            )
+        request = ExploreRequest(
+            query=query,
+            repo_root=root,
+            max_turns=effective_settings.max_turns,
+            citation=True if citation is None else citation,
+            format="json",
+        )
+        return core(request, effective_settings).to_dict()
+    except ExplorerError as exc:
+        return {"error": exc.to_dict()}
+
+
+def create_server() -> Any:
+    try:
+        fastmcp_module = import_module("mcp.server.fastmcp")
+    except ImportError as exc:
+        raise ExplorerError(
+            "CONFIG_INVALID",
+            "MCP SDK is not installed; install flash-context[mcp]",
+        ) from exc
+    fast_mcp = fastmcp_module.FastMCP
+    server = fast_mcp("flash-context")
+
+    @server.tool()  # type: ignore[untyped-decorator]
+    def explore_repository(
+        query: str,
+        repo_root: str | None = None,
+        max_turns: int | None = None,
+        citation: bool | None = None,
+    ) -> dict[str, object]:
+        return explore_repository_handler(
+            query=query,
+            repo_root=repo_root,
+            max_turns=max_turns,
+            citation=citation,
+        )
+
+    return server
+
+
+def run_mcp(*, transport: str = "stdio") -> None:
+    server = create_server()
+    server.run(transport=transport)

flash_context/tools/__init__.py

@@ -0,0 +1,5 @@
+from flash_context.tools.glob import repo_glob
+from flash_context.tools.grep import repo_grep
+from flash_context.tools.read import read_file
+
+__all__ = ["read_file", "repo_glob", "repo_grep"]

flash_context/tools/glob.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from flash_context.tools.safety import PathSafety
+from flash_context.types import ExplorerError, ToolObservation
+
+
+def repo_glob(
+    *,
+    repo_root: Path,
+    pattern: str,
+    max_results: int,
+    ignore: list[str] | None = None,
+) -> ToolObservation:
+    if not pattern:
+        raise ExplorerError("INVALID_TOOL_ARGUMENTS", "pattern is required")
+    if max_results <= 0:
+        raise ExplorerError("CONFIG_INVALID", "max_results must be positive")
+    safety = PathSafety(repo_root, tuple(ignore or ()))
+    matches: list[str] = []
+    truncated = False
+    for path in safety.iter_safe_files(pattern):
+        if len(matches) >= max_results:
+            truncated = True
+            break
+        matches.append(safety.relative_or_raw(path))
+    return ToolObservation(matches=matches, truncated=truncated)
+

flash_context/tools/grep.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterator
+from pathlib import Path
+
+from flash_context.tools.safety import PathSafety
+from flash_context.types import ExplorerError, SearchHit, ToolObservation
+
+MAX_HIT_TEXT = 240
+
+
+def repo_grep(
+    *,
+    repo_root: Path,
+    pattern: str,
+    max_results: int,
+    glob: str = "**/*",
+    configured_max_results: int | None = None,
+    ignore: list[str] | None = None,
+) -> ToolObservation:
+    if not pattern:
+        raise ExplorerError("INVALID_TOOL_ARGUMENTS", "pattern is required")
+    if max_results <= 0:
+        raise ExplorerError("CONFIG_INVALID", "max_results must be positive")
+    limit = min(max_results, configured_max_results or max_results)
+    try:
+        regex = re.compile(pattern)
+    except re.error as exc:
+        raise ExplorerError(
+            "INVALID_GREP_PATTERN",
+            "Grep pattern is not a valid regular expression",
+            details={"error": str(exc)},
+        ) from exc
+
+    safety = PathSafety(repo_root, tuple(ignore or ()))
+    hits: list[SearchHit] = []
+    truncated = False
+    for path in safety.iter_safe_files(glob or "**/*"):
+        for line_number, line in _iter_text_lines(path):
+            if not regex.search(line):
+                continue
+            if len(hits) >= limit:
+                truncated = True
+                return ToolObservation(hits=hits, truncated=truncated)
+            hits.append(
+                SearchHit(
+                    path=safety.relative_or_raw(path),
+                    line=line_number,
+                    text=_cap_text(line.rstrip("\n")),
+                )
+            )
+    return ToolObservation(hits=hits, truncated=truncated)
+
+
+def _iter_text_lines(path: Path) -> Iterator[tuple[int, str]]:
+    with path.open("r", encoding="utf-8", errors="replace") as handle:
+        for line_number, line in enumerate(handle, start=1):
+            if "\x00" in line:
+                continue
+            yield line_number, line
+
+
+def _cap_text(text: str) -> str:
+    if len(text) <= MAX_HIT_TEXT:
+        return text
+    return f"{text[: MAX_HIT_TEXT - 1]}..."

flash_context/tools/read.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from flash_context.tools.safety import PathSafety
+from flash_context.types import ExplorerError, ToolObservation
+
+
+def read_file(
+    *,
+    repo_root: Path,
+    path: str,
+    max_bytes: int,
+    start_line: int | None = None,
+    end_line: int | None = None,
+    ignore: list[str] | None = None,
+) -> ToolObservation:
+    if not path:
+        raise ExplorerError("INVALID_TOOL_ARGUMENTS", "path is required")
+    if start_line is not None and start_line <= 0:
+        raise ExplorerError("INVALID_TOOL_ARGUMENTS", "start_line must be positive")
+    if end_line is not None and end_line <= 0:
+        raise ExplorerError("INVALID_TOOL_ARGUMENTS", "end_line must be positive")
+    if start_line is not None and end_line is not None and end_line < start_line:
+        raise ExplorerError(
+            "INVALID_TOOL_ARGUMENTS",
+            "end_line must be greater than or equal to start_line",
+        )
+    if max_bytes <= 0:
+        raise ExplorerError("CONFIG_INVALID", "max_bytes must be positive")
+
+    safety = PathSafety(repo_root, tuple(ignore or ()))
+    resolved = safety.resolve_existing_file(path)
+    rel_path = safety.relative_or_raw(resolved)
+    first_line = start_line or 1
+    last_line = end_line
+    content_bytes = bytearray()
+    actual_last: int | None = None
+    total_lines = 0
+    truncated = False
+
+    with resolved.open("rb") as handle:
+        for line_number, raw_line in enumerate(handle, start=1):
+            total_lines = line_number
+            if line_number < first_line:
+                continue
+            if last_line is not None and line_number > last_line:
+                break
+            if len(content_bytes) + len(raw_line) > max_bytes:
+                remaining = max_bytes - len(content_bytes)
+                if remaining > 0:
+                    content_bytes.extend(raw_line[:remaining])
+                actual_last = line_number
+                truncated = True
+                break
+            content_bytes.extend(raw_line)
+            actual_last = line_number
+
+    if actual_last is None:
+        raise ExplorerError(
+            "INVALID_TOOL_ARGUMENTS",
+            "Requested start_line is beyond the end of file",
+            details={
+                "path": rel_path,
+                "start_line": first_line,
+                "total_lines": total_lines,
+            },
+        )
+
+    content = content_bytes.decode("utf-8", errors="replace")
+    line_range = f"{first_line}-{actual_last}"
+    return ToolObservation(
+        path=rel_path,
+        line_range=line_range,
+        content=content,
+        truncated=truncated,
+    )

flash_context/tools/safety.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import fnmatch
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from flash_context.types import ExplorerError
+
+DEFAULT_DENY_PATTERNS = (
+    ".git/**",
+    ".env",
+    ".env.*",
+    "*.pem",
+    "*.key",
+    "id_rsa",
+    "id_ed25519",
+    "**/secrets/**",
+    "**/.aws/**",
+    "**/.ssh/**",
+    ".venv/**",
+    "venv/**",
+    "node_modules/**",
+    "dist/**",
+    "build/**",
+)
+
+DENY_DIR_NAMES = {
+    ".git",
+    ".aws",
+    ".ssh",
+    ".venv",
+    "venv",
+    "node_modules",
+    "dist",
+    "build",
+    "secrets",
+}
+
+DENY_FILE_NAMES = {"id_rsa", "id_ed25519", ".env"}
+DENY_FILE_SUFFIXES = (".pem", ".key")
+
+
+@dataclass(frozen=True, slots=True)
+class PathSafety:
+    repo_root: Path
+    ignore: tuple[str, ...] = field(default_factory=tuple)
+
+    def __post_init__(self) -> None:
+        root = self.repo_root.resolve()
+        object.__setattr__(self, "repo_root", root)
+        if not root.exists() or not root.is_dir():
+            raise ExplorerError(
+                "REPO_NOT_FOUND",
+                "Repository root must exist and be a directory",
+                details={"repo_root": str(root)},
+            )
+
+    def resolve_existing_file(self, requested: str | Path) -> Path:
+        raw = Path(requested)
+        if ".." in raw.parts:
+            raise ExplorerError(
+                "PATH_OUTSIDE_ROOT",
+                "Requested path is outside the repository root",
+                details={"path": str(requested)},
+            )
+        candidate = raw if raw.is_absolute() else self.repo_root / raw
+        resolved = candidate.resolve(strict=False)
+        if not _is_relative_to(resolved, self.repo_root):
+            raise ExplorerError(
+                "PATH_OUTSIDE_ROOT",
+                "Requested path is outside the repository root",
+                details={"path": str(requested)},
+            )
+        if not resolved.exists():
+            raise ExplorerError(
+                "PATH_NOT_FOUND",
+                "Requested path does not exist",
+                details={"path": self.relative_or_raw(resolved)},
+            )
+        if not resolved.is_file():
+            raise ExplorerError(
+                "PATH_DENIED",
+                "Requested path is not a regular file",
+                details={"path": self.relative_or_raw(resolved)},
+            )
+        self.ensure_allowed(resolved)
+        return resolved
+
+    def ensure_allowed(self, path: Path) -> None:
+        rel = self.relative_or_raw(path)
+        parts = rel.split("/")
+        if any(part in DENY_DIR_NAMES for part in parts[:-1]):
+            raise ExplorerError(
+                "PATH_DENIED",
+                "Requested path is denied by repository safety policy",
+                details={"path": rel},
+            )
+        name = parts[-1] if parts else rel
+        if (
+            name in DENY_FILE_NAMES
+            or name.startswith(".env.")
+            or name.endswith(DENY_FILE_SUFFIXES)
+            or _matches_patterns(rel, DEFAULT_DENY_PATTERNS)
+            or _matches_patterns(rel, self.ignore)
+        ):
+            raise ExplorerError(
+                "PATH_DENIED",
+                "Requested path is denied by repository safety policy",
+                details={"path": rel},
+            )
+
+    def relative_or_raw(self, path: Path) -> str:
+        try:
+            return path.resolve(strict=False).relative_to(self.repo_root).as_posix()
+        except ValueError:
+            return path.as_posix()
+
+    def validate_glob_pattern(self, pattern: str) -> None:
+        raw = Path(pattern)
+        if raw.is_absolute() or ".." in raw.parts:
+            raise ExplorerError(
+                "PATH_OUTSIDE_ROOT",
+                "Glob pattern must stay inside the repository root",
+                details={"pattern": pattern},
+            )
+
+    def iter_safe_files(self, pattern: str = "**/*") -> list[Path]:
+        self.validate_glob_pattern(pattern)
+        files: list[Path] = []
+        for candidate in sorted(self.repo_root.glob(pattern)):
+            try:
+                resolved = candidate.resolve(strict=False)
+                if not _is_relative_to(resolved, self.repo_root):
+                    continue
+                if not resolved.is_file():
+                    continue
+                self.ensure_allowed(resolved)
+            except ExplorerError:
+                continue
+            files.append(resolved)
+        return files
+
+
+def _matches_patterns(rel_path: str, patterns: tuple[str, ...]) -> bool:
+    return any(
+        _matches_pattern(rel_path, pattern.strip())
+        for pattern in patterns
+        if pattern.strip()
+    )
+
+
+def _matches_pattern(rel_path: str, pattern: str) -> bool:
+    normalized = pattern.removeprefix("./")
+    if normalized.endswith("/**"):
+        prefix = normalized[:-3].rstrip("/")
+        return rel_path == prefix or rel_path.startswith(f"{prefix}/")
+    if "/" not in normalized:
+        parts = rel_path.split("/")
+        return any(fnmatch.fnmatchcase(part, normalized) for part in parts)
+    return fnmatch.fnmatchcase(rel_path, normalized)
+
+
+def _is_relative_to(path: Path, root: Path) -> bool:
+    try:
+        path.relative_to(root)
+    except ValueError:
+        return False
+    return True

flash_context/types.py

@@ -0,0 +1,210 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Literal
+
+ErrorCode = Literal[
+    "CONFIG_MISSING_ENDPOINT",
+    "CONFIG_INVALID",
+    "REPO_NOT_FOUND",
+    "QUERY_EMPTY",
+    "PATH_OUTSIDE_ROOT",
+    "PATH_DENIED",
+    "PATH_NOT_FOUND",
+    "TOOL_OUTPUT_TRUNCATED",
+    "ENDPOINT_UNREACHABLE",
+    "ENDPOINT_TIMEOUT",
+    "ENDPOINT_BAD_RESPONSE",
+    "UNSUPPORTED_TOOL_CALL",
+    "MAX_TURNS_EXCEEDED",
+    "INVALID_TOOL_ARGUMENTS",
+    "INVALID_GREP_PATTERN",
+]
+
+
+@dataclass(slots=True)
+class ExplorerError(Exception):
+    code: ErrorCode
+    message: str
+    retryable: bool = False
+    details: dict[str, object] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        super().__init__(f"{self.code}: {self.message}")
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "code": self.code,
+            "message": self.message,
+            "retryable": self.retryable,
+            "details": self.details,
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class Citation:
+    path: str
+    start_line: int | None = None
+    end_line: int | None = None
+    reason: str | None = None
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "path": self.path,
+            "start_line": self.start_line,
+            "end_line": self.end_line,
+            "reason": self.reason,
+        }
+
+    def label(self) -> str:
+        if self.start_line is None:
+            return self.path
+        if self.end_line is None or self.end_line == self.start_line:
+            return f"{self.path}:{self.start_line}"
+        return f"{self.path}:{self.start_line}-{self.end_line}"
+
+
+@dataclass(frozen=True, slots=True)
+class RawLocation:
+    path: str
+    start_line: int
+    end_line: int
+    text: str
+    truncated: bool = False
+
+    def to_dict(self, *, include_text: bool = True) -> dict[str, object]:
+        data: dict[str, object] = {
+            "path": self.path,
+            "start_line": self.start_line,
+            "end_line": self.end_line,
+            "truncated": self.truncated,
+        }
+        if include_text:
+            data["text"] = self.text
+        else:
+            data["text_length"] = len(self.text)
+        return data
+
+
+@dataclass(frozen=True, slots=True)
+class SearchHit:
+    path: str
+    line: int
+    text: str
+
+    def to_dict(self) -> dict[str, object]:
+        return {"path": self.path, "line": self.line, "text": self.text}
+
+
+@dataclass(slots=True)
+class ToolCall:
+    id: str
+    name: Literal["read_file", "repo_glob", "repo_grep"]
+    arguments: dict[str, object]
+
+
+@dataclass(slots=True)
+class ToolObservation:
+    tool_call_id: str | None = None
+    ok: bool = True
+    path: str | None = None
+    line_range: str | None = None
+    content: str | None = None
+    matches: list[str] = field(default_factory=list)
+    hits: list[SearchHit] = field(default_factory=list)
+    truncated: bool = False
+    error: ExplorerError | None = None
+
+    def to_dict(self, *, include_content: bool = True) -> dict[str, object]:
+        data: dict[str, object] = {
+            "ok": self.ok,
+            "truncated": self.truncated,
+        }
+        if self.tool_call_id is not None:
+            data["tool_call_id"] = self.tool_call_id
+        if self.path is not None:
+            data["path"] = self.path
+        if self.line_range is not None:
+            data["line_range"] = self.line_range
+        if include_content and self.content is not None:
+            data["content"] = self.content
+        elif self.content is not None:
+            data["content_length"] = len(self.content)
+        if self.matches:
+            data["matches"] = self.matches
+        if self.hits:
+            data["hits"] = [hit.to_dict() for hit in self.hits]
+        if self.error is not None:
+            data["error"] = self.error.to_dict()
+        return data
+
+    @classmethod
+    def from_error(
+        cls, error: ExplorerError, *, tool_call_id: str | None = None
+    ) -> ToolObservation:
+        return cls(tool_call_id=tool_call_id, ok=False, error=error)
+
+
+@dataclass(frozen=True, slots=True)
+class ExploreRequest:
+    query: str
+    repo_root: Path
+    max_turns: int = 6
+    citation: bool = True
+    format: Literal["text", "json"] = "text"
+
+    def validate(self) -> None:
+        if not self.query.strip():
+            raise ExplorerError("QUERY_EMPTY", "Query must not be blank")
+        if self.max_turns <= 0:
+            raise ExplorerError(
+                "CONFIG_INVALID",
+                "max_turns must be positive",
+                details={"max_turns": self.max_turns},
+            )
+        if not self.repo_root.exists() or not self.repo_root.is_dir():
+            raise ExplorerError(
+                "REPO_NOT_FOUND",
+                "Repository root must exist and be a directory",
+                details={"repo_root": str(self.repo_root)},
+            )
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "query": self.query,
+            "repo_root": str(self.repo_root),
+            "max_turns": self.max_turns,
+            "citation": self.citation,
+            "format": self.format,
+        }
+
+
+@dataclass(slots=True)
+class ExploreResult:
+    query: str
+    repo_root: str
+    answer: str
+    citations: list[Citation]
+    turns_used: int
+    raw_locations: list[RawLocation] = field(default_factory=list)
+    truncated: bool = False
+    warnings: list[str] = field(default_factory=list)
+
+    def to_dict(self, *, include_raw_location_text: bool = True) -> dict[str, object]:
+        return {
+            "query": self.query,
+            "repo_root": self.repo_root,
+            "answer": self.answer,
+            "citations": [citation.to_dict() for citation in self.citations],
+            "raw_locations": [
+                raw_location.to_dict(include_text=include_raw_location_text)
+                for raw_location in self.raw_locations
+            ],
+            "turns_used": self.turns_used,
+            "truncated": self.truncated,
+            "warnings": self.warnings,
+        }
+
+
+JsonObject = dict[str, Any]