flintai-cli 1.0.0__py3-none-any.whl

flintai/scan/trace_logger.py

@@ -0,0 +1,118 @@
+"""
+trace_logger.py — Structured Observability for the Agentic Reasoning Loop
+=========================================================================
+Writes a JSONL trace log of every tool call made during an agentic scan session.
+
+WHY THIS EXISTS:
+  The agentic scanner itself exhibits the behaviours it scans for in customer agents:
+  autonomous tool-calling, adaptive control flow, incremental decision-making.
+  Per ASI10 (missing_agent_monitoring), any agent that takes consequential actions
+  without an audit trail is a security risk. This module is the scanner's own
+  compliance with that rule.
+
+OUTPUT FORMAT:
+  One JSON object per line, written to <output_path>.trace.jsonl
+  Each line represents one tool call with timing, token estimates, and a result preview.
+  The final line is a session summary record.
+
+USAGE:
+  logger = FileTraceLogger(session_id="AGT-SCAN-abc123", output_path="report.json")
+  logger.start(provider_model="claude-sonnet-4-6")
+
+  with logger.record_call("fetch_file", {"path": "tools/exec.py"}) as ctx:
+      result = do_fetch(...)
+      ctx.set_result(result)
+
+  logger.finish(findings_count=5, exit_reason="completed")
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Iterator
+from contextlib import contextmanager
+from datetime import datetime, timezone
+
+import tiktoken
+
+# ── Token estimation ──────────────────────────────────────────────────────────
+
+_TOKEN_REFERENCE_MODEL = "gpt-4o"
+
+
+def estimate_tokens(text: str) -> int:
+    """Estimate token count using tiktoken (accurate) or char heuristic (fallback).
+
+    Returns:
+        Estimated number of tokens (minimum 1).
+    """
+    try:
+        return max(
+            1, len(tiktoken.encoding_for_model(_TOKEN_REFERENCE_MODEL).encode(text))
+        )
+    except Exception:
+        return max(1, len(text) // 4)
+
+
+def now_iso() -> str:
+    """Return the current UTC time as an ISO 8601 string."""
+    return datetime.now(timezone.utc).isoformat()
+
+
+class CallContext:
+    """Mutable context shared between the record_call context manager and the caller."""
+
+    def __init__(self):
+        self.result: str | None = None
+        self.success: bool = True
+        self.error: str | None = None
+
+    def set_result(self, result: str) -> None:
+        self.result = result
+
+    def set_error(self, error: str) -> None:
+        self.error = error
+        self.success = False
+        self.result = f"ERROR: {error}"
+
+
+# ── Abstract interface ────────────────────────────────────────────────────────
+
+
+class TraceLogger(abc.ABC):
+    """
+    Abstract interface for structured trace logging of tool calls
+    during an agentic scan session.
+
+    Thread-safety: single-threaded use only (scanner runs synchronously).
+    """
+
+    @abc.abstractmethod
+    def start(self, provider_model: str) -> None:
+        """Begin a trace session. Call once before the reasoning loop starts."""
+
+    @abc.abstractmethod
+    @contextmanager
+    def record_call(
+        self, tool_name: str, tool_args: dict, iteration: int
+    ) -> Iterator[CallContext]:
+        """Context manager that times a tool call and writes the trace entry.
+
+        Usage:
+            with logger.record_call("fetch_file", {"path": "x.py"}, iteration=2) as ctx:
+                result = fetch_file(...)
+                ctx.set_result(result)
+        """
+        yield CallContext()  # pragma: no cover
+
+    @abc.abstractmethod
+    def set_iteration(self, n: int) -> None:
+        """Update the current iteration count."""
+
+    @abc.abstractmethod
+    def finish(self, findings_count: int, exit_reason: str = "completed") -> None:
+        """Write the session summary and close. Call once after the reasoning loop exits."""
+
+    @abc.abstractmethod
+    def as_dict(self) -> dict:
+        """Return a summary dict suitable for embedding in ScanReport.scan_metadata."""

flintai/scan/trace_logger_file.py

@@ -0,0 +1,143 @@
+"""FileTraceLogger — writes a JSONL trace file to disk."""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+import uuid
+from contextlib import contextmanager
+
+from flintai.scan.trace_logger import (
+    CallContext,
+    TraceLogger,
+    estimate_tokens,
+    now_iso,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class FileTraceLogger(TraceLogger):
+    """
+    Writes a structured JSONL trace of all tool calls made during an agentic scan.
+
+    Output is written to ``<output_path>.trace.jsonl``.
+    """
+
+    def __init__(self, session_id: str | None = None, output_path: str | None = None):
+        self.session_id = session_id or f"AGT-SCAN-{uuid.uuid4().hex[:8].upper()}"
+        self.output_path = output_path
+        self._trace_path = f"{output_path}.trace.jsonl" if output_path else None
+        self._start_time: float | None = None
+        self._provider = "unknown"
+        self._iterations = 0
+        self._call_count = 0
+        self._total_tokens = 0
+        self._tool_calls: list[dict] = []
+        self._file = None
+
+    def start(self, provider_model: str) -> None:
+        self._start_time = time.monotonic()
+        self._provider = provider_model
+        self._write_line(
+            {
+                "event": "session_start",
+                "session_id": self.session_id,
+                "model": provider_model,
+                "timestamp": now_iso(),
+            }
+        )
+        logger.info(
+            "Trace session started: %s | model: %s", self.session_id, provider_model
+        )
+
+    @contextmanager
+    def record_call(self, tool_name: str, tool_args: dict, iteration: int):
+        call_start = time.monotonic()
+        ctx = CallContext()
+        try:
+            yield ctx
+        finally:
+            duration_ms = int((time.monotonic() - call_start) * 1000)
+            result_text = ctx.result or ""
+            tokens = estimate_tokens(result_text)
+            self._total_tokens += tokens
+            self._call_count += 1
+
+            entry = {
+                "event": "tool_call",
+                "session_id": self.session_id,
+                "iteration": iteration,
+                "call_number": self._call_count,
+                "tool": tool_name,
+                "args": tool_args,
+                "result_preview": result_text[:200],
+                "result_length": len(result_text),
+                "tokens_consumed": tokens,
+                "tokens_total": self._total_tokens,
+                "duration_ms": duration_ms,
+                "timestamp": now_iso(),
+                "success": ctx.success,
+                "error": ctx.error,
+            }
+            self._tool_calls.append(entry)
+            self._write_line(entry)
+
+    def set_iteration(self, n: int) -> None:
+        self._iterations = n
+
+    def finish(self, findings_count: int, exit_reason: str = "completed") -> None:
+        if self._start_time is None:
+            return
+        wall_ms = int((time.monotonic() - self._start_time) * 1000)
+        summary = {
+            "event": "session_end",
+            "session_id": self.session_id,
+            "model": self._provider,
+            "exit_reason": exit_reason,
+            "total_iterations": self._iterations,
+            "total_tool_calls": self._call_count,
+            "total_tokens": self._total_tokens,
+            "wall_clock_ms": wall_ms,
+            "findings_count": findings_count,
+            "timestamp": now_iso(),
+        }
+        self._write_line(summary)
+        if self._file:
+            self._file.close()
+            self._file = None
+        logger.info(
+            "Trace session complete: %s | iterations=%d | calls=%d | "
+            "tokens≈%d | %dms | exit=%s",
+            self.session_id,
+            self._iterations,
+            self._call_count,
+            self._total_tokens,
+            wall_ms,
+            exit_reason,
+        )
+        if self._trace_path:
+            logger.info("Trace written to: %s", self._trace_path)
+
+    def as_dict(self) -> dict:
+        return {
+            "session_id": self.session_id,
+            "provider_model": self._provider,
+            "total_iterations": self._iterations,
+            "total_tool_calls": self._call_count,
+            "total_tokens": self._total_tokens,
+        }
+
+    def _write_line(self, obj: dict) -> None:
+        """Append a JSON line to the trace file. Opens on first write."""
+        if self._trace_path is None:
+            return
+        try:
+            if self._file is None:
+                self._file = open(self._trace_path, "w", encoding="utf-8")
+            self._file.write(json.dumps(obj) + "\n")
+            self._file.flush()
+        except Exception as e:
+            # Trace writes are best-effort — never crash the scanner
+            logger.warning("Could not write trace entry: %s", e)

flintai/scan/trace_logger_log.py

@@ -0,0 +1,100 @@
+"""LogTraceLogger — emits trace events via Python logging."""
+
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from contextlib import contextmanager
+
+from flintai.scan.trace_logger import (
+    CallContext,
+    TraceLogger,
+    estimate_tokens,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class LogTraceLogger(TraceLogger):
+    """
+    Trace logger that emits structured information via Python's logging module
+    instead of writing to a file. Useful for environments where file output
+    is unavailable or when trace data should flow through the standard
+    log pipeline.
+    """
+
+    def __init__(self, session_id: str | None = None, log_level: int = logging.INFO):
+        self.session_id = session_id or f"AGT-SCAN-{uuid.uuid4().hex[:8].upper()}"
+        self._log_level = log_level
+        self._start_time: float | None = None
+        self._provider = "unknown"
+        self._iterations = 0
+        self._call_count = 0
+        self._total_tokens = 0
+
+    def start(self, provider_model: str) -> None:
+        self._start_time = time.monotonic()
+        self._provider = provider_model
+        logger.log(
+            self._log_level,
+            "[trace:%s] session_start | model=%s",
+            self.session_id,
+            provider_model,
+        )
+
+    @contextmanager
+    def record_call(self, tool_name: str, tool_args: dict, iteration: int):
+        call_start = time.monotonic()
+        ctx = CallContext()
+        try:
+            yield ctx
+        finally:
+            duration_ms = int((time.monotonic() - call_start) * 1000)
+            result_text = ctx.result or ""
+            tokens = estimate_tokens(result_text)
+            self._total_tokens += tokens
+            self._call_count += 1
+
+            logger.log(
+                self._log_level,
+                "[trace:%s] tool_call #%d | iter=%d tool=%s "
+                "duration=%dms tokens=%d success=%s%s",
+                self.session_id,
+                self._call_count,
+                iteration,
+                tool_name,
+                duration_ms,
+                tokens,
+                ctx.success,
+                f" error={ctx.error}" if ctx.error else "",
+            )
+
+    def set_iteration(self, n: int) -> None:
+        self._iterations = n
+
+    def finish(self, findings_count: int, exit_reason: str = "completed") -> None:
+        if self._start_time is None:
+            return
+        wall_ms = int((time.monotonic() - self._start_time) * 1000)
+        logger.log(
+            self._log_level,
+            "[trace:%s] session_end | exit=%s iterations=%d calls=%d "
+            "tokens≈%d wall=%dms findings=%d",
+            self.session_id,
+            exit_reason,
+            self._iterations,
+            self._call_count,
+            self._total_tokens,
+            wall_ms,
+            findings_count,
+        )
+
+    def as_dict(self) -> dict:
+        return {
+            "session_id": self.session_id,
+            "provider_model": self._provider,
+            "total_iterations": self._iterations,
+            "total_tool_calls": self._call_count,
+            "total_tokens": self._total_tokens,
+        }