diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..ea2b674 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,38 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + branches: [main] + +permissions: + contents: read + +jobs: + python-sdk: + name: Python SDK (py${{ matrix.python-version }}) + runs-on: ubuntu-latest + defaults: + run: + working-directory: sdk/python + strategy: + fail-fast: false + matrix: + python-version: ["3.10", "3.11", "3.12", "3.13"] + + steps: + - uses: actions/checkout@v4 + + - name: Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install -e ".[dev]" + + - name: Run tests + run: python -m pytest --tb=short -q diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..255bfa2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,9 @@ +.claude/settings.local.json +.idea +.vscode +dist +node_modules +.DS_Store +__pycache__ +*.pyc +.env diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..52a1b3f --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,69 @@ +[Problem Statement](docs/problem-statement.md) | +[Use Cases](docs/use-cases.md) | +[Approaches](docs/approaches.md) | +[Open Questions](docs/open-questions.md) | +[Experimental Findings](docs/experimental-findings.md) | +[Related Work](docs/related-work.md) | +Contributing + +# Contributing + +## How to Participate + +This Interest Group welcomes contributions from anyone interested in tool annotations, data classification, and sensitivity metadata for MCP. You can participate by: + +- Opening or commenting on GitHub Discussions in this repo +- Sharing experimental findings from your own implementations +- Contributing to documentation and pattern evaluation + +## Communication Channels + +| Channel | Purpose | Response Expectation | +|---------|---------|----------------------| +| [Discord: #tool-annotations-ig](https://discord.com/channels/1358869848138059966/1482836798517543073) | Quick questions, coordination, async discussion | Best effort | +| GitHub Discussions | Long-form technical proposals, experimental findings | Weekly triage | +| This repository | Living reference for approaches, findings, and decisions | Updated after meetings | + +## Meetings + +*To be scheduled* — likely biweekly working sessions once the group establishes momentum. + +Meeting norms: + +- Agendas published 24 hours in advance +- Notes published within 48 hours + +## Decision-Making + +As an Interest Group, we operate by **rough consensus** — we're exploring and recommending, not deciding. Outputs include: + +- Documented requirements and use cases +- Evaluated approaches with findings +- Recommendations to relevant WGs or as SEP proposals + +## Contribution Guidelines + +### Documenting Approaches and Findings + +When adding experimental findings or new approaches: + +- Include enough detail for others to reproduce or evaluate +- Note which clients and servers were tested +- Be explicit about what worked, what didn't, and what remains untested +- Attribute community input with GitHub handles and link to the source where possible + +### Community Input + +When adding quotes or input from community discussions: + +- Attribute to the contributor by name and GitHub handle +- Link to the original source (Discord thread, GitHub comment, etc.) where possible +- Present input as blockquotes to distinguish it from editorial content + +### Filing Issues + +Use GitHub Issues for: + +- Proposing new approaches or use cases +- Reporting gaps in documentation +- Tracking action items from meetings diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..04b3c0c --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 MCP Trust Contributors + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index dc87aeb..c7c4ad9 100644 --- a/README.md +++ b/README.md @@ -1 +1,86 @@ # experimental-ext-tool-annotations + +> ⚠️ **Experimental** — This repository is an incubation space for the Tool Annotations Interest Group. Contents are exploratory and do not represent official MCP specifications or recommendations. + +## Mission + +This Interest Group explores how MCP tool annotations can be enhanced to support data classification, sensitivity labeling, and provenance tracking — enabling hosts and clients to make informed, policy-driven decisions about data flow across tool boundaries. + +As a starting point, we are contributing a reference SDK implementation based on [SEP-1913](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/1913) (Trust & Sensitivity Annotations). + +## Scope + +### In Scope + +- **Requirements gathering:** Documenting use cases and constraints for tool-level data sensitivity metadata +- **Pattern exploration:** Testing and evaluating annotation schemas, policy enforcement models, and sensitivity propagation strategies +- **Proof of concepts:** Maintaining a shared repo of reference implementations and experimental findings (starting with the SEP-1913 Python SDK) + +### Out of Scope + +- **Approving spec changes:** This IG does not have authority to approve protocol changes; recommendations flow through the SEP process +- **Implementation mandates:** We can document patterns but not require specific client or server behavior + +## Problem Statement + +MCP tools today are semantically opaque when it comes to data sensitivity. A tool's definition includes its name, description, and input schema — but nothing about the nature of the data it returns or processes. There is no machine-readable signal that allows a host to differentiate a benign health check from a HIPAA-regulated patient record lookup. + +- **No sensitivity metadata** — Tools carry no data classification labels; hosts cannot distinguish PII from public data +- **No destination awareness** — The protocol does not express where tool outputs may flow (storage, third-party APIs, user display) +- **No policy enforcement surface** — Without structured annotations, security policies cannot be evaluated at the protocol level +- **LLM-dependent safety** — Data flow decisions rely entirely on the LLM's interpretation of tool descriptions, which can be bypassed by prompt injection, hallucination, or inadequate descriptions + +See the [Problem Statement](docs/problem-statement.md) for full details. + +## Repository Contents + +| Document | Description | +| :--- | :--- | +| [Problem Statement](docs/problem-statement.md) | Current limitations and gaps | +| [Use Cases](docs/use-cases.md) | Key use cases driving this work | +| [Approaches](docs/approaches.md) | Approaches being explored (not mutually exclusive) | +| [Open Questions](docs/open-questions.md) | Unresolved questions with community input | +| [Experimental Findings](docs/experimental-findings.md) | Results from implementations and testing | +| [Related Work](docs/related-work.md) | SEPs, implementations, and external resources | +| [Contributing](CONTRIBUTING.md) | How to participate | +| [Python SDK](sdk/python/) | Reference implementation (SEP-1913, 138 tests) | + +## Facilitators + +| Role | Name | Organization | GitHub | +| :--- | :--- | :--- | :--- | +| TO BE ADDED | TO BE ADDED | TO BE ADDED| TO BE ADDED | + +## Lifecycle + +**Current Status: Active Exploration** + +### Graduation Criteria (IG → WG) + +This IG may propose becoming a Working Group if: + +- Clear consensus emerges on an annotation schema requiring sustained spec work +- Cross-cutting coordination requires formal authority delegation +- At least two Core Maintainers sponsor WG formation + +### Retirement Criteria + +- Problem space resolved (conventions established, absorbed into other WGs) +- Insufficient participation to maintain momentum +- Community consensus that tool annotations don't belong in MCP protocol scope + +## Work Tracking + +| Item | Status | Champion | Notes | +| :--- | :--- | :--- | :--- | +| Repository scaffolding | Done | All facilitators | Align structure with other experimental-ext repos | +| Problem statement & use cases | Done | All facilitators | Document motivating scenarios and constraints | +| SEP-1913 reference SDK | Done | TBD | Python SDK implementing trust & sensitivity annotations (138 tests passing) | +| Experimental findings | Proposed | TBD | Usability study results and implementation learnings | + +## Success Criteria + +- **Short-term:** Documented consensus on requirements and evaluation of existing annotation approaches +- **Medium-term:** Clear recommendation (annotation schema convention vs. protocol extension vs. both) +- **Long-term:** Interoperable tool annotation convention across MCP servers and clients + diff --git a/docs/approaches.md b/docs/approaches.md new file mode 100644 index 0000000..7f8bda8 --- /dev/null +++ b/docs/approaches.md @@ -0,0 +1,11 @@ +[Problem Statement](problem-statement.md) | +[Use Cases](use-cases.md) | +Approaches | +[Open Questions](open-questions.md) | +[Experimental Findings](experimental-findings.md) | +[Related Work](related-work.md) | +[Contributing](../CONTRIBUTING.md) + +# Approaches Being Explored + +*This section will document annotation schema designs, policy enforcement models, and sensitivity propagation strategies being considered by the Interest Group.* diff --git a/docs/experimental-findings.md b/docs/experimental-findings.md new file mode 100644 index 0000000..7afb76f --- /dev/null +++ b/docs/experimental-findings.md @@ -0,0 +1,11 @@ +[Problem Statement](problem-statement.md) | +[Use Cases](use-cases.md) | +[Approaches](approaches.md) | +[Open Questions](open-questions.md) | +Experimental Findings | +[Related Work](related-work.md) | +[Contributing](../CONTRIBUTING.md) + +# Experimental Findings + +*This section will document results from implementations and testing.* diff --git a/docs/open-questions.md b/docs/open-questions.md new file mode 100644 index 0000000..ba0363a --- /dev/null +++ b/docs/open-questions.md @@ -0,0 +1,11 @@ +[Problem Statement](problem-statement.md) | +[Use Cases](use-cases.md) | +[Approaches](approaches.md) | +Open Questions | +[Experimental Findings](experimental-findings.md) | +[Related Work](related-work.md) | +[Contributing](../CONTRIBUTING.md) + +# Open Questions + +*This section will track unresolved questions requiring community input.* diff --git a/docs/problem-statement.md b/docs/problem-statement.md new file mode 100644 index 0000000..0038329 --- /dev/null +++ b/docs/problem-statement.md @@ -0,0 +1,56 @@ +Problem Statement | +[Use Cases](use-cases.md) | +[Approaches](approaches.md) | +[Open Questions](open-questions.md) | +[Experimental Findings](experimental-findings.md) | +[Related Work](related-work.md) | +[Contributing](../CONTRIBUTING.md) + +# Problem Statement + +MCP tools today are semantically opaque when it comes to data sensitivity. A tool's definition includes its name, description, and input schema — but nothing about the nature of the data it returns or processes. Consider these two tools: + +```jsonc +// Tool 1: System health check +{ + "name": "health_check", + "description": "Check if the system is running", + "inputSchema": {} +} + +// Tool 2: Patient record lookup +{ + "name": "patient_lookup", + "description": "Look up a patient record from the EHR", + "inputSchema": { + "type": "object", + "properties": { "patient_id": { "type": "string" } } + } +} +``` + +At the protocol level, these tools are indistinguishable in terms of data sensitivity. Both are just callable functions. Yet one returns benign system status, while the other returns Protected Health Information (PHI) regulated under HIPAA law. There is no machine-readable signal that allows a host to differentiate between them. + +## Key Gaps + +### No Sensitivity Metadata + +Tools carry no data classification labels. A host cannot distinguish a tool returning public system metrics from one returning PII, financial records, or credentials. Without classification, all tool outputs are treated equally — which means sensitive data receives no additional protection. + +### No Destination Awareness + +The protocol does not express where tool outputs may flow. A tool that sends an email to an external address, writes to a third-party API, or stores data in a public bucket looks identical to one that only returns data to the user. Hosts have no structured way to assess data egress risk. + +### No Policy Enforcement Surface + +Without structured annotations, security policies cannot be evaluated at the protocol level. Organizations cannot express rules like "never send credentials to external destinations" or "require user confirmation before forwarding PII" — because the primitives to describe sensitivity and destination don't exist. + +### LLM-Dependent Safety + +Data flow decisions currently rely entirely on the LLM's interpretation of tool descriptions. This can be bypassed by: + +- **Prompt injection** — Malicious content in tool outputs can instruct the model to misroute data +- **Hallucination** — The model may fabricate a safe interpretation of an ambiguous description +- **Inadequate descriptions** — Tool authors may not describe sensitivity in human-readable text + +⚠️ The common thread: the host has no structured metadata to make informed decisions about data flow. \ No newline at end of file diff --git a/docs/related-work.md b/docs/related-work.md new file mode 100644 index 0000000..8eca88d --- /dev/null +++ b/docs/related-work.md @@ -0,0 +1,11 @@ +[Problem Statement](problem-statement.md) | +[Use Cases](use-cases.md) | +[Approaches](approaches.md) | +[Open Questions](open-questions.md) | +[Experimental Findings](experimental-findings.md) | +Related Work | +[Contributing](../CONTRIBUTING.md) + +# Related Work + +*This section will track SEPs, implementations, and external resources related to tool annotations.* diff --git a/docs/use-cases.md b/docs/use-cases.md new file mode 100644 index 0000000..1a2af2b --- /dev/null +++ b/docs/use-cases.md @@ -0,0 +1,39 @@ +[Problem Statement](problem-statement.md) | +Use Cases | +[Approaches](approaches.md) | +[Open Questions](open-questions.md) | +[Experimental Findings](experimental-findings.md) | +[Related Work](related-work.md) | +[Contributing](../CONTRIBUTING.md) + +# Use Cases + +## UC-1: Healthcare Data Protection + +A healthcare MCP server exposes tools for patient lookup, insurance processing, and staff directory queries. A host needs to enforce HIPAA-compliant data handling: patient data must not be forwarded to external services, and tools returning Protected Health Information (PHI) must be distinguishable from those returning non-sensitive data. + +**Requires:** Data classification (`regulated(HIPAA)`), destination constraints, policy enforcement. + +## UC-2: Multi-Agent Data Leak Prevention + +In a multi-agent architecture, Agent A retrieves sensitive credentials from a vault tool and Agent B has access to an external email tool. Without sensitivity metadata and propagation tracking, there is no protocol-level mechanism to prevent Agent B from exfiltrating credential data it received from Agent A's context. + +**Requires:** Sensitivity propagation across tool calls, session-level tracking, policy rules blocking credential egress. + +## UC-3: Audit Logging for Compliance + +An organization must maintain audit trails showing which tools accessed sensitive data, what classification the data had, and whether policy rules permitted or blocked the action. Current MCP provides no structured metadata to include in audit events. + +**Requires:** Structured annotations on tool calls and results, policy decision logging. + +## UC-4: User Consent for Sensitive Operations + +A tool sends a notification to a patient via SMS. The host should be able to detect — from annotations, not just the tool description — that this tool has an external destination and involves PII, and therefore should prompt the user for confirmation before execution. + +**Requires:** Destination metadata (`external`), sensitivity labels, client-side policy evaluation. + +## UC-5: Progressive Adoption + +Not all MCP server authors need full trust annotations. A simple internal tool may need no annotations at all, while a healthcare integration requires full classification, policy, and audit. The annotation system must support progressive adoption — from zero annotations to full enforcement — without requiring all-or-nothing commitment. + +**Requires:** Optional annotations, graceful degradation, layered adoption levels. \ No newline at end of file diff --git a/sdk/python/README.md b/sdk/python/README.md new file mode 100644 index 0000000..f0f15ca --- /dev/null +++ b/sdk/python/README.md @@ -0,0 +1,230 @@ +# mcp-trust-annotations + +> SEP-1913 Trust & Sensitivity Annotations SDK for the Model Context Protocol + +[![Python 3.10+](https://img.shields.io/badge/python-3.10%2B-blue.svg)]() +[![Zero Dependencies](https://img.shields.io/badge/dependencies-zero-green.svg)]() +[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)]() + +The first implementation of [SEP-1913: Trust and Sensitivity Annotations](https://github.com/modelcontextprotocol/modelcontextprotocol/pull/1913) for Python. Annotate your MCP tools with data classification metadata, track sensitivity propagation across agent sessions, and enforce security policies — all with zero external dependencies. + +--- + +## Why This Exists + +MCP tools can read patient records, send emails, deploy services, and process payments — but there's no standard way to tell a client **what kind of data a tool handles**. SEP-1913 proposes adding trust and sensitivity annotations to the MCP spec. This SDK implements those types so you can start using them today. + +**Without annotations:** +``` +Agent reads patient record → Agent sends email with patient data → 💥 HIPAA violation +``` + +**With annotations + policy enforcement:** +``` +Agent reads patient record → Session marked as HIPAA-regulated → +Agent tries to send email to public → ❌ BLOCKED by policy engine +``` + +--- + +## Installation + +```bash +cd sdk/python +pip install -e ".[dev]" +``` + +**Requirements:** Python 3.10+. No external dependencies. + +--- + +## Quick Start + +### 1. Annotate a Tool + +```python +from annotate import trust_annotated +from trust_types import ReturnMetadata, Source, Regulated + +@trust_annotated( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + attribution=("Epic-EHR",), +) +async def patient_lookup(patient_id: str) -> dict: + """Look up a patient record from the EHR system.""" + return await ehr.get_patient(patient_id) +``` + +### 2. Serialize to MCP Wire Format + +```python +from annotate import get_trust_annotations, to_wire + +ann = get_trust_annotations(patient_lookup) +wire = to_wire(ann) + +# Use in your MCP tools/list response: +tool_def = { + "name": "patient_lookup", + "description": "Look up a patient record", + "inputSchema": { ... }, + "annotations": { + "readOnlyHint": True, + **wire, # ← SEP-1913 fields injected here + }, +} +``` + +### 3. Track Session Propagation + +```python +from propagate import SessionTracker +from trust_types import ResultAnnotations, SimpleDataClass, Regulated + +tracker = SessionTracker(session_id="session-001") + +# After calling health_check (no sensitive data) +tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.NONE), "health_check") + +# After calling patient_lookup (HIPAA-regulated) +tracker.merge( + ResultAnnotations(sensitivity=Regulated.of("HIPAA"), attribution=("Epic-EHR",)), + "patient_lookup", +) +print(tracker.session.max_sensitivity) # → Regulated(HIPAA) + +# Sensitivity never de-escalates +tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.NONE), "health_check") +print(tracker.session.max_sensitivity) # → still Regulated(HIPAA) +``` + +### 4. Enforce Policies + +```python +from policy import PolicyEngine + +engine = PolicyEngine(mode="enforce") # "audit" | "warn" | "enforce" +engine.register_tool("patient_lookup", wire) + +decision = engine.evaluate( + "patient_lookup", + action="call", + target_destination="public", +) +print(decision.allowed) # → False +print(decision.reason) # → "Regulated data (HIPAA) cannot leave organization" +``` + +### 5. Enable Audit Logging + +```python +from emit import enable_logging + +enable_logging(agent_id="urn:agent:my-app") +``` + +--- + +## Running Tests + +```bash +cd sdk/python +pip install -e ".[dev]" +python -m pytest tests/ -v +``` + +--- + +## Architecture + +``` +src/ +├── trust_types.py # SEP-1913 type definitions (enums, dataclasses) +├── annotate.py # @trust_annotated decorator, to_wire/from_wire +├── propagate.py # SessionTracker: sensitivity escalation +├── policy.py # PolicyEngine: audit/warn/enforce modes +└── emit.py # Structured JSON audit logging +``` + +**Dependency rule:** Modules only import downward. `emit.py` ↔ `annotate.py` circular dependency is resolved via lazy imports inside function bodies. + +--- + +## Examples + +The `examples/` directory contains working demonstrations of SEP-1913 annotations, organized by use case: + +``` +examples/ + _shared/ # Shared MCP server used by all examples + healthcare/ # UC-1: HIPAA policy enforcement + multi-agent/ # UC-2: Cross-agent PHI leak prevention + dashboard/ # Interactive web UI for all scenarios +``` + +**Prerequisites:** `pip install mcp` (the official MCP Python SDK, for stdio transport) + +### Shared Server + +The Healthcare Clinic MCP server in `examples/_shared/` is used by all examples. It exposes 8 annotated tools (health check, patient lookup, staff directory, insurance claims, API key rotation, notifications) with varying sensitivity levels, destinations, and outcomes. + +```bash +cd sdk/python +PYTHONPATH=src python examples/_shared/mcp_server.py +``` + +### UC-1: Healthcare Data Protection + +Demonstrates HIPAA-compliant data handling — patient data must not be forwarded to external services, and tools returning PHI must be distinguishable from non-sensitive tools. + +```bash +cd sdk/python + +# Client: connects via stdio, calls tools, shows policy enforcement +PYTHONPATH=src python examples/healthcare/mcp_client.py + +# Host: builds a PolicyEngine from annotations, demonstrates ALLOW / BLOCK / ESCALATE / REDACT +PYTHONPATH=src python examples/healthcare/host.py +``` + +### UC-2: Multi-Agent Data Leak Prevention + +Three agents (front-desk → analytics → external reporting) handling HIPAA patient data. Shows how PHI leaks freely without annotations vs. how SEP-1913 policy enforcement blocks the leak deterministically. + +```bash +cd sdk/python +PYTHONPATH=src python examples/multi-agent/data_leak_prevention.py +``` + +### Web Dashboard + +```bash +cd sdk/python +pip install starlette uvicorn +PYTHONPATH=src python examples/dashboard/app.py +``` + +Opens at http://localhost:8913. Interactive web UI for: +- Viewing all server tools and their SEP-1913 annotations +- Calling tools interactively +- Running UC-1, UC-2, and host scenarios +- Executing usability tests +- Building custom policy enforcement tests + +**Note**: Start mcp server and dashboard using script: + +On Linux/macOS: +```bash +cd sdk/python/examples/dashboard +chmod +x start_dashboard.sh +./start_dashboard.sh +``` + +On Windows (PowerShell): +```powershell +cd sdk\python\examples\dashboard +.\start_dashboard.ps1 +``` diff --git a/sdk/python/examples/README.md b/sdk/python/examples/README.md new file mode 100644 index 0000000..a19adf1 --- /dev/null +++ b/sdk/python/examples/README.md @@ -0,0 +1,31 @@ +# Examples + +Working demonstrations of SEP-1913 trust annotations, organized by use case. + +``` +_shared/ Shared Healthcare Clinic MCP server (used by all examples) +healthcare/ UC-1: Healthcare Data Protection — HIPAA policy enforcement +multi-agent/ UC-2: Multi-Agent Leak Prevention — cross-agent PHI containment +dashboard/ Interactive web UI for all scenarios +``` + +## Prerequisites + +```bash +pip install mcp # MCP Python SDK (stdio transport) +pip install starlette uvicorn # only for dashboard +``` + +## Quick Start + +```bash +cd sdk/python + +# Run any example: +PYTHONPATH=src python examples/healthcare/mcp_client.py +PYTHONPATH=src python examples/healthcare/host.py +PYTHONPATH=src python examples/multi-agent/data_leak_prevention.py +PYTHONPATH=src python examples/dashboard/app.py +``` + +See each subfolder's README for details. diff --git a/sdk/python/examples/_shared/README.md b/sdk/python/examples/_shared/README.md new file mode 100644 index 0000000..f5fd67e --- /dev/null +++ b/sdk/python/examples/_shared/README.md @@ -0,0 +1,23 @@ +# Shared: Healthcare Clinic MCP Server + +Stdio-based MCP server with 8 SEP-1913 annotated tools covering a range of sensitivity levels: + +| Tool | Sensitivity | Key Annotations | +|------|-------------|-----------------| +| `health_check` | none | read-only | +| `patient_lookup` | regulated(HIPAA) | PHI, internal source | +| `search_patients` | regulated(HIPAA) | PHI, internal source | +| `update_patient_record` | regulated(HIPAA) | destructive, consequential | +| `staff_directory` | PII | internal source | +| `process_insurance_claim` | financial | consequential | +| `rotate_api_key` | credentials | system source | +| `send_notification` | none | public destination, malicious-activity hint | + +## Usage + +```bash +cd sdk/python +PYTHONPATH=src python examples/_shared/mcp_server.py +``` + +This server is started automatically by the other examples via stdio transport. diff --git a/sdk/python/examples/_shared/mcp_server.py b/sdk/python/examples/_shared/mcp_server.py new file mode 100644 index 0000000..f1b2963 --- /dev/null +++ b/sdk/python/examples/_shared/mcp_server.py @@ -0,0 +1,341 @@ +"""Healthcare Clinic — Real MCP Server with SEP-1913 trust annotations. + +Runs as a stdio MCP server using the official MCP Python SDK (FastMCP). +Every tool carries SEP-1913 trust annotations via ToolAnnotations extra fields. + +Usage: + # Run directly (stdio transport): + PYTHONPATH=sdk/python/src python sdk/python/examples/_shared/mcp_server.py + + # Or via the client: + PYTHONPATH=sdk/python/src python sdk/python/examples/healthcare/mcp_client.py +""" + +from __future__ import annotations + +import sys +import os + +# Ensure src/ is importable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) + +from mcp.server import FastMCP +from mcp.types import ToolAnnotations + +from trust_types import ( + Destination, + InputMetadata, + Outcome, + Regulated, + ReturnMetadata, + SimpleDataClass, + Source, + TrustAnnotations, +) +from annotate import to_wire + + +# =================================================================== +# In-memory data stores (shared across the server session) +# =================================================================== + +PATIENTS: dict[str, dict] = { + "P-12345": { + "name": "Jane Doe", + "dob": "1985-03-15", + "ssn": "***-**-6789", + "diagnoses": ["Type 2 Diabetes", "Hypertension"], + "medications": ["Metformin 500mg", "Lisinopril 10mg"], + }, + "P-67890": { + "name": "John Smith", + "dob": "1972-11-02", + "ssn": "***-**-1234", + "diagnoses": ["Asthma"], + "medications": ["Albuterol inhaler"], + }, +} + +CLAIMS: list[dict] = [] +NOTIFICATIONS: list[dict] = [] +_audit_counter = 0 +_claim_counter = 0 +_msg_counter = 0 + + +def _next_audit_id() -> str: + global _audit_counter + _audit_counter += 1 + return f"AUD-{_audit_counter:05d}" + + +def _next_claim_id() -> str: + global _claim_counter + _claim_counter += 1 + return f"CLM-{_claim_counter:05d}" + + +def _next_msg_id() -> str: + global _msg_counter + _msg_counter += 1 + return f"MSG-{_msg_counter:05d}" + + +def _make_annotations( + *, + read_only: bool = False, + destructive: bool = False, + open_world: bool = False, + trust: TrustAnnotations | None = None, +) -> ToolAnnotations: + """Build ToolAnnotations with standard MCP hints + SEP-1913 extensions.""" + base = { + "readOnlyHint": read_only, + "destructiveHint": destructive, + "openWorldHint": open_world, + } + if trust: + base.update(to_wire(trust)) + return ToolAnnotations(**base) + + +# =================================================================== +# Server +# =================================================================== + +server = FastMCP( + "Healthcare Clinic", + instructions=( + "A healthcare clinic MCP server with SEP-1913 trust annotations. " + "Tools dealing with patient data are HIPAA-regulated. " + "Always check tool annotations before forwarding results." + ), +) + + +# ── Tool 1: health_check ────────────────────────────────────────── + +@server.tool( + name="health_check", + description="Returns server health status. No sensitive data involved.", + annotations=_make_annotations(read_only=True), +) +def health_check() -> dict: + return {"status": "ok", "uptime_seconds": 86400, "version": "1.2.0"} + + +# ── Tool 2: patient_lookup ──────────────────────────────────────── + +@server.tool( + name="patient_lookup", + description="Look up a patient by ID. Returns HIPAA-regulated PHI.", + annotations=_make_annotations( + read_only=True, + trust=TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + attribution=("Epic-EHR",), + ), + ), +) +def patient_lookup(patient_id: str) -> dict: + patient = PATIENTS.get(patient_id) + if patient is None: + return {"error": f"Patient {patient_id} not found"} + return {"patient_id": patient_id, **patient} + + +# ── Tool 3: search_patients ────────────────────────────────────── + +@server.tool( + name="search_patients", + description="Search patients by name or diagnosis. Returns HIPAA-regulated PHI.", + annotations=_make_annotations( + read_only=True, + trust=TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + attribution=("Epic-EHR",), + ), + ), +) +def search_patients(query: str, limit: int = 10) -> dict: + q = query.lower() + results = [] + for pid, rec in PATIENTS.items(): + if ( + q in pid.lower() + or q in rec["name"].lower() + or any(q in d.lower() for d in rec.get("diagnoses", [])) + ): + results.append({"patient_id": pid, "name": rec["name"]}) + return {"results": results[:limit], "total": len(results)} + + +# ── Tool 4: update_patient_record ──────────────────────────────── + +@server.tool( + name="update_patient_record", + description="Update a patient record field. Writes HIPAA-regulated PHI.", + annotations=_make_annotations( + destructive=True, + trust=TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + outcomes=Outcome.CONSEQUENTIAL, + ), + attribution=("Epic-EHR",), + ), + ), +) +def update_patient_record(patient_id: str, field: str, value: str) -> dict: + patient = PATIENTS.get(patient_id) + if patient is None: + return {"error": f"Patient {patient_id} not found"} + allowed_fields = {"name", "dob", "diagnoses", "medications"} + if field not in allowed_fields: + return {"error": f"Field '{field}' is not updatable. Allowed: {sorted(allowed_fields)}"} + if field in ("diagnoses", "medications"): + # Append to list fields + patient[field].append(value) + else: + patient[field] = value + return { + "patient_id": patient_id, + "field": field, + "new_value": patient[field], + "status": "updated", + "audit_id": _next_audit_id(), + } + + +# ── Tool 5: staff_directory ────────────────────────────────────── + +@server.tool( + name="staff_directory", + description="List clinic staff. Contains PII (names, emails, phone numbers).", + annotations=_make_annotations( + read_only=True, + trust=TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=SimpleDataClass.PII, + ), + ), + ), +) +def staff_directory(department: str = "") -> dict: + staff = [ + {"name": "Dr. Sarah Chen", "email": "s.chen@clinic.org", + "phone": "555-0101", "department": "Cardiology"}, + {"name": "Nurse Mike Johnson", "email": "m.johnson@clinic.org", + "phone": "555-0102", "department": "Emergency"}, + ] + if department: + staff = [s for s in staff if s["department"].lower() == department.lower()] + return {"staff": staff} + + +# ── Tool 6: process_insurance_claim ────────────────────────────── + +@server.tool( + name="process_insurance_claim", + description="Submit an insurance claim. Financial data + patient reference.", + annotations=_make_annotations( + trust=TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=SimpleDataClass.FINANCIAL, + ), + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=SimpleDataClass.FINANCIAL, + outcomes=Outcome.CONSEQUENTIAL, + ), + ), + ), +) +def process_insurance_claim(patient_id: str, amount: float, code: str) -> dict: + if patient_id not in PATIENTS: + return {"error": f"Patient {patient_id} not found"} + claim = { + "claim_id": _next_claim_id(), + "patient_id": patient_id, + "patient_name": PATIENTS[patient_id]["name"], + "amount": amount, + "procedure_code": code, + "status": "submitted", + } + CLAIMS.append(claim) + return claim + + +# ── Tool 7: rotate_api_key ────────────────────────────────────── + +@server.tool( + name="rotate_api_key", + description="Rotate an API key for an external service. Returns credentials.", + annotations=_make_annotations( + trust=TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.SYSTEM, + sensitivity=SimpleDataClass.CREDENTIALS, + ), + ), + ), +) +def rotate_api_key(service: str) -> dict: + return { + "service": service, + "new_key": "sk-REDACTED-FOR-DEMO", + "expires_at": "2025-03-01T00:00:00Z", + } + + +# ── Tool 8: send_notification ──────────────────────────────────── + +@server.tool( + name="send_notification", + description=( + "Send a notification to a patient or external party. " + "Marked with maliciousActivityHint because the message content " + "could contain prompt-injected text. Destination is PUBLIC." + ), + annotations=_make_annotations( + destructive=True, + trust=TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.PUBLIC, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.IRREVERSIBLE, + ), + malicious_activity_hint=True, + ), + ), +) +def send_notification(recipient: str, message: str) -> dict: + notification = { + "message_id": _next_msg_id(), + "recipient": recipient, + "message": message, + "status": "sent", + } + NOTIFICATIONS.append(notification) + return { + "recipient": recipient, + "status": "sent", + "message_id": notification["message_id"], + } + + +# =================================================================== +# Entry point — stdio transport +# =================================================================== + +if __name__ == "__main__": + server.run(transport="stdio") diff --git a/sdk/python/examples/dashboard/app.py b/sdk/python/examples/dashboard/app.py new file mode 100644 index 0000000..ce53e18 --- /dev/null +++ b/sdk/python/examples/dashboard/app.py @@ -0,0 +1,591 @@ +"""SEP-1913 Trust Annotations Test Dashboard. + +Web-based UI for testers to: + 1. View all server tools and their SEP-1913 trust annotations + 2. Run single-client demo scenarios (mcp_client.py equivalent) + 3. Run multi-agent data leak prevention scenarios + 4. Execute policy enforcement tests interactively + 5. View session propagation state in real time + +Usage: + cd /sdk/python + PYTHONPATH=src python examples/dashboard/app.py + +Opens at http://localhost:8913 +""" + +from __future__ import annotations + +import asyncio +import io +import json +import os +import re +import sys +import traceback + +# Ensure src/ is importable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) + +from starlette.applications import Starlette +from starlette.responses import HTMLResponse, JSONResponse +from starlette.routing import Route + +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +from policy import PolicyEngine +from propagate import SessionTracker +from annotate import from_wire, to_wire +from trust_types import ( + ResultAnnotations, + SessionAnnotations, + SimpleDataClass, + Regulated, + TrustAnnotations, + InputMetadata, + ReturnMetadata, + Source, + Destination, + Outcome, + sensitivity_level, + max_sensitivity, +) + +SDK_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +SERVER_SCRIPT = os.path.join(SDK_ROOT, "examples", "_shared", "mcp_server.py") +PYTHON_SRC = os.path.join(SDK_ROOT, "src") + + +# ── Helpers ─────────────────────────────────────────────────────── + +def _strip_ansi(text: str) -> str: + return re.sub(r"\033\[[0-9;]*m", "", text) + + +def _fmt_sensitivity(val) -> str: + if val is None: + return "none" + if isinstance(val, str): + return val + if isinstance(val, dict) and "regulated" in val: + scopes = val["regulated"].get("scopes", []) + return f"regulated({', '.join(scopes)})" + if isinstance(val, list): + return " | ".join(_fmt_sensitivity(v) for v in val) + return str(val) + + +def _wire_to_typed(val): + if isinstance(val, str): + return SimpleDataClass(val) + if isinstance(val, dict) and "regulated" in val: + return Regulated.of(*val["regulated"].get("scopes", [])) + if isinstance(val, list): + r = SimpleDataClass.NONE + for v in val: + r = max_sensitivity(r, _wire_to_typed(v)) + return r + return SimpleDataClass.NONE + + +def _extract_sens(ann: dict): + rm = ann.get("returnMetadata", {}) + im = ann.get("inputMetadata", {}) + return rm.get("sensitivity") or im.get("sensitivity") or "none" + + +def _redact_values(data: dict) -> dict: + out = {} + for k, v in data.items(): + if isinstance(v, dict): + out[k] = _redact_values(v) + elif isinstance(v, list): + out[k] = ["[REDACTED]" if isinstance(i, str) and len(i) > 1 else i for i in v] + elif isinstance(v, str) and k not in ("status", "patient_id", "audit_id", "claim_id"): + out[k] = "[REDACTED]" + else: + out[k] = v + return out + + +def _server_params(): + return StdioServerParameters( + command=sys.executable, + args=[SERVER_SCRIPT], + env={**os.environ, "PYTHONPATH": PYTHON_SRC}, + cwd=SDK_ROOT, + ) + + +# ── API Endpoints ───────────────────────────────────────────────── + +async def api_tools(request): + """List all tools with full SEP-1913 annotations.""" + try: + async with stdio_client(_server_params(), errlog=open(os.devnull, "w")) as (rs, ws): + async with ClientSession(rs, ws) as session: + await session.initialize() + result = await session.list_tools() + tools = [] + for t in result.tools: + ann = t.annotations.model_dump(exclude_none=True) if t.annotations else {} + sens_raw = _extract_sens(ann) + tools.append({ + "name": t.name, + "description": t.description, + "inputSchema": t.inputSchema if hasattr(t, "inputSchema") else {}, + "annotations": ann, + "sensitivity": _fmt_sensitivity(sens_raw), + "source": ann.get("returnMetadata", {}).get("source", "—"), + "destination": ann.get("inputMetadata", {}).get("destination", "—"), + "outcomes": ann.get("inputMetadata", {}).get("outcomes", "—"), + "attribution": ann.get("attribution", []), + "maliciousActivityHint": ann.get("maliciousActivityHint", False), + "readOnlyHint": ann.get("readOnlyHint", False), + "destructiveHint": ann.get("destructiveHint", False), + "openWorldHint": ann.get("openWorldHint", False), + }) + return JSONResponse({"tools": tools}) + except Exception as e: + return JSONResponse({"error": str(e)}, status_code=500) + + +async def api_call_tool(request): + """Call a single tool and return result + annotations.""" + body = await request.json() + tool_name = body.get("tool") + args = body.get("args", {}) + + try: + async with stdio_client(_server_params(), errlog=open(os.devnull, "w")) as (rs, ws): + async with ClientSession(rs, ws) as session: + await session.initialize() + cr = await session.call_tool(tool_name, args) + text = cr.content[0].text if cr.content else "" + try: + result_data = json.loads(text) + except (json.JSONDecodeError, TypeError): + result_data = {"raw": text} + return JSONResponse({"tool": tool_name, "args": args, "result": result_data}) + except Exception as e: + return JSONResponse({"error": str(e)}, status_code=500) + + +async def api_run_single_client(request): + """Run the single-client demo (mcp_client.py equivalent).""" + steps = [] + try: + async with stdio_client(_server_params(), errlog=open(os.devnull, "w")) as (rs, ws): + async with ClientSession(rs, ws) as session: + await session.initialize() + steps.append({"step": "connect", "status": "ok", "message": "Connected to Healthcare Clinic server"}) + + # Step 1: List tools + result = await session.list_tools() + tool_ann = {} + tool_list = [] + for t in result.tools: + ann = t.annotations.model_dump(exclude_none=True) if t.annotations else {} + tool_ann[t.name] = ann + tool_list.append({ + "name": t.name, + "sensitivity": _fmt_sensitivity(_extract_sens(ann)), + "source": ann.get("returnMetadata", {}).get("source", "—"), + "destination": ann.get("inputMetadata", {}).get("destination", "—"), + }) + steps.append({"step": "list_tools", "status": "ok", "tools": tool_list}) + + # Step 2: Policy engine + engine = PolicyEngine(mode="enforce") + for name, ann in tool_ann.items(): + engine.register_tool(name, ann) + steps.append({"step": "policy_init", "status": "ok", "mode": "enforce", "tools_registered": len(tool_ann)}) + + # Step 3: Call tools + track session + tracker = SessionTracker(session_id="dashboard-single-client") + calls = [ + ("health_check", {}), + ("patient_lookup", {"patient_id": "P-12345"}), + ("staff_directory", {"department": "Cardiology"}), + ("process_insurance_claim", {"patient_id": "P-12345", "amount": 1500.0, "code": "99213"}), + ("rotate_api_key", {"service": "lab-integration"}), + ] + + call_results = [] + for tool_name, args in calls: + cr = await session.call_tool(tool_name, args) + text = cr.content[0].text if cr.content else "" + try: + result_data = json.loads(text) + except (json.JSONDecodeError, TypeError): + result_data = {"raw": text} + + ann = tool_ann.get(tool_name, {}) + sens_raw = _extract_sens(ann) + typed_sens = _wire_to_typed(sens_raw) + attribution = tuple(ann.get("attribution", [])) + + result_ann = ResultAnnotations( + sensitivity=typed_sens, + attribution=attribution, + malicious_activity_hint=ann.get("maliciousActivityHint", False), + open_world_hint=ann.get("openWorldHint", False), + ) + tracker.merge(result_ann, tool_name) + + call_results.append({ + "tool": tool_name, + "args": args, + "result": result_data, + "session_sensitivity": str(tracker.session.max_sensitivity), + "session_attribution": sorted(tracker.session.attribution), + }) + + steps.append({"step": "call_tools", "status": "ok", "calls": call_results}) + + # Step 4: Policy enforcement scenarios + scenarios = [ + {"desc": "Forward patient data to internal analytics", "tool": "patient_lookup", "dest": "internal"}, + {"desc": "Forward patient data to PUBLIC endpoint", "tool": "patient_lookup", "dest": "public"}, + {"desc": "Send notification (malicious hint) with HIPAA session", "tool": "send_notification", "dest": "public"}, + {"desc": "Forward credentials externally", "tool": "rotate_api_key", "dest": "public"}, + {"desc": "Read health check (benign)", "tool": "health_check", "dest": None}, + ] + + policy_results = [] + for sc in scenarios: + decision = engine.evaluate( + tool=sc["tool"], action="call", + target_destination=sc["dest"], + session=tracker.session, + ) + policy_results.append({ + "description": sc["desc"], + "tool": sc["tool"], + "destination": sc["dest"] or "any", + "allowed": decision.allowed, + "effect": decision.effect, + "rule": decision.rule, + "reason": decision.reason, + }) + steps.append({"step": "policy_enforcement", "status": "ok", "scenarios": policy_results}) + + # Step 5: Session summary + steps.append({ + "step": "session_summary", + "status": "ok", + "summary": _strip_ansi(tracker.summary()), + "wire_format": tracker.to_wire(), + }) + + return JSONResponse({"status": "ok", "steps": steps}) + except Exception as e: + steps.append({"step": "error", "status": "error", "message": str(e), "traceback": traceback.format_exc()}) + return JSONResponse({"status": "error", "steps": steps}, status_code=500) + + +async def api_run_multi_agent(request): + """Run the multi-agent data leak prevention scenario.""" + results = {"without_annotations": [], "with_annotations": []} + + # ── Part 1: WITHOUT Trust Annotations (simulated) ── + patient_data = { + "patient_id": "P-12345", "name": "Jane Doe", "dob": "1985-03-15", + "ssn": "123-45-6789", "diagnoses": ["Type 2 Diabetes", "Hypertension"], + "medications": ["Metformin 500mg", "Lisinopril 10mg"], + } + + results["without_annotations"] = [ + {"agent": "A", "action": "patient_lookup('P-12345')", "result": "Full PHI returned", + "data_preview": patient_data, "status": "ok", "issue": None}, + {"agent": "A→B", "action": "Forward to Internal Analytics", "result": "PHI forwarded in full", + "data_preview": patient_data, "status": "warning", "issue": "No signal it was HIPAA-regulated"}, + {"agent": "B→C", "action": "Forward to External Reporting (PUBLIC)", + "result": "PHI forwarded in full — HIPAA VIOLATION", + "data_preview": {"name": "Jane Doe", "ssn": "123-45-6789"}, "status": "danger", + "issue": "Regulated data reached external system"}, + {"agent": "B", "action": "Log to analytics pipeline", "result": "PHI in plaintext logs", + "data_preview": {"name": "Jane Doe", "ssn": "123-45-6789"}, "status": "danger", + "issue": "No redaction signal"}, + ] + + # ── Part 2: WITH Trust Annotations (real MCP) ── + try: + async with stdio_client(_server_params(), errlog=open(os.devnull, "w")) as (rs, ws): + async with ClientSession(rs, ws) as session: + await session.initialize() + + tools_result = await session.list_tools() + tool_ann = {} + for t in tools_result.tools: + ann = t.annotations.model_dump(exclude_none=True) if t.annotations else {} + tool_ann[t.name] = ann + + engine = PolicyEngine(mode="enforce") + for name, ann in tool_ann.items(): + engine.register_tool(name, ann) + + tracker_a = SessionTracker(session_id="agent-A") + tracker_b = SessionTracker(session_id="agent-B") + + # Step 1: Agent A calls patient_lookup + cr = await session.call_tool("patient_lookup", {"patient_id": "P-12345"}) + raw = json.loads(cr.content[0].text) if cr.content else {} + + tracker_a.merge( + ResultAnnotations(sensitivity=Regulated.of("HIPAA"), attribution=("Epic-EHR",)), + "patient_lookup", + ) + + results["with_annotations"].append({ + "agent": "A", "action": "patient_lookup('P-12345')", + "result": "PHI returned with trust metadata", + "data_preview": raw, + "trust_metadata": { + "returnMetadata": {"source": "internal", "sensitivity": "regulated(HIPAA)"}, + "attribution": ["Epic-EHR"], + }, + "status": "ok", "issue": None, + }) + + # Step 2: A → B (internal) — ALLOW + decision_ab = engine.evaluate("patient_lookup", action="forward", target_destination="internal") + tracker_b.merge( + ResultAnnotations(sensitivity=Regulated.of("HIPAA"), attribution=("Epic-EHR",)), + "patient_lookup (forwarded)", + ) + + results["with_annotations"].append({ + "agent": "A→B", "action": "Forward to Internal Analytics", + "result": f"{'ALLOWED' if decision_ab.allowed else 'BLOCKED'} — internal destination compatible", + "policy_decision": { + "allowed": decision_ab.allowed, "effect": decision_ab.effect, + "rule": decision_ab.rule, "reason": decision_ab.reason, + }, + "status": "ok", "issue": None, + }) + + # Step 3: B → C (public) — BLOCKED + decision_bc = engine.evaluate("patient_lookup", action="forward", target_destination="public") + + results["with_annotations"].append({ + "agent": "B→C", "action": "Forward to External Reporting (PUBLIC)", + "result": f"{'ALLOWED' if decision_bc.allowed else 'BLOCKED'} — regulated(HIPAA) cannot reach public", + "policy_decision": { + "allowed": decision_bc.allowed, "effect": decision_bc.effect, + "rule": decision_bc.rule, "reason": decision_bc.reason, + }, + "status": "success" if not decision_bc.allowed else "danger", + "issue": None if not decision_bc.allowed else "Should have been blocked!", + }) + + # Step 4: B logs — REDACT + redacted = _redact_values(raw) + results["with_annotations"].append({ + "agent": "B", "action": "Log to analytics pipeline", + "result": "REDACTED — regulated data, plaintext logging blocked", + "data_preview": redacted, + "status": "success", "issue": None, + }) + + results["session_state"] = { + "agent_a": {"sensitivity": str(tracker_a.session.max_sensitivity), + "attribution": sorted(tracker_a.session.attribution)}, + "agent_b": {"sensitivity": str(tracker_b.session.max_sensitivity), + "attribution": sorted(tracker_b.session.attribution)}, + } + + return JSONResponse({"status": "ok", "results": results}) + except Exception as e: + return JSONResponse({"status": "error", "error": str(e), "results": results}, status_code=500) + + +async def api_run_host_scenarios(request): + """Run all host policy enforcement scenarios.""" + scenarios = [] + try: + async with stdio_client(_server_params(), errlog=open(os.devnull, "w")) as (rs, ws): + async with ClientSession(rs, ws) as session: + await session.initialize() + + tools_result = await session.list_tools() + tool_ann = {} + for t in tools_result.tools: + ann = t.annotations.model_dump(exclude_none=True) if t.annotations else {} + tool_ann[t.name] = ann + + engine = PolicyEngine(mode="enforce") + for name, ann in tool_ann.items(): + engine.register_tool(name, ann) + + tracker = SessionTracker(session_id="host-dashboard") + + # Run through all tools to build session state + tool_calls = [ + ("health_check", {}), + ("patient_lookup", {"patient_id": "P-12345"}), + ("search_patients", {"query": "diabetes"}), + ("staff_directory", {}), + ("process_insurance_claim", {"patient_id": "P-12345", "amount": 1500.0, "code": "99213"}), + ("rotate_api_key", {"service": "lab-integration"}), + ] + + for tool_name, args in tool_calls: + cr = await session.call_tool(tool_name, args) + text = cr.content[0].text if cr.content else "" + try: + result_data = json.loads(text) + except (json.JSONDecodeError, TypeError): + result_data = {"raw": text} + + ann = tool_ann.get(tool_name, {}) + typed_sens = _wire_to_typed(_extract_sens(ann)) + tracker.merge( + ResultAnnotations( + sensitivity=typed_sens, + attribution=tuple(ann.get("attribution", [])), + malicious_activity_hint=ann.get("maliciousActivityHint", False), + open_world_hint=ann.get("openWorldHint", False), + ), + tool_name, + ) + + # Policy evaluation scenarios + policy_scenarios = [ + {"desc": "health_check (benign read)", "tool": "health_check", "dest": None, + "category": "Basic"}, + {"desc": "patient_lookup read (HIPAA)", "tool": "patient_lookup", "dest": None, + "category": "HIPAA"}, + {"desc": "Forward patient data → internal", "tool": "patient_lookup", "dest": "internal", + "category": "HIPAA"}, + {"desc": "Forward patient data → public", "tool": "patient_lookup", "dest": "public", + "category": "HIPAA"}, + {"desc": "Forward credentials → public", "tool": "rotate_api_key", "dest": "public", + "category": "Credentials"}, + {"desc": "send_notification (malicious hint)", "tool": "send_notification", "dest": "public", + "category": "Malicious"}, + {"desc": "send_notification → internal", "tool": "send_notification", "dest": "internal", + "category": "Malicious"}, + {"desc": "process_insurance_claim (financial)", "tool": "process_insurance_claim", "dest": None, + "category": "Financial"}, + {"desc": "Forward insurance claim → public", "tool": "process_insurance_claim", "dest": "public", + "category": "Financial"}, + {"desc": "staff_directory (PII) → public", "tool": "staff_directory", "dest": "public", + "category": "PII"}, + ] + + for sc in policy_scenarios: + decision = engine.evaluate( + tool=sc["tool"], action="call", + target_destination=sc["dest"], + session=tracker.session, + ) + scenarios.append({ + "description": sc["desc"], + "category": sc["category"], + "tool": sc["tool"], + "destination": sc["dest"] or "any", + "allowed": decision.allowed, + "effect": decision.effect, + "rule": decision.rule, + "reason": decision.reason, + }) + + return JSONResponse({ + "status": "ok", + "scenarios": scenarios, + "session": { + "sensitivity": str(tracker.session.max_sensitivity), + "attribution": sorted(tracker.session.attribution), + "open_world_hint": tracker.session.open_world_hint, + "malicious_activity_hint": tracker.session.malicious_activity_hint, + "call_count": tracker.call_count, + }, + "wire_format": tracker.to_wire(), + }) + except BaseException as e: + detail = str(e) + if hasattr(e, 'exceptions'): + detail = "; ".join(str(sub) for sub in e.exceptions) + return JSONResponse({"status": "error", "error": detail, "traceback": traceback.format_exc()}, status_code=500) + + +async def api_run_usability_tests(request): + """Run pytest usability scenario tests and return results.""" + import subprocess + test_file = os.path.join(SDK_ROOT, "tests", "test_usability_scenarios.py") + env = {**os.environ, "PYTHONPATH": os.path.join(SDK_ROOT, "src")} + + proc = subprocess.run( + [sys.executable, "-m", "pytest", test_file, "-v", "--tb=short", "--no-header"], + capture_output=True, text=True, env=env, cwd=SDK_ROOT, + timeout=60, + ) + + lines = proc.stdout.strip().split("\n") + tests = [] + for line in lines: + if "::" in line and ("PASSED" in line or "FAILED" in line or "ERROR" in line): + parts = line.strip().split("::") + status = "passed" if "PASSED" in line else "failed" if "FAILED" in line else "error" + test_path = "::".join(parts[1:]).split(" ")[0] if len(parts) > 1 else line + class_name = parts[1] if len(parts) > 1 else "" + test_name = parts[2].split(" ")[0] if len(parts) > 2 else "" + # Extract scenario number from class name + scenario_match = re.search(r"Scenario(\d+)", class_name) + scenario_num = int(scenario_match.group(1)) if scenario_match else 0 + tests.append({ + "scenario": scenario_num, + "class": class_name, + "test": test_name, + "status": status, + "line": _strip_ansi(line.strip()), + }) + + # Summary line + summary = "" + for line in reversed(lines): + if "passed" in line or "failed" in line: + summary = _strip_ansi(line.strip()) + break + + return JSONResponse({ + "status": "ok" if proc.returncode == 0 else "failed", + "exit_code": proc.returncode, + "tests": tests, + "summary": summary, + "stdout": _strip_ansi(proc.stdout), + "stderr": _strip_ansi(proc.stderr), + }) + + +# ── HTML Dashboard ──────────────────────────────────────────────── + +async def homepage(request): + html_path = os.path.join(os.path.dirname(__file__), "index.html") + with open(html_path, "r", encoding="utf-8") as f: + return HTMLResponse(f.read()) + + +# ── App ─────────────────────────────────────────────────────────── + +app = Starlette( + debug=True, + routes=[ + Route("/", homepage), + Route("/api/tools", api_tools), + Route("/api/call-tool", api_call_tool, methods=["POST"]), + Route("/api/run/single-client", api_run_single_client, methods=["POST"]), + Route("/api/run/multi-agent", api_run_multi_agent, methods=["POST"]), + Route("/api/run/host-scenarios", api_run_host_scenarios, methods=["POST"]), + Route("/api/run/usability-tests", api_run_usability_tests, methods=["POST"]) + ], +) + +if __name__ == "__main__": + import uvicorn + port = int(os.environ.get("DASHBOARD_PORT", "8913")) + print(f"\n SEP-1913 Trust Annotations Test Dashboard") + print(f" http://localhost:{port}\n") + uvicorn.run(app, host="127.0.0.1", port=port, log_level="info") diff --git a/sdk/python/examples/dashboard/index.html b/sdk/python/examples/dashboard/index.html new file mode 100644 index 0000000..59dcf1b --- /dev/null +++ b/sdk/python/examples/dashboard/index.html @@ -0,0 +1,788 @@ + + + + + +SEP-1913 Trust Annotations — Test Dashboard + + + + +
+

🛡️ SEP-1913 Trust Annotations

+ Test Dashboard + Checking server... +
+ +
+ +
+
🔧 Tools & Annotations
+
📡 Call Tool
+
▶ Single Client Demo
+
🔀 Multi-Agent Demo
+
🏛️ Host Policy Scenarios
+
📝 Usability Tests
+
+ + +
+
+

Server Tools with SEP-1913 Trust Annotations

+

+ All tools registered on the Healthcare Clinic MCP server. Trust annotations enable + deterministic policy enforcement without relying on the AI model. +

+
Loading tools...
+ +
+ +
+ + +
+
+

Interactive Tool Caller

+

+ Call any server tool directly and see the result with trust annotation context. +

+
+
+ + +
+
+ +
+
+
+
+ + +
+
+

▶ Single Client Demo

+

+ Equivalent to mcp_client.py — connects a single client to the server, + lists tools, calls them sequentially, tracks session propagation, and evaluates policy. +

+
+ Clients: 1  |  + Tools called: 5 (health_check, patient_lookup, staff_directory, process_insurance_claim, rotate_api_key)  |  + Policy scenarios: 5 +
+ +
+
+
+ + +
+
+

🔀 Multi-Agent Data Leak Prevention

+

+ Equivalent to data_leak_prevention.py — three agents cooperate: + Agent A (front-desk) → Agent B (internal analytics) → Agent C (external reporting). + Compares behavior WITHOUT vs WITH trust annotations. +

+
+ Agents: 3 (A: front-desk, B: analytics, C: external)  |  + Patient: Jane Doe (P-12345, HIPAA-regulated PHI)  |  + Scenarios: Side-by-side comparison +
+ +
+
+
+ + +
+
+

🏛️ Host Policy Enforcement Scenarios

+

+ Equivalent to host.py — runs all 10 policy enforcement scenarios. + Calls every tool to build up session state, then evaluates each scenario + against the policy engine in ENFORCE mode. +

+
+ Mode: ENFORCE (blocks violating calls)  |  + Scenarios: 10 across 5 categories (Basic, HIPAA, Credentials, Malicious, Financial, PII) +
+ +
+
+
+ + +
+
+

📝 Usability Scenario Tests

+

+ Runs test_usability_scenarios.py — 10 scenarios testing whether + developers reading SEP-1913 can correctly predict trust annotation behavior. + Each scenario has oracle tests that reveal the correct answers. +

+
+ Scenarios: 10  |  + Categories: Classification, Propagation, Policy, Spec Gaps  |  + Expected difficulty: Easy → Hard +
+ +
+
+
+ +
+ + + + diff --git a/sdk/python/examples/dashboard/start_dashboard.ps1 b/sdk/python/examples/dashboard/start_dashboard.ps1 new file mode 100644 index 0000000..eb7e56c --- /dev/null +++ b/sdk/python/examples/dashboard/start_dashboard.ps1 @@ -0,0 +1,132 @@ +<# +.SYNOPSIS + Launch the SEP-1913 Trust Annotations Test Dashboard. + +.DESCRIPTION + - Activates the project venv + - Checks that required packages are installed + - Kills any process already on port 8913 + - Verifies the MCP server script can start (stdio health check) + - Launches the dashboard at http://localhost:8913 + +.EXAMPLE + .\start_dashboard.ps1 + .\start_dashboard.ps1 -Port 9000 +#> + +param( + [int]$Port = 8913 +) + +$ErrorActionPreference = "Stop" +$sdkRoot = Split-Path -Parent (Split-Path -Parent $PSScriptRoot) + +# If run from sdk root directly, adjust +if (-not (Test-Path "$sdkRoot\examples\dashboard\app.py")) { + $sdkRoot = $PSScriptRoot | Split-Path -Parent | Split-Path -Parent + if (-not (Test-Path "$sdkRoot\examples\dashboard\app.py")) { + $sdkRoot = Get-Location + } +} + +$venvPython = Join-Path $sdkRoot ".venv\Scripts\python.exe" +# Fall back to repo-level venv +if (-not (Test-Path $venvPython)) { + $repoRoot = Split-Path -Parent (Split-Path -Parent $sdkRoot) + $venvPython = Join-Path $repoRoot ".venv\Scripts\python.exe" +} +$dashboardApp = Join-Path $sdkRoot "examples\dashboard\app.py" +$serverScript = Join-Path $sdkRoot "examples\_shared\mcp_server.py" +$pythonSrc = Join-Path $sdkRoot "src" + +Write-Host "" +Write-Host " SEP-1913 Trust Annotations Test Dashboard" -ForegroundColor Cyan +Write-Host " ==========================================" -ForegroundColor Cyan +Write-Host "" + +# ── Step 1: Check venv ────────────────────────────────────── +Write-Host "[1/5] Checking Python virtual environment..." -ForegroundColor Yellow +if (-not (Test-Path $venvPython)) { + Write-Host " ERROR: Virtual environment not found" -ForegroundColor Red + Write-Host " Run: python -m venv .venv && .venv\Scripts\pip install -e '.[dev]' mcp starlette uvicorn" -ForegroundColor Red + exit 1 +} +$pyVersion = & $venvPython --version 2>&1 +Write-Host " OK: $pyVersion" -ForegroundColor Green + +# ── Step 2: Check required packages ───────────────────────── +Write-Host "[2/5] Checking required packages..." -ForegroundColor Yellow +$missing = @() +foreach ($pkg in @("mcp", "starlette", "uvicorn")) { + $check = & $venvPython -c "import $pkg" 2>&1 + if ($LASTEXITCODE -ne 0) { $missing += $pkg } +} +if ($missing.Count -gt 0) { + Write-Host " Missing packages: $($missing -join ', ')" -ForegroundColor Red + Write-Host " Installing..." -ForegroundColor Yellow + & $venvPython -m pip install $missing --quiet + if ($LASTEXITCODE -ne 0) { + Write-Host " ERROR: Failed to install packages" -ForegroundColor Red + exit 1 + } +} +Write-Host " OK: mcp, starlette, uvicorn" -ForegroundColor Green + +# ── Step 3: Verify MCP server can start ───────────────────── +Write-Host "[3/5] Verifying MCP server health..." -ForegroundColor Yellow +if (-not (Test-Path $serverScript)) { + Write-Host " ERROR: Server script not found: $serverScript" -ForegroundColor Red + exit 1 +} + +# Quick syntax check on the server script +$syntaxCheck = & $venvPython -c " +import sys, os +sys.path.insert(0, r'$pythonSrc') +os.environ['PYTHONPATH'] = r'$pythonSrc' +import importlib.util +spec = importlib.util.spec_from_file_location('mcp_server', r'$serverScript') +mod = importlib.util.module_from_spec(spec) +print('Server module loadable') +" 2>&1 +if ($LASTEXITCODE -ne 0) { + Write-Host " ERROR: MCP server has import errors:" -ForegroundColor Red + Write-Host " $syntaxCheck" -ForegroundColor Red + exit 1 +} +Write-Host " OK: Server script is valid and importable" -ForegroundColor Green + +# ── Step 4: Free the port ──────────────────────────────────── +Write-Host "[4/5] Checking port $Port..." -ForegroundColor Yellow +try { + $existing = Get-NetTCPConnection -LocalPort $Port -ErrorAction Stop +} catch { + $existing = $null +} +if ($existing) { + $pids = $existing | Select-Object -ExpandProperty OwningProcess -Unique + foreach ($procId in $pids) { + try { + $proc = Get-Process -Id $procId -ErrorAction Stop + Write-Host " Stopping existing process on port ${Port}: $($proc.ProcessName) (PID $procId)" -ForegroundColor Yellow + Stop-Process -Id $procId -Force + } catch {} + } + Start-Sleep -Seconds 1 + Write-Host " OK: Port $Port freed" -ForegroundColor Green +} else { + Write-Host " OK: Port $Port is available" -ForegroundColor Green +} + +# ── Step 5: Launch dashboard ───────────────────────────────── +Write-Host "[5/5] Starting dashboard..." -ForegroundColor Yellow +Write-Host "" + +$env:PYTHONPATH = $pythonSrc +$env:DASHBOARD_PORT = $Port + +Write-Host " Dashboard URL: http://localhost:$Port" -ForegroundColor Cyan +Write-Host " Press Ctrl+C to stop" -ForegroundColor DarkGray +Write-Host "" + +& $venvPython $dashboardApp diff --git a/sdk/python/examples/dashboard/start_dashboard.sh b/sdk/python/examples/dashboard/start_dashboard.sh new file mode 100644 index 0000000..ea25ac1 --- /dev/null +++ b/sdk/python/examples/dashboard/start_dashboard.sh @@ -0,0 +1,118 @@ +#!/usr/bin/env bash +# +# Launch the SEP-1913 Trust Annotations Test Dashboard. +# +# - Checks for a Python venv +# - Verifies required packages are installed +# - Kills any process already on the dashboard port +# - Verifies the MCP server script can load +# - Launches the dashboard at http://localhost:8913 +# +# Usage: +# cd sdk/python/examples/dashboard +# chmod +x start_dashboard.sh +# ./start_dashboard.sh +# ./start_dashboard.sh 9000 + +set -euo pipefail + +PORT="${1:-8913}" + +# ── Resolve paths ──────────────────────────────────────────── +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SDK_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" + +DASHBOARD_APP="$SDK_ROOT/examples/dashboard/app.py" +SERVER_SCRIPT="$SDK_ROOT/examples/_shared/mcp_server.py" +PYTHON_SRC="$SDK_ROOT/src" + +# Locate venv python: sdk-level first, then repo-level +VENV_PYTHON="$SDK_ROOT/.venv/bin/python" +if [ ! -f "$VENV_PYTHON" ]; then + REPO_ROOT="$(cd "$SDK_ROOT/../.." && pwd)" + VENV_PYTHON="$REPO_ROOT/.venv/bin/python" +fi + +# ANSI colors +CYAN='\033[96m' +YELLOW='\033[93m' +GREEN='\033[92m' +RED='\033[91m' +DIM='\033[2m' +RST='\033[0m' + +echo "" +echo -e " ${CYAN}SEP-1913 Trust Annotations Test Dashboard${RST}" +echo -e " ${CYAN}==========================================${RST}" +echo "" + +# ── Step 1: Check venv ────────────────────────────────────── +echo -e "${YELLOW}[1/5] Checking Python virtual environment...${RST}" +if [ ! -f "$VENV_PYTHON" ]; then + echo -e " ${RED}ERROR: Virtual environment not found${RST}" + echo -e " ${RED}Run: python3 -m venv .venv && .venv/bin/pip install -e '.[dev]' mcp starlette uvicorn${RST}" + exit 1 +fi +PY_VERSION=$("$VENV_PYTHON" --version 2>&1) +echo -e " ${GREEN}OK: $PY_VERSION${RST}" + +# ── Step 2: Check required packages ───────────────────────── +echo -e "${YELLOW}[2/5] Checking required packages...${RST}" +MISSING=() +for pkg in mcp starlette uvicorn; do + if ! "$VENV_PYTHON" -c "import $pkg" 2>/dev/null; then + MISSING+=("$pkg") + fi +done +if [ ${#MISSING[@]} -gt 0 ]; then + echo -e " ${RED}Missing packages: ${MISSING[*]}${RST}" + echo -e " ${YELLOW}Installing...${RST}" + "$VENV_PYTHON" -m pip install "${MISSING[@]}" --quiet +fi +echo -e " ${GREEN}OK: mcp, starlette, uvicorn${RST}" + +# ── Step 3: Verify MCP server can start ───────────────────── +echo -e "${YELLOW}[3/5] Verifying MCP server health...${RST}" +if [ ! -f "$SERVER_SCRIPT" ]; then + echo -e " ${RED}ERROR: Server script not found: $SERVER_SCRIPT${RST}" + exit 1 +fi + +if ! "$VENV_PYTHON" -c " +import sys, os +sys.path.insert(0, '$PYTHON_SRC') +os.environ['PYTHONPATH'] = '$PYTHON_SRC' +import importlib.util +spec = importlib.util.spec_from_file_location('mcp_server', '$SERVER_SCRIPT') +mod = importlib.util.module_from_spec(spec) +print('Server module loadable') +" 2>&1; then + echo -e " ${RED}ERROR: MCP server has import errors${RST}" + exit 1 +fi +echo -e " ${GREEN}OK: Server script is valid and importable${RST}" + +# ── Step 4: Free the port ──────────────────────────────────── +echo -e "${YELLOW}[4/5] Checking port $PORT...${RST}" +PIDS=$(lsof -ti :"$PORT" 2>/dev/null || true) +if [ -n "$PIDS" ]; then + echo -e " ${YELLOW}Stopping existing process(es) on port $PORT: $PIDS${RST}" + echo "$PIDS" | xargs kill -9 2>/dev/null || true + sleep 1 + echo -e " ${GREEN}OK: Port $PORT freed${RST}" +else + echo -e " ${GREEN}OK: Port $PORT is available${RST}" +fi + +# ── Step 5: Launch dashboard ───────────────────────────────── +echo -e "${YELLOW}[5/5] Starting dashboard...${RST}" +echo "" + +export PYTHONPATH="$PYTHON_SRC" +export DASHBOARD_PORT="$PORT" + +echo -e " ${CYAN}Dashboard URL: http://localhost:$PORT${RST}" +echo -e " ${DIM}Press Ctrl+C to stop${RST}" +echo "" + +exec "$VENV_PYTHON" "$DASHBOARD_APP" diff --git a/sdk/python/examples/healthcare/README.md b/sdk/python/examples/healthcare/README.md new file mode 100644 index 0000000..a284a97 --- /dev/null +++ b/sdk/python/examples/healthcare/README.md @@ -0,0 +1,24 @@ +# UC-1: Healthcare Data Protection + +Demonstrates HIPAA-compliant data handling using SEP-1913 trust annotations. + +## What it shows + +- **Data classification**: Tools returning PHI are tagged `regulated(HIPAA)`, distinguishable from benign tools +- **Policy enforcement**: ALLOW / BLOCK / ESCALATE / REDACT decisions based on annotations +- **Destination constraints**: Patient data blocked from reaching external/public destinations + +## Scripts + +| Script | Description | +|--------|-------------| +| `mcp_client.py` | Connects to the server, calls all tools, applies policy enforcement | +| `host.py` | Policy-enforcing host with 8 scenarios (HIPAA forwarding, credential blocking, session escalation) | + +## Usage + +```bash +cd sdk/python +PYTHONPATH=src python examples/healthcare/mcp_client.py +PYTHONPATH=src python examples/healthcare/host.py +``` diff --git a/sdk/python/examples/healthcare/host.py b/sdk/python/examples/healthcare/host.py new file mode 100644 index 0000000..e9e547e --- /dev/null +++ b/sdk/python/examples/healthcare/host.py @@ -0,0 +1,332 @@ +"""D8: Policy-Enforcing Host — connects to Healthcare Clinic MCP server +and demonstrates SEP-1913 trust annotation enforcement in a real host. + +This is the "smart host" that: + 1. Connects to the Healthcare Clinic server via stdio + 2. Calls tools/list and extracts SEP-1913 annotations + 3. Registers every tool with a PolicyEngine + 4. Calls tools, tracks session propagation, and enforces policy + 5. Shows ALLOW / BLOCK / ESCALATE / REDACT decisions + +Usage: + cd /sdk/python + PYTHONPATH=src python examples/healthcare/host.py +""" + +from __future__ import annotations + +import asyncio +import io +import json +import os +import re +import sys + +# Force UTF-8 output on Windows to support Unicode symbols +if sys.stdout.encoding != "utf-8": + sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace") +if sys.stderr.encoding != "utf-8": + sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace") + +# Ensure src/ is importable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) + +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +from policy import PolicyEngine +from propagate import SessionTracker +from annotate import from_wire +from trust_types import ( + ResultAnnotations, + SimpleDataClass, + Regulated, + sensitivity_level, + max_sensitivity, +) + + +# ── ANSI helpers ────────────────────────────────────────────────── + +class C: + RST = "\033[0m" + B = "\033[1m" + DIM = "\033[2m" + R = "\033[91m" + G = "\033[92m" + Y = "\033[93m" + BL = "\033[94m" + CY = "\033[96m" + + +def _tag(label: str, color: str, width: int = 8) -> str: + return f"{color}[{label:^{width}}]{C.RST}" + +HOST = lambda m: f"{_tag('HOST', C.BL)} {m}" +POLICY = lambda m: f"{_tag('POLICY', C.CY)} {m}" +LOG = lambda m: f"{_tag('LOG', C.DIM, 8)} {m}" +OK = lambda m: f"{C.G}✓ ALLOW{C.RST} {m}" +BLOCK = lambda m: f"{C.R}✗ BLOCK{C.RST} {m}" +ESCAL = lambda m: f"{C.Y}⚠ ESCALATE{C.RST} {m}" +REDACT = lambda m: f"{C.Y}◉ REDACT{C.RST} {m}" + + +def _fmt_sens(val) -> str: + if val is None: + return "none" + if isinstance(val, str): + return val + if isinstance(val, dict) and "regulated" in val: + scopes = val["regulated"].get("scopes", []) + return f"regulated({','.join(scopes)})" + if isinstance(val, list): + return "|".join(_fmt_sens(v) for v in val) + return str(val) + + +def _wire_to_typed(val): + if isinstance(val, str): + return SimpleDataClass(val) + if isinstance(val, dict) and "regulated" in val: + return Regulated.of(*val["regulated"].get("scopes", [])) + if isinstance(val, list): + r = SimpleDataClass.NONE + for v in val: + r = max_sensitivity(r, _wire_to_typed(v)) + return r + return SimpleDataClass.NONE + + +def _extract_sens(ann: dict): + rm = ann.get("returnMetadata", {}) + im = ann.get("inputMetadata", {}) + return rm.get("sensitivity") or im.get("sensitivity") or "none" + + +def _redact_values(data: dict) -> dict: + """Replace string values that look like PII/PHI with [REDACTED].""" + out = {} + for k, v in data.items(): + if isinstance(v, dict): + out[k] = _redact_values(v) + elif isinstance(v, list): + out[k] = ["[REDACTED]" if isinstance(i, str) and len(i) > 1 else i for i in v] + elif isinstance(v, str) and k not in ("status", "patient_id", "audit_id", "claim_id"): + out[k] = "[REDACTED]" + else: + out[k] = v + return out + + +# ── Main ────────────────────────────────────────────────────────── + +async def run_host(): + sdk_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + server_script = os.path.join(sdk_root, "examples", "_shared", "mcp_server.py") + python_src = os.path.join(sdk_root, "src") + + server_params = StdioServerParameters( + command=sys.executable, + args=[server_script], + env={**os.environ, "PYTHONPATH": python_src}, + cwd=sdk_root, + ) + + print() + print("=" * 72) + print(f" {C.B}{C.BL}D8: Policy-Enforcing Host — SEP-1913 Demo{C.RST}") + print("=" * 72) + print() + + async with stdio_client(server_params, errlog=open(os.devnull, "w")) as (rs, ws): + async with ClientSession(rs, ws) as session: + await session.initialize() + print(HOST("Connected to Healthcare Clinic server")) + + # ── tools/list ──────────────────────────────────── + result = await session.list_tools() + tools = result.tools + tool_ann: dict[str, dict] = {} + for t in tools: + ann = t.annotations.model_dump(exclude_none=True) if t.annotations else {} + tool_ann[t.name] = ann + + print(HOST(f"Loaded {len(tools)} tools with SEP-1913 trust annotations")) + + # ── Policy engine ───────────────────────────────── + engine = PolicyEngine(mode="enforce") + for name, ann in tool_ann.items(): + engine.register_tool(name, ann) + print(HOST(f"Policy mode: {C.R}ENFORCE{C.RST}")) + print() + + tracker = SessionTracker(session_id="host-demo") + + # ═══════════════════════════════════════════════════ + # Scenario 1: health_check (benign) + # ═══════════════════════════════════════════════════ + print(HOST(f"Calling {C.B}health_check{C.RST}...")) + ann = tool_ann["health_check"] + sens = _fmt_sens(_extract_sens(ann)) + source = ann.get("returnMetadata", {}).get("source", "system") + decision = engine.evaluate("health_check") + print(POLICY(f"tool=health_check sensitivity={sens} source={source} → {OK('')}")) + + cr = await session.call_tool("health_check") + text = cr.content[0].text if cr.content else "" + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.NONE), "health_check") + print(LOG(f'result={text}')) + print() + + # ═══════════════════════════════════════════════════ + # Scenario 2: patient_lookup — HIPAA read + REDACT log + # ═══════════════════════════════════════════════════ + print(HOST(f'Calling {C.B}patient_lookup{C.RST}(patient_id="P-12345")')) + ann = tool_ann["patient_lookup"] + sens = _fmt_sens(_extract_sens(ann)) + source = ann.get("returnMetadata", {}).get("source", "—") + decision = engine.evaluate("patient_lookup") + print(POLICY(f"tool=patient_lookup sensitivity={sens} source={source} → {OK('(read)')}")) + + cr = await session.call_tool("patient_lookup", {"patient_id": "P-12345"}) + raw_result = json.loads(cr.content[0].text) if cr.content else {} + + # Merge into session + tracker.merge( + ResultAnnotations( + sensitivity=Regulated.of("HIPAA"), + attribution=("Epic-EHR",), + ), + "patient_lookup", + ) + + # Policy: log action → REDACT (HIPAA data) + decision_log = engine.evaluate("patient_lookup", action="call", target_destination=None) + print(POLICY(f"action=log → {REDACT('regulated data, plaintext logging blocked')}")) + redacted = _redact_values(raw_result) + print(LOG(f"result={json.dumps(redacted)}")) + print() + + # ═══════════════════════════════════════════════════ + # Scenario 3: Forward patient data to internal (ALLOW) + # ═══════════════════════════════════════════════════ + print(HOST(f"Attempting to forward patient_lookup result to {C.B}Agent B (internal analytics){C.RST}...")) + decision = engine.evaluate( + "patient_lookup", action="forward", target_destination="internal", + ) + if decision.allowed: + print(POLICY(f"action=forward target=internal → {OK('internal destination compatible')}")) + else: + print(POLICY(f"action=forward target=internal → {BLOCK(decision.reason)}")) + print() + + # ═══════════════════════════════════════════════════ + # Scenario 4: Forward patient data to public (BLOCK!) + # ═══════════════════════════════════════════════════ + print(HOST(f"Attempting to forward patient_lookup result to {C.B}Agent C (external reporting){C.RST}...")) + decision = engine.evaluate( + "patient_lookup", action="forward", target_destination="public", + ) + if not decision.allowed: + print(POLICY(f"action=forward target=public → {BLOCK(decision.reason)}")) + print(HOST(f"{C.R}✗ BLOCKED{C.RST} — regulated(HIPAA) → public destination")) + else: + print(POLICY(f"action=forward target=public → {OK('')}")) + print() + + # ═══════════════════════════════════════════════════ + # Scenario 5: search_patients — open world propagation + # ═══════════════════════════════════════════════════ + print(HOST(f'Calling {C.B}search_patients{C.RST}(query="diabetes treatment")')) + + cr = await session.call_tool("search_patients", {"query": "diabetes treatment"}) + text = cr.content[0].text if cr.content else "" + + ann = tool_ann["search_patients"] + source = ann.get("returnMetadata", {}).get("source", "—") + print(POLICY(f"tool=search_patients source={source} → {OK('(read)')}")) + + # Merge — this has HIPAA too, and we add attribution + tracker.merge( + ResultAnnotations( + sensitivity=Regulated.of("HIPAA"), + attribution=("Epic-EHR",), + ), + "search_patients", + ) + + # Preview + try: + preview = json.dumps(json.loads(text), indent=None) + if len(preview) > 70: + preview = preview[:67] + "..." + except Exception: + preview = text[:70] + print(LOG(f'result="{preview}" attribution=["Epic-EHR"]')) + print() + + # ═══════════════════════════════════════════════════ + # Scenario 6: send_notification (malicious hint) + # ═══════════════════════════════════════════════════ + print(HOST(f"Attempting to {C.B}send_notification{C.RST} to external recipient...")) + decision = engine.evaluate( + "send_notification", action="call", target_destination="public", + session=tracker.session, + ) + if not decision.allowed: + print(POLICY(f"maliciousActivityHint=true → {BLOCK(decision.reason)}")) + print(HOST(f"{C.R}✗ BLOCKED{C.RST} — tool flagged for potential malicious activity")) + print() + + # ═══════════════════════════════════════════════════ + # Scenario 7: Session escalation — PII+ to public + # ═══════════════════════════════════════════════════ + print(HOST(f"Attempting to call external reporting tool with HIPAA-tainted session...")) + decision = engine.evaluate( + "health_check", + action="call", + target_destination="public", + session=tracker.session, + ) + if decision.effect in ("escalate", "block") and not decision.allowed: + print(POLICY(f"session.max_sensitivity=regulated(HIPAA) + dest=public → " + f"{ESCAL(decision.reason)}")) + print(HOST(f"{C.Y}⚠ ESCALATE{C.RST} — regulated data in session, requires user confirmation")) + elif decision.allowed: + print(POLICY(f"→ {OK('')}")) + print() + + # ═══════════════════════════════════════════════════ + # Scenario 8: Credentials to public (BLOCK) + # ═══════════════════════════════════════════════════ + print(HOST(f"Attempting to forward {C.B}rotate_api_key{C.RST} result to public API...")) + decision = engine.evaluate( + "rotate_api_key", action="forward", target_destination="public", + ) + if not decision.allowed: + print(POLICY(f"sensitivity=credentials target=public → {BLOCK(decision.reason)}")) + print(HOST(f"{C.R}✗ BLOCKED{C.RST} — credentials cannot reach public endpoint")) + print() + + # ═══════════════════════════════════════════════════ + # Session summary + # ═══════════════════════════════════════════════════ + print("─" * 72) + print(HOST("Session summary:")) + print() + print(tracker.summary()) + print() + print("─" * 72) + print(HOST("Wire format (for outbound request annotations):")) + print() + print(json.dumps(tracker.to_wire(), indent=2)) + print() + + print("=" * 72) + print(f" {C.B}{C.BL}Host demo complete — all enforcement used real MCP transport.{C.RST}") + print("=" * 72) + print() + + +if __name__ == "__main__": + asyncio.run(run_host()) diff --git a/sdk/python/examples/healthcare/mcp_client.py b/sdk/python/examples/healthcare/mcp_client.py new file mode 100644 index 0000000..5b124f2 --- /dev/null +++ b/sdk/python/examples/healthcare/mcp_client.py @@ -0,0 +1,339 @@ +"""MCP Client — connects to the Healthcare Clinic server and demonstrates +SEP-1913 trust annotations in action. + +Performs: + 1. Connects to the server via stdio + 2. Lists all tools and extracts SEP-1913 annotations + 3. Calls each tool and tracks session propagation + 4. Applies policy enforcement (block/allow/escalate) + 5. Prints a full session summary + +Usage: + cd /sdk/python + PYTHONPATH=src python examples/healthcare/mcp_client.py +""" + +from __future__ import annotations + +import asyncio +import json +import os +import sys + +# Ensure src/ is importable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) + +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +from policy import PolicyEngine +from propagate import SessionTracker +from annotate import from_wire +from trust_types import ( + ResultAnnotations, + SimpleDataClass, + Regulated, + sensitivity_level, + max_sensitivity, +) + + +# ── Helpers ─────────────────────────────────────────────────────── + +def _format_sensitivity(val) -> str: + """Pretty-print a sensitivity value from wire format.""" + if val is None: + return "none" + if isinstance(val, str): + return val + if isinstance(val, dict) and "regulated" in val: + scopes = val["regulated"].get("scopes", []) + return f"regulated({', '.join(scopes)})" + if isinstance(val, list): + return " | ".join(_format_sensitivity(v) for v in val) + return str(val) + + +def _extract_sensitivity_from_wire(ann_dict: dict): + """Extract the highest sensitivity DataClass from a wire-format annotations dict.""" + rm = ann_dict.get("returnMetadata", {}) + im = ann_dict.get("inputMetadata", {}) + sens = rm.get("sensitivity") or im.get("sensitivity") or "none" + return sens + + +def _wire_sensitivity_to_typed(val): + """Convert wire sensitivity to typed DataClass.""" + if isinstance(val, str): + return SimpleDataClass(val) + if isinstance(val, dict) and "regulated" in val: + scopes = val["regulated"].get("scopes", []) + return Regulated.of(*scopes) + if isinstance(val, list): + result = SimpleDataClass.NONE + for v in val: + result = max_sensitivity(result, _wire_sensitivity_to_typed(v)) + return result + return SimpleDataClass.NONE + + +# ── Colors (ANSI) ──────────────────────────────────────────────── + +class C: + RESET = "\033[0m" + BOLD = "\033[1m" + RED = "\033[91m" + GREEN = "\033[92m" + YELLOW = "\033[93m" + BLUE = "\033[94m" + CYAN = "\033[96m" + DIM = "\033[2m" + + +def _ok(msg: str) -> str: + return f"{C.GREEN}✓ {msg}{C.RESET}" + +def _blocked(msg: str) -> str: + return f"{C.RED}✗ BLOCKED{C.RESET} — {msg}" + +def _escalate(msg: str) -> str: + return f"{C.YELLOW}⚠ ESCALATE{C.RESET} — {msg}" + +def _warn(msg: str) -> str: + return f"{C.YELLOW}⚠ WARN{C.RESET} — {msg}" + +def _info(msg: str) -> str: + return f"{C.CYAN}ℹ {msg}{C.RESET}" + +def _header(msg: str) -> str: + return f"{C.BOLD}{C.BLUE}{msg}{C.RESET}" + + +# ── Main client logic ──────────────────────────────────────────── + +async def run_client(): + sdk_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + server_script = os.path.join(sdk_root, "examples", "_shared", "mcp_server.py") + python_src = os.path.join(sdk_root, "src") + + # Build env with PYTHONPATH so server can import from src/ + env = {**os.environ, "PYTHONPATH": python_src} + + server_params = StdioServerParameters( + command=sys.executable, + args=[server_script], + env=env, + cwd=sdk_root, + ) + + print() + print("=" * 70) + print(_header(" MCP Client — SEP-1913 Trust Annotations Demo")) + print("=" * 70) + print() + + async with stdio_client(server_params, errlog=open(os.devnull, "w")) as (read_stream, write_stream): + async with ClientSession(read_stream, write_stream) as session: + # Initialize the session + await session.initialize() + print(_info("Connected to Healthcare Clinic server")) + print() + + # ── Step 1: List tools ──────────────────────────── + print(_header("─── Step 1: tools/list — Discover tools & SEP-1913 annotations ───")) + print() + + result = await session.list_tools() + tools = result.tools + + print(f" Found {len(tools)} tools:\n") + tool_annotations: dict[str, dict] = {} + + for tool in tools: + ann_dict = {} + if tool.annotations: + ann_dict = tool.annotations.model_dump(exclude_none=True) + tool_annotations[tool.name] = ann_dict + + sens_raw = _extract_sensitivity_from_wire(ann_dict) + sens_str = _format_sensitivity(sens_raw) + source = ann_dict.get("returnMetadata", {}).get("source", "—") + dest = ann_dict.get("inputMetadata", {}).get("destination", "—") + malicious = ann_dict.get("maliciousActivityHint", False) + + flags = [] + if malicious: + flags.append(f"{C.RED}⚡malicious{C.RESET}") + if ann_dict.get("readOnlyHint"): + flags.append("readonly") + if ann_dict.get("destructiveHint"): + flags.append(f"{C.YELLOW}destructive{C.RESET}") + + flag_str = f" [{', '.join(flags)}]" if flags else "" + print(f" {C.BOLD}{tool.name:30s}{C.RESET} " + f"sensitivity={sens_str:20s} source={str(source):18s} " + f"dest={str(dest)}{flag_str}") + + print() + + # ── Step 2: Set up policy engine ────────────────── + print(_header("─── Step 2: Initialize policy engine (enforce mode) ───")) + print() + + engine = PolicyEngine(mode="enforce") + for name, ann in tool_annotations.items(): + engine.register_tool(name, ann) + print(f" Registered {len(tool_annotations)} tools with policy engine") + print(f" Mode: {C.RED}ENFORCE{C.RESET} (blocks violating calls)") + print() + + # ── Step 3: Call tools + track session ──────────── + print(_header("─── Step 3: Call tools & track session propagation ───")) + print() + + tracker = SessionTracker(session_id="client-demo-001") + + # Call sequence + calls = [ + ("health_check", {}), + ("patient_lookup", {"patient_id": "P-12345"}), + ("staff_directory", {"department": "Cardiology"}), + ("process_insurance_claim", {"patient_id": "P-12345", "amount": 1500.0, "code": "99213"}), + ("rotate_api_key", {"service": "lab-integration"}), + ] + + for tool_name, args in calls: + print(f" {C.BOLD}Calling {tool_name}{C.RESET}({_format_args(args)})...") + + call_result = await session.call_tool(tool_name, args) + + # Extract text content from result + content_text = "" + for content in call_result.content: + if hasattr(content, "text"): + content_text = content.text + break + + # Parse the result to show a preview + try: + result_data = json.loads(content_text) + preview = json.dumps(result_data, indent=None) + if len(preview) > 80: + preview = preview[:77] + "..." + except (json.JSONDecodeError, TypeError): + preview = content_text[:80] if content_text else "(empty)" + + print(f" Result: {C.DIM}{preview}{C.RESET}") + + # Merge into session tracker + ann = tool_annotations.get(tool_name, {}) + sens_raw = _extract_sensitivity_from_wire(ann) + typed_sens = _wire_sensitivity_to_typed(sens_raw) + attribution = tuple(ann.get("attribution", [])) + + result_ann = ResultAnnotations( + sensitivity=typed_sens, + attribution=attribution, + malicious_activity_hint=ann.get("maliciousActivityHint", False), + open_world_hint=ann.get("openWorldHint", False), + ) + tracker.merge(result_ann, tool_name) + + print(f" Session: max_sensitivity={tracker.session.max_sensitivity}, " + f"attribution={sorted(tracker.session.attribution)}") + print() + + # ── Step 4: Policy enforcement scenarios ────────── + print(_header("─── Step 4: Policy enforcement scenarios ───")) + print() + + scenarios = [ + { + "desc": "Forward patient data to internal analytics", + "tool": "patient_lookup", + "action": "call", + "dest": "internal", + }, + { + "desc": "Forward patient data to PUBLIC endpoint", + "tool": "patient_lookup", + "action": "call", + "dest": "public", + }, + { + "desc": "Send notification (malicious hint) with HIPAA session", + "tool": "send_notification", + "action": "call", + "dest": "public", + }, + { + "desc": "Forward credentials externally", + "tool": "rotate_api_key", + "action": "call", + "dest": "public", + }, + { + "desc": "Read health check (benign)", + "tool": "health_check", + "action": "call", + "dest": None, + }, + ] + + for scenario in scenarios: + decision = engine.evaluate( + tool=scenario["tool"], + action=scenario["action"], + target_destination=scenario["dest"], + session=tracker.session, + ) + + status_str = "" + if not decision.allowed: + status_str = _blocked(decision.reason) + elif decision.effect == "escalate": + status_str = _escalate(decision.reason) + elif decision.effect == "warn": + status_str = _warn(decision.reason) + else: + status_str = _ok("ALLOWED") + + dest_str = scenario['dest'] or 'any' + print(f" Scenario: {scenario['desc']}") + print(f" {scenario['tool']} → {dest_str}") + print(f" {status_str}") + print() + + # ── Step 5: Session summary ─────────────────────── + print(_header("─── Step 5: Session summary ───")) + print() + print(tracker.summary()) + print() + + # ── Step 6: Wire format ─────────────────────────── + print(_header("─── Step 6: Wire format (session state for outbound propagation) ───")) + print() + print(json.dumps(tracker.to_wire(), indent=2)) + print() + + print("=" * 70) + print(_header(" Demo complete — all calls used real MCP stdio transport.")) + print("=" * 70) + print() + + +def _format_args(args: dict) -> str: + """Format call arguments for display.""" + if not args: + return "" + parts = [] + for k, v in args.items(): + if isinstance(v, str): + parts.append(f'{k}="{v}"') + else: + parts.append(f"{k}={v}") + return ", ".join(parts) + + +if __name__ == "__main__": + asyncio.run(run_client()) diff --git a/sdk/python/examples/multi-agent/README.md b/sdk/python/examples/multi-agent/README.md new file mode 100644 index 0000000..e98bd73 --- /dev/null +++ b/sdk/python/examples/multi-agent/README.md @@ -0,0 +1,21 @@ +# UC-2: Multi-Agent Data Leak Prevention + +Demonstrates how SEP-1913 annotations prevent PHI from leaking across agent boundaries. + +## What it shows + +Three agents cooperate on patient data: + +``` +Agent A (front-desk) → Agent B (analytics) → Agent C (external reporting) +``` + +- **Without annotations**: PHI flows freely A → B → C → leaked to public +- **With annotations**: Policy blocks B → C, redacts sensitive data in logs + +## Usage + +```bash +cd sdk/python +PYTHONPATH=src python examples/multi-agent/data_leak_prevention.py +``` diff --git a/sdk/python/examples/multi-agent/data_leak_prevention.py b/sdk/python/examples/multi-agent/data_leak_prevention.py new file mode 100644 index 0000000..555b016 --- /dev/null +++ b/sdk/python/examples/multi-agent/data_leak_prevention.py @@ -0,0 +1,319 @@ +"""D9: Multi-Agent Data Leak Prevention Demo. + +Side-by-side comparison showing how PHI leaks across agents WITHOUT +trust annotations (current MCP) vs. how SEP-1913 prevents it. + +The demo uses a real MCP server (Healthcare Clinic) for the "WITH" +scenario and simulates the "WITHOUT" scenario to show the contrast. + +Three agents cooperate: + Agent A — Front-desk (calls patient_lookup) + Agent B — Internal analytics + Agent C — External reporting (public) + +WITHOUT annotations: PHI flows freely A → B → C → leaked. +WITH annotations: Policy blocks B → C, redacts B → logs. + +Usage: + cd /sdk/python + PYTHONPATH=src python examples/multi-agent/data_leak_prevention.py +""" + +from __future__ import annotations + +import asyncio +import io +import json +import os +import sys +import time + +# Force UTF-8 output on Windows to support Unicode symbols +if sys.stdout.encoding != "utf-8": + sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace") +if sys.stderr.encoding != "utf-8": + sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding="utf-8", errors="replace") + +# Ensure src/ is importable +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src")) + +from mcp import ClientSession +from mcp.client.stdio import StdioServerParameters, stdio_client + +from policy import PolicyEngine +from propagate import SessionTracker +from trust_types import ( + ResultAnnotations, + SimpleDataClass, + Regulated, + sensitivity_level, +) + + +# ── ANSI helpers ────────────────────────────────────────────────── + +class C: + RST = "\033[0m" + B = "\033[1m" + DIM = "\033[2m" + R = "\033[91m" + G = "\033[92m" + Y = "\033[93m" + BL = "\033[94m" + CY = "\033[96m" + MAG = "\033[95m" + +PATIENT_DATA = { + "patient_id": "P-12345", + "name": "Jane Doe", + "dob": "1985-03-15", + "ssn": "123-45-6789", + "diagnoses": ["Type 2 Diabetes", "Hypertension"], + "medications": ["Metformin 500mg", "Lisinopril 10mg"], +} + +META_EMPTY = {} + +META_WITH_ANNOTATIONS = { + "annotations": { + "attribution": ["mcp://clinic-server/patients/P-12345"], + "openWorldHint": False, + } +} + +TOOL_ANNOTATIONS = { + "returnMetadata": { + "source": "internal", + "sensitivity": {"regulated": {"scopes": ["HIPAA"]}}, + } +} + + +def _hr(char: str = "═", width: int = 72) -> str: + return char * width + + +def _pause(seconds: float = 0.4): + time.sleep(seconds) + + +def _json_preview(data: dict, indent: int = 2) -> str: + return json.dumps(data, indent=indent) + + +def _redact(data: dict) -> dict: + out = {} + for k, v in data.items(): + if isinstance(v, dict): + out[k] = _redact(v) + elif isinstance(v, list): + out[k] = ["[REDACTED]" if isinstance(i, str) else i for i in v] + elif isinstance(v, str) and k not in ("patient_id", "status"): + out[k] = "[REDACTED]" + else: + out[k] = v + return out + + +# ── WITHOUT Trust Annotations (simulated) ──────────────────────── + +def run_without(): + print() + print(_hr()) + print(f" {C.R}{C.B}WITHOUT Trust Annotations (current MCP){C.RST}") + print(_hr()) + print() + + # Step 1: Agent A calls patient_lookup + print(f" {C.B}Agent A{C.RST} calls patient_lookup(\"P-12345\")") + _pause() + print(f" → Result: {C.DIM}{_json_preview(PATIENT_DATA)}{C.RST}") + print(f" → _meta: {C.DIM}{{}}{C.RST} ← {C.Y}empty, no annotations{C.RST}") + print() + _pause() + + # Step 2: Agent A forwards to Agent B (internal analytics) + print(f" {C.B}Agent A{C.RST} forwards result to {C.B}Agent B{C.RST} (Internal Analytics)") + _pause() + print(f" → Result forwarded {C.R}in full{C.RST}") + print(f" {C.Y}⚠ PHI forwarded — no signal it was HIPAA-regulated data{C.RST}") + print() + _pause() + + # Step 3: Agent B forwards to Agent C (external reporting) + print(f" {C.B}Agent B{C.RST} forwards result to {C.B}Agent C{C.RST} (External Reporting — {C.R}PUBLIC{C.RST})") + _pause() + print(f" → Result forwarded {C.R}in full{C.RST}: name=\"Jane Doe\", ssn=\"123-45-6789\"...") + print(f" {C.R}⚠ HIPAA VIOLATION — regulated data reached external system{C.RST}") + print() + _pause() + + # Step 4: Agent B logs result + print(f" {C.B}Agent B{C.RST} logs result to analytics pipeline") + _pause() + print(f" → Logged: {C.DIM}{{\"name\": \"Jane Doe\", \"ssn\": \"123-45-6789\", ...}}{C.RST}") + print(f" {C.R}⚠ PHI in plaintext logs — no redaction signal{C.RST}") + print() + + +# ── WITH Trust Annotations (real MCP transport) ────────────────── + +async def run_with(): + sdk_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + server_script = os.path.join(sdk_root, "examples", "_shared", "mcp_server.py") + python_src = os.path.join(sdk_root, "src") + + server_params = StdioServerParameters( + command=sys.executable, + args=[server_script], + env={**os.environ, "PYTHONPATH": python_src}, + cwd=sdk_root, + ) + + print() + print(_hr()) + print(f" {C.G}{C.B}WITH Trust Annotations (SEP-1913){C.RST}") + print(_hr()) + print() + + async with stdio_client(server_params, errlog=open(os.devnull, "w")) as (rs, ws): + async with ClientSession(rs, ws) as session: + await session.initialize() + + # Discover tools + tools_result = await session.list_tools() + tool_ann: dict[str, dict] = {} + for t in tools_result.tools: + ann = t.annotations.model_dump(exclude_none=True) if t.annotations else {} + tool_ann[t.name] = ann + + # Set up policy engine + session tracker + engine = PolicyEngine(mode="enforce") + for name, ann in tool_ann.items(): + engine.register_tool(name, ann) + + tracker_a = SessionTracker(session_id="agent-A") + tracker_b = SessionTracker(session_id="agent-B") + + # ── Step 1: Agent A calls patient_lookup (via real MCP) ── + print(f" {C.B}Agent A{C.RST} calls patient_lookup(\"P-12345\")") + _pause() + + cr = await session.call_tool("patient_lookup", {"patient_id": "P-12345"}) + raw = json.loads(cr.content[0].text) if cr.content else {} + + print(f" → Result: {C.DIM}{json.dumps(raw)}{C.RST}") + print(f" → _meta.annotations: {C.CY}{_json_preview(META_WITH_ANNOTATIONS['annotations'])}{C.RST}") + print(f" → tool.annotations.returnMetadata: {C.CY}{_json_preview(TOOL_ANNOTATIONS['returnMetadata'])}{C.RST}") + print() + _pause() + + # Merge into Agent A's session + tracker_a.merge( + ResultAnnotations( + sensitivity=Regulated.of("HIPAA"), + attribution=("mcp://clinic-server/patients/P-12345",), + ), + "patient_lookup", + ) + + # ── Step 2: Agent A forwards to Agent B (internal) ── + print(f" {C.B}Agent A{C.RST} forwards result to {C.B}Agent B{C.RST} (Internal Analytics — {C.BL}internal{C.RST})") + _pause() + + decision = engine.evaluate( + "patient_lookup", action="forward", target_destination="internal", + ) + print(f" → Host policy: regulated(HIPAA) + target=internal → {C.G}✓ ALLOW{C.RST}") + print(f" → Request _meta.annotations: {C.CY}{{\"attribution\": [\"mcp://clinic-server/patients/P-12345\"]}}{C.RST}") + print(f" → Session tracks: sensitivity=regulated(HIPAA), attribution accumulated") + print() + _pause() + + # Agent B now has the data in its session + tracker_b.merge( + ResultAnnotations( + sensitivity=Regulated.of("HIPAA"), + attribution=("mcp://clinic-server/patients/P-12345",), + ), + "patient_lookup (forwarded)", + ) + + # ── Step 3: Agent B forwards to Agent C (public) — BLOCKED ── + print(f" {C.B}Agent B{C.RST} forwards result to {C.B}Agent C{C.RST} (External Reporting — {C.R}PUBLIC{C.RST})") + _pause() + + decision = engine.evaluate( + "patient_lookup", action="forward", target_destination="public", + ) + assert not decision.allowed, "Expected BLOCK" + print(f" → Host policy: regulated(HIPAA) + target=public → {C.R}✗ BLOCK{C.RST}") + print(f" → Agent C receives: {C.DIM}\"Policy denied: HIPAA-regulated data cannot leave trust boundary\"{C.RST}") + print(f" {C.G}✓ PHI PROTECTED — deterministic enforcement, no model involved{C.RST}") + print() + _pause() + + # ── Step 4: Agent B logs result — REDACT ── + print(f" {C.B}Agent B{C.RST} logs result to analytics pipeline") + _pause() + + # Policy says regulated data can't be logged in plaintext + redacted = _redact(raw) + print(f" → Host policy: regulated(HIPAA) + action=log → {C.Y}◉ REDACT{C.RST}") + print(f" → Logged: {C.DIM}{json.dumps(redacted)}{C.RST}") + print(f" → Audit metadata: {C.CY}{{\"source\": \"internal\", \"sensitivity\": \"regulated:HIPAA\"}}{C.RST}") + print(f" {C.G}✓ COMPLIANT — PHI redacted, audit trail preserved{C.RST}") + print() + + +# ── Comparison summary ──────────────────────────────────────────── + +def print_summary(): + print(_hr()) + print(f" {C.B}Summary: What Changed?{C.RST}") + print(_hr()) + print() + print(f" ┌─────────────────────────────┬─────────────────┬──────────────────┐") + print(f" │ Scenario │ {C.R}Without SEP-1913{C.RST} │ {C.G}With SEP-1913{C.RST} │") + print(f" ├─────────────────────────────┼─────────────────┼──────────────────┤") + print(f" │ A → B (internal forward) │ {C.Y}No check{C.RST} │ {C.G}✓ ALLOW{C.RST} │") + print(f" │ B → C (public forward) │ {C.R}PHI leaked{C.RST} │ {C.G}✗ BLOCKED{C.RST} │") + print(f" │ B → logs (analytics) │ {C.R}PHI in plaintext{C.RST} │ {C.G}◉ REDACTED{C.RST} │") + print(f" │ Audit trail │ {C.R}None{C.RST} │ {C.G}✓ Full provenance{C.RST} │") + print(f" │ Model involvement needed │ {C.R}Yes (unreliable){C.RST} │ {C.G}No (deterministic){C.RST}│") + print(f" └─────────────────────────────┴─────────────────┴──────────────────┘") + print() + print(f" {C.B}Key insight:{C.RST} SEP-1913 annotations give hosts the metadata they") + print(f" need to enforce data boundaries {C.B}without relying on the model{C.RST} to") + print(f" understand sensitivity. The policy engine is deterministic —") + print(f" HIPAA data {C.B}cannot{C.RST} reach a public endpoint regardless of the prompt.") + print() + + +# ── Main ────────────────────────────────────────────────────────── + +async def main(): + print() + print(f" {C.B}{C.MAG}{'=' * 68}{C.RST}") + print(f" {C.B}{C.MAG} D9: Multi-Agent Data Leak Prevention{C.RST}") + print(f" {C.B}{C.MAG}{'=' * 68}{C.RST}") + print(f" {C.DIM}Three agents: A (front-desk) → B (analytics) → C (external reporting)") + print(f" Patient: Jane Doe (P-12345), HIPAA-regulated PHI{C.RST}") + + # Part 1: WITHOUT (simulated — no server needed) + run_without() + + # Part 2: WITH (real MCP server) + await run_with() + + # Comparison table + print_summary() + + print(f" {C.B}{C.MAG}{'=' * 68}{C.RST}") + print(f" {C.B}{C.MAG} Demo complete — the \"WITH\" scenario used real MCP stdio transport.{C.RST}") + print(f" {C.B}{C.MAG}{'=' * 68}{C.RST}") + print() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/sdk/python/pyproject.toml b/sdk/python/pyproject.toml new file mode 100644 index 0000000..2ce8cda --- /dev/null +++ b/sdk/python/pyproject.toml @@ -0,0 +1,43 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "mcp-trust-annotations" +version = "0.1.0" +description = "SEP-1913 Trust & Sensitivity Annotations SDK for Model Context Protocol" +readme = "README.md" +license = "MIT" +requires-python = ">=3.10" +authors = [{ name = "MCP Trust Contributors" }] +keywords = ["mcp", "trust", "annotations", "security", "data-classification"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Topic :: Security", + "Topic :: Software Development :: Libraries", + "Typing :: Typed", +] + +# Zero runtime dependencies — stdlib only +dependencies = [] + +[project.optional-dependencies] +dev = ["pytest>=7.0", "pytest-asyncio>=0.21"] + +[project.urls] +Repository = "https://github.com/modelcontextprotocol/experimental-ext-tool-annotations" + +[tool.hatch.build.targets.wheel] +packages = ["src"] + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] +asyncio_mode = "auto" diff --git a/sdk/python/pyrightconfig.json b/sdk/python/pyrightconfig.json new file mode 100644 index 0000000..6b262cd --- /dev/null +++ b/sdk/python/pyrightconfig.json @@ -0,0 +1,3 @@ +{ + "extraPaths": ["src"] +} diff --git a/sdk/python/src/annotate.py b/sdk/python/src/annotate.py new file mode 100644 index 0000000..325438f --- /dev/null +++ b/sdk/python/src/annotate.py @@ -0,0 +1,321 @@ +"""Declarative trust annotation decorator and wire-format serialization. + +Attach SEP-1913 trust annotations to any function with a single decorator. +Read them back at runtime for policy evaluation, logging, and MCP wire output. + +Usage: + from trust_annotated, get_trust_annotations, to_wire + from ReturnMetadata, Source, Regulated + + @trust_annotated( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + attribution=("EHR-System",), + ) + async def patient_lookup(pid: str) -> dict: + ... + + # Read metadata + ann = get_trust_annotations(patient_lookup) + assert ann.return_metadata.source == Source.INTERNAL + + # Serialize to MCP wire format + wire = to_wire(ann) + # → {"attribution": ["EHR-System"], "returnMetadata": {"source": "internal", ...}} +""" + +from __future__ import annotations + +import asyncio +import functools +import time +from typing import Any, Callable, TypeVar + +from trust_types import ( + DataClass, + Destination, + InputMetadata, + Outcome, + Regulated, + ReturnMetadata, + SimpleDataClass, + Source, + TrustAnnotations, +) + +_TRUST_ATTR = "__mcp_trust_annotations__" +F = TypeVar("F", bound=Callable[..., Any]) + + +# --------------------------------------------------------------------------- +# Decorator +# --------------------------------------------------------------------------- + +def trust_annotated( + *, + malicious_activity_hint: bool | None = None, + attribution: tuple[str, ...] = (), + input_metadata: InputMetadata | None = None, + return_metadata: ReturnMetadata | None = None, +) -> Callable[[F], F]: + """Attach SEP-1913 trust annotations to a tool function. + + This is a pure metadata decorator — it does NOT wrap the function + at runtime unless logging is enabled (see emit.py). When logging + is disabled, overhead is effectively zero. + + Args: + malicious_activity_hint: Tool may encounter malicious content. + attribution: Source identifiers for the data this tool handles. + input_metadata: Where inputs go and their security implications. + return_metadata: Where outputs come from and their sensitivity. + + Returns: + Decorator that attaches TrustAnnotations to the function. + """ + annotations = TrustAnnotations( + malicious_activity_hint=malicious_activity_hint, + attribution=attribution, + input_metadata=input_metadata, + return_metadata=return_metadata, + ) + + def decorator(fn: F) -> F: + setattr(fn, _TRUST_ATTR, annotations) + + @functools.wraps(fn) + def sync_wrapper(*args: Any, **kwargs: Any) -> Any: + from emit import _emit_call, _emit_result + _emit_call(fn.__name__, annotations) + start = time.monotonic() + try: + result = fn(*args, **kwargs) + elapsed = (time.monotonic() - start) * 1000 + _emit_result(fn.__name__, annotations, elapsed, "ok") + return result + except Exception as exc: + elapsed = (time.monotonic() - start) * 1000 + _emit_result(fn.__name__, annotations, elapsed, "error", str(exc)) + raise + + @functools.wraps(fn) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + from emit import _emit_call, _emit_result + _emit_call(fn.__name__, annotations) + start = time.monotonic() + try: + result = await fn(*args, **kwargs) + elapsed = (time.monotonic() - start) * 1000 + _emit_result(fn.__name__, annotations, elapsed, "ok") + return result + except Exception as exc: + elapsed = (time.monotonic() - start) * 1000 + _emit_result(fn.__name__, annotations, elapsed, "error", str(exc)) + raise + + wrapper = async_wrapper if asyncio.iscoroutinefunction(fn) else sync_wrapper + setattr(wrapper, _TRUST_ATTR, annotations) + return wrapper # type: ignore[return-value] + + return decorator + + +# --------------------------------------------------------------------------- +# Runtime inspection +# --------------------------------------------------------------------------- + +def get_trust_annotations(fn: Callable[..., Any]) -> TrustAnnotations | None: + """Read trust annotations from a decorated function. + + Returns None if the function is not annotated. + """ + return getattr(fn, _TRUST_ATTR, None) + + +def is_trust_annotated(fn: Callable[..., Any]) -> bool: + """Check whether a function has trust annotations.""" + return hasattr(fn, _TRUST_ATTR) + + +# --------------------------------------------------------------------------- +# Wire-format serialization (Python → MCP JSON) +# --------------------------------------------------------------------------- + +def to_wire(annotations: TrustAnnotations) -> dict[str, Any]: + """Serialize TrustAnnotations to MCP wire format dict. + + Output can be spread into ToolAnnotations in a tools/list response: + + tool = { + "name": "patient_lookup", + "description": "...", + "annotations": { + "readOnlyHint": True, + **to_wire(trust_annotations), + } + } + """ + result: dict[str, Any] = {} + + if annotations.malicious_activity_hint is not None: + result["maliciousActivityHint"] = annotations.malicious_activity_hint + + if annotations.attribution: + result["attribution"] = list(annotations.attribution) + + if annotations.input_metadata: + result["inputMetadata"] = _input_meta_to_wire(annotations.input_metadata) + + if annotations.return_metadata: + result["returnMetadata"] = _return_meta_to_wire(annotations.return_metadata) + + return result + + +def from_wire(data: dict[str, Any]) -> TrustAnnotations: + """Deserialize MCP wire format dict to TrustAnnotations. + + Parses the trust-related fields from a ToolAnnotations dict. + Raises ``ValueError`` on malformed input. + + wire = {"maliciousActivityHint": True, "attribution": ["EHR"], + "returnMetadata": {"source": "internal", "sensitivity": "pii"}} + ann = from_wire(wire) + """ + if not isinstance(data, dict): + raise ValueError(f"from_wire() expects a dict, got {type(data).__name__}") + + # Validate maliciousActivityHint + hint = data.get("maliciousActivityHint") + if hint is not None and not isinstance(hint, bool): + raise ValueError( + f"maliciousActivityHint must be bool or null, got {type(hint).__name__}" + ) + + # Validate attribution + raw_attr = data.get("attribution", ()) + if not isinstance(raw_attr, (list, tuple)): + raise ValueError( + f"attribution must be an array, got {type(raw_attr).__name__}" + ) + for i, item in enumerate(raw_attr): + if not isinstance(item, str): + raise ValueError( + f"attribution[{i}] must be a string, got {type(item).__name__}: {item!r}" + ) + + return TrustAnnotations( + malicious_activity_hint=hint, + attribution=tuple(raw_attr), + input_metadata=_input_meta_from_wire(data["inputMetadata"]) + if "inputMetadata" in data + else None, + return_metadata=_return_meta_from_wire(data["returnMetadata"]) + if "returnMetadata" in data + else None, + ) + + +# --------------------------------------------------------------------------- +# Internal: serialization helpers +# --------------------------------------------------------------------------- + +def _sensitivity_to_wire(s: DataClass) -> str | dict[str, Any]: + if isinstance(s, SimpleDataClass): + return s.value + elif isinstance(s, Regulated): + return {"regulated": {"scopes": list(s.regulated.scopes)}} + raise TypeError(f"Unknown DataClass type: {type(s)}") + + +def _sensitivity_from_wire(raw: str | dict[str, Any]) -> DataClass: + if isinstance(raw, str): + try: + return SimpleDataClass(raw) + except ValueError: + valid = [e.value for e in SimpleDataClass] + raise ValueError( + f"Unknown sensitivity value {raw!r}; " + f"expected one of {valid} or a regulated object" + ) from None + if isinstance(raw, dict) and "regulated" in raw: + reg = raw["regulated"] + if not isinstance(reg, dict) or "scopes" not in reg: + raise ValueError( + f"regulated must be an object with 'scopes', got: {reg!r}" + ) + raw_scopes = reg["scopes"] + if not isinstance(raw_scopes, (list, tuple)): + raise ValueError( + f"regulated.scopes must be an array, got {type(raw_scopes).__name__}" + ) + for i, s in enumerate(raw_scopes): + if not isinstance(s, str): + raise ValueError( + f"regulated.scopes[{i}] must be a string, got {type(s).__name__}" + ) + scopes = tuple(raw_scopes) + from trust_types import RegulatoryScope + return Regulated(RegulatoryScope(scopes)) + raise ValueError(f"Cannot parse DataClass from: {raw!r}") + + +def _enum_or_array_to_wire(val: Any) -> str | list[str]: + """Convert a single enum or tuple of enums to wire format.""" + if isinstance(val, tuple): + return [v.value for v in val] + return val.value + + +def _sensitivity_or_array_to_wire( + val: DataClass | tuple[DataClass, ...], +) -> str | dict | list: + if isinstance(val, tuple): + return [_sensitivity_to_wire(v) for v in val] + return _sensitivity_to_wire(val) + + +def _input_meta_to_wire(m: InputMetadata) -> dict[str, Any]: + return { + "destination": _enum_or_array_to_wire(m.destination), + "sensitivity": _sensitivity_or_array_to_wire(m.sensitivity), + "outcomes": _enum_or_array_to_wire(m.outcomes), + } + + +def _return_meta_to_wire(m: ReturnMetadata) -> dict[str, Any]: + return { + "source": _enum_or_array_to_wire(m.source), + "sensitivity": _sensitivity_or_array_to_wire(m.sensitivity), + } + + +def _enum_or_array_from_wire(raw: str | list[str], enum_cls: type) -> Any: + if isinstance(raw, list): + return tuple(enum_cls(v) for v in raw) + return enum_cls(raw) + + +def _sensitivity_or_array_from_wire( + raw: str | dict | list, +) -> DataClass | tuple[DataClass, ...]: + if isinstance(raw, list): + return tuple(_sensitivity_from_wire(v) for v in raw) + return _sensitivity_from_wire(raw) + + +def _input_meta_from_wire(data: dict[str, Any]) -> InputMetadata: + return InputMetadata( + destination=_enum_or_array_from_wire(data.get("destination", "ephemeral"), Destination), + sensitivity=_sensitivity_or_array_from_wire(data.get("sensitivity", "none")), + outcomes=_enum_or_array_from_wire(data.get("outcomes", "benign"), Outcome), + ) + + +def _return_meta_from_wire(data: dict[str, Any]) -> ReturnMetadata: + return ReturnMetadata( + source=_enum_or_array_from_wire(data.get("source", "system"), Source), + sensitivity=_sensitivity_or_array_from_wire(data.get("sensitivity", "none")), + ) diff --git a/sdk/python/src/emit.py b/sdk/python/src/emit.py new file mode 100644 index 0000000..cc313f3 --- /dev/null +++ b/sdk/python/src/emit.py @@ -0,0 +1,148 @@ +"""Structured audit log emitter — zero-infrastructure observability. + +Emits JSON lines to stderr (default) for every tool call, recording +SEP-1913 trust annotations alongside timing and status information. +Works with any log aggregator: ELK, Datadog, Splunk, CloudWatch, etc. + +Usage: + from emit import enable_logging + + enable_logging() # JSON to stderr + enable_logging(stream=sys.stdout) # JSON to stdout + enable_logging(callback=send_to_siem) # custom handler + enable_logging(agent_id="urn:agent:my-app") # with agent ID +""" + +from __future__ import annotations + +import json +import sys +import threading +from datetime import datetime, timezone +from typing import Any, Callable, TextIO + +from trust_types import TrustAnnotations + +# --------------------------------------------------------------------------- +# Module-level emitter state +# --------------------------------------------------------------------------- + +_lock = threading.Lock() +_emitter: Callable[[dict[str, Any]], None] | None = None +_agent_id: str | None = None + + +# --------------------------------------------------------------------------- +# Public configuration +# --------------------------------------------------------------------------- + +def enable_logging( + *, + stream: TextIO | None = None, + callback: Callable[[dict[str, Any]], None] | None = None, + agent_id: str | None = None, + pretty: bool = False, +) -> None: + """Enable structured audit logging. Call once at startup. + + Args: + stream: Output stream (default: stderr). + callback: Custom handler receiving event dicts. Overrides stream. + agent_id: Agent identifier included in every event. + pretty: Indent JSON for readability (development only). + """ + global _emitter, _agent_id + _agent_id = agent_id + + if callback is not None: + _emitter = callback + else: + target = stream or sys.stderr + indent = 2 if pretty else None + + def _stream_emit(event: dict[str, Any]) -> None: + line = json.dumps(event, default=str, indent=indent) + with _lock: + target.write(line + "\n") + target.flush() + + _emitter = _stream_emit + + +def disable_logging() -> None: + """Disable structured logging.""" + global _emitter, _agent_id + _emitter = None + _agent_id = None + + +# --------------------------------------------------------------------------- +# Internal: emit events +# --------------------------------------------------------------------------- + +def _emit_call(tool: str, annotations: TrustAnnotations) -> None: + """Emit a 'tool.call' event when a tool is invoked.""" + if _emitter is None: + return + from annotate import to_wire + event = { + "ts": _now_iso(), + "event": "tool.call", + "tool": tool, + "annotations": to_wire(annotations), + } + if _agent_id: + event["agent"] = _agent_id + _emitter(event) + + +def _emit_result( + tool: str, + annotations: TrustAnnotations, + duration_ms: float, + status: str, + error: str | None = None, +) -> None: + """Emit a 'tool.result' event when a tool call completes.""" + if _emitter is None: + return + from annotate import to_wire + event: dict[str, Any] = { + "ts": _now_iso(), + "event": "tool.result", + "tool": tool, + "annotations": to_wire(annotations), + "duration_ms": round(duration_ms, 2), + "status": status, + } + if error: + event["error"] = error + if _agent_id: + event["agent"] = _agent_id + _emitter(event) + + +def _emit_policy( + tool: str, + effect: str, + reason: str, + **extra: Any, +) -> None: + """Emit a 'policy.decision' event when a policy rule fires.""" + if _emitter is None: + return + event: dict[str, Any] = { + "ts": _now_iso(), + "event": "policy.decision", + "tool": tool, + "effect": effect, + "reason": reason, + } + event.update(extra) + if _agent_id: + event["agent"] = _agent_id + _emitter(event) + + +def _now_iso() -> str: + return datetime.now(timezone.utc).isoformat() diff --git a/sdk/python/src/policy.py b/sdk/python/src/policy.py new file mode 100644 index 0000000..3eebd99 --- /dev/null +++ b/sdk/python/src/policy.py @@ -0,0 +1,449 @@ +"""Policy engine — evaluate trust annotations against configurable rules. + +Provides three enforcement modes: + - audit: log decisions, never block (default) + - warn: log decisions, surface warnings to the user + - enforce: block tool calls that violate policy + +The engine ships with sensible default rules but is fully extensible. + +Usage: + from policy import PolicyEngine + from trust_types import TrustAnnotations + + engine = PolicyEngine(mode="enforce") + + # Register tools with their annotations + engine.register_tool("patient_lookup", {"returnMetadata": {"sensitivity": "pii"}}) + + # Evaluate before calling + decision = engine.evaluate("patient_lookup", action="call") + if not decision.allowed: + print(f"Blocked: {decision.reason}") + + # Or use the convenience method + engine.register_trust_annotations("send_email", trust_ann) + decision = engine.evaluate( + "send_email", + action="call", + target_destination="public", + session=session_tracker.session, + ) +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any + +from emit import _emit_policy +from trust_types import ( + DataClass, + Destination, + Regulated, + SessionAnnotations, + SimpleDataClass, + TrustAnnotations, + sensitivity_level, +) + + +# --------------------------------------------------------------------------- +# Policy decision +# --------------------------------------------------------------------------- + +@dataclass(frozen=True) +class PolicyDecision: + """Result of evaluating a policy rule.""" + allowed: bool + effect: str # "allow" | "block" | "escalate" | "redact" | "warn" + reason: str + rule: str = "" # which rule triggered + regulations: tuple[str, ...] = () + + +# --------------------------------------------------------------------------- +# Rule interface +# --------------------------------------------------------------------------- + +class PolicyRule(ABC): + """Base class for policy rules. + + Implement evaluate() to return a PolicyDecision if the rule triggers, + or None to pass through to the next rule. + """ + name: str = "unnamed" + + @abstractmethod + def evaluate( + self, + annotations: dict[str, Any], + action: str, + target_destination: str | None, + session: SessionAnnotations | None, + ) -> PolicyDecision | None: + ... + + +# --------------------------------------------------------------------------- +# Default rules +# --------------------------------------------------------------------------- + +class BlockCredentialsToPublic(PolicyRule): + """Block sending credential-level data to public destinations.""" + name = "block-credentials-to-public" + + def evaluate( + self, + annotations: dict[str, Any], + action: str, + target_destination: str | None, + session: SessionAnnotations | None, + ) -> PolicyDecision | None: + if target_destination != "public": + return None + sensitivity = _extract_sensitivity(annotations) + if sensitivity and sensitivity_level(sensitivity) >= sensitivity_level(SimpleDataClass.CREDENTIALS): + return PolicyDecision( + allowed=False, + effect="block", + reason=f"Cannot send {sensitivity} data to public destination", + rule=self.name, + ) + return None + + +class BlockRegulatedToExternal(PolicyRule): + """Block sending regulated data to public/external destinations.""" + name = "block-regulated-to-external" + + def evaluate( + self, + annotations: dict[str, Any], + action: str, + target_destination: str | None, + session: SessionAnnotations | None, + ) -> PolicyDecision | None: + if target_destination not in ("public", "external"): + return None + sensitivity = _extract_sensitivity(annotations) + if isinstance(sensitivity, Regulated): + return PolicyDecision( + allowed=False, + effect="block", + reason=f"Regulated data ({', '.join(sensitivity.regulated.scopes)}) cannot leave organization", + rule=self.name, + regulations=sensitivity.regulated.scopes, + ) + return None + + +class WarnOnPIIForward(PolicyRule): + """Warn when PII data is being forwarded to another tool.""" + name = "warn-pii-forward" + + def evaluate( + self, + annotations: dict[str, Any], + action: str, + target_destination: str | None, + session: SessionAnnotations | None, + ) -> PolicyDecision | None: + if action != "forward": + return None + sensitivity = _extract_sensitivity(annotations) + if sensitivity and sensitivity_level(sensitivity) >= sensitivity_level(SimpleDataClass.PII): + return PolicyDecision( + allowed=True, + effect="warn", + reason=f"Forwarding {sensitivity} data — ensure recipient has need-to-know", + rule=self.name, + ) + return None + + +class BlockMaliciousActivity(PolicyRule): + """Block tool calls flagged as potentially malicious.""" + name = "block-malicious-activity" + + def evaluate( + self, + annotations: dict[str, Any], + action: str, + target_destination: str | None, + session: SessionAnnotations | None, + ) -> PolicyDecision | None: + if annotations.get("maliciousActivityHint"): + return PolicyDecision( + allowed=False, + effect="block", + reason="Tool flagged for potential malicious activity", + rule=self.name, + ) + return None + + +class EscalateSessionSensitivity(PolicyRule): + """Require escalation when session sensitivity exceeds threshold. + + When the session has accumulated sensitive data (e.g. PII) and + the tool targets a less-restricted destination, escalate for + human approval. + """ + name = "escalate-session-sensitivity" + + def __init__(self, threshold: DataClass = SimpleDataClass.PII) -> None: + self._threshold = threshold + + def evaluate( + self, + annotations: dict[str, Any], + action: str, + target_destination: str | None, + session: SessionAnnotations | None, + ) -> PolicyDecision | None: + if session is None: + return None + if target_destination not in ("public", "external"): + return None + if sensitivity_level(session.max_sensitivity) >= sensitivity_level(self._threshold): + return PolicyDecision( + allowed=False, + effect="escalate", + reason=( + f"Session contains {session.max_sensitivity} " + f"data — outbound to '{target_destination}' requires approval" + ), + rule=self.name, + ) + return None + + +class BlockOpenWorldToSensitiveWrite(PolicyRule): + """Block writes to sensitive tools when session has seen untrusted content. + + Prevents prompt-injection-driven data corruption: if the session has + consumed open-world (untrusted) content and the target tool is + destructive with sensitive inputs, the call is blocked. + + This closes the gap where an attacker injects instructions into + untrusted tool output (e.g. search results) that trick the LLM + into calling a write tool with attacker-controlled arguments. + """ + name = "block-openworld-to-sensitive-write" + + def __init__(self, threshold: DataClass = SimpleDataClass.PII) -> None: + self._threshold = threshold + + def evaluate( + self, + annotations: dict[str, Any], + action: str, + target_destination: str | None, + session: SessionAnnotations | None, + ) -> PolicyDecision | None: + if session is None or not session.open_world_hint: + return None + if not annotations.get("destructiveHint"): + return None + sensitivity = _extract_sensitivity(annotations) + if sensitivity and sensitivity_level(sensitivity) >= sensitivity_level(self._threshold): + return PolicyDecision( + allowed=False, + effect="block", + reason=( + "Write to sensitive tool blocked: session contains " + "untrusted open-world content" + ), + rule=self.name, + ) + return None + + +# Default rule set — ordered from most to least severe +DEFAULT_RULES: list[PolicyRule] = [ + BlockMaliciousActivity(), + BlockCredentialsToPublic(), + BlockRegulatedToExternal(), + BlockOpenWorldToSensitiveWrite(), + EscalateSessionSensitivity(), + WarnOnPIIForward(), +] + + +# --------------------------------------------------------------------------- +# Engine +# --------------------------------------------------------------------------- + +class PolicyEngine: + """Evaluates tool calls against a configurable set of policy rules. + + Modes: + - audit: decisions are logged but never enforced + - warn: "block" decisions become "warn" (logged, not enforced) — **default** + - enforce: decisions are enforced as-is + """ + + def __init__( + self, + mode: str = "warn", + rules: list[PolicyRule] | None = None, + fail_closed: bool = False, + ) -> None: + if mode not in ("audit", "warn", "enforce"): + raise ValueError(f"Invalid mode: {mode!r} — must be 'audit', 'warn', or 'enforce'") + self.mode = mode + self.fail_closed = fail_closed + self._tools: dict[str, dict[str, Any]] = {} + self._rules = rules if rules is not None else DEFAULT_RULES.copy() + + # ----------------------------------------------------------------- + # Registration + # ----------------------------------------------------------------- + + def register_tool(self, name: str, annotations: dict[str, Any]) -> None: + """Register a tool's wire-format annotations dict.""" + self._tools[name] = annotations + + def register_trust_annotations( + self, name: str, annotations: TrustAnnotations + ) -> None: + """Register a tool using typed TrustAnnotations.""" + from annotate import to_wire + self._tools[name] = to_wire(annotations) + + def add_rule(self, rule: PolicyRule) -> None: + """Add a custom policy rule.""" + self._rules.append(rule) + + # ----------------------------------------------------------------- + # Evaluation + # ----------------------------------------------------------------- + + def evaluate( + self, + tool: str, + action: str = "call", + target_destination: str | None = None, + session: SessionAnnotations | None = None, + ) -> PolicyDecision: + """Evaluate policy rules for a tool call. + + Args: + tool: Tool name (must be registered). + action: "call" | "forward" | "store" | etc. + target_destination: Where the output will be sent. + session: Current session annotations (for escalation rules). + + Returns: + PolicyDecision with the first matching rule's verdict, + adjusted for the engine's mode. + """ + if tool not in self._tools: + if self.fail_closed: + decision = PolicyDecision( + allowed=False, + effect="block", + reason=f"Unregistered tool '{tool}' blocked by fail-closed policy", + rule="fail-closed", + ) + _emit_policy( + tool=tool, + effect=decision.effect, + reason=decision.reason, + rule=decision.rule, + mode=self.mode, + ) + return decision + annotations: dict[str, Any] = {} + else: + annotations = self._tools[tool] + + for rule in self._rules: + decision = rule.evaluate(annotations, action, target_destination, session) + if decision is not None: + final = self._apply_mode(decision) + _emit_policy( + tool=tool, + effect=final.effect, + reason=final.reason, + rule=final.rule, + mode=self.mode, + original_effect=decision.effect, + ) + return final + + return PolicyDecision( + allowed=True, + effect="allow", + reason="No policy rule triggered", + ) + + # ----------------------------------------------------------------- + # Inspection + # ----------------------------------------------------------------- + + @property + def tools(self) -> dict[str, dict[str, Any]]: + """Copy of registered tool annotations.""" + return dict(self._tools) + + @property + def rules(self) -> list[PolicyRule]: + """Current rule list.""" + return list(self._rules) + + # ----------------------------------------------------------------- + # Internal + # ----------------------------------------------------------------- + + def _apply_mode(self, decision: PolicyDecision) -> PolicyDecision: + if self.mode == "enforce": + return decision + if self.mode == "warn": + if decision.effect == "block" or decision.effect == "escalate": + return PolicyDecision( + allowed=True, + effect="warn", + reason=f"[WARN] {decision.reason}", + rule=decision.rule, + regulations=decision.regulations, + ) + return decision + # audit mode — everything is allowed, just logged + return PolicyDecision( + allowed=True, + effect="audit", + reason=f"[AUDIT] {decision.reason}", + rule=decision.rule, + regulations=decision.regulations, + ) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _extract_sensitivity(annotations: dict[str, Any]) -> DataClass | None: + """Extract the highest sensitivity from a wire-format annotations dict.""" + from annotate import _sensitivity_from_wire + from trust_types import max_sensitivity as _max + + best: DataClass | None = None + + for key in ("inputMetadata", "returnMetadata"): + meta = annotations.get(key, {}) + raw_sens = meta.get("sensitivity") + if raw_sens is None: + continue + + items = raw_sens if isinstance(raw_sens, list) else [raw_sens] + for item in items: + parsed = _sensitivity_from_wire(item) + if best is None: + best = parsed + else: + best = _max(best, parsed) + + return best diff --git a/sdk/python/src/propagate.py b/sdk/python/src/propagate.py new file mode 100644 index 0000000..2f48ceb --- /dev/null +++ b/sdk/python/src/propagate.py @@ -0,0 +1,215 @@ +"""Session propagation tracker — implements SEP-1913 propagation rules. + +Tracks accumulated annotations across an agent session, applying the +monotonic escalation rules defined in SEP-1913: + + - openWorldHint: boolean union (once true, stays true) + - maliciousActivityHint: boolean union + - attribution: set union (accumulates sources) + - sensitivity: session-max ordering (never decreases) + +Usage: + from propagate import SessionTracker + from trust_types import ResultAnnotations, SimpleDataClass + + tracker = SessionTracker() + + # After each tool call, merge the result annotations + tracker.merge(ResultAnnotations( + sensitivity=SimpleDataClass.PII, + attribution=("EHR-System",), + )) + + # Check session state before allowing an outbound call + if tracker.session.max_sensitivity != SimpleDataClass.NONE: + print("Session has seen sensitive data — restrict outbound tools") + + # Get session annotations as wire format + wire = tracker.to_wire() +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass, field +from typing import Any + +from trust_types import ( + DataClass, + Regulated, + ResultAnnotations, + RegulatoryScope, + SessionAnnotations, + SimpleDataClass, + max_sensitivity, + sensitivity_level, +) + + +class SessionTracker: + """Tracks accumulated trust annotations across an agent session. + + Thread-safe: all mutations are serialized through a + ``threading.Lock``. + + The tracker implements the monotonic escalation principle: session + state can only become MORE restrictive, never less. + """ + + def __init__( + self, + session_id: str | None = None, + allow_reset: bool = False, + ) -> None: + self._lock = threading.Lock() + self._session = SessionAnnotations() + self._session_id = session_id + self._call_count = 0 + self._history: list[_MergeEvent] = [] + self._allow_reset = allow_reset + + @property + def session(self) -> SessionAnnotations: + """Read-only snapshot of current session state. + + Returns a defensive copy — mutations on the returned object + do NOT affect the tracker's internal state. + """ + with self._lock: + s = self._session + return SessionAnnotations( + open_world_hint=s.open_world_hint, + malicious_activity_hint=s.malicious_activity_hint, + attribution=set(s.attribution), + max_sensitivity=s.max_sensitivity, + ) + + @property + def session_id(self) -> str | None: + return self._session_id + + @property + def call_count(self) -> int: + """Number of tool calls merged into this session.""" + return self._call_count + + # ----------------------------------------------------------------- + # Core: merge tool result into session + # ----------------------------------------------------------------- + + def merge(self, result: ResultAnnotations, tool_name: str = "") -> None: + """Merge a tool call's result annotations into the session. + + Applies SEP-1913 propagation rules: + - Boolean union for openWorldHint / maliciousActivityHint + - Set union for attribution + - Max-sensitivity for DataClass + + Args: + result: Annotations from a CallToolResult. + tool_name: Optional tool name for audit history. + """ + with self._lock: + prev_level = sensitivity_level(self._session.max_sensitivity) + + self._session.merge(result) + self._call_count += 1 + + new_level = sensitivity_level(self._session.max_sensitivity) + escalated = new_level > prev_level + + self._history.append(_MergeEvent( + tool=tool_name, + sensitivity=result.sensitivity, + escalated=escalated, + )) + + # ----------------------------------------------------------------- + # Queries + # ----------------------------------------------------------------- + + def has_seen_sensitive_data(self) -> bool: + """Return True if session has seen anything above 'none'.""" + return sensitivity_level(self._session.max_sensitivity) > 0 + + def sensitivity_at_least(self, threshold: DataClass) -> bool: + """Return True if session sensitivity >= the given threshold.""" + return sensitivity_level(self._session.max_sensitivity) >= sensitivity_level(threshold) + + def get_regulatory_scopes(self) -> tuple[str, ...]: + """Return all regulatory scopes accumulated in the session. + + Returns empty tuple if no regulated data has been seen. + """ + ms = self._session.max_sensitivity + if isinstance(ms, Regulated): + return ms.regulated.scopes + return () + + # ----------------------------------------------------------------- + # Serialization + # ----------------------------------------------------------------- + + def to_wire(self) -> dict[str, Any]: + """Serialize session state to wire format dict. + + Suitable for including in request annotations when calling + outbound tools. + """ + result: dict[str, Any] = {} + if self._session.open_world_hint: + result["openWorldHint"] = True + if self._session.malicious_activity_hint: + result["maliciousActivityHint"] = True + if self._session.attribution: + result["attribution"] = sorted(self._session.attribution) + if sensitivity_level(self._session.max_sensitivity) > 0: + from annotate import _sensitivity_to_wire + result["maxSensitivity"] = _sensitivity_to_wire( + self._session.max_sensitivity + ) + return result + + def summary(self) -> str: + """Human-readable session state summary.""" + lines = [ + f"Session: {self._session_id or '(anonymous)'}", + f" Calls: {self._call_count}", + f" Open-world: {self._session.open_world_hint}", + f" Malicious activity: {self._session.malicious_activity_hint}", + f" Max sensitivity: {self._session.max_sensitivity}", + f" Attribution: {sorted(self._session.attribution) or '(none)'}", + ] + if self._history: + lines.append(" History:") + for evt in self._history: + esc = " [ESCALATED]" if evt.escalated else "" + lines.append(f" - {evt.tool or '?'}: {evt.sensitivity}{esc}") + return "\n".join(lines) + + def reset(self) -> None: + """Reset session state. Violates monotonicity. + + Only available when the tracker was created with + ``allow_reset=True`` (intended for testing). Raises + ``RuntimeError`` otherwise to prevent accidental session + laundering. + """ + if not self._allow_reset: + raise RuntimeError( + "reset() is disabled by default to preserve monotonic " + "escalation. Pass allow_reset=True to the SessionTracker " + "constructor if you need this (e.g. in tests)." + ) + with self._lock: + self._session = SessionAnnotations() + self._call_count = 0 + self._history.clear() + + +@dataclass(frozen=True) +class _MergeEvent: + """Internal audit record of a merge operation.""" + tool: str + sensitivity: DataClass + escalated: bool diff --git a/sdk/python/src/trust_types.py b/sdk/python/src/trust_types.py new file mode 100644 index 0000000..9792039 --- /dev/null +++ b/sdk/python/src/trust_types.py @@ -0,0 +1,224 @@ +"""SEP-1913 Trust & Sensitivity Annotation types — 1:1 mapping to the spec. + +All types are immutable (frozen dataclasses / enums). This module is the +single source of truth for SEP-1913's schema in Python. If the SEP changes, +update HERE first, then propagate to the rest of the SDK. + +Wire format mapping (Python → JSON): + Destination.EPHEMERAL → "ephemeral" + SimpleDataClass.PII → "pii" + Regulated(("HIPAA",)) → {"regulated": {"scopes": ["HIPAA"]}} +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Union + + +# --------------------------------------------------------------------------- +# Enums — map directly to SEP-1913 string unions +# --------------------------------------------------------------------------- + +class Destination(str, Enum): + """Where tool inputs are sent / stored (InputMetadata.destination).""" + EPHEMERAL = "ephemeral" + SYSTEM = "system" + USER = "user" + INTERNAL = "internal" + PUBLIC = "public" + + +class Source(str, Enum): + """Where tool outputs originate from (ReturnMetadata.source).""" + UNTRUSTED_PUBLIC = "untrustedPublic" + TRUSTED_PUBLIC = "trustedPublic" + INTERNAL = "internal" + USER = "user" + SYSTEM = "system" + + +class Outcome(str, Enum): + """Consequence severity of a tool invocation (InputMetadata.outcomes).""" + BENIGN = "benign" + CONSEQUENTIAL = "consequential" + IRREVERSIBLE = "irreversible" + + +class SimpleDataClass(str, Enum): + """Built-in data classification categories.""" + NONE = "none" + USER = "user" + PII = "pii" + FINANCIAL = "financial" + CREDENTIALS = "credentials" + + +# --------------------------------------------------------------------------- +# Composite types +# --------------------------------------------------------------------------- + +@dataclass(frozen=True) +class RegulatoryScope: + """A set of regulatory compliance scopes (e.g. HIPAA, GDPR, PCI-DSS). + + Stored as a tuple for immutability and hashability. + """ + scopes: tuple[str, ...] + + def __post_init__(self) -> None: + if not self.scopes: + raise ValueError("RegulatoryScope must have at least one scope") + + +@dataclass(frozen=True) +class Regulated: + """DataClass variant for regulated data — wraps RegulatoryScope. + + Wire format: {"regulated": {"scopes": ["HIPAA", "SOX"]}} + """ + regulated: RegulatoryScope + + @classmethod + def of(cls, *scopes: str) -> Regulated: + """Convenience constructor: Regulated.of("HIPAA", "SOX").""" + return cls(RegulatoryScope(scopes)) + + +# The union type matching SEP-1913's DataClass +DataClass = Union[SimpleDataClass, Regulated] + + +# --------------------------------------------------------------------------- +# Sensitivity ordering (for propagation escalation) +# --------------------------------------------------------------------------- + +_SIMPLE_ORDER: dict[SimpleDataClass, int] = { + SimpleDataClass.NONE: 0, + SimpleDataClass.USER: 1, + SimpleDataClass.PII: 2, + SimpleDataClass.FINANCIAL: 3, + SimpleDataClass.CREDENTIALS: 4, +} + +_REGULATED_ORDER = 5 # regulated(*) is always highest + + +def sensitivity_level(dc: DataClass) -> int: + """Return a numeric ordering for DataClass values. + + Used by the session propagation tracker to compute + session-max sensitivity. + + Ordering: none(0) < user(1) < pii(2) < financial(3) + < credentials(4) < regulated(*)(5) + """ + if isinstance(dc, SimpleDataClass): + return _SIMPLE_ORDER[dc] + return _REGULATED_ORDER + + +def max_sensitivity(a: DataClass, b: DataClass) -> DataClass: + """Return the higher-sensitivity DataClass value. + + For two Regulated values, merges their scopes. + """ + la, lb = sensitivity_level(a), sensitivity_level(b) + if la > lb: + return a + if lb > la: + return b + # Equal level — if both regulated, merge scopes + if isinstance(a, Regulated) and isinstance(b, Regulated): + merged = set(a.regulated.scopes) | set(b.regulated.scopes) + return Regulated(RegulatoryScope(tuple(sorted(merged)))) + return a # same simple level, same value + + +# --------------------------------------------------------------------------- +# Action Security Metadata +# --------------------------------------------------------------------------- + +@dataclass(frozen=True) +class InputMetadata: + """Describes where tool inputs go and their security implications. + + Corresponds to SEP-1913's InputMetadata interface. + Fields accept either a single value or a tuple of values + (representing a possible set for tools/list declarations). + """ + destination: Destination | tuple[Destination, ...] = Destination.EPHEMERAL + sensitivity: DataClass | tuple[DataClass, ...] = SimpleDataClass.NONE + outcomes: Outcome | tuple[Outcome, ...] = Outcome.BENIGN + + +@dataclass(frozen=True) +class ReturnMetadata: + """Describes where tool outputs come from and their sensitivity. + + Corresponds to SEP-1913's ReturnMetadata interface. + """ + source: Source | tuple[Source, ...] = Source.SYSTEM + sensitivity: DataClass | tuple[DataClass, ...] = SimpleDataClass.NONE + + +# --------------------------------------------------------------------------- +# Trust Annotations (extends ToolAnnotations) +# --------------------------------------------------------------------------- + +@dataclass(frozen=True) +class TrustAnnotations: + """SEP-1913 trust extension fields for ToolAnnotations. + + These are declared on tools/list and narrowed on tools/resolve. + """ + malicious_activity_hint: bool | None = None + attribution: tuple[str, ...] = () + input_metadata: InputMetadata | None = None + return_metadata: ReturnMetadata | None = None + + +# --------------------------------------------------------------------------- +# Result & Session Annotations (propagation layer) +# --------------------------------------------------------------------------- + +@dataclass(frozen=True) +class ResultAnnotations: + """Annotations returned in CallToolResult._meta.annotations. + + Immutable — produced per tool call, then merged into + SessionAnnotations via propagation rules. + """ + open_world_hint: bool = False + malicious_activity_hint: bool = False + attribution: tuple[str, ...] = () + sensitivity: DataClass = SimpleDataClass.NONE + + +@dataclass +class SessionAnnotations: + """Accumulated annotations across an agent session. + + Mutable — grows monotonically as the session progresses. + Implements SEP-1913 propagation rules: + - Boolean union for openWorldHint + - Boolean union for maliciousActivityHint + - Attribution accumulation (set union) + - Sensitivity escalation (session max) + """ + open_world_hint: bool = False + malicious_activity_hint: bool = False + attribution: set[str] = field(default_factory=set) + max_sensitivity: DataClass = SimpleDataClass.NONE + + def merge(self, result: ResultAnnotations) -> None: + """Apply propagation rules from a tool call result.""" + if result.open_world_hint: + self.open_world_hint = True + if result.malicious_activity_hint: + self.malicious_activity_hint = True + self.attribution.update(result.attribution) + self.max_sensitivity = max_sensitivity( + self.max_sensitivity, result.sensitivity + ) diff --git a/sdk/python/tests/__init__.py b/sdk/python/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sdk/python/tests/test_annotate.py b/sdk/python/tests/test_annotate.py new file mode 100644 index 0000000..b0483f9 --- /dev/null +++ b/sdk/python/tests/test_annotate.py @@ -0,0 +1,283 @@ +"""Tests for mcp_trust.annotate — decorator and wire-format serialization.""" + +import asyncio +import pytest +from trust_types import ( + Destination, + InputMetadata, + Outcome, + Regulated, + ReturnMetadata, + SimpleDataClass, + Source, + TrustAnnotations, +) +from annotate import ( + from_wire, + get_trust_annotations, + is_trust_annotated, + to_wire, + trust_annotated, +) + + +# --------------------------------------------------------------------------- +# Decorator basics +# --------------------------------------------------------------------------- + +class TestDecorator: + def test_sync_function_preserves_behavior(self): + @trust_annotated(attribution=("test",)) + def greet(name: str) -> str: + return f"Hello, {name}" + + assert greet("World") == "Hello, World" + + def test_async_function_preserves_behavior(self): + @trust_annotated(attribution=("test",)) + async def greet(name: str) -> str: + return f"Hello, {name}" + + result = asyncio.get_event_loop().run_until_complete(greet("World")) + assert result == "Hello, World" + + def test_annotations_attached(self): + @trust_annotated( + malicious_activity_hint=True, + attribution=("EHR",), + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=SimpleDataClass.PII, + ), + ) + def patient_lookup(pid: str) -> dict: + return {} + + ann = get_trust_annotations(patient_lookup) + assert ann is not None + assert ann.malicious_activity_hint is True + assert ann.attribution == ("EHR",) + assert ann.return_metadata.source == Source.INTERNAL + assert ann.return_metadata.sensitivity == SimpleDataClass.PII + + def test_is_trust_annotated(self): + @trust_annotated() + def foo(): + pass + + def bar(): + pass + + assert is_trust_annotated(foo) is True + assert is_trust_annotated(bar) is False + + def test_unannotated_returns_none(self): + def plain(): + pass + + assert get_trust_annotations(plain) is None + + def test_preserves_name_and_docstring(self): + @trust_annotated() + def my_tool(): + """My docstring.""" + pass + + assert my_tool.__name__ == "my_tool" + assert my_tool.__doc__ == "My docstring." + + def test_exception_propagation(self): + @trust_annotated() + def failing(): + raise ValueError("boom") + + with pytest.raises(ValueError, match="boom"): + failing() + + +# --------------------------------------------------------------------------- +# Wire-format serialization +# --------------------------------------------------------------------------- + +class TestToWire: + def test_empty_annotations(self): + ta = TrustAnnotations() + assert to_wire(ta) == {} + + def test_malicious_hint(self): + ta = TrustAnnotations(malicious_activity_hint=True) + assert to_wire(ta) == {"maliciousActivityHint": True} + + def test_attribution(self): + ta = TrustAnnotations(attribution=("EHR", "Billing")) + wire = to_wire(ta) + assert wire == {"attribution": ["EHR", "Billing"]} + + def test_return_metadata_simple(self): + ta = TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=SimpleDataClass.PII, + ) + ) + wire = to_wire(ta) + assert wire == { + "returnMetadata": { + "source": "internal", + "sensitivity": "pii", + } + } + + def test_return_metadata_regulated(self): + ta = TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA", "SOX"), + ) + ) + wire = to_wire(ta) + assert wire["returnMetadata"]["sensitivity"] == { + "regulated": {"scopes": ["HIPAA", "SOX"]} + } + + def test_input_metadata_simple(self): + ta = TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.PUBLIC, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.IRREVERSIBLE, + ) + ) + wire = to_wire(ta) + assert wire == { + "inputMetadata": { + "destination": "public", + "sensitivity": "none", + "outcomes": "irreversible", + } + } + + def test_input_metadata_arrays(self): + ta = TrustAnnotations( + input_metadata=InputMetadata( + destination=(Destination.INTERNAL, Destination.PUBLIC), + sensitivity=(SimpleDataClass.PII, SimpleDataClass.FINANCIAL), + outcomes=(Outcome.CONSEQUENTIAL, Outcome.IRREVERSIBLE), + ) + ) + wire = to_wire(ta) + im = wire["inputMetadata"] + assert im["destination"] == ["internal", "public"] + assert im["sensitivity"] == ["pii", "financial"] + assert im["outcomes"] == ["consequential", "irreversible"] + + def test_full_annotations(self): + ta = TrustAnnotations( + malicious_activity_hint=False, + attribution=("EHR-v3",), + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=SimpleDataClass.PII, + outcomes=Outcome.CONSEQUENTIAL, + ), + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + ) + wire = to_wire(ta) + assert "maliciousActivityHint" in wire + assert "attribution" in wire + assert "inputMetadata" in wire + assert "returnMetadata" in wire + + +# --------------------------------------------------------------------------- +# Wire-format deserialization +# --------------------------------------------------------------------------- + +class TestFromWire: + def test_empty(self): + ta = from_wire({}) + assert ta == TrustAnnotations() + + def test_simple_sensitivity(self): + ta = from_wire({ + "returnMetadata": { + "source": "internal", + "sensitivity": "pii", + } + }) + assert ta.return_metadata is not None + assert ta.return_metadata.source == Source.INTERNAL + assert ta.return_metadata.sensitivity == SimpleDataClass.PII + + def test_regulated_sensitivity(self): + ta = from_wire({ + "returnMetadata": { + "source": "system", + "sensitivity": {"regulated": {"scopes": ["HIPAA"]}}, + } + }) + assert isinstance(ta.return_metadata.sensitivity, Regulated) + assert ta.return_metadata.sensitivity.regulated.scopes == ("HIPAA",) + + def test_roundtrip(self): + original = TrustAnnotations( + malicious_activity_hint=True, + attribution=("EHR",), + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=SimpleDataClass.FINANCIAL, + outcomes=Outcome.CONSEQUENTIAL, + ), + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("PCI-DSS"), + ), + ) + wire = to_wire(original) + restored = from_wire(wire) + + assert restored.malicious_activity_hint == original.malicious_activity_hint + assert restored.attribution == original.attribution + assert restored.input_metadata.destination == original.input_metadata.destination + assert restored.return_metadata.sensitivity == original.return_metadata.sensitivity + + +# --------------------------------------------------------------------------- +# Wire-format validation +# --------------------------------------------------------------------------- + +class TestFromWireValidation: + def test_rejects_non_dict(self): + with pytest.raises(ValueError, match="expects a dict"): + from_wire("not a dict") + + def test_rejects_non_bool_malicious_hint(self): + with pytest.raises(ValueError, match="maliciousActivityHint must be bool"): + from_wire({"maliciousActivityHint": "yes"}) + + def test_rejects_non_array_attribution(self): + with pytest.raises(ValueError, match="attribution must be an array"): + from_wire({"attribution": "not-an-array"}) + + def test_rejects_non_string_attribution_items(self): + with pytest.raises(ValueError, match="attribution\\[1\\] must be a string"): + from_wire({"attribution": ["ok", 42]}) + + def test_rejects_unknown_sensitivity_value(self): + with pytest.raises(ValueError, match="Unknown sensitivity value"): + from_wire({"returnMetadata": {"sensitivity": "top_secret"}}) + + def test_rejects_malformed_regulated_no_scopes(self): + with pytest.raises(ValueError, match="regulated must be an object with 'scopes'"): + from_wire({"returnMetadata": {"sensitivity": {"regulated": "bad"}}}) + + def test_rejects_non_array_regulated_scopes(self): + with pytest.raises(ValueError, match="regulated.scopes must be an array"): + from_wire({"returnMetadata": {"sensitivity": {"regulated": {"scopes": "HIPAA"}}}}) + + def test_rejects_non_string_scope_items(self): + with pytest.raises(ValueError, match="regulated.scopes\\[0\\] must be a string"): + from_wire({"returnMetadata": {"sensitivity": {"regulated": {"scopes": [123]}}}}) diff --git a/sdk/python/tests/test_emit.py b/sdk/python/tests/test_emit.py new file mode 100644 index 0000000..fc6469b --- /dev/null +++ b/sdk/python/tests/test_emit.py @@ -0,0 +1,78 @@ +"""Tests for mcp_trust.emit — structured audit logging.""" + +import json +from trust_types import ( + ReturnMetadata, + SimpleDataClass, + Source, + TrustAnnotations, +) +from emit import disable_logging, enable_logging + + +class TestEmitLogging: + def test_enable_disable_cycle(self): + """Logging can be enabled and disabled without error.""" + events = [] + enable_logging(callback=lambda e: events.append(e)) + disable_logging() + assert events == [] + + def test_callback_receives_events(self): + events = [] + enable_logging(callback=lambda e: events.append(e)) + try: + from emit import _emit_call, _emit_result + ann = TrustAnnotations(attribution=("test",)) + _emit_call("my_tool", ann) + _emit_result("my_tool", ann, 42.5, "ok") + assert len(events) == 2 + assert events[0]["event"] == "tool.call" + assert events[0]["tool"] == "my_tool" + assert events[1]["event"] == "tool.result" + assert events[1]["duration_ms"] == 42.5 + finally: + disable_logging() + + def test_agent_id_in_events(self): + events = [] + enable_logging(callback=lambda e: events.append(e), agent_id="urn:agent:test") + try: + from emit import _emit_call + _emit_call("tool", TrustAnnotations()) + assert events[0]["agent"] == "urn:agent:test" + finally: + disable_logging() + + def test_no_events_when_disabled(self): + disable_logging() + from emit import _emit_call + # Should not raise + _emit_call("tool", TrustAnnotations()) + + def test_stream_output(self): + import io + buf = io.StringIO() + enable_logging(stream=buf) + try: + from emit import _emit_call + _emit_call("my_tool", TrustAnnotations()) + output = buf.getvalue() + parsed = json.loads(output.strip()) + assert parsed["event"] == "tool.call" + assert parsed["tool"] == "my_tool" + finally: + disable_logging() + + def test_policy_event(self): + events = [] + enable_logging(callback=lambda e: events.append(e)) + try: + from emit import _emit_policy + _emit_policy("tool", "block", "test reason", rule="test-rule") + assert len(events) == 1 + assert events[0]["event"] == "policy.decision" + assert events[0]["effect"] == "block" + assert events[0]["rule"] == "test-rule" + finally: + disable_logging() diff --git a/sdk/python/tests/test_policy.py b/sdk/python/tests/test_policy.py new file mode 100644 index 0000000..c9baf71 --- /dev/null +++ b/sdk/python/tests/test_policy.py @@ -0,0 +1,346 @@ +"""Tests for mcp_trust.policy — policy engine and default rules.""" + +import pytest +from trust_types import ( + Regulated, + SessionAnnotations, + SimpleDataClass, + TrustAnnotations, + ReturnMetadata, + InputMetadata, + Source, + Destination, + Outcome, +) +from annotate import to_wire +from policy import ( + BlockCredentialsToPublic, + BlockMaliciousActivity, + BlockOpenWorldToSensitiveWrite, + BlockRegulatedToExternal, + EscalateSessionSensitivity, + PolicyDecision, + PolicyEngine, + WarnOnPIIForward, +) + + +# --------------------------------------------------------------------------- +# Individual rules +# --------------------------------------------------------------------------- + +class TestBlockCredentialsToPublic: + def test_blocks_credentials_to_public(self): + rule = BlockCredentialsToPublic() + ann = to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.SYSTEM, + sensitivity=SimpleDataClass.CREDENTIALS, + ) + )) + decision = rule.evaluate(ann, "call", "public", None) + assert decision is not None + assert decision.allowed is False + assert decision.effect == "block" + + def test_allows_credentials_to_internal(self): + rule = BlockCredentialsToPublic() + ann = to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.SYSTEM, + sensitivity=SimpleDataClass.CREDENTIALS, + ) + )) + decision = rule.evaluate(ann, "call", "internal", None) + assert decision is None # no opinion + + +class TestBlockRegulatedToExternal: + def test_blocks_regulated_to_public(self): + rule = BlockRegulatedToExternal() + ann = to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ) + )) + decision = rule.evaluate(ann, "call", "public", None) + assert decision is not None + assert decision.allowed is False + assert "HIPAA" in decision.regulations + + def test_allows_regulated_to_internal(self): + rule = BlockRegulatedToExternal() + ann = to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ) + )) + decision = rule.evaluate(ann, "call", "internal", None) + assert decision is None + + +class TestBlockMaliciousActivity: + def test_blocks_malicious(self): + rule = BlockMaliciousActivity() + decision = rule.evaluate( + {"maliciousActivityHint": True}, "call", None, None + ) + assert decision is not None + assert decision.allowed is False + + def test_allows_non_malicious(self): + rule = BlockMaliciousActivity() + decision = rule.evaluate({}, "call", None, None) + assert decision is None + + +class TestWarnOnPIIForward: + def test_warns_on_pii_forward(self): + rule = WarnOnPIIForward() + ann = to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=SimpleDataClass.PII, + ) + )) + decision = rule.evaluate(ann, "forward", None, None) + assert decision is not None + assert decision.effect == "warn" + assert decision.allowed is True + + def test_no_warn_on_call(self): + rule = WarnOnPIIForward() + ann = to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=SimpleDataClass.PII, + ) + )) + decision = rule.evaluate(ann, "call", None, None) + assert decision is None + + +class TestBlockOpenWorldToSensitiveWrite: + def test_blocks_destructive_write_after_open_world(self): + """Core scenario: untrusted content in session + destructive + sensitive input → BLOCK.""" + rule = BlockOpenWorldToSensitiveWrite() + ann = to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + outcomes=Outcome.CONSEQUENTIAL, + ), + )) + ann["destructiveHint"] = True + session = SessionAnnotations() + session.open_world_hint = True + decision = rule.evaluate(ann, "call", "internal", session) + assert decision is not None + assert decision.allowed is False + assert decision.effect == "block" + assert "open-world" in decision.reason + + def test_allows_write_without_open_world(self): + """No open-world content in session → no opinion.""" + rule = BlockOpenWorldToSensitiveWrite() + ann = to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + outcomes=Outcome.CONSEQUENTIAL, + ), + )) + ann["destructiveHint"] = True + session = SessionAnnotations() + session.open_world_hint = False + decision = rule.evaluate(ann, "call", "internal", session) + assert decision is None + + def test_allows_read_only_tool_after_open_world(self): + """Open-world session but tool is read-only (not destructive) → no opinion.""" + rule = BlockOpenWorldToSensitiveWrite() + ann = to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + )) + # No destructiveHint + session = SessionAnnotations() + session.open_world_hint = True + decision = rule.evaluate(ann, "call", "internal", session) + assert decision is None + + def test_allows_destructive_below_threshold(self): + """Destructive + open-world but sensitivity below threshold → no opinion.""" + rule = BlockOpenWorldToSensitiveWrite() + ann = to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=SimpleDataClass.USER, + outcomes=Outcome.CONSEQUENTIAL, + ), + )) + ann["destructiveHint"] = True + session = SessionAnnotations() + session.open_world_hint = True + decision = rule.evaluate(ann, "call", "internal", session) + assert decision is None # USER < PII threshold + + def test_custom_threshold(self): + """Custom threshold still blocks at that level.""" + rule = BlockOpenWorldToSensitiveWrite(threshold=SimpleDataClass.FINANCIAL) + ann = to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=SimpleDataClass.PII, + outcomes=Outcome.CONSEQUENTIAL, + ), + )) + ann["destructiveHint"] = True + session = SessionAnnotations() + session.open_world_hint = True + decision = rule.evaluate(ann, "call", "internal", session) + assert decision is None # PII < FINANCIAL threshold + + def test_no_session(self): + """No session provided → no opinion.""" + rule = BlockOpenWorldToSensitiveWrite() + ann = {"destructiveHint": True, "inputMetadata": {"sensitivity": "credentials"}} + decision = rule.evaluate(ann, "call", "internal", None) + assert decision is None + + +class TestEscalateSessionSensitivity: + def test_escalates_when_session_has_pii_and_target_public(self): + rule = EscalateSessionSensitivity() + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.PII + decision = rule.evaluate({}, "call", "public", session) + assert decision is not None + assert decision.effect == "escalate" + assert decision.allowed is False + + def test_no_escalation_for_internal(self): + rule = EscalateSessionSensitivity() + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.PII + decision = rule.evaluate({}, "call", "internal", session) + assert decision is None + + def test_no_escalation_below_threshold(self): + rule = EscalateSessionSensitivity() + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.USER + decision = rule.evaluate({}, "call", "public", session) + assert decision is None + + def test_custom_threshold(self): + rule = EscalateSessionSensitivity(threshold=SimpleDataClass.FINANCIAL) + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.PII + decision = rule.evaluate({}, "call", "public", session) + assert decision is None # PII < FINANCIAL + + +# --------------------------------------------------------------------------- +# PolicyEngine modes +# --------------------------------------------------------------------------- + +class TestPolicyEngine: + def _make_engine(self, mode: str) -> PolicyEngine: + engine = PolicyEngine(mode=mode) + engine.register_tool("send_email", to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.PUBLIC, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.CONSEQUENTIAL, + ), + ))) + engine.register_tool("patient_lookup", to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + ))) + return engine + + def test_enforce_blocks(self): + engine = self._make_engine("enforce") + decision = engine.evaluate( + "patient_lookup", action="call", target_destination="public" + ) + assert decision.allowed is False + assert decision.effect == "block" + + def test_warn_converts_block_to_warn(self): + engine = self._make_engine("warn") + decision = engine.evaluate( + "patient_lookup", action="call", target_destination="public" + ) + assert decision.allowed is True + assert decision.effect == "warn" + + def test_audit_allows_everything(self): + engine = self._make_engine("audit") + decision = engine.evaluate( + "patient_lookup", action="call", target_destination="public" + ) + assert decision.allowed is True + assert decision.effect == "audit" + + def test_no_rule_triggered(self): + engine = self._make_engine("enforce") + decision = engine.evaluate("send_email", action="call", target_destination="internal") + assert decision.allowed is True + assert decision.effect == "allow" + + def test_invalid_mode_raises(self): + with pytest.raises(ValueError, match="Invalid mode"): + PolicyEngine(mode="invalid") + + def test_register_trust_annotations(self): + engine = PolicyEngine(mode="enforce") + ta = TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.SYSTEM, + sensitivity=SimpleDataClass.CREDENTIALS, + ) + ) + engine.register_trust_annotations("secret_tool", ta) + decision = engine.evaluate("secret_tool", target_destination="public") + assert decision.allowed is False + + def test_session_escalation(self): + engine = self._make_engine("enforce") + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.PII + + # send_email to public with PII session → escalate + decision = engine.evaluate( + "send_email", action="call", + target_destination="public", + session=session, + ) + assert decision.effect == "escalate" + + def test_default_mode_is_warn(self): + engine = PolicyEngine() + assert engine.mode == "warn" + + def test_fail_closed_blocks_unregistered_tool(self): + engine = PolicyEngine(mode="enforce", fail_closed=True) + engine.register_tool("known_tool", {}) + decision = engine.evaluate("unknown_tool", action="call") + assert decision.allowed is False + assert decision.effect == "block" + assert "Unregistered tool" in decision.reason + assert decision.rule == "fail-closed" + + def test_fail_open_allows_unregistered_tool(self): + engine = PolicyEngine(mode="enforce", fail_closed=False) + decision = engine.evaluate("unknown_tool", action="call") + assert decision.allowed is True + assert decision.effect == "allow" diff --git a/sdk/python/tests/test_propagate.py b/sdk/python/tests/test_propagate.py new file mode 100644 index 0000000..f45c539 --- /dev/null +++ b/sdk/python/tests/test_propagate.py @@ -0,0 +1,132 @@ +"""Tests for mcp_trust.propagate — session propagation tracker.""" + +import pytest +from trust_types import ( + Regulated, + ResultAnnotations, + SimpleDataClass, +) +from propagate import SessionTracker + + +class TestSessionTracker: + def test_initial_state(self): + tracker = SessionTracker(session_id="test-session") + assert tracker.session_id == "test-session" + assert tracker.call_count == 0 + assert tracker.has_seen_sensitive_data() is False + + def test_merge_escalates_sensitivity(self): + tracker = SessionTracker() + tracker.merge( + ResultAnnotations(sensitivity=SimpleDataClass.PII), + tool_name="patient_lookup", + ) + assert tracker.session.max_sensitivity == SimpleDataClass.PII + assert tracker.call_count == 1 + assert tracker.has_seen_sensitive_data() is True + + def test_cannot_deescalate(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.CREDENTIALS)) + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.NONE)) + assert tracker.session.max_sensitivity == SimpleDataClass.CREDENTIALS + + def test_escalation_to_regulated(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII)) + tracker.merge(ResultAnnotations(sensitivity=Regulated.of("HIPAA"))) + assert isinstance(tracker.session.max_sensitivity, Regulated) + + def test_attribution_accumulates(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(attribution=("EHR",))) + tracker.merge(ResultAnnotations(attribution=("Billing", "EHR"))) + assert tracker.session.attribution == {"EHR", "Billing"} + + def test_open_world_hint_sticky(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(open_world_hint=True)) + tracker.merge(ResultAnnotations(open_world_hint=False)) + assert tracker.session.open_world_hint is True + + def test_sensitivity_at_least(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII)) + assert tracker.sensitivity_at_least(SimpleDataClass.USER) is True + assert tracker.sensitivity_at_least(SimpleDataClass.PII) is True + assert tracker.sensitivity_at_least(SimpleDataClass.CREDENTIALS) is False + + def test_get_regulatory_scopes_empty(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII)) + assert tracker.get_regulatory_scopes() == () + + def test_get_regulatory_scopes(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(sensitivity=Regulated.of("HIPAA", "SOX"))) + assert set(tracker.get_regulatory_scopes()) == {"HIPAA", "SOX"} + + def test_to_wire(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations( + open_world_hint=True, + attribution=("EHR",), + sensitivity=SimpleDataClass.PII, + )) + wire = tracker.to_wire() + assert wire["openWorldHint"] is True + assert wire["attribution"] == ["EHR"] + assert wire["maxSensitivity"] == "pii" + + def test_to_wire_empty(self): + tracker = SessionTracker() + assert tracker.to_wire() == {} + + def test_summary(self): + tracker = SessionTracker(session_id="s1") + tracker.merge( + ResultAnnotations(sensitivity=SimpleDataClass.PII), + tool_name="patient_lookup", + ) + summary = tracker.summary() + assert "s1" in summary + assert "patient_lookup" in summary + assert "pii" in summary.lower() or "PII" in summary + + def test_reset_blocked_by_default(self): + tracker = SessionTracker() + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.CREDENTIALS)) + with pytest.raises(RuntimeError, match="reset\\(\\) is disabled"): + tracker.reset() + + def test_reset_allowed_when_opted_in(self): + tracker = SessionTracker(allow_reset=True) + tracker.merge(ResultAnnotations( + sensitivity=SimpleDataClass.CREDENTIALS, + open_world_hint=True, + )) + tracker.reset() + assert tracker.call_count == 0 + assert tracker.has_seen_sensitive_data() is False + assert tracker.session.open_world_hint is False + + def test_call_count(self): + tracker = SessionTracker() + for i in range(5): + tracker.merge(ResultAnnotations()) + assert tracker.call_count == 5 + + def test_session_returns_snapshot_not_reference(self): + """Mutations on the returned session must not affect the tracker.""" + tracker = SessionTracker() + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII)) + snapshot = tracker.session + # Mutate the snapshot + snapshot.max_sensitivity = SimpleDataClass.NONE + snapshot.open_world_hint = True + snapshot.attribution.add("injected") + # Tracker state must be unchanged + assert tracker.session.max_sensitivity == SimpleDataClass.PII + assert tracker.session.open_world_hint is False + assert "injected" not in tracker.session.attribution diff --git a/sdk/python/tests/test_types.py b/sdk/python/tests/test_types.py new file mode 100644 index 0000000..6ebda1d --- /dev/null +++ b/sdk/python/tests/test_types.py @@ -0,0 +1,213 @@ +"""Tests for mcp_trust.types — SEP-1913 type definitions.""" + +import pytest +from trust_types import ( + DataClass, + Destination, + InputMetadata, + Outcome, + Regulated, + RegulatoryScope, + ResultAnnotations, + ReturnMetadata, + SessionAnnotations, + SimpleDataClass, + Source, + TrustAnnotations, + max_sensitivity, + sensitivity_level, +) + + +# --------------------------------------------------------------------------- +# Enum values match SEP-1913 wire format +# --------------------------------------------------------------------------- + +class TestEnumValues: + def test_destination_values(self): + assert Destination.EPHEMERAL.value == "ephemeral" + assert Destination.SYSTEM.value == "system" + assert Destination.USER.value == "user" + assert Destination.INTERNAL.value == "internal" + assert Destination.PUBLIC.value == "public" + + def test_source_values(self): + assert Source.UNTRUSTED_PUBLIC.value == "untrustedPublic" + assert Source.TRUSTED_PUBLIC.value == "trustedPublic" + assert Source.INTERNAL.value == "internal" + assert Source.USER.value == "user" + assert Source.SYSTEM.value == "system" + + def test_outcome_values(self): + assert Outcome.BENIGN.value == "benign" + assert Outcome.CONSEQUENTIAL.value == "consequential" + assert Outcome.IRREVERSIBLE.value == "irreversible" + + def test_simple_data_class_values(self): + assert SimpleDataClass.NONE.value == "none" + assert SimpleDataClass.USER.value == "user" + assert SimpleDataClass.PII.value == "pii" + assert SimpleDataClass.FINANCIAL.value == "financial" + assert SimpleDataClass.CREDENTIALS.value == "credentials" + + +# --------------------------------------------------------------------------- +# Immutability +# --------------------------------------------------------------------------- + +class TestImmutability: + def test_regulatory_scope_frozen(self): + rs = RegulatoryScope(("HIPAA",)) + with pytest.raises(AttributeError): + rs.scopes = ("GDPR",) # type: ignore + + def test_regulated_frozen(self): + r = Regulated.of("HIPAA") + with pytest.raises(AttributeError): + r.regulated = RegulatoryScope(("GDPR",)) # type: ignore + + def test_input_metadata_frozen(self): + im = InputMetadata() + with pytest.raises(AttributeError): + im.destination = Destination.PUBLIC # type: ignore + + def test_return_metadata_frozen(self): + rm = ReturnMetadata() + with pytest.raises(AttributeError): + rm.source = Source.UNTRUSTED_PUBLIC # type: ignore + + def test_trust_annotations_frozen(self): + ta = TrustAnnotations() + with pytest.raises(AttributeError): + ta.malicious_activity_hint = True # type: ignore + + def test_result_annotations_frozen(self): + ra = ResultAnnotations() + with pytest.raises(AttributeError): + ra.open_world_hint = True # type: ignore + + +# --------------------------------------------------------------------------- +# Regulated convenience constructor +# --------------------------------------------------------------------------- + +class TestRegulated: + def test_of_single_scope(self): + r = Regulated.of("HIPAA") + assert r.regulated.scopes == ("HIPAA",) + + def test_of_multiple_scopes(self): + r = Regulated.of("HIPAA", "SOX", "GDPR") + assert r.regulated.scopes == ("HIPAA", "SOX", "GDPR") + + def test_empty_scope_raises(self): + with pytest.raises(ValueError, match="at least one scope"): + RegulatoryScope(()) + + +# --------------------------------------------------------------------------- +# Sensitivity ordering +# --------------------------------------------------------------------------- + +class TestSensitivityOrdering: + def test_simple_ordering(self): + levels = [sensitivity_level(dc) for dc in ( + SimpleDataClass.NONE, + SimpleDataClass.USER, + SimpleDataClass.PII, + SimpleDataClass.FINANCIAL, + SimpleDataClass.CREDENTIALS, + )] + assert levels == [0, 1, 2, 3, 4] + + def test_regulated_is_highest(self): + assert sensitivity_level(Regulated.of("HIPAA")) == 5 + assert sensitivity_level(Regulated.of("HIPAA")) > sensitivity_level(SimpleDataClass.CREDENTIALS) + + def test_max_sensitivity_picks_higher(self): + assert max_sensitivity(SimpleDataClass.NONE, SimpleDataClass.PII) == SimpleDataClass.PII + assert max_sensitivity(SimpleDataClass.PII, SimpleDataClass.NONE) == SimpleDataClass.PII + assert max_sensitivity(SimpleDataClass.CREDENTIALS, SimpleDataClass.USER) == SimpleDataClass.CREDENTIALS + + def test_max_sensitivity_regulated_wins(self): + result = max_sensitivity(SimpleDataClass.CREDENTIALS, Regulated.of("HIPAA")) + assert isinstance(result, Regulated) + + def test_max_sensitivity_merges_regulated_scopes(self): + a = Regulated.of("HIPAA") + b = Regulated.of("GDPR") + result = max_sensitivity(a, b) + assert isinstance(result, Regulated) + assert set(result.regulated.scopes) == {"GDPR", "HIPAA"} + + +# --------------------------------------------------------------------------- +# SessionAnnotations propagation +# --------------------------------------------------------------------------- + +class TestSessionAnnotations: + def test_merge_open_world(self): + session = SessionAnnotations() + assert session.open_world_hint is False + session.merge(ResultAnnotations(open_world_hint=True)) + assert session.open_world_hint is True + # Once true, stays true + session.merge(ResultAnnotations(open_world_hint=False)) + assert session.open_world_hint is True + + def test_merge_malicious(self): + session = SessionAnnotations() + session.merge(ResultAnnotations(malicious_activity_hint=True)) + assert session.malicious_activity_hint is True + + def test_merge_attribution(self): + session = SessionAnnotations() + session.merge(ResultAnnotations(attribution=("EHR",))) + session.merge(ResultAnnotations(attribution=("Billing", "EHR"))) + assert session.attribution == {"EHR", "Billing"} + + def test_merge_sensitivity_escalation(self): + session = SessionAnnotations() + assert session.max_sensitivity == SimpleDataClass.NONE + + session.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII)) + assert session.max_sensitivity == SimpleDataClass.PII + + # Can't de-escalate + session.merge(ResultAnnotations(sensitivity=SimpleDataClass.NONE)) + assert session.max_sensitivity == SimpleDataClass.PII + + # Can escalate further + session.merge(ResultAnnotations(sensitivity=SimpleDataClass.CREDENTIALS)) + assert session.max_sensitivity == SimpleDataClass.CREDENTIALS + + def test_merge_sensitivity_to_regulated(self): + session = SessionAnnotations() + session.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII)) + session.merge(ResultAnnotations(sensitivity=Regulated.of("HIPAA"))) + assert isinstance(session.max_sensitivity, Regulated) + assert session.max_sensitivity.regulated.scopes == ("HIPAA",) + + +# --------------------------------------------------------------------------- +# Defaults +# --------------------------------------------------------------------------- + +class TestDefaults: + def test_input_metadata_defaults(self): + im = InputMetadata() + assert im.destination == Destination.EPHEMERAL + assert im.sensitivity == SimpleDataClass.NONE + assert im.outcomes == Outcome.BENIGN + + def test_return_metadata_defaults(self): + rm = ReturnMetadata() + assert rm.source == Source.SYSTEM + assert rm.sensitivity == SimpleDataClass.NONE + + def test_trust_annotations_defaults(self): + ta = TrustAnnotations() + assert ta.malicious_activity_hint is None + assert ta.attribution == () + assert ta.input_metadata is None + assert ta.return_metadata is None diff --git a/sdk/python/tests/test_usability_scenarios.py b/sdk/python/tests/test_usability_scenarios.py new file mode 100644 index 0000000..cbceb29 --- /dev/null +++ b/sdk/python/tests/test_usability_scenarios.py @@ -0,0 +1,762 @@ +"""SEP-1913 Usability Scenario Test Suite. + +Each scenario class combines a *participant worksheet* (the class docstring) +with *SDK oracle tests* (the test methods). To run a spec-clarity study: + +1. Hand participants USABILITY_SCENARIOS.md + the SEP-1913 spec. +2. Collect their answers. +3. Run this file to produce the correct answers. +4. Compute the delta — low agreement highlights spec ambiguities. + +Run: + cd packages/python + PYTHONPATH=src pytest tests/test_usability_scenarios.py -v +""" + +from trust_types import ( + Destination, + InputMetadata, + Outcome, + Regulated, + ResultAnnotations, + ReturnMetadata, + SessionAnnotations, + SimpleDataClass, + Source, + TrustAnnotations, + sensitivity_level, +) +from propagate import SessionTracker +from policy import PolicyEngine, PolicyDecision +from annotate import to_wire, from_wire + + +# =================================================================== +# Scenario 1 — Classification: Simple Email Send Tool +# Spec section: InputMetadata, ReturnMetadata, Destination, Outcome, Source +# Expected agreement: HIGH +# =================================================================== + +class TestScenario01_SimpleToolClassification: + """SCENARIO 1 — Simple: Email Send Tool + + Context: + You are implementing an MCP server that exposes a `send_email` tool. + This tool accepts a recipient address and message body, sends the + email via SMTP, and returns a delivery receipt from the mail server. + + Questions: + Q1. What Destination should `send_email` use? + a) ephemeral b) system c) user d) internal e) public + + Q2. What Outcome should `send_email` use? + a) benign b) consequential c) irreversible + + Q3. What Source should the return metadata use? + a) untrustedPublic b) trustedPublic c) internal d) user e) system + + Your answers: + Q1: ___ Q2: ___ Q3: ___ + """ + + def test_destination_is_public(self): + """Email leaves the organization boundary → public.""" + meta = InputMetadata( + destination=Destination.PUBLIC, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.IRREVERSIBLE, + ) + assert meta.destination == Destination.PUBLIC + + def test_outcome_is_irreversible(self): + """Once sent, an email cannot be recalled → irreversible.""" + meta = InputMetadata( + destination=Destination.PUBLIC, + outcomes=Outcome.IRREVERSIBLE, + ) + assert meta.outcomes == Outcome.IRREVERSIBLE + + def test_source_is_system(self): + """Delivery receipt comes from the mail server (trusted infra) → system.""" + meta = ReturnMetadata( + source=Source.SYSTEM, + sensitivity=SimpleDataClass.NONE, + ) + assert meta.source == Source.SYSTEM + + def test_full_annotation_roundtrip(self): + """Verify the complete annotation serializes correctly.""" + ann = TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.PUBLIC, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.IRREVERSIBLE, + ), + return_metadata=ReturnMetadata( + source=Source.SYSTEM, + sensitivity=SimpleDataClass.NONE, + ), + ) + wire = to_wire(ann) + restored = from_wire(wire) + assert restored.input_metadata.destination == Destination.PUBLIC + assert restored.input_metadata.outcomes == Outcome.IRREVERSIBLE + assert restored.return_metadata.source == Source.SYSTEM + + +# =================================================================== +# Scenario 2 — Classification: Ambiguous Save to Workspace +# Spec section: Destination enum boundary +# Expected agreement: LOW (deliberate ambiguity) +# =================================================================== + +class TestScenario02_AmbiguousSaveToWorkspace: + """SCENARIO 2 — Ambiguous: Save to Workspace + + Context: + Your tool `save_to_workspace` writes a document to a shared + team workspace. The workspace is accessible to all team + members (not just the invoking user) but is NOT accessible + outside the organization. + + Questions: + Q1. What Destination should this tool use? + a) ephemeral b) system c) user d) internal e) public + + Q2. Justify your choice — what distinguishes `user` from `internal` + in this scenario? + + Your answers: + Q1: ___ + Q2: _______________________________________________ + """ + + def test_destination_is_internal(self): + """Team-visible but org-internal → internal (not user, not public).""" + meta = InputMetadata( + destination=Destination.INTERNAL, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.CONSEQUENTIAL, + ) + assert meta.destination == Destination.INTERNAL + + def test_not_user_destination(self): + """'user' implies single-user visibility; team workspace is broader.""" + assert Destination.INTERNAL != Destination.USER + + def test_not_public_destination(self): + """Team workspace is inside the org boundary → not public.""" + assert Destination.INTERNAL != Destination.PUBLIC + + +# =================================================================== +# Scenario 3 — Propagation: Sensitivity Monotonic Escalation +# Spec section: Propagation Rules (sensitivity — GAP: not explicitly listed) +# Expected agreement: HIGH (intuitive, but spec may not state it) +# =================================================================== + +class TestScenario03_SensitivityMonotonicEscalation: + """SCENARIO 3 — Sensitivity Monotonic Escalation + + Context: + In an agent session, Tool A returns sensitivity="pii". + Then Tool B returns sensitivity="none". + + Questions: + Q1. After both calls, what is the session's max_sensitivity? + a) none b) pii + + Q2. Can a later tool call ever DECREASE the session sensitivity? + a) yes b) no + + Q3. Does the spec's Propagation Rules section explicitly define + sensitivity tracking? + a) yes b) no + + Your answers: + Q1: ___ Q2: ___ Q3: ___ + """ + + def test_session_retains_highest_sensitivity(self): + """PII followed by NONE → session max stays PII.""" + tracker = SessionTracker(session_id="test-s03") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII), "tool_a") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.NONE), "tool_b") + assert tracker.session.max_sensitivity == SimpleDataClass.PII + + def test_sensitivity_never_decreases(self): + """Monotonic escalation: sensitivity level can only go up.""" + tracker = SessionTracker(session_id="test-s03-mono") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.CREDENTIALS), "tool_x") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.USER), "tool_y") + # CREDENTIALS (level 4) > USER (level 1) → stays CREDENTIALS + assert tracker.session.max_sensitivity == SimpleDataClass.CREDENTIALS + assert sensitivity_level(tracker.session.max_sensitivity) == 4 + + def test_escalation_across_multiple_levels(self): + """NONE → USER → PII → FINANCIAL → session tracks the max at each step.""" + tracker = SessionTracker(session_id="test-s03-multi") + for sdc in [SimpleDataClass.NONE, SimpleDataClass.USER, + SimpleDataClass.PII, SimpleDataClass.FINANCIAL]: + tracker.merge(ResultAnnotations(sensitivity=sdc), "tool") + assert tracker.session.max_sensitivity == SimpleDataClass.FINANCIAL + + +# =================================================================== +# Scenario 4 — Propagation: openWorldHint Persistence +# Spec section: Propagation Rules §1 — boolean union +# Expected agreement: HIGH +# =================================================================== + +class TestScenario04_OpenWorldHintPersistence: + """SCENARIO 4 — openWorldHint Persistence + + Context: + Tool A returns openWorldHint=true. + Tool B returns openWorldHint=false. + + Questions: + Q1. After both calls, what is the session's openWorldHint? + a) true b) false + + Q2. Which propagation rule governs this? + a) boolean union b) set union c) max-sensitivity d) not defined + + Your answers: + Q1: ___ Q2: ___ + """ + + def test_open_world_hint_sticky_true(self): + """Boolean union: once true, stays true regardless of later false.""" + tracker = SessionTracker(session_id="test-s04") + tracker.merge(ResultAnnotations(open_world_hint=True), "tool_a") + tracker.merge(ResultAnnotations(open_world_hint=False), "tool_b") + assert tracker.session.open_world_hint is True + + def test_false_then_true_also_sticky(self): + """Order doesn't matter — boolean union is commutative.""" + tracker = SessionTracker(session_id="test-s04-reverse") + tracker.merge(ResultAnnotations(open_world_hint=False), "tool_b") + tracker.merge(ResultAnnotations(open_world_hint=True), "tool_a") + assert tracker.session.open_world_hint is True + + def test_all_false_stays_false(self): + """If no tool ever sets openWorldHint=true, session stays false.""" + tracker = SessionTracker(session_id="test-s04-allfalse") + tracker.merge(ResultAnnotations(open_world_hint=False), "tool_1") + tracker.merge(ResultAnnotations(open_world_hint=False), "tool_2") + assert tracker.session.open_world_hint is False + + +# =================================================================== +# Scenario 5 — Policy: Data Exfiltration (PII Session + Public Dest) +# Spec section: EscalateSessionSensitivity rule with session context +# Expected agreement: MEDIUM +# =================================================================== + +class TestScenario05_DataExfiltrationPIIToPublic: + """SCENARIO 5 — Data Exfiltration: PII Session → Public Destination + + Context: + An agent session has accumulated max_sensitivity="pii" from + earlier tool calls. The agent now wants to call `send_email`, + which is annotated with: + - inputMetadata.destination = "public" + - inputMetadata.sensitivity = "none" + - inputMetadata.outcomes = "consequential" + + The PolicyEngine is in "enforce" mode with default rules. + + Questions: + Q1. Does BlockCredentialsToPublic fire? + a) yes b) no — why? + + Q2. Does EscalateSessionSensitivity fire? + a) yes b) no — why? + + Q3. What is the expected PolicyDecision? + allowed: ___ effect: ___ rule: ___ + + Your answers: + Q1: ___ reason: _______________________________ + Q2: ___ reason: _______________________________ + Q3: allowed=___ effect="___" rule="___" + """ + + def _make_engine(self): + engine = PolicyEngine(mode="enforce") + engine.register_tool("send_email", to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.PUBLIC, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.CONSEQUENTIAL, + ), + ))) + return engine + + def test_block_credentials_does_not_fire(self): + """send_email tool sensitivity is NONE (level 0) < CREDENTIALS (level 4).""" + engine = self._make_engine() + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.PII + decision = engine.evaluate( + "send_email", action="call", + target_destination="public", session=session, + ) + # If BlockCredentialsToPublic had fired, effect would be "block" + assert decision.rule != "block-credentials-to-public" + + def test_escalate_session_sensitivity_fires(self): + """Session PII (level 2) >= threshold PII (level 2), dest=public → escalate.""" + engine = self._make_engine() + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.PII + decision = engine.evaluate( + "send_email", action="call", + target_destination="public", session=session, + ) + assert decision.allowed is False + assert decision.effect == "escalate" + assert decision.rule == "escalate-session-sensitivity" + + def test_key_insight_tool_vs_session_sensitivity(self): + """The critical distinction: tool sensitivity (NONE) vs session + sensitivity (PII). + + BlockCredentialsToPublic checks _extract_sensitivity (tool-level). + EscalateSessionSensitivity checks session.max_sensitivity. + """ + engine = self._make_engine() + # With no session, no escalation rule fires either + decision_no_session = engine.evaluate( + "send_email", action="call", target_destination="public", + ) + assert decision_no_session.allowed is True + assert decision_no_session.effect == "allow" + + +# =================================================================== +# Scenario 6 — Policy: Prompt Injection (Open World + Destructive) +# Spec section: BlockOpenWorldToSensitiveWrite rule +# Expected agreement: MEDIUM +# =================================================================== + +class TestScenario06_PromptInjectionOpenWorldDestructive: + """SCENARIO 6 — Prompt Injection: Open World + Destructive Write + + Context: + An agent session has openWorldHint=true (it read untrusted web + content earlier). The agent now wants to call + `update_patient_record`, annotated with: + - destructiveHint = true + - inputMetadata.sensitivity = regulated(["HIPAA"]) + - inputMetadata.destination = "internal" + + The PolicyEngine is in "enforce" mode with default rules. + + Questions: + Q1. Which rule fires first? + a) block-credentials-to-public + b) block-regulated-to-external + c) block-openworld-to-sensitive-write + d) escalate-session-sensitivity + + Q2. What is the expected PolicyDecision? + allowed: ___ effect: ___ rule: ___ + + Q3. Why don't the "to-public" rules fire? + + Your answers: + Q1: ___ + Q2: allowed=___ effect="___" rule="___" + Q3: _______________________________________________ + """ + + def _make_engine(self): + engine = PolicyEngine(mode="enforce") + ann = to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + outcomes=Outcome.CONSEQUENTIAL, + ), + )) + ann["destructiveHint"] = True + engine.register_tool("update_patient_record", ann) + return engine + + def test_block_openworld_fires(self): + """Open-world session + destructive + HIPAA input → blocked.""" + engine = self._make_engine() + session = SessionAnnotations() + session.open_world_hint = True + decision = engine.evaluate( + "update_patient_record", action="call", + target_destination="internal", session=session, + ) + assert decision.allowed is False + assert decision.effect == "block" + assert decision.rule == "block-openworld-to-sensitive-write" + + def test_no_block_without_open_world(self): + """Same tool but session has no open-world content → allowed.""" + engine = self._make_engine() + session = SessionAnnotations() + session.open_world_hint = False + decision = engine.evaluate( + "update_patient_record", action="call", + target_destination="internal", session=session, + ) + assert decision.allowed is True + assert decision.effect == "allow" + + def test_credentials_rule_skipped_for_internal_dest(self): + """BlockCredentialsToPublic only checks target_destination='public'.""" + engine = self._make_engine() + session = SessionAnnotations() + session.open_world_hint = True + decision = engine.evaluate( + "update_patient_record", action="call", + target_destination="internal", session=session, + ) + assert decision.rule != "block-credentials-to-public" + assert decision.rule != "block-regulated-to-external" + + +# =================================================================== +# Scenario 7 — Policy: Mode Behavior (enforce vs warn vs audit) +# Spec section: PolicyEngine mode transformation (_apply_mode) +# Expected agreement: MEDIUM +# =================================================================== + +class TestScenario07_ModeBehavior: + """SCENARIO 7 — Mode Behavior: enforce vs warn vs audit + + Context: + A `patient_lookup` tool is annotated with: + - returnMetadata.source = "internal" + - returnMetadata.sensitivity = regulated(["HIPAA"]) + + An agent calls it with target_destination="public". + The default rules are loaded. + + Questions: + Q1. Which default rule fires first? + a) block-credentials-to-public + b) block-regulated-to-external + (Hint: consider the sensitivity ordering and rule eval order) + + Q2. In enforce mode: allowed=___ effect="___" + Q3. In warn mode: allowed=___ effect="___" + Q4. In audit mode: allowed=___ effect="___" + + Your answers: + Q1: ___ + Q2: allowed=___ effect="___" + Q3: allowed=___ effect="___" + Q4: allowed=___ effect="___" + """ + + def _make_engine(self, mode): + engine = PolicyEngine(mode=mode) + engine.register_tool("patient_lookup", to_wire(TrustAnnotations( + return_metadata=ReturnMetadata( + source=Source.INTERNAL, + sensitivity=Regulated.of("HIPAA"), + ), + ))) + return engine + + def test_credentials_rule_fires_first(self): + """Regulated (level 5) >= CREDENTIALS (level 4) → BlockCredentialsToPublic fires. + + BlockCredentialsToPublic is evaluated BEFORE BlockRegulatedToExternal + in DEFAULT_RULES order, and Regulated.of("HIPAA") has level 5 which + is >= level 4 (CREDENTIALS threshold). + """ + engine = self._make_engine("enforce") + decision = engine.evaluate( + "patient_lookup", action="call", target_destination="public", + ) + assert decision.rule == "block-credentials-to-public" + + def test_enforce_blocks(self): + """Enforce mode: decision enforced as-is → blocked.""" + engine = self._make_engine("enforce") + decision = engine.evaluate( + "patient_lookup", action="call", target_destination="public", + ) + assert decision.allowed is False + assert decision.effect == "block" + + def test_warn_converts_to_warn(self): + """Warn mode: 'block' becomes 'warn', allowed=True.""" + engine = self._make_engine("warn") + decision = engine.evaluate( + "patient_lookup", action="call", target_destination="public", + ) + assert decision.allowed is True + assert decision.effect == "warn" + assert decision.rule == "block-credentials-to-public" + + def test_audit_converts_to_audit(self): + """Audit mode: everything becomes 'audit', allowed=True.""" + engine = self._make_engine("audit") + decision = engine.evaluate( + "patient_lookup", action="call", target_destination="public", + ) + assert decision.allowed is True + assert decision.effect == "audit" + assert decision.rule == "block-credentials-to-public" + + def test_warn_also_converts_escalate(self): + """Warn mode converts BOTH 'block' AND 'escalate' to 'warn'.""" + engine = PolicyEngine(mode="warn") + engine.register_tool("send_email", to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.PUBLIC, + sensitivity=SimpleDataClass.NONE, + outcomes=Outcome.CONSEQUENTIAL, + ), + ))) + session = SessionAnnotations() + session.max_sensitivity = SimpleDataClass.PII + decision = engine.evaluate( + "send_email", action="call", + target_destination="public", session=session, + ) + assert decision.allowed is True + assert decision.effect == "warn" + assert decision.rule == "escalate-session-sensitivity" + + +# =================================================================== +# Scenario 8 — Gap: Sensitivity Propagation Missing from Spec +# Spec section: Propagation Rules completeness +# Expected agreement: LOW (tests a spec gap) +# =================================================================== + +class TestScenario08_SensitivityPropagationGap: + """SCENARIO 8 — Sensitivity Propagation Gap + + Context: + Read the Propagation Rules section of SEP-1913 carefully. + It defines propagation for: + - openWorldHint (boolean union) + - attribution (set union) + + Questions: + Q1. Does the Propagation Rules section explicitly define how + sensitivity propagates across tool calls? + a) yes b) no + + Q2. Despite Q1, does the SDK implement sensitivity propagation? + a) yes b) no + + Q3. If there is a gap, what propagation rule SHOULD apply + to sensitivity? + a) boolean union b) set union c) max ordering d) not applicable + + Your answers: + Q1: ___ Q2: ___ Q3: ___ + """ + + def test_sdk_implements_sensitivity_propagation(self): + """The SDK tracks max_sensitivity even though the spec's + Propagation Rules section does not explicitly list it.""" + tracker = SessionTracker(session_id="test-s08") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII), "tool_a") + assert tracker.session.max_sensitivity == SimpleDataClass.PII + + def test_sensitivity_uses_max_ordering(self): + """SDK uses max-ordering (not boolean union or set union).""" + tracker = SessionTracker(session_id="test-s08-max") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.USER), "t1") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.FINANCIAL), "t2") + tracker.merge(ResultAnnotations(sensitivity=SimpleDataClass.PII), "t3") + # FINANCIAL (level 3) > PII (level 2) > USER (level 1) + assert tracker.session.max_sensitivity == SimpleDataClass.FINANCIAL + + def test_propagation_rules_do_define_boolean_union(self): + """Confirm that boolean union IS defined for openWorldHint + (to contrast with sensitivity which is NOT defined).""" + tracker = SessionTracker(session_id="test-s08-owh") + tracker.merge(ResultAnnotations(open_world_hint=True), "t1") + tracker.merge(ResultAnnotations(open_world_hint=False), "t2") + assert tracker.session.open_world_hint is True + + def test_propagation_rules_do_define_attribution_union(self): + """Confirm that attribution set union IS defined.""" + tracker = SessionTracker(session_id="test-s08-attr") + tracker.merge(ResultAnnotations(attribution=("A",)), "t1") + tracker.merge(ResultAnnotations(attribution=("B",)), "t2") + assert tracker.session.attribution == {"A", "B"} + + +# =================================================================== +# Scenario 9 — Gap: Host-Detected PII (Server Says "none") +# Spec section: Client Responsibilities #5, Security Implications +# Expected agreement: LOW (tests a spec gap) +# =================================================================== + +class TestScenario09_HostDetectedPII: + """SCENARIO 9 — Host-Detected PII (Server says "none") + + Context: + A tool server returns sensitivity="none" in its result annotations. + However, the host application's own PII scanner detects Social + Security Numbers in the tool's output. + + Questions: + Q1. Should the client trust the server's sensitivity="none"? + a) yes (server is authoritative) + b) no (defense-in-depth) + + Q2. Does the spec define a mechanism for the host to override + or augment the server's sensitivity annotation? + a) yes b) no + + Q3. If the host rewrites sensitivity from "none" to "pii", + what is lost? + a) nothing b) provenance (who classified it) + c) attribution d) all of the above + + Q4. What would a clean mechanism look like? + _______________________________________________ + + Your answers: + Q1: ___ Q2: ___ Q3: ___ + Q4: _______________________________________________ + """ + + def test_result_annotations_have_no_host_override_field(self): + """ResultAnnotations has no field for host-detected sensitivity. + The only sensitivity field is the server-provided one.""" + result = ResultAnnotations(sensitivity=SimpleDataClass.NONE) + assert result.sensitivity == SimpleDataClass.NONE + assert not hasattr(result, "host_sensitivity") + assert not hasattr(result, "override_sensitivity") + + def test_host_can_construct_higher_sensitivity_result(self): + """A host CAN create a new ResultAnnotations with higher sensitivity, + but this loses provenance (no way to indicate 'host-detected').""" + server_result = ResultAnnotations(sensitivity=SimpleDataClass.NONE) + # Host scanner finds PII — constructs override + host_override = ResultAnnotations(sensitivity=SimpleDataClass.PII) + assert sensitivity_level(host_override.sensitivity) > sensitivity_level( + server_result.sensitivity + ) + + def test_session_tracker_accepts_host_override(self): + """SessionTracker merges whatever ResultAnnotations it receives, + regardless of who constructed them.""" + tracker = SessionTracker(session_id="test-s09") + # Merge the host-overridden result instead of the server result + tracker.merge( + ResultAnnotations(sensitivity=SimpleDataClass.PII), + "tool_with_ssns", + ) + assert tracker.session.max_sensitivity == SimpleDataClass.PII + + +# =================================================================== +# Scenario 10 — Gap: maliciousActivityHint — HITL vs Propagation +# Spec section: maliciousActivityHint semantics, boolean union +# Expected agreement: LOW (tests debatable design choice) +# =================================================================== + +class TestScenario10_MaliciousActivityHintSemantics: + """SCENARIO 10 — maliciousActivityHint: HITL vs Propagation + + Context: + A tool server returns maliciousActivityHint=true in a + CallToolResult's annotations. + + Questions: + Q1. What should the client do with this hint? + a) Surface it to the human user (HITL) + b) Propagate it into session state via boolean union + c) Both a and b + d) Ignore it + + Q2. Should maliciousActivityHint propagate into session state + the same way openWorldHint does (boolean union)? + a) yes b) no — why not? + + Q3. If it propagates, what happens when a LATER tool call + in the same session is evaluated by the policy engine? + _______________________________________________ + + Q4. Does the spec explicitly define propagation behavior + for maliciousActivityHint? + a) yes b) no + + Your answers: + Q1: ___ Q2: ___ Q4: ___ + Q3: _______________________________________________ + """ + + def test_sdk_propagates_malicious_hint(self): + """The SDK DOES propagate maliciousActivityHint via boolean union. + This is the current behavior, even though one can argue it + should only trigger HITL and not persist in session state.""" + tracker = SessionTracker(session_id="test-s10") + tracker.merge( + ResultAnnotations(malicious_activity_hint=True), "suspicious_tool" + ) + assert tracker.session.malicious_activity_hint is True + + def test_malicious_hint_is_sticky(self): + """Once true, maliciousActivityHint stays true (boolean union).""" + tracker = SessionTracker(session_id="test-s10-sticky") + tracker.merge( + ResultAnnotations(malicious_activity_hint=True), "bad_tool" + ) + tracker.merge( + ResultAnnotations(malicious_activity_hint=False), "good_tool" + ) + assert tracker.session.malicious_activity_hint is True + + def test_malicious_session_does_not_block_clean_tools(self): + """BlockMaliciousActivity checks the TOOL annotation, not session. + A clean tool in a 'malicious' session is NOT blocked by this rule.""" + engine = PolicyEngine(mode="enforce") + engine.register_tool("clean_tool", to_wire(TrustAnnotations())) + engine.register_tool("suspect_tool", {"maliciousActivityHint": True}) + + session = SessionAnnotations() + session.malicious_activity_hint = True + + # Clean tool is NOT blocked (rule checks tool annotation, not session) + d1 = engine.evaluate("clean_tool", session=session) + assert d1.allowed is True + + # Suspect tool IS blocked (tool annotation has maliciousActivityHint) + d2 = engine.evaluate("suspect_tool", session=session) + assert d2.allowed is False + assert d2.rule == "block-malicious-activity" + + def test_no_session_policy_for_malicious_session(self): + """There is no default rule that checks session.malicious_activity_hint. + + BlockMaliciousActivity only checks the tool's annotations dict, + not the session. This is a potential gap — should there be a + session-level malicious activity rule?""" + engine = PolicyEngine(mode="enforce") + engine.register_tool("innocent_tool", to_wire(TrustAnnotations( + input_metadata=InputMetadata( + destination=Destination.INTERNAL, + sensitivity=SimpleDataClass.NONE, + ), + ))) + + session = SessionAnnotations() + session.malicious_activity_hint = True + + decision = engine.evaluate( + "innocent_tool", action="call", + target_destination="internal", session=session, + ) + # Session is "malicious" but no rule checks session for this + assert decision.allowed is True + assert decision.effect == "allow"