@jaguilar87/gaia 5.0.0-rc.2

package/dist/gaia-security/hooks/modules/context/context_writer.py

@@ -0,0 +1,530 @@
+"""
+ContextWriter module for gaia-ops progressive context enrichment.
+
+Parses CONTEXT_UPDATE blocks from agent output, validates write permissions
+against contracts, and applies deep-merge updates to project-context.json.
+
+Public API:
+    - parse_context_update(agent_output) -> Optional[dict]
+    - validate_permissions(update, agent_type, contracts) -> (dict, list)
+    - apply_update(context_path, update, agent_type) -> dict
+    - load_contracts(provider, config_dir) -> dict
+    - process_agent_output(agent_output, task_info) -> dict
+"""
+
+import json
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Dynamic import: deep_merge from tools/context/deep_merge.py
+# ---------------------------------------------------------------------------
+def _import_deep_merge():
+    """Import deep_merge function from tools/context/deep_merge.py."""
+    import importlib.util
+    import sys
+
+    # Resolve paths relative to this file:
+    # this file:   hooks/modules/context/context_writer.py
+    # deep_merge:  tools/context/deep_merge.py
+    hooks_dir = Path(__file__).resolve().parents[2]  # hooks/
+    repo_root = hooks_dir.parent                      # repo root
+    deep_merge_path = repo_root / "tools" / "context" / "deep_merge.py"
+
+    if not deep_merge_path.exists():
+        raise ImportError(f"deep_merge.py not found at {deep_merge_path}")
+
+    spec = importlib.util.spec_from_file_location("deep_merge", str(deep_merge_path))
+    module = importlib.util.module_from_spec(spec)
+    sys.modules["deep_merge"] = module
+    spec.loader.exec_module(module)
+    return module.deep_merge
+
+
+_deep_merge = _import_deep_merge()
+
+
+# ---------------------------------------------------------------------------
+# LEGACY_AGENT_CONTRACTS (fallback when no contracts file exists)
+# In legacy mode, write permissions = read permissions.
+# ---------------------------------------------------------------------------
+LEGACY_AGENT_CONTRACTS: Dict[str, List[str]] = {
+    "terraform-architect": [
+        "project_identity",
+        "stack",
+        "git",
+        "environment",
+        "infrastructure",
+        "terraform_infrastructure",
+        "infrastructure_topology",
+        "operational_guidelines",
+    ],
+    "gitops-operator": [
+        "project_identity",
+        "stack",
+        "git",
+        "environment",
+        "infrastructure",
+        "gitops_configuration",
+        "infrastructure_topology",
+        "cluster_details",
+        "operational_guidelines",
+    ],
+    "cloud-troubleshooter": [
+        "project_identity",
+        "stack",
+        "git",
+        "environment",
+        "infrastructure",
+        "infrastructure_topology",
+        "terraform_infrastructure",
+        "gitops_configuration",
+        "application_services",
+        "monitoring_observability",
+        "cluster_details",
+    ],
+    "developer": [
+        "project_identity",
+        "stack",
+        "git",
+        "environment",
+        "infrastructure",
+        "application_services",
+        "operational_guidelines",
+    ],
+    "gaia-operator": [
+        "project_identity",
+        "stack",
+        "git",
+        "environment",
+        "infrastructure",
+        "infrastructure_topology",
+        "application_services",
+        "architecture_overview",
+        "operational_guidelines",
+        "workspace_repos",
+    ],
+}
+
+
+# ---------------------------------------------------------------------------
+# Module-level cache for load_contracts
+# ---------------------------------------------------------------------------
+_contracts_cache: Dict[str, dict] = {}
+
+
+# ============================================================================
+# 1. parse_context_update
+# ============================================================================
+
+def parse_context_update(agent_output: str) -> Optional[dict]:
+    """Extract and parse a CONTEXT_UPDATE block from agent output.
+
+    Searches for the ``CONTEXT_UPDATE:`` marker (case-sensitive, on its own
+    line), extracts the JSON that follows until end-of-output or the next
+    known marker, and returns the parsed dict.
+
+    Returns None when:
+    - No marker is found
+    - The JSON is malformed
+    - The parsed value is not a dict
+    """
+    # Find the CONTEXT_UPDATE: marker on its own line
+    marker = "CONTEXT_UPDATE:"
+    lines = agent_output.split("\n")
+
+    marker_idx = None
+    for i, line in enumerate(lines):
+        if line.strip() == marker:
+            marker_idx = i
+            break
+
+    if marker_idx is None:
+        return None
+
+    # Collect all text after the marker
+    remaining = "\n".join(lines[marker_idx + 1:]).strip()
+
+    if not remaining:
+        return None
+
+    # Strip markdown code fences — LLMs reading SKILL.md documentation
+    # often wrap the JSON in ```json ... ``` or ``` ... ``` blocks.
+    if remaining.startswith("```"):
+        fence_lines = remaining.split("\n")
+        # Remove opening fence (```json, ```JSON, ```, etc.)
+        fence_lines.pop(0)
+        # Remove closing fence if present
+        for i in range(len(fence_lines) - 1, -1, -1):
+            if fence_lines[i].strip() == "```":
+                fence_lines.pop(i)
+                break
+        remaining = "\n".join(fence_lines).strip()
+
+    if not remaining:
+        return None
+
+    # Use raw_decode to extract the first complete JSON value, ignoring
+    # any trailing text (summaries, AGENT_STATUS blocks, etc.)
+    decoder = json.JSONDecoder()
+    try:
+        parsed, _ = decoder.raw_decode(remaining)
+    except (json.JSONDecodeError, ValueError) as exc:
+        logger.warning("Malformed JSON in CONTEXT_UPDATE block: %s", exc)
+        return None
+
+    if not isinstance(parsed, dict):
+        return None
+
+    return parsed
+
+
+# ============================================================================
+# 2. validate_permissions
+# ============================================================================
+
+def validate_permissions(
+    update: dict,
+    agent_type: str,
+    contracts: dict,
+) -> tuple:
+    """Validate which sections the agent is allowed to write.
+
+    Returns ``(allowed_updates, rejected_sections)`` where:
+    - ``allowed_updates``: dict with only permitted sections
+    - ``rejected_sections``: list of section names that were rejected
+    """
+    allowed: dict = {}
+    rejected: List[str] = []
+
+    # Determine writable sections for this agent
+    agents_map = contracts.get("agents", {})
+
+    if agent_type in agents_map:
+        # Use explicit write list from contracts file
+        writable = set(agents_map[agent_type].get("write", []))
+    elif agent_type in LEGACY_AGENT_CONTRACTS:
+        # Fallback: in legacy mode, write = read
+        writable = set(LEGACY_AGENT_CONTRACTS[agent_type])
+    else:
+        # Unknown agent: no permissions
+        writable = set()
+
+    for section, data in update.items():
+        if section in writable:
+            allowed[section] = data
+        else:
+            rejected.append(section)
+
+    return allowed, rejected
+
+
+# ============================================================================
+# 3. apply_update
+# ============================================================================
+
+def apply_update(
+    context_path: Path,
+    update: dict,
+    agent_type: str,
+) -> dict:
+    """Apply a validated update to project-context.json using deep merge.
+
+    Performs an atomic write (write to .tmp, then rename) and appends an
+    audit entry to ``context-audit.jsonl`` in the same directory.
+
+    Returns an audit entry dict. Never raises on I/O errors.
+    """
+    context_path = Path(context_path)
+    timestamp = datetime.now(timezone.utc).isoformat()
+    sections_updated = list(update.keys())
+
+    audit_entry = {
+        "timestamp": timestamp,
+        "agent": agent_type,
+        "sections_updated": sections_updated,
+        "changes": {},
+        "success": False,
+        "error": None,
+    }
+
+    try:
+        # Read current context
+        current = json.loads(context_path.read_text())
+
+        # Deep merge each section
+        all_changes = {}
+        for section, section_data in update.items():
+            current_section = current.get("sections", {}).get(section, {})
+            merged_section, diff = _deep_merge(current_section, section_data)
+
+            # Ensure sections dict exists
+            if "sections" not in current:
+                current["sections"] = {}
+            current["sections"][section] = merged_section
+
+            if diff:
+                all_changes[section] = diff
+
+        # Update metadata timestamp
+        if "metadata" not in current:
+            current["metadata"] = {}
+        current["metadata"]["last_updated"] = timestamp
+
+        # Atomic write: write to .tmp, then rename
+        tmp_path = context_path.with_suffix(".tmp")
+        tmp_path.write_text(json.dumps(current, indent=2))
+        tmp_path.rename(context_path)
+
+        audit_entry["changes"] = all_changes
+        audit_entry["success"] = True
+
+    except Exception as exc:
+        logger.error("Failed to apply context update: %s", exc)
+        audit_entry["error"] = str(exc)
+        audit_entry["success"] = False
+        return audit_entry
+
+    # Append audit entry to context-audit.jsonl
+    try:
+        audit_file = context_path.parent / "context-audit.jsonl"
+        with open(audit_file, "a") as f:
+            f.write(json.dumps(audit_entry) + "\n")
+    except Exception as exc:
+        logger.warning("Failed to write audit entry: %s", exc)
+        # Audit write failure doesn't affect the main result
+
+    return audit_entry
+
+
+# ============================================================================
+# 4. load_contracts
+# ============================================================================
+
+def load_contracts(provider: str, config_dir: Path) -> dict:
+    """Load agent contracts using the base+cloud merge strategy, with caching.
+
+    Strategy (in priority order):
+    1. Load base contracts from context-contracts.json (cloud-agnostic)
+    2. Merge cloud overrides from cloud/{provider}.json (extend read/write lists)
+    3. Fallback: try legacy context-contracts.{provider}.json
+    4. Final fallback: LEGACY_AGENT_CONTRACTS hardcoded dict (write = read)
+
+    Results are cached by provider string.
+    """
+    config_dir = Path(config_dir)
+
+    # Check cache first
+    if provider in _contracts_cache:
+        return _contracts_cache[provider]
+
+    result = _merge_base_and_cloud(provider, config_dir)
+    _contracts_cache[provider] = result
+    return result
+
+
+def _merge_base_and_cloud(provider: str, config_dir: Path) -> dict:
+    """Load and merge base + cloud/{provider}.json contracts.
+
+    Returns a merged contracts dict with 'agents' keyed by agent name,
+    each containing 'read' and 'write' lists.
+    """
+    base_file = config_dir / "context-contracts.json"
+    cloud_file = config_dir / "cloud" / f"{provider}.json"
+
+    # Step 1: Load base contracts
+    base_contracts = None
+    if base_file.exists():
+        try:
+            base_contracts = json.loads(base_file.read_text())
+        except (json.JSONDecodeError, OSError) as exc:
+            logger.warning("Failed to load base contracts from %s: %s", base_file, exc)
+
+    # Step 2: Final fallback to hardcoded LEGACY_AGENT_CONTRACTS
+    if base_contracts is None:
+        logger.warning("No contract files found in %s, using hardcoded legacy contracts", config_dir)
+        return {
+            "version": "legacy",
+            "provider": provider,
+            "agents": {
+                agent: {"read": sections, "write": list(sections)}
+                for agent, sections in LEGACY_AGENT_CONTRACTS.items()
+            },
+        }
+
+    # Step 4: Merge cloud-specific overrides
+    if cloud_file.exists():
+        try:
+            cloud_overrides = json.loads(cloud_file.read_text())
+            for agent_name, agent_overrides in cloud_overrides.get("agents", {}).items():
+                if agent_name in base_contracts.get("agents", {}):
+                    existing_read = base_contracts["agents"][agent_name].get("read", [])
+                    existing_write = base_contracts["agents"][agent_name].get("write", [])
+                    extra_read = [s for s in agent_overrides.get("read", []) if s not in existing_read]
+                    extra_write = [s for s in agent_overrides.get("write", []) if s not in existing_write]
+                    base_contracts["agents"][agent_name]["read"] = existing_read + extra_read
+                    base_contracts["agents"][agent_name]["write"] = existing_write + extra_write
+                else:
+                    base_contracts["agents"][agent_name] = agent_overrides
+            logger.info("Merged %s cloud overrides from %s", provider.upper(), cloud_file)
+        except (json.JSONDecodeError, OSError) as exc:
+            logger.warning("Failed to load cloud overrides from %s: %s — skipping", cloud_file, exc)
+
+    return base_contracts
+
+
+# ============================================================================
+# 5. process_agent_output
+# ============================================================================
+
+def process_agent_output(agent_output: str, task_info: dict) -> dict:
+    """Orchestrate the full context-update flow.
+
+    Steps: parse -> detect provider -> load contracts -> validate -> apply.
+
+    Parameters
+    ----------
+    agent_output : str
+        Full agent output string.
+    task_info : dict
+        Must contain: ``agent_type``, ``context_path``, ``config_dir``.
+
+    Returns
+    -------
+    dict
+        ``{updated, sections/sections_updated, rejected, error}``
+    """
+    result = {
+        "updated": False,
+        "sections_updated": [],
+        "rejected": [],
+        "error": None,
+    }
+
+    # 1. Parse CONTEXT_UPDATE
+    update = parse_context_update(agent_output)
+    if update is None:
+        return result
+
+    agent_type = task_info.get("agent_type", "unknown")
+    context_path = Path(task_info.get("context_path", ""))
+    config_dir = Path(task_info.get("config_dir", ""))
+
+    # 2. Detect cloud provider from existing context metadata
+    provider = "gcp"  # default
+    try:
+        if context_path.exists():
+            ctx = json.loads(context_path.read_text())
+            provider = ctx.get("metadata", {}).get("cloud_provider", "gcp").lower()
+    except Exception:
+        pass
+
+    # 3. Load contracts
+    # Clear cache to avoid cross-test pollution (keyed by provider+config_dir)
+    cache_key = f"{provider}:{config_dir}"
+    contracts = _load_contracts_with_dir(provider, config_dir)
+
+    # 4. Validate permissions
+    allowed, rejected = validate_permissions(update, agent_type, contracts)
+    result["rejected"] = rejected
+
+    # 5. If nothing allowed, return early
+    if not allowed:
+        return result
+
+    # 6. Apply update
+    audit = apply_update(context_path, allowed, agent_type)
+
+    if audit.get("success"):
+        result["updated"] = True
+        result["sections_updated"] = list(allowed.keys())
+    else:
+        result["error"] = audit.get("error")
+
+    return result
+
+
+def _load_contracts_with_dir(provider: str, config_dir: Path) -> dict:
+    """Load contracts, bypassing the module-level cache for process_agent_output.
+
+    This ensures each call with a different config_dir gets fresh results
+    while still allowing load_contracts to cache for repeated calls with
+    the same provider. Uses the same base+cloud merge strategy as load_contracts.
+    """
+    return _merge_base_and_cloud(provider, Path(config_dir))
+
+
+# ============================================================================
+# 6. process_context_updates (thin wrapper for subagent_stop integration)
+# ============================================================================
+
+def process_context_updates(
+    agent_output: str,
+    task_info: dict,
+    find_claude_dir_fn=None,
+) -> Optional[dict]:
+    """
+    Process CONTEXT_UPDATE blocks from agent output via context_writer.
+
+    Loads project-context.json, resolves config_dir from .claude, and calls
+    process_agent_output() to apply progressive enrichment.
+
+    This function MUST NOT break the existing hook flow -- all errors are caught
+    and logged, returning None on failure.
+
+    Args:
+        agent_output: Complete output from agent execution
+        task_info: Task metadata (agent, description, task_id)
+        find_claude_dir_fn: Callable that returns the .claude Path. Defaults to
+            modules.core.paths.find_claude_dir if not provided.
+
+    Returns:
+        Result dict from process_agent_output, or None on error
+    """
+    try:
+        if find_claude_dir_fn is None:
+            from ..core.paths import find_claude_dir
+            find_claude_dir_fn = find_claude_dir
+
+        # Find project-context.json via find_claude_dir
+        claude_dir = find_claude_dir_fn()
+        context_path = claude_dir / "project-context" / "project-context.json"
+
+        if not context_path.exists():
+            logger.debug("project-context.json not found at %s, skipping context updates", context_path)
+            return None
+
+        # Determine config_dir (inside .claude directory)
+        config_dir = claude_dir / "config"
+
+        # Build task_info dict for process_agent_output
+        agent_type = task_info.get("agent", "unknown")
+        task_info_for_writer = {
+            "agent_type": agent_type,
+            "context_path": str(context_path),
+            "config_dir": str(config_dir),
+        }
+
+        result = process_agent_output(agent_output, task_info_for_writer)
+
+        if result and result.get("updated"):
+            logger.info(
+                "Context updated by %s: sections=%s",
+                agent_type, result.get("sections_updated", []),
+            )
+        if result and result.get("rejected"):
+            logger.debug(
+                "Context sections rejected for %s: %s",
+                agent_type, result.get("rejected", []),
+            )
+
+        return result
+
+    except Exception as e:
+        logger.debug("Context update processing failed (non-fatal): %s", e)
+        return None

package/dist/gaia-security/hooks/modules/context/contracts_loader.py

@@ -0,0 +1,161 @@
+"""Load context contracts and merge agent permissions.
+
+Subsystem 2 of the pre_tool_use Task/Agent path.
+
+Loads context-contracts.json + cloud overlays, merges agent permissions,
+finds empty writable sections, and builds a reminder string.
+
+Cloud provider detection (formerly infrastructure_reader) is internal
+to the contract loading process.
+"""
+
+import json
+import logging
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def _detect_cloud_from_infrastructure(sections: dict) -> str:
+    """Extract cloud provider name from v2 infrastructure.cloud_providers section.
+
+    Args:
+        sections: The sections dict from project-context.json.
+
+    Returns:
+        Cloud provider name string (e.g. aws, gcp) or empty string.
+    """
+    infra = sections.get("infrastructure", {})
+    if isinstance(infra, dict):
+        providers = infra.get("cloud_providers", [])
+        if isinstance(providers, list) and providers:
+            primary = providers[0]
+            if isinstance(primary, dict):
+                return primary.get("name", "")
+    return ""
+
+
+def build_context_update_reminder(
+    subagent_type: str,
+    project_agents: list,
+    hooks_dir: Path = None,
+) -> str:
+    """Check which writable sections are empty and build a reminder.
+
+    Reads the context contracts to find writable sections for this agent,
+    then checks project-context.json to see which are empty.
+
+    Args:
+        subagent_type: The agent type string (e.g. developer).
+        project_agents: List of valid project agent names.
+        hooks_dir: Path to the hooks directory (for fallback config lookup).
+            Defaults to Path(__file__).parent.parent.parent if None.
+
+    Returns:
+        Reminder string or empty string if no empty sections.
+    """
+    if subagent_type not in project_agents:
+        return ""
+
+    if hooks_dir is None:
+        hooks_dir = Path(__file__).parent.parent.parent
+
+    # Load contracts to find writable sections.
+    # Strategy: load context-contracts.json (base) then merge cloud/{provider}.json.
+    # Fallback to legacy per-provider files for backward compatibility.
+    # We detect the cloud provider from project-context.json first.
+    cloud_provider = "gcp"  # default
+    pc_paths_for_provider = [
+        Path(".claude/project-context/project-context.json"),
+        Path("project-context.json"),
+    ]
+    for pp in pc_paths_for_provider:
+        if pp.exists():
+            try:
+                pc_data = json.loads(pp.read_text())
+                detected = (
+                    pc_data.get("metadata", {}).get("cloud_provider", "")
+                    or _detect_cloud_from_infrastructure(pc_data.get("sections", {}))
+                )
+                if detected:
+                    cloud_provider = detected.lower()
+                break
+            except Exception:
+                continue
+
+    # Candidate config directories (installed project first, package fallback)
+    config_dirs = [
+        Path(".claude/config"),
+        hooks_dir.parent / "config",
+    ]
+
+    writable = []
+    for config_dir in config_dirs:
+        if not config_dir.is_dir():
+            continue
+        # Load base contracts
+        base_file = config_dir / "context-contracts.json"
+        cloud_file = config_dir / "cloud" / f"{cloud_provider}.json"
+
+        merged_agents = {}
+        if base_file.exists():
+            try:
+                data = json.loads(base_file.read_text())
+                merged_agents = data.get("agents", {})
+            except Exception:
+                pass
+
+        # Merge cloud overrides
+        if merged_agents and cloud_file.exists():
+            try:
+                cloud_data = json.loads(cloud_file.read_text())
+                agent_cloud = cloud_data.get("agents", {}).get(subagent_type, {})
+                base_write = merged_agents.get(subagent_type, {}).get("write", [])
+                extra_write = [s for s in agent_cloud.get("write", []) if s not in base_write]
+                if subagent_type in merged_agents:
+                    merged_agents[subagent_type]["write"] = base_write + extra_write
+            except Exception:
+                pass
+
+        if merged_agents:
+            agent_perms = merged_agents.get(subagent_type, {})
+            writable = agent_perms.get("write", [])
+            if writable:
+                break
+
+    if not writable:
+        return ""
+
+    # Load project-context.json to find empty sections
+    pc_paths = [
+        Path(".claude/project-context/project-context.json"),
+        Path("project-context.json"),
+    ]
+
+    sections = {}
+    for pp in pc_paths:
+        if pp.exists():
+            try:
+                pc = json.loads(pp.read_text())
+                sections = pc.get("sections", {})
+                break
+            except Exception:
+                continue
+
+    # Find empty writable sections
+    empty = []
+    for section_name in writable:
+        section_data = sections.get(section_name, {})
+        if not section_data or section_data == {}:
+            empty.append(section_name)
+
+    if not empty:
+        return ""
+
+    empty_list = ", ".join(f"`{s}`" for s in empty)
+    return (
+        f"\n**CONTEXT_UPDATE REQUIRED:** Your writable sections {empty_list} "
+        f"are currently EMPTY. After completing your task, you MUST emit a "
+        f"CONTEXT_UPDATE block with any data you discovered. "
+        f"See \"Context Updater Protocol\" above for the format.\n\n"
+    )