gemini-cli/tools/caretaker-agent/cloudrun/triage-worker/utils/agent_logger.py

import json
import datetime
import os
import uuid
from google.cloud import storage
from google.antigravity.types import Text

BUCKET_NAME = os.environ.get("TRIAGE_DEBUG_LOGS_BUCKET")
_storage_client = None


def _get_storage_client() -> storage.Client:
    global _storage_client
    if _storage_client is None:
        _storage_client = storage.Client()
    return _storage_client


def upload_to_bucket(repository: str, issue_number: str | int, payload: str) -> None:
    """
    Uploads a string payload directly to the triage debug logs GCS bucket.
    """
    if not payload or not BUCKET_NAME:
        if not BUCKET_NAME:
            print("[LOGIC] Warning: Missing TRIAGE_DEBUG_LOGS_BUCKET, skipping GCS upload.")
        return
    try:
        storage_client = _get_storage_client()
        bucket = storage_client.bucket(BUCKET_NAME)
        
        timestamp = datetime.datetime.now(datetime.timezone.utc).strftime("%Y%m%d_%H%M%S")
        safe_repo = str(repository).replace("/", "_") if repository else "unknown"
        unique_id = uuid.uuid4().hex[:8]
        blob_name = f"{safe_repo}/issue_{issue_number}_{timestamp}_{unique_id}_debug.log"
        
        blob = bucket.blob(blob_name)
        blob.upload_from_string(payload, content_type="text/plain")
        print(f"[LOGIC] Uploaded debug logs to gs://{BUCKET_NAME}/{blob_name}")
    except Exception as e:
        print(f"[LOGIC] Error uploading debug logs to GCS: {e}")

def _format_debug_trajectory(resolved_chunks: list) -> str:
    """
    Formats Antigravity Agent resolved stream chunks (thoughts, tool calls,
    tool outputs) into a structured, human-readable JSON string.
    """
    if not resolved_chunks:
        return "[]"
    serializable = []
    for chunk in resolved_chunks:
        try:
            dumped = chunk.model_dump()
        except AttributeError:
            dumped = chunk.dict()
        dumped["chunk_type"] = chunk.__class__.__name__
        serializable.append(dumped)
    return json.dumps(serializable, indent=2, default=str)


def log_agent_run(
    repository: str,
    issue_number: str | int,
    resolved_chunks: list,
    mode: str = "GCS",
) -> None:
    """
    Logs the agent execution trajectory based on the mode parameter.

    Modes:
      - "LOCAL": Saves log file to local disk (requires LOCAL_LOG_DIR env var).
      - "GCS": Uploads log file to the GCS bucket.
      - Any other mode (e.g. "OFF"): Skips logging.
    """
    mode_upper = str(mode).upper()
    if not resolved_chunks or mode_upper not in ("LOCAL", "GCS"):
        return

    try:
        debug_log_str = _format_debug_trajectory(resolved_chunks)

        if mode_upper == "LOCAL":
            local_dir = os.environ.get("LOCAL_LOG_DIR")
            if not local_dir:
                print(
                    "[LOGIC] Error: LOCAL_LOG_DIR env var is not configured."
                )
                return
            os.makedirs(local_dir, exist_ok=True)
            log_path = os.path.join(
                local_dir, f"gemini_cli_{issue_number}_debug.json"
            )
            with open(log_path, "w", encoding="utf-8") as f:
                f.write(debug_log_str)
            print(f"[LOGIC] 📄 Saved local agent trajectory log to: {log_path}")

        elif mode_upper == "GCS":
            upload_to_bucket(repository, issue_number, debug_log_str)

    except Exception as e:
        print(f"[LOGIC] Error logging agent run: {e}")

def extract_final_output(resolved_chunks: list) -> str:
    """
    Extracts the final response text generated by the agent during the last execution step.
    
    Since the agent outputs conversational planning text during its tool-calling turns,
    this function isolates the final turn (highest step_index) and filters for text chunks,
    stripping away any intermediate planning text, thoughts, or markdown backticks around the JSON.
    """
    if not resolved_chunks:
        return ""
    # Filter for Text chunks (skip Thought and ToolCall chunks)
    text_chunks = [c for c in resolved_chunks if isinstance(c, Text)]
    if not text_chunks:
        return ""
    
    # Final output is in the last step
    last_step = text_chunks[-1].step_index
    return "".join(c.text for c in text_chunks if c.step_index == last_step).strip()