agent harness initialize repo

agent-security/tools/audit.py

@@ -0,0 +1,189 @@
+"""Append-only JSONL audit log for agent sessions.
+
+Every event is written as one JSON object per line, with an auto-added
+ISO timestamp and session id.  The log is flushed after every record so
+that a crash still leaves a complete trail.  The file alone is enough
+to reconstruct what the agent did, in order.
+
+§2.4: tool results are truncated to ``MAX_RESULT_BYTES`` in the log to
+prevent unbounded growth; a SHA-256 and full length are stored alongside
+so the truncated entry is self-describing and tamper-evident.
+
+§6.1: dedicated events for every decision step —
+``log_permission_decision``, ``log_llm_request``, ``log_llm_response``,
+``log_session_abort`` — give full forensic replay without guessing.
+"""
+
+import hashlib
+import json
+import uuid
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+
+# Tool results larger than this are stored truncated in the audit log.
+MAX_RESULT_BYTES = 8 * 1024
+
+
+def _truncate_for_log(result: str) -> dict:
+    """Return a dict describing *result*, truncating if large.
+
+    The dict always has ``size`` (bytes) and ``sha256``.  If the result
+    is short enough it is included verbatim under ``result``; otherwise
+    only the first ``MAX_RESULT_BYTES`` are kept under ``result_truncated``.
+    """
+    raw = result if isinstance(result, str) else str(result)
+    encoded = raw.encode("utf-8", errors="replace")
+    digest = hashlib.sha256(encoded).hexdigest()
+    size = len(encoded)
+    if size <= MAX_RESULT_BYTES:
+        return {"result": raw, "size": size, "sha256": digest}
+    truncated = encoded[:MAX_RESULT_BYTES].decode("utf-8", errors="replace")
+    return {
+        "result_truncated": truncated,
+        "truncated_from_size": size,
+        "sha256": digest,
+    }
+
+
+class AuditLog:
+    """Append-only JSONL audit log of agent activity."""
+
+    def __init__(self, log_dir: Path, session_id: str | None = None):
+        self.log_dir = Path(log_dir)
+        self.log_dir.mkdir(parents=True, exist_ok=True)
+        self.session_id = session_id or uuid.uuid4().hex[:12]
+        ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+        self.path = self.log_dir / f"audit-{self.session_id}-{ts}.jsonl"
+        self._fh = self.path.open("a", encoding="utf-8")
+        self.log("session_start", log_file=str(self.path))
+
+    def log(self, event: str, **fields: Any) -> None:
+        record = {
+            "ts": datetime.now(timezone.utc).isoformat(timespec="milliseconds"),
+            "session": self.session_id,
+            "event": event,
+        }
+        record.update(fields)
+        self._fh.write(json.dumps(record, ensure_ascii=False, default=str) + "\n")
+        self._fh.flush()
+
+    # -- convenience wrappers ------------------------------------------
+
+    def log_config(self, **fields: Any) -> None:
+        """Record the session configuration (mode, sandbox root, etc.)."""
+        self.log("config", **fields)
+
+    def log_user_message(self, content: str) -> None:
+        self.log("user_message", content=content)
+
+    def log_assistant_message(self, content: str | None, tool_calls) -> None:
+        calls = []
+        if tool_calls:
+            for tc in tool_calls:
+                try:
+                    args = json.loads(tc.function.arguments)
+                except Exception:
+                    args = tc.function.arguments
+                calls.append({"id": tc.id, "name": tc.function.name, "args": args})
+        self.log("assistant_message", content=content, tool_calls=calls)
+
+    def log_tool_result(
+        self,
+        tool_call_id: str,
+        name: str,
+        args: dict,
+        permission_allowed: bool,
+        container_error: bool,
+        container_reason: str | None,
+        result: str,
+        intent_drift: bool = False,
+        intent_reason: str | None = None,
+        permission_reason: str | None = None,
+    ) -> None:
+        self.log(
+            "tool_result",
+            tool_call_id=tool_call_id,
+            tool=name,
+            args=args,
+            permission_allowed=permission_allowed,
+            permission_reason=permission_reason,
+            container_error=container_error,
+            container_reason=container_reason,
+            intent_drift=intent_drift,
+            intent_reason=intent_reason,
+            **_truncate_for_log(result),
+        )
+
+    # -- §6.1 decision-step logging ----------------------------------
+
+    def log_permission_decision(
+        self,
+        tool: str,
+        args: dict,
+        mode: str,
+        allowed: bool,
+        reason: str | None = None,
+    ) -> None:
+        """Record a standalone permission decision (§6.1).
+
+        This is emitted *before* the tool runs (or is refused), so the
+        audit trail shows the decision and its reason even if the
+        subsequent tool execution crashes.
+        """
+        self.log(
+            "permission_decision",
+            tool=tool,
+            args=args,
+            mode=mode,
+            allowed=allowed,
+            reason=reason,
+        )
+
+    def log_llm_request(
+        self,
+        model: str,
+        message_count: int,
+        token_estimate: int,
+        has_tools: bool,
+    ) -> None:
+        """Record that an LLM request is about to be sent (§6.1)."""
+        self.log(
+            "llm_request",
+            model=model,
+            message_count=message_count,
+            token_estimate=token_estimate,
+            has_tools=has_tools,
+        )
+
+    def log_llm_response(
+        self,
+        model: str,
+        finish_reason: str | None,
+        usage: dict | None = None,
+        message_hash: str | None = None,
+        tool_call_count: int = 0,
+    ) -> None:
+        """Record the LLM response metadata (§6.1).
+
+        ``message_hash`` is a SHA-256 of the assistant message content
+        so the full conversation can be verified for forensic replay
+        without storing every token in the log.
+        """
+        self.log(
+            "llm_response",
+            model=model,
+            finish_reason=finish_reason,
+            usage=usage,
+            message_hash=message_hash,
+            tool_call_count=tool_call_count,
+        )
+
+    def log_session_abort(self, reason: str) -> None:
+        """Record that the session was aborted (§6.3)."""
+        self.log("session_abort", reason=reason)
+
+    def close(self) -> None:
+        self.log("session_end")
+        self._fh.close()