From aa391835af6985dee3032db1281f697825a3f062 Mon Sep 17 00:00:00 2001 From: tn-pisama Date: Wed, 29 Jul 2026 14:07:50 -0700 Subject: [PATCH 1/2] feat: fold pisama-auto and pisama-agent-sdk into pisama as submodules Bumps pisama to 0.6.0. Adds pisama.auto (in-package original of the standalone pisama-auto distribution: zero-code OTEL auto-instrumentation for the Anthropic/OpenAI SDKs) and pisama.agents (in-package original of pisama-agent-sdk: real-time hooks, tools, and the pisama-openhands-monitor console script for Claude Agent SDK / Harbor integrations), each behind its own optional extra so a bare `pip install pisama` pulls in no new required dependencies. Also fixes two gaps found in adversarial verification of the port: - pisama.agents was missing a __version__ attribute (the original pisama_agent_sdk exported one, 58 __all__ symbols vs the ported 57). Since this code no longer releases independently, it now aliases pisama.__version__ instead of carrying a separate hardcoded number. - The `auto` extra omitted opentelemetry-exporter-otlp-proto-http, which setup_tracer() needs for the self-hosted-collector OTLP path (it falls back to console-only tracing via a graceful ImportError otherwise, so the primary Pisama-platform-ingest path was unaffected either way). And one gap found while re-verifying "pytest with all extras installed" from a clean install: the `dev` extra was missing anthropic/openai, which tests/test_auto_sdk_instrumentation.py imports directly to exercise the patches. Carried over from pisama-auto's own dev extra for the same reason; not a runtime dependency of the `auto` extra itself. pisama-auto and pisama-agent-sdk keep publishing as standalone distributions in this release; they become thin shims re-exporting these modules in a later, separate step once this lands and pisama>=0.6.0 is live on PyPI. Verified: 196/196 tests pass, ruff clean, mypy clean on src/pisama/auto and src/pisama/agents, wheel contains both submodules, both console scripts work, bare install (no extras) still imports cleanly with opentelemetry/wrapt/posthog/anthropic/openai all absent. Co-Authored-By: Claude Sonnet 5 --- CHANGELOG.md | 23 + pyproject.toml | 39 +- src/pisama/__init__.py | 2 +- src/pisama/agents/__init__.py | 212 +++++ src/pisama/agents/atif.py | 699 +++++++++++++++++ src/pisama/agents/bridge.py | 592 ++++++++++++++ src/pisama/agents/chaos/__init__.py | 23 + src/pisama/agents/chaos/config.py | 43 + src/pisama/agents/chaos/experiments.py | 166 ++++ src/pisama/agents/check.py | 247 ++++++ src/pisama/agents/check_compliance.py | 280 +++++++ src/pisama/agents/clarification.py | 238 ++++++ src/pisama/agents/cli/__init__.py | 8 + src/pisama/agents/cli/openhands_monitor.py | 129 +++ src/pisama/agents/config.py | 210 +++++ src/pisama/agents/converter.py | 196 +++++ src/pisama/agents/evaluator.py | 197 +++++ src/pisama/agents/heal.py | 245 ++++++ src/pisama/agents/hooks/__init__.py | 28 + src/pisama/agents/hooks/matchers.py | 122 +++ src/pisama/agents/hooks/post_tool_use.py | 117 +++ src/pisama/agents/hooks/pre_tool_use.py | 244 ++++++ src/pisama/agents/indication.py | 300 +++++++ src/pisama/agents/openhands_adapter.py | 482 ++++++++++++ src/pisama/agents/session.py | 250 ++++++ src/pisama/agents/tools.py | 144 ++++ src/pisama/agents/types.py | 144 ++++ src/pisama/auto/__init__.py | 91 +++ src/pisama/auto/_tracer.py | 264 +++++++ src/pisama/auto/patches/__init__.py | 70 ++ src/pisama/auto/patches/anthropic_patch.py | 154 ++++ src/pisama/auto/patches/openai_patch.py | 83 ++ tests/agents/__init__.py | 1 + tests/agents/fixtures/harbor/PROVENANCE.toml | 21 + ...tion-linear-history.trajectory.cont-1.json | 118 +++ ...mmarization-linear-history.trajectory.json | 120 +++ ...on.trajectory.summarization-1-summary.json | 103 +++ tests/agents/test_atif.py | 733 ++++++++++++++++++ tests/agents/test_bridge.py | 382 +++++++++ tests/agents/test_chaos_hooks.py | 168 ++++ tests/agents/test_check_compliance.py | 210 +++++ tests/agents/test_heal.py | 45 ++ tests/agents/test_openhands_adapter.py | 356 +++++++++ tests/agents/test_real_public_workflows.py | 348 +++++++++ tests/test_auto_platform_exporter.py | 225 ++++++ tests/test_auto_sdk_instrumentation.py | 299 +++++++ 46 files changed, 9168 insertions(+), 3 deletions(-) create mode 100644 src/pisama/agents/__init__.py create mode 100644 src/pisama/agents/atif.py create mode 100644 src/pisama/agents/bridge.py create mode 100644 src/pisama/agents/chaos/__init__.py create mode 100644 src/pisama/agents/chaos/config.py create mode 100644 src/pisama/agents/chaos/experiments.py create mode 100644 src/pisama/agents/check.py create mode 100644 src/pisama/agents/check_compliance.py create mode 100644 src/pisama/agents/clarification.py create mode 100644 src/pisama/agents/cli/__init__.py create mode 100644 src/pisama/agents/cli/openhands_monitor.py create mode 100644 src/pisama/agents/config.py create mode 100644 src/pisama/agents/converter.py create mode 100644 src/pisama/agents/evaluator.py create mode 100644 src/pisama/agents/heal.py create mode 100644 src/pisama/agents/hooks/__init__.py create mode 100644 src/pisama/agents/hooks/matchers.py create mode 100644 src/pisama/agents/hooks/post_tool_use.py create mode 100644 src/pisama/agents/hooks/pre_tool_use.py create mode 100644 src/pisama/agents/indication.py create mode 100644 src/pisama/agents/openhands_adapter.py create mode 100644 src/pisama/agents/session.py create mode 100644 src/pisama/agents/tools.py create mode 100644 src/pisama/agents/types.py create mode 100644 src/pisama/auto/__init__.py create mode 100644 src/pisama/auto/_tracer.py create mode 100644 src/pisama/auto/patches/__init__.py create mode 100644 src/pisama/auto/patches/anthropic_patch.py create mode 100644 src/pisama/auto/patches/openai_patch.py create mode 100644 tests/agents/__init__.py create mode 100644 tests/agents/fixtures/harbor/PROVENANCE.toml create mode 100644 tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.cont-1.json create mode 100644 tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.json create mode 100644 tests/agents/fixtures/harbor/hello-world-context-summarization.trajectory.summarization-1-summary.json create mode 100644 tests/agents/test_atif.py create mode 100644 tests/agents/test_bridge.py create mode 100644 tests/agents/test_chaos_hooks.py create mode 100644 tests/agents/test_check_compliance.py create mode 100644 tests/agents/test_heal.py create mode 100644 tests/agents/test_openhands_adapter.py create mode 100644 tests/agents/test_real_public_workflows.py create mode 100644 tests/test_auto_platform_exporter.py create mode 100644 tests/test_auto_sdk_instrumentation.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a62794..7f77a0c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,29 @@ All notable changes to the `pisama` meta-package are documented here. The packag ## [Unreleased] +## [0.6.0] - 2026-07-29 + +### Added + +- `pisama.auto` -- consolidated in-package original of the standalone + `pisama-auto` distribution (zero-code auto-instrumentation for the + Anthropic/OpenAI SDKs, tracing to either the Pisama platform or a + self-hosted OTLP collector). New `auto` extra: + `pip install "pisama[auto]"`. +- `pisama.agents` -- consolidated in-package original of the standalone + `pisama-agent-sdk` distribution (real-time hooks, tools, and a + `pisama-openhands-monitor` console script for Claude Agent SDK / Harbor + integrations). New `agents` and `telemetry` extras: + `pip install "pisama[agents]"`. +- `pisama.agents.__version__`, aliased to `pisama.__version__` since this + code no longer releases independently. + +Both submodules keep publishing as standalone distributions +(`pisama-auto`, `pisama-agent-sdk`) for this release; those packages become +thin shims re-exporting this module in a later release. A bare +`pip install pisama` pulls in no new required dependencies -- `pisama.auto` +and `pisama.agents` stay behind their respective extras. + ## [0.5.7] - 2026-07-29 ### Changed diff --git a/pyproject.toml b/pyproject.toml index 96d2618..11482f7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "pisama" -version = "0.5.7" +version = "0.6.0" description = "Multi-agent failure detection for production AI systems" readme = "README.md" license = "MIT" @@ -46,8 +46,33 @@ dependencies = [ # server.py:138 with AttributeError: 'Server' object has no attribute # 'list_tools'. Lift this only once the server is ported to the 2.x API. mcp = ["mcp>=1.0.0,<2"] -auto = ["pisama-auto>=0.1.0", "opentelemetry-api>=1.20.0", "opentelemetry-sdk>=1.20.0"] +# pisama.auto is now original code inside this package (see src/pisama/auto/), +# not an import of the standalone pisama-auto distribution -- pisama-auto stays +# published separately and, in a later release, becomes a thin shim re-exporting +# this module. wrapt is the function-wrapping library pisama.auto's patches +# submodule needs to monkeypatch the anthropic/openai clients. +auto = [ + "opentelemetry-api>=1.20.0", + "opentelemetry-sdk>=1.20.0", + "opentelemetry-exporter-otlp-proto-http>=1.20.0", + "wrapt>=1.15.0", +] langgraph = ["langgraph>=0.2.0", "langchain-core>=0.3.0"] +# pisama.agents mirrors pisama-agent-sdk's own optional-extra split: httpx is +# only needed for the HTTP-backed calls (analyze_atif, PisamaEvaluator), and +# posthog only for opt-in SDK-side telemetry in DetectionBridge. Neither is +# imported at module load time (see pisama/agents/atif.py, evaluator.py, +# bridge.py) so `import pisama` and `import pisama.agents` both work without +# either installed; only the specific calls that need them raise ImportError +# with a `pip install pisama[agents]` hint until the extra is added. httpx is +# already a hard dependency of the base `pisama` install (see `dependencies` +# above) for unrelated reasons, so this extra is declared for documentation +# and for behavioral parity with pisama-agent-sdk's `evaluator` extra, not +# because it changes what a bare `pip install pisama` pulls in today. +agents = ["httpx>=0.24"] +# Pass-through of pisama-agent-sdk's own `telemetry` extra (opt-in PostHog +# telemetry in pisama.agents.bridge.DetectionBridge). +telemetry = ["posthog>=3.0"] dev = [ "build>=1.2", "mypy>=1.10", @@ -56,10 +81,20 @@ dev = [ "pytest-cov>=5.0", "ruff>=0.4", "twine>=5.0", + # Test-only: tests/test_auto_sdk_instrumentation.py patches these SDK + # clients directly. Carried over from pisama-auto's own `dev` extra; + # not a runtime dependency of the `auto` extra itself, since users + # bring their own anthropic/openai client. + "anthropic>=0.40", + "openai>=1.50", ] [project.scripts] pisama = "pisama.cli.main:main" +# From pisama-agent-sdk: batch analysis of a completed OpenHands / +# Harbor-on-OpenHands session directory. Needs the `agents` extra installed +# to reach the network call inside analyze_atif(); `--help` works bare. +pisama-openhands-monitor = "pisama.agents.cli.openhands_monitor:main" [project.urls] Homepage = "https://pisama.ai" diff --git a/src/pisama/__init__.py b/src/pisama/__init__.py index 995a527..9ac6d5f 100644 --- a/src/pisama/__init__.py +++ b/src/pisama/__init__.py @@ -14,7 +14,7 @@ __version__ = _pkg_version("pisama") except PackageNotFoundError: # running from a source checkout without an install - __version__ = "0.5.7" + __version__ = "0.6.0" from pisama._analyze import AnalyzeResult, Issue, analyze, async_analyze from pisama._http import PisamaAuthError diff --git a/src/pisama/agents/__init__.py b/src/pisama/agents/__init__.py new file mode 100644 index 0000000..b127e77 --- /dev/null +++ b/src/pisama/agents/__init__.py @@ -0,0 +1,212 @@ +"""Pisama Agents -- real-time hooks and tools for agent runtimes. + +Optional submodule of the ``pisama`` package (extra: ``pisama[agents]``). +Provides hooks and tools for Claude Agent SDK that connect to Pisama's +detection infrastructure for real-time failure prevention. + +Mirrors the standalone ``pisama-agent-sdk`` distribution's public API. +That package keeps shipping independently; this module is the +consolidated, in-package original (not a re-export or a dependency on +the standalone package -- see the module-level docstrings for the +per-symbol origin). + +Quick Start (passive monitoring): + from pisama.agents import pre_tool_use_hook, post_tool_use_hook + + agent.hooks.pre_tool_use = pre_tool_use_hook + agent.hooks.post_tool_use = post_tool_use_hook + +Agent Self-Check (active verification): + from pisama.agents import check + + result = await check( + output="The server is healthy based on the metrics.", + context={"query": "Is auth-service down?", "sources": [...]} + ) + if not result["passed"]: + # Revise output based on result["issues"] + +Claude Agent SDK Custom Tool: + from pisama.agents import create_check_tool + from claude_agent_sdk import ClaudeAgentOptions + + options = ClaudeAgentOptions( + custom_tools=[create_check_tool()], + ) +""" + +# This code no longer releases independently (it shipped standalone as +# pisama-agent-sdk through 0.3.2), so __version__ tracks the installed +# base-package version rather than a separate hardcoded number. +from pisama import __version__ + +# Hook functions (primary API) +# ATIF (Harbor) trajectory analysis +from pisama.agents.atif import ( + AtifAnalyzeResult, + AtifDetection, + analyze_atif, + analyze_atif_batch, +) + +# Auto-verify was never mirrored here: pisama-agent-sdk removed it in 0.3.0 +# because it vendored private backend verification primitives into an MIT +# package. Verification stays a hosted-service concern. +# Configuration +# Bridge (for advanced use) +from pisama.agents.bridge import DetectionBridge, configure_bridge, create_bridge, get_bridge + +# Chaos engineering (SDK-level failure injection) +from pisama.agents.chaos import ( + ChaosConfig, + ContextTruncation, + ErrorInjection, + LatencyInjection, + OutputCorruption, + ToolFailure, +) + +# Agent self-check +from pisama.agents.check import check, configure_check + +# Specification compliance (beta, gated by PISAMA_ENABLE_CHECK_COMPLIANCE) +from pisama.agents.check_compliance import ( + BehavioralRule, + ComplianceResult, + PisamaFeatureNotEnabledError, + Violation, + check_compliance, +) + +# Clarification primitive -- pause/ask/resume for entity_confusion etc. +from pisama.agents.clarification import ( + ClarificationPrimitive, + ClarificationRequest, + register_clarification_builder, +) +from pisama.agents.clarification import ( + Resolution as ClarificationResolution, +) +from pisama.agents.config import BridgeConfig, load_config + +# Evaluator client (Pisama-as-evaluator for multi-agent harnesses) +from pisama.agents.evaluator import EvalFailure, EvalResult, PisamaEvaluator + +# In-loop healing +from pisama.agents.heal import HealingResult, heal_now + +# Matchers +from pisama.agents.hooks.matchers import ( + AGENT_TOOLS, + ALL_TOOLS, + DANGEROUS_COMMANDS, + FILE_TOOLS, + SHELL_TOOLS, + HookMatcher, + create_matcher, +) +from pisama.agents.hooks.post_tool_use import PostToolUseHook, post_tool_use_hook +from pisama.agents.hooks.pre_tool_use import PreToolUseHook, pre_tool_use_hook + +# Indication channel -- out-of-band signal for the developer running the +# agent. Wire on_indication(callable) to receive structured notifications +# on every healing outcome. +from pisama.agents.indication import ( + SDKIndication, + clear_indication_callbacks, + on_indication, +) + +# OpenHands event-stream adapter +from pisama.agents.openhands_adapter import ( + OpenHandsEventStreamAdapter, + StreamingCallback, + StreamingDetection, +) + +# Session management +from pisama.agents.session import SessionManager, session_manager + +# Custom tools for Claude Agent SDK +from pisama.agents.tools import create_check_tool, pisama_check_handler + +# Types +from pisama.agents.types import BridgeResult, HookContext, HookInput, HookJSONOutput + +__all__ = [ + "__version__", + # Hook functions + "pre_tool_use_hook", + "post_tool_use_hook", + # Hook classes + "PreToolUseHook", + "PostToolUseHook", + # Configuration + "configure_bridge", + "create_bridge", + "get_bridge", + "BridgeConfig", + "load_config", + # Bridge + "DetectionBridge", + # Types + "BridgeResult", + "HookInput", + "HookContext", + "HookJSONOutput", + # Matchers + "HookMatcher", + "ALL_TOOLS", + "FILE_TOOLS", + "SHELL_TOOLS", + "DANGEROUS_COMMANDS", + "AGENT_TOOLS", + "create_matcher", + # Session + "SessionManager", + "session_manager", + # Agent self-check + "check", + "configure_check", + # In-loop healing + "heal_now", + "HealingResult", + # Clarification primitive + "ClarificationPrimitive", + "ClarificationRequest", + "ClarificationResolution", + "register_clarification_builder", + # Specification compliance (beta) + "check_compliance", + "ComplianceResult", + "BehavioralRule", + "Violation", + "PisamaFeatureNotEnabledError", + # Indication channel + "SDKIndication", + "on_indication", + "clear_indication_callbacks", + # Custom tools + "create_check_tool", + "pisama_check_handler", + # Evaluator + "PisamaEvaluator", + "EvalResult", + "EvalFailure", + # ATIF (Harbor) trajectory analysis + "analyze_atif", + "analyze_atif_batch", + "AtifAnalyzeResult", + "AtifDetection", + # OpenHands event-stream adapter + "OpenHandsEventStreamAdapter", + "StreamingDetection", + "StreamingCallback", + # Chaos engineering + "ChaosConfig", + "ToolFailure", + "LatencyInjection", + "ErrorInjection", + "OutputCorruption", + "ContextTruncation", +] diff --git a/src/pisama/agents/atif.py b/src/pisama/agents/atif.py new file mode 100644 index 0000000..02993c9 --- /dev/null +++ b/src/pisama/agents/atif.py @@ -0,0 +1,699 @@ +"""ATIF trajectory analysis client. + +Wraps the Pisama backend's ``POST /api/v1/atif/analyze`` endpoint so +notebook/script users can analyze Harbor-emitted ATIF trajectories +without going through the CLI. + +Usage:: + + from pisama.agents import analyze_atif + + result = analyze_atif("./trajectories/run-001.json") + if not result.analysis_complete: + raise RuntimeError("ATIF analysis was incomplete") + elif result.has_failures: + for d in result.failures: + print(d.detector, d.severity, d.title) + +For directories or many files:: + + from pisama.agents import analyze_atif_batch + + results = analyze_atif_batch("./trajectories/") + for r in results: + print(r.trace_id, r.failure_count) +""" + +from __future__ import annotations + +import json +import logging +import os +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable + +logger = logging.getLogger(__name__) + +try: + import httpx + + _HTTPX_AVAILABLE = True +except ImportError: + _HTTPX_AVAILABLE = False + + +DEFAULT_API_URL = "https://api.pisama.ai" +DEFAULT_TIMEOUT_SECONDS = 60.0 + + +# Keep in lockstep with backend/app/ingestion/atif_models.py — bump when the +# vendor pin updates. +SUPPORTED_SCHEMA_VERSIONS = frozenset( + { + "ATIF-v1.0", + "ATIF-v1.1", + "ATIF-v1.2", + "ATIF-v1.3", + "ATIF-v1.4", + "ATIF-v1.5", + "ATIF-v1.6", + "ATIF-v1.7", + } +) + + +@dataclass +class AtifDetection: + """One detection surfaced by the orchestrator.""" + + detector: str + confidence: float + severity: str + title: str + description: str + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "AtifDetection": + # Backend's DetectionResult.to_dict() returns the detector under + # `category` (the DetectionCategory enum value). Accept the older + # `detector` / `detection_type` keys too in case the response + # schema is normalized later. + name = ( + data.get("category") + or data.get("detector") + or data.get("detection_type") + or "unknown" + ) + return cls( + detector=str(name), + confidence=float(data.get("confidence", 0.0)), + severity=str(data.get("severity", "low")), + title=str(data.get("title", "")), + description=str(data.get("description", "")), + ) + + +@dataclass +class AtifAnalyzeResult: + """One trajectory's diagnosis.""" + + trace_id: str + source_path: str | None + span_count: int + total_tokens: int + total_duration_ms: int + atif_schema_version: str + atif_session_id: str | None + has_failures: bool + failure_count: int + detection_status: str + failures: list[AtifDetection] = field(default_factory=list) + detectors_run: list[str] = field(default_factory=list) + detectors_failed: dict[str, str] = field(default_factory=dict) + # False when the backend could not see every trajectory segment referenced + # by this ATIF document. Detector execution can still be complete over the + # visible spans, so callers should use ``analysis_complete`` for a single + # trustworthy success/failure gate. + topology_complete: bool = True + unresolved_trajectory_refs: list[str] = field(default_factory=list) + # Subset of ``unresolved_trajectory_refs`` whose safe, explicit + # continuation targets were submitted as separate documents in the same + # batch. The backend fields above remain unchanged so callers can inspect + # each response exactly as the server reported it. + client_resolved_trajectory_refs: list[str] = field(default_factory=list) + # Sum of tokens in spans parsed from the submitted document. This can be + # lower than ``total_tokens`` when ATIF ``final_metrics`` includes external + # subagent or continuation trajectories that were not submitted inline. + span_token_total: int = 0 + # Scalar PROCESS-QUALITY score in [0.0, 1.0]; 1.0 = clean, lower = worse. + # + # NOT a task-success score. Pisama Bench v1 (May 2026, n=270 across 5 + # corpora) shows the same scalar correlates differently with task + # outcome depending on corpus type: + # - multi-agent reasoning (m500): Pearson r = +0.45 with correctness + # - role-play dialogue (sotopia): r = -0.44 (process noise IS the goal) + # - single-agent web tasks (ARB): r ≈ -0.03 (uncorrelated) + # + # If you need "did the agent solve the task?", combine this with an + # outcome signal from your own evaluator (test pass, reward, completion + # check). The Tier 2 follow-up adds a separate task_completion_score + # scalar for the outcome axis. + # + # Multiplicative composite of orchestration_quality_overall and the max + # severity-weighted detector penalty. See orchestrator.DetectionOrchestrator + # ._compute_trajectory_score for the formula. Defaults to 1.0 for backwards + # compat with older server versions that don't yet emit the field. + trajectory_score: float = 1.0 + # Scalar OUTCOME-QUALITY score in [0.0, 1.0] — "did the agent solve the + # user's task?". Orthogonal to trajectory_score (process quality): a run can + # execute cleanly (high trajectory_score) yet fail the user's job (low + # task_completion_score). This is the AgentRewardBench gap — web agents + # whose runs are clean while cum_reward ≈ 0.11. + # + # Driven by the Tier 2 outcome detector family (task_failure / + # silent_failure / objective_unmet); 1.0 when that family is disabled + # server-side or none fire. Defaults to 1.0 for backwards compat with + # older servers that don't yet emit the field. See orchestrator + # ._compute_task_completion_score for the multiplicative formula. + task_completion_score: float = 1.0 + score_breakdown: dict[str, Any] = field(default_factory=dict) + raw: dict[str, Any] = field(default_factory=dict) + + @property + def remaining_unresolved_trajectory_refs(self) -> list[str]: + """Server-reported references not reconciled by this client batch.""" + resolved = set(self.client_resolved_trajectory_refs) + return [ref for ref in self.unresolved_trajectory_refs if ref not in resolved] + + @property + def reconciled_topology_complete(self) -> bool: + """Whether server topology plus safe batch submissions are complete.""" + if self.remaining_unresolved_trajectory_refs: + return False + return self.topology_complete or bool(self.unresolved_trajectory_refs) + + @property + def analysis_complete(self) -> bool: + """Whether detector execution and trajectory topology are complete.""" + return ( + self.reconciled_topology_complete + and self.detection_status.lower() == "complete" + and not self.detectors_failed + ) + + @classmethod + def from_response( + cls, body: dict[str, Any], source_path: str | Path | None = None + ) -> "AtifAnalyzeResult": + diag = body.get("diagnosis", {}) + trace = body.get("trace", {}) + failures = [AtifDetection.from_dict(d) for d in diag.get("all_detections", [])] + total_tokens = int(trace.get("total_tokens", 0)) + unresolved_refs = trace.get("unresolved_trajectory_refs", []) + if not isinstance(unresolved_refs, list): + unresolved_refs = [] + return cls( + trace_id=str(trace.get("trace_id", diag.get("trace_id", ""))), + source_path=str(source_path) if source_path is not None else None, + span_count=int(trace.get("span_count", 0)), + total_tokens=total_tokens, + total_duration_ms=int(trace.get("total_duration_ms") or 0), + atif_schema_version=str(trace.get("atif_schema_version", "")), + atif_session_id=trace.get("atif_session_id"), + has_failures=bool(diag.get("has_failures", False)), + failure_count=int(diag.get("failure_count", 0)), + detection_status=str(diag.get("detection_status", "unknown")), + failures=failures, + detectors_run=list(diag.get("detectors_run", [])), + detectors_failed=dict(diag.get("detectors_failed", {})), + topology_complete=bool(trace.get("topology_complete", True)), + unresolved_trajectory_refs=[str(ref) for ref in unresolved_refs], + span_token_total=int(trace.get("span_token_total", total_tokens)), + trajectory_score=float(diag.get("trajectory_score", 1.0)), + task_completion_score=float(diag.get("task_completion_score", 1.0)), + score_breakdown=dict(diag.get("score_breakdown", {})), + raw=body, + ) + + +def _load_trajectory(source: str | Path | dict[str, Any]) -> dict[str, Any]: + """Resolve ``source`` to a trajectory dict. + + Accepts a path (``str``/``Path``) or an already-parsed ATIF dict so + callers can pre-process the JSON if needed. + """ + if isinstance(source, dict): + traj = source + else: + path = Path(source) + with path.open("r", encoding="utf-8") as f: + traj = json.load(f) + schema = traj.get("schema_version") + if not schema or schema not in SUPPORTED_SCHEMA_VERSIONS: + raise ValueError( + f"Unsupported ATIF schema_version {schema!r}. " + f"Expected one of: {sorted(SUPPORTED_SCHEMA_VERSIONS)}" + ) + continued_ref = traj.get("continued_trajectory_ref") + if isinstance(continued_ref, str) and Path(continued_ref).is_absolute(): + raise ValueError( + f"ATIF continued_trajectory_ref must be relative: {continued_ref!r}" + ) + return traj + + +def _require_httpx() -> None: + if not _HTTPX_AVAILABLE: + raise ImportError( + "httpx is required for analyze_atif: `pip install pisama[agents]`" + ) + + +def _resolve_api_key(api_key: str | None) -> str | None: + """Resolve an explicit API key or the SDK-wide environment fallback. + + An explicit empty string disables the environment fallback. This keeps + unauthenticated local and test endpoints usable even when the parent + process has ``PISAMA_API_KEY`` configured. + """ + candidate = api_key if api_key is not None else os.getenv("PISAMA_API_KEY") + if candidate is None: + return None + candidate = candidate.strip() + return candidate or None + + +def _authorization_headers( + client: Any, + *, + base_url: str, + api_key: str | None, +) -> dict[str, str]: + """Exchange an API key for the JWT accepted by authenticated endpoints.""" + resolved_key = _resolve_api_key(api_key) + if resolved_key is None: + return {} + + response = client.post( + f"{base_url}/api/v1/auth/token", + json={"api_key": resolved_key, "scope": "full"}, + ) + response.raise_for_status() + token_body = response.json() + access_token = ( + token_body.get("access_token") if isinstance(token_body, dict) else None + ) + if not isinstance(access_token, str) or not access_token.strip(): + raise ValueError( + "Pisama authentication response did not include a valid access_token" + ) + return {"Authorization": f"Bearer {access_token.strip()}"} + + +def _analyze_with_client( + client: Any, + trajectory: dict[str, Any], + *, + source_path: str | None, + project_id: str | None, + base_url: str, + headers: dict[str, str], +) -> AtifAnalyzeResult: + """Analyze one trajectory using an already-authenticated HTTP client.""" + payload: dict[str, Any] = {"trajectory": trajectory} + if project_id is not None: + payload["project_id"] = project_id + + response = client.post( + f"{base_url}/api/v1/atif/analyze", + json=payload, + headers=headers, + ) + response.raise_for_status() + + return AtifAnalyzeResult.from_response(response.json(), source_path=source_path) + + +def analyze_atif( + source: str | Path | dict[str, Any], + *, + project_id: str | None = None, + api_url: str | None = None, + api_key: str | None = None, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> AtifAnalyzeResult: + """Analyze a single ATIF trajectory with Pisama's detectors. + + Args: + source: Path to a ``.json`` ATIF trajectory file, or a pre-parsed + trajectory dict. + project_id: Optional Pisama project id (ps_...) to correlate the + run with. Reserved for future persistence; not used in v0. + api_url: Override the Pisama backend URL. Default + ``https://api.pisama.ai``. + api_key: Pisama API key. When omitted, ``PISAMA_API_KEY`` is used. + The key is exchanged for a short-lived JWT before analysis. Pass + an empty string to disable the environment fallback for an + unauthenticated local or test endpoint. + timeout: HTTP timeout in seconds. + + Returns: + ``AtifAnalyzeResult`` with diagnosis details. + + Raises: + ValueError: If the trajectory's ``schema_version`` is unsupported. + ImportError: If httpx is not installed (`pip install pisama[agents]`). + httpx.HTTPStatusError: If the backend returns a non-2xx response. + """ + _require_httpx() + trajectory = _load_trajectory(source) + source_path = str(source) if isinstance(source, (str, Path)) else None + base_url = (api_url or DEFAULT_API_URL).rstrip("/") + + with httpx.Client(timeout=timeout) as client: + headers = _authorization_headers( + client, + base_url=base_url, + api_key=api_key, + ) + return _analyze_with_client( + client, + trajectory, + source_path=source_path, + project_id=project_id, + base_url=base_url, + headers=headers, + ) + + +def _discover_trajectory_files(root: Path) -> list[Path]: + """Find trajectory JSON files under ``root``. + + Three modes, tried in order: + 1. Harbor trial dir: ``root/agent/trajectory.json``. + 2. Harbor job-output dir: recursive exact ``**/trajectory.json``. + 3. Flat dir of ``*.json`` trajectories. + + Harbor job roots contain top-level ``config.json``, ``lock.json``, and + ``result.json`` files. Recursive trajectory discovery must therefore run + before the flat-directory fallback, or those metadata files are mistaken + for ATIF trajectories. + + Exact ``trajectory.json`` files are Harbor trial roots. Each root's + explicit ``continued_trajectory_ref`` chain is followed inside that + root's agent directory. Files such as + ``trajectory.summarization-1-summary.json`` are therefore not submitted + as independent trajectories. + + This ordering is the canonical Harbor discovery behavior for SDK clients: + nested trajectory files take precedence over job-level metadata. + """ + try: + directory_boundary = root.resolve(strict=True) + except OSError as exc: + raise ValueError(f"ATIF directory could not be resolved: {root}") from exc + if not directory_boundary.is_dir(): + raise ValueError(f"ATIF directory is not a directory: {root}") + _reject_escaping_agent_symlinks(root, directory_boundary) + + trial_file = root / "agent" / "trajectory.json" + if trial_file.is_file(): + return _expand_continuation_chains( + [trial_file], directory_boundary=directory_boundary + ) + recursive = sorted(root.rglob("trajectory.json")) + if recursive: + return _expand_continuation_chains( + recursive, directory_boundary=directory_boundary + ) + flat = sorted( + path for path in root.glob("*.json") if not _is_harbor_helper_trajectory(path) + ) + if flat: + return _expand_continuation_chains(flat, directory_boundary=directory_boundary) + return [] + + +def _reject_escaping_agent_symlinks(root: Path, directory_boundary: Path) -> None: + """Reject Harbor agent directory links that leave the selected tree.""" + for agent_dir in root.rglob("agent"): + if not agent_dir.is_symlink(): + continue + try: + resolved = agent_dir.resolve(strict=True) + except OSError as exc: + raise ValueError( + f"ATIF agent directory symlink could not be resolved: {agent_dir}" + ) from exc + if not resolved.is_dir() or not _is_relative_to(resolved, directory_boundary): + raise ValueError( + f"ATIF agent directory symlink escapes selected directory: {agent_dir}" + ) + + +def _is_harbor_helper_trajectory(path: Path) -> bool: + """Exclude Harbor's non-root context summarization trajectory files.""" + name = path.name.lower() + return ( + ".summarization-" in name + or ".trajectory.cont-" in name + or name.startswith("trajectory.cont-") + ) + + +def _expand_continuation_chains( + roots: Iterable[Path], *, directory_boundary: Path | None = None +) -> list[Path]: + """Return roots plus their safe, explicit continuation chains.""" + discovered: list[Path] = [] + emitted: set[Path] = set() + for root in roots: + for path in _follow_continuation_chain( + root, directory_boundary=directory_boundary + ): + canonical = path.resolve() + if canonical not in emitted: + discovered.append(path) + emitted.add(canonical) + return discovered + + +def _follow_continuation_chain( + root: Path, *, directory_boundary: Path | None = None +) -> list[Path]: + """Follow one ATIF continuation chain without leaving its agent directory.""" + agent_dir, current = _resolve_chain_root(root, directory_boundary) + chain: list[Path] = [] + seen: set[Path] = set() + + while True: + _validate_chain_member(current, agent_dir, seen) + seen.add(current) + chain.append(current) + candidate = _next_continuation(current, agent_dir) + if candidate is None: + return chain + current = candidate + + +def _resolve_chain_root( + root: Path, + directory_boundary: Path | None, +) -> tuple[Path, Path]: + try: + agent_dir = root.parent.resolve(strict=True) + current = root.resolve(strict=True) + except OSError as exc: + raise ValueError(f"ATIF trajectory file could not be resolved: {root}") from exc + + if directory_boundary is not None and ( + not _is_relative_to(agent_dir, directory_boundary) + or not _is_relative_to(current, directory_boundary) + ): + raise ValueError( + f"ATIF trajectory root or agent directory escapes selected directory: " + f"{root}" + ) + return agent_dir, current + + +def _validate_chain_member( + current: Path, + agent_dir: Path, + seen: set[Path], +) -> None: + if not _is_relative_to(current, agent_dir): + raise ValueError( + f"ATIF continued_trajectory_ref escapes agent directory: {current}" + ) + if current in seen: + raise ValueError( + f"ATIF continued_trajectory_ref cycle detected at: {current}" + ) + if not current.is_file(): + raise ValueError(f"ATIF continuation is not a file: {current}") + + +def _next_continuation(current: Path, agent_dir: Path) -> Path | None: + continued_ref = _read_continued_trajectory_ref(current) + if continued_ref is None: + return None + if Path(continued_ref).is_absolute(): + raise ValueError( + f"ATIF continued_trajectory_ref must be relative: " + f"{continued_ref!r} in {current}" + ) + + try: + candidate = (current.parent / continued_ref).resolve(strict=False) + except (OSError, ValueError) as exc: + raise ValueError( + f"Invalid ATIF continued_trajectory_ref {continued_ref!r} in {current}" + ) from exc + if not _is_relative_to(candidate, agent_dir): + raise ValueError( + f"ATIF continued_trajectory_ref escapes agent directory: " + f"{continued_ref!r} in {current}" + ) + if not candidate.exists(): + raise ValueError( + f"Missing ATIF continued_trajectory_ref {continued_ref!r} in {current}" + ) + return candidate.resolve(strict=True) + + +def _read_continued_trajectory_ref(path: Path) -> str | None: + """Read the sole filesystem relationship clients are allowed to follow.""" + try: + with path.open("r", encoding="utf-8") as f: + trajectory = json.load(f) + except (OSError, json.JSONDecodeError) as exc: + raise ValueError(f"Could not read ATIF trajectory JSON: {path}") from exc + + if not isinstance(trajectory, dict): + raise ValueError(f"ATIF trajectory must be a JSON object: {path}") + continued_ref = trajectory.get("continued_trajectory_ref") + if continued_ref is None: + return None + if not isinstance(continued_ref, str) or not continued_ref.strip(): + raise ValueError( + f"ATIF continued_trajectory_ref must be a non-empty string: {path}" + ) + return continued_ref + + +def _is_relative_to(path: Path, root: Path) -> bool: + """Python 3.10-compatible containment check.""" + try: + path.relative_to(root) + except ValueError: + return False + return True + + +def _reconcile_submitted_continuations( + submissions: list[tuple[str | Path | dict[str, Any], dict[str, Any]]], + results: list[AtifAnalyzeResult], +) -> None: + """Mark server-unresolved continuation refs submitted in this batch.""" + submitted_paths = { + Path(source).resolve(strict=True) + for source, _trajectory in submissions + if isinstance(source, (str, Path)) + } + + for (source, trajectory), result in zip(submissions, results): + if not isinstance(source, (str, Path)): + continue + continued_ref = trajectory.get("continued_trajectory_ref") + if not isinstance(continued_ref, str): + continue + if continued_ref in _subagent_trajectory_ref_targets(trajectory): + continue + + source_path = Path(source).resolve(strict=True) + agent_dir = source_path.parent + target = (agent_dir / continued_ref).resolve(strict=False) + if ( + _is_relative_to(target, agent_dir) + and target in submitted_paths + and continued_ref in result.unresolved_trajectory_refs + ): + result.client_resolved_trajectory_refs.append(continued_ref) + + +def _subagent_trajectory_ref_targets(value: Any) -> set[str]: + """Collect path and id targets from nested ATIF subagent references.""" + targets: set[str] = set() + + def collect(current: Any) -> None: + if isinstance(current, dict): + refs = current.get("subagent_trajectory_ref") + if isinstance(refs, list): + for ref in refs: + if not isinstance(ref, dict): + continue + for key in ("trajectory_path", "trajectory_id"): + target = ref.get(key) + if isinstance(target, str): + targets.add(target) + for nested in current.values(): + collect(nested) + elif isinstance(current, list): + for nested in current: + collect(nested) + + collect(value) + return targets + + +def analyze_atif_batch( + sources: str | Path | Iterable[str | Path | dict[str, Any]], + *, + project_id: str | None = None, + api_url: str | None = None, + api_key: str | None = None, + timeout: float = DEFAULT_TIMEOUT_SECONDS, +) -> list[AtifAnalyzeResult]: + """Analyze many ATIF trajectories sequentially. + + Args: + sources: One of: + - A file path (string or ``Path``); its explicit continuation + chain is included. + - A directory: a single Harbor trial dir + (``/agent/trajectory.json``), a flat dir of + ``*.json`` trajectories, or a Harbor job-output dir (walked + recursively for exact ``trajectory.json`` roots). Explicit + continuation chains are included; summarization helpers are not. + - An iterable of file paths and/or pre-parsed dicts. + + Other args mirror :func:`analyze_atif`. Returns one + :class:`AtifAnalyzeResult` per source in input order. + """ + paths_or_dicts: list[str | Path | dict[str, Any]] + if isinstance(sources, (str, Path)): + p = Path(sources) + if p.is_dir(): + paths_or_dicts = list(_discover_trajectory_files(p)) + else: + paths_or_dicts = list(_expand_continuation_chains([p])) + else: + paths_or_dicts = list(sources) + + if not paths_or_dicts: + return [] + + submissions = [(src, _load_trajectory(src)) for src in paths_or_dicts] + _require_httpx() + base_url = (api_url or DEFAULT_API_URL).rstrip("/") + with httpx.Client(timeout=timeout) as client: + headers = _authorization_headers( + client, + base_url=base_url, + api_key=api_key, + ) + results = [ + _analyze_with_client( + client, + trajectory, + source_path=str(src) if isinstance(src, (str, Path)) else None, + project_id=project_id, + base_url=base_url, + headers=headers, + ) + for src, trajectory in submissions + ] + _reconcile_submitted_continuations(submissions, results) + return results + + +__all__ = [ + "AtifAnalyzeResult", + "AtifDetection", + "DEFAULT_API_URL", + "SUPPORTED_SCHEMA_VERSIONS", + "analyze_atif", + "analyze_atif_batch", +] diff --git a/src/pisama/agents/bridge.py b/src/pisama/agents/bridge.py new file mode 100644 index 0000000..0264a8e --- /dev/null +++ b/src/pisama/agents/bridge.py @@ -0,0 +1,592 @@ +"""Detection Bridge - connects Agent SDK hooks to MAO detection.""" + +import asyncio +import logging +import re +import time +from typing import TYPE_CHECKING, Any, Optional + +from pisama_core.detection.orchestrator import DetectionOrchestrator, RealtimeResult +from pisama_core.detection.registry import DetectorRegistry +from pisama_core.detection.registry import registry as global_registry + +from .config import BridgeConfig +from .converter import HookInputConverter +from .session import SessionManager, session_manager +from .types import BridgeResult, HookInput + +if TYPE_CHECKING: + from .chaos.experiments import ChaosResult + +logger = logging.getLogger(__name__) + + +class DetectionBridge: + """Bridge between Agent SDK hooks and MAO detection. + + This is the core integration layer that: + 1. Converts HookInput to Span format + 2. Maintains session context + 3. Runs real-time detection with timeout + 4. Generates appropriate hook outputs + + Example: + bridge = DetectionBridge() + + # PreToolUse hook + result = await bridge.analyze_pre_tool(hook_input, tool_use_id) + if result.should_block: + return {"permissionDecision": "block", ...} + + # PostToolUse hook + result = await bridge.analyze_post_tool(hook_input, tool_use_id) + if result.system_message: + return {"systemMessage": result.system_message} + """ + + def __init__( + self, + config: Optional[BridgeConfig] = None, + detector_registry: Optional[DetectorRegistry] = None, + session_mgr: Optional[SessionManager] = None, + ) -> None: + """Initialize the detection bridge. + + Args: + config: Bridge configuration (defaults to BridgeConfig()) + detector_registry: Detector registry (defaults to global) + session_mgr: Session manager (defaults to global) + """ + self.config = config or BridgeConfig() + self.registry = detector_registry or global_registry + self.sessions = session_mgr or session_manager + + self.converter = HookInputConverter() + self.orchestrator = DetectionOrchestrator( + registry=self.registry, + severity_threshold=self.config.warning_threshold, + block_threshold=self.config.block_threshold, + parallel=True, + ) + + # Optional telemetry (opt-in only) + self._posthog = None + if self.config.enable_telemetry and self.config.telemetry_api_key: + try: + import posthog + + posthog.project_api_key = self.config.telemetry_api_key + posthog.host = self.config.telemetry_host + self._posthog = posthog + except ImportError: + logger.debug("posthog not installed — telemetry disabled") + + # Compile tool patterns + self._include_patterns = [re.compile(p) for p in self.config.tool_patterns] + self._exclude_patterns = [ + re.compile(f"^{re.escape(t)}$") for t in self.config.excluded_tools + ] + + async def analyze_pre_tool( + self, + hook_input: HookInput, + tool_use_id: Optional[str] = None, + ) -> BridgeResult: + """Analyze tool call before execution (PreToolUse). + + This runs detection with strict timeout to decide whether + to block the tool call. + + Args: + hook_input: Input data from Agent SDK + tool_use_id: Unique tool invocation ID + + Returns: + BridgeResult with blocking decision and messages + """ + start_time = time.perf_counter() + tool_name = hook_input.get("tool_name", "") + + if not self._should_analyze(tool_name): + return BridgeResult(execution_time_ms=0) + + session_id = hook_input.get("session_id", "unknown") + if self.sessions.is_blocked(session_id): + return self._blocked_session_result(session_id) + + hook_input, chaos_block = await self._prepare_pre_tool_input( + hook_input, + tool_name=tool_name, + start_time=start_time, + ) + if chaos_block is not None: + return chaos_block + + span = self.converter.to_span(hook_input, tool_use_id, is_post=False) + context = self.sessions.get_context( + session_id, window=self.config.context_window + ) + result = await self._run_pre_detection( + span, + context, + tool_name=tool_name, + start_time=start_time, + ) + if isinstance(result, BridgeResult): + return result + + self.sessions.add_span(session_id, span) + execution_time_ms = (time.perf_counter() - start_time) * 1000 + bridge_result = self._build_pre_tool_result(result, execution_time_ms) + self._record_pre_tool_outcome( + session_id=session_id, + tool_name=tool_name, + result=result, + bridge_result=bridge_result, + ) + return bridge_result + + def _blocked_session_result(self, session_id: str) -> BridgeResult: + return BridgeResult( + should_block=True, + severity=100, + issues=["Session is blocked due to previous violations"], + block_reason=self.sessions.get_block_reason(session_id), + system_message=self._format_blocked_message(session_id), + ) + + async def _prepare_pre_tool_input( + self, + hook_input: HookInput, + *, + tool_name: str, + start_time: float, + ) -> tuple[HookInput, Optional[BridgeResult]]: + chaos = self.config.chaos + if not chaos or not chaos.is_active: + return hook_input, None + + chaos_result = self._apply_pre_chaos(tool_name, hook_input) + if chaos_result and chaos_result.block: + blocked = BridgeResult( + should_block=True, + severity=0, + block_reason=chaos_result.message, + system_message=chaos_result.message, + execution_time_ms=(time.perf_counter() - start_time) * 1000, + ) + return hook_input, blocked + if chaos_result and chaos_result.delay_ms: + await asyncio.sleep(chaos_result.delay_ms / 1000) + if chaos_result and chaos_result.modified_input: + hook_input = {**hook_input, "tool_input": chaos_result.modified_input} + return hook_input, None + + async def _run_pre_detection( + self, + span: Any, + context: Any, + *, + tool_name: str, + start_time: float, + ) -> RealtimeResult | BridgeResult: + try: + return await asyncio.wait_for( + self.orchestrator.analyze_realtime(span, context), + timeout=self.config.detection_timeout_ms / 1000, + ) + except asyncio.TimeoutError: + logger.warning( + f"Detection timeout for tool {tool_name} " + f"(>{self.config.detection_timeout_ms}ms)" + ) + return BridgeResult( + timed_out=True, + execution_time_ms=(time.perf_counter() - start_time) * 1000, + ) + + def _build_pre_tool_result( + self, + result: RealtimeResult, + execution_time_ms: float, + ) -> BridgeResult: + should_block = ( + self.config.enable_blocking + and result.should_block + and result.severity >= self.config.block_threshold + ) + bridge_result = BridgeResult( + should_block=should_block, + severity=result.severity, + issues=result.issues, + recommendations=self._extract_recommendations(result), + block_reason=result.block_reason, + execution_time_ms=execution_time_ms, + ) + if result.severity >= self.config.warning_threshold: + bridge_result.system_message = self._format_pre_tool_message( + result.severity, + result.issues, + should_block, + ) + return bridge_result + + def _record_pre_tool_outcome( + self, + *, + session_id: str, + tool_name: str, + result: RealtimeResult, + bridge_result: BridgeResult, + ) -> None: + if bridge_result.should_block and result.block_reason: + self.sessions.block(session_id, result.block_reason) + + if self.config.log_detections and result.severity > 0: + logger.info( + f"PreToolUse detection: tool={tool_name} " + f"severity={result.severity} block={bridge_result.should_block} " + f"time={bridge_result.execution_time_ms:.1f}ms" + ) + + if self._posthog: + try: + self._posthog.capture( + distinct_id=session_id, + event="sdk_pre_tool_analyzed", + properties={ + "tool_name": tool_name, + "severity": result.severity, + "blocked": bridge_result.should_block, + "execution_time_ms": round(bridge_result.execution_time_ms, 1), + }, + ) + except Exception: + pass + + async def analyze_post_tool( + self, + hook_input: HookInput, + tool_use_id: Optional[str] = None, + ) -> BridgeResult: + """Analyze tool call after execution (PostToolUse). + + This captures the result and may inject recovery messages + if issues are detected. + + Args: + hook_input: Input data from Agent SDK (includes tool_response) + tool_use_id: Unique tool invocation ID + + Returns: + BridgeResult with recovery message if needed + """ + start_time = time.perf_counter() + tool_name = hook_input.get("tool_name", "") + + if not self._should_analyze(tool_name): + return BridgeResult(execution_time_ms=0) + + session_id = hook_input.get("session_id", "unknown") + + # Apply chaos experiments to tool output + if self.config.chaos and self.config.chaos.is_active: + chaos_result = self._apply_post_chaos(tool_name, hook_input) + if chaos_result and chaos_result.modified_output is not None: + hook_input = {**hook_input, "tool_response": chaos_result.modified_output} + + # Convert to span (with response) + span = self.converter.to_span(hook_input, tool_use_id, is_post=True) + + # Always add to session history + self.sessions.add_span(session_id, span) + + if not self.config.enable_recovery: + return BridgeResult( + execution_time_ms=(time.perf_counter() - start_time) * 1000 + ) + + # Get context for analysis + context = self.sessions.get_context( + session_id, window=self.config.context_window + ) + + # Run detection (with timeout) + try: + result = await asyncio.wait_for( + self.orchestrator.analyze_realtime(span, context), + timeout=self.config.detection_timeout_ms / 1000, + ) + except asyncio.TimeoutError: + return BridgeResult( + timed_out=True, + execution_time_ms=(time.perf_counter() - start_time) * 1000, + ) + + execution_time_ms = (time.perf_counter() - start_time) * 1000 + + # Generate recovery message if issues detected + system_message = None + if result.severity >= self.config.warning_threshold: + system_message = self._format_post_tool_message( + result.severity, + result.issues, + result.recommendations, + ) + + if self.config.log_detections and result.severity > 0: + logger.info( + f"PostToolUse detection: tool={tool_name} " + f"severity={result.severity} time={execution_time_ms:.1f}ms" + ) + + return BridgeResult( + should_block=False, # Post-tool never blocks + severity=result.severity, + issues=result.issues, + recommendations=self._extract_recommendations(result), + execution_time_ms=execution_time_ms, + system_message=system_message, + ) + + def _apply_pre_chaos(self, tool_name: str, hook_input: HookInput) -> Optional["ChaosResult"]: + """Apply pre-tool chaos experiments. Returns first matching result.""" + chaos = self.config.chaos + if not chaos: + return None + tool_input = hook_input.get("tool_input", {}) + for exp in chaos.experiments: + if exp.matches(tool_name) and exp.should_trigger(): + result = exp.apply_pre(tool_name, tool_input) + if result.applied: + chaos.record_affected() + logger.info(result.message) + return result + return None + + def _apply_post_chaos(self, tool_name: str, hook_input: HookInput) -> Optional["ChaosResult"]: + """Apply post-tool chaos experiments. Returns first matching result.""" + chaos = self.config.chaos + if not chaos: + return None + tool_output = hook_input.get("tool_response") + for exp in chaos.experiments: + if exp.matches(tool_name) and exp.should_trigger(): + result = exp.apply_post(tool_name, tool_output) + if result.applied: + chaos.record_affected() + logger.info(result.message) + return result + return None + + def _should_analyze(self, tool_name: str) -> bool: + """Check if tool should be analyzed. + + Args: + tool_name: Name of the tool + + Returns: + True if tool should be analyzed + """ + # Check exclusions first + for pattern in self._exclude_patterns: + if pattern.match(tool_name): + return False + + # Check inclusions + for pattern in self._include_patterns: + if pattern.match(tool_name): + return True + + return False + + def _extract_recommendations(self, result: RealtimeResult) -> list[str]: + """Extract recommendation strings from result. + + Args: + result: Detection result + + Returns: + List of recommendation strings + """ + recommendations = [] + for rec in result.recommendations: + if isinstance(rec, dict) and "fix_instruction" in rec: + recommendations.append(rec["fix_instruction"]) + elif isinstance(rec, str): + recommendations.append(rec) + return recommendations + + def _format_pre_tool_message( + self, + severity: int, + issues: list[str], + blocked: bool, + ) -> str: + """Format system message for PreToolUse. + + Args: + severity: Detection severity + issues: List of issues detected + blocked: Whether the tool was blocked + + Returns: + Formatted message string + """ + issue_text = "\n".join(f"- {i}" for i in issues[:3]) + + if blocked: + return f"""[MAO Detection: BLOCKED] +Severity: {severity}/100 + +Issues detected: +{issue_text} + +This tool call has been blocked. Please try a different approach. +Consider: stopping repetitive patterns, changing strategy, or asking the user for guidance.""" + else: + return f"""[MAO Detection: Warning] +Severity: {severity}/100 + +Issues detected: +{issue_text} + +Consider adjusting your approach to avoid potential failure patterns.""" + + def _format_post_tool_message( + self, + severity: int, + issues: list[str], + recommendations: list[Any], + ) -> str: + """Format system message for PostToolUse recovery. + + Args: + severity: Detection severity + issues: List of issues detected + recommendations: List of recommendations + + Returns: + Formatted message string + """ + issue_text = "\n".join(f"- {i}" for i in issues[:3]) + + rec_text = "" + if recommendations: + rec_lines = [] + for r in recommendations[:2]: + if isinstance(r, dict) and "fix_instruction" in r: + rec_lines.append(r["fix_instruction"]) + elif isinstance(r, str): + rec_lines.append(r) + if rec_lines: + rec_text = "\n\nRecommended actions:\n" + "\n".join( + f"- {r}" for r in rec_lines + ) + + return f"""[MAO Detection: Recovery Guidance] +Severity: {severity}/100 + +Pattern detected: +{issue_text} +{rec_text} + +Adjust your approach to prevent this pattern from continuing.""" + + def _format_blocked_message(self, session_id: str) -> str: + """Format message for blocked session. + + Args: + session_id: Session identifier + + Returns: + Formatted message string + """ + reason = self.sessions.get_block_reason(session_id) or "repeated violations" + + return f"""[MAO Detection: Session Blocked] +This session has been blocked due to: {reason} + +To continue, the user must acknowledge and reset the session.""" + + +# Module-level bridge instance +_default_bridge: Optional[DetectionBridge] = None + + +def get_bridge() -> DetectionBridge: + """Get or create the default detection bridge. + + Returns: + DetectionBridge instance + """ + global _default_bridge + if _default_bridge is None: + _default_bridge = DetectionBridge() + return _default_bridge + + +def configure_bridge( + warning_threshold: int | BridgeConfig = 40, + block_threshold: int = 60, + timeout_ms: float = 80, + enable_blocking: bool = True, + enable_recovery: bool = True, +) -> DetectionBridge: + """Configure and return the default detection bridge. + + Call this before using hooks to customize behavior. + + Args: + warning_threshold: Severity to trigger warnings, or a complete + :class:`BridgeConfig` instance. + block_threshold: Severity to trigger blocking + timeout_ms: Detection timeout in milliseconds + enable_blocking: Whether to allow blocking + enable_recovery: Whether to inject recovery messages + + Returns: + Configured DetectionBridge instance + """ + global _default_bridge + if isinstance(warning_threshold, BridgeConfig): + config = warning_threshold + else: + config = BridgeConfig( + warning_threshold=warning_threshold, + block_threshold=block_threshold, + detection_timeout_ms=timeout_ms, + enable_blocking=enable_blocking, + enable_recovery=enable_recovery, + ) + _default_bridge = DetectionBridge(config=config) + return _default_bridge + + +def create_bridge( + warning_threshold: int = 40, + block_threshold: int = 60, + timeout_ms: float = 80, + enable_blocking: bool = True, +) -> DetectionBridge: + """Create a new detection bridge with custom configuration. + + Unlike configure_bridge, this creates a new instance without + affecting the default bridge. + + Args: + warning_threshold: Severity to trigger warnings + block_threshold: Severity to trigger blocking + timeout_ms: Detection timeout in milliseconds + enable_blocking: Whether to allow blocking + + Returns: + New DetectionBridge instance + """ + config = BridgeConfig( + warning_threshold=warning_threshold, + block_threshold=block_threshold, + detection_timeout_ms=timeout_ms, + enable_blocking=enable_blocking, + ) + return DetectionBridge(config=config) diff --git a/src/pisama/agents/chaos/__init__.py b/src/pisama/agents/chaos/__init__.py new file mode 100644 index 0000000..2e98185 --- /dev/null +++ b/src/pisama/agents/chaos/__init__.py @@ -0,0 +1,23 @@ +"""SDK-level chaos engineering — inject failures during agent execution.""" + +from .config import ChaosConfig +from .experiments import ( + ChaosExperiment, + ChaosResult, + ContextTruncation, + ErrorInjection, + LatencyInjection, + OutputCorruption, + ToolFailure, +) + +__all__ = [ + "ChaosConfig", + "ChaosExperiment", + "ChaosResult", + "ToolFailure", + "LatencyInjection", + "ErrorInjection", + "OutputCorruption", + "ContextTruncation", +] diff --git a/src/pisama/agents/chaos/config.py b/src/pisama/agents/chaos/config.py new file mode 100644 index 0000000..6c99924 --- /dev/null +++ b/src/pisama/agents/chaos/config.py @@ -0,0 +1,43 @@ +"""Chaos configuration for the SDK bridge.""" + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from .experiments import ChaosExperiment + + +@dataclass +class ChaosConfig: + """Configure chaos experiments for the SDK bridge. + + Example: + from pisama.agents.chaos import ChaosConfig, ToolFailure, LatencyInjection + + config = ChaosConfig( + experiments=[ + ToolFailure(tools=["database_query"], probability=0.3), + LatencyInjection(min_ms=500, max_ms=3000, probability=0.2), + ], + ) + """ + + experiments: list["ChaosExperiment"] = field(default_factory=list) + safety_max_affected: int = 100 + enabled: bool = True + + # Runtime state (not config) + _affected_count: int = field(default=0, repr=False) + + @property + def is_active(self) -> bool: + """Check if chaos is enabled and safety limit not exceeded.""" + return self.enabled and self._affected_count < self.safety_max_affected + + def record_affected(self) -> None: + """Increment affected count. Disables chaos when safety limit reached.""" + self._affected_count += 1 + + def reset(self) -> None: + """Reset affected count.""" + self._affected_count = 0 diff --git a/src/pisama/agents/chaos/experiments.py b/src/pisama/agents/chaos/experiments.py new file mode 100644 index 0000000..b929ca5 --- /dev/null +++ b/src/pisama/agents/chaos/experiments.py @@ -0,0 +1,166 @@ +"""Chaos experiments for SDK-level failure injection. + +Each experiment can target specific tools/agents and fires with configurable +probability. Applied during pre_tool_use or post_tool_use hooks. +""" + +import random +from dataclasses import dataclass, field +from typing import Any, Optional + + +@dataclass +class ChaosResult: + """Result of applying a chaos experiment.""" + + applied: bool = False + block: bool = False + delay_ms: int = 0 + modified_input: Optional[dict] = None + modified_output: Any = None + message: str = "" + + +@dataclass +class ChaosExperiment: + """Base class for chaos experiments.""" + + probability: float = 1.0 + tools: list[str] = field(default_factory=list) + agents: list[str] = field(default_factory=list) + + def matches(self, tool_name: str, agent_name: str = "") -> bool: + """Check if this experiment targets the given tool/agent.""" + if self.tools and tool_name not in self.tools: + return False + if self.agents and agent_name not in self.agents: + return False + return True + + def should_trigger(self) -> bool: + """Probabilistic trigger check.""" + return random.random() < self.probability + + def apply_pre(self, tool_name: str, tool_input: dict) -> ChaosResult: + """Apply chaos before tool execution. Override in subclasses.""" + return ChaosResult() + + def apply_post(self, tool_name: str, tool_output: Any) -> ChaosResult: + """Apply chaos after tool execution. Override in subclasses.""" + return ChaosResult() + + +@dataclass +class ToolFailure(ChaosExperiment): + """Block tool calls — simulates tool being unavailable. + + Example: + ToolFailure(tools=["database_query"], probability=0.3) + """ + + error_message: str = "Tool unavailable (chaos experiment)" + + def apply_pre(self, tool_name: str, tool_input: dict) -> ChaosResult: + return ChaosResult( + applied=True, + block=True, + message=f"Chaos: {tool_name} blocked — {self.error_message}", + ) + + +@dataclass +class LatencyInjection(ChaosExperiment): + """Add delay before tool execution — simulates slow tools. + + Example: + LatencyInjection(min_ms=500, max_ms=3000, probability=0.2) + """ + + min_ms: int = 100 + max_ms: int = 5000 + + def apply_pre(self, tool_name: str, tool_input: dict) -> ChaosResult: + delay = random.randint(self.min_ms, self.max_ms) + return ChaosResult( + applied=True, + delay_ms=delay, + message=f"Chaos: {tool_name} delayed {delay}ms", + ) + + +@dataclass +class ErrorInjection(ChaosExperiment): + """Return error response — simulates tool errors. + + Example: + ErrorInjection(tools=["search"], error_code=500, probability=0.1) + """ + + error_code: int = 500 + error_message: str = "Internal server error (chaos experiment)" + + def apply_pre(self, tool_name: str, tool_input: dict) -> ChaosResult: + return ChaosResult( + applied=True, + block=True, + message=f"Chaos: {tool_name} error {self.error_code} — {self.error_message}", + ) + + +@dataclass +class OutputCorruption(ChaosExperiment): + """Corrupt tool output — simulates malformed responses. + + Applied in post_tool_use. Truncates, empties, or breaks JSON in the response. + + Example: + OutputCorruption(tools=["search"], corruption="truncate", probability=0.2) + """ + + corruption: str = "truncate" # truncate | empty | json_break + + def apply_post(self, tool_name: str, tool_output: Any) -> ChaosResult: + output_str = str(tool_output) if tool_output else "" + + if self.corruption == "truncate": + corrupted = output_str[: len(output_str) // 2] if output_str else "" + elif self.corruption == "empty": + corrupted = "" + elif self.corruption == "json_break": + corrupted = output_str + '{"incomplete": true' + else: + corrupted = output_str + + return ChaosResult( + applied=True, + modified_output=corrupted, + message=f"Chaos: {tool_name} output corrupted ({self.corruption})", + ) + + +@dataclass +class ContextTruncation(ChaosExperiment): + """Truncate tool input — simulates context window pressure. + + Applied in pre_tool_use. Truncates string values in tool_input. + + Example: + ContextTruncation(truncation_pct=0.5, probability=0.1) + """ + + truncation_pct: float = 0.5 + + def apply_pre(self, tool_name: str, tool_input: dict) -> ChaosResult: + truncated = {} + for key, value in tool_input.items(): + if isinstance(value, str) and len(value) > 100: + keep = int(len(value) * (1 - self.truncation_pct)) + truncated[key] = value[:keep] + else: + truncated[key] = value + + return ChaosResult( + applied=True, + modified_input=truncated, + message=f"Chaos: {tool_name} input truncated {self.truncation_pct:.0%}", + ) diff --git a/src/pisama/agents/check.py b/src/pisama/agents/check.py new file mode 100644 index 0000000..1d2ea13 --- /dev/null +++ b/src/pisama/agents/check.py @@ -0,0 +1,247 @@ +"""Agent-initiated self-check — pisama.check(). + +Allows agents to verify their own output mid-execution by running +Pisama's detection pipeline on arbitrary input. Returns a confidence +score, detected issues, and suggested fixes. + +Usage: + from pisama.agents import check + + result = await check( + output="The server is healthy based on the metrics I found.", + context={"query": "Is auth-service down?", "sources": ["..."]} + ) + if not result["passed"]: + # Agent can retry, adjust, or escalate + print(result["issues"]) +""" + +import asyncio +import json +import logging +import time +from typing import Any, Dict, List, Optional +from urllib.error import URLError +from urllib.request import Request, urlopen + +logger = logging.getLogger(__name__) + +# Default API URL (overridden by bridge config or env) +_api_url: Optional[str] = None + + +def configure_check(api_url: str) -> None: + """Configure the check function to use a specific API URL. + + Args: + api_url: Base URL of the Pisama backend (e.g., "https://api.pisama.ai") + """ + global _api_url + _api_url = api_url.rstrip("/") + + +async def check( + output: str, + context: Optional[Dict[str, Any]] = None, + detectors: Optional[List[str]] = None, + timeout_ms: float = 2000, +) -> Dict[str, Any]: + """Run Pisama detectors on arbitrary output. Agent-initiated self-check. + + This is the primary API for agent self-verification. The agent calls + this when uncertain about its output, and Pisama returns a verdict. + + Args: + output: The agent's output text to verify. + context: Optional context dict with keys: + - query: The original query/task + - sources: List of source documents + - task: Task description + - previous_state: Previous state (for corruption detection) + detectors: Optional list of specific detector names to run. + Defaults to ["hallucination", "derailment", "specification", "completion"]. + timeout_ms: Maximum time to wait for detection (default 2000ms). + + Returns: + Dict with: + passed (bool): True if no issues found + score (float): 0.0-1.0 confidence score (1.0 = clean) + issues (list): Detected problems with details + detectors_run (list[str]): Which detectors were executed + check_time_ms (int): How long the check took + """ + start_ms = time.monotonic_ns() // 1_000_000 + + # Try bridge first (local detection, fastest) + try: + from .bridge import get_bridge + bridge = get_bridge() + if bridge is not None: + result = await _check_via_bridge(bridge, output, context, detectors, timeout_ms) + result["check_time_ms"] = (time.monotonic_ns() // 1_000_000) - start_ms + return result + except Exception: + pass + + # Fallback: call backend API + result = await _check_via_api(output, context, detectors, timeout_ms) + result["check_time_ms"] = (time.monotonic_ns() // 1_000_000) - start_ms + return result + + +async def _check_via_bridge( + bridge: Any, + output: str, + context: Optional[Dict[str, Any]], + detectors: Optional[List[str]], + timeout_ms: float, +) -> Dict[str, Any]: + """Run check using the local detection bridge (no network call).""" + from pisama_core.traces.enums import Platform, SpanKind, SpanStatus + from pisama_core.traces.models import Span + + ctx = context or {} + + # Build a synthetic span from the output + span = Span( + span_id="check-" + str(int(time.time() * 1000)), + name="pisama_check", + kind=SpanKind.AGENT, + status=SpanStatus.OK, + platform=Platform.CLAUDE_CODE, + input_data={ + "content": ctx.get("query", ctx.get("task", "")), + "context": ctx, + }, + output_data={"content": output, "text": output}, + attributes={"check_context": ctx}, + ) + + # Run realtime detection with timeout + try: + result = await asyncio.wait_for( + bridge.orchestrator.analyze_realtime(span, ctx), + timeout=timeout_ms / 1000, + ) + except asyncio.TimeoutError: + return { + "passed": True, + "score": 1.0, + "issues": [], + "detectors_run": [], + "timed_out": True, + } + + # Convert to check result format + issues = [] + for issue_str in (result.issues or []): + if result.severity >= 60: + issue_severity = "high" + elif result.severity >= 30: + issue_severity = "medium" + else: + issue_severity = "low" + issues.append({ + "detector": "realtime", + "description": issue_str, + "severity": issue_severity, + }) + + for rec in (result.recommendations or []): + if isinstance(rec, dict) and "fix_instruction" in rec: + if issues: + issues[-1]["fix"] = rec["fix_instruction"] + + score = max(0.0, 1.0 - (result.severity / 100)) + + return { + "passed": result.severity < 30, + "score": round(score, 3), + "issues": issues, + "detectors_run": ["realtime"], + } + + +async def _check_via_api( + output: str, + context: Optional[Dict[str, Any]], + detectors: Optional[List[str]], + timeout_ms: float, +) -> Dict[str, Any]: + """Run check via the backend /api/v1/evaluate endpoint.""" + ctx = context or {} + + api_base = _api_url + if not api_base: + import os + api_base = os.environ.get("PISAMA_API_URL", "http://localhost:8000") + + url = f"{api_base}/api/v1/evaluate" + + # Build evaluate request + payload = { + "specification": { + "text": ctx.get("task", ctx.get("query", "")), + "user_intent": ctx.get("query", ""), + "sources": ctx.get("sources", []), + "subtasks": ctx.get("subtasks", []), + "success_criteria": ctx.get("success_criteria", []), + }, + "output": { + "text": output, + }, + "agent_role": "generator", + } + if detectors: + payload["detectors"] = detectors + + body = json.dumps(payload).encode() + headers = { + "Content-Type": "application/json", + } + + # Add API key if available + import os + api_key = os.environ.get("PISAMA_API_KEY") + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + try: + req = Request(url, data=body, headers=headers, method="POST") + timeout_sec = timeout_ms / 1000 + + raw = await asyncio.to_thread( + lambda: urlopen(req, timeout=timeout_sec).read() + ) + data = json.loads(raw.decode()) + + # Convert evaluate response to check format + issues = [] + for failure in data.get("failures", []): + issue = { + "detector": failure.get("detector", "unknown"), + "confidence": failure.get("confidence", 0.0), + "severity": failure.get("severity", "medium"), + "description": failure.get("description", ""), + } + if failure.get("suggested_fix"): + issue["fix"] = failure["suggested_fix"] + issues.append(issue) + + return { + "passed": data.get("passed", True), + "score": data.get("score", 1.0), + "issues": issues, + "detectors_run": data.get("detectors_run", []), + } + + except (URLError, TimeoutError, json.JSONDecodeError) as exc: + logger.warning("pisama.check() API call failed: %s", exc) + # Fail open — don't block the agent if check fails + return { + "passed": True, + "score": 1.0, + "issues": [], + "detectors_run": [], + "error": str(exc), + } diff --git a/src/pisama/agents/check_compliance.py b/src/pisama/agents/check_compliance.py new file mode 100644 index 0000000..7e735c8 --- /dev/null +++ b/src/pisama/agents/check_compliance.py @@ -0,0 +1,280 @@ +"""Specification-compliance check (beta, feature-flagged). + +Exposes the backend ``specification_compliance`` detector as a top-level +SDK call. Given an agent's system prompt and a trace of events, the +analyzer extracts behavioural rules from the system prompt and reports +any rule violations detected in the trace. + +The call is gated behind ``PISAMA_ENABLE_CHECK_COMPLIANCE`` so the +detector stays opt-in while the API shape stabilises. When the flag is +not set, ``check_compliance`` raises ``PisamaFeatureNotEnabledError``. + +Usage:: + + from pisama.agents import check_compliance + + result = await check_compliance( + system_prompt="You are a careful assistant. Never call drop_table.", + trace_events=[{"type": "tool_call", "name": "drop_table", "args": {}}], + ) + if result.detected: + for v in result.violations: + print(v.rule_id, v.explanation) + +This is a beta API. The dataclass field set will stay stable, but new +fields may be added before GA. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import time +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional +from urllib.error import URLError +from urllib.request import Request, urlopen + +logger = logging.getLogger(__name__) + + +# --------------------------------------------------------------------------- +# Errors +# --------------------------------------------------------------------------- + + +class PisamaFeatureNotEnabledError(RuntimeError): + """Raised when an opt-in SDK feature is called without its flag set. + + Subclassing ``RuntimeError`` (not ``Exception``) keeps callers that + blanket-catch ``RuntimeError`` from missing the gate, while still + being narrow enough that targeted ``except`` clauses can match it. + """ + + +# --------------------------------------------------------------------------- +# Result dataclasses (public API surface) +# --------------------------------------------------------------------------- + + +@dataclass +class BehavioralRule: + """A single behavioural rule extracted from the system prompt. + + Mirrors the backend's ``BehavioralRule`` shape but is redefined here + so SDK consumers don't take a dependency on the backend package. + """ + + rule_id: str + description: str + trigger: str + required_action: Optional[str] = None + forbidden_action: Optional[str] = None + severity: str = "medium" + + +@dataclass +class Violation: + """A single detected violation of a behavioural rule.""" + + rule_id: str + evidence: str + explanation: str + confidence: float = 0.5 + + +@dataclass +class ComplianceResult: + """Return type for ``check_compliance``. + + ``extracted_rules`` is included for transparency so callers can + inspect what the detector treated as the rule set without a second + round-trip. ``tokens_used`` and ``cost_usd`` aggregate both stages + of the two-stage pipeline (rule extraction + per-rule check). + """ + + detected: bool + confidence: float + violations: List[Violation] = field(default_factory=list) + extracted_rules: List[BehavioralRule] = field(default_factory=list) + tokens_used: int = 0 + cost_usd: float = 0.0 + + +# --------------------------------------------------------------------------- +# Feature flag +# --------------------------------------------------------------------------- + + +_FEATURE_FLAG_ENV = "PISAMA_ENABLE_CHECK_COMPLIANCE" +_TRUTHY = {"1", "true", "yes"} + + +def _is_compliance_enabled() -> bool: + """Return True if ``PISAMA_ENABLE_CHECK_COMPLIANCE`` is set truthy.""" + raw = os.environ.get(_FEATURE_FLAG_ENV, "") + return raw.strip().lower() in _TRUTHY + + +# --------------------------------------------------------------------------- +# HTTP transport +# --------------------------------------------------------------------------- + + +_DEFAULT_TIMEOUT_MS = 60_000 # rule extraction + per-rule LLM calls; minutes-class + + +def _resolve_api_base() -> str: + """Pick the backend base URL the SDK should target. + + Order of precedence: + 1. Whatever ``configure_check`` set on the existing check module. + 2. ``PISAMA_API_URL`` environment variable. + 3. ``http://localhost:8000`` for local dev. + """ + try: + from . import check as _check_module + configured = getattr(_check_module, "_api_url", None) + if configured: + return configured.rstrip("/") + except Exception: + pass + return os.environ.get("PISAMA_API_URL", "http://localhost:8000").rstrip("/") + + +def _build_request(payload: Dict[str, Any]) -> Request: + api_base = _resolve_api_base() + url = f"{api_base}/api/v1/evaluate/detect/specification-compliance" + body = json.dumps(payload).encode() + headers = {"Content-Type": "application/json"} + api_key = os.environ.get("PISAMA_API_KEY") + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + return Request(url, data=body, headers=headers, method="POST") + + +def _parse_response(data: Dict[str, Any]) -> ComplianceResult: + """Convert a backend JSON response into a ComplianceResult. + + Defensive against missing keys: the backend response is the source + of truth for shape, but the SDK should not crash if a future field + is removed or renamed — instead it falls back to safe defaults. + """ + rules = [ + BehavioralRule( + rule_id=str(r.get("rule_id", "")), + description=str(r.get("description", "")), + trigger=str(r.get("trigger", "always")), + required_action=r.get("required_action"), + forbidden_action=r.get("forbidden_action"), + severity=str(r.get("severity", "medium")), + ) + for r in (data.get("extracted_rules") or []) + if isinstance(r, dict) + ] + violations = [ + Violation( + rule_id=str(v.get("rule_id", "")), + evidence=str(v.get("evidence", "")), + explanation=str(v.get("explanation", "")), + confidence=float(v.get("confidence", 0.5)), + ) + for v in (data.get("violations") or []) + if isinstance(v, dict) + ] + return ComplianceResult( + detected=bool(data.get("detected", False)), + confidence=float(data.get("confidence", 0.0)), + violations=violations, + extracted_rules=rules, + tokens_used=int(data.get("tokens_used", 0) or 0), + cost_usd=float(data.get("cost_usd", 0.0) or 0.0), + ) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +async def check_compliance( + system_prompt: str, + trace_events: List[Dict[str, Any]], + timeout_ms: float = _DEFAULT_TIMEOUT_MS, +) -> ComplianceResult: + """Run the specification-compliance detector against a trace. + + Args: + system_prompt: The agent's system prompt. Behavioural rules are + extracted from this text via an LLM call (cached per-prompt + on the backend). + trace_events: Chronological list of trace event dicts. Each event + should have a ``type`` field; supported types are + ``tool_call``, ``agent_message``, ``user_message`` and + ``tool_result``. Other shapes are accepted but only + generically scanned. + timeout_ms: Maximum total wait for the backend call. Defaults to + 60 seconds because the analyzer makes multiple LLM calls. + + Returns: + ComplianceResult with ``detected``, ``confidence``, ``violations`` + (per detected rule violation), ``extracted_rules`` (everything + the rule extractor pulled from the system prompt, included for + transparency), ``tokens_used`` and ``cost_usd``. + + Raises: + PisamaFeatureNotEnabledError: When ``PISAMA_ENABLE_CHECK_COMPLIANCE`` + is not set to a truthy value. + TimeoutError: When the backend call exceeds ``timeout_ms``. + URLError: When the backend is unreachable. + """ + if not _is_compliance_enabled(): + raise PisamaFeatureNotEnabledError( + "check_compliance is gated behind a feature flag. " + f"Set {_FEATURE_FLAG_ENV}=1 to enable. " + "This is a beta API; the response shape may change before GA." + ) + + payload: Dict[str, Any] = { + "system_prompt": system_prompt or "", + "trace_events": list(trace_events or []), + } + req = _build_request(payload) + timeout_sec = max(0.001, timeout_ms / 1000) + + start = time.monotonic() + try: + raw = await asyncio.to_thread( + lambda: urlopen(req, timeout=timeout_sec).read() + ) + except (URLError, TimeoutError) as exc: + elapsed_ms = (time.monotonic() - start) * 1000 + logger.warning( + "check_compliance backend call failed after %.0fms: %s", + elapsed_ms, exc, + ) + raise + + try: + data = json.loads(raw.decode()) + except (json.JSONDecodeError, ValueError) as exc: + logger.warning("check_compliance: malformed JSON response: %s", exc) + # Surface a deterministic empty result rather than crashing the + # caller — the agent has no compliance signal but at least the + # control flow stays intact. This is consistent with the + # fail-open posture the rest of the SDK takes for transport + # errors that aren't the caller's fault. + return ComplianceResult(detected=False, confidence=0.0) + + return _parse_response(data) + + +__all__ = [ + "PisamaFeatureNotEnabledError", + "BehavioralRule", + "Violation", + "ComplianceResult", + "check_compliance", +] diff --git a/src/pisama/agents/clarification.py b/src/pisama/agents/clarification.py new file mode 100644 index 0000000..5f0190c --- /dev/null +++ b/src/pisama/agents/clarification.py @@ -0,0 +1,238 @@ +"""ClarificationPrimitive — pause, ask, resume. + +Most agent frameworks today have no native way for an agent to ask +its human a clarifying question mid-turn. This module provides one: +when `entity_confusion` or another ambiguity detector fires, the +agent emits a `ClarificationRequest`; the host (Claude SDK, LangGraph +node, FastAPI route) blocks until a human answers; the agent then +resumes with the answer in context. + +Why this lives in the SDK rather than the backend: +- The pause needs to happen *inside the agent's loop* — only the SDK + has access to the agent's execution state. +- The resume needs to inject the user's answer into the next prompt, + which is framework-specific (system message vs assistant turn vs + tool result). The primitive defers that to a per-framework adapter. + +This unlocks `entity_confusion` (which was insight-only before) and +provides a foundation for other "ask a human" patterns. + +Contract: +- Caller hands the primitive a `ClarificationRequest` describing what + to ask. The primitive returns a `Resolution` once the human responds + (or a timeout fires). +- `await_response` is async-friendly but works synchronously when + the host has a blocking-IO answer channel (most CLI / Slack / web + hooks). +""" + +from __future__ import annotations + +import logging +import time +import uuid +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Mapping, Optional, Protocol, Union + +logger = logging.getLogger(__name__) + + +@dataclass +class ClarificationRequest: + """A structured 'agent needs human input' payload. + + Fields are deliberately narrow — the primitive forces the agent to + ask a single well-formed question with bounded answer choices, not + a free-form prompt. This keeps the resume-with-answer step + predictable. + """ + + question: str + options: List[str] = field(default_factory=list) + detection_type: str = "" + evidence: Dict[str, Any] = field(default_factory=dict) + request_id: str = field(default_factory=lambda: f"clar_{uuid.uuid4().hex[:12]}") + created_at: float = field(default_factory=time.time) + timeout_seconds: float = 300.0 # 5min default; tune per host + + def to_dict(self) -> Dict[str, Any]: + return { + "request_id": self.request_id, + "question": self.question, + "options": list(self.options), + "detection_type": self.detection_type, + "evidence": self.evidence, + "created_at": self.created_at, + "timeout_seconds": self.timeout_seconds, + } + + +@dataclass +class Resolution: + """Outcome of awaiting a clarification answer.""" + + request_id: str + answered: bool + answer: Optional[str] = None + answer_index: Optional[int] = None + timed_out: bool = False + elapsed_seconds: float = 0.0 + error: Optional[str] = None + + +# Type alias for the answer-channel callable. Hosts wire this to: +# - CLI: prompt() with timeout +# - Slack: post-and-wait on thread reply +# - Web: store request, return URL, poll DB for answer +# Function returns (answer_text, index_into_options) or None on timeout. +class AnswerResult(Protocol): + """Object-form answer accepted from host integrations.""" + + text: str + index: Optional[int] + + +AnswerProviderValue = Union[AnswerResult, Mapping[str, Any]] +AnswerProvider = Callable[[ClarificationRequest], Optional[AnswerProviderValue]] + + +# --------------------------------------------------------------------- +# Detection → ClarificationRequest builders. Each maps a detector's +# evidence to a well-formed question with bounded options. +# --------------------------------------------------------------------- + + +def build_entity_confusion_request(detection: Dict[str, Any]) -> Optional[ClarificationRequest]: + """Build a clarification question for `entity_confusion`. + + Picks the two confused entities and asks the human which one the + agent should have referenced. Returns None if the detection lacks + the entities (shouldn't happen post-Track-D but guard anyway). + """ + details = detection.get("details") or {} + entities = details.get("confused_entities") or details.get("entities") or [] + if len(entities) < 2: + return None + a, b = entities[0], entities[1] + return ClarificationRequest( + question=( + f"Pisama caught a possible entity confusion — the agent " + f"referred to both {a!r} and {b!r} as if they were the same. " + "Which one did you mean?" + ), + options=[ + str(a), + str(b), + "Both (and I want to disambiguate explicitly)", + "Neither — the agent misread me", + ], + detection_type="entity_confusion", + evidence={ + "entities": entities, + "context_snippet": details.get("context_snippet", ""), + }, + ) + + +# --------------------------------------------------------------------- +# The primitive itself. +# --------------------------------------------------------------------- + + +class ClarificationPrimitive: + """Pause the agent, ask the human, resume with the answer.""" + + def __init__(self, *, answer_provider: AnswerProvider) -> None: + self._answer_provider = answer_provider + + @classmethod + def for_detection( + cls, + detection: Dict[str, Any], + *, + answer_provider: AnswerProvider, + ) -> Optional["ClarificationPrimitive"]: + """Build a primitive iff we have a request-builder for this detector.""" + det_type = detection.get("detection_type", "") + if det_type not in _DETECTION_BUILDERS: + return None + return cls(answer_provider=answer_provider) + + def request_from_detection( + self, + detection: Dict[str, Any], + ) -> Optional[ClarificationRequest]: + det_type = detection.get("detection_type", "") + builder = _DETECTION_BUILDERS.get(det_type) + if builder is None: + return None + return builder(detection) + + def await_response(self, request: ClarificationRequest) -> Resolution: + """Block until the human answers or the request times out.""" + started = time.monotonic() + try: + result = self._answer_provider(request) + except Exception as exc: + return Resolution( + request_id=request.request_id, + answered=False, + elapsed_seconds=time.monotonic() - started, + error=f"answer_provider raised: {exc}", + ) + + elapsed = time.monotonic() - started + if result is None: + return Resolution( + request_id=request.request_id, + answered=False, + timed_out=True, + elapsed_seconds=elapsed, + ) + + text = getattr(result, "text", None) or ( + result.get("text") if isinstance(result, Mapping) else None + ) + index = getattr(result, "index", None) + if isinstance(result, Mapping): + index = result.get("index", index) + + return Resolution( + request_id=request.request_id, + answered=True, + answer=text, + answer_index=index, + elapsed_seconds=elapsed, + ) + + def resolve_for_detection( + self, + detection: Dict[str, Any], + ) -> Resolution: + """Build the request and await the response in one call.""" + request = self.request_from_detection(detection) + if request is None: + return Resolution( + request_id="", + answered=False, + error=f"No clarification builder for {detection.get('detection_type','?')!r}", + ) + return self.await_response(request) + + +_DETECTION_BUILDERS: Dict[str, Callable[[Dict[str, Any]], Optional[ClarificationRequest]]] = { + "entity_confusion": build_entity_confusion_request, +} + + +def register_clarification_builder( + detection_type: str, + builder: Callable[[Dict[str, Any]], Optional[ClarificationRequest]], +) -> None: + """Register a builder for a new detection type. + + Adding a new clarifiable detector is a one-line change at the + registration point plus the builder function — no SDK release + required. + """ + _DETECTION_BUILDERS[detection_type] = builder diff --git a/src/pisama/agents/cli/__init__.py b/src/pisama/agents/cli/__init__.py new file mode 100644 index 0000000..50bd87a --- /dev/null +++ b/src/pisama/agents/cli/__init__.py @@ -0,0 +1,8 @@ +"""CLI entry points for pisama.agents. + +Currently exposes: + +- ``pisama-openhands-monitor`` — analyze a completed OpenHands / + Harbor-on-OpenHands session directory through the Pisama backend. + See ``cli/openhands_monitor.py``. +""" diff --git a/src/pisama/agents/cli/openhands_monitor.py b/src/pisama/agents/cli/openhands_monitor.py new file mode 100644 index 0000000..0a98b45 --- /dev/null +++ b/src/pisama/agents/cli/openhands_monitor.py @@ -0,0 +1,129 @@ +"""``pisama-openhands-monitor`` CLI. + +Wraps :class:`pisama.agents.OpenHandsEventStreamAdapter` for batch +analysis of a completed session directory. Looks for the Harbor-emitted +``agent/trajectory.json`` (or ``trajectory.json`` at session root), +POSTs it to the Pisama backend, prints a concise diagnosis. + +Usage:: + + pisama-openhands-monitor + pisama-openhands-monitor --api-url https://api.pisama.ai + pisama-openhands-monitor --json # full diagnosis as JSON + +Exit codes: + 0 — analysis succeeded, no failures detected + 1 — analysis succeeded, at least one failure detected + 2 — usage / runtime error (e.g., session_dir missing trajectory.json) +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Optional, Sequence + +from ..openhands_adapter import OpenHandsEventStreamAdapter + + +def _format_human_summary(result) -> str: + lines: list[str] = [] + lines.append(f"trace_id: {result.trace_id}") + if result.source_path: + lines.append(f"source: {result.source_path}") + lines.append(f"spans: {result.span_count}") + lines.append(f"tokens: {result.total_tokens}") + lines.append( + f"score: {result.trajectory_score:.3f} " + f"(1.0 = clean; lower = worse)" + ) + lines.append("") + if not result.has_failures: + lines.append("✓ No failures detected.") + if result.detectors_run: + lines.append(f" ran {len(result.detectors_run)} detectors: " + f"{', '.join(sorted(result.detectors_run))}") + return "\n".join(lines) + + lines.append(f"✗ {result.failure_count} failure(s) detected:") + for d in result.failures: + lines.append( + f" [{d.severity.upper():<8}] {d.detector:<24} " + f"conf={d.confidence:.2f} {d.title}" + ) + if d.description: + lines.append(f" {d.description}") + if result.detectors_failed: + lines.append("") + lines.append("Detectors that raised during the run:") + for name, err in result.detectors_failed.items(): + lines.append(f" - {name}: {err}") + return "\n".join(lines) + + +def main(argv: Optional[Sequence[str]] = None) -> int: + parser = argparse.ArgumentParser( + prog="pisama-openhands-monitor", + description=( + "Analyze a completed OpenHands session through the Pisama " + "backend. Reads the Harbor-emitted trajectory.json from the " + "session dir; expects the trajectory in ATIF format." + ), + ) + parser.add_argument( + "session_dir", + type=Path, + help="Path to the completed session directory.", + ) + parser.add_argument( + "--api-url", + type=str, + default=None, + help="Override the Pisama backend URL (default: https://api.pisama.ai).", + ) + parser.add_argument( + "--timeout", + type=float, + default=60.0, + help="HTTP timeout in seconds (default: 60).", + ) + parser.add_argument( + "--project-id", + type=str, + default=None, + help="Optional Pisama project id (ps_...) to correlate the run with.", + ) + parser.add_argument( + "--json", + action="store_true", + help="Emit the full diagnosis as JSON instead of the human summary.", + ) + args = parser.parse_args(argv) + + if not args.session_dir.exists(): + print(f"error: session_dir does not exist: {args.session_dir}", file=sys.stderr) + return 2 + + adapter = OpenHandsEventStreamAdapter( + api_url=args.api_url, + mode="batch", + project_id=args.project_id, + timeout_seconds=args.timeout, + ) + try: + result = adapter.on_session_complete(args.session_dir) + except ValueError as e: + print(f"error: {e}", file=sys.stderr) + return 2 + + if args.json: + print(json.dumps(result.raw, indent=2, default=str)) + else: + print(_format_human_summary(result)) + + return 1 if result.has_failures else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/pisama/agents/config.py b/src/pisama/agents/config.py new file mode 100644 index 0000000..f065e76 --- /dev/null +++ b/src/pisama/agents/config.py @@ -0,0 +1,210 @@ +"""Configuration management for Agent SDK integration.""" + +from __future__ import annotations + +import json +import os +from dataclasses import asdict, dataclass, field, fields +from pathlib import Path +from typing import TYPE_CHECKING, Optional + +if TYPE_CHECKING: + from .chaos.config import ChaosConfig + + +@dataclass +class BridgeConfig: + """Configuration for the detection bridge. + + Can be loaded from: + - Environment variables (PISAMA_*) + - JSON config file + - Direct instantiation + + Example: + # From environment + config = BridgeConfig.from_env() + + # From file + config = BridgeConfig.from_file(Path("~/.pisama/config.json")) + + # Direct + config = BridgeConfig(warning_threshold=30, block_threshold=50) + """ + + # Detection thresholds + warning_threshold: int = 40 # Severity to trigger warnings + block_threshold: int = 60 # Severity to trigger blocking + + # Timeout settings (in milliseconds) + detection_timeout_ms: float = 80 # Max time for detection (leave 20ms buffer) + total_timeout_ms: float = 100 # Hard limit for hook execution + + # Feature flags + enable_blocking: bool = True # Allow blocking tool calls + enable_recovery: bool = True # Enable PostToolUse recovery messages + fail_open: bool = True # Allow execution on detection errors + + # Session config + context_window: int = 10 # Recent spans to include in context + max_sessions: int = 100 # Maximum concurrent sessions + session_ttl_seconds: int = 3600 # Session expiry time + + # Detector filtering + enabled_detectors: list[str] = field( + default_factory=lambda: ["loop", "repetition", "coordination", "cost"] + ) + disabled_detectors: list[str] = field(default_factory=list) + + # Tool patterns (regex) + tool_patterns: list[str] = field(default_factory=lambda: [".*"]) + excluded_tools: list[str] = field( + default_factory=lambda: ["AskUserQuestion", "user_input"] + ) + + # Logging + log_level: str = "INFO" + log_detections: bool = True + + # Chaos engineering (optional, off by default) + chaos: Optional["ChaosConfig"] = None + + # Telemetry (opt-in, off by default) + enable_telemetry: bool = False + telemetry_api_key: str = "" + telemetry_host: str = "https://us.i.posthog.com" + + @classmethod + def from_env(cls) -> "BridgeConfig": + """Create config from environment variables. + + Environment variables: + PISAMA_WARNING_THRESHOLD: Warning severity threshold + PISAMA_BLOCK_THRESHOLD: Blocking severity threshold + PISAMA_TIMEOUT_MS: Detection timeout + PISAMA_ENABLE_BLOCKING: Enable/disable blocking + PISAMA_ENABLE_RECOVERY: Enable/disable recovery messages + PISAMA_CONTEXT_WINDOW: Context window size + PISAMA_LOG_LEVEL: Logging level + """ + return cls( + warning_threshold=int(os.getenv("PISAMA_WARNING_THRESHOLD", "40")), + block_threshold=int(os.getenv("PISAMA_BLOCK_THRESHOLD", "60")), + detection_timeout_ms=float(os.getenv("PISAMA_TIMEOUT_MS", "80")), + enable_blocking=os.getenv("PISAMA_ENABLE_BLOCKING", "true").lower() + == "true", + enable_recovery=os.getenv("PISAMA_ENABLE_RECOVERY", "true").lower() + == "true", + fail_open=os.getenv("PISAMA_FAIL_OPEN", "true").lower() == "true", + context_window=int(os.getenv("PISAMA_CONTEXT_WINDOW", "10")), + log_level=os.getenv("PISAMA_LOG_LEVEL", "INFO"), + enable_telemetry=os.getenv("PISAMA_TELEMETRY", "false").lower() == "true", + telemetry_api_key=os.getenv("PISAMA_TELEMETRY_API_KEY", ""), + telemetry_host=os.getenv("PISAMA_TELEMETRY_HOST", "https://us.i.posthog.com"), + ) + + @classmethod + def from_file(cls, path: Path) -> "BridgeConfig": + """Create config from JSON file. + + Expected format: + { + "detection": { + "warning_threshold": 40, + "block_threshold": 60, + "timeout_ms": 80 + }, + "session": { + "context_window": 10, + "max_sessions": 100 + }, + "logging": { + "level": "INFO" + } + } + """ + with open(path) as f: + data = json.load(f) + + # ``save()`` writes the dataclass's flat public shape. Older config + # files used nested detection/session/logging sections, so accept both + # formats and keep save/from_file a true round trip. + if not any(section in data for section in ("detection", "session", "logging")): + allowed = {item.name for item in fields(cls)} + return cls(**{key: value for key, value in data.items() if key in allowed}) + + # Handle nested structure + detection = data.get("detection", {}) + session = data.get("session", {}) + logging_cfg = data.get("logging", {}) + + return cls( + warning_threshold=detection.get("warning_threshold", 40), + block_threshold=detection.get("block_threshold", 60), + detection_timeout_ms=detection.get("timeout_ms", 80), + enable_blocking=detection.get("enable_blocking", True), + enable_recovery=detection.get("enable_recovery", True), + fail_open=detection.get("fail_open", True), + enabled_detectors=detection.get("enabled_detectors", []), + disabled_detectors=detection.get("disabled_detectors", []), + context_window=session.get("context_window", 10), + max_sessions=session.get("max_sessions", 100), + session_ttl_seconds=session.get("ttl_seconds", 3600), + tool_patterns=data.get("tool_patterns", [".*"]), + excluded_tools=data.get("excluded_tools", []), + log_level=logging_cfg.get("level", "INFO"), + log_detections=logging_cfg.get("log_detections", True), + ) + + def to_dict(self) -> dict: + """Convert to dictionary.""" + return asdict(self) + + def save(self, path: Path) -> None: + """Save config to JSON file.""" + with open(path, "w") as f: + json.dump(self.to_dict(), f, indent=2) + + +def load_config( + config_path: Optional[Path] = None, + use_env: bool = True, +) -> BridgeConfig: + """Load configuration with fallback chain. + + Priority: + 1. Explicit config_path + 2. Environment variable PISAMA_CONFIG_PATH + 3. Default ~/.pisama/agent_sdk_config.json + 4. Environment variables (PISAMA_*) + 5. Default values + + Args: + config_path: Explicit path to config file + use_env: Whether to fall back to environment variables + + Returns: + BridgeConfig instance + """ + # Check explicit path + if config_path and config_path.exists(): + return BridgeConfig.from_file(config_path) + + # Check env var for path + env_path = os.getenv("PISAMA_CONFIG_PATH") + if env_path: + path = Path(env_path) + if path.exists(): + return BridgeConfig.from_file(path) + + # Check default location + default_path = Path.home() / ".pisama" / "agent_sdk_config.json" + if default_path.exists(): + return BridgeConfig.from_file(default_path) + + # Fall back to env vars + if use_env: + return BridgeConfig.from_env() + + # Default config + return BridgeConfig() diff --git a/src/pisama/agents/converter.py b/src/pisama/agents/converter.py new file mode 100644 index 0000000..2c3d494 --- /dev/null +++ b/src/pisama/agents/converter.py @@ -0,0 +1,196 @@ +"""Converts Claude Agent SDK HookInput to pisama-core Span format.""" + +import uuid +from datetime import datetime, timezone +from typing import Any, Optional + +from pisama_core.traces.enums import Platform, SpanKind, SpanStatus +from pisama_core.traces.models import Span + +from .types import HookInput + + +class HookInputConverter: + """Converts Agent SDK HookInput to universal Span format. + + Similar to pisama-claude-code's TraceConverter but optimized for + the Agent SDK's hook data structure. + + Example: + converter = HookInputConverter() + + # Convert PreToolUse input + span = converter.to_span(hook_input, tool_use_id, is_post=False) + + # Convert PostToolUse input (includes response) + span = converter.to_span(hook_input, tool_use_id, is_post=True) + """ + + # Tool name to SpanKind mapping + TOOL_KIND_MAP: dict[str, SpanKind] = { + # Computer use tools + "computer": SpanKind.TOOL, + "bash": SpanKind.TOOL, + "text_editor": SpanKind.TOOL, + "file_read": SpanKind.TOOL, + "file_write": SpanKind.TOOL, + "web_search": SpanKind.TOOL, + # Agent tools + "agent": SpanKind.AGENT, + "spawn": SpanKind.AGENT, + # User interaction + "user_input": SpanKind.USER_INPUT, + # Claude Code tools (capitalized) + "Bash": SpanKind.TOOL, + "Read": SpanKind.TOOL, + "Write": SpanKind.TOOL, + "Edit": SpanKind.TOOL, + "Glob": SpanKind.TOOL, + "Grep": SpanKind.TOOL, + "Task": SpanKind.AGENT, + "WebFetch": SpanKind.TOOL, + "WebSearch": SpanKind.TOOL, + "AskUserQuestion": SpanKind.USER_INPUT, + "TodoWrite": SpanKind.TOOL, + "NotebookEdit": SpanKind.TOOL, + } + + def __init__(self) -> None: + """Initialize the converter.""" + self._session_traces: dict[str, str] = {} + + def to_span( + self, + hook_input: HookInput, + tool_use_id: Optional[str] = None, + is_post: bool = False, + ) -> Span: + """Convert HookInput to Span. + + Args: + hook_input: Data from Agent SDK hook + tool_use_id: Unique tool invocation ID + is_post: Whether this is PostToolUse (has response) + + Returns: + Universal Span object + """ + tool_name = hook_input.get("tool_name", "unknown") + tool_input = hook_input.get("tool_input", {}) + session_id = hook_input.get("session_id", "unknown") + + # Get or create trace ID for session + if session_id not in self._session_traces: + self._session_traces[session_id] = str(uuid.uuid4()) + trace_id = self._session_traces[session_id] + + # Use tool_use_id as span_id if available + span_id = tool_use_id or str(uuid.uuid4()) + + # Determine kind and status + kind = self._get_span_kind(tool_name) + status = SpanStatus.OK if is_post else SpanStatus.IN_PROGRESS + + # Check for errors in response + tool_response = hook_input.get("tool_response") + error_message = None + if is_post and isinstance(tool_response, dict): + if tool_response.get("is_error"): + status = SpanStatus.ERROR + error_message = str(tool_response.get("output", "Unknown error")) + + # Build attributes + attributes: dict[str, Any] = { + "session_id": session_id, + "hook_type": "post" if is_post else "pre", + } + if "conversation_id" in hook_input: + attributes["conversation_id"] = hook_input["conversation_id"] + if "model" in hook_input: + attributes["model"] = hook_input["model"] + if "usage" in hook_input: + attributes["usage"] = hook_input["usage"] + + return Span( + span_id=span_id, + parent_id=None, + trace_id=trace_id, + name=tool_name, + kind=kind, + platform=Platform.CLAUDE_CODE, + start_time=datetime.now(timezone.utc), + end_time=datetime.now(timezone.utc) if is_post else None, + status=status, + attributes=attributes, + input_data=self._normalize_input(tool_input), + output_data=self._normalize_output(tool_response) if is_post else None, + events=[], + error_message=error_message, + ) + + def _get_span_kind(self, tool_name: str) -> SpanKind: + """Determine SpanKind from tool name. + + Args: + tool_name: Name of the tool + + Returns: + Appropriate SpanKind + """ + # Direct match + if tool_name in self.TOOL_KIND_MAP: + return self.TOOL_KIND_MAP[tool_name] + + # MCP tools + if tool_name.startswith("mcp__") or tool_name.startswith("mcp:"): + return SpanKind.TOOL + + return SpanKind.TOOL + + def _normalize_input(self, tool_input: Any) -> dict[str, Any]: + """Normalize tool input to dict. + + Args: + tool_input: Raw tool input + + Returns: + Normalized dict + """ + if isinstance(tool_input, dict): + return tool_input + elif isinstance(tool_input, str): + return {"value": tool_input} + elif tool_input is None: + return {} + return {"value": str(tool_input)} + + def _normalize_output(self, tool_output: Any) -> dict[str, Any]: + """Normalize tool output to dict. + + Args: + tool_output: Raw tool output + + Returns: + Normalized dict + """ + if isinstance(tool_output, dict): + return tool_output + elif isinstance(tool_output, str): + return {"value": tool_output} + elif isinstance(tool_output, list): + return {"content": tool_output} + elif tool_output is None: + return {} + return {"value": str(tool_output)} + + def reset_session(self, session_id: str) -> None: + """Reset trace tracking for a session. + + Args: + session_id: Session to reset + """ + self._session_traces.pop(session_id, None) + + def reset_all(self) -> None: + """Reset all session trace tracking.""" + self._session_traces.clear() diff --git a/src/pisama/agents/evaluator.py b/src/pisama/agents/evaluator.py new file mode 100644 index 0000000..938a071 --- /dev/null +++ b/src/pisama/agents/evaluator.py @@ -0,0 +1,197 @@ +"""Pisama Evaluator — drop-in evaluator for multi-agent harnesses. + +Usage: + from pisama.agents import PisamaEvaluator + + evaluator = PisamaEvaluator(api_key="psk_...", base_url="https://api.pisama.ai") + + result = evaluator.evaluate( + specification={"text": "Build a login page with OAuth"}, + output={"text": generator_output}, + ) + if not result.passed: + for failure in result.failures: + print(f"{failure.detector}: {failure.description}") +""" + +import logging +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +logger = logging.getLogger(__name__) + +try: + import httpx + _HTTPX_AVAILABLE = True +except ImportError: + _HTTPX_AVAILABLE = False + + +@dataclass +class EvalFailure: + detector: str + confidence: float + severity: str + title: str + description: str + suggested_fix: Optional[str] = None + + +@dataclass +class EvalResult: + passed: bool + score: float + failures: List[EvalFailure] + suggestions: List[str] + detectors_run: List[str] + evaluation_time_ms: int + + +class PisamaEvaluator: + """Client for the Pisama evaluation API. + + Args: + api_key: Pisama API key (psk_...) + base_url: Backend URL (default: https://api.pisama.ai) + timeout: Request timeout in seconds + """ + + def __init__( + self, + api_key: str, + base_url: str = "https://api.pisama.ai", + timeout: float = 30.0, + ): + if not _HTTPX_AVAILABLE: + raise ImportError( + "httpx is required for PisamaEvaluator: `pip install pisama[agents]`" + ) + + self.api_key = api_key + self.base_url = base_url.rstrip("/") + self.timeout = timeout + self._client = httpx.Client( + base_url=self.base_url, + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + timeout=timeout, + ) + + def evaluate( + self, + specification: Dict[str, Any], + output: Dict[str, Any], + agent_role: str = "generator", + detectors: Optional[List[str]] = None, + context_limit: Optional[int] = None, + ) -> EvalResult: + """Evaluate generator output against a specification. + + Args: + specification: Sprint contract or task spec. + output: Generator output to evaluate. + agent_role: Role of the producing agent (generator/evaluator/planner). + detectors: Specific detectors to run (default: auto-select). + context_limit: Model context window for pressure detection. + + Returns: + EvalResult with pass/fail verdict and failure details. + """ + payload: Dict[str, Any] = { + "specification": specification, + "output": output, + "agent_role": agent_role, + } + if detectors: + payload["detectors"] = detectors + if context_limit: + payload["context_limit"] = context_limit + + response = self._client.post("/api/v1/evaluate", json=payload) + response.raise_for_status() + data = response.json() + + failures = [ + EvalFailure( + detector=f["detector"], + confidence=f["confidence"], + severity=f["severity"], + title=f["title"], + description=f["description"], + suggested_fix=f.get("suggested_fix"), + ) + for f in data.get("failures", []) + ] + + return EvalResult( + passed=data["passed"], + score=data["score"], + failures=failures, + suggestions=data.get("suggestions", []), + detectors_run=data.get("detectors_run", []), + evaluation_time_ms=data.get("evaluation_time_ms", 0), + ) + + async def evaluate_async( + self, + specification: Dict[str, Any], + output: Dict[str, Any], + agent_role: str = "generator", + detectors: Optional[List[str]] = None, + context_limit: Optional[int] = None, + ) -> EvalResult: + """Async version of evaluate().""" + payload: Dict[str, Any] = { + "specification": specification, + "output": output, + "agent_role": agent_role, + } + if detectors: + payload["detectors"] = detectors + if context_limit: + payload["context_limit"] = context_limit + + async with httpx.AsyncClient( + base_url=self.base_url, + headers={ + "Authorization": f"Bearer {self.api_key}", + "Content-Type": "application/json", + }, + timeout=self.timeout, + ) as client: + response = await client.post("/api/v1/evaluate", json=payload) + response.raise_for_status() + data = response.json() + + failures = [ + EvalFailure( + detector=f["detector"], + confidence=f["confidence"], + severity=f["severity"], + title=f["title"], + description=f["description"], + suggested_fix=f.get("suggested_fix"), + ) + for f in data.get("failures", []) + ] + + return EvalResult( + passed=data["passed"], + score=data["score"], + failures=failures, + suggestions=data.get("suggestions", []), + detectors_run=data.get("detectors_run", []), + evaluation_time_ms=data.get("evaluation_time_ms", 0), + ) + + def close(self): + """Close the underlying HTTP client.""" + self._client.close() + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() diff --git a/src/pisama/agents/heal.py b/src/pisama/agents/heal.py new file mode 100644 index 0000000..f35a5df --- /dev/null +++ b/src/pisama/agents/heal.py @@ -0,0 +1,245 @@ +"""Sync in-loop healing for the Pisama Agent SDK. + +Most of the SDK's value is observe-only: detectors fire, the bridge +emits warnings or blocks. `heal_now()` is the bridge into Pisama's +healing pipeline for the SAFE-fix slice — the in-loop primitive behind +"self-heals where safe, escalates the rest." + +The call pattern is intentionally synchronous and short. It is invoked +from inside `pre_tool_use_hook` when `auto_heal=True`, so it must come +back fast enough not to disrupt the agent's tool-decision latency. + +Usage: + from pisama.agents.heal import heal_now + + result = heal_now( + detection_type="loop", + details={"states": [...]}, + framework="claude_sdk", + ) + if result.applied and result.prompt_patch: + # Re-issue the agent's next step with the patch. + ... + elif result.escalated: + # Block and route to human. + ... +""" + +from __future__ import annotations + +import json +import logging +import os +from dataclasses import dataclass +from typing import Any, Dict, Optional +from urllib import error as urllib_error +from urllib import request as urllib_request + +logger = logging.getLogger(__name__) + +_DEFAULT_API_URL = "https://api.pisama.ai" +_DEFAULT_TIMEOUT_S = 5.0 + + +@dataclass +class HealingResult: + """Outcome of a sync healing call. + + `applied`: a SAFE-risk fix (or a MEDIUM fix that passed inline + verification, see D7) was returned and is ready to apply inline. + `escalated`: no fix passed the inline gates; the top fix requires + human approval. + Both false means no fix was available at all (e.g. unrecognised + detection type) — caller should fall back to its observe-only path. + + `verification_passed` is populated when the server ran inline + verification (MEDIUM-tier fixes only): True iff the simulated apply + predicted improvement, False iff verification ran and failed, None + iff no verification was attempted (SAFE fixes skip it; DANGEROUS + fixes escalate without it). + """ + + applied: bool + escalated: bool + risk_level: str + fix: Optional[Dict[str, Any]] = None + approval_url: Optional[str] = None + message: str = "" + verification_passed: Optional[bool] = None + verification_reason: Optional[str] = None + # Track F: when the backend's verification gate couldn't run (needs + # an agent_callable only the SDK has), this hint names the primitive + # the SDK should instantiate locally and which inputs it needs. + recommended_verification: Optional[Dict[str, Any]] = None + + @property + def prompt_patch(self) -> Optional[str]: + """Convenience: a system-message patch derived from the fix. + + Resolution order: + 1. `metadata.framework_specific_code` (the LLM enricher's stack- + aware output, when enabled). + 2. The first `code_changes[].suggested_code` from the template + generator — production agents want the snippet, not the prose. + 3. `rationale` / `description` / `title` as last-resort prose. + + The SDK's PreToolUse hook surfaces this directly as the + `systemMessage` on a `permissionDecision: deny` payload. + """ + if not self.fix: + return None + meta = self.fix.get("metadata") or {} + snippet = meta.get("framework_specific_code") + if isinstance(snippet, str) and snippet.strip(): + return snippet + code_changes = self.fix.get("code_changes") or [] + if code_changes and isinstance(code_changes, list): + first = code_changes[0] + if isinstance(first, dict): + code = first.get("suggested_code") + if isinstance(code, str) and code.strip(): + return code + for key in ("rationale", "description", "title"): + value = self.fix.get(key) + if isinstance(value, str) and value.strip(): + return value + return None + + @property + def ide_patch(self) -> Optional[Dict[str, Any]]: + """IDE-native output for this fix, when the server rendered one. + + A dict with a ready-to-paste CLAUDE.md/.cursorrules instruction + block (`instructions`) and a verification command block + (`verification`), plus `target_files`, `apply_mode`, and + `framework`. This is what the /pisama-diagnose skill pastes into a + user's repo. None when no fix was returned or the server predates + IDE-patch rendering. + """ + if not self.fix: + return None + patch = self.fix.get("ide_patch") + return patch if isinstance(patch, dict) else None + + +def heal_now( + *, + detection_type: str, + details: Optional[Dict[str, Any]] = None, + framework: str = "", + method: Optional[str] = None, + api_url: Optional[str] = None, + api_key: Optional[str] = None, + timeout_s: float = _DEFAULT_TIMEOUT_S, +) -> HealingResult: + """Synchronously request a SAFE remediation for a fresh detection. + + Posts to `POST {api_url}/api/v1/healing/trigger/sync`. The endpoint + is side-effect-free — no HealingRecord is persisted — so this can be + invoked in-loop without growing audit-trail noise. + + Failure modes (network error, timeout, server 5xx, missing API key) + all return an empty `HealingResult` rather than raising; the caller + falls back to its observe-only path. + """ + base_url = ( + api_url or os.getenv("PISAMA_API_URL") or _DEFAULT_API_URL + ).rstrip("/") + key = api_key or os.getenv("PISAMA_API_KEY", "") + if not key: + logger.warning("heal_now: PISAMA_API_KEY not set; skipping healing") + return HealingResult( + applied=False, + escalated=False, + risk_level="unknown", + message="PISAMA_API_KEY not configured; healing skipped.", + ) + + payload = { + "detection_type": detection_type, + "details": details or {}, + "framework": framework or "", + "method": method or "", + } + body = json.dumps(payload).encode("utf-8") + req = urllib_request.Request( + url=f"{base_url}/api/v1/healing/trigger/sync", + data=body, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {key}", + "User-Agent": "pisama-python/agents.heal_now", + }, + method="POST", + ) + + try: + with urllib_request.urlopen(req, timeout=timeout_s) as resp: + raw = resp.read() + except urllib_error.HTTPError as exc: + # WARNING (not INFO): a non-2xx here means the call silently degraded + # to observe-only. Keeping it at INFO previously masked a route- + # shadowing 422 + a 404 from a mis-built URL. Surface it. + logger.warning("heal_now: HTTP %s from healing endpoint", exc.code) + return HealingResult( + applied=False, + escalated=False, + risk_level="unknown", + message=f"healing endpoint returned HTTP {exc.code}", + ) + except (urllib_error.URLError, TimeoutError) as exc: + logger.info("heal_now: network failure (%s)", exc) + return HealingResult( + applied=False, + escalated=False, + risk_level="unknown", + message=f"healing endpoint unreachable: {exc}", + ) + except Exception as exc: # pragma: no cover - belt and suspenders + logger.warning("heal_now: unexpected failure (%s)", exc) + return HealingResult( + applied=False, + escalated=False, + risk_level="unknown", + message=f"healing call failed: {exc}", + ) + + try: + data = json.loads(raw.decode("utf-8")) + except (json.JSONDecodeError, UnicodeDecodeError) as exc: + logger.warning("heal_now: invalid JSON response (%s)", exc) + return HealingResult( + applied=False, + escalated=False, + risk_level="unknown", + message="invalid JSON from healing endpoint", + ) + + result = HealingResult( + applied=bool(data.get("applied")), + escalated=bool(data.get("escalated")), + risk_level=str(data.get("risk_level") or "unknown"), + fix=data.get("fix"), + approval_url=data.get("approval_url"), + message=str(data.get("message") or ""), + verification_passed=data.get("verification_passed"), + verification_reason=data.get("verification_reason"), + recommended_verification=data.get("recommended_verification"), + ) + + # Track G: fire any registered SDK-side on_indication callbacks so + # the developer running the agent gets an out-of-band signal even + # when nothing surfaces via systemMessage. Best-effort. + try: + from .indication import SDKIndication + from .indication import _fire as _fire_indication + indication = SDKIndication.from_healing_result( + result, + detection_type=detection_type, + confidence=float((details or {}).get("confidence") or 0.0), + ) + _fire_indication(indication) + except Exception as exc: + logger.warning("heal_now: indication dispatch failed (%s)", exc) + + return result diff --git a/src/pisama/agents/hooks/__init__.py b/src/pisama/agents/hooks/__init__.py new file mode 100644 index 0000000..628d0d0 --- /dev/null +++ b/src/pisama/agents/hooks/__init__.py @@ -0,0 +1,28 @@ +"""Hook implementations for Claude Agent SDK integration.""" + +from .matchers import ( + AGENT_TOOLS, + ALL_TOOLS, + DANGEROUS_COMMANDS, + FILE_TOOLS, + SHELL_TOOLS, + HookMatcher, +) +from .post_tool_use import PostToolUseHook, post_tool_use_hook +from .pre_tool_use import PreToolUseHook, pre_tool_use_hook + +__all__ = [ + # Hook functions + "pre_tool_use_hook", + "post_tool_use_hook", + # Hook classes + "PreToolUseHook", + "PostToolUseHook", + # Matchers + "HookMatcher", + "ALL_TOOLS", + "FILE_TOOLS", + "SHELL_TOOLS", + "DANGEROUS_COMMANDS", + "AGENT_TOOLS", +] diff --git a/src/pisama/agents/hooks/matchers.py b/src/pisama/agents/hooks/matchers.py new file mode 100644 index 0000000..7aefaa4 --- /dev/null +++ b/src/pisama/agents/hooks/matchers.py @@ -0,0 +1,122 @@ +"""Tool matching patterns for hook filtering.""" + +import re +from dataclasses import dataclass, field +from typing import Any, Optional + + +@dataclass +class HookMatcher: + """Pattern matcher for filtering which tools to analyze. + + Use matchers to selectively apply detection to specific tools + or input patterns. + + Example: + # Match all file operations + matcher = HookMatcher(tool_name_pattern="^(Read|Write|Edit)$") + + # Match with input conditions + matcher = HookMatcher( + tool_name_pattern="Bash", + input_pattern=r"rm\\s+-rf", # Dangerous commands + ) + + # Check if tool matches + if matcher.matches("Bash", {"command": "rm -rf /tmp"}): + # Analyze this tool + pass + """ + + tool_name_pattern: Optional[str] = None + input_pattern: Optional[str] = None + exclude_tools: list[str] = field(default_factory=list) + + def __post_init__(self) -> None: + """Compile regex patterns.""" + self._name_re = ( + re.compile(self.tool_name_pattern) if self.tool_name_pattern else None + ) + self._input_re = ( + re.compile(self.input_pattern) if self.input_pattern else None + ) + + def matches(self, tool_name: str, tool_input: Optional[dict[str, Any]] = None) -> bool: + """Check if tool matches this pattern. + + Args: + tool_name: Name of the tool + tool_input: Tool input dictionary + + Returns: + True if tool matches all criteria + """ + # Check exclusions + if tool_name in self.exclude_tools: + return False + + # Check name pattern + if self._name_re and not self._name_re.match(tool_name): + return False + + # Check input pattern + if self._input_re and tool_input: + input_str = str(tool_input) + if not self._input_re.search(input_str): + return False + + return True + + +# Pre-built matchers for common use cases + +ALL_TOOLS = HookMatcher(tool_name_pattern=".*") +"""Match all tools.""" + +FILE_TOOLS = HookMatcher( + tool_name_pattern="^(Read|Write|Edit|Glob|Grep|file_read|file_write)$" +) +"""Match file operation tools.""" + +SHELL_TOOLS = HookMatcher( + tool_name_pattern="^(Bash|bash|computer|shell)$" +) +"""Match shell/command execution tools.""" + +DANGEROUS_COMMANDS = HookMatcher( + tool_name_pattern="^(Bash|bash)$", + input_pattern=r"(rm\s+-rf|sudo|chmod\s+777|curl.*\|\s*sh)", +) +"""Match shell tools with dangerous command patterns.""" + +AGENT_TOOLS = HookMatcher( + tool_name_pattern="^(Task|agent|spawn)$" +) +"""Match agent/subagent invocation tools.""" + + +def create_matcher( + tools: Optional[list[str]] = None, + exclude: Optional[list[str]] = None, + input_pattern: Optional[str] = None, +) -> HookMatcher: + """Create a custom matcher. + + Args: + tools: List of tool names to match (OR logic) + exclude: List of tool names to exclude + input_pattern: Regex pattern to match in tool input + + Returns: + Configured HookMatcher + """ + tool_pattern = None + if tools: + escaped = [re.escape(t) for t in tools] + tool_pattern = f"^({'|'.join(escaped)})$" + + return HookMatcher( + tool_name_pattern=tool_pattern, + input_pattern=input_pattern, + exclude_tools=exclude or [], + ) diff --git a/src/pisama/agents/hooks/post_tool_use.py b/src/pisama/agents/hooks/post_tool_use.py new file mode 100644 index 0000000..025efc7 --- /dev/null +++ b/src/pisama/agents/hooks/post_tool_use.py @@ -0,0 +1,117 @@ +"""PostToolUse hook implementation.""" + +import logging +from typing import Any, Optional + +from ..bridge import DetectionBridge, get_bridge +from ..types import HookContext, HookInput + +logger = logging.getLogger(__name__) + + +async def post_tool_use_hook( + input_data: HookInput, + tool_use_id: Optional[str], + context: HookContext, +) -> dict[str, Any]: + """PostToolUse hook for failure capture and recovery. + + This hook is called after each tool execution and can: + - Capture trace data for analysis + - Inject recovery guidance via systemMessage + + Args: + input_data: Contains tool_name, tool_input, tool_response, session_id + tool_use_id: Unique identifier for this tool invocation + context: Hook context with signal + + Returns: + Hook output dict with recovery message if needed + + Example: + from pisama.agents.hooks import post_tool_use_hook + + # Register with Claude Agent SDK + agent.hooks.post_tool_use = post_tool_use_hook + """ + if not tool_use_id: + logger.debug("PostToolUse called without tool_use_id, skipping") + return {} + + bridge = get_bridge() + + try: + result = await bridge.analyze_post_tool(input_data, tool_use_id) + + output: dict[str, Any] = {} + if result.system_message: + output["systemMessage"] = result.system_message + + return output + + except Exception as e: + logger.error(f"PostToolUse hook error: {e}", exc_info=True) + return {} + + +class PostToolUseHook: + """Class-based PostToolUse hook with configuration. + + Use this when you need more control over the hook behavior, + such as custom bridge configuration or disabling recovery messages. + + Example: + from pisama.agents.hooks import PostToolUseHook + from pisama.agents import DetectionBridge + + # Create hook with custom bridge + hook = PostToolUseHook(bridge=my_bridge, inject_recovery=True) + + # Register with agent + agent.hooks.post_tool_use = hook + """ + + def __init__( + self, + bridge: Optional[DetectionBridge] = None, + inject_recovery: bool = True, + ) -> None: + """Initialize the hook. + + Args: + bridge: Custom detection bridge (defaults to global) + inject_recovery: If True, inject recovery messages on issues + """ + self.bridge = bridge or get_bridge() + self.inject_recovery = inject_recovery + + async def __call__( + self, + input_data: HookInput, + tool_use_id: Optional[str], + context: HookContext, + ) -> dict[str, Any]: + """Handle PostToolUse event. + + Args: + input_data: Hook input data + tool_use_id: Tool use identifier + context: Hook context + + Returns: + Hook output dict + """ + if not tool_use_id: + return {} + + try: + result = await self.bridge.analyze_post_tool(input_data, tool_use_id) + + output: dict[str, Any] = {} + if self.inject_recovery and result.system_message: + output["systemMessage"] = result.system_message + + return output + except Exception as e: + logger.error(f"PostToolUse error: {e}") + return {} diff --git a/src/pisama/agents/hooks/pre_tool_use.py b/src/pisama/agents/hooks/pre_tool_use.py new file mode 100644 index 0000000..bd0fced --- /dev/null +++ b/src/pisama/agents/hooks/pre_tool_use.py @@ -0,0 +1,244 @@ +"""PreToolUse hook implementation.""" + +import asyncio +import logging +import os +from typing import Any, Optional + +from ..bridge import DetectionBridge, get_bridge +from ..heal import HealingResult, heal_now +from ..types import BridgeResult, HookContext, HookInput +from .matchers import HookMatcher + +logger = logging.getLogger(__name__) + + +def _auto_heal_enabled_from_env() -> bool: + return os.getenv("PISAMA_AUTO_HEAL", "false").lower() in ("1", "true", "yes", "on") + + +def _attempt_inline_heal( + result: BridgeResult, *, timeout_s: float = 2.0 +) -> Optional[HealingResult]: + """Synchronous call to `heal_now()` with the bridge's primary detection. + + Returns None when there is nothing to heal (no detection metadata, + no network connectivity, no fix available). The caller falls back + to the standard block path in that case. + """ + if not result.primary_detection_type: + return None + try: + healing = heal_now( + detection_type=result.primary_detection_type, + details=result.primary_details or {}, + framework=result.framework or "", + timeout_s=timeout_s, + ) + except Exception as exc: # pragma: no cover - heal_now is already defensive + logger.warning("auto_heal: heal_now raised (%s); falling back", exc) + return None + return healing + + +def _heal_output(healing: HealingResult) -> dict[str, Any]: + """Build a `permissionDecision: deny` hook payload that carries the patch. + + `deny` (vs `block`) signals to Claude SDK and LangGraph that the + current step is rejected but the agent may proceed after taking the + `systemMessage` into account — which is exactly the in-loop self-heal + semantics. If we returned `block`, the agent would stop hard. + """ + patch = ( + healing.prompt_patch + or "Pisama applied a SAFE in-loop fix; retry with the adjusted context." + ) + return { + "hookSpecificOutput": { + "hookEventName": "PreToolUse", + "permissionDecision": "deny", + "permissionDecisionReason": ( + "Pisama auto-healed in-loop (SAFE fix). Retry with the patched system message." + ), + }, + "systemMessage": patch, + } + + +def _escalation_output(healing: HealingResult, fallback: dict[str, Any]) -> dict[str, Any]: + """Surface escalation via the block path so a human can approve.""" + output = dict(fallback) + extra = ( + f"\n\nPisama escalated this failure (risk: {healing.risk_level}). " + f"Review at {healing.approval_url}" if healing.approval_url else "" + ) + existing_msg = output.get("systemMessage") or "" + output["systemMessage"] = (existing_msg + extra).strip() or healing.message + return output + + +async def pre_tool_use_hook( + input_data: HookInput, + tool_use_id: Optional[str], + context: HookContext, + *, + auto_heal: Optional[bool] = None, +) -> dict[str, Any]: + """PreToolUse hook for real-time failure prevention. + + Default behaviour is observe-and-block: detection runs, severe + findings return `permissionDecision: block`. When `auto_heal=True` + (or `PISAMA_AUTO_HEAL=1` in env), a SAFE-risk fix is fetched + synchronously via `heal_now()` and surfaced as a `permissionDecision: + deny` plus systemMessage patch — Claude SDK and LangGraph naturally + retry the step with the patched context, breaking the loop in-loop. + DANGEROUS fixes still block; the systemMessage carries an approval URL. + + Args: + input_data: Contains tool_name, tool_input, session_id + tool_use_id: Unique identifier for this tool invocation + context: Hook context with signal + auto_heal: Override the env-based default for this call + + Returns: + Hook output dict with blocking decision or message + """ + if not tool_use_id: + logger.debug("PreToolUse called without tool_use_id, skipping") + return {} + + bridge = get_bridge() + + try: + result = await bridge.analyze_pre_tool(input_data, tool_use_id) + + if result.timed_out: + logger.warning( + f"Detection timeout for {input_data.get('tool_name')}, " + f"allowing to proceed" + ) + return {} + + output = result.to_hook_output() + if not result.should_block: + return output + + heal_flag = auto_heal if auto_heal is not None else _auto_heal_enabled_from_env() + if heal_flag: + healing = await asyncio.to_thread(_attempt_inline_heal, result) + if healing is not None: + if healing.applied: + logger.info( + "auto_heal: applied SAFE fix for %s (%s)", + input_data.get("tool_name"), + result.primary_detection_type, + ) + return _heal_output(healing) + if healing.escalated: + logger.info( + "auto_heal: escalating %s (risk=%s)", + result.primary_detection_type, + healing.risk_level, + ) + return _escalation_output(healing, output) + + logger.info( + f"Blocking tool {input_data.get('tool_name')} " + f"(severity={result.severity})" + ) + return output + + except Exception as e: + logger.error(f"PreToolUse hook error: {e}", exc_info=True) + # Fail open - don't block on errors + return {} + + +class PreToolUseHook: + """Class-based PreToolUse hook with configuration. + + Use this when you need more control over the hook behavior, + such as custom bridge configuration or fail behavior. + + Example: + from pisama.agents.hooks import PreToolUseHook + from pisama.agents import create_bridge, BridgeConfig + + # Custom configuration + config = BridgeConfig(warning_threshold=30, block_threshold=50) + bridge = DetectionBridge(config=config) + + # Create hook with custom bridge + hook = PreToolUseHook(bridge=bridge, fail_open=True) + + # Register with agent + agent.hooks.pre_tool_use = hook + """ + + def __init__( + self, + bridge: Optional[DetectionBridge] = None, + fail_open: bool = True, + auto_heal: Optional[bool] = None, + matcher: Optional[HookMatcher] = None, + ) -> None: + """Initialize the hook. + + Args: + bridge: Custom detection bridge (defaults to global) + fail_open: If True, allow execution on hook errors + auto_heal: Enable in-loop SAFE-risk healing. Defaults to the + `PISAMA_AUTO_HEAL` env flag when None. + matcher: Optional tool and input matcher. Unmatched calls pass + through without running detection. + """ + self.bridge = bridge or get_bridge() + self.fail_open = fail_open + self.auto_heal = ( + auto_heal if auto_heal is not None else _auto_heal_enabled_from_env() + ) + self.matcher = matcher + + async def __call__( + self, + input_data: HookInput, + tool_use_id: Optional[str], + context: HookContext, + ) -> dict[str, Any]: + """Handle PreToolUse event. + + Args: + input_data: Hook input data + tool_use_id: Tool use identifier + context: Hook context + + Returns: + Hook output dict + """ + if not tool_use_id: + return {} + if self.matcher and not self.matcher.matches( + input_data.get("tool_name", ""), + input_data.get("tool_input"), + ): + return {} + + try: + result = await self.bridge.analyze_pre_tool(input_data, tool_use_id) + output = result.to_hook_output() + if not result.should_block or not self.auto_heal: + return output + + healing = await asyncio.to_thread(_attempt_inline_heal, result) + if healing is None: + return output + if healing.applied: + return _heal_output(healing) + if healing.escalated: + return _escalation_output(healing, output) + return output + except Exception as e: + logger.error(f"PreToolUse error: {e}") + if self.fail_open: + return {} + raise diff --git a/src/pisama/agents/indication.py b/src/pisama/agents/indication.py new file mode 100644 index 0000000..a6ed86f --- /dev/null +++ b/src/pisama/agents/indication.py @@ -0,0 +1,300 @@ +"""SDK-side indication channel — out-of-band signal for the developer. + +The PreToolUse hook injects `systemMessage` payloads for the *agent* +to consume. That signals the agent but leaves the *developer* running +the agent blind unless they tail logs. This module fills the gap: a +registered `on_indication` callback fires once per healing outcome +with the same `UserIndication` structure the backend produces, so the +developer can wire it to their own logging / observability / desktop +notification / paging surface. + +Usage: + from pisama.agents import on_indication + + @on_indication + def alert_me(indication): + print(f"[Pisama] {indication.severity}: {indication.headline}") + if indication.action_required: + send_to_pagerduty(indication.to_dict()) + + # Or pass a callable directly: + on_indication(my_callback) + +The SDK auto-fires the registered callback after every `heal_now()` +call. Multiple callbacks supported — first registered, first fired. +Exceptions in a callback are swallowed and logged; one bad handler +doesn't break the others. +""" + +from __future__ import annotations + +import logging +import threading +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class _IndicationMessage: + severity: str + category: str + headline: str + detail: str + action_required: bool + action_summary: Optional[str] = None + + +def _safe_message(detection_type: str, fix: Dict[str, Any]) -> _IndicationMessage: + return _IndicationMessage( + severity="info", + category="auto_healed_safe", + headline=f"Pisama auto-healed {detection_type} ({fix.get('fix_type') or 'SAFE'})", + detail="SAFE fix applied inline — agent retried with patched config.", + action_required=False, + ) + + +def _verified_message(detection_type: str) -> _IndicationMessage: + return _IndicationMessage( + severity="info", + category="auto_healed_verified", + headline=f"Pisama auto-healed {detection_type} (verified MEDIUM)", + detail="MEDIUM fix applied after inline verification predicted improvement.", + action_required=False, + ) + + +def _governance_message( + detection_type: str, + handoff: Dict[str, Any], +) -> _IndicationMessage: + page = handoff.get("page") or "owner" + sla = handoff.get("sla_hours") + sla_str = f" within {sla}h" if sla else "" + return _IndicationMessage( + severity="critical", + category="escalated_governance", + headline=f"GOVERNANCE: {detection_type} — page {page}{sla_str}", + detail="Governance failure — auto-healing disabled by design.", + action_required=True, + action_summary=f"Page {page}{sla_str}; attach evidence.", + ) + + +def _dangerous_message( + detection_type: str, + risk: str, + rec: Dict[str, Any], +) -> _IndicationMessage: + primitive = rec.get("primitive") + detail = "DANGEROUS fix blocked pending review." + if primitive: + detail += f" SDK can verify locally via {primitive}." + action_summary = ( + f"Approve in dashboard OR wire {primitive} locally." + if primitive + else "Review at the approval URL." + ) + return _IndicationMessage( + severity="warning", + category="escalated_dangerous", + headline=f"{detection_type} requires approval (risk: {risk})", + detail=detail, + action_required=True, + action_summary=action_summary, + ) + + +def _observation_message(detection_type: str, result: Any) -> _IndicationMessage: + return _IndicationMessage( + severity="observation", + category="insight_only", + headline=f"{detection_type}: observation logged", + detail=result.message or "No remediation; insight-only.", + action_required=False, + ) + + +def _message_for_healing_result( + result: Any, + *, + detection_type: str, + fix: Dict[str, Any], + rec: Dict[str, Any], + risk: str, + handoff: Dict[str, Any] | None, +) -> _IndicationMessage: + if result.applied and risk == "safe": + return _safe_message(detection_type, fix) + if result.applied and result.verification_passed is True: + return _verified_message(detection_type) + if result.escalated and handoff: + return _governance_message(detection_type, handoff) + if result.escalated: + return _dangerous_message(detection_type, risk, rec) + return _observation_message(detection_type, result) + + +@dataclass +class SDKIndication: + """SDK-side mirror of the backend's UserIndication. + + Kept in the SDK rather than imported from backend so the SDK has + no backend dependency. Field names match exactly so developers can + treat the two interchangeably. + """ + + severity: str # "critical" | "warning" | "info" | "observation" + category: str + detection_type: str + headline: str + detail: str + confidence: float = 0.0 + action_required: bool = False + action_summary: Optional[str] = None + approval_url: Optional[str] = None + evidence_summary: Optional[str] = None + fix_title: Optional[str] = None + fix_type: Optional[str] = None + risk_level: Optional[str] = None + recommendation_source: Optional[str] = None + verification_passed: Optional[bool] = None + sdk_primitive: Optional[str] = None + handoff: Dict[str, Any] = field(default_factory=dict) + tags: List[str] = field(default_factory=list) + + @classmethod + def from_healing_result( + cls, + result: Any, + *, + detection_type: str = "", + confidence: float = 0.0, + ) -> "SDKIndication": + """Build an SDKIndication directly from a HealingResult. + + Mirrors `backend/app/notifications/indication.py::build_indication` + but with the smaller data the SDK has access to (no + baseline_detections, etc.). + """ + fix = result.fix or {} + rec = result.recommended_verification or {} + risk = result.risk_level or "unknown" + verification_passed = result.verification_passed + metadata = fix.get("metadata") if isinstance(fix, dict) else None + handoff = (metadata or {}).get("handoff") if isinstance(metadata, dict) else None + message = _message_for_healing_result( + result, + detection_type=detection_type, + fix=fix, + rec=rec, + risk=risk, + handoff=handoff, + ) + + return cls( + severity=message.severity, + category=message.category, + detection_type=detection_type, + headline=message.headline, + detail=message.detail, + confidence=confidence, + action_required=message.action_required, + action_summary=message.action_summary, + approval_url=result.approval_url, + evidence_summary=_short_evidence(fix), + fix_title=fix.get("title") if isinstance(fix, dict) else None, + fix_type=fix.get("fix_type") if isinstance(fix, dict) else None, + risk_level=risk, + recommendation_source=( + fix.get("recommendation_source") if isinstance(fix, dict) else None + ), + verification_passed=verification_passed, + sdk_primitive=rec.get("primitive") if isinstance(rec, dict) else None, + handoff=handoff or {}, + tags=[message.category], + ) + + def to_dict(self) -> Dict[str, Any]: + return { + "severity": self.severity, + "category": self.category, + "detection_type": self.detection_type, + "headline": self.headline, + "detail": self.detail, + "confidence": round(self.confidence, 3), + "action_required": self.action_required, + "action_summary": self.action_summary, + "approval_url": self.approval_url, + "evidence_summary": self.evidence_summary, + "fix_title": self.fix_title, + "fix_type": self.fix_type, + "risk_level": self.risk_level, + "recommendation_source": self.recommendation_source, + "verification_passed": self.verification_passed, + "sdk_primitive": self.sdk_primitive, + "handoff": dict(self.handoff), + "tags": list(self.tags), + } + + +# --------------------------------------------------------------------- +# Callback registry. Thread-safe. +# --------------------------------------------------------------------- + + +_CALLBACKS_LOCK = threading.Lock() +_CALLBACKS: List[Callable[[SDKIndication], None]] = [] + + +def on_indication(callback: Callable[[SDKIndication], None]) -> Callable[[SDKIndication], None]: + """Register a callback. Usable as a decorator OR as a function call. + + Returns the callback unchanged so it can be used inline: + + @on_indication + def my_handler(indication): ... + """ + with _CALLBACKS_LOCK: + if callback not in _CALLBACKS: + _CALLBACKS.append(callback) + return callback + + +def clear_indication_callbacks() -> None: + """Test hook: drop all registered callbacks.""" + with _CALLBACKS_LOCK: + _CALLBACKS.clear() + + +def _fire(indication: SDKIndication) -> None: + """Dispatch an indication to every registered callback. + + Used internally by `heal_now()` after a healing response is built. + Exceptions are swallowed and logged so one bad handler can't break + the agent loop. + """ + with _CALLBACKS_LOCK: + callbacks = list(_CALLBACKS) + for cb in callbacks: + try: + cb(indication) + except Exception as exc: + logger.warning("on_indication callback %r raised: %s", cb, exc) + + +def _short_evidence(fix: Any) -> Optional[str]: + if not isinstance(fix, dict): + return None + meta = fix.get("metadata") or {} + handoff = meta.get("handoff") if isinstance(meta, dict) else None + if isinstance(handoff, dict): + ev = handoff.get("evidence_summary") + if isinstance(ev, str) and ev: + return ev + desc = fix.get("description") + if isinstance(desc, str): + return desc[:160] + ("…" if len(desc) > 160 else "") + return None diff --git a/src/pisama/agents/openhands_adapter.py b/src/pisama/agents/openhands_adapter.py new file mode 100644 index 0000000..323046c --- /dev/null +++ b/src/pisama/agents/openhands_adapter.py @@ -0,0 +1,482 @@ +"""OpenHands event-stream → Pisama bridge. + +Phase C of plan-to-all-5-unified-glade.md. Stream OpenHands agent +events into Pisama detectors. Two operating modes: + +- **batch** (default): buffer all events through the session, then on + ``on_session_complete()`` either (a) read the Harbor-emitted ATIF + trajectory from the session directory, or (b) synthesize an ATIF + trajectory from the buffered events and POST it to the Pisama backend + via ``analyze_atif``. Returns a full ``AtifAnalyzeResult``. + +- **streaming**: opt-in. ``on_action()`` runs the cheap pattern detectors + (``destructive_command``, plus a lightweight loop heuristic) on the + in-progress span and invokes an optional ``on_streaming_detection`` + callback. Heavy LLM detectors stay in batch mode. v0 only wires + ``destructive_command`` because it's the canonical safety-critical + signal; the rest of the streaming detector list rolls in once their + pattern catalogs are vendored alongside the SDK. + +Usage:: + + from pathlib import Path + from pisama.agents import OpenHandsEventStreamAdapter + + adapter = OpenHandsEventStreamAdapter() # batch mode + # ... wire adapter.on_action / on_observation into the OpenHands + # EventStream subscriber ... + result = adapter.on_session_complete(Path("/path/to/session")) + if result.has_failures: + for d in result.failures: + print(d.detector, d.severity, d.title) +""" +from __future__ import annotations + +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable, Dict, Literal, Optional + +from .atif import ( + DEFAULT_API_URL, + DEFAULT_TIMEOUT_SECONDS, + AtifAnalyzeResult, + analyze_atif, +) + +# Default Harbor-on-OpenHands layout. Used by on_session_complete() to +# locate the agent trajectory inside a session directory. +_DEFAULT_TRAJECTORY_RELPATHS: tuple[str, ...] = ( + "agent/trajectory.json", + "trajectory.json", +) + + +# --------------------------------------------------------------------------- +# Vendored destructive_command patterns +# --------------------------------------------------------------------------- +# Kept in sync with backend/app/detection/destructive_command.py. SDK +# vendors its own copy so streaming mode can match without an HTTP +# round-trip per action. Re-vendor when the backend catalog changes. + +_DESTRUCTIVE_PATTERN_SOURCES: tuple[tuple[str, str, str], ...] = ( + ("rm_rf_root", r"\brm\s+-[a-zA-Z]*[rRf][a-zA-Z]*\s+/(?:\s|$)", "critical"), + ( + "rm_rf_home", + r"\brm\s+-[a-zA-Z]*[rRf][a-zA-Z]*\s+(?:~|\$HOME)(?:/|\s|$)", + "critical", + ), + ("kill_all", r"\bkill\s+-9\s+-1\b", "critical"), + ( + "pkill_runtime", + r"\bpkill\s+(?:-[0-9A-Z]+\s+)*(?:python|node|ruby|sh|bash|perl|php|java)\b", + "high", + ), + ( + "killall_runtime", + r"\bkillall\s+(?:-[0-9A-Z]+\s+)*(?:python|node|ruby|sh|bash|perl|php|java)\b", + "high", + ), + ( + "dd_to_device", + r"\bdd\s+(?:[a-z=]+=\S+\s+)*of=/dev/(?:sd[a-z]|nvme\d+n\d+|hd[a-z]|vd[a-z]|xvd[a-z])", + "critical", + ), + ("mkfs_device", r"\bmkfs(?:\.\w+)?\s+(?:-\S+\s+)*/dev/", "critical"), + ("fork_bomb", r":\(\)\s*\{\s*:\s*\|\s*:\s*&\s*\}\s*;\s*:", "critical"), + ( + "chmod_root_recursive", + r"\bchmod\s+-R\s+\d+\s+/(?:bin|etc|usr|lib|var|root)(?:\s|$|/)", + "high", + ), +) + +_DESTRUCTIVE_PATTERNS: tuple[tuple[str, re.Pattern[str], str], ...] = tuple( + (name, re.compile(src, re.IGNORECASE), sev) + for name, src, sev in _DESTRUCTIVE_PATTERN_SOURCES +) + + +def _match_destructive(command: str) -> tuple[str, str, str] | None: + """Return (pattern_name, severity, matched_text) for the first hit.""" + if not command: + return None + for name, pat, sev in _DESTRUCTIVE_PATTERNS: + m = pat.search(command) + if m: + return name, sev, m.group(0) + return None + + +# --------------------------------------------------------------------------- +# Streaming detection callback shape +# --------------------------------------------------------------------------- + + +@dataclass +class StreamingDetection: + """A streaming-mode hit surfaced by the SDK adapter. + + Mirrors the shape of an orchestrator ``DetectionResult`` but stays + self-contained so the SDK doesn't need a backend round-trip. + """ + detector: str # e.g., "destructive_command" + pattern_name: str # e.g., "rm_rf_root" + severity: str # critical | high | medium | low + confidence: float + command: str # the offending command (truncated) + step_index: int # 1-based index of the action in the session + suggested_action: str # operator-facing remediation + + +StreamingCallback = Callable[[StreamingDetection], None] + + +# --------------------------------------------------------------------------- +# Adapter +# --------------------------------------------------------------------------- + + +@dataclass +class _Event: + """An OpenHands event captured during the session.""" + kind: Literal["action", "observation"] + payload: Dict[str, Any] + + +class OpenHandsEventStreamAdapter: + """Stream OpenHands events into Pisama detectors. + + See module docstring for the streaming vs batch contract. + """ + + def __init__( + self, + *, + api_url: Optional[str] = None, + mode: Literal["streaming", "batch"] = "batch", + on_streaming_detection: Optional[StreamingCallback] = None, + project_id: Optional[str] = None, + agent_name: str = "openhands", + agent_version: str = "unknown", + timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS, + ): + if mode not in ("streaming", "batch"): + raise ValueError(f"mode must be 'streaming' or 'batch', got {mode!r}") + self.api_url = api_url or DEFAULT_API_URL + self.mode = mode + self.on_streaming_detection = on_streaming_detection + self.project_id = project_id + self.agent_name = agent_name + self.agent_version = agent_version + self.timeout_seconds = timeout_seconds + self._events: list[_Event] = [] + + # --------------------------------------------------------------------- + # Event handlers — wire these into OpenHands EventStream subscribers + # --------------------------------------------------------------------- + + def on_action(self, action: Dict[str, Any]) -> Optional[StreamingDetection]: + """Record an OpenHands ``Action`` event. + + Returns a ``StreamingDetection`` when streaming mode is enabled + AND a destructive-command pattern matched the action's command + payload. The user's on_streaming_detection callback is invoked + before this returns. + """ + if not isinstance(action, dict): + return None + self._events.append(_Event("action", action)) + + if self.mode != "streaming": + return None + command = _extract_command_from_openhands_action(action) + if not command: + return None + hit = _match_destructive(command) + if hit is None: + return None + pattern_name, severity, matched_text = hit + detection = StreamingDetection( + detector="destructive_command", + pattern_name=pattern_name, + severity=severity, + confidence=0.99 if severity == "critical" else 0.85, + command=command[:240], + step_index=sum(1 for e in self._events if e.kind == "action"), + suggested_action=( + "Block the action; investigate the agent's reasoning. " + "Scope the operation (specific subpath / pid / non-device " + "target) instead of the blanket form." + ), + ) + if self.on_streaming_detection is not None: + try: + self.on_streaming_detection(detection) + except Exception: + # Never let a user callback crash the adapter; the + # underlying agent stream must keep flowing. + pass + return detection + + def on_observation(self, observation: Dict[str, Any]) -> None: + """Record an OpenHands ``Observation`` event. + + Observations don't currently surface streaming detections (the + cheap pattern detectors all key off the action / command). + """ + if not isinstance(observation, dict): + return + self._events.append(_Event("observation", observation)) + + # --------------------------------------------------------------------- + # Session complete — batch analysis + # --------------------------------------------------------------------- + + def on_session_complete( + self, session_dir: Optional[Path] = None + ) -> AtifAnalyzeResult: + """Run full-orchestrator analysis on the completed session. + + Resolution order: + + 1. If ``session_dir`` is a path and an ATIF trajectory file is + present at one of the conventional Harbor relative paths + (``agent/trajectory.json`` or ``trajectory.json``), POST that + trajectory to ``api_url`` via ``analyze_atif``. This is the + recommended call path: Harbor already writes ATIF. + + 2. Else, synthesize a minimal ATIF trajectory from the buffered + events and POST that. The synthesized trajectory carries the + agent's actions as agent-source ATIF Steps and the + observations as the producing step's ``observation`` field. + + Buffered events are preserved across calls so the adapter can + be re-analyzed multiple times. + """ + if session_dir is not None: + path = Path(session_dir) + traj_path = _find_session_trajectory(path) + if traj_path is not None: + return analyze_atif( + traj_path, + project_id=self.project_id, + api_url=self.api_url, + timeout=self.timeout_seconds, + ) + + trajectory = self._synthesize_trajectory() + if not trajectory.get("steps"): + raise ValueError( + "OpenHandsEventStreamAdapter.on_session_complete: no trajectory " + "found at session_dir and no buffered events to synthesize from" + ) + return analyze_atif( + trajectory, + project_id=self.project_id, + api_url=self.api_url, + timeout=self.timeout_seconds, + ) + + # --------------------------------------------------------------------- + # Helpers + # --------------------------------------------------------------------- + + def event_count(self) -> int: + """Return the number of buffered events.""" + return len(self._events) + + def reset(self) -> None: + """Drop the buffered events. Use between independent sessions.""" + self._events.clear() + + def _synthesize_trajectory(self) -> Dict[str, Any]: + """Build a minimal ATIF v1.7 trajectory dict from buffered events. + + Each action becomes one agent step. The next observation in the + event sequence (if any) attaches to that step as the observation + result. Used as a fallback when no on-disk Harbor trajectory is + available. + """ + steps: list[Dict[str, Any]] = [] + + # Iterate events, pairing each action with the next observation. + i = 0 + step_id = 1 + events = self._events + while i < len(events): + ev = events[i] + if ev.kind == "action": + step = _action_event_to_step(ev.payload, step_id) + # Look ahead for the next observation (may be the very next event). + if ( + i + 1 < len(events) + and events[i + 1].kind == "observation" + and step.get("tool_calls") + ): + obs_payload = events[i + 1].payload + obs = _observation_event_to_observation( + obs_payload, step["tool_calls"][0]["tool_call_id"] + ) + if obs: + step["observation"] = obs + i += 2 + else: + i += 1 + steps.append(step) + step_id += 1 + elif ev.kind == "observation": + # Lone observation with no preceding action — emit as a + # system step so the orchestrator still sees the data. + fallback_message = ev.payload.get("content") or ev.payload.get("message") or "" + steps.append({ + "step_id": step_id, + "source": "system", + "message": str(fallback_message)[:4000], + }) + step_id += 1 + i += 1 + else: + i += 1 + + if not steps: + return {"schema_version": "ATIF-v1.7", "steps": []} + + return { + "schema_version": "ATIF-v1.7", + "agent": { + "name": self.agent_name, + "version": self.agent_version, + }, + "steps": steps, + } + + +# --------------------------------------------------------------------------- +# Module-level conversion helpers +# --------------------------------------------------------------------------- + + +def _find_session_trajectory(session_dir: Path) -> Optional[Path]: + """Resolve a Harbor-emitted trajectory inside a session directory. + + Looks for ``agent/trajectory.json`` first (Harbor multi-trial layout), + falls back to ``trajectory.json`` at the session root. Returns None + when neither is present. + """ + if not session_dir.is_dir(): + return None + for relpath in _DEFAULT_TRAJECTORY_RELPATHS: + candidate = session_dir / relpath + if candidate.exists() and candidate.is_file(): + return candidate + return None + + +# OpenHands action types that involve a shell-style command. The +# command payload is sometimes nested under ``args`` and sometimes at +# the top level (older versions); both shapes are handled. +_OPENHANDS_SHELL_ACTIONS: frozenset[str] = frozenset({ + "run", "execute_bash", "cmd_run", "bash", "shell", +}) + + +def _extract_command_from_openhands_action(action: Dict[str, Any]) -> str: + """Pull a shell-style command string out of an OpenHands action dict. + + Returns empty string when the action is not a shell command (the + destructive_command detector only matches against shell commands). + """ + action_type = str(action.get("action") or action.get("type") or "").strip().lower() + if action_type and action_type not in _OPENHANDS_SHELL_ACTIONS: + return "" + # Common nesting: args.command / args.cmd. Fall back to top-level. + args = action.get("args") or {} + if isinstance(args, dict): + for key in ("command", "cmd", "code", "shell_command"): + val = args.get(key) + if val: + return val if isinstance(val, str) else str(val) + # Top-level fallbacks (older OpenHands shapes). + for key in ("command", "cmd", "code"): + val = action.get(key) + if val: + return val if isinstance(val, str) else str(val) + return "" + + +def _action_event_to_step(action: Dict[str, Any], step_id: int) -> Dict[str, Any]: + """Convert one OpenHands action event into an ATIF Step dict. + + For shell actions the command becomes a ``tool_calls`` entry. For + message actions ("message", "talk") the content becomes the step's + ``message`` field. For other action types we still emit a step with + the action serialized as the message so the trace is complete. + """ + action_type = str(action.get("action") or action.get("type") or "agent").strip().lower() + thought = str(action.get("thought") or (action.get("args") or {}).get("thought") or "") + args = action.get("args") or {} + + if action_type in _OPENHANDS_SHELL_ACTIONS: + command = _extract_command_from_openhands_action(action) + tool_call_id = f"oh_{step_id:03d}" + return { + "step_id": step_id, + "source": "agent", + "message": thought or "", + "reasoning_content": thought or None, + "tool_calls": [{ + "tool_call_id": tool_call_id, + "function_name": action_type, + "arguments": ( + {"command": command} if command else (args if isinstance(args, dict) else {}) + ), + }], + } + if action_type in ("message", "talk"): + # User vs agent attribution: OpenHands uses source="user" / + # source="agent" on the MessageAction itself. + source = str(action.get("source") or "agent").lower() + if source not in ("user", "agent", "system"): + source = "agent" + return { + "step_id": step_id, + "source": source, + "message": str( + action.get("content") + or args.get("content") if isinstance(args, dict) else "" + ) or "", + } + # Generic action — preserve as agent step with serialized payload. + payload_text = thought or str(args)[:4000] + return { + "step_id": step_id, + "source": "agent", + "message": payload_text, + } + + +def _observation_event_to_observation( + observation: Dict[str, Any], source_call_id: str +) -> Optional[Dict[str, Any]]: + """Convert an OpenHands observation event into an ATIF Observation.""" + content = observation.get("content") + if content is None: + extras = observation.get("extras") or {} + if isinstance(extras, dict): + content = extras.get("output") or extras.get("result") + if content is None: + return None + return { + "results": [{ + "source_call_id": source_call_id, + "content": content if isinstance(content, str) else str(content), + }], + } + + +__all__ = [ + "OpenHandsEventStreamAdapter", + "StreamingDetection", + "StreamingCallback", +] diff --git a/src/pisama/agents/session.py b/src/pisama/agents/session.py new file mode 100644 index 0000000..4ca5469 --- /dev/null +++ b/src/pisama/agents/session.py @@ -0,0 +1,250 @@ +"""Session state management for detection context.""" + +import threading +from collections import deque +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Optional + +from pisama_core.traces.models import Span + + +@dataclass +class SessionState: + """State for a single agent session. + + Maintains the context needed for real-time detection, + including recent spans, tool usage statistics, and + blocking state. + """ + + session_id: str + recent_spans: deque = field(default_factory=lambda: deque(maxlen=50)) + tool_counts: dict[str, int] = field(default_factory=dict) + total_cost: float = 0.0 + created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + last_activity: datetime = field(default_factory=lambda: datetime.now(timezone.utc)) + + # Harness-aware fields (for multi-agent orchestration tracing) + agent_role: Optional[str] = None # planner/generator/evaluator/orchestrator/tool + sprint_id: Optional[str] = None # Groups spans into sprint boundaries + context_reset: bool = False # True if context was cleared at session start + + # Blocking state + blocked: bool = False + block_reason: Optional[str] = None + + def add_span(self, span: Span) -> None: + """Add a span to session history. + + Args: + span: The span to add + """ + self.recent_spans.appendleft(span) + self.tool_counts[span.name] = self.tool_counts.get(span.name, 0) + 1 + self.last_activity = datetime.now(timezone.utc) + + def get_context(self, window: int = 10) -> dict[str, Any]: + """Get context for detection. + + Returns: + Context dict with: + - recent_spans: List of recent Span objects + - tool_counts: Dict of tool name -> call count + - total_tools: Total number of tool calls + - session_duration_s: Session duration in seconds + """ + recent = list(self.recent_spans)[:window] + ctx = { + "recent_spans": recent, + "tool_counts": dict(self.tool_counts), + "total_tools": len(self.recent_spans), + "session_duration_s": ( + datetime.now(timezone.utc) - self.created_at + ).total_seconds(), + } + # Include harness-aware fields if set + if self.agent_role: + ctx["agent_role"] = self.agent_role + if self.sprint_id: + ctx["sprint_id"] = self.sprint_id + if self.context_reset: + ctx["context_reset"] = self.context_reset + return ctx + + def get_recent_tool_sequence(self, n: int = 5) -> list[str]: + """Get the sequence of recent tool names. + + Args: + n: Number of recent tools to return + + Returns: + List of tool names, most recent first + """ + return [span.name for span in list(self.recent_spans)[:n]] + + +class SessionManager: + """Thread-safe session state manager. + + Maintains state for multiple sessions, enabling context-aware + detection across tool calls within a session. + + Example: + manager = SessionManager() + + # Get or create session + session = manager.get_or_create("session-123") + + # Add span to session + manager.add_span("session-123", span) + + # Get detection context + context = manager.get_context("session-123", window=10) + """ + + def __init__( + self, + max_sessions: int = 100, + session_ttl_seconds: int = 3600, + ) -> None: + """Initialize session manager. + + Args: + max_sessions: Maximum concurrent sessions to track + session_ttl_seconds: Session expiry time in seconds + """ + self._sessions: dict[str, SessionState] = {} + self._lock = threading.RLock() + self._max_sessions = max_sessions + self._session_ttl = session_ttl_seconds + + def get_or_create(self, session_id: str) -> SessionState: + """Get or create session state. + + Args: + session_id: Unique session identifier + + Returns: + SessionState for the session + """ + with self._lock: + self._cleanup_expired() + + if session_id not in self._sessions: + if len(self._sessions) >= self._max_sessions: + # Remove oldest session + oldest = min( + self._sessions.items(), key=lambda x: x[1].last_activity + ) + del self._sessions[oldest[0]] + + self._sessions[session_id] = SessionState(session_id=session_id) + + return self._sessions[session_id] + + def add_span(self, session_id: str, span: Span) -> None: + """Add span to session. + + Args: + session_id: Session identifier + span: Span to add + """ + session = self.get_or_create(session_id) + session.add_span(span) + + def get_context(self, session_id: str, window: int = 10) -> dict[str, Any]: + """Get detection context for session. + + Args: + session_id: Session identifier + window: Number of recent spans to include + + Returns: + Context dictionary for detectors + """ + session = self.get_or_create(session_id) + return session.get_context(window) + + def is_blocked(self, session_id: str) -> bool: + """Check if session is blocked. + + Args: + session_id: Session identifier + + Returns: + True if session is blocked + """ + with self._lock: + session = self._sessions.get(session_id) + return session.blocked if session else False + + def get_block_reason(self, session_id: str) -> Optional[str]: + """Get the reason a session is blocked. + + Args: + session_id: Session identifier + + Returns: + Block reason or None + """ + with self._lock: + session = self._sessions.get(session_id) + return session.block_reason if session else None + + def block(self, session_id: str, reason: str) -> None: + """Block a session. + + Args: + session_id: Session identifier + reason: Reason for blocking + """ + session = self.get_or_create(session_id) + session.blocked = True + session.block_reason = reason + + def unblock(self, session_id: str) -> None: + """Unblock a session. + + Args: + session_id: Session identifier + """ + with self._lock: + if session_id in self._sessions: + self._sessions[session_id].blocked = False + self._sessions[session_id].block_reason = None + + def clear(self, session_id: str) -> None: + """Clear session state. + + Args: + session_id: Session identifier + """ + with self._lock: + self._sessions.pop(session_id, None) + + def clear_all(self) -> None: + """Clear all sessions.""" + with self._lock: + self._sessions.clear() + + def _cleanup_expired(self) -> None: + """Remove expired sessions.""" + now = datetime.now(timezone.utc) + expired = [ + sid + for sid, state in self._sessions.items() + if (now - state.last_activity).total_seconds() > self._session_ttl + ] + for sid in expired: + del self._sessions[sid] + + @property + def session_count(self) -> int: + """Number of active sessions.""" + with self._lock: + return len(self._sessions) + + +# Global session manager instance +session_manager = SessionManager() diff --git a/src/pisama/agents/tools.py b/src/pisama/agents/tools.py new file mode 100644 index 0000000..a5f319f --- /dev/null +++ b/src/pisama/agents/tools.py @@ -0,0 +1,144 @@ +"""Custom tools for Claude Agent SDK integration. + +Provides pisama_check as a custom tool that agents can call +for self-verification during execution. + +Usage with Claude Agent SDK: + from claude_agent_sdk import ClaudeSDKClient, ClaudeAgentOptions + from pisama.agents import create_check_tool + + options = ClaudeAgentOptions( + custom_tools=[create_check_tool()], + ) + + async with ClaudeSDKClient(options=options) as client: + await client.query("Analyze the auth service incident") + async for message in client.receive_response(): + print(message) +""" + +import logging +from typing import Any, Dict, Optional + +from .check import check + +logger = logging.getLogger(__name__) + +# Tool definition for Claude Agent SDK +PISAMA_CHECK_TOOL_SCHEMA = { + "type": "object", + "properties": { + "output": { + "type": "string", + "description": "The output text you want to verify for issues", + }, + "context": { + "type": "object", + "description": "Context about the task: query, sources, task description", + "properties": { + "query": { + "type": "string", + "description": "The original query or question being answered", + }, + "sources": { + "type": "array", + "items": {"type": "string"}, + "description": "Source documents the output should be grounded in", + }, + "task": { + "type": "string", + "description": "The task description or specification", + }, + }, + }, + "detectors": { + "type": "array", + "items": {"type": "string"}, + "description": ( + "Specific detectors to run (optional). Available: hallucination, " + "derailment, specification, completion, corruption, persona_drift" + ), + }, + }, + "required": ["output"], +} + +PISAMA_CHECK_DESCRIPTION = ( + "Check your output for potential issues before returning it to the user. " + "Use this when you're uncertain about accuracy, when making claims " + "based on retrieved data, or when the task is high-stakes. " + "Returns a confidence score (0-1, higher is better) and any detected " + "issues with suggested fixes. If score > 0.8, the output is likely fine. " + "If score < 0.5, consider revising based on the suggested fixes." +) + + +async def pisama_check_handler( + input_data: Dict[str, Any], + tool_use_id: Optional[str] = None, + context: Any = None, +) -> Dict[str, Any]: + """Handler for the pisama_check custom tool. + + Called by Claude Agent SDK when the agent invokes pisama_check. + + Args: + input_data: Tool input with "output", optional "context" and "detectors" + tool_use_id: Unique tool invocation ID + context: SDK context (signal, etc.) + + Returns: + Check result dict with passed, score, issues, detectors_run + """ + output_text = input_data.get("output", "") + check_context = input_data.get("context") + detectors = input_data.get("detectors") + + if not output_text: + return { + "passed": True, + "score": 1.0, + "issues": [], + "detectors_run": [], + "check_time_ms": 0, + "error": "No output provided to check", + } + + result = await check( + output=output_text, + context=check_context, + detectors=detectors, + ) + + logger.info( + "pisama_check: passed=%s score=%.2f issues=%d time=%dms", + result.get("passed"), + result.get("score", 0), + len(result.get("issues", [])), + result.get("check_time_ms", 0), + ) + + return result + + +def create_check_tool() -> Dict[str, Any]: + """Create a pisama_check custom tool definition for Claude Agent SDK. + + Returns a dict that can be passed to ClaudeAgentOptions.custom_tools. + + Returns: + Tool definition dict with name, description, input_schema, and handler. + + Usage: + from pisama.agents import create_check_tool + + options = ClaudeAgentOptions( + custom_tools=[create_check_tool()], + ) + """ + return { + "name": "pisama_check", + "description": PISAMA_CHECK_DESCRIPTION, + "input_schema": PISAMA_CHECK_TOOL_SCHEMA, + "handler": pisama_check_handler, + } diff --git a/src/pisama/agents/types.py b/src/pisama/agents/types.py new file mode 100644 index 0000000..702abd3 --- /dev/null +++ b/src/pisama/agents/types.py @@ -0,0 +1,144 @@ +"""Type definitions for Claude Agent SDK integration.""" + +from dataclasses import dataclass, field +from typing import Any, Literal, Optional, TypedDict + +# ───────────────────────────────────────────────────────────── +# Agent SDK Types (matching Claude Agent SDK API) +# ───────────────────────────────────────────────────────────── + + +class HookInput(TypedDict, total=False): + """Input data passed to hooks by Claude Agent SDK. + + This matches the structure provided by the Agent SDK's hook system. + """ + + # Tool information + tool_name: str + tool_input: dict[str, Any] + tool_response: Any # PostToolUse only + + # Session info + session_id: str + conversation_id: str + tool_use_id: str + + # Hook metadata + hook_event_name: str + transcript_path: str + cwd: str + + # Extended fields + model: str + usage: dict[str, int] + + +class HookContext(TypedDict, total=False): + """Context provided to hooks by Agent SDK. + + Currently minimal in Python SDK, reserved for future use. + """ + + signal: Any # AbortSignal for cancellation + session_id: str + + +# Permission decision for PreToolUse +PermissionDecision = Literal["allow", "block"] + + +class HookSpecificOutput(TypedDict, total=False): + """Hook-specific output fields.""" + + hookEventName: str + permissionDecision: PermissionDecision + permissionDecisionReason: str + updatedInput: dict[str, Any] + + +class HookJSONOutput(TypedDict, total=False): + """Output returned from hooks to Agent SDK. + + This is the response format expected by the Agent SDK. + """ + + hookSpecificOutput: HookSpecificOutput + systemMessage: str # Inject message into conversation + error: str # Error message if hook fails + continue_: bool # Whether agent should continue (use 'continue' in actual dict) + suppressOutput: bool # Hide from transcript + + +# ───────────────────────────────────────────────────────────── +# Bridge Types +# ───────────────────────────────────────────────────────────── + + +@dataclass +class BridgeResult: + """Result from detection bridge analysis. + + This is the internal result format used by the bridge, + which gets converted to HookJSONOutput for the Agent SDK. + """ + + # Detection decision + should_block: bool = False + severity: int = 0 + issues: list[str] = field(default_factory=list) + recommendations: list[str] = field(default_factory=list) + + # For blocking + block_reason: Optional[str] = None + + # Timing + execution_time_ms: float = 0.0 + timed_out: bool = False + + # Output message + system_message: Optional[str] = None + + # Primary detection metadata for in-loop healing (Track B). The bridge + # populates these when a specific detector fires; the PreToolUse hook + # uses them to call `heal_now()`. When unset, the hook falls back to + # the standard block-and-warn path. + primary_detection_type: Optional[str] = None + primary_details: Optional[dict[str, Any]] = None + framework: Optional[str] = None + + def to_hook_output(self) -> dict[str, Any]: + """Convert to Agent SDK hook output format. + + Returns: + Dictionary in HookJSONOutput format + """ + output: dict[str, Any] = {} + + if self.should_block: + output["hookSpecificOutput"] = { + "hookEventName": "PreToolUse", + "permissionDecision": "block", + "permissionDecisionReason": self.block_reason or "Blocked by MAO detection", + } + + if self.system_message: + output["systemMessage"] = self.system_message + + return output + + def to_dict(self) -> dict[str, Any]: + """Convert to dictionary for logging/serialization.""" + return { + "should_block": self.should_block, + "severity": self.severity, + "issues": self.issues, + "recommendations": self.recommendations, + "block_reason": self.block_reason, + "execution_time_ms": self.execution_time_ms, + "timed_out": self.timed_out, + "system_message": self.system_message, + "primary_detection_type": self.primary_detection_type, + "primary_details": self.primary_details, + "framework": self.framework, + } diff --git a/src/pisama/auto/__init__.py b/src/pisama/auto/__init__.py new file mode 100644 index 0000000..9045486 --- /dev/null +++ b/src/pisama/auto/__init__.py @@ -0,0 +1,91 @@ +"""Pisama Auto-Instrumentation. + +Zero-code instrumentation for LLM applications. +Automatically patches supported libraries to emit OTEL traces +that Pisama can analyze for failure detection. + +This is the in-package home of the auto-instrumentation that used to ship +only as the standalone ``pisama-auto`` distribution. The public API +(``init()`` signature and behavior, the ``patches`` submodule) is +unchanged; only the import path moved. ``pisama-auto`` on PyPI stays +published and, in a later release, becomes a thin shim re-exporting this +module so ``import pisama_auto`` keeps working. + +Usage: + import pisama.auto + pisama.auto.init(api_key="ps_...") + + # All subsequent LLM calls are automatically traced + import anthropic + client = anthropic.Anthropic() + response = client.messages.create(...) # <-- automatically traced + +Nothing here is imported by ``import pisama`` or by ``import pisama.auto`` +alone -- ``opentelemetry`` and ``wrapt`` (the ``pisama[auto]`` extra) are +only imported inside ``init()``, once a caller actually opts in. +""" + +import logging +from typing import Optional + +logger = logging.getLogger("pisama.auto") + +__version__ = "0.2.1" +_initialized = False + + +def init( + api_key: Optional[str] = None, + endpoint: Optional[str] = None, + service_name: str = "pisama-auto", + auto_patch: bool = True, +) -> None: + """Initialize Pisama auto-instrumentation. + + Sets up OTEL tracing and patches supported LLM libraries to automatically + emit traces that Pisama can analyze. + + Args: + api_key: Pisama API key (ps_...). Also reads PISAMA_API_KEY env var. + endpoint: Pisama OTEL ingestion endpoint. Also reads PISAMA_ENDPOINT env var. + If not set, defaults to the Pisama platform ingest endpoint. + service_name: Service name for OTEL resource. + auto_patch: If True, automatically patch all detected libraries. + """ + global _initialized + if _initialized: + logger.debug("Pisama auto-instrumentation already initialized") + return + + import os + api_key = api_key or os.environ.get("PISAMA_API_KEY") + endpoint = endpoint or os.environ.get("PISAMA_ENDPOINT") + + if not api_key: + logger.warning( + "No Pisama API key provided. Set PISAMA_API_KEY or pass api_key to init(). " + "Traces will be generated but not exported." + ) + + # Set up OTEL tracer + from pisama.auto._tracer import setup_tracer + if endpoint: + setup_tracer(api_key=api_key, endpoint=endpoint, service_name=service_name) + else: + setup_tracer(api_key=api_key, service_name=service_name) + + # Auto-patch detected libraries + if auto_patch: + from pisama.auto.patches import patch_all + patched = patch_all() + if patched: + logger.info(f"Pisama: auto-instrumented {', '.join(patched)}") + else: + logger.info("Pisama: initialized (no patchable libraries detected yet)") + + _initialized = True + + +def is_initialized() -> bool: + """Check if Pisama auto-instrumentation is initialized.""" + return _initialized diff --git a/src/pisama/auto/_tracer.py b/src/pisama/auto/_tracer.py new file mode 100644 index 0000000..310fa21 --- /dev/null +++ b/src/pisama/auto/_tracer.py @@ -0,0 +1,264 @@ +"""OTEL tracer setup for Pisama auto-instrumentation.""" + +import json +import logging +import os +import urllib.error +import urllib.request +from typing import Optional + +from opentelemetry import trace +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import ( + BatchSpanProcessor, + ConsoleSpanExporter, + SpanExporter, + SpanExportResult, +) +from opentelemetry.trace import SpanKind, format_span_id, format_trace_id + +from pisama.auto import __version__ + +logger = logging.getLogger("pisama.auto") + +_tracer: Optional[trace.Tracer] = None + +# The platform ingest route (the default endpoint, or a self-hosted instance +# with the same path). It requires a JWT bearer minted from the API key via +# POST /auth/token plus an OTLP JSON body — a stock OTLP exporter sends +# protobuf with the raw key, which the route rejects (401/422) and the spans +# are silently lost. Endpoints without this suffix (e.g. a local OTLP +# collector like pisama-watch's /v1/traces) keep the stock exporter path. +_PLATFORM_INGEST_SUFFIX = "/api/v1/traces/ingest" + +# OTel SDK SpanKind -> OTLP wire enum. +_OTLP_SPAN_KINDS = { + SpanKind.INTERNAL: 1, + SpanKind.SERVER: 2, + SpanKind.CLIENT: 3, + SpanKind.PRODUCER: 4, + SpanKind.CONSUMER: 5, +} + + +def _encode_attr_value(value) -> dict: + if isinstance(value, bool): + return {"boolValue": value} + if isinstance(value, int): + # OTLP/JSON carries int64 as a string. + return {"intValue": str(value)} + if isinstance(value, float): + return {"doubleValue": value} + if isinstance(value, (list, tuple)): + return {"arrayValue": {"values": [_encode_attr_value(v) for v in value]}} + return {"stringValue": str(value)} + + +def _encode_attributes(attributes) -> list: + return [ + {"key": key, "value": _encode_attr_value(value)} + for key, value in (attributes or {}).items() + ] + + +def _encode_span(span) -> dict: + """Encode a ReadableSpan as an OTLP/JSON span dict.""" + context = span.get_span_context() + status = {"code": span.status.status_code.value} + if span.status.description: + status["message"] = span.status.description + encoded = { + "traceId": format_trace_id(context.trace_id), + "spanId": format_span_id(context.span_id), + "name": span.name, + "kind": _OTLP_SPAN_KINDS.get(span.kind, 1), + "startTimeUnixNano": str(span.start_time or 0), + "endTimeUnixNano": str(span.end_time or span.start_time or 0), + "attributes": _encode_attributes(span.attributes), + "status": status, + "events": [ + { + "name": event.name, + "timeUnixNano": str(event.timestamp), + "attributes": _encode_attributes(event.attributes), + } + for event in span.events or [] + ], + } + if span.parent is not None: + encoded["parentSpanId"] = format_span_id(span.parent.span_id) + return encoded + + +def _encode_resource_spans(spans) -> dict: + """Wrap ReadableSpans in the OTLP/JSON resourceSpans envelope.""" + grouped: list = [] # [(resource, [span, ...])] — one provider, usually one + for span in spans: + for resource, bucket in grouped: + if resource is span.resource: + bucket.append(span) + break + else: + grouped.append((span.resource, [span])) + return { + "resourceSpans": [ + { + "resource": {"attributes": _encode_attributes(resource.attributes)}, + "scopeSpans": [ + { + "scope": {"name": "pisama_auto", "version": __version__}, + "spans": [_encode_span(s) for s in bucket], + } + ], + } + for resource, bucket in grouped + ] + } + + +def _post_json(url: str, payload: dict, headers: dict, timeout: float): + """POST JSON, returning (status_code, body). 4xx/5xx return, not raise.""" + request = urllib.request.Request( + url, + data=json.dumps(payload).encode("utf-8"), + headers={"Content-Type": "application/json", **headers}, + method="POST", + ) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + return response.status, response.read().decode("utf-8", "replace") + except urllib.error.HTTPError as error: + return error.code, error.read().decode("utf-8", "replace") + + +class PisamaPlatformExporter(SpanExporter): + """Span exporter for the Pisama platform ingest API. + + Exchanges the API key for a JWT (cached; one re-exchange on 401, mirroring + pisama's PlatformSession) and POSTs OTLP JSON the ingest route's + TraceIngestRequest schema accepts. Stdlib HTTP only, so it needs nothing + beyond opentelemetry-sdk. + """ + + def __init__(self, endpoint: str, api_key: str, timeout: float = 30.0, _post=None): + self._endpoint = endpoint + self._token_url = endpoint[: -len("/traces/ingest")] + "/auth/token" + self._api_key = api_key + self._timeout = timeout + # `_post` is a test seam: (url, payload, headers, timeout) -> (status, body). + self._post = _post or _post_json + self._jwt: Optional[str] = None + + def export(self, spans) -> SpanExportResult: + if not spans: + return SpanExportResult.SUCCESS + payload = _encode_resource_spans(spans) + try: + status, body = self._send(payload) + except Exception as exc: + logger.warning(f"Pisama span export failed: {exc}") + return SpanExportResult.FAILURE + if 200 <= status < 300: + return SpanExportResult.SUCCESS + logger.warning(f"Pisama span export rejected: HTTP {status} {body[:200]}") + return SpanExportResult.FAILURE + + def _send(self, payload: dict): + status, body = self._post( + self._endpoint, payload, self._auth_headers(), self._timeout + ) + if status == 401: + # JWT expired mid-run — re-exchange the API key once. + self._jwt = None + status, body = self._post( + self._endpoint, payload, self._auth_headers(), self._timeout + ) + return status, body + + def _auth_headers(self) -> dict: + if self._jwt is None: + status, body = self._post( + self._token_url, + {"api_key": self._api_key, "scope": "ingest"}, + {}, + self._timeout, + ) + if not 200 <= status < 300: + raise RuntimeError( + f"API key exchange failed: HTTP {status} {body[:200]}" + ) + self._jwt = json.loads(body)["access_token"] + return {"Authorization": f"Bearer {self._jwt}"} + + def shutdown(self) -> None: + pass + + def force_flush(self, timeout_millis: int = 30_000) -> bool: + return True + + +def setup_tracer( + api_key: Optional[str] = None, + endpoint: str = "https://api.pisama.ai/api/v1/traces/ingest", + service_name: str = "pisama-auto", +) -> trace.Tracer: + """Set up the OTEL tracer with Pisama exporter.""" + global _tracer + + resource = Resource.create({ + "service.name": service_name, + "pisama.sdk": "pisama-auto", + "pisama.sdk.version": __version__, + }) + + provider = TracerProvider(resource=resource) + + if api_key: + if endpoint.rstrip("/").endswith(_PLATFORM_INGEST_SUFFIX): + # Platform ingest: JWT-authenticated OTLP JSON. The stock OTLP + # exporters can't speak this route (see _PLATFORM_INGEST_SUFFIX). + exporter = PisamaPlatformExporter( + endpoint=endpoint.rstrip("/"), api_key=api_key + ) + provider.add_span_processor(BatchSpanProcessor(exporter)) + logger.debug(f"Pisama platform exporter configured: {endpoint}") + else: + # Honor OTEL_EXPORTER_OTLP_PROTOCOL=grpc for environments that require + # gRPC transport. Defaults to HTTP/protobuf which works without extra + # opentelemetry-exporter-otlp-proto-grpc deps in the common case. + protocol = os.environ.get("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf").lower() + try: + if protocol in ("grpc", "otlp/grpc"): + from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( + OTLPSpanExporter, + ) + else: + from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( + OTLPSpanExporter, + ) + + exporter = OTLPSpanExporter( + endpoint=endpoint, + headers={"Authorization": f"Bearer {api_key}"}, + ) + provider.add_span_processor(BatchSpanProcessor(exporter)) + logger.debug(f"Pisama OTEL exporter configured: {endpoint} (protocol={protocol})") + except ImportError: + logger.warning("OTLP exporter not available, using console exporter") + provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter())) + else: + # No API key — still trace locally for debugging + provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter())) + + trace.set_tracer_provider(provider) + _tracer = trace.get_tracer("pisama_auto", __version__) + return _tracer + + +def get_tracer() -> trace.Tracer: + """Get the Pisama tracer. Sets up a default if not initialized.""" + global _tracer + if _tracer is None: + _tracer = trace.get_tracer("pisama_auto", __version__) + return _tracer diff --git a/src/pisama/auto/patches/__init__.py b/src/pisama/auto/patches/__init__.py new file mode 100644 index 0000000..3d79984 --- /dev/null +++ b/src/pisama/auto/patches/__init__.py @@ -0,0 +1,70 @@ +"""Auto-patching for LLM libraries. + +Detects installed libraries and patches them to emit OTEL spans +with gen_ai.* semantic conventions. +""" + +import importlib +import logging +from typing import List + +logger = logging.getLogger("pisama.auto") + +# Map of library name -> patch module +_PATCHABLE = { + "anthropic": ".anthropic_patch", + "openai": ".openai_patch", +} + +_patched: List[str] = [] + + +def patch_all() -> List[str]: + """Patch all detected LLM libraries. + + Returns: + List of library names that were successfully patched. + """ + for lib_name, patch_module in _PATCHABLE.items(): + if lib_name in _patched: + continue + try: + importlib.import_module(lib_name) + except ImportError: + continue + + try: + mod = importlib.import_module(patch_module, package="pisama.auto.patches") + mod.patch() + _patched.append(lib_name) + logger.debug(f"Patched {lib_name}") + except Exception as e: + logger.warning(f"Failed to patch {lib_name}: {e}") + + return list(_patched) + + +def patch(library: str) -> bool: + """Patch a specific library. + + Args: + library: Library name (anthropic, openai) + + Returns: + True if patched successfully + """ + if library in _patched: + return True + + if library not in _PATCHABLE: + logger.warning(f"Unknown library: {library}. Supported: {list(_PATCHABLE.keys())}") + return False + + try: + mod = importlib.import_module(_PATCHABLE[library], package="pisama.auto.patches") + mod.patch() + _patched.append(library) + return True + except Exception as e: + logger.warning(f"Failed to patch {library}: {e}") + return False diff --git a/src/pisama/auto/patches/anthropic_patch.py b/src/pisama/auto/patches/anthropic_patch.py new file mode 100644 index 0000000..79ad707 --- /dev/null +++ b/src/pisama/auto/patches/anthropic_patch.py @@ -0,0 +1,154 @@ +"""Auto-instrumentation patch for the Anthropic Python SDK. + +Wraps anthropic.Anthropic.messages.create() and .stream() to emit +OTEL spans with gen_ai.* semantic conventions. +""" + +import logging +from typing import Any + +import wrapt +from opentelemetry.trace import StatusCode + +logger = logging.getLogger("pisama.auto") + +_original_create = None +_original_stream = None + + +def patch() -> None: + """Patch the Anthropic SDK to emit OTEL spans.""" + import anthropic + + messages_cls = anthropic.resources.Messages + + global _original_create, _original_stream + _original_create = messages_cls.create + _original_stream = getattr(messages_cls, "stream", None) + + wrapt.wrap_function_wrapper(messages_cls, "create", _traced_create) + + if _original_stream: + wrapt.wrap_function_wrapper(messages_cls, "stream", _traced_stream) + + logger.debug("Anthropic SDK patched") + + +def _traced_create(wrapped, instance, args, kwargs) -> Any: + """Wrapper for messages.create that adds OTEL tracing.""" + from pisama.auto._tracer import get_tracer + + tracer = get_tracer() + model = kwargs.get("model", "unknown") + span_name = f"gen_ai.chat {model}" + + with tracer.start_as_current_span(span_name) as span: + # Set gen_ai.* attributes + span.set_attribute("gen_ai.system", "anthropic") + span.set_attribute("gen_ai.request.model", model) + span.set_attribute("gen_ai.operation.name", "chat") + + max_tokens = kwargs.get("max_tokens") + if max_tokens: + span.set_attribute("gen_ai.request.max_tokens", max_tokens) + + temperature = kwargs.get("temperature") + if temperature is not None: + span.set_attribute("gen_ai.request.temperature", temperature) + + # Record input messages count + messages = kwargs.get("messages", []) + span.set_attribute("gen_ai.request.messages_count", len(messages)) + + system = kwargs.get("system") + if system: + span.set_attribute("gen_ai.request.has_system", True) + + try: + response = wrapped(*args, **kwargs) + + # Record response attributes + span.set_attribute("gen_ai.response.model", getattr(response, "model", model)) + span.set_attribute("gen_ai.response.stop_reason", getattr(response, "stop_reason", "")) + + usage = getattr(response, "usage", None) + if usage: + input_tokens = getattr(usage, "input_tokens", 0) + output_tokens = getattr(usage, "output_tokens", 0) + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens) + + return response + + except Exception as e: + span.set_attribute("error.type", type(e).__name__) + span.set_attribute("error.message", str(e)[:500]) + span.set_status(StatusCode.ERROR, str(e)[:200]) + raise + + +class _TracedStream: + """Wrapper that ends the OTEL span when the stream is consumed or closed.""" + + def __init__(self, stream, span): + self._stream = stream + self._span = span + + def __enter__(self): + self._stream.__enter__() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + try: + result = self._stream.__exit__(exc_type, exc_val, exc_tb) + if exc_type: + self._span.set_attribute("error.type", exc_type.__name__) + self._span.set_status(StatusCode.ERROR, str(exc_val)[:200]) + return result + finally: + self._span.end() + + def __iter__(self): + return self + + def __next__(self): + try: + return next(self._stream) + except StopIteration: + self._span.end() + raise + except Exception as e: + self._span.set_attribute("error.type", type(e).__name__) + self._span.set_status(StatusCode.ERROR, str(e)[:200]) + self._span.end() + raise + + def __getattr__(self, name): + return getattr(self._stream, name) + + +def _traced_stream(wrapped, instance, args, kwargs) -> Any: + """Wrapper for messages.stream that adds OTEL tracing.""" + from pisama.auto._tracer import get_tracer + + tracer = get_tracer() + model = kwargs.get("model", "unknown") + span_name = f"gen_ai.chat.stream {model}" + + span = tracer.start_span(span_name) + span.set_attribute("gen_ai.system", "anthropic") + span.set_attribute("gen_ai.request.model", model) + span.set_attribute("gen_ai.operation.name", "chat.stream") + + messages = kwargs.get("messages", []) + span.set_attribute("gen_ai.request.messages_count", len(messages)) + + try: + result = wrapped(*args, **kwargs) + return _TracedStream(result, span) + except Exception as e: + span.set_attribute("error.type", type(e).__name__) + span.set_status(StatusCode.ERROR, str(e)[:200]) + span.end() + raise diff --git a/src/pisama/auto/patches/openai_patch.py b/src/pisama/auto/patches/openai_patch.py new file mode 100644 index 0000000..dbb5aa3 --- /dev/null +++ b/src/pisama/auto/patches/openai_patch.py @@ -0,0 +1,83 @@ +"""Auto-instrumentation patch for the OpenAI Python SDK. + +Wraps openai.ChatCompletion.create() and the v1+ client to emit +OTEL spans with gen_ai.* semantic conventions. +""" + +import logging +from typing import Any + +import wrapt + +logger = logging.getLogger("pisama.auto") + + +def patch() -> None: + """Patch the OpenAI SDK to emit OTEL spans.""" + import openai + + # OpenAI v1+ uses client.chat.completions.create() + if hasattr(openai, "resources"): + completions_cls = openai.resources.chat.Completions + wrapt.wrap_function_wrapper(completions_cls, "create", _traced_create) + logger.debug("OpenAI SDK v1+ patched") + else: + logger.warning("OpenAI SDK version not supported for auto-patching") + + +def _traced_create(wrapped, instance, args, kwargs) -> Any: + """Wrapper for chat.completions.create that adds OTEL tracing.""" + from pisama.auto._tracer import get_tracer + + tracer = get_tracer() + model = kwargs.get("model", "unknown") + span_name = f"gen_ai.chat {model}" + + with tracer.start_as_current_span(span_name) as span: + span.set_attribute("gen_ai.system", "openai") + span.set_attribute("gen_ai.request.model", model) + span.set_attribute("gen_ai.operation.name", "chat") + + max_tokens = kwargs.get("max_tokens") + if max_tokens: + span.set_attribute("gen_ai.request.max_tokens", max_tokens) + + temperature = kwargs.get("temperature") + if temperature is not None: + span.set_attribute("gen_ai.request.temperature", temperature) + + messages = kwargs.get("messages", []) + span.set_attribute("gen_ai.request.messages_count", len(messages)) + + stream = kwargs.get("stream", False) + if stream: + span.set_attribute("gen_ai.operation.name", "chat.stream") + + try: + response = wrapped(*args, **kwargs) + + if not stream and response: + resp_model = getattr(response, "model", model) + span.set_attribute("gen_ai.response.model", resp_model) + + choices = getattr(response, "choices", []) + if choices: + finish_reason = getattr(choices[0], "finish_reason", "") + span.set_attribute("gen_ai.response.finish_reason", finish_reason or "") + + usage = getattr(response, "usage", None) + if usage: + input_tokens = getattr(usage, "prompt_tokens", 0) + output_tokens = getattr(usage, "completion_tokens", 0) + span.set_attribute("gen_ai.usage.input_tokens", input_tokens) + span.set_attribute("gen_ai.usage.output_tokens", output_tokens) + span.set_attribute("gen_ai.usage.total_tokens", input_tokens + output_tokens) + + return response + + except Exception as e: + span.set_attribute("error.type", type(e).__name__) + span.set_attribute("error.message", str(e)[:500]) + from opentelemetry.trace import StatusCode + span.set_status(StatusCode.ERROR, str(e)[:200]) + raise diff --git a/tests/agents/__init__.py b/tests/agents/__init__.py new file mode 100644 index 0000000..2f602b7 --- /dev/null +++ b/tests/agents/__init__.py @@ -0,0 +1 @@ +"""Tests for pisama.agents (mirrored from pisama-agent-sdk).""" diff --git a/tests/agents/fixtures/harbor/PROVENANCE.toml b/tests/agents/fixtures/harbor/PROVENANCE.toml new file mode 100644 index 0000000..c184777 --- /dev/null +++ b/tests/agents/fixtures/harbor/PROVENANCE.toml @@ -0,0 +1,21 @@ +format_version = 1 +source_repository = "https://github.com/harbor-framework/harbor" +source_root = "tests/golden/terminus_2" +source_license = "Apache-2.0" +source_revision = "00c19fe2a9c1b9b7ed07efc270412007ac4cb3da" +imported_on = "2026-07-23" + +[[fixtures]] +path = "hello-world-context-summarization-linear-history.trajectory.cont-1.json" +schema_version = "ATIF-v1.7" +sha256 = "054dfbe7e9835830dca2d7c58dc01a8c7129362538b49df4f50468a875d69082" + +[[fixtures]] +path = "hello-world-context-summarization-linear-history.trajectory.json" +schema_version = "ATIF-v1.7" +sha256 = "06b56ab49c627ae3410a0dcb8650a41bb1b30f113d601d10060e590a5018c2bb" + +[[fixtures]] +path = "hello-world-context-summarization.trajectory.summarization-1-summary.json" +schema_version = "ATIF-v1.7" +sha256 = "4b2fa414287428143d74673fef51f2e6f32b22ea9118714b8208cd15c3c34c27" diff --git a/tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.cont-1.json b/tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.cont-1.json new file mode 100644 index 0000000..f5784c3 --- /dev/null +++ b/tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.cont-1.json @@ -0,0 +1,118 @@ +{ + "schema_version": "ATIF-v1.7", + "session_id": "NORMALIZED_SESSION_ID", + "agent": { + "name": "terminus-2", + "version": "2.0.0", + "model_name": "openai/gpt-4o", + "extra": { + "parser": "json", + "continuation_index": 1 + } + }, + "steps": [ + { + "step_id": 1, + "source": "user", + "message": "You are an AI assistant tasked with solving command-line tasks in a Linux environment. You will be given a task description and the output from previously executed commands. Your goal is to solve the task by providing batches of shell commands.\n\nFormat your response as JSON with the following structure:\n\n{\n \"analysis\": \"Analyze the current state based on the terminal output provided. What do you see? What has been accomplished? What still needs to be done?\",\n \"plan\": \"Describe your plan for the next steps. What commands will you run and why? Be specific about what you expect each command to accomplish.\",\n \"commands\": [\n {\n \"keystrokes\": \"ls -la\\n\",\n \"duration\": 0.1\n },\n {\n \"keystrokes\": \"cd project\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": true\n}\n\nRequired fields:\n- \"analysis\": Your analysis of the current situation\n- \"plan\": Your plan for the next steps\n- \"commands\": Array of command objects to execute\n\nOptional fields:\n- \"task_complete\": Boolean indicating if the task is complete (defaults to false if not present)\n\nCommand object structure:\n- \"keystrokes\": String containing the exact keystrokes to send to the terminal (required)\n- \"duration\": Number of seconds to wait for the command to complete before the next command will be executed (defaults to 1.0 if not present)\n\nIMPORTANT: The text inside \"keystrokes\" will be used completely verbatim as keystrokes. Write commands exactly as you want them sent to the terminal:\n- You must end every command with a newline (\\n) or it will not execute.\n- For special key sequences, use tmux-style escape sequences:\n - C-c for Ctrl+C\n - C-d for Ctrl+D\n\nThe \"duration\" attribute specifies the number of seconds to wait for the command to complete (default: 1.0) before the next command will be executed. On immediate tasks (e.g., cd, ls, echo, cat) set a duration of 0.1 seconds. On commands (e.g., gcc, find, rustc) set a duration of 1.0 seconds. On slow commands (e.g., make, python3 [long running script], wget [file]) set an appropriate duration as you determine necessary.\n\nIt is better to set a smaller duration than a longer duration. It is always possible to wait again if the prior output has not finished, by running {\"keystrokes\": \"\", \"duration\": 10.0} on subsequent requests to wait longer. Never wait longer than 60 seconds; prefer to poll to see intermediate result status.\n\nImportant notes:\n- Each command's keystrokes are sent exactly as written to the terminal\n- Do not include extra whitespace before or after the keystrokes unless it's part of the intended command\n- Extra text before or after the JSON will generate warnings but be tolerated\n- The JSON must be valid - use proper escaping for quotes and special characters within strings\n- Commands array can be empty if you want to wait without taking action\n\nTask Description:\nCreate a file called hello.txt with \"Hello, world!\" as the content.\n\nCurrent terminal state:\nCurrent Terminal Screen:\n\n", + "is_copied_context": true + }, + { + "step_id": 2, + "source": "user", + "message": "You are picking up work from a previous AI agent on this task:\n\n**Original Task:** Create a file called hello.txt with \"Hello, world!\" as the content.\n\n**Summary from Previous Agent:**\nI have been working on creating hello.txt. Here's what I've accomplished:\n\n1. **Major Actions Completed**:\n - Created a test directory with `mkdir test_dir`\n - Created file1.txt, file2.txt in the test directory\n - Attempted to list files to verify creation\n\n2. **Important Information Learned**:\n - All test files are located in /app/test_dir/\n - The working directory is /app/\n - The terminal environment is working correctly\n\n3. **Challenging Problems Addressed**:\n - None so far, everything has been straightforward\n\n4. **Current Status**:\n - Test setup is partially complete\n - Still need to create the main hello.txt file with \"Hello, world!\" content in /app/\n\n**Current Terminal Screen:**\nroot@CONTAINER_ID:/app# mkdir test_dir\nroot@CONTAINER_ID:/app# echo 'test1' > test_dir/file1.txt\nroot@CONTAINER_ID:/app# echo 'test2' > test_dir/file2.txt\n\nPlease begin by asking several questions (at least five, more if necessary)\nabout the current state of the solution that are not answered in the summary\nfrom the prior agent. After you ask these questions you will be on your own,\nso ask everything you need to know.", + "is_copied_context": true + }, + { + "step_id": 3, + "source": "agent", + "model_name": "gpt-4o", + "message": "I have the following questions about the current state:\n\n1. What exactly is in the test_dir directory? Which files have been created?\n2. What are the contents of file1.txt and file2.txt?\n3. Has the main hello.txt file been created yet in /app/?\n4. Are there any errors or issues I should be aware of?\n5. What is the exact task requirement again for hello.txt?", + "is_copied_context": true + }, + { + "step_id": 4, + "source": "user", + "message": "Here are the answers the other agent provided.\n\nHere are the answers:\n\n1. The test_dir contains file1.txt and file2.txt\n2. file1.txt contains \"test1\" and file2.txt contains \"test2\"\n3. No, hello.txt has not been created yet\n4. No errors, everything is working fine\n5. The task is to create a file called hello.txt with \"Hello, world!\" as the content in /app/\n\nContinue working on this task from where the previous agent left off. You can no longer ask questions. Please follow the spec to interact with the terminal.", + "is_copied_context": true + }, + { + "step_id": 5, + "source": "agent", + "model_name": "gpt-4o", + "message": "{\n \"analysis\": \"Based on the handoff, I need to create hello.txt with 'Hello, world!' content.\",\n \"plan\": \"I will use printf to create hello.txt with the required content.\",\n \"commands\": [\n {\n \"keystrokes\": \"printf 'Hello, world!\\\\n' > hello.txt\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": false\n}", + "observation": { + "results": [ + { + "content": "New Terminal Output:\n\nroot@CONTAINER_ID:/app# printf 'Hello, world!\\n' > hello.txt\n\n" + } + ] + }, + "metrics": { + "prompt_tokens": 1700, + "completion_tokens": 420, + "cost_usd": 0.00845 + } + }, + { + "step_id": 6, + "source": "agent", + "model_name": "gpt-4o", + "message": "{\n \"analysis\": \"Created hello.txt. Let me verify it.\",\n \"plan\": \"Use cat to check the content.\",\n \"commands\": [\n {\n \"keystrokes\": \"cat hello.txt\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": false\n}", + "observation": { + "results": [ + { + "content": "New Terminal Output:\n\nroot@CONTAINER_ID:/app# cat hello.txt\nHello, world!\n\n" + } + ] + }, + "metrics": { + "prompt_tokens": 850, + "completion_tokens": 40, + "cost_usd": 0.0025249999999999995 + } + }, + { + "step_id": 7, + "source": "agent", + "model_name": "gpt-4o", + "message": "{\n \"analysis\": \"Verified hello.txt has the correct content.\",\n \"plan\": \"Task is complete.\",\n \"commands\": [],\n \"task_complete\": true\n}", + "observation": { + "results": [ + { + "content": "Current terminal state:\nNew Terminal Output:\n\nAre you sure you want to mark the task as complete? This will trigger your solution to be graded and you won't be able to make any further corrections. If so, include \"task_complete\": true in your JSON response again." + } + ] + }, + "metrics": { + "prompt_tokens": 850, + "completion_tokens": 40, + "cost_usd": 0.0025249999999999995 + } + }, + { + "step_id": 8, + "source": "agent", + "model_name": "gpt-4o", + "message": "{\n \"analysis\": \"Yes, confirming task completion.\",\n \"plan\": \"Final confirmation.\",\n \"commands\": [],\n \"task_complete\": true\n}", + "observation": { + "results": [ + { + "content": "New Terminal Output:\n\n" + } + ] + }, + "metrics": { + "prompt_tokens": 850, + "completion_tokens": 30, + "cost_usd": 0.002425 + } + } + ], + "final_metrics": { + "total_prompt_tokens": 7802, + "total_completion_tokens": 1030, + "total_cached_tokens": 0, + "total_cost_usd": 0.029804999999999998 + } +} diff --git a/tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.json b/tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.json new file mode 100644 index 0000000..0a7f72b --- /dev/null +++ b/tests/agents/fixtures/harbor/hello-world-context-summarization-linear-history.trajectory.json @@ -0,0 +1,120 @@ +{ + "schema_version": "ATIF-v1.7", + "session_id": "NORMALIZED_SESSION_ID", + "agent": { + "name": "terminus-2", + "version": "2.0.0", + "model_name": "openai/gpt-4o", + "extra": { + "parser": "json" + } + }, + "steps": [ + { + "step_id": 1, + "source": "user", + "message": "You are an AI assistant tasked with solving command-line tasks in a Linux environment. You will be given a task description and the output from previously executed commands. Your goal is to solve the task by providing batches of shell commands.\n\nFormat your response as JSON with the following structure:\n\n{\n \"analysis\": \"Analyze the current state based on the terminal output provided. What do you see? What has been accomplished? What still needs to be done?\",\n \"plan\": \"Describe your plan for the next steps. What commands will you run and why? Be specific about what you expect each command to accomplish.\",\n \"commands\": [\n {\n \"keystrokes\": \"ls -la\\n\",\n \"duration\": 0.1\n },\n {\n \"keystrokes\": \"cd project\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": true\n}\n\nRequired fields:\n- \"analysis\": Your analysis of the current situation\n- \"plan\": Your plan for the next steps\n- \"commands\": Array of command objects to execute\n\nOptional fields:\n- \"task_complete\": Boolean indicating if the task is complete (defaults to false if not present)\n\nCommand object structure:\n- \"keystrokes\": String containing the exact keystrokes to send to the terminal (required)\n- \"duration\": Number of seconds to wait for the command to complete before the next command will be executed (defaults to 1.0 if not present)\n\nIMPORTANT: The text inside \"keystrokes\" will be used completely verbatim as keystrokes. Write commands exactly as you want them sent to the terminal:\n- You must end every command with a newline (\\n) or it will not execute.\n- For special key sequences, use tmux-style escape sequences:\n - C-c for Ctrl+C\n - C-d for Ctrl+D\n\nThe \"duration\" attribute specifies the number of seconds to wait for the command to complete (default: 1.0) before the next command will be executed. On immediate tasks (e.g., cd, ls, echo, cat) set a duration of 0.1 seconds. On commands (e.g., gcc, find, rustc) set a duration of 1.0 seconds. On slow commands (e.g., make, python3 [long running script], wget [file]) set an appropriate duration as you determine necessary.\n\nIt is better to set a smaller duration than a longer duration. It is always possible to wait again if the prior output has not finished, by running {\"keystrokes\": \"\", \"duration\": 10.0} on subsequent requests to wait longer. Never wait longer than 60 seconds; prefer to poll to see intermediate result status.\n\nImportant notes:\n- Each command's keystrokes are sent exactly as written to the terminal\n- Do not include extra whitespace before or after the keystrokes unless it's part of the intended command\n- Extra text before or after the JSON will generate warnings but be tolerated\n- The JSON must be valid - use proper escaping for quotes and special characters within strings\n- Commands array can be empty if you want to wait without taking action\n\nTask Description:\nCreate a file called hello.txt with \"Hello, world!\" as the content.\n\nCurrent terminal state:\nCurrent Terminal Screen:\n\n" + }, + { + "step_id": 2, + "source": "agent", + "model_name": "gpt-4o", + "message": "{\n \"analysis\": \"Terminal is ready. Let me create a test directory first.\",\n \"plan\": \"Create test_dir to set up some context.\",\n \"commands\": [\n {\n \"keystrokes\": \"mkdir test_dir\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": false\n}", + "observation": { + "results": [ + { + "content": "New Terminal Output:\n\nroot@CONTAINER_ID:/app# mkdir test_dir\n\n" + } + ] + }, + "metrics": { + "prompt_tokens": 682, + "completion_tokens": 60, + "cost_usd": 0.002305 + } + }, + { + "step_id": 3, + "source": "agent", + "model_name": "gpt-4o", + "message": "{\n \"analysis\": \"Created test directory. Let me add some files.\",\n \"plan\": \"Create test files.\",\n \"commands\": [\n {\n \"keystrokes\": \"echo 'test1' > test_dir/file1.txt\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": false\n}", + "observation": { + "results": [ + { + "content": "New Terminal Output:\n\nroot@CONTAINER_ID:/app# echo 'test1' > test_dir/file1.txt\n\n" + } + ] + }, + "metrics": { + "prompt_tokens": 750, + "completion_tokens": 50, + "cost_usd": 0.002375 + } + }, + { + "step_id": 4, + "source": "agent", + "model_name": "gpt-4o", + "message": "{\n \"analysis\": \"Created file1.txt. Let me create another file.\",\n \"plan\": \"Create file2.txt.\",\n \"commands\": [\n {\n \"keystrokes\": \"echo 'test2' > test_dir/file2.txt\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": false\n}", + "observation": { + "results": [ + { + "content": "New Terminal Output:\n\nroot@CONTAINER_ID:/app# echo 'test2' > test_dir/file2.txt\n\n" + } + ] + }, + "metrics": { + "prompt_tokens": 820, + "completion_tokens": 50, + "cost_usd": 0.00255 + } + }, + { + "step_id": 5, + "source": "system", + "message": "Performed context summarization and handoff to continue task.", + "observation": { + "results": [ + { + "subagent_trajectory_ref": [ + { + "session_id": "test-session-linear-history-summarization-1-summary", + "trajectory_path": "trajectory.summarization-1-summary.json", + "extra": { + "summary": "Context summarization 1: Step 1 - Summary generation" + } + }, + { + "session_id": "test-session-linear-history-summarization-1-questions", + "trajectory_path": "trajectory.summarization-1-questions.json", + "extra": { + "summary": "Context summarization 1: Step 2 - Question asking" + } + }, + { + "session_id": "test-session-linear-history-summarization-1-answers", + "trajectory_path": "trajectory.summarization-1-answers.json", + "extra": { + "summary": "Context summarization 1: Step 3 - Answer providing" + } + } + ] + } + ] + }, + "extra": { + "context_management": { + "type": "compaction", + "boundary": "replace" + } + } + } + ], + "final_metrics": { + "total_prompt_tokens": 2252, + "total_completion_tokens": 160, + "total_cached_tokens": 0, + "total_cost_usd": 0.00723 + }, + "continued_trajectory_ref": "trajectory.cont-1.json" +} diff --git a/tests/agents/fixtures/harbor/hello-world-context-summarization.trajectory.summarization-1-summary.json b/tests/agents/fixtures/harbor/hello-world-context-summarization.trajectory.summarization-1-summary.json new file mode 100644 index 0000000..c5cb537 --- /dev/null +++ b/tests/agents/fixtures/harbor/hello-world-context-summarization.trajectory.summarization-1-summary.json @@ -0,0 +1,103 @@ +{ + "schema_version": "ATIF-v1.7", + "session_id": "test-session-context-summarization-summarization-1-summary", + "agent": { + "name": "terminus-2-summarization-summary", + "version": "2.0.0", + "model_name": "openai/gpt-4o", + "extra": { + "parent_session_id": "NORMALIZED_SESSION_ID", + "summarization_index": 1 + } + }, + "steps": [ + { + "step_id": 1, + "source": "user", + "message": "You are an AI assistant tasked with solving command-line tasks in a Linux environment. You will be given a task description and the output from previously executed commands. Your goal is to solve the task by providing batches of shell commands.\n\nFormat your response as JSON with the following structure:\n\n{\n \"analysis\": \"Analyze the current state based on the terminal output provided. What do you see? What has been accomplished? What still needs to be done?\",\n \"plan\": \"Describe your plan for the next steps. What commands will you run and why? Be specific about what you expect each command to accomplish.\",\n \"commands\": [\n {\n \"keystrokes\": \"ls -la\\n\",\n \"duration\": 0.1\n },\n {\n \"keystrokes\": \"cd project\\n\",\n \"duration\": 0.1\n }\n ],\n \"task_complete\": true\n}\n\nRequired fields:\n- \"analysis\": Your analysis of the current situation\n- \"plan\": Your plan for the next steps\n- \"commands\": Array of command objects to execute\n\nOptional fields:\n- \"task_complete\": Boolean indicating if the task is complete (defaults to false if not present)\n\nCommand object structure:\n- \"keystrokes\": String containing the exact keystrokes to send to the terminal (required)\n- \"duration\": Number of seconds to wait for the command to complete before the next command will be executed (defaults to 1.0 if not present)\n\nIMPORTANT: The text inside \"keystrokes\" will be used completely verbatim as keystrokes. Write commands exactly as you want them sent to the terminal:\n- You must end every command with a newline (\\n) or it will not execute.\n- For special key sequences, use tmux-style escape sequences:\n - C-c for Ctrl+C\n - C-d for Ctrl+D\n\nThe \"duration\" attribute specifies the number of seconds to wait for the command to complete (default: 1.0) before the next command will be executed. On immediate tasks (e.g., cd, ls, echo, cat) set a duration of 0.1 seconds. On commands (e.g., gcc, find, rustc) set a duration of 1.0 seconds. On slow commands (e.g., make, python3 [long running script], wget [file]) set an appropriate duration as you determine necessary.\n\nIt is better to set a smaller duration than a longer duration. It is always possible to wait again if the prior output has not finished, by running {\"keystrokes\": \"\", \"duration\": 10.0} on subsequent requests to wait longer. Never wait longer than 60 seconds; prefer to poll to see intermediate result status.\n\nImportant notes:\n- Each command's keystrokes are sent exactly as written to the terminal\n- Do not include extra whitespace before or after the keystrokes unless it's part of the intended command\n- Extra text before or after the JSON will generate warnings but be tolerated\n- The JSON must be valid - use proper escaping for quotes and special characters within strings\n- Commands array can be empty if you want to wait without taking action\n\nTask Description:\nCreate a file called hello.txt with \"Hello, world!\" as the content.\n\nCurrent terminal state:\nCurrent Terminal Screen:\n\n", + "is_copied_context": true + }, + { + "step_id": 2, + "source": "agent", + "model_name": "gpt-4o", + "message": "Analysis: Terminal is ready. Let me create a test directory first.\nPlan: Create test_dir to set up some context.", + "tool_calls": [ + { + "tool_call_id": "call_0_1", + "function_name": "bash_command", + "arguments": { + "keystrokes": "mkdir test_dir\n", + "duration": 0.1 + } + } + ], + "observation": { + "results": [ + { + "source_call_id": "call_0_1", + "content": "New Terminal Output:\n\nroot@CONTAINER_ID:/app# mkdir test_dir\n\n" + } + ] + }, + "is_copied_context": true, + "extra": { + "note": "Metrics omitted to avoid duplication - already recorded in parent trajectory" + } + }, + { + "step_id": 3, + "source": "agent", + "model_name": "gpt-4o", + "message": "Analysis: Created test directory. Let me add some files.\nPlan: Create test files.", + "tool_calls": [ + { + "tool_call_id": "call_1_1", + "function_name": "bash_command", + "arguments": { + "keystrokes": "echo 'test1' > test_dir/file1.txt\n", + "duration": 0.1 + } + } + ], + "observation": { + "results": [ + { + "source_call_id": "call_1_1", + "content": "New Terminal Output:\n\nroot@CONTAINER_ID:/app# echo 'test1' > test_dir/file1.txt\n\n" + } + ] + }, + "is_copied_context": true, + "extra": { + "note": "Metrics omitted to avoid duplication - already recorded in parent trajectory" + } + }, + { + "step_id": 4, + "source": "user", + "message": "You are about to hand off your work to another AI agent.\n Please provide a comprehensive summary of what you have\n accomplished so far on this task:\n\nOriginal Task: Create a file called hello.txt with \"Hello, world!\" as the content.\n\nBased on the conversation history, please provide a detailed summary covering:\n1. **Major Actions Completed** - List each significant command you executed\n and what you learned from it.\n2. **Important Information Learned** - A summary of crucial findings, file\n locations, configurations, error messages, or system state discovered.\n3. **Challenging Problems Addressed** - Any significant issues you\n encountered and how you resolved them.\n4. **Current Status** - Exactly where you are in the task completion process.\n\nBe comprehensive and detailed. The next agent needs to understand everything\n that has happened so far in order to continue." + }, + { + "step_id": 5, + "source": "agent", + "model_name": "gpt-4o", + "message": "I have been working on creating hello.txt. Here's what I've accomplished:\n\n1. **Major Actions Completed**:\n - Created a test directory with `mkdir test_dir`\n - Created file1.txt, file2.txt in the test directory\n - Attempted to list files to verify creation\n\n2. **Important Information Learned**:\n - All test files are located in /app/test_dir/\n - The working directory is /app/\n - The terminal environment is working correctly\n\n3. **Challenging Problems Addressed**:\n - None so far, everything has been straightforward\n\n4. **Current Status**:\n - Test setup is partially complete\n - Still need to create the main hello.txt file with \"Hello, world!\" content in /app/", + "metrics": { + "prompt_tokens": 500, + "completion_tokens": 200, + "cached_tokens": 0, + "cost_usd": 0.0032500000000000003, + "prompt_token_ids": [100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 128, 129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146, 147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, 173, 174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191, 192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209, 210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227, 228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 256, 257, 258, 259, 260, 261, 262, 263, 264, 265, 266, 267, 268, 269, 270, 271, 272, 273, 274, 275, 276, 277, 278, 279, 280, 281, 282, 283, 284, 285, 286, 287, 288, 289, 290, 291, 292, 293, 294, 295, 296, 297, 298, 299, 300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311, 312, 313, 314, 315, 316, 317, 318, 319, 320, 321, 322, 323, 324, 325, 326, 327, 328, 329, 330, 331, 332, 333, 334, 335, 336, 337, 338, 339, 340, 341, 342, 343, 344, 345, 346, 347, 348, 349, 350, 351, 352, 353, 354, 355, 356, 357, 358, 359, 360, 361, 362, 363, 364, 365, 366, 367, 368, 369, 370, 371, 372, 373, 374, 375, 376, 377, 378, 379, 380, 381, 382, 383, 384, 385, 386, 387, 388, 389, 390, 391, 392, 393, 394, 395, 396, 397, 398, 399, 400, 401, 402, 403, 404, 405, 406, 407, 408, 409, 410, 411, 412, 413, 414, 415, 416, 417, 418, 419, 420, 421, 422, 423, 424, 425, 426, 427, 428, 429, 430, 431, 432, 433, 434, 435, 436, 437, 438, 439, 440, 441, 442, 443, 444, 445, 446, 447, 448, 449, 450, 451, 452, 453, 454, 455, 456, 457, 458, 459, 460, 461, 462, 463, 464, 465, 466, 467, 468, 469, 470, 471, 472, 473, 474, 475, 476, 477, 478, 479, 480, 481, 482, 483, 484, 485, 486, 487, 488, 489, 490, 491, 492, 493, 494, 495, 496, 497, 498, 499, 500, 501, 502, 503, 504, 505, 506, 507, 508, 509, 510, 511, 512, 513, 514, 515, 516, 517, 518, 519, 520, 521, 522, 523, 524, 525, 526, 527, 528, 529, 530, 531, 532, 533, 534, 535, 536, 537, 538, 539, 540, 541, 542, 543, 544, 545, 546, 547, 548, 549, 550, 551, 552, 553, 554, 555, 556, 557, 558, 559, 560, 561, 562, 563, 564, 565, 566, 567, 568, 569, 570, 571, 572, 573, 574, 575, 576, 577, 578, 579, 580, 581, 582, 583, 584, 585, 586, 587, 588, 589, 590, 591, 592, 593, 594, 595, 596, 597, 598, 599], + "completion_token_ids": [1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010, 1011, 1012, 1013, 1014, 1015, 1016, 1017, 1018, 1019, 1020, 1021, 1022, 1023, 1024, 1025, 1026, 1027, 1028, 1029, 1030, 1031, 1032, 1033, 1034, 1035, 1036, 1037, 1038, 1039, 1040, 1041, 1042, 1043, 1044, 1045, 1046, 1047, 1048, 1049, 1050, 1051, 1052, 1053, 1054, 1055, 1056, 1057, 1058, 1059, 1060, 1061, 1062, 1063, 1064, 1065, 1066, 1067, 1068, 1069, 1070, 1071, 1072, 1073, 1074, 1075, 1076, 1077, 1078, 1079, 1080, 1081, 1082, 1083, 1084, 1085, 1086, 1087, 1088, 1089, 1090, 1091, 1092, 1093, 1094, 1095, 1096, 1097, 1098, 1099, 1100, 1101, 1102, 1103, 1104, 1105, 1106, 1107, 1108, 1109, 1110, 1111, 1112, 1113, 1114, 1115, 1116, 1117, 1118, 1119, 1120, 1121, 1122, 1123, 1124, 1125, 1126, 1127, 1128, 1129, 1130, 1131, 1132, 1133, 1134, 1135, 1136, 1137, 1138, 1139, 1140, 1141, 1142, 1143, 1144, 1145, 1146, 1147, 1148, 1149, 1150, 1151, 1152, 1153, 1154, 1155, 1156, 1157, 1158, 1159, 1160, 1161, 1162, 1163, 1164, 1165, 1166, 1167, 1168, 1169, 1170, 1171, 1172, 1173, 1174, 1175, 1176, 1177, 1178, 1179, 1180, 1181, 1182, 1183, 1184, 1185, 1186, 1187, 1188, 1189, 1190, 1191, 1192, 1193, 1194, 1195, 1196, 1197, 1198, 1199], + "logprobs": [-0.1, -0.2, -0.30000000000000004, -0.4, -0.5, -0.6000000000000001, -0.7000000000000001, -0.8, -0.9, -1.0, -1.1, -1.2000000000000002, -1.3, -1.4000000000000001, -1.5, -1.6, -1.7000000000000002, -1.8, -1.9000000000000001, -2.0, -2.1, -2.2, -2.3000000000000003, -2.4000000000000004, -2.5, -2.6, -2.7, -2.8000000000000003, -2.9000000000000004, -3.0, -3.1, -3.2, -3.3000000000000003, -3.4000000000000004, -3.5, -3.6, -3.7, -3.8000000000000003, -3.9000000000000004, -4.0, -4.1000000000000005, -4.2, -4.3, -4.4, -4.5, -4.6000000000000005, -4.7, -4.800000000000001, -4.9, -5.0, -5.1000000000000005, -5.2, -5.300000000000001, -5.4, -5.5, -5.6000000000000005, -5.7, -5.800000000000001, -5.9, -6.0, -6.1000000000000005, -6.2, -6.300000000000001, -6.4, -6.5, -6.6000000000000005, -6.7, -6.800000000000001, -6.9, -7.0, -7.1000000000000005, -7.2, -7.300000000000001, -7.4, -7.5, -7.6000000000000005, -7.7, -7.800000000000001, -7.9, -8.0, -8.1, -8.200000000000001, -8.3, -8.4, -8.5, -8.6, -8.700000000000001, -8.8, -8.9, -9.0, -9.1, -9.200000000000001, -9.3, -9.4, -9.5, -9.600000000000001, -9.700000000000001, -9.8, -9.9, -10.0, -10.100000000000001, -10.200000000000001, -10.3, -10.4, -10.5, -10.600000000000001, -10.700000000000001, -10.8, -10.9, -11.0, -11.100000000000001, -11.200000000000001, -11.3, -11.4, -11.5, -11.600000000000001, -11.700000000000001, -11.8, -11.9, -12.0, -12.100000000000001, -12.200000000000001, -12.3, -12.4, -12.5, -12.600000000000001, -12.700000000000001, -12.8, -12.9, -13.0, -13.100000000000001, -13.200000000000001, -13.3, -13.4, -13.5, -13.600000000000001, -13.700000000000001, -13.8, -13.9, -14.0, -14.100000000000001, -14.200000000000001, -14.3, -14.4, -14.5, -14.600000000000001, -14.700000000000001, -14.8, -14.9, -15.0, -15.100000000000001, -15.200000000000001, -15.3, -15.4, -15.5, -15.600000000000001, -15.700000000000001, -15.8, -15.9, -16.0, -16.1, -16.2, -16.3, -16.400000000000002, -16.5, -16.6, -16.7, -16.8, -16.900000000000002, -17.0, -17.1, -17.2, -17.3, -17.400000000000002, -17.5, -17.6, -17.7, -17.8, -17.900000000000002, -18.0, -18.1, -18.2, -18.3, -18.400000000000002, -18.5, -18.6, -18.7, -18.8, -18.900000000000002, -19.0, -19.1, -19.200000000000003, -19.3, -19.400000000000002, -19.5, -19.6, -19.700000000000003, -19.8, -19.900000000000002, -20.0] + } + } + ], + "final_metrics": { + "total_prompt_tokens": 500, + "total_completion_tokens": 200, + "total_cached_tokens": 0, + "total_cost_usd": 0.0032500000000000003 + } +} \ No newline at end of file diff --git a/tests/agents/test_atif.py b/tests/agents/test_atif.py new file mode 100644 index 0000000..6dc4e25 --- /dev/null +++ b/tests/agents/test_atif.py @@ -0,0 +1,733 @@ +"""Tests for the Python SDK's analyze_atif client.""" + +from __future__ import annotations + +import json +import shutil +from pathlib import Path + +import httpx +import pytest + +from pisama.agents import ( + AtifAnalyzeResult, + AtifDetection, + analyze_atif, + analyze_atif_batch, +) +from pisama.agents.atif import _discover_trajectory_files, _load_trajectory + +HARBOR_TERMINUS_FIXTURES = Path(__file__).parent / "fixtures" / "harbor" + + +@pytest.fixture(autouse=True) +def _clear_pisama_api_key(monkeypatch): + """Keep unit tests independent from a developer's shell environment.""" + monkeypatch.delenv("PISAMA_API_KEY", raising=False) + + +# A minimal but realistic ATIF v1.7 trajectory matching what the backend's +# vendored Pydantic models accept. +def _sample_trajectory() -> dict: + return { + "schema_version": "ATIF-v1.7", + "session_id": "test-sdk-001", + "agent": {"name": "claude-code", "version": "1.0", "model_name": "sonnet-4-6"}, + "steps": [ + {"step_id": 1, "source": "user", "message": "do thing"}, + { + "step_id": 2, + "source": "agent", + "message": "done", + "model_name": "sonnet-4-6", + }, + ], + } + + +def _fake_backend_response() -> dict: + return { + "diagnosis": { + "trace_id": "abc123", + "has_failures": True, + "failure_count": 2, + "detection_status": "complete", + "all_detections": [ + { + "detector": "hallucination", + "confidence": 0.91, + "severity": "high", + "title": "Unsupported claim", + "description": "Output claims X without source", + }, + { + "detector": "loop", + "confidence": 0.6, + "severity": "medium", + "title": "Repeated tool call", + "description": "Same call 3x", + }, + ], + "detectors_run": ["hallucination", "loop", "completion"], + "detectors_failed": {}, + }, + "trace": { + "trace_id": "abc123", + "source_format": "atif", + "span_count": 2, + "total_tokens": 0, + "total_duration_ms": 1500, + "topology_complete": True, + "unresolved_trajectory_refs": [], + "span_token_total": 0, + "atif_schema_version": "ATIF-v1.7", + "atif_session_id": "test-sdk-001", + "atif_trajectory_id": None, + }, + } + + +def _mock_transport(response_body: dict, status_code: int = 200) -> httpx.MockTransport: + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/v1/atif/analyze" + assert request.method == "POST" + body = json.loads(request.content) + assert "trajectory" in body + return httpx.Response(status_code, json=response_body) + + return httpx.MockTransport(handler) + + +def _install_transport(monkeypatch, transport: httpx.MockTransport) -> None: + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + + +def test_load_trajectory_accepts_dict(): + traj = _load_trajectory(_sample_trajectory()) + assert traj["schema_version"] == "ATIF-v1.7" + + +def test_load_trajectory_rejects_unknown_schema(): + bad = _sample_trajectory() | {"schema_version": "ATIF-v9.9"} + with pytest.raises(ValueError, match="Unsupported ATIF schema_version"): + _load_trajectory(bad) + + +def test_load_trajectory_reads_file(tmp_path: Path): + p = tmp_path / "trajectory.json" + p.write_text(json.dumps(_sample_trajectory())) + traj = _load_trajectory(p) + assert traj["session_id"] == "test-sdk-001" + + +def test_analyze_atif_parses_response(monkeypatch): + transport = _mock_transport(_fake_backend_response()) + # Patch httpx.Client to use the mock transport. + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + + result = analyze_atif(_sample_trajectory()) + assert isinstance(result, AtifAnalyzeResult) + assert result.trace_id == "abc123" + assert result.has_failures is True + assert result.failure_count == 2 + assert result.total_duration_ms == 1500 + assert result.atif_schema_version == "ATIF-v1.7" + assert result.atif_session_id == "test-sdk-001" + assert result.topology_complete is True + assert result.unresolved_trajectory_refs == [] + assert result.span_token_total == 0 + assert result.analysis_complete is True + assert len(result.failures) == 2 + assert isinstance(result.failures[0], AtifDetection) + assert result.failures[0].severity == "high" + assert result.detectors_run == ["hallucination", "loop", "completion"] + + +def test_analyze_result_marks_partial_topology_incomplete(): + response = _fake_backend_response() + response["diagnosis"]["has_failures"] = False + response["diagnosis"]["failure_count"] = 0 + response["diagnosis"]["all_detections"] = [] + response["trace"].update( + { + "total_tokens": 8832, + "span_token_total": 7192, + "topology_complete": False, + "unresolved_trajectory_refs": [ + "trajectory.summarization-1-summary.json", + "trajectory.summarization-1-questions.json", + ], + } + ) + + result = AtifAnalyzeResult.from_response(response) + + assert result.detection_status == "complete" + assert result.detectors_failed == {} + assert result.topology_complete is False + assert result.unresolved_trajectory_refs == [ + "trajectory.summarization-1-summary.json", + "trajectory.summarization-1-questions.json", + ] + assert result.total_tokens == 8832 + assert result.span_token_total == 7192 + assert result.analysis_complete is False + + +def test_analyze_result_defaults_old_server_topology_to_complete(): + response = _fake_backend_response() + response["trace"].pop("topology_complete") + response["trace"].pop("unresolved_trajectory_refs") + response["trace"].pop("span_token_total") + + result = AtifAnalyzeResult.from_response(response) + + assert result.topology_complete is True + assert result.unresolved_trajectory_refs == [] + assert result.span_token_total == result.total_tokens + assert result.analysis_complete is True + + +def test_analyze_atif_exchanges_explicit_api_key_for_jwt(monkeypatch): + paths: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + paths.append(request.url.path) + if request.url.path == "/api/v1/auth/token": + assert request.headers.get("Authorization") is None + assert json.loads(request.content) == { + "api_key": "psk_explicit", + "scope": "full", + } + return httpx.Response(200, json={"access_token": "jwt-explicit"}) + + assert request.url.path == "/api/v1/atif/analyze" + assert request.headers["Authorization"] == "Bearer jwt-explicit" + assert b"psk_explicit" not in request.content + return httpx.Response(200, json=_fake_backend_response()) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + result = analyze_atif(_sample_trajectory(), api_key="psk_explicit") + + assert result.trace_id == "abc123" + assert paths == ["/api/v1/auth/token", "/api/v1/atif/analyze"] + + +def test_analyze_atif_uses_api_key_environment_fallback(monkeypatch): + captured_key: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/api/v1/auth/token": + captured_key.append(json.loads(request.content)["api_key"]) + return httpx.Response(200, json={"access_token": "jwt-from-env"}) + + assert request.headers["Authorization"] == "Bearer jwt-from-env" + return httpx.Response(200, json=_fake_backend_response()) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + monkeypatch.setenv("PISAMA_API_KEY", "psk_from_env") + + analyze_atif(_sample_trajectory()) + + assert captured_key == ["psk_from_env"] + + +def test_analyze_atif_empty_key_allows_unauthenticated_local_endpoint(monkeypatch): + monkeypatch.setenv("PISAMA_API_KEY", "psk_ambient") + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/v1/atif/analyze" + assert request.headers.get("Authorization") is None + return httpx.Response(200, json=_fake_backend_response()) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + result = analyze_atif( + _sample_trajectory(), + api_url="http://localhost:8000", + api_key="", + ) + + assert result.trace_id == "abc123" + + +def test_analyze_atif_stops_when_token_exchange_fails(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/v1/auth/token" + return httpx.Response(401, json={"detail": "invalid API key"}) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + with pytest.raises(httpx.HTTPStatusError): + analyze_atif(_sample_trajectory(), api_key="psk_invalid") + + +def test_analyze_atif_rejects_token_response_without_access_token(monkeypatch): + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/v1/auth/token" + return httpx.Response(200, json={"token_type": "bearer"}) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + with pytest.raises(ValueError, match="valid access_token"): + analyze_atif(_sample_trajectory(), api_key="psk_invalid_response") + + +def test_analyze_atif_batch_handles_directory(monkeypatch, tmp_path: Path): + transport = _mock_transport(_fake_backend_response()) + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + + (tmp_path / "a.json").write_text(json.dumps(_sample_trajectory())) + (tmp_path / "b.json").write_text(json.dumps(_sample_trajectory())) + results = analyze_atif_batch(tmp_path) + assert len(results) == 2 + assert all(r.trace_id == "abc123" for r in results) + # source_path is set when given a file path. + assert results[0].source_path is not None + assert results[0].source_path.endswith("a.json") + + +def test_analyze_atif_batch_exchanges_key_once_and_reuses_jwt( + monkeypatch, tmp_path: Path +): + auth_calls = 0 + analyze_calls = 0 + + def handler(request: httpx.Request) -> httpx.Response: + nonlocal auth_calls, analyze_calls + if request.url.path == "/api/v1/auth/token": + auth_calls += 1 + assert json.loads(request.content) == { + "api_key": "psk_batch", + "scope": "full", + } + return httpx.Response(200, json={"access_token": "jwt-batch"}) + + analyze_calls += 1 + assert request.url.path == "/api/v1/atif/analyze" + assert request.headers["Authorization"] == "Bearer jwt-batch" + return httpx.Response(200, json=_fake_backend_response()) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + (tmp_path / "a.json").write_text(json.dumps(_sample_trajectory())) + (tmp_path / "b.json").write_text(json.dumps(_sample_trajectory())) + + results = analyze_atif_batch(tmp_path, api_key="psk_batch") + + assert len(results) == 2 + assert auth_calls == 1 + assert analyze_calls == 2 + + +def test_analyze_atif_batch_empty_directory_skips_authentication( + monkeypatch, tmp_path: Path +): + def handler(request: httpx.Request) -> httpx.Response: + pytest.fail(f"unexpected request to {request.url.path}") + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + assert analyze_atif_batch(tmp_path, api_key="psk_unused") == [] + + +def test_analyze_atif_batch_discovers_harbor_job_output(monkeypatch, tmp_path: Path): + """Harbor's job output layout is //agent/trajectory.json.""" + transport = _mock_transport(_fake_backend_response()) + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + + # A real Harbor job root has metadata JSON beside its trial directories. + # These files must not shadow the nested ATIF trajectories. + for name in ("config.json", "lock.json", "result.json"): + (tmp_path / name).write_text("{}") + + # Build the real //agent/trajectory.json layout. + for trial in ("trial-1", "trial-2"): + agent_dir = tmp_path / trial / "agent" + agent_dir.mkdir(parents=True) + (agent_dir / "trajectory.json").write_text(json.dumps(_sample_trajectory())) + + results = analyze_atif_batch(tmp_path) + assert len(results) == 2 + assert {Path(r.source_path).parent.parent.name for r in results} == { + "trial-1", + "trial-2", + } + + +def test_analyze_atif_batch_uses_single_trial_layout(monkeypatch, tmp_path: Path): + """If /agent/trajectory.json exists, just use that one file.""" + transport = _mock_transport(_fake_backend_response()) + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + + agent_dir = tmp_path / "agent" + agent_dir.mkdir() + (agent_dir / "trajectory.json").write_text(json.dumps(_sample_trajectory())) + # Decoy .json that should be ignored when the trial layout is detected. + (tmp_path / "result.json").write_text("{}") + + results = analyze_atif_batch(tmp_path) + assert len(results) == 1 + assert Path(results[0].source_path).name == "trajectory.json" + + +def _copy_real_harbor_continuation_layout(root: Path) -> Path: + """Recreate Harbor's agent directory names from committed golden files.""" + agent_dir = root / "agent" + agent_dir.mkdir(parents=True) + shutil.copy( + HARBOR_TERMINUS_FIXTURES + / "hello-world-context-summarization-linear-history.trajectory.json", + agent_dir / "trajectory.json", + ) + shutil.copy( + HARBOR_TERMINUS_FIXTURES + / "hello-world-context-summarization-linear-history.trajectory.cont-1.json", + agent_dir / "trajectory.cont-1.json", + ) + shutil.copy( + HARBOR_TERMINUS_FIXTURES + / "hello-world-context-summarization.trajectory.summarization-1-summary.json", + agent_dir / "trajectory.summarization-1-summary.json", + ) + return agent_dir + + +def test_harbor_trial_discovery_follows_real_continuation_and_skips_helper( + tmp_path: Path, +): + agent_dir = _copy_real_harbor_continuation_layout(tmp_path) + + discovered = _discover_trajectory_files(tmp_path) + + assert discovered == [ + (agent_dir / "trajectory.json").resolve(), + (agent_dir / "trajectory.cont-1.json").resolve(), + ] + + +def test_harbor_trial_discovery_rejects_missing_continuation(tmp_path: Path): + agent_dir = _copy_real_harbor_continuation_layout(tmp_path) + (agent_dir / "trajectory.cont-1.json").unlink() + + with pytest.raises(ValueError, match="Missing ATIF continued_trajectory_ref"): + _discover_trajectory_files(tmp_path) + + +def test_harbor_trial_discovery_rejects_continuation_cycle(tmp_path: Path): + agent_dir = _copy_real_harbor_continuation_layout(tmp_path) + continuation_path = agent_dir / "trajectory.cont-1.json" + continuation = json.loads(continuation_path.read_text()) + continuation["continued_trajectory_ref"] = "trajectory.json" + continuation_path.write_text(json.dumps(continuation)) + + with pytest.raises(ValueError, match="continued_trajectory_ref cycle"): + _discover_trajectory_files(tmp_path) + + +def test_harbor_trial_discovery_rejects_continuation_escape(tmp_path: Path): + agent_dir = _copy_real_harbor_continuation_layout(tmp_path) + trajectory_path = agent_dir / "trajectory.json" + trajectory = json.loads(trajectory_path.read_text()) + trajectory["continued_trajectory_ref"] = "../outside.json" + trajectory_path.write_text(json.dumps(trajectory)) + + with pytest.raises(ValueError, match="escapes agent directory"): + _discover_trajectory_files(tmp_path) + + +def test_harbor_trial_discovery_rejects_absolute_continuation_inside_agent( + tmp_path: Path, +): + agent_dir = _copy_real_harbor_continuation_layout(tmp_path) + trajectory_path = agent_dir / "trajectory.json" + trajectory = json.loads(trajectory_path.read_text()) + trajectory["continued_trajectory_ref"] = str( + (agent_dir / "trajectory.cont-1.json").resolve() + ) + trajectory_path.write_text(json.dumps(trajectory)) + + with pytest.raises(ValueError, match="continued_trajectory_ref must be relative"): + _discover_trajectory_files(tmp_path) + + +def test_harbor_trial_discovery_rejects_agent_symlink_outside_directory( + tmp_path: Path, +): + selected = tmp_path / "selected" + selected.mkdir() + external_agent = tmp_path / "external-agent" + external_agent.mkdir() + (external_agent / "trajectory.json").write_text(json.dumps(_sample_trajectory())) + (selected / "agent").symlink_to(external_agent, target_is_directory=True) + + with pytest.raises(ValueError, match="escapes selected directory"): + _discover_trajectory_files(selected) + + +def test_harbor_job_discovery_rejects_nested_agent_symlink_outside_directory( + tmp_path: Path, +): + selected = tmp_path / "selected" + trial_dir = selected / "trial-1" + trial_dir.mkdir(parents=True) + external_agent = tmp_path / "external-agent" + external_agent.mkdir() + (external_agent / "trajectory.json").write_text(json.dumps(_sample_trajectory())) + (trial_dir / "agent").symlink_to(external_agent, target_is_directory=True) + + with pytest.raises(ValueError, match="agent directory symlink escapes"): + _discover_trajectory_files(selected) + + +def test_harbor_trial_discovery_rejects_root_symlink_outside_directory( + tmp_path: Path, +): + selected = tmp_path / "selected" + agent_dir = selected / "agent" + agent_dir.mkdir(parents=True) + external_trajectory = tmp_path / "external-trajectory.json" + external_trajectory.write_text(json.dumps(_sample_trajectory())) + (agent_dir / "trajectory.json").symlink_to(external_trajectory) + + with pytest.raises(ValueError, match="escapes selected directory"): + _discover_trajectory_files(selected) + + +def test_analyze_atif_batch_single_file_follows_real_continuation( + monkeypatch, + tmp_path: Path, +): + agent_dir = _copy_real_harbor_continuation_layout(tmp_path) + submitted_kinds: list[str] = [] + + def handler(request: httpx.Request) -> httpx.Response: + trajectory = json.loads(request.content)["trajectory"] + extra = trajectory.get("agent", {}).get("extra", {}) + submitted_kinds.append( + "continuation" if extra.get("continuation_index") == 1 else "root" + ) + return httpx.Response(200, json=_fake_backend_response()) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + results = analyze_atif_batch(agent_dir / "trajectory.json") + + assert len(results) == 2 + assert submitted_kinds == ["root", "continuation"] + assert [Path(result.source_path).name for result in results] == [ + "trajectory.json", + "trajectory.cont-1.json", + ] + + +def test_batch_reconciles_submitted_continuation_without_mutating_api_data( + monkeypatch, + tmp_path: Path, +): + _copy_real_harbor_continuation_layout(tmp_path) + + def handler(request: httpx.Request) -> httpx.Response: + trajectory = json.loads(request.content)["trajectory"] + response = _fake_backend_response() + continued_ref = trajectory.get("continued_trajectory_ref") + if continued_ref: + response["trace"].update( + { + "topology_complete": False, + "unresolved_trajectory_refs": [continued_ref], + } + ) + return httpx.Response(200, json=response) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + root_result, continuation_result = analyze_atif_batch(tmp_path) + + assert root_result.topology_complete is False + assert root_result.unresolved_trajectory_refs == ["trajectory.cont-1.json"] + assert root_result.raw["trace"]["topology_complete"] is False + assert root_result.raw["trace"]["unresolved_trajectory_refs"] == [ + "trajectory.cont-1.json" + ] + assert root_result.client_resolved_trajectory_refs == ["trajectory.cont-1.json"] + assert root_result.remaining_unresolved_trajectory_refs == [] + assert root_result.reconciled_topology_complete is True + assert root_result.analysis_complete is True + assert continuation_result.client_resolved_trajectory_refs == [] + assert continuation_result.analysis_complete is True + + +def test_batch_reconciliation_keeps_other_refs_incomplete(monkeypatch, tmp_path: Path): + _copy_real_harbor_continuation_layout(tmp_path) + + def handler(request: httpx.Request) -> httpx.Response: + trajectory = json.loads(request.content)["trajectory"] + response = _fake_backend_response() + continued_ref = trajectory.get("continued_trajectory_ref") + if continued_ref: + response["trace"].update( + { + "topology_complete": False, + "unresolved_trajectory_refs": [ + continued_ref, + "external-subagent-trajectory", + ], + } + ) + return httpx.Response(200, json=response) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + root_result, _continuation_result = analyze_atif_batch(tmp_path) + + assert root_result.client_resolved_trajectory_refs == ["trajectory.cont-1.json"] + assert root_result.remaining_unresolved_trajectory_refs == [ + "external-subagent-trajectory" + ] + assert root_result.reconciled_topology_complete is False + assert root_result.analysis_complete is False + + +def test_batch_reconciliation_does_not_hide_same_string_subagent_ref( + monkeypatch, + tmp_path: Path, +): + agent_dir = _copy_real_harbor_continuation_layout(tmp_path) + root_path = agent_dir / "trajectory.json" + trajectory = json.loads(root_path.read_text()) + trajectory["steps"].append( + { + "step_id": 999, + "source": "system", + "message": "subagent ref collision", + "observation": { + "results": [ + { + "subagent_trajectory_ref": [ + {"trajectory_path": "trajectory.cont-1.json"} + ] + } + ] + }, + } + ) + root_path.write_text(json.dumps(trajectory)) + + def handler(request: httpx.Request) -> httpx.Response: + submitted = json.loads(request.content)["trajectory"] + response = _fake_backend_response() + continued_ref = submitted.get("continued_trajectory_ref") + if continued_ref: + response["trace"].update( + { + "topology_complete": False, + # The backend de-duplicates refs, so this value could mean + # the continuation, the subagent, or both. + "unresolved_trajectory_refs": [continued_ref], + } + ) + return httpx.Response(200, json=response) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + root_result, _continuation_result = analyze_atif_batch(tmp_path) + + assert root_result.client_resolved_trajectory_refs == [] + assert root_result.remaining_unresolved_trajectory_refs == [ + "trajectory.cont-1.json" + ] + assert root_result.analysis_complete is False + + +def test_batch_reconciliation_keeps_detector_failure_incomplete( + monkeypatch, + tmp_path: Path, +): + _copy_real_harbor_continuation_layout(tmp_path) + + def handler(request: httpx.Request) -> httpx.Response: + trajectory = json.loads(request.content)["trajectory"] + response = _fake_backend_response() + continued_ref = trajectory.get("continued_trajectory_ref") + if continued_ref: + response["trace"].update( + { + "topology_complete": False, + "unresolved_trajectory_refs": [continued_ref], + } + ) + response["diagnosis"]["detectors_failed"] = { + "hallucination": "detector timed out" + } + return httpx.Response(200, json=response) + + _install_transport(monkeypatch, httpx.MockTransport(handler)) + + root_result, _continuation_result = analyze_atif_batch(tmp_path) + + assert root_result.reconciled_topology_complete is True + assert root_result.detectors_failed == {"hallucination": "detector timed out"} + assert root_result.analysis_complete is False + + +def test_analyze_atif_raises_on_non_2xx(monkeypatch): + transport = _mock_transport({"detail": "bad"}, status_code=422) + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + with pytest.raises(httpx.HTTPStatusError): + analyze_atif(_sample_trajectory()) + + +def test_analyze_atif_passes_project_id_when_set(monkeypatch): + captured: dict = {} + + def handler(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + return httpx.Response(200, json=_fake_backend_response()) + + transport = httpx.MockTransport(handler) + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + analyze_atif(_sample_trajectory(), project_id="ps_test_123") + assert captured["body"]["project_id"] == "ps_test_123" diff --git a/tests/agents/test_bridge.py b/tests/agents/test_bridge.py new file mode 100644 index 0000000..4628416 --- /dev/null +++ b/tests/agents/test_bridge.py @@ -0,0 +1,382 @@ +"""Tests for DetectionBridge and hooks.""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from pisama.agents import ( + BridgeConfig, + BridgeResult, + DetectionBridge, + configure_bridge, + create_bridge, +) +from pisama.agents.converter import HookInputConverter +from pisama.agents.hooks.matchers import HookMatcher, create_matcher +from pisama.agents.session import SessionManager + + +class TestBridgeConfig: + """Tests for BridgeConfig.""" + + def test_default_config(self): + """Should have sensible defaults.""" + config = BridgeConfig() + assert config.warning_threshold == 40 + assert config.block_threshold == 60 + assert config.detection_timeout_ms == 80 + assert config.enable_blocking is True + assert config.fail_open is True + + def test_custom_config(self): + """Should accept custom values.""" + config = BridgeConfig( + warning_threshold=30, + block_threshold=50, + detection_timeout_ms=60, + ) + assert config.warning_threshold == 30 + assert config.block_threshold == 50 + assert config.detection_timeout_ms == 60 + + def test_config_to_dict(self): + """Should convert to dictionary.""" + config = BridgeConfig() + d = config.to_dict() + assert "warning_threshold" in d + assert "block_threshold" in d + assert d["warning_threshold"] == 40 + + +class TestBridgeResult: + """Tests for BridgeResult.""" + + def test_default_result(self): + """Should have sensible defaults.""" + result = BridgeResult() + assert result.should_block is False + assert result.severity == 0 + assert result.issues == [] + assert result.timed_out is False + + def test_to_hook_output_no_block(self): + """Should return empty dict when not blocking.""" + result = BridgeResult(severity=30) + output = result.to_hook_output() + assert output == {} + + def test_to_hook_output_block(self): + """Should return block decision when blocking.""" + result = BridgeResult( + should_block=True, + block_reason="Loop detected", + ) + output = result.to_hook_output() + assert "hookSpecificOutput" in output + assert output["hookSpecificOutput"]["permissionDecision"] == "block" + + def test_to_hook_output_with_message(self): + """Should include system message when present.""" + result = BridgeResult( + system_message="Warning: pattern detected", + ) + output = result.to_hook_output() + assert output["systemMessage"] == "Warning: pattern detected" + + +class TestSessionManager: + """Tests for SessionManager.""" + + def test_get_or_create(self): + """Should create new session.""" + manager = SessionManager() + session = manager.get_or_create("test-session") + assert session.session_id == "test-session" + assert manager.session_count == 1 + + def test_add_span(self): + """Should add span to session.""" + manager = SessionManager() + mock_span = MagicMock() + mock_span.name = "Bash" + + manager.add_span("test-session", mock_span) + + session = manager.get_or_create("test-session") + assert len(session.recent_spans) == 1 + + def test_get_context(self): + """Should return context with recent spans.""" + manager = SessionManager() + mock_span = MagicMock() + mock_span.name = "Bash" + + manager.add_span("test-session", mock_span) + context = manager.get_context("test-session", window=10) + + assert "recent_spans" in context + assert "tool_counts" in context + assert context["tool_counts"]["Bash"] == 1 + + def test_block_session(self): + """Should block session.""" + manager = SessionManager() + manager.block("test-session", "Too many loops") + + assert manager.is_blocked("test-session") + assert manager.get_block_reason("test-session") == "Too many loops" + + def test_unblock_session(self): + """Should unblock session.""" + manager = SessionManager() + manager.block("test-session", "reason") + manager.unblock("test-session") + + assert not manager.is_blocked("test-session") + + +class TestHookInputConverter: + """Tests for HookInputConverter.""" + + def test_convert_pre_tool(self): + """Should convert PreToolUse input to span.""" + converter = HookInputConverter() + hook_input = { + "tool_name": "Bash", + "tool_input": {"command": "ls -la"}, + "session_id": "test-session", + } + + span = converter.to_span(hook_input, "tool-123", is_post=False) + + assert span.name == "Bash" + assert span.span_id == "tool-123" + assert span.input_data == {"command": "ls -la"} + assert span.end_time is None # PreToolUse hasn't completed + + def test_convert_post_tool(self): + """Should convert PostToolUse input with response.""" + converter = HookInputConverter() + hook_input = { + "tool_name": "Read", + "tool_input": {"file_path": "/tmp/test.txt"}, + "tool_response": {"content": "file contents"}, + "session_id": "test-session", + } + + span = converter.to_span(hook_input, "tool-456", is_post=True) + + assert span.name == "Read" + assert span.output_data == {"content": "file contents"} + assert span.end_time is not None # PostToolUse is complete + + +class TestHookMatcher: + """Tests for HookMatcher.""" + + def test_match_tool_name(self): + """Should match by tool name pattern.""" + matcher = HookMatcher(tool_name_pattern="^Bash$") + + assert matcher.matches("Bash") + assert not matcher.matches("Read") + + def test_match_multiple_tools(self): + """Should match multiple tool patterns.""" + matcher = HookMatcher(tool_name_pattern="^(Bash|Read|Write)$") + + assert matcher.matches("Bash") + assert matcher.matches("Read") + assert matcher.matches("Write") + assert not matcher.matches("Glob") + + def test_match_with_input_pattern(self): + """Should match input patterns.""" + matcher = HookMatcher( + tool_name_pattern="^Bash$", + input_pattern=r"rm\s+-rf", + ) + + assert matcher.matches("Bash", {"command": "rm -rf /tmp"}) + assert not matcher.matches("Bash", {"command": "ls -la"}) + + def test_exclude_tools(self): + """Should exclude specific tools.""" + matcher = HookMatcher( + tool_name_pattern=".*", + exclude_tools=["AskUserQuestion"], + ) + + assert matcher.matches("Bash") + assert not matcher.matches("AskUserQuestion") + + def test_create_matcher(self): + """Should create matcher from tool list.""" + matcher = create_matcher( + tools=["Bash", "Read"], + exclude=["AskUserQuestion"], + ) + + assert matcher.matches("Bash") + assert matcher.matches("Read") + assert not matcher.matches("Write") + + +class TestDetectionBridge: + """Tests for DetectionBridge.""" + + @pytest.fixture + def mock_orchestrator(self): + """Create mock orchestrator.""" + with patch( + "pisama.agents.bridge.DetectionOrchestrator" + ) as mock_class: + mock_instance = MagicMock() + mock_instance.analyze_realtime = AsyncMock() + mock_class.return_value = mock_instance + yield mock_instance + + @pytest.fixture + def bridge(self, mock_orchestrator): + """Create bridge with mocked orchestrator.""" + config = BridgeConfig( + tool_patterns=[".*"], + excluded_tools=[], + ) + return DetectionBridge(config=config) + + @pytest.mark.asyncio + async def test_analyze_pre_tool_no_detection(self, bridge, mock_orchestrator): + """Should allow tool when no issues detected.""" + mock_orchestrator.analyze_realtime.return_value = MagicMock( + should_block=False, + severity=0, + issues=[], + recommendations=[], + block_reason=None, + ) + + hook_input = { + "tool_name": "Read", + "tool_input": {"file_path": "/tmp/test.txt"}, + "session_id": "test-session", + } + + result = await bridge.analyze_pre_tool(hook_input, "tool-123") + + assert result.should_block is False + assert result.severity == 0 + + @pytest.mark.asyncio + async def test_analyze_pre_tool_with_warning(self, bridge, mock_orchestrator): + """Should return warning when severity above threshold.""" + mock_orchestrator.analyze_realtime.return_value = MagicMock( + should_block=False, + severity=45, + issues=["Repetitive pattern detected"], + recommendations=[], + block_reason=None, + ) + + hook_input = { + "tool_name": "Bash", + "tool_input": {"command": "ls"}, + "session_id": "test-session", + } + + result = await bridge.analyze_pre_tool(hook_input, "tool-456") + + assert result.should_block is False + assert result.severity == 45 + assert result.system_message is not None + assert "Warning" in result.system_message + + @pytest.mark.asyncio + async def test_analyze_pre_tool_with_block(self, bridge, mock_orchestrator): + """Should block when severity above block threshold.""" + mock_orchestrator.analyze_realtime.return_value = MagicMock( + should_block=True, + severity=75, + issues=["Loop detected"], + recommendations=[], + block_reason="Repetitive loop pattern", + ) + + hook_input = { + "tool_name": "Bash", + "tool_input": {"command": "ls"}, + "session_id": "test-session", + } + + result = await bridge.analyze_pre_tool(hook_input, "tool-789") + + assert result.should_block is True + assert result.severity == 75 + assert "BLOCKED" in result.system_message + + @pytest.mark.asyncio + async def test_analyze_pre_tool_excluded_tool(self, bridge, mock_orchestrator): + """Should skip excluded tools.""" + bridge.config.excluded_tools = ["AskUserQuestion"] + bridge._exclude_patterns = [__import__("re").compile("^AskUserQuestion$")] + + hook_input = { + "tool_name": "AskUserQuestion", + "tool_input": {}, + "session_id": "test-session", + } + + result = await bridge.analyze_pre_tool(hook_input, "tool-000") + + # Should not call orchestrator + mock_orchestrator.analyze_realtime.assert_not_called() + assert result.severity == 0 + + @pytest.mark.asyncio + async def test_analyze_post_tool(self, bridge, mock_orchestrator): + """Should analyze post-tool and inject recovery message.""" + mock_orchestrator.analyze_realtime.return_value = MagicMock( + should_block=False, + severity=50, + issues=["Pattern detected"], + recommendations=[{"fix_instruction": "Try a different approach"}], + block_reason=None, + ) + + hook_input = { + "tool_name": "Bash", + "tool_input": {"command": "ls"}, + "tool_response": {"output": "file.txt"}, + "session_id": "test-session", + } + + result = await bridge.analyze_post_tool(hook_input, "tool-111") + + assert result.should_block is False # PostToolUse never blocks + assert result.system_message is not None + assert "Recovery" in result.system_message + + +class TestConfigureBridge: + """Tests for bridge configuration functions.""" + + def test_configure_bridge(self): + """Should configure global bridge.""" + bridge = configure_bridge( + warning_threshold=30, + block_threshold=50, + timeout_ms=60, + ) + + assert bridge.config.warning_threshold == 30 + assert bridge.config.block_threshold == 50 + assert bridge.config.detection_timeout_ms == 60 + + def test_create_bridge(self): + """Should create independent bridge.""" + bridge1 = create_bridge(warning_threshold=30) + bridge2 = create_bridge(warning_threshold=50) + + assert bridge1.config.warning_threshold == 30 + assert bridge2.config.warning_threshold == 50 + assert bridge1 is not bridge2 diff --git a/tests/agents/test_chaos_hooks.py b/tests/agents/test_chaos_hooks.py new file mode 100644 index 0000000..0931e4f --- /dev/null +++ b/tests/agents/test_chaos_hooks.py @@ -0,0 +1,168 @@ +"""Tests for SDK-level chaos hooks.""" + +import sys +from pathlib import Path + +# Add src to path so we can import chaos module without full SDK deps. +# This file lives at tests/agents/, two levels below the repo root (the +# original pisama-agent-sdk had tests/ directly at repo root). +sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src")) + +import importlib.util + +# Load chaos modules directly from file paths to avoid triggering +# pisama.agents's __init__.py (which requires pisama_core). +_repo_root = Path(__file__).parent.parent.parent +_src = _repo_root / "src" / "pisama" / "agents" / "chaos" + + +def _load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +_config_mod = _load("chaos_config", _src / "config.py") +_exp_mod = _load("chaos_experiments", _src / "experiments.py") + +ChaosConfig = _config_mod.ChaosConfig +ChaosExperiment = _exp_mod.ChaosExperiment +ChaosResult = _exp_mod.ChaosResult +ToolFailure = _exp_mod.ToolFailure +LatencyInjection = _exp_mod.LatencyInjection +ErrorInjection = _exp_mod.ErrorInjection +OutputCorruption = _exp_mod.OutputCorruption +ContextTruncation = _exp_mod.ContextTruncation + + +class TestToolFailure: + def test_blocks_targeted_tool(self): + exp = ToolFailure(tools=["search"], probability=1.0) + result = exp.apply_pre("search", {"query": "test"}) + assert result.applied + assert result.block + assert "search" in result.message + + def test_does_not_block_non_targeted_tool(self): + exp = ToolFailure(tools=["search"], probability=1.0) + assert not exp.matches("database_query") + + def test_matches_all_when_no_tools_specified(self): + exp = ToolFailure(probability=1.0) + assert exp.matches("anything") + assert exp.matches("search") + + def test_matches_targeted_agent(self): + exp = ToolFailure(agents=["researcher"], probability=1.0) + assert exp.matches("search", agent_name="researcher") + assert not exp.matches("search", agent_name="planner") + + +class TestLatencyInjection: + def test_returns_delay(self): + exp = LatencyInjection(min_ms=100, max_ms=200, probability=1.0) + result = exp.apply_pre("search", {}) + assert result.applied + assert 100 <= result.delay_ms <= 200 + assert not result.block + + def test_no_block(self): + exp = LatencyInjection(min_ms=500, max_ms=500, probability=1.0) + result = exp.apply_pre("search", {}) + assert result.delay_ms == 500 + assert not result.block + + +class TestErrorInjection: + def test_blocks_with_error(self): + exp = ErrorInjection(error_code=503, probability=1.0) + result = exp.apply_pre("search", {}) + assert result.applied + assert result.block + assert "503" in result.message + + +class TestOutputCorruption: + def test_truncates_output(self): + exp = OutputCorruption(corruption="truncate", probability=1.0) + original = "This is a long search result with lots of content" + result = exp.apply_post("search", original) + assert result.applied + assert result.modified_output is not None + assert len(result.modified_output) < len(original) + + def test_empties_output(self): + exp = OutputCorruption(corruption="empty", probability=1.0) + result = exp.apply_post("search", "Some output") + assert result.modified_output == "" + + def test_json_break(self): + exp = OutputCorruption(corruption="json_break", probability=1.0) + result = exp.apply_post("search", '{"data": "valid"}') + assert '{"incomplete": true' in result.modified_output + + +class TestContextTruncation: + def test_truncates_long_strings(self): + exp = ContextTruncation(truncation_pct=0.5, probability=1.0) + long_input = {"query": "x" * 200, "short_field": "hi"} + result = exp.apply_pre("search", long_input) + assert result.applied + assert len(result.modified_input["query"]) == 100 # 50% of 200 + assert result.modified_input["short_field"] == "hi" # short strings unchanged + + def test_preserves_non_strings(self): + exp = ContextTruncation(truncation_pct=0.5, probability=1.0) + result = exp.apply_pre("search", {"count": 42, "flag": True}) + assert result.modified_input["count"] == 42 + assert result.modified_input["flag"] is True + + +class TestChaosConfig: + def test_is_active_when_enabled(self): + config = ChaosConfig(experiments=[ToolFailure()], enabled=True) + assert config.is_active + + def test_not_active_when_disabled(self): + config = ChaosConfig(experiments=[ToolFailure()], enabled=False) + assert not config.is_active + + def test_safety_limit_disables(self): + config = ChaosConfig(experiments=[ToolFailure()], safety_max_affected=3) + assert config.is_active + config.record_affected() + config.record_affected() + config.record_affected() + assert not config.is_active # 3 >= safety_max_affected + + def test_reset_re_enables(self): + config = ChaosConfig(experiments=[ToolFailure()], safety_max_affected=1) + config.record_affected() + assert not config.is_active + config.reset() + assert config.is_active + + +class TestProbability: + def test_zero_probability_never_triggers(self): + exp = ToolFailure(probability=0.0) + triggered = sum(1 for _ in range(100) if exp.should_trigger()) + assert triggered == 0 + + def test_one_probability_always_triggers(self): + exp = ToolFailure(probability=1.0) + triggered = sum(1 for _ in range(100) if exp.should_trigger()) + assert triggered == 100 + + +class TestChaosExperimentBase: + def test_default_apply_pre_is_noop(self): + exp = ChaosExperiment() + result = exp.apply_pre("tool", {}) + assert not result.applied + + def test_default_apply_post_is_noop(self): + exp = ChaosExperiment() + result = exp.apply_post("tool", "output") + assert not result.applied diff --git a/tests/agents/test_check_compliance.py b/tests/agents/test_check_compliance.py new file mode 100644 index 0000000..8f56ca4 --- /dev/null +++ b/tests/agents/test_check_compliance.py @@ -0,0 +1,210 @@ +"""Tests for the gated check_compliance SDK helper.""" + +from __future__ import annotations + +# Tests need the module object, not the function. ``from pisama.agents +# import check_compliance`` above imports the function, so we re-resolve +# the module via importlib for unambiguous patching. +import importlib +import io +import json +from unittest.mock import patch + +import pytest + +from pisama.agents import ( + BehavioralRule, + ComplianceResult, + PisamaFeatureNotEnabledError, + Violation, + check_compliance, +) + +check_compliance_module = importlib.import_module( + "pisama.agents.check_compliance" +) + + +_FLAG = "PISAMA_ENABLE_CHECK_COMPLIANCE" + + +# --------------------------------------------------------------------------- +# Feature-flag gating +# --------------------------------------------------------------------------- + + +class TestFeatureFlag: + """check_compliance must refuse to run unless the flag is set truthy.""" + + @pytest.mark.asyncio + async def test_flag_unset_raises(self, monkeypatch): + monkeypatch.delenv(_FLAG, raising=False) + with pytest.raises(PisamaFeatureNotEnabledError) as exc: + await check_compliance("system prompt", []) + # Error message should tell the caller exactly how to enable it. + assert _FLAG in str(exc.value) + + @pytest.mark.asyncio + async def test_flag_empty_raises(self, monkeypatch): + monkeypatch.setenv(_FLAG, "") + with pytest.raises(PisamaFeatureNotEnabledError): + await check_compliance("system prompt", []) + + @pytest.mark.asyncio + async def test_flag_zero_raises(self, monkeypatch): + monkeypatch.setenv(_FLAG, "0") + with pytest.raises(PisamaFeatureNotEnabledError): + await check_compliance("system prompt", []) + + @pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "Yes"]) + def test_flag_truthy_values_enabled(self, monkeypatch, value): + monkeypatch.setenv(_FLAG, value) + assert check_compliance_module._is_compliance_enabled() is True + + @pytest.mark.parametrize("value", ["", "0", "no", "false", "off"]) + def test_flag_falsy_values_disabled(self, monkeypatch, value): + monkeypatch.setenv(_FLAG, value) + assert check_compliance_module._is_compliance_enabled() is False + + +# --------------------------------------------------------------------------- +# Bridge / HTTP path with the flag enabled +# --------------------------------------------------------------------------- + + +def _fake_http_response(payload: dict): + """Return a context-manager-shaped object that urlopen() expects.""" + return io.BytesIO(json.dumps(payload).encode()) + + +class TestCheckComplianceCall: + """When the flag is set, check_compliance hits the backend endpoint.""" + + @pytest.mark.asyncio + async def test_passes_through_to_backend(self, monkeypatch): + monkeypatch.setenv(_FLAG, "1") + + captured = {} + + def fake_urlopen(request, timeout=None): + captured["url"] = request.full_url + captured["body"] = request.data + captured["timeout"] = timeout + return _fake_http_response({ + "detected": True, + "confidence": 0.9, + "violations": [ + { + "rule_id": "no_drop_table", + "evidence": "tool_call: drop_table()", + "explanation": "Trace contains forbidden action.", + "confidence": 0.9, + } + ], + "extracted_rules": [ + { + "rule_id": "no_drop_table", + "description": "Never call drop_table.", + "trigger": "always", + "required_action": None, + "forbidden_action": "call_tool: drop_table", + "severity": "critical", + } + ], + "tokens_used": 1234, + "cost_usd": 0.0042, + }) + + with patch.object(check_compliance_module, "urlopen", fake_urlopen): + result = await check_compliance( + system_prompt="You are a careful assistant. Never call drop_table.", + trace_events=[{"type": "tool_call", "name": "drop_table", "args": {}}], + ) + + # URL is correct + assert captured["url"].endswith("/api/v1/evaluate/detect/specification-compliance") + # Body shape matches the analyzer's signature + body = json.loads(captured["body"].decode()) + assert body["system_prompt"].startswith("You are a careful assistant") + assert body["trace_events"] == [ + {"type": "tool_call", "name": "drop_table", "args": {}} + ] + + # Result deserialization + assert isinstance(result, ComplianceResult) + assert result.detected is True + assert result.confidence == pytest.approx(0.9) + assert result.tokens_used == 1234 + assert result.cost_usd == pytest.approx(0.0042) + + assert len(result.violations) == 1 + v = result.violations[0] + assert isinstance(v, Violation) + assert v.rule_id == "no_drop_table" + assert v.evidence == "tool_call: drop_table()" + + assert len(result.extracted_rules) == 1 + r = result.extracted_rules[0] + assert isinstance(r, BehavioralRule) + assert r.severity == "critical" + assert r.forbidden_action == "call_tool: drop_table" + + @pytest.mark.asyncio + async def test_handles_no_violations(self, monkeypatch): + monkeypatch.setenv(_FLAG, "1") + + def fake_urlopen(request, timeout=None): + return _fake_http_response({ + "detected": False, + "confidence": 0.5, + "violations": [], + "extracted_rules": [], + "tokens_used": 200, + "cost_usd": 0.0001, + }) + + with patch.object(check_compliance_module, "urlopen", fake_urlopen): + result = await check_compliance("be careful", []) + + assert result.detected is False + assert result.violations == [] + assert result.extracted_rules == [] + + @pytest.mark.asyncio + async def test_malformed_json_returns_empty_result(self, monkeypatch): + """Bad JSON should not crash the agent — we return a safe default.""" + monkeypatch.setenv(_FLAG, "1") + + def fake_urlopen(request, timeout=None): + return io.BytesIO(b"not json at all") + + with patch.object(check_compliance_module, "urlopen", fake_urlopen): + result = await check_compliance("sp", []) + + assert isinstance(result, ComplianceResult) + assert result.detected is False + assert result.confidence == 0.0 + + @pytest.mark.asyncio + async def test_uses_pisama_api_url_env(self, monkeypatch): + monkeypatch.setenv(_FLAG, "1") + monkeypatch.setenv("PISAMA_API_URL", "https://api.example.test") + + # Reset any previously configured _api_url on the check module + from pisama.agents import check as _check_module + monkeypatch.setattr(_check_module, "_api_url", None, raising=False) + + captured = {} + + def fake_urlopen(request, timeout=None): + captured["url"] = request.full_url + return _fake_http_response({ + "detected": False, "confidence": 0.0, + "violations": [], "extracted_rules": [], + "tokens_used": 0, "cost_usd": 0.0, + }) + + with patch.object(check_compliance_module, "urlopen", fake_urlopen): + await check_compliance("sp", []) + + assert captured["url"] == "https://api.example.test/api/v1/evaluate/detect/specification-compliance" diff --git a/tests/agents/test_heal.py b/tests/agents/test_heal.py new file mode 100644 index 0000000..e1f28ac --- /dev/null +++ b/tests/agents/test_heal.py @@ -0,0 +1,45 @@ +"""Tests for HealingResult convenience accessors (no network).""" + +from pisama.agents import HealingResult + + +class TestIdePatchProperty: + """HealingResult.ide_patch surfaces the server-rendered IDE patch.""" + + def test_returns_patch_dict_when_present(self): + patch = { + "target_files": ["CLAUDE.md", ".cursorrules"], + "apply_mode": "suggested", + "framework": None, + "instructions": "## Add retry limit\n...", + "verification": "pytest -q\n", + } + result = HealingResult( + applied=True, + escalated=False, + risk_level="safe", + fix={"id": "fix_1", "ide_patch": patch}, + ) + assert result.ide_patch == patch + + def test_none_when_no_fix(self): + result = HealingResult(applied=False, escalated=False, risk_level="unknown") + assert result.ide_patch is None + + def test_none_when_fix_lacks_patch(self): + result = HealingResult( + applied=True, + escalated=False, + risk_level="safe", + fix={"id": "fix_1"}, + ) + assert result.ide_patch is None + + def test_none_when_patch_not_a_dict(self): + result = HealingResult( + applied=True, + escalated=False, + risk_level="safe", + fix={"id": "fix_1", "ide_patch": "oops"}, + ) + assert result.ide_patch is None diff --git a/tests/agents/test_openhands_adapter.py b/tests/agents/test_openhands_adapter.py new file mode 100644 index 0000000..0fba1a0 --- /dev/null +++ b/tests/agents/test_openhands_adapter.py @@ -0,0 +1,356 @@ +"""Tests for OpenHandsEventStreamAdapter (Phase C). + +Covers: +- batch-mode analysis on a Harbor-emitted session_dir (real fixture) +- batch-mode analysis on synthesized trajectories from buffered events +- streaming-mode destructive_command detection on the action stream +- CLI smoke test on a Harbor-real fixture +""" +from __future__ import annotations + +import json +import shutil +from pathlib import Path + +import httpx +import pytest + +from pisama.agents import ( + AtifAnalyzeResult, + OpenHandsEventStreamAdapter, + StreamingDetection, +) +from pisama.agents.cli.openhands_monitor import main as cli_main + + +@pytest.fixture(autouse=True) +def _clear_pisama_api_key(monkeypatch): + """Keep HTTP request assertions independent from the shell environment.""" + monkeypatch.delenv("PISAMA_API_KEY", raising=False) + + +# Captured Harbor trajectory shipped with this standalone package. Keeping the +# fixture local makes the published repository's test suite self-contained. +HARBOR_TRAJECTORY = ( + Path(__file__).parent + / "fixtures" + / "harbor" + / "hello-world-context-summarization-linear-history.trajectory.json" +) + + +def _fake_clean_response() -> dict: + return { + "diagnosis": { + "trace_id": "oh-test-001", + "has_failures": False, + "failure_count": 0, + "detection_status": "complete", + "all_detections": [], + "detectors_run": ["loop", "completion"], + "detectors_failed": {}, + }, + "trace": { + "trace_id": "oh-test-001", + "source_format": "atif", + "span_count": 4, + "total_tokens": 200, + "atif_schema_version": "ATIF-v1.7", + "atif_session_id": None, + "atif_trajectory_id": None, + }, + } + + +def _fake_failure_response() -> dict: + return { + "diagnosis": { + "trace_id": "oh-test-002", + "has_failures": True, + "failure_count": 1, + "detection_status": "complete", + "all_detections": [{ + "detector": "runtime_error", + "confidence": 0.92, + "severity": "high", + "title": "Runtime / Environment Error", + "description": "Agent terminated with a model_unknown meta-error", + }], + "detectors_run": ["runtime_error"], + "detectors_failed": {}, + }, + "trace": { + "trace_id": "oh-test-002", + "source_format": "atif", + "span_count": 2, + "total_tokens": 50, + "atif_schema_version": "ATIF-v1.7", + "atif_session_id": None, + "atif_trajectory_id": None, + }, + } + + +def _mock_transport(response_body: dict, status_code: int = 200) -> httpx.MockTransport: + captured: dict[str, object] = {} + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/v1/atif/analyze" + assert request.method == "POST" + body = json.loads(request.content) + captured["payload"] = body + return httpx.Response(status_code, json=response_body) + + transport = httpx.MockTransport(handler) + transport.captured = captured # type: ignore[attr-defined] + return transport + + +@pytest.fixture +def patch_httpx_client(monkeypatch): + """Factory that swaps httpx.Client for one with a MockTransport.""" + + def _apply(transport: httpx.MockTransport): + real_client = httpx.Client + + def patched_client(*args, **kwargs): + kwargs["transport"] = transport + return real_client(*args, **kwargs) + + monkeypatch.setattr(httpx, "Client", patched_client) + return transport + + return _apply + + +# --------------------------------------------------------------------------- +# Constructor + mode validation +# --------------------------------------------------------------------------- + + +def test_adapter_defaults_to_batch_mode(): + a = OpenHandsEventStreamAdapter() + assert a.mode == "batch" + assert a.event_count() == 0 + + +def test_adapter_rejects_unknown_mode(): + with pytest.raises(ValueError, match="mode must be"): + OpenHandsEventStreamAdapter(mode="not-a-mode") # type: ignore[arg-type] + + +def test_event_buffer_grows_and_resets(): + a = OpenHandsEventStreamAdapter() + a.on_action({"action": "run", "args": {"command": "ls"}}) + a.on_observation({"observation": "run", "content": "ok"}) + assert a.event_count() == 2 + a.reset() + assert a.event_count() == 0 + + +# --------------------------------------------------------------------------- +# Streaming mode — destructive_command +# --------------------------------------------------------------------------- + + +def test_streaming_destructive_command_fires_callback(): + captured: list[StreamingDetection] = [] + adapter = OpenHandsEventStreamAdapter( + mode="streaming", + on_streaming_detection=captured.append, + ) + result = adapter.on_action({"action": "run", "args": {"command": "rm -rf /"}}) + assert result is not None + assert result.detector == "destructive_command" + assert result.pattern_name == "rm_rf_root" + assert result.severity == "critical" + assert result.confidence >= 0.99 + assert captured == [result] + + +def test_streaming_safe_command_no_detection(): + adapter = OpenHandsEventStreamAdapter(mode="streaming") + result = adapter.on_action({"action": "run", "args": {"command": "rm -rf ./node_modules"}}) + assert result is None + + +def test_streaming_non_shell_action_does_not_match(): + adapter = OpenHandsEventStreamAdapter(mode="streaming") + result = adapter.on_action({ + "action": "browse", + "args": {"url": "https://example.com"}, + }) + assert result is None + + +def test_batch_mode_does_not_fire_streaming_detection(): + """In batch mode, on_action buffers the event but does not invoke + the streaming-detection path even on a destructive command.""" + captured: list[StreamingDetection] = [] + adapter = OpenHandsEventStreamAdapter( + mode="batch", + on_streaming_detection=captured.append, + ) + result = adapter.on_action({"action": "run", "args": {"command": "rm -rf /"}}) + assert result is None + assert captured == [] + assert adapter.event_count() == 1 + + +def test_streaming_callback_exception_does_not_propagate(): + """A user callback that raises must not crash the adapter (the + underlying agent stream must keep flowing).""" + def boom(_d): + raise RuntimeError("user callback exploded") + + adapter = OpenHandsEventStreamAdapter( + mode="streaming", + on_streaming_detection=boom, + ) + # Should not raise. + result = adapter.on_action({"action": "run", "args": {"command": "rm -rf /"}}) + assert result is not None + + +# --------------------------------------------------------------------------- +# Batch mode — analyze on session_dir +# --------------------------------------------------------------------------- + + +def test_on_session_complete_uses_session_trajectory_json( + patch_httpx_client, tmp_path: Path, +): + """Given a session_dir containing agent/trajectory.json, the adapter + POSTs that trajectory and returns the AtifAnalyzeResult.""" + # Copy a real harbor-real fixture into the temp session_dir layout. + (tmp_path / "agent").mkdir() + shutil.copy(HARBOR_TRAJECTORY, tmp_path / "agent" / "trajectory.json") + + transport = _mock_transport(_fake_clean_response()) + patch_httpx_client(transport) + + adapter = OpenHandsEventStreamAdapter() + result = adapter.on_session_complete(tmp_path) + assert isinstance(result, AtifAnalyzeResult) + assert not result.has_failures + # The POSTed payload should include the trajectory dict. + assert "trajectory" in transport.captured["payload"] + + +def test_on_session_complete_uses_session_trajectory_json_at_root( + patch_httpx_client, tmp_path: Path, +): + """Falls back to trajectory.json at the session root when there is + no agent/ subdir.""" + shutil.copy(HARBOR_TRAJECTORY, tmp_path / "trajectory.json") + + transport = _mock_transport(_fake_clean_response()) + patch_httpx_client(transport) + + adapter = OpenHandsEventStreamAdapter() + result = adapter.on_session_complete(tmp_path) + assert result.has_failures is False + + +def test_on_session_complete_synthesizes_from_buffer( + patch_httpx_client, tmp_path: Path, +): + """No session_dir + buffered events → adapter synthesizes a trajectory + from the events and POSTs it.""" + adapter = OpenHandsEventStreamAdapter() + adapter.on_action({ + "action": "message", + "source": "user", + "content": "Please list the files in /tmp", + }) + adapter.on_action({"action": "run", "args": {"command": "ls /tmp", "thought": "Listing"}}) + adapter.on_observation({"observation": "run", "content": "file1\nfile2"}) + adapter.on_action({ + "action": "message", + "source": "agent", + "content": "Found two files: file1 and file2.", + }) + + transport = _mock_transport(_fake_clean_response()) + patch_httpx_client(transport) + + result = adapter.on_session_complete() + assert isinstance(result, AtifAnalyzeResult) + payload = transport.captured["payload"]["trajectory"] + assert payload["schema_version"] == "ATIF-v1.7" + assert payload["agent"]["name"] == "openhands" + # 4 events become 4 steps (the bash action absorbs the next observation + # into its tool_calls[0].observation rather than a separate step). + # Synthesized step ids must start at 1 and be sequential. + step_ids = [s["step_id"] for s in payload["steps"]] + assert step_ids == list(range(1, len(step_ids) + 1)) + # Run step should carry a tool call. + run_steps = [ + s for s in payload["steps"] + if s.get("tool_calls") and s["tool_calls"][0]["function_name"] == "run" + ] + assert len(run_steps) == 1 + assert run_steps[0]["tool_calls"][0]["arguments"]["command"] == "ls /tmp" + # Observation should be attached to that step. + assert run_steps[0].get("observation") is not None + + +def test_on_session_complete_raises_when_nothing_available(tmp_path: Path): + """No buffered events + no trajectory.json in session_dir → raise.""" + adapter = OpenHandsEventStreamAdapter() + with pytest.raises(ValueError, match="no trajectory found"): + adapter.on_session_complete(tmp_path) + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def test_cli_clean_session_returns_zero( + patch_httpx_client, tmp_path: Path, capsys, +): + (tmp_path / "agent").mkdir() + shutil.copy(HARBOR_TRAJECTORY, tmp_path / "agent" / "trajectory.json") + + patch_httpx_client(_mock_transport(_fake_clean_response())) + + exit_code = cli_main([str(tmp_path)]) + assert exit_code == 0 + out = capsys.readouterr().out + assert "No failures detected" in out + + +def test_cli_failure_session_returns_one( + patch_httpx_client, tmp_path: Path, capsys, +): + shutil.copy(HARBOR_TRAJECTORY, tmp_path / "trajectory.json") + + patch_httpx_client(_mock_transport(_fake_failure_response())) + + exit_code = cli_main([str(tmp_path)]) + assert exit_code == 1 + out = capsys.readouterr().out + assert "runtime_error" in out + assert "Runtime / Environment Error" in out + + +def test_cli_missing_session_dir_returns_two(capsys): + exit_code = cli_main([str(Path("/nonexistent/path/qwerty"))]) + assert exit_code == 2 + err = capsys.readouterr().err + assert "does not exist" in err + + +def test_cli_json_mode_emits_diagnosis_json( + patch_httpx_client, tmp_path: Path, capsys, +): + shutil.copy(HARBOR_TRAJECTORY, tmp_path / "trajectory.json") + patch_httpx_client(_mock_transport(_fake_clean_response())) + + exit_code = cli_main([str(tmp_path), "--json"]) + assert exit_code == 0 + out = capsys.readouterr().out + parsed = json.loads(out) + assert "diagnosis" in parsed + assert parsed["diagnosis"]["has_failures"] is False diff --git a/tests/agents/test_real_public_workflows.py b/tests/agents/test_real_public_workflows.py new file mode 100644 index 0000000..3b932af --- /dev/null +++ b/tests/agents/test_real_public_workflows.py @@ -0,0 +1,348 @@ +"""Public SDK workflows using real detectors and captured Harbor data.""" + +from __future__ import annotations + +import json +from pathlib import Path +from queue import Queue +from typing import Any + +import pytest + +from pisama.agents import ( + BridgeConfig, + ClarificationPrimitive, + DetectionBridge, + HealingResult, + HookMatcher, + PostToolUseHook, + PreToolUseHook, + SessionManager, + check, + configure_bridge, + heal_now, + post_tool_use_hook, + pre_tool_use_hook, +) +from pisama.agents.config import load_config +from pisama.agents.indication import SDKIndication +from pisama.agents.tools import create_check_tool, pisama_check_handler + +HARBOR_FIXTURES = Path(__file__).parent / "fixtures" / "harbor" + + +def _captured_bash_call() -> tuple[str, dict[str, Any]]: + trajectory = json.loads( + ( + HARBOR_FIXTURES + / "hello-world-context-summarization.trajectory.summarization-1-summary.json" + ).read_text(encoding="utf-8") + ) + for step in trajectory["steps"]: + calls = step.get("tool_calls") or [] + if calls: + call = calls[0] + return call["function_name"], call["arguments"] + raise AssertionError("captured Harbor trajectory has no tool call") + + +@pytest.mark.asyncio +async def test_real_bridge_blocks_a_repeated_captured_harbor_tool_call() -> None: + tool_name, tool_input = _captured_bash_call() + config = BridgeConfig( + warning_threshold=30, + block_threshold=60, + detection_timeout_ms=1000, + tool_patterns=[".*"], + excluded_tools=[], + ) + bridge = DetectionBridge(config=config, session_mgr=SessionManager()) + + results = [] + for index in range(5): + results.append( + await bridge.analyze_pre_tool( + { + "tool_name": tool_name, + "tool_input": tool_input, + "session_id": "harbor-context-summarization", + }, + f"captured-call-{index}", + ) + ) + + # The captured shell call also triggers approval-bypass guidance at 40. + assert results[0].severity == 40 + assert results[2].severity >= 40 + assert results[2].should_block is True + assert results[2].severity >= 60 + assert "repeated 3x" in results[2].issues[0] + assert results[-1].issues == ["Session is blocked due to previous violations"] + assert bridge.sessions.is_blocked("harbor-context-summarization") + + +@pytest.mark.asyncio +async def test_config_instance_and_matcher_are_supported_by_public_hook_api() -> None: + config = BridgeConfig( + detection_timeout_ms=1000, + tool_patterns=[".*"], + excluded_tools=[], + ) + bridge = configure_bridge(config) + assert bridge.config is config + bridge.sessions.clear_all() + + hook = PreToolUseHook( + bridge=bridge, + matcher=HookMatcher(tool_name_pattern=r"^Read$"), + ) + unmatched = await hook( + { + "tool_name": "bash_command", + "tool_input": {"keystrokes": "mkdir test_dir\n", "duration": 0.1}, + "session_id": "matcher-session", + }, + "captured-bash", + {}, + ) + assert unmatched == {} + baseline_sessions = bridge.sessions.session_count + assert baseline_sessions == 0 + + matched = await hook( + { + "tool_name": "Read", + "tool_input": {"file_path": "PROVENANCE.toml"}, + "session_id": "matcher-session", + }, + "read-provenance", + {}, + ) + assert matched == {} + assert bridge.sessions.session_count == baseline_sessions + 1 + + +@pytest.mark.asyncio +async def test_post_hook_and_self_check_run_real_local_detection() -> None: + bridge = DetectionBridge( + BridgeConfig( + detection_timeout_ms=1000, + tool_patterns=[".*"], + excluded_tools=[], + ), + session_mgr=SessionManager(), + ) + post_hook = PostToolUseHook(bridge=bridge) + output = await post_hook( + { + "tool_name": "bash_command", + "tool_input": {"keystrokes": "mkdir test_dir\n", "duration": 0.1}, + "tool_response": {"content": "command completed"}, + "session_id": "harbor-post-tool", + }, + "captured-post-call", + {}, + ) + assert isinstance(output, dict) + assert bridge.sessions.session_count == 1 + + configure_bridge( + BridgeConfig( + detection_timeout_ms=1000, + tool_patterns=[".*"], + excluded_tools=[], + ) + ) + result = await check( + output="Created test_dir and wrote test_dir/file1.txt.", + context={ + "query": "Create a directory and put a file in it.", + "sources": ["Harbor context-summarization trajectory"], + }, + timeout_ms=1000, + ) + assert result["detectors_run"] == ["realtime"] + assert 0.0 <= result["score"] <= 1.0 + assert result["check_time_ms"] >= 0 + + +@pytest.mark.asyncio +async def test_function_hooks_process_captured_calls_through_configured_bridge() -> None: + tool_name, tool_input = _captured_bash_call() + bridge = configure_bridge( + BridgeConfig( + warning_threshold=30, + block_threshold=60, + detection_timeout_ms=1000, + tool_patterns=[".*"], + excluded_tools=[], + ) + ) + bridge.sessions.clear_all() + session_id = "harbor-function-hooks" + + outputs = [] + for index in range(3): + outputs.append( + await pre_tool_use_hook( + { + "tool_name": tool_name, + "tool_input": tool_input, + "session_id": session_id, + }, + f"captured-function-{index}", + {}, + auto_heal=False, + ) + ) + + assert outputs[0].get("systemMessage") + assert outputs[-1]["hookSpecificOutput"]["permissionDecision"] == "block" + assert "repeated 3x" in outputs[-1]["systemMessage"] + + post = await post_tool_use_hook( + { + "tool_name": "Read", + "tool_input": {"file_path": "PROVENANCE.toml"}, + "tool_response": {"content": "captured provenance"}, + "session_id": "harbor-post-function", + }, + "captured-post-function", + {}, + ) + assert isinstance(post, dict) + + assert await pre_tool_use_hook({}, None, {}) == {} + assert await post_tool_use_hook({}, None, {}) == {} + + +@pytest.mark.asyncio +async def test_custom_check_tool_runs_real_handler_and_validates_empty_input() -> None: + tool = create_check_tool() + assert tool["name"] == "pisama_check" + assert tool["handler"] is pisama_check_handler + assert tool["input_schema"]["required"] == ["output"] + + empty = await pisama_check_handler({}) + assert empty["passed"] is True + assert empty["error"] == "No output provided to check" + + checked = await pisama_check_handler( + { + "output": "Created test_dir and wrote test_dir/file1.txt.", + "context": {"query": "Create a directory and a file."}, + } + ) + assert "passed" in checked + assert "score" in checked + + +def test_config_round_trip_and_environment_loading( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + path = tmp_path / "agent-sdk.json" + expected = BridgeConfig( + warning_threshold=25, + block_threshold=70, + detection_timeout_ms=125, + context_window=14, + tool_patterns=[r"^(Read|bash_command)$"], + excluded_tools=["AskUserQuestion"], + ) + expected.save(path) + loaded = BridgeConfig.from_file(path) + assert loaded.warning_threshold == 25 + assert loaded.block_threshold == 70 + assert loaded.context_window == 14 + assert loaded.tool_patterns == [r"^(Read|bash_command)$"] + + monkeypatch.setenv("PISAMA_CONFIG_PATH", str(path)) + assert load_config().block_threshold == 70 + monkeypatch.delenv("PISAMA_CONFIG_PATH") + monkeypatch.setenv("PISAMA_WARNING_THRESHOLD", "35") + monkeypatch.setenv("PISAMA_ENABLE_BLOCKING", "false") + from_environment = BridgeConfig.from_env() + assert from_environment.warning_threshold == 35 + assert from_environment.enable_blocking is False + + +def test_healing_and_indication_outcomes_cover_safe_and_governance_paths( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("PISAMA_API_KEY", raising=False) + missing_key = heal_now(detection_type="loop") + assert missing_key.applied is False + assert "not configured" in missing_key.message + + safe = HealingResult( + applied=True, + escalated=False, + risk_level="safe", + fix={ + "title": "Bound retries", + "fix_type": "configuration", + "description": "Set a bounded retry count.", + "metadata": {"framework_specific_code": "max_retries = 3"}, + }, + ) + safe_indication = SDKIndication.from_healing_result( + safe, + detection_type="loop", + confidence=0.91, + ) + assert safe.prompt_patch == "max_retries = 3" + assert safe_indication.category == "auto_healed_safe" + assert safe_indication.action_required is False + assert safe_indication.to_dict()["confidence"] == 0.91 + + governance = HealingResult( + applied=False, + escalated=True, + risk_level="dangerous", + fix={ + "metadata": { + "handoff": { + "page": "security-owner", + "sla_hours": 2, + "evidence_summary": "Approval bypass in the captured run.", + } + } + }, + approval_url="https://pisama.ai/review/fix", + ) + governance_indication = SDKIndication.from_healing_result( + governance, + detection_type="approval_bypass", + ) + assert governance_indication.category == "escalated_governance" + assert governance_indication.action_required is True + assert governance_indication.evidence_summary == "Approval bypass in the captured run." + + +def test_clarification_primitive_resolves_entity_confusion() -> None: + answers: Queue[dict[str, Any]] = Queue() + answers.put({"text": "test_dir", "index": 0}) + + primitive = ClarificationPrimitive.for_detection( + { + "detection_type": "entity_confusion", + "details": { + "confused_entities": ["test_dir", "test_dir/file1.txt"], + "context_snippet": "Create a directory and put a file in it.", + }, + }, + answer_provider=lambda request: answers.get(timeout=request.timeout_seconds), + ) + assert primitive is not None + resolution = primitive.resolve_for_detection( + { + "detection_type": "entity_confusion", + "details": {"confused_entities": ["test_dir", "test_dir/file1.txt"]}, + } + ) + assert resolution.answered is True + assert resolution.answer == "test_dir" + assert resolution.answer_index == 0 + + diff --git a/tests/test_auto_platform_exporter.py b/tests/test_auto_platform_exporter.py new file mode 100644 index 0000000..7400479 --- /dev/null +++ b/tests/test_auto_platform_exporter.py @@ -0,0 +1,225 @@ +"""Wire-conformance tests for PisamaPlatformExporter. + +The platform ingest route (``POST /api/v1/traces/ingest``) requires a JWT +bearer minted from the API key via ``POST /api/v1/auth/token`` plus an OTLP +JSON body. These tests run the exporter against a real local HTTP server that +implements that wire contract: token exchange first, OTLP JSON body shape, +one re-exchange + retry on 401, FAILURE on persistent rejection. + +This is the portable subset of the monorepo's +``backend/tests/test_pisama_auto_platform_export.py``, which additionally +verifies the same flow end-to-end against the real backend + Postgres. +""" +from __future__ import annotations + +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer + +import pytest +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor, SpanExportResult +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.trace import SpanKind + +from pisama.auto import __version__ +from pisama.auto._tracer import PisamaPlatformExporter + +INGEST_PATH = "/api/v1/traces/ingest" +TOKEN_PATH = "/api/v1/auth/token" + + +def _set_genai_attrs(span, i: int) -> None: + span.set_attribute("gen_ai.system", "anthropic") + span.set_attribute("gen_ai.request.model", "claude-sonnet-4-6") + span.set_attribute("gen_ai.operation.name", "chat") + span.set_attribute("gen_ai.response.stop_reason", "end_turn") + span.set_attribute("gen_ai.usage.input_tokens", 150 + i) + span.set_attribute("gen_ai.usage.output_tokens", 350 + i) + span.set_attribute("gen_ai.usage.total_tokens", 500 + 2 * i) + span.set_attribute("pisama.test.enabled", True) + span.set_attribute("pisama.test.temperature", 0.25) + span.set_attribute("pisama.test.tags", ["wire", "agent"]) + + +def _make_agent_spans() -> list: + """Two real ReadableSpans (parent + child, one trace) shaped like + pisama_auto's anthropic_patch emits.""" + provider = TracerProvider(resource=Resource.create({ + "service.name": "pisama-auto-export-test", + "pisama.sdk": "pisama-auto", + "pisama.sdk.version": "0.1.0", + })) + memory = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(memory)) + tracer = provider.get_tracer("pisama_auto", "0.1.0") + + with tracer.start_as_current_span( + "gen_ai.chat claude-sonnet-4-6", kind=SpanKind.CLIENT + ) as parent: + _set_genai_attrs(parent, 0) + with tracer.start_as_current_span( + "gen_ai.chat claude-sonnet-4-6", kind=SpanKind.CLIENT + ) as child: + _set_genai_attrs(child, 1) + return list(memory.get_finished_spans()) + + +class _PlatformHandler(BaseHTTPRequestHandler): + """Implements the platform's token + keyless-ingest wire contract.""" + + def log_message(self, *_args) -> None: + pass + + def do_POST(self) -> None: + state = self.server.state # type: ignore[attr-defined] + length = int(self.headers.get("Content-Length", 0)) + body = json.loads(self.rfile.read(length)) if length else {} + state["requests"].append({ + "path": self.path, + "authorization": self.headers.get("Authorization"), + "content_type": self.headers.get("Content-Type"), + "body": body, + }) + if self.path == TOKEN_PATH: + state["tokens_minted"] += 1 + self._reply(200, {"access_token": f"jwt-{state['tokens_minted']}"}) + elif self.path == INGEST_PATH: + if state["ingest_401s_remaining"] > 0: + state["ingest_401s_remaining"] -= 1 + self._reply(401, {"detail": "Token has expired"}) + else: + self._reply(202, {"accepted": 1}) + else: + self._reply(404, {"detail": "not found"}) + + def _reply(self, code: int, payload: dict) -> None: + data = json.dumps(payload).encode("utf-8") + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(data))) + self.end_headers() + self.wfile.write(data) + + +@pytest.fixture +def platform_stand_in(): + server = HTTPServer(("127.0.0.1", 0), _PlatformHandler) + server.state = { # type: ignore[attr-defined] + "requests": [], + "tokens_minted": 0, + "ingest_401s_remaining": 0, + } + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", server.state # type: ignore[attr-defined] + finally: + server.shutdown() + thread.join(timeout=5) + + +def test_exporter_wire_contract(platform_stand_in): + """Token exchange first, then OTLP JSON ingest with the minted JWT.""" + base_url, state = platform_stand_in + spans = _make_agent_spans() + + exporter = PisamaPlatformExporter( + endpoint=f"{base_url}{INGEST_PATH}", api_key="pisama_wire_test_key" + ) + assert exporter.export(spans) is SpanExportResult.SUCCESS + + token_req, ingest_req = state["requests"] + assert token_req["path"] == TOKEN_PATH + assert token_req["body"] == {"api_key": "pisama_wire_test_key", "scope": "ingest"} + + assert ingest_req["path"] == INGEST_PATH + assert ingest_req["authorization"] == "Bearer jwt-1" + assert ingest_req["content_type"] == "application/json" + sent_spans = ingest_req["body"]["resourceSpans"][0]["scopeSpans"][0]["spans"] + assert len(sent_spans) == 2 + for sent in sent_spans: + attrs = {a["key"]: a["value"] for a in sent["attributes"]} + assert attrs["gen_ai.system"] == {"stringValue": "anthropic"} + assert len(sent["traceId"]) == 32 + assert { + a["value"]["intValue"] + for sent in sent_spans + for a in sent["attributes"] + if a["key"] == "gen_ai.usage.input_tokens" + } == {"150", "151"} + # One trace: the child span carries its parent's id. + assert len({sent["traceId"] for sent in sent_spans}) == 1 + assert sum("parentSpanId" in sent for sent in sent_spans) == 1 + + # The cached JWT is reused on the next batch — no second exchange. + assert exporter.export(spans) is SpanExportResult.SUCCESS + assert state["tokens_minted"] == 1 + + +def test_exporter_reexchanges_once_on_401(platform_stand_in): + """An expired JWT gets one re-exchange + retry, mirroring PlatformSession.""" + base_url, state = platform_stand_in + state["ingest_401s_remaining"] = 1 + spans = _make_agent_spans() + + exporter = PisamaPlatformExporter( + endpoint=f"{base_url}{INGEST_PATH}", api_key="pisama_wire_test_key" + ) + assert exporter.export(spans) is SpanExportResult.SUCCESS + assert [r["path"] for r in state["requests"]] == [ + TOKEN_PATH, INGEST_PATH, TOKEN_PATH, INGEST_PATH, + ] + assert state["requests"][-1]["authorization"] == "Bearer jwt-2" + + +def test_exporter_reports_failure_on_persistent_rejection(platform_stand_in): + """Persistent 401s surface as FAILURE (no infinite re-exchange loop).""" + base_url, state = platform_stand_in + state["ingest_401s_remaining"] = 99 + spans = _make_agent_spans() + + exporter = PisamaPlatformExporter( + endpoint=f"{base_url}{INGEST_PATH}", api_key="pisama_wire_test_key" + ) + assert exporter.export(spans) is SpanExportResult.FAILURE + # One initial exchange + exactly one re-exchange. + assert [r["path"] for r in state["requests"]] == [ + TOKEN_PATH, INGEST_PATH, TOKEN_PATH, INGEST_PATH, + ] + + +def test_empty_export_is_a_success_without_network(platform_stand_in): + base_url, state = platform_stand_in + exporter = PisamaPlatformExporter( + endpoint=f"{base_url}{INGEST_PATH}", api_key="pisama_wire_test_key" + ) + + assert exporter.export([]) is SpanExportResult.SUCCESS + assert state["requests"] == [] + + +def test_wire_scope_reports_installed_package_version(platform_stand_in): + base_url, state = platform_stand_in + spans = _make_agent_spans() + exporter = PisamaPlatformExporter( + endpoint=f"{base_url}{INGEST_PATH}", api_key="pisama_wire_test_key" + ) + + assert exporter.export(spans) is SpanExportResult.SUCCESS + scope = state["requests"][-1]["body"]["resourceSpans"][0]["scopeSpans"][0]["scope"] + assert scope == {"name": "pisama_auto", "version": __version__} + + +def test_exporter_reports_connection_failure(): + spans = _make_agent_spans() + exporter = PisamaPlatformExporter( + endpoint="http://127.0.0.1:1/api/v1/traces/ingest", + api_key="pisama_wire_test_key", + timeout=0.1, + ) + + assert exporter.export(spans) is SpanExportResult.FAILURE diff --git a/tests/test_auto_sdk_instrumentation.py b/tests/test_auto_sdk_instrumentation.py new file mode 100644 index 0000000..0956cde --- /dev/null +++ b/tests/test_auto_sdk_instrumentation.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +import json +import os +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer + +import anthropic +import openai +import pytest +from opentelemetry.sdk.resources import Resource +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +import pisama.auto +from pisama.auto import _tracer +from pisama.auto.patches import _patched, patch, patch_all +from pisama.auto.patches.anthropic_patch import _traced_stream, _TracedStream + + +class _SdkHandler(BaseHTTPRequestHandler): + def log_message(self, *_args) -> None: + pass + + def do_POST(self) -> None: + length = int(self.headers.get("Content-Length", 0)) + request = json.loads(self.rfile.read(length)) + self.server.requests.append((self.path, request)) # type: ignore[attr-defined] + + if request.get("model") == "failure-test": + self._reply( + {"error": {"type": "server_error", "message": "deliberate failure"}}, + status=500, + ) + return + + if self.path.endswith("/chat/completions"): + self._reply( + { + "id": "chatcmpl-pisama", + "object": "chat.completion", + "created": 1_721_692_800, + "model": request["model"], + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Paris"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 1, + "total_tokens": 10, + }, + } + ) + return + + if self.path.endswith("/messages"): + self._reply( + { + "id": "msg_pisama", + "type": "message", + "role": "assistant", + "model": request["model"], + "content": [{"type": "text", "text": "Paris"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 9, "output_tokens": 1}, + } + ) + return + + self._reply({"error": {"message": "unknown route"}}, status=404) + + def _reply(self, payload: dict, status: int = 200) -> None: + body = json.dumps(payload).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + +@pytest.fixture +def sdk_server(): + server = HTTPServer(("127.0.0.1", 0), _SdkHandler) + server.requests = [] # type: ignore[attr-defined] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}", server.requests # type: ignore[attr-defined] + finally: + server.shutdown() + thread.join(timeout=5) + + +@pytest.fixture +def captured_spans(): + exporter = InMemorySpanExporter() + provider = TracerProvider(resource=Resource.create({"service.name": "sdk-test"})) + provider.add_span_processor(SimpleSpanProcessor(exporter)) + _tracer._tracer = provider.get_tracer("pisama_auto") + yield exporter + _tracer._tracer = None + + +def test_real_openai_and_anthropic_clients_emit_semantic_spans(sdk_server, captured_spans): + base_url, requests = sdk_server + _patched.clear() + + assert set(patch_all()) == {"openai", "anthropic"} + assert patch("openai") + assert patch("unknown-sdk") is False + + openai_client = openai.OpenAI(api_key="test-key", base_url=f"{base_url}/v1") + openai_response = openai_client.chat.completions.create( + model="gpt-5.5", + messages=[{"role": "user", "content": "Capital of France?"}], + max_tokens=8, + temperature=0, + ) + anthropic_client = anthropic.Anthropic(api_key="test-key", base_url=base_url) + anthropic_response = anthropic_client.messages.create( + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Capital of France?"}], + max_tokens=8, + temperature=0, + system="Answer with one city.", + ) + + assert openai_response.choices[0].message.content == "Paris" + assert anthropic_response.content[0].text == "Paris" + assert len(requests) == 2 + + spans = captured_spans.get_finished_spans() + assert [span.name for span in spans] == [ + "gen_ai.chat gpt-5.5", + "gen_ai.chat claude-sonnet-4-6", + ] + openai_attrs = dict(spans[0].attributes) + anthropic_attrs = dict(spans[1].attributes) + assert openai_attrs["gen_ai.usage.total_tokens"] == 10 + assert openai_attrs["gen_ai.response.finish_reason"] == "stop" + assert anthropic_attrs["gen_ai.usage.total_tokens"] == 10 + assert anthropic_attrs["gen_ai.request.has_system"] is True + + +def test_anthropic_stream_wrapper_ends_span_on_real_file_iteration(tmp_path, captured_spans): + stream_path = tmp_path / "events.txt" + stream_path.write_text("first\nsecond\n", encoding="utf-8") + + with stream_path.open(encoding="utf-8") as stream: + wrapped = _traced_stream( + lambda **_kwargs: stream, + None, + (), + {"model": "claude-sonnet-4-6", "messages": [{"role": "user"}]}, + ) + assert isinstance(wrapped, _TracedStream) + assert wrapped.name == str(stream_path) + with wrapped as active: + assert list(active) == ["first\n", "second\n"] + + spans = captured_spans.get_finished_spans() + assert len(spans) == 1 + assert spans[0].name == "gen_ai.chat.stream claude-sonnet-4-6" + assert spans[0].attributes["gen_ai.operation.name"] == "chat.stream" + + +def test_init_without_credentials_is_idempotent(): + previous_api_key = os.environ.pop("PISAMA_API_KEY", None) + previous_endpoint = os.environ.pop("PISAMA_ENDPOINT", None) + pisama.auto._initialized = False + + try: + pisama.auto.init(service_name="integration-test", auto_patch=True) + first_tracer = _tracer._tracer + pisama.auto.init(service_name="ignored", auto_patch=True) + + assert pisama.auto.is_initialized() + assert _tracer._tracer is first_tracer + finally: + if previous_api_key is not None: + os.environ["PISAMA_API_KEY"] = previous_api_key + if previous_endpoint is not None: + os.environ["PISAMA_ENDPOINT"] = previous_endpoint + + +def test_tracer_setup_selects_platform_and_collector_transports(monkeypatch): + platform = _tracer.setup_tracer( + api_key="platform-key", + endpoint="https://api.pisama.ai/api/v1/traces/ingest", + service_name="platform-test", + ) + collector = _tracer.setup_tracer( + api_key="collector-key", + endpoint="http://127.0.0.1:4318/v1/traces", + service_name="collector-test", + ) + monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "grpc") + grpc_fallback = _tracer.setup_tracer( + api_key="collector-key", + endpoint="http://127.0.0.1:4317", + service_name="grpc-fallback-test", + ) + + assert platform is not None + assert collector is not None + assert grpc_fallback is not None + + +def test_stream_wrapper_records_iteration_errors(tmp_path, captured_spans): + stream_path = tmp_path / "closed.txt" + stream_path.write_text("event\n", encoding="utf-8") + stream = stream_path.open(encoding="utf-8") + stream.close() + wrapped = _TracedStream(stream, _tracer.get_tracer().start_span("closed-stream")) + + with pytest.raises(ValueError, match="closed file"): + next(wrapped) + + span = captured_spans.get_finished_spans()[0] + assert span.attributes["error.type"] == "ValueError" + + +def test_real_sdk_errors_are_recorded_and_propagated(sdk_server, captured_spans): + base_url, _requests = sdk_server + openai_client = openai.OpenAI( + api_key="test-key", + base_url=f"{base_url}/v1", + max_retries=0, + ) + anthropic_client = anthropic.Anthropic( + api_key="test-key", + base_url=base_url, + max_retries=0, + ) + + with pytest.raises(openai.InternalServerError): + openai_client.chat.completions.create( + model="failure-test", + messages=[{"role": "user", "content": "Trigger the error contract."}], + ) + with pytest.raises(anthropic.InternalServerError): + anthropic_client.messages.create( + model="failure-test", + messages=[{"role": "user", "content": "Trigger the error contract."}], + max_tokens=8, + ) + + spans = captured_spans.get_finished_spans() + assert len(spans) == 2 + assert all(span.attributes["error.type"] == "InternalServerError" for span in spans) + + +def test_stream_context_records_body_errors(tmp_path, captured_spans): + stream_path = tmp_path / "events.txt" + stream_path.write_text("event\n", encoding="utf-8") + stream = stream_path.open(encoding="utf-8") + wrapped = _TracedStream(stream, _tracer.get_tracer().start_span("body-error")) + + with pytest.raises(RuntimeError, match="consumer failed"): + with wrapped: + raise RuntimeError("consumer failed") + + span = captured_spans.get_finished_spans()[0] + assert span.attributes["error.type"] == "RuntimeError" + + +def test_stream_call_error_ends_span(tmp_path, captured_spans): + missing = tmp_path / "missing-events.txt" + + with pytest.raises(FileNotFoundError): + _traced_stream( + lambda **_kwargs: missing.open(encoding="utf-8"), + None, + (), + {"model": "claude-sonnet-4-6", "messages": []}, + ) + + span = captured_spans.get_finished_spans()[0] + assert span.attributes["error.type"] == "FileNotFoundError" + + +def test_patch_registry_handles_absent_and_broken_integrations(monkeypatch): + from pisama.auto import patches + + monkeypatch.setitem(patches._PATCHABLE, "definitely_not_installed", ".missing") + monkeypatch.setitem(patches._PATCHABLE, "json", ".missing") + patches._patched.clear() + + patched = patch_all() + + assert "definitely_not_installed" not in patched + assert "json" not in patched + assert patch("json") is False From 8f3818e1d7fe49a1c83c40ec2d845b040ebb8711 Mon Sep 17 00:00:00 2001 From: tn-pisama Date: Wed, 29 Jul 2026 14:21:47 -0700 Subject: [PATCH 2/2] ci: install auto extra so pisama.auto tests run The quality and test jobs installed -e ".[dev,mcp]", which never pulled in the new auto extra this PR adds. tests/test_auto_platform_exporter.py and tests/test_auto_sdk_instrumentation.py import opentelemetry.sdk modules unconditionally at collection time, so pytest failed to collect and the whole run aborted for every job on that install line (quality, test on 3.10-3.13). Switch both to -e ".[dev,mcp,auto]" so pisama.auto is actually exercised. distribution is untouched -- it deliberately installs the bare wheel with no extras. Co-Authored-By: Claude Sonnet 5 --- .github/workflows/ci.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b66ef2a..f498c07 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -19,7 +19,7 @@ jobs: python-version: "3.12" cache: pip - run: python -m pip install "pisama-core @ git+https://github.com/Pisama-AI/pisama-core.git@v1.8.2" - - run: python -m pip install -e ".[dev,mcp]" + - run: python -m pip install -e ".[dev,mcp,auto]" - run: ruff check src tests - run: mypy src/pisama - run: pytest -q --cov=pisama --cov-report=term --cov-fail-under=60 @@ -37,7 +37,7 @@ jobs: python-version: ${{ matrix.python-version }} cache: pip - run: python -m pip install "pisama-core @ git+https://github.com/Pisama-AI/pisama-core.git@v1.8.2" - - run: python -m pip install -e ".[dev,mcp]" + - run: python -m pip install -e ".[dev,mcp,auto]" - run: pytest -q distribution: