diff --git a/README.md b/README.md index c034e6cc..ccc6b5b5 100644 --- a/README.md +++ b/README.md @@ -8,7 +8,7 @@ **Runtime guardrails for AI agents — configurable, extensible, and production-ready.** -AI agents interact with users, tools, and external systems in unpredictable ways. **Agent Control** provides an extensible, policy-based runtime layer that evaluates inputs and outputs against configurable rules — blocking prompt injections, PII leakage, and other risks without modifying your agent's code. +AI agents interact with users, tools, and external systems in unpredictable ways. **Agent Control** provides an extensible, control-based runtime layer that evaluates inputs and outputs against configurable rules — blocking prompt injections, PII leakage, and other risks without modifying your agent's code. ![Agent Control Architecture](docs/images/Architecture.png) @@ -23,7 +23,7 @@ Traditional guardrails embedded inside your agent code have critical limitations **Agent Control gives you runtime control over what your agents can and cannot do:** - **For developers:** Centralize safety logic and adapt to emerging threats instantly without redeployment - **For non-technical teams:** Intuitive UI to configure and monitor agent safety without touching code -- **For organizations:** Reusable policies across agents with comprehensive audit trails +- **For organizations:** Reusable controls across agents with comprehensive audit trails --- @@ -31,7 +31,7 @@ Traditional guardrails embedded inside your agent code have critical limitations - **Safety Without Code Changes** — Add guardrails with a `@control()` decorator - **Runtime Configuration** — Update controls instantly via API or UI without having to re-deploy your agentic applications -- **Centralized Policies** — Define controls once, apply to multiple agents +- **Centralized Controls** — Define controls once, apply to multiple agents - **Web Dashboard** — Visual interface for managing agents, controls, and viewing analytics - **Pluggable Evaluators** — Built-in (regex, list matching, Luna-2 AI) or custom evaluators - **Fail-Safe Defaults** — Deny controls fail closed on error with configurable error handling @@ -119,12 +119,16 @@ cd agent-control # Install dependencies make sync -# Start server (automatically starts Postgres + runs migrations + starts server) +# Start PostgreSQL database +cd server && docker-compose up -d && cd .. + +# Run database migrations +make server-alembic-upgrade + +# Start the Agent Control server make server-run ``` -> 💡 **First time?** The command above handles everything: starts Postgres, runs migrations, and starts the server. Migrations are idempotent - safe to run multiple times. - **Server is now running at `http://localhost:8000`** ✅ > 💡 **Verify it's working:** Open http://localhost:8000/health in your browser - you should see `{"status": "ok"}` @@ -155,16 +159,16 @@ Create controls to protect your agent's operations: # setup.py - Run once to configure everything import asyncio from datetime import datetime, UTC -from agent_control import AgentControlClient, controls, policies, agents +from agent_control import AgentControlClient, controls, agents from agent_control_models import Agent async def setup(): async with AgentControlClient() as client: # Defaults to localhost:8000 - # 1. Register agent first (required before assigning policy) + # 1. Register agent first agent = Agent( # Your agent's UUID agent_name="550e8400-e29b-41d4-a716-446655440000", - agent_name="My Chatbot", + agent_description="My Chatbot", agent_created_at=datetime.now(UTC).isoformat() ) await agents.register_agent(client, agent, steps=[]) @@ -185,26 +189,15 @@ async def setup(): "action": {"decision": "deny"} } ) - # 3. Create policy - policy = await policies.create_policy(client, name="production-policy") - - # 4. Add control to policy - await policies.add_control_to_policy( - client, - policy_id=policy["policy_id"], - control_id=control["control_id"] - ) - - # 5. Assign policy to agent - await policies.assign_policy_to_agent( + # 3. Associate control directly with agent + await agents.add_agent_control( client, - agent_name=AGENT_ID, - policy_id=policy["policy_id"] + agent_name=agent.agent_name, + control_id=control["control_id"], ) print("✅ Setup complete!") print(f" Control ID: {control['control_id']}") - print(f" Policy ID: {policy['policy_id']}") asyncio.run(setup()) ``` @@ -231,8 +224,8 @@ from agent_control import control, ControlViolationError # Initialize your agent agent_control.init( - agent_name="My Chatbot", - agent_name="550e8400-e29b-41d4-a716-446655440000" + agent_name="550e8400-e29b-41d4-a716-446655440000", # Agent identifier (UUID recommended) + agent_description="My Chatbot", ) # Protect any function (like LLM calls) @@ -289,6 +282,28 @@ uv run my_agent.py > **💡 Pro Tip:** Start with simple regex controls, then graduate to AI-powered evaluators for complex safety checks! +### 5. Assign Controls + +Controls can be associated directly with agents. An agent's **active controls** are the controls currently linked to that agent. + +**Direct controls** — attach individual controls to an agent: + +```bash +# Add a control directly to an agent +curl -X POST http://localhost:8000/api/v1/agents/support-agent-v1/controls/3 + +# Remove a direct control +curl -X DELETE http://localhost:8000/api/v1/agents/support-agent-v1/controls/3 +``` + +**List all active controls**: + +```bash +curl http://localhost:8000/api/v1/agents/support-agent-v1/controls +``` + +> Controls are optional. An agent can operate with no controls configured. + --- ## Configuration @@ -299,7 +314,7 @@ uv run my_agent.py |----------|---------|-------------| | `AGENT_CONTROL_URL` | `http://localhost:8000` | Server URL for SDK | | `AGENT_CONTROL_API_KEY` | — | API key for authentication (if enabled) | -| `DATABASE_URL` or `DB_URL` | `postgresql+psycopg://agent_control:agent_control@localhost:5432/agent_control` | Database connection string (`DATABASE_URL` preferred for Docker, `DB_URL` for local dev. SQLite: `sqlite+aiosqlite:///./agent_control.db`) | +| `DB_URL` | `postgresql+psycopg://agent_control:agent_control@localhost:5432/agent_control` | Database connection string (SQLite: `sqlite+aiosqlite:///./agent_control.db`) | | `GALILEO_API_KEY` | — | Required for Luna-2 AI evaluator | ### Server Configuration @@ -334,8 +349,9 @@ Agent Control is built as a monorepo with these components: ▼ ┌──────────────────────────────────────────────────────────────────┐ │ Agent Control Server │ +│ │ │ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │ -│ │ Controls │ │ Policies │ │ Evaluators │ │ Agents │ │ +│ │ Controls │ │Control Link│ │ Evaluators │ │ Agents │ │ │ │ API │ │ API │ │ Registry │ │ API │ │ │ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │ └──────────────────────────────────────────────────────────────────┘ @@ -353,6 +369,7 @@ Agent Control is built as a monorepo with these components: | Package | Description | |:--------|:------------| | `agent-control-sdk` | Python SDK with `@control()` decorator | +| `agent-control` (npm) | TypeScript SDK (generated from OpenAPI) | | `agent-control-server` | FastAPI server with Control Management API | | `agent-control-engine` | Core evaluation logic and evaluator system | | `agent-control-models` | Shared Pydantic v2 models | @@ -368,10 +385,12 @@ Agent Control is built as a monorepo with these components: ``` agent-control/ ├── sdks/python/ # Python SDK (agent-control) +├── sdks/typescript/ # TypeScript SDK (generated) ├── server/ # FastAPI server (agent-control-server) ├── engine/ # Evaluation engine (agent-control-engine) ├── models/ # Shared models (agent-control-models) ├── evaluators/ # Evaluator implementations (agent-control-evaluators) +├── ui/ # Next.js web dashboard └── examples/ # Usage examples ``` diff --git a/docs/OVERVIEW.md b/docs/OVERVIEW.md index bc427af9..67f31de9 100644 --- a/docs/OVERVIEW.md +++ b/docs/OVERVIEW.md @@ -2,7 +2,7 @@ **Runtime guardrails for AI agents — configurable, extensible, and production-ready.** -Agent Control provides a policy-based control layer that sits between your AI agents and the outside world. It evaluates inputs and outputs against configurable rules, blocking harmful content, prompt injections, PII leakage, and other risks — all without changing your agent's code. +Agent Control provides a control-based runtime layer that sits between your AI agents and the outside world. It evaluates inputs and outputs against configurable rules, blocking harmful content, prompt injections, PII leakage, and other risks — all without changing your agent's code. --- @@ -57,15 +57,12 @@ A **Control Set** is a named group of related controls. Use them to organize con | `compliance-controls` | block-pii, block-phi, audit-logging | | `quality-controls` | check-hallucination, verify-sources | -### 📜 Policies +### 🔗 Control Associations -A **Policy** combines one or more Control Sets and is assigned to agents. Policies let you: -- Reuse control sets across multiple agents -- Version and audit your safety rules -- Apply different policies to different environments (dev/staging/prod) +Controls can be linked directly to agents and reused across multiple agents. This keeps rollout and tuning fast without redeploying application code. ``` -Policy → Control Sets → Controls → Agents +Controls → Agents ``` ### 🎯 Selectors @@ -335,14 +332,14 @@ async def chat(message: str) -> str: Update controls without redeploying your application. Critical for: - Responding to emerging threats - Tuning thresholds based on real-world data -- A/B testing different safety policies +- A/B testing different safety controls -### 🎯 Centralized Policy Management -Define controls once, apply them to multiple agents. Security teams can manage policies independently from development teams. +### 🎯 Centralized Control Management +Define controls once, apply them to multiple agents. Security teams can manage controls independently from development teams. ``` ┌─────────────────────────────────────────────────────┐ -│ Policy │ +│ Control Collection │ │ ┌─────────────┐ ┌─────────────┐ ┌─────────────┐ │ │ │ Control Set │ │ Control Set │ │ Control Set │ │ │ │ (Safety) │ │ (Compliance)│ │ (Quality) │ │ @@ -398,7 +395,7 @@ Choose how to handle failures: ┌──────────────────────────────────────────────────────────────────┐ │ Agent Control Server │ │ ┌────────────┐ ┌────────────┐ ┌────────────┐ ┌────────────┐ │ -│ │ Controls │ │ Policies │ │ Evaluators │ │ Agents │ │ +│ │ Controls │ │Control Link│ │ Evaluators │ │ Agents │ │ │ │ API │ │ API │ │ Registry │ │ API │ │ │ └────────────┘ └────────────┘ └────────────┘ └────────────┘ │ └──────────────────────────────────────────────────────────────────┘ diff --git a/docs/REFERENCE.md b/docs/REFERENCE.md index 10145f22..66e8c2e7 100644 --- a/docs/REFERENCE.md +++ b/docs/REFERENCE.md @@ -20,7 +20,7 @@ This document provides comprehensive technical reference for Agent Control. Each ## Introduction -Agent Control provides a policy-based control layer that sits between your AI agents and the outside world. It evaluates inputs and outputs against configurable rules, blocking harmful content, prompt injections, PII leakage, and other risks. +Agent Control provides a control-based runtime layer that sits between your AI agents and the outside world. It evaluates inputs and outputs against configurable rules, blocking harmful content, prompt injections, PII leakage, and other risks. ### Why Agent Control? @@ -64,16 +64,12 @@ Example: *"If the output contains an SSN pattern, block the response."* } ``` -### Policies +### Control Associations -A **Policy** is a named collection of controls assigned to agents. Policies enable you to: - -- Reuse control sets across multiple agents -- Version and audit your safety rules -- Apply different policies to different environments (dev/staging/prod) +Controls can be assigned directly to agents and reused across multiple agents. ``` -Policy → Controls → Agents +Controls → Agents ``` ### Check Stages @@ -252,7 +248,7 @@ graph TB ### Agent Initialization 1. **Agent Registration**: Agent initializes with `agent_control.init()`, registering with the server -2. **Policy Assignment**: Server returns the agent's assigned policy and active controls +2. **Control Resolution**: Server returns the agent's active controls ### Control Execution Flow 1. **Function Invocation**: User calls a function decorated with `@control()` @@ -590,8 +586,8 @@ The Python SDK provides decorator-based protection and programmatic control mana import agent_control agent_control.init( - agent_name="my-agent", # Required: human-readable name - agent_name="550e8400-e29b-41d4-a716-446655440000", # Required: UUID + agent_name="my-agent", # Required: unique identifier + agent_description="My Agent", # Optional: human-readable description server_url="http://localhost:8000", # Optional: defaults to env var policy_refresh_interval_seconds=60, # Optional: set 0 to disable background refresh steps=[ # Optional: register available steps @@ -605,12 +601,12 @@ agent_control.init( ) ``` -When enabled, background refresh fetches controls via `GET /agents/{agent_id}/controls`. +When enabled, background refresh fetches controls via `GET /agents/{agent_name}/controls`. Refresh failures are fail-open: the SDK keeps the last successful local cache snapshot. ### The @control Decorator -The `@control()` decorator applies server-side policies to any function. +The `@control()` decorator applies server-side controls to any function. ```python from agent_control import control @@ -632,7 +628,6 @@ async def chat(message: str) -> str: **Parameters**: -- `policy` (str, optional): Policy name for documentation purposes. The agent's assigned policy is automatically used. - `step_name` (str, optional): Custom name for this step. If not provided, uses the function name. Useful for: - Overriding auto-detected names when they don't match your control configuration - Applying the same controls to functions with different names @@ -779,22 +774,14 @@ async with AgentControlClient() as client: | `agent_control.get_control()` | Get control by ID | | `agent_control.update_control()` | Update control properties | | `agent_control.delete_control()` | Delete a control | -| `agent_control.add_control_to_policy()` | Add control to policy | -| `agent_control.remove_control_from_policy()` | Remove control from policy | -| `agent_control.list_policy_controls()` | List controls in a policy | - -**Policy management** (via `agent_control.policies` module): - -| Function | Description | -|----------|-------------| -| `policies.create_policy()` | Create a new policy | -| `policies.assign_policy_to_agent()` | Assign policy to agent | +| `agent_control.add_control_to_agent()` | Add control to agent | +| `agent_control.remove_control_from_agent()` | Remove control from agent | --- ## Server API -The Agent Control server exposes a RESTful API for managing agents, controls, and policies. +The Agent Control server exposes a RESTful API for managing agents and controls. ### Base URL @@ -811,9 +798,8 @@ Default: `http://localhost:8000/api/v1` | `GET` | `/agents/{agent_name}` | Get agent details | | `PATCH` | `/agents/{agent_name}` | Update agent | | `GET` | `/agents/{agent_name}/controls` | List controls for agent | -| `GET` | `/agents/{agent_name}/policy` | Get agent's policy | -| `POST` | `/agents/{agent_name}/policy/{policy_id}` | Assign policy | -| `DELETE` | `/agents/{agent_name}/policy` | Remove policy | +| `POST` | `/agents/{agent_name}/controls/{control_id}` | Add control to agent | +| `DELETE` | `/agents/{agent_name}/controls/{control_id}` | Remove control from agent | **Controls**: @@ -825,15 +811,6 @@ Default: `http://localhost:8000/api/v1` | `PATCH` | `/controls/{control_id}` | Update control | | `DELETE` | `/controls/{control_id}` | Delete control | -**Policies**: - -| Method | Endpoint | Description | -|--------|----------|-------------| -| `PUT` | `/policies` | Create policy | -| `GET` | `/policies/{policy_id}/controls` | List controls in policy | -| `POST` | `/policies/{policy_id}/controls/{control_id}` | Add control to policy | -| `DELETE` | `/policies/{policy_id}/controls/{control_id}` | Remove control | - **System**: | Method | Endpoint | Description | @@ -928,7 +905,7 @@ Agent Control supports multiple API keys for zero-downtime rotation: | `DB_PASSWORD` | `agent_control` | Database password | | `DB_DATABASE` | `agent_control` | Database name | | `DB_DRIVER` | `psycopg` | Database driver | -| `DATABASE_URL` or `DB_URL` | — | Full database URL (overrides above). `DATABASE_URL` is preferred for Docker environments. | +| `DB_URL` | — | Full database URL (overrides above) | **Authentication**: diff --git a/docs/evaluators/sql.md b/docs/evaluators/sql.md index 5ca18ae3..74fe2b77 100644 --- a/docs/evaluators/sql.md +++ b/docs/evaluators/sql.md @@ -6,7 +6,7 @@ A practical guide for configuring SQL validation controls in your AI agent. ## What is the SQL Evaluator? -The SQL Evaluator validates SQL query strings (e.g., from LLM responses) before they execute against your database. It acts as a security and safety layer, preventing dangerous operations, enforcing access policies, and ensuring data isolation. +The SQL Evaluator validates SQL query strings (e.g., from LLM responses) before they execute against your database. It acts as a security and safety layer, preventing dangerous operations, enforcing access rules, and ensuring data isolation. **Technical Foundation**: Uses [sqlglot](https://github.com/tobymao/sqlglot) with the Rust-accelerated parser (`sqlglot[rs]`) for high-performance SQL parsing and AST-based validation. @@ -36,7 +36,7 @@ The SQL Evaluator validates SQL query strings (e.g., from LLM responses) before > **⚠️ Important Security Note** > -> This evaluator validates query structure and enforces access rules, but **it is not a complete defense against SQL injection**. The primary defense against SQL injection is using **prepared statements** (parameterized queries) at the database layer. This evaluator provides an additional security layer by validating query syntax and enforcing policies, but should not be relied upon as the sole protection mechanism. +> This evaluator validates query structure and enforces access rules, but **it is not a complete defense against SQL injection**. The primary defense against SQL injection is using **prepared statements** (parameterized queries) at the database layer. This evaluator provides an additional security layer by validating query syntax and enforcing controls, but should not be relied upon as the sole protection mechanism. > > **Best Practice**: Always use prepared statements/parameterized queries when executing SQL from untrusted sources. This evaluator complements, but does not replace, proper database security practices. diff --git a/examples/README.md b/examples/README.md index 5183f326..fbf1f37a 100644 --- a/examples/README.md +++ b/examples/README.md @@ -47,7 +47,7 @@ uv run python examples/agent_control_demo/update_controls.py --block-ssn **Files:** - `setup_controls.py` - Create and configure controls via SDK -- `demo_agent.py` - Agent that uses `@control` decorator with server-side policies +- `demo_agent.py` - Agent that uses `@control` decorator with server-side controls - `update_controls.py` - Dynamically update controls without code changes - `agent_luna_demo.py` - Luna-2 evaluator integration for AI safety checks @@ -126,13 +126,13 @@ See [steer_action_demo/README.md](steer_action_demo/README.md) for details. import agent_control from agent_control import control, ControlViolationError -# Initialize agent (connects to server, loads policy) +# Initialize agent (connects to server, loads control associations) agent_control.init( agent_name="my-bot", - agent_name="550e8400-e29b-41d4-a716-446655440000", + agent_description="My Bot", ) -# Apply the agent's assigned policy +# Apply controls associated with the agent @control() async def chat(message: str) -> str: return await assistant.respond(message) @@ -150,7 +150,7 @@ except ControlViolationError as e: import agent_control from agent_control import control, ControlSteerError, ControlViolationError -agent_control.init(agent_name="my-bot", agent_id="...") +agent_control.init(agent_name="my-bot", agent_description="My Bot") @control() async def generate_content(prompt: str) -> str: diff --git a/examples/agent_control_demo/demo_agent.py b/examples/agent_control_demo/demo_agent.py index f0fd4e10..6e139698 100644 --- a/examples/agent_control_demo/demo_agent.py +++ b/examples/agent_control_demo/demo_agent.py @@ -77,12 +77,12 @@ def simulate_database_query(sql: str) -> str: # AGENT FUNCTIONS WITH CONTROLS # ============================================================================= -@control(policy="demo-policy") +@control() async def chat(message: str) -> str: """ - Chat function protected by the demo policy. - - The policy's 'block-ssn-output' control checks OUTPUT for SSN patterns. + Chat function protected by agent-associated controls. + + The 'block-ssn-output' control checks OUTPUT for SSN patterns. This prevents PII leakage in responses. """ print(f" [Agent] Processing message: {message}") @@ -91,12 +91,12 @@ async def chat(message: str) -> str: return response -@control(policy="demo-policy") +@control() async def execute_query(query: str) -> str: """ - Database query function protected by the demo policy. - - The policy's 'block-dangerous-sql' control checks INPUT for dangerous + Database query function protected by agent-associated controls. + + The 'block-dangerous-sql' control checks INPUT for dangerous SQL keywords and blocks before execution. """ print(f" [Agent] Executing query: {query}") @@ -105,12 +105,12 @@ async def execute_query(query: str) -> str: return result -@control() # Apply agent's assigned policy +@control() # Apply agent-associated controls async def process_request(input: str) -> str: """ - General processing function with the agent's policy applied. - - All controls in the policy are evaluated (both pre and post stage). + General processing function with agent-associated controls applied. + + Controls directly associated to the agent are evaluated. """ print(f" [Agent] Processing request: {input}") response = simulate_llm_response(input) @@ -162,8 +162,8 @@ async def run_demo(): logger.info(f"Initializing agent: {AGENT_NAME}") agent_control.init( agent_name=AGENT_ID, + agent_description=AGENT_NAME, server_url=SERVER_URL, - agent_description="Demo chatbot for testing controls" ) logger.info("Agent initialized successfully") except Exception as e: diff --git a/examples/agent_control_demo/setup_controls.py b/examples/agent_control_demo/setup_controls.py index 658d0dda..7bf9b9fd 100644 --- a/examples/agent_control_demo/setup_controls.py +++ b/examples/agent_control_demo/setup_controls.py @@ -5,13 +5,12 @@ This script demonstrates the full control lifecycle: 1. Create an agent 2. Create controls (regex and list) -3. Create a policy and add controls to it -4. Assign the policy to the agent -5. List controls -6. Update a control +3. Associate controls directly with the agent +4. List controls +5. Update a control API Structure: - Agent → Policy → Controls + Agent → Controls Prerequisites: - Agent Control server running at http://localhost:8000 @@ -53,7 +52,6 @@ async def create_agent(client: AgentControlClient) -> str: json={ "agent": { "agent_name": str(agent_name), - "agent_name": AGENT_NAME, "agent_description": "Demo chatbot for testing controls", }, "steps": [] @@ -185,74 +183,27 @@ async def create_list_control(client: AgentControlClient) -> int: -async def create_policy(client: AgentControlClient, name: str) -> int: - """Create a policy.""" - print("\n" + "=" * 60) - print("STEP 4: Creating Policy") - print("=" * 60) - - try: - response = await client.http_client.put( - "/api/v1/policies", - json={"name": name} - ) - - if response.status_code == 409: - print(f" ℹ️ Policy '{name}' already exists") - return -1 - - response.raise_for_status() - policy_id = response.json().get("policy_id") - - print(f"✓ Created policy '{name}' with ID: {policy_id}") - return policy_id - - except Exception as e: - print(f"✗ Failed to create policy: {e}") - raise - - -async def add_control_to_policy( - client: AgentControlClient, - policy_id: int, - control_id: int -) -> bool: - """Add a control directly to a policy.""" - try: - response = await client.http_client.post( - f"/api/v1/policies/{policy_id}/controls/{control_id}" - ) - response.raise_for_status() - data = response.json() - print(f" ✓ Added control {control_id} to policy {policy_id}") - print(f" Response: {data}") - return True - except Exception as e: - print(f" ✗ Failed to add control to policy: {e}") - return False - - -async def assign_policy_to_agent( +async def add_control_to_agent( client: AgentControlClient, agent_name: str, - policy_id: int + control_id: int, ) -> bool: - """Assign a policy to an agent.""" + """Associate a control directly with an agent.""" print("\n" + "=" * 60) - print("STEP 5: Assigning Policy to Agent") + print("STEP 4: Associating Control to Agent") print("=" * 60) try: response = await client.http_client.post( - f"/api/v1/agents/{agent_name}/policy/{policy_id}" + f"/api/v1/agents/{agent_name}/controls/{control_id}" ) response.raise_for_status() data = response.json() - print(f"✓ Assigned policy {policy_id} to agent {agent_name}") + print(f"✓ Associated control {control_id} to agent {agent_name}") print(f" Response: {data}") return True except Exception as e: - print(f"✗ Failed to assign policy: {e}") + print(f"✗ Failed to associate control to agent: {e}") return False @@ -379,33 +330,8 @@ async def verify_full_chain(client: AgentControlClient, agent_name: str) -> None except Exception as e: print(f" Error: {e}") - # 2. Get agent's policy - print("\n2. Agent's Policy:") - try: - resp = await client.http_client.get(f"/api/v1/agents/{agent_name}/policy") - if resp.status_code == 404: - print(" No policy assigned to agent") - policy_id = None - else: - resp.raise_for_status() - policy_data = resp.json() - policy_id = policy_data.get("policy_id") - print(f" Policy ID: {policy_id}") - - if policy_id: - # 3. Get policy's controls - print("\n3. Policy's Controls:") - resp = await client.http_client.get(f"/api/v1/policies/{policy_id}/controls") - resp.raise_for_status() - ctrl_data = resp.json() - control_ids = ctrl_data.get("control_ids", []) - print(f" Control IDs: {control_ids}") - except Exception as e: - print(f" Error: {e}") - policy_id = None - - # 4. Final: List agent controls (the API we're testing) - print("\n4. Final Agent Controls (via /agents/{id}/controls):") + # 2. Final: List agent controls (the API we're testing) + print("\n2. Final Agent Controls (via /agents/{id}/controls):") try: resp = await client.http_client.get(f"/api/v1/agents/{agent_name}/controls") resp.raise_for_status() @@ -460,52 +386,29 @@ async def main(): await verify_full_chain(client, agent_name) return - # 3. Create policy - policy_id = await create_policy(client, "demo-policy") - if policy_id == -1: - print("\n⚠️ Policy already exists. Running verification...") - await verify_full_chain(client, agent_name) - return - - # 4. Add controls to policy - print("\n Adding controls to policy...") - ok1 = await add_control_to_policy(client, policy_id, regex_control_id) - ok2 = await add_control_to_policy(client, policy_id, list_control_id) + # 3. Associate controls directly with the agent + print("\n Associating controls directly to agent...") + ok1 = await add_control_to_agent(client, agent_name, regex_control_id) + ok2 = await add_control_to_agent(client, agent_name, list_control_id) if not (ok1 and ok2): - print("\n⚠️ Failed to add controls to policy!") - - # Verify: List controls in policy - print("\n Verifying policy contents...") - try: - resp = await client.http_client.get( - f"/api/v1/policies/{policy_id}/controls" - ) - resp.raise_for_status() - policy_controls = resp.json() - print(f" Policy {policy_id} has controls: {policy_controls}") - except Exception as e: - print(f" Failed to verify policy: {e}") - - # 5. Assign policy to agent - ok3 = await assign_policy_to_agent(client, agent_name, policy_id) - if not ok3: - print("\n⚠️ Failed to assign policy to agent!") + print("\n⚠️ Failed to associate one or more controls to agent!") - # Verify: Get agent's policy - print("\n Verifying agent policy assignment...") + # Verify: ensure both controls are active on the agent + print("\n Verifying direct control associations...") try: - resp = await client.http_client.get( - f"/api/v1/agents/{agent_name}/policy" - ) + resp = await client.http_client.get(f"/api/v1/agents/{agent_name}/controls") resp.raise_for_status() - agent_policy = resp.json() - assigned_policy_id = agent_policy.get("policy_id") - if assigned_policy_id == policy_id: - print(f" ✓ Agent correctly assigned to policy {policy_id}") + active_controls = resp.json().get("controls", []) + active_ids = {control.get("id") for control in active_controls} + expected_ids = {regex_control_id, list_control_id} + if expected_ids.issubset(active_ids): + print(f" ✓ Agent has expected control IDs: {sorted(expected_ids)}") else: - print(f" ⚠️ Agent assigned to policy {assigned_policy_id}, expected {policy_id}") + print( + f" ⚠️ Active control IDs are {sorted(active_ids)}, expected at least {sorted(expected_ids)}" + ) except Exception as e: - print(f" ✗ Failed to verify agent policy: {e}") + print(f" ✗ Failed to verify direct control associations: {e}") # 6. List controls await list_agent_controls(client, agent_name) @@ -533,7 +436,7 @@ async def main(): - Keywords: DROP, DELETE, TRUNCATE, ALTER, GRANT, REVOKE, EXECUTE, SHUTDOWN, BACKUP API Flow: - Agent → Policy → Controls + Agent → Controls Now run the agent demo: uv run python examples/agent_control_demo/demo_agent.py diff --git a/examples/crewai/README.md b/examples/crewai/README.md index 6540cf89..72cd9114 100644 --- a/examples/crewai/README.md +++ b/examples/crewai/README.md @@ -109,8 +109,7 @@ This creates: - Unauthorized access control (blocks requests for other users' data - PRE-execution) - PII detection control for tool outputs (blocks SSN, credit cards, emails, phones - POST-execution) - Final output validation control (catches agent-generated PII - POST-execution) -- Policy with all three controls -- Assigns policy to the customer support crew agent +- Direct association of all three controls with the customer support crew agent ## Running the Example @@ -157,7 +156,7 @@ uv run content_agent_protection.py 🚫 [LAYER 2: Agent Control POST] BLOCKED Reason: Control 'pii-detection-output' matched Tool executed but output contained violations - LLM generated content that violated policies + LLM generated content that violated controls 🚫 SECURITY VIOLATION (POST-execution): [Details...] ``` @@ -390,7 +389,7 @@ This control catches PII in the final crew output, protecting against orchestrat Agent Control works seamlessly **with** CrewAI's agent orchestration: 1. **CrewAI Agent Layer**: Plans tasks, selects tools, manages conversation flow -2. **Agent Control Layer**: Enforces policies and business rules at tool boundaries +2. **Agent Control Layer**: Enforces controls and business rules at tool boundaries ``` User Request @@ -475,7 +474,7 @@ Return result or raise ControlViolationError ## Files - `content_agent_protection.py` - Main CrewAI crew with @control() -- `setup_content_controls.py` - One-time setup for controls/policy +- `setup_content_controls.py` - One-time setup for controls/direct associations - `pyproject.toml` - Dependencies - `README.md` - This file diff --git a/examples/crewai/content_agent_protection.py b/examples/crewai/content_agent_protection.py index 2827d4e3..a2d4a2f3 100644 --- a/examples/crewai/content_agent_protection.py +++ b/examples/crewai/content_agent_protection.py @@ -298,7 +298,7 @@ def handle_ticket_tool(ticket: str) -> str: print("\n🚫 [LAYER 2: Agent Control POST] BLOCKED") print(f" Reason: {e.message}") print(" Tool executed but output contained violations") - print(" LLM generated content that violated policies") + print(" LLM generated content that violated controls") stage = "POST-execution" error_msg = f"🚫 SECURITY VIOLATION ({stage}): {e.message}\n\nThis request has been logged for security review." @@ -377,7 +377,7 @@ def create_support_crew(): goal="Provide helpful customer support while protecting user privacy and data security", backstory=( "You are an experienced customer support agent who helps customers with their questions. " - "You are friendly, professional, and always respect customer privacy and data security policies." + "You are friendly, professional, and always respect customer privacy and data security controls." ), tools=[ticket_handler_tool], verbose=True diff --git a/examples/crewai/setup_content_controls.py b/examples/crewai/setup_content_controls.py index 779ae814..77501310 100644 --- a/examples/crewai/setup_content_controls.py +++ b/examples/crewai/setup_content_controls.py @@ -11,14 +11,14 @@ import asyncio import os -from agent_control import Agent, AgentControlClient, agents, controls, policies +from agent_control import Agent, AgentControlClient, agents, controls AGENT_ID = "989d84f0-9afe-4fb2-9e9e-e9d076271e29" SERVER_URL = os.getenv("AGENT_CONTROL_URL", "http://localhost:8000") async def setup_content_controls(): - """Create PII protection and unauthorized access controls, policy, and assign to agent.""" + """Create PII protection/unauthorized-access controls and add them directly to the agent.""" async with AgentControlClient(base_url=SERVER_URL) as client: # 1. Register Agent agent_name = AGENT_ID @@ -172,65 +172,35 @@ async def setup_content_controls(): else: raise - # 6. Create Policy + # 6. Associate controls directly with agent try: - policy_result = await policies.create_policy( - client, name="support-pii-protection-policy" - ) - policy_id = policy_result["policy_id"] - print(f"✓ Policy created (ID: {policy_id})") - except Exception as e: - if "409" in str(e): - print(f"⚠️ Policy 'support-pii-protection-policy' already exists.") - print(" Cannot proceed - SDK doesn't support looking up policies by name.") - print("\n To fix this, run one of these commands:") - print(" 1. Delete via server API:") - print(f" curl -X DELETE {SERVER_URL}/api/v1/policies/") - print(" 2. Or use the server UI to delete the policy") - print("\n Then re-run this script.") - raise SystemExit(1) - raise - - # 7. Add Controls to Policy - try: - await policies.add_control_to_policy(client, policy_id, unauthorized_control_id) - print(f"✓ Added unauthorized access control to policy") - except Exception as e: - if "409" in str(e) or "already" in str(e).lower(): - print(f"ℹ️ Unauthorized access control already in policy (OK)") - else: - print(f"❌ Failed to add control to policy: {e}") - raise - - try: - await policies.add_control_to_policy(client, policy_id, pii_control_id) - print(f"✓ Added PII detection control to policy") + await agents.add_agent_control(client, agent_name, unauthorized_control_id) + print("✓ Added unauthorized access control to agent") except Exception as e: if "409" in str(e) or "already" in str(e).lower(): - print(f"ℹ️ PII detection control already in policy (OK)") + print("ℹ️ Unauthorized access control already associated with agent (OK)") else: - print(f"❌ Failed to add control to policy: {e}") + print(f"❌ Failed to add unauthorized access control to agent: {e}") raise try: - await policies.add_control_to_policy(client, policy_id, final_output_control_id) - print(f"✓ Added final output validation control to policy") + await agents.add_agent_control(client, agent_name, pii_control_id) + print("✓ Added PII detection control to agent") except Exception as e: if "409" in str(e) or "already" in str(e).lower(): - print(f"ℹ️ Final output validation control already in policy (OK)") + print("ℹ️ PII detection control already associated with agent (OK)") else: - print(f"❌ Failed to add control to policy: {e}") + print(f"❌ Failed to add PII detection control to agent: {e}") raise - # 8. Assign Policy to Agent try: - await policies.assign_policy_to_agent(client, agent_name, policy_id) - print(f"✓ Assigned policy to agent") + await agents.add_agent_control(client, agent_name, final_output_control_id) + print("✓ Added final output validation control to agent") except Exception as e: if "409" in str(e) or "already" in str(e).lower(): - print(f"ℹ️ Policy already assigned to agent (OK)") + print("ℹ️ Final output validation control already associated with agent (OK)") else: - print(f"❌ Failed to assign policy: {e}") + print(f"❌ Failed to add final output validation control to agent: {e}") raise print("\n✅ Setup complete! You can now run content_agent_protection.py") diff --git a/examples/customer_support_agent/README.md b/examples/customer_support_agent/README.md index 30571390..3d3adf1b 100644 --- a/examples/customer_support_agent/README.md +++ b/examples/customer_support_agent/README.md @@ -160,7 +160,6 @@ Initialize once at application startup: import agent_control agent_control.init( - agent_name="Customer Support Agent", agent_name="646d5dea-c2e6-4453-b446-7035482b38e4", agent_description="AI-powered customer support assistant", ) @@ -168,7 +167,7 @@ agent_control.init( This: - Registers the agent with the server -- Fetches the assigned policy and controls +- Fetches controls associated with the agent - Enables the `@control()` decorator ### 2. Protecting Functions @@ -209,10 +208,10 @@ except ControlViolationError as e: **Important**: Controls are defined on the server via the UI, not in code. This design provides: -- **Centralized management**: Security team controls policies without code changes +- **Centralized management**: Security team controls safeguards without code changes - **Instant updates**: Change controls without redeploying agents - **Audit trail**: Server logs all control evaluations -- **Separation of concerns**: Developers focus on features, security team on policies +- **Separation of concerns**: Developers focus on features, security team on safeguards ## Project Structure diff --git a/examples/customer_support_agent/demo.sh b/examples/customer_support_agent/demo.sh index f71c0630..5959589d 100755 --- a/examples/customer_support_agent/demo.sh +++ b/examples/customer_support_agent/demo.sh @@ -222,13 +222,13 @@ setup_demo_controls() { python setup_demo_controls.py } -# Agent UUID used by the demo (must match support_agent.py) -get_agent_uuid() { +# Agent name used by the demo (must match support_agent.py) +get_agent_name() { echo "646d5dea-c2e6-4453-b446-7035482b38e4" } show_observability_stats() { - local agent_name=$(get_agent_uuid) + local agent_name=$(get_agent_name) local server_url="${AGENT_CONTROL_URL:-http://localhost:8000}" echo "" diff --git a/examples/customer_support_agent/run_demo.py b/examples/customer_support_agent/run_demo.py index 34ad71ef..5b612233 100644 --- a/examples/customer_support_agent/run_demo.py +++ b/examples/customer_support_agent/run_demo.py @@ -7,7 +7,7 @@ Usage: python run_demo.py # Interactive chat mode (default) python run_demo.py --automated # Run automated test scenarios - python run_demo.py --reset # Reset agent (remove policy/controls) and exit + python run_demo.py --reset # Reset agent (remove control associations) and exit Test Commands (in interactive mode): /test-safe Run safe message tests @@ -31,7 +31,7 @@ import os import sys -from agent_control import AgentControlClient, agents +from agent_control import AgentControlClient, agents, controls # Configure logging to see SDK debug output logging.basicConfig( @@ -59,7 +59,7 @@ async def reset_agent(): - """Reset the agent by removing its policy (which disconnects all controls).""" + """Reset the agent by removing direct control associations.""" agent_name = AGENT_ID server_url = os.getenv("AGENT_CONTROL_URL", "http://localhost:8000") @@ -81,22 +81,47 @@ async def reset_agent(): print(f"Error checking agent: {e}") return - # Remove policy from agent (disconnects all controls) - try: - await agents.remove_agent_policy(client, agent_name) - logger.info("Successfully removed policy from agent") - print("Removed policy from agent (all controls disconnected).") - except Exception as e: - if "404" in str(e): - logger.info("Agent has no policy attached") - print("Agent has no policy - already clean.") - return - logger.error(f"Error removing policy: {e}") - print(f"Error removing policy: {e}") - return + # Remove direct control associations (idempotent per control ID). + removed_direct_associations = 0 + cursor: int | None = None + while True: + controls_page = await controls.list_controls(client, cursor=cursor, limit=100) + control_summaries = controls_page.get("controls", []) + + for summary in control_summaries: + control_id = summary.get("id") + if control_id is None: + continue + try: + remove_result = await agents.remove_agent_control( + client, + agent_name, + control_id, + ) + if remove_result.get("removed_direct_association"): + removed_direct_associations += 1 + except Exception as e: + # Ignore transient 404s while iterating controls; keep best-effort cleanup. + if "404" not in str(e): + logger.debug( + "Error removing direct control %s from %s: %s", + control_id, + agent_name, + e, + ) + + pagination = controls_page.get("pagination", {}) + next_cursor = pagination.get("next_cursor") + if next_cursor is None: + break + cursor = int(next_cursor) + + print( + f"Removed {removed_direct_associations} direct control association(s) from agent." + ) print() - print("Reset complete. The agent now has no controls.") + print("Reset complete. The agent now has no direct control associations.") print("Run the demo again and add controls via the UI to test.") @@ -575,7 +600,7 @@ def main(): parser.add_argument( "--reset", action="store_true", - help="Reset the agent (remove policy/controls) and exit" + help="Reset the agent (remove control associations) and exit" ) args = parser.parse_args() diff --git a/examples/customer_support_agent/setup_demo_controls.py b/examples/customer_support_agent/setup_demo_controls.py index bce590fe..7b7b8a3c 100644 --- a/examples/customer_support_agent/setup_demo_controls.py +++ b/examples/customer_support_agent/setup_demo_controls.py @@ -5,15 +5,14 @@ This script: 1. Registers the agent with the server 2. Creates demo controls (PII detection, prompt injection) -3. Creates a policy and attaches controls -4. Assigns the policy to the agent +3. Directly associates controls to the agent Run this after starting the server to have a working demo out of the box. """ import asyncio import os -from agent_control import Agent, AgentControlClient, agents, controls, policies +from agent_control import Agent, AgentControlClient, agents, controls # Same agent ID as in support_agent.py AGENT_ID = "646d5dea-c2e6-4453-b446-7035482b38e4" @@ -252,41 +251,7 @@ async def setup_demo(quiet: bool = False): print(f" Error registering agent: {e}") return False - # Get or create a policy for the agent - policy_name = f"policy-{AGENT_ID}" - policy_id = None - - # Check if agent already has a policy - try: - policy_info = await agents.get_agent_policy(client, agent_name) - policy_id = policy_info.get("policy_id") - except Exception: - policy_id = None # No policy yet - - # Create policy if needed - if not policy_id: - try: - policy_result = await policies.create_policy(client, policy_name) - policy_id = policy_result["policy_id"] - print(f" Created policy: {policy_name}") - except Exception as e: - if "409" in str(e): - import time - policy_name = f"policy-{AGENT_ID}-{int(time.time())}" - policy_result = await policies.create_policy(client, policy_name) - policy_id = policy_result["policy_id"] - print(f" Created policy: {policy_name}") - else: - print(f" Error setting up policy: {e}") - return False - - try: - await policies.assign_policy_to_agent(client, agent_name, policy_id) - except Exception as e: - print(f" Error assigning policy: {e}") - return False - - # Create controls and add to policy + # Create controls and directly associate them to the agent controls_created = 0 for control_spec in DEMO_CONTROLS: control_name = control_spec["name"] @@ -312,9 +277,11 @@ async def setup_demo(quiet: bool = False): continue try: - await policies.add_control_to_policy(client, policy_id, control_id) + await agents.add_agent_control(client, agent_name, control_id) except Exception as e: - print(f" Error adding control '{control_name}' to policy: {e}") + if "409" in str(e) or "already" in str(e).lower(): + continue + print(f" Error adding control '{control_name}' to agent: {e}") continue if controls_created > 0: diff --git a/examples/customer_support_agent/support_agent.py b/examples/customer_support_agent/support_agent.py index 82273005..aa6fb56d 100644 --- a/examples/customer_support_agent/support_agent.py +++ b/examples/customer_support_agent/support_agent.py @@ -10,7 +10,7 @@ 4. Realistic enterprise patterns (mock services, multiple tools) NOTE: Controls are defined on the server via the UI, not in code. -This keeps security policies centrally managed and separate from code. +This keeps security controls centrally managed and separate from code. """ import asyncio @@ -26,7 +26,7 @@ # SDK INITIALIZATION # ============================================================================= # Call this once at the start of your application. -# The agent registers with the server and loads its assigned policy. +# The agent registers with the server and loads associated controls. agent_control.init( agent_name="646d5dea-c2e6-4453-b446-7035482b38e4", @@ -49,7 +49,7 @@ class MockLLM: RESPONSES = { "greeting": "Hello! I'm your customer support assistant. How can I help you today?", "refund": "I understand you'd like a refund. Let me look into your order. " - "Our refund policy allows returns within 30 days of purchase.", + "Our refund guidelines allow returns within 30 days of purchase.", "technical": "I can help with technical issues. Could you describe the problem " "you're experiencing in more detail?", "status": "I'll check the status of your order right away. " @@ -176,7 +176,7 @@ def create(cls, subject: str, description: str, priority: str = "medium") -> dic # PROTECTED AGENT FUNCTIONS # ============================================================================= # These functions are protected by the @control() decorator. -# The server evaluates controls and blocks/allows based on policy. +# The server evaluates controls and blocks/allows based on agent associations. @control() diff --git a/examples/deepeval/README.md b/examples/deepeval/README.md index e98e8ef4..12768a8c 100644 --- a/examples/deepeval/README.md +++ b/examples/deepeval/README.md @@ -380,7 +380,7 @@ python -m build # Creates dist/*.whl - Deploy your own agent-control server instance - Install custom evaluator packages (wheel, source, or private PyPI) - Your agents connect to this server via the SDK - - Complete control over evaluators and policies + - Complete control over evaluators and controls 2. **Managed Service (If Available)** - Use a hosted agent-control service @@ -455,7 +455,7 @@ You can create specialized evaluators for specific use cases: 2. **Extensibility**: The `Evaluator` base class makes it easy to integrate any evaluation library 3. **Configuration**: Pydantic models provide type-safe, validated configuration 4. **Registration**: The `@register_evaluator` decorator handles registration automatically -5. **Integration**: Evaluators work seamlessly with agent-control's policy system +5. **Integration**: Evaluators work seamlessly with agent-control's control system 6. **Control Logic**: `matched=True` triggers the action (deny/allow), so invert when quality passes ## Troubleshooting diff --git a/examples/deepeval/qa_agent.py b/examples/deepeval/qa_agent.py index b4552326..91bf4bda 100755 --- a/examples/deepeval/qa_agent.py +++ b/examples/deepeval/qa_agent.py @@ -38,7 +38,7 @@ agent_control.init( agent_name="qa-agent-deepeval", - agent_description="Question answering agent with DeepEval quality controls", + agent_description="Q&A Agent with DeepEval", agent_version="1.0.0", ) diff --git a/examples/deepeval/setup_controls.py b/examples/deepeval/setup_controls.py index f3178cb5..cbf2eda1 100755 --- a/examples/deepeval/setup_controls.py +++ b/examples/deepeval/setup_controls.py @@ -5,8 +5,7 @@ This script: 1. Registers the agent with the server 2. Creates DeepEval GEval evaluator controls for quality checks -3. Creates a policy and attaches controls -4. Assigns the policy to the agent +3. Directly associates controls to the agent The controls demonstrate using DeepEval's LLM-as-a-judge to enforce: - Response coherence @@ -19,8 +18,6 @@ import asyncio import os import sys -import uuid - import httpx # Add the current directory to the path so we can import the evaluator @@ -148,12 +145,11 @@ async def setup_demo(quiet: bool = False): """Set up the demo agent with DeepEval controls.""" - # Generate the same UUID5 that the SDK generates - agent_name = str(uuid.uuid5(uuid.NAMESPACE_DNS, AGENT_ID)) + agent_name = AGENT_ID print(f"Setting up agent: {AGENT_NAME}") print(f"Agent ID: {AGENT_ID}") - print(f"Agent UUID: {agent_name}") + print(f"Agent Name: {agent_name}") print(f"Server URL: {SERVER_URL}") print() @@ -176,10 +172,10 @@ async def setup_demo(quiet: bool = False): json={ "agent": { "agent_name": agent_name, - "agent_name": AGENT_NAME, "agent_description": AGENT_DESCRIPTION, + "agent_version": "1.0.0", }, - "tools": [], + "steps": [], }, ) resp.raise_for_status() @@ -190,48 +186,7 @@ async def setup_demo(quiet: bool = False): print(f"❌ Error registering agent: {e}") return False - # Get or create a policy for the agent - policy_name = f"policy-{AGENT_ID}" - policy_id = None - - # Check if agent already has a policy - try: - resp = await client.get(f"/api/v1/agents/{agent_name}/policy") - if resp.status_code == 200: - policy_id = resp.json().get("policy_id") - print(f"✓ Found existing policy: {policy_id}") - except httpx.HTTPError: - pass # No policy yet - - # Create policy if needed - if not policy_id: - try: - resp = await client.put( - "/api/v1/policies", - json={"name": policy_name}, - ) - if resp.status_code == 409: - # Policy name exists but not assigned - create with unique name - import time - - policy_name = f"policy-{AGENT_ID}-{int(time.time())}" - resp = await client.put( - "/api/v1/policies", - json={"name": policy_name}, - ) - resp.raise_for_status() - policy_id = resp.json()["policy_id"] - print(f"✓ Created policy: {policy_name}") - - # Assign policy to agent - resp = await client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") - resp.raise_for_status() - print(f"✓ Assigned policy to agent") - except httpx.HTTPError as e: - print(f"❌ Error setting up policy: {e}") - return False - - # Create controls and add to policy + # Create controls and associate them directly with the agent print() print("Creating DeepEval controls...") controls_created = 0 @@ -270,9 +225,10 @@ async def setup_demo(quiet: bool = False): ) resp.raise_for_status() - # Add control to policy - resp = await client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") - resp.raise_for_status() + # Associate control directly with the agent + resp = await client.post(f"/api/v1/agents/{agent_name}/controls/{control_id}") + if resp.status_code not in (200, 409): + resp.raise_for_status() status = "✓" if definition.get("enabled") else "○" enabled_text = "enabled" if definition.get("enabled") else "disabled" diff --git a/examples/galileo/README.md b/examples/galileo/README.md index 96e141d0..b87170db 100644 --- a/examples/galileo/README.md +++ b/examples/galileo/README.md @@ -41,7 +41,7 @@ The demo tests various inputs against a pre-configured toxicity detection stage: ### Central vs Local Stages -- **Central Stage** (used in this demo): Rulesets and policies are pre-configured on the Galileo server. Simply reference the stage by name. +- **Central Stage** (used in this demo): Rulesets are pre-configured on the Galileo server. Simply reference the stage by name. - **Local Stage**: Define rulesets at runtime in your code (see evaluator documentation). ### Expected Output @@ -85,4 +85,3 @@ Testing toxicity detection with Central Stage... - [Galileo Protect Overview](https://v2docs.galileo.ai/concepts/protect/overview) - [Luna-2 Python API Reference](https://v2docs.galileo.ai/sdk-api/python/reference/protect) - [Agent Control Luna-2 Evaluator](../../evaluators/extra/galileo/) - diff --git a/examples/langchain/README.md b/examples/langchain/README.md index d9169459..56db9793 100644 --- a/examples/langchain/README.md +++ b/examples/langchain/README.md @@ -41,8 +41,7 @@ uv run setup_sql_controls.py This creates: - SQL safety control (blocks DROP, DELETE, TRUNCATE, ALTER, GRANT) -- Policy with the control -- Assigns policy to the SQL agent +- Direct association of the control with the SQL agent > For local execution, create the control with `execution: "sdk"` in > `setup_sql_controls.py` (see `sql_control_data_sdk`) and enable @@ -145,17 +144,17 @@ pkill -f "uvicorn agent_control_server" cd server && make run ``` -### "Policy 'sql-protection-policy' already exists" +### "Control already associated with agent" -**Cause:** Setup script was run multiple times. +**Cause:** Setup script was run multiple times and the direct association already exists. -**Fix:** Either delete the policy via the API or use a different name in `setup_sql_controls.py`. +**Fix:** This is expected/idempotent; you can ignore and continue. ### DROP TABLE still executes **Causes:** 1. Server not running or evaluators not loaded (remote mode) -2. Control not assigned to agent's policy +2. Control not associated directly with the agent 3. Control data missing/invalid (control not returned to agent) 4. Local mode enabled but control is still `execution: "server"` @@ -192,7 +191,7 @@ Raise ControlViolationError (DENY) or Execute (ALLOW) ## Files - `sql_agent_protection.py` - Main SQL agent with `@control()` decorator -- `setup_sql_controls.py` - One-time setup script for controls/policy +- `setup_sql_controls.py` - One-time setup script for controls/direct associations - `pyproject.toml` - Dependencies and configuration - `README.md` - This file diff --git a/examples/langchain/langgraph_auto_schema_agent.py b/examples/langchain/langgraph_auto_schema_agent.py index b077d0df..b6afb26a 100644 --- a/examples/langchain/langgraph_auto_schema_agent.py +++ b/examples/langchain/langgraph_auto_schema_agent.py @@ -264,7 +264,7 @@ async def main() -> None: final_message = result["messages"][-1] print(f"Assistant: {final_message.content}") except ControlViolationError as exc: - print(f"Assistant: Request blocked by control policy: {exc.message}") + print(f"Assistant: Request blocked by control rules: {exc.message}") except RuntimeError as exc: print( "Assistant: Control evaluation is unavailable. " diff --git a/examples/langchain/setup_sql_controls.py b/examples/langchain/setup_sql_controls.py index 28512907..88ad4fbe 100644 --- a/examples/langchain/setup_sql_controls.py +++ b/examples/langchain/setup_sql_controls.py @@ -1,7 +1,7 @@ """ Setup script for SQL agent controls. -This script creates the SQL control and policy on the server. +This script creates the SQL control and associates it directly to the agent. Run this once before running sql_agent_protection.py. NOTE: This script is designed to run once. If resources already exist, @@ -16,7 +16,7 @@ import requests -from agent_control import Agent, AgentControlClient, agents, controls, policies +from agent_control import Agent, AgentControlClient, agents, controls AGENT_ID = "edf66504-0db5-4ee8-9e09-3ef37bbb8faa" SERVER_URL = os.getenv("AGENT_CONTROL_URL", "http://localhost:8000") @@ -48,7 +48,7 @@ def setup_database(): async def setup_sql_controls(): - """Create SQL control, policy, and assign to agent.""" + """Create SQL control and associate it directly with the agent.""" async with AgentControlClient(base_url=SERVER_URL) as client: # 1. Register Agent agent_name = AGENT_ID @@ -128,55 +128,15 @@ async def setup_sql_controls(): await controls.set_control_data(client, control_id, sql_control_data) print("✓ Control configuration updated") - # 3. Create Policy + # 3. Associate control directly with agent try: - policy_result = await policies.create_policy( - client, name="sql-protection-policy" - ) - policy_id = policy_result["policy_id"] - print(f"✓ Policy created (ID: {policy_id})") - except Exception as e: - if "409" in str(e): - print("ℹ️ Policy 'sql-protection-policy' already exists, checking agent...") - try: - policy_info = await agents.get_agent_policy(client, str(agent_name)) - policy_id = policy_info.get("policy_id") - if policy_id is None: - raise ValueError("No policy assigned to agent.") - print(f"ℹ️ Using agent's existing policy (ID: {policy_id})") - except Exception: - # Create a new policy name to avoid conflicts and keep script idempotent. - unique_name = f"sql-protection-policy-{uuid.uuid4().hex[:8]}" - print("⚠️ Policy exists but could not resolve its ID.") - print(f" Creating a new policy '{unique_name}' instead.") - policy_result = await policies.create_policy( - client, name=unique_name - ) - policy_id = policy_result["policy_id"] - print(f"✓ Policy created (ID: {policy_id})") - else: - raise - - # 4. Add Control to Policy - try: - await policies.add_control_to_policy(client, policy_id, control_id) - print(f"✓ Added control to policy") - except Exception as e: - if "409" in str(e) or "already" in str(e).lower(): - print(f"ℹ️ Control already in policy (OK)") - else: - print(f"❌ Failed to add control to policy: {e}") - raise - - # 5. Assign Policy to Agent - try: - await policies.assign_policy_to_agent(client, agent_name, policy_id) - print(f"✓ Assigned policy to agent") + await agents.add_agent_control(client, agent_name, control_id) + print("✓ Associated control directly with agent") except Exception as e: if "409" in str(e) or "already" in str(e).lower(): - print(f"ℹ️ Policy already assigned to agent (OK)") + print("ℹ️ Control already associated with agent (OK)") else: - print(f"❌ Failed to assign policy: {e}") + print(f"❌ Failed to associate control with agent: {e}") raise print("\n✅ Setup complete! You can now run sql_agent_protection.py") @@ -192,5 +152,5 @@ async def setup_sql_controls(): setup_database() print() - # Step 2: Setup controls and policies + # Step 2: Setup controls and direct agent associations asyncio.run(setup_sql_controls()) diff --git a/examples/langchain/sql_agent_protection.py b/examples/langchain/sql_agent_protection.py index b0fd84f6..eb051085 100644 --- a/examples/langchain/sql_agent_protection.py +++ b/examples/langchain/sql_agent_protection.py @@ -5,7 +5,7 @@ with graceful error handling. PREREQUISITE: - Run setup_sql_controls.py FIRST to create the SQL control and policy: + Run setup_sql_controls.py FIRST to create the SQL control and direct agent association: $ uv run setup_sql_controls.py diff --git a/examples/steer_action_demo/README.md b/examples/steer_action_demo/README.md index dca3b531..73abdf6e 100644 --- a/examples/steer_action_demo/README.md +++ b/examples/steer_action_demo/README.md @@ -13,7 +13,7 @@ This example shows all three AgentControl action types in a real-world banking s ## Understanding Steer Actions -**Steer is a non-fatal control signal** - unlike DENY which blocks execution, STEER provides corrective guidance to help agents satisfy policy requirements: +**Steer is a non-fatal control signal** - unlike DENY which blocks execution, STEER provides corrective guidance to help agents satisfy control requirements: - **Philosophy**: Agents are expected to correct the issue and retry - **Behavior**: Raises `ControlSteerError` with structured guidance diff --git a/examples/steer_action_demo/autonomous_agent_demo.py b/examples/steer_action_demo/autonomous_agent_demo.py index 20bfd7db..a2078f68 100644 --- a/examples/steer_action_demo/autonomous_agent_demo.py +++ b/examples/steer_action_demo/autonomous_agent_demo.py @@ -134,7 +134,7 @@ def check_fraud_score(amount: float, destination: str) -> float: """Check fraud risk based on destination. Note: Amount thresholds are handled by steer controls, not here. - Controls are the source of truth for policy decisions. + Controls are the source of truth for transfer decisions. """ agent_think("Running fraud detection...") if "north korea" in destination.lower() or "iran" in destination.lower(): @@ -311,7 +311,7 @@ async def process_transfer_node(state: AgentState) -> AgentState: else: fraud_score = state["fraud_score"] - agent_think("Checking compliance and policy controls...") + agent_think("Checking compliance controls...") log_trace("agent-control", "Initiating pre-execution control evaluation") agent_reason("All transfers must pass control checks before execution") @@ -341,7 +341,7 @@ async def process_transfer_node(state: AgentState) -> AgentState: # Success! log_trace("agent-control", "All control checks passed ✓") log_trace("execution", "Wire transfer executed successfully") - agent_reason("Transfer complies with all policies - proceeding with execution") + agent_reason("Transfer complies with all controls - proceeding with execution") # Show if this was a successful retry after steer correction if state.get("verified_2fa") or state.get("manager_approved"): @@ -372,8 +372,8 @@ async def process_transfer_node(state: AgentState) -> AgentState: except ControlViolationError as e: # DENY - Hard block log_trace("agent-control", f"DENY control triggered: {e.control_name}") - log_trace("compliance", "Transaction blocked by compliance policy") - agent_reason("This transaction violates a hard policy rule - cannot proceed") + log_trace("compliance", "Transaction blocked by compliance control") + agent_reason("This transaction violates a hard control rule - cannot proceed") agent_say(f"❌ I cannot process this transfer.") # Control Evaluation Summary @@ -393,7 +393,7 @@ async def process_transfer_node(state: AgentState) -> AgentState: print(f" \033[90m{e.metadata}\033[0m") log_trace("audit", f"Blocked transaction: Control={e.control_name}, Amount=${request['amount']:,.2f}") - agent_say("This transaction violates compliance policies and cannot be approved under any circumstances.") + agent_say("This transaction violates compliance controls and cannot be approved under any circumstances.") return { **state, @@ -407,7 +407,7 @@ async def process_transfer_node(state: AgentState) -> AgentState: log_trace("control", f"Steer action triggered by control '{e.control_name}'") log_trace("decision", "Transfer cannot proceed without additional approvals") - agent_reason("Control policy requires verification before large transactions") + agent_reason("Control rules require verification before large transactions") # Parse structured steering context (deterministic, no LLM needed) agent_think("Parsing steering context...") @@ -492,7 +492,7 @@ async def process_transfer_node(state: AgentState) -> AgentState: elif action == "approval" and not state.get("manager_approved"): log_escalation("ESCALATING TO: Manager Approval System") agent_reason(f"Control requires manager approval for this ${request['amount']:,.2f} transfer") - log_trace("compliance", "Manager approval required per control policy") + log_trace("compliance", "Manager approval required per control rules") # Get justification if not state.get("justification"): @@ -544,7 +544,7 @@ async def process_transfer_node(state: AgentState) -> AgentState: if approval.lower() in ['yes', 'y']: log_trace("approval-system", "Manager approved the transfer") log_trace("audit", f"Approval logged: Manager ID: MGR-001, Timestamp: {time.strftime('%Y-%m-%d %H:%M:%S')}") - agent_reason("Manager authorization received - transfer now complies with policy") + agent_reason("Manager authorization received - transfer now complies with controls") agent_say("✅ Manager approved!") print(f"\n 🔄 \033[1mSTEER CORRECTION COMPLETE - RETRYING TRANSFER\033[0m") diff --git a/examples/steer_action_demo/setup_controls.py b/examples/steer_action_demo/setup_controls.py index c57b5d4a..f6d8a0a2 100644 --- a/examples/steer_action_demo/setup_controls.py +++ b/examples/steer_action_demo/setup_controls.py @@ -2,7 +2,7 @@ Setup script for Banking Transaction Agent controls. Demonstrates all AgentControl action types in a realistic banking scenario: -- ALLOW: Simple transfers within policy +- ALLOW: Simple transfers that satisfy controls - DENY: Compliance violations (sanctioned countries, fraud) - STEER: Large transfers requiring verification and approval @@ -12,7 +12,7 @@ import asyncio import os -from agent_control import Agent, AgentControlClient, agents, controls, policies +from agent_control import Agent, AgentControlClient, agents, controls AGENT_ID = "f8e5d3c2-4b1a-4e7f-9c8d-2a3b4c5d6e7f" SERVER_URL = os.getenv("AGENT_CONTROL_URL", "http://localhost:8000") @@ -238,59 +238,18 @@ async def setup_banking_controls(): print(f" Response: {e.response.text if hasattr(e.response, 'text') else e.response}") raise - # 3. Create Policy - print("\n📋 Creating policy...") - try: - policy_result = await policies.create_policy( - client, - name="banking-transaction-policy" - ) - policy_id = policy_result["policy_id"] - print(f" ✓ Created policy: banking-transaction-policy (ID: {policy_id})") - except Exception as e: - if "409" in str(e): - print(f" ℹ️ Policy 'banking-transaction-policy' already exists") - try: - policy_info = await agents.get_agent_policy(client, agent.agent_name) - policy_id = policy_info.get("policy_id") - if policy_id: - print(f" ℹ️ Using agent's existing policy (ID: {policy_id})") - else: - raise ValueError("No policy assigned") - except Exception: - import uuid as uuid_module - unique_name = f"banking-policy-{uuid_module.uuid4().hex[:8]}" - print(f" ⚠️ Creating new policy '{unique_name}' instead") - policy_result = await policies.create_policy(client, name=unique_name) - policy_id = policy_result["policy_id"] - print(f" ✓ Created policy (ID: {policy_id})") - else: - print(f" ❌ Error creating policy: {e}") - raise - - # 4. Add Controls to Policy - print("\n📋 Adding controls to policy...") + # 3. Associate controls directly with the agent + print("\n📋 Associating controls directly with agent...") for control_id in control_ids: try: - await policies.add_control_to_policy(client, policy_id, control_id) - print(f" ✓ Added control {control_id} to policy") + await agents.add_agent_control(client, agent.agent_name, control_id) + print(f" ✓ Added control {control_id} to agent") except Exception as e: if "409" in str(e) or "already" in str(e).lower(): - print(f" ℹ️ Control {control_id} already in policy (OK)") + print(f" ℹ️ Control {control_id} already associated with agent (OK)") else: print(f" ⚠️ Failed to add control: {e}") - # 5. Assign Policy to Agent - print("\n📋 Assigning policy to agent...") - try: - await policies.assign_policy_to_agent(client, agent.agent_name, policy_id) - print(f" ✓ Assigned policy {policy_id} to agent {agent.agent_name}") - except Exception as e: - if "409" in str(e) or "already" in str(e).lower(): - print(f" ℹ️ Policy already assigned to agent (OK)") - else: - print(f" ⚠️ Failed to assign policy: {e}") - print("\n✅ Setup complete!") print(f"\nAgent ID: {AGENT_ID}") print(f"Server URL: {SERVER_URL}") diff --git a/models/src/agent_control_models/server.py b/models/src/agent_control_models/server.py index 693c439a..f32c9f44 100644 --- a/models/src/agent_control_models/server.py +++ b/models/src/agent_control_models/server.py @@ -176,7 +176,7 @@ class InitAgentResponse(BaseModel): ) controls: list[Control] = Field( default_factory=list, - description="Active protection controls for the agent (if policy assigned)", + description="Active protection controls for the agent", ) overwrite_applied: bool = Field( default=False, @@ -201,24 +201,37 @@ class CreatePolicyResponse(BaseModel): policy_id: int = Field(description="Identifier of the created policy") +class GetAgentPoliciesResponse(BaseModel): + policy_ids: list[int] = Field( + default_factory=list, description="IDs of policies associated with the agent" + ) + + class SetPolicyResponse(BaseModel): - success: bool = Field(description="Whether the policy was successfully assigned") + """Compatibility response for singular policy assignment endpoint.""" + + success: bool = Field(description="Whether the request succeeded") old_policy_id: int | None = Field( - default=None, description="Previous policy id if one was replaced" + default=None, + description="Previously associated policy ID, if any", ) class GetPolicyResponse(BaseModel): - policy_id: int = Field(description="Identifier of the policy assigned to the agent") + """Compatibility response for singular policy retrieval endpoint.""" + + policy_id: int = Field(description="Associated policy ID") class DeletePolicyResponse(BaseModel): - success: bool = Field(description="Whether the policy was successfully removed") + """Compatibility response for singular policy deletion endpoint.""" + + success: bool = Field(description="Whether the request succeeded") class AgentControlsResponse(BaseModel): controls: list[Control] = Field( - description="List of controls associated with the agent via its policy" + description="List of active controls associated with the agent" ) @@ -248,6 +261,18 @@ class AssocResponse(BaseModel): success: bool = Field(description="Whether the association change succeeded") +class RemoveAgentControlResponse(BaseModel): + """Response for removing a direct agent-control association.""" + + success: bool = Field(description="Whether the request succeeded") + removed_direct_association: bool = Field( + description="True if a direct agent-control link was removed" + ) + control_still_active: bool = Field( + description="True if the control remains active via policy association(s)" + ) + + class GetControlDataResponse(BaseModel): data: ControlDefinition = Field(description="Control data payload") @@ -310,12 +335,14 @@ class AgentSummary(BaseModel): """Summary of an agent for list responses.""" agent_name: str = Field(..., description="Unique identifier of the agent") - policy_id: int | None = Field(None, description="ID of assigned policy, if any") + policy_ids: list[int] = Field( + default_factory=list, description="IDs of policies associated with the agent" + ) created_at: str | None = Field(None, description="ISO 8601 timestamp when agent was created") step_count: int = Field(0, description="Number of steps registered with the agent") evaluator_count: int = Field(0, description="Number of evaluators registered with the agent") active_controls_count: int = Field( - 0, description="Number of active controls from agent's policy" + 0, description="Number of active controls for this agent" ) @@ -345,7 +372,7 @@ class ListAgentsResponse(BaseModel): class AgentRef(BaseModel): """Reference to an agent (for listing which agents use a control).""" - agent_name: str = Field(..., description="Agent identifier") + agent_name: str = Field(..., description="Agent name") class ControlSummary(BaseModel): @@ -360,6 +387,10 @@ class ControlSummary(BaseModel): stages: list[str] | None = Field(None, description="Evaluation stages in scope") tags: list[str] = Field(default_factory=list, description="Control tags") used_by_agent: AgentRef | None = Field(None, description="Agent using this control") + # TODO: Follow-up with full `used_by_agents` list for richer attribution. + used_by_agents_count: int = Field( + 0, description="Number of unique agents using this control" + ) class ListControlsResponse(BaseModel): @@ -374,9 +405,17 @@ class DeleteControlResponse(BaseModel): success: bool = Field(..., description="Whether the control was deleted") dissociated_from: list[int] = Field( + default_factory=list, + description="Deprecated: policy IDs the control was removed from before deletion", + ) + dissociated_from_policies: list[int] = Field( default_factory=list, description="Policy IDs the control was removed from before deletion", ) + dissociated_from_agents: list[str] = Field( + default_factory=list, + description="Agent names the control was removed from before deletion", + ) class PatchControlRequest(BaseModel): diff --git a/sdks/python/README.md b/sdks/python/README.md index 179306bc..ced0c4df 100644 --- a/sdks/python/README.md +++ b/sdks/python/README.md @@ -17,8 +17,8 @@ import agent_control # Initialize at the base of your agent agent_control.init( - agent_name="My Customer Service Bot", - agent_name="550e8400-e29b-41d4-a716-446655440000" + agent_name="550e8400-e29b-41d4-a716-446655440000", + agent_description="My Customer Service Bot", ) # Use the control decorator @@ -35,7 +35,6 @@ async def handle_message(message: str): import agent_control agent_control.init( - agent_name="Customer Service Bot", agent_name="550e8400-e29b-41d4-a716-446655440000", agent_description="Handles customer inquiries and support", agent_version="2.1.0", @@ -54,8 +53,8 @@ One line to set up your agent with full protection: ```python agent_control.init( - agent_name="...", agent_name="550e8400-e29b-41d4-a716-446655440000", + agent_description="...", ) ``` @@ -103,7 +102,7 @@ Access your agent information: ```python agent = agent_control.current_agent() print(f"Agent: {agent.agent_name}") -print(f"ID: {agent.agent_name}") +print(f"Name: {agent.agent_name}") print(f"Version: {agent.agent_version}") ``` @@ -116,8 +115,8 @@ from agent_control import control, ControlViolationError # Initialize agent_control.init( - agent_name="Customer Support Bot", agent_name="550e8400-e29b-41d4-a716-446655440000", + agent_description="Customer Support Bot", agent_version="1.0.0" ) @@ -158,11 +157,15 @@ asyncio.run(main()) ```python def init( agent_name: str, - agent_name: str | UUID, agent_description: Optional[str] = None, agent_version: Optional[str] = None, server_url: Optional[str] = None, + api_key: Optional[str] = None, controls_file: Optional[str] = None, + steps: Optional[list[dict]] = None, + conflict_mode: Literal["strict", "overwrite"] = "overwrite", + observability_enabled: Optional[bool] = None, + log_config: Optional[dict[str, Any]] = None, policy_refresh_interval_seconds: int = 60, **kwargs ) -> Agent: @@ -171,12 +174,16 @@ def init( Initialize Agent Control with your agent's information. **Parameters:** -- `agent_name`: Human-readable name -- `agent_name`: UUID string (or UUID instance) +- `agent_name`: Unique identifier for the agent - `agent_description`: Optional description - `agent_version`: Optional version string - `server_url`: Optional server URL (defaults to `AGENT_CONTROL_URL` env var) +- `api_key`: Optional API key (defaults to `AGENT_CONTROL_API_KEY` env var) - `controls_file`: Optional controls file path (auto-discovered if not provided) +- `steps`: Optional step schemas to register with initAgent +- `conflict_mode`: initAgent registration conflict mode (`"strict"` or `"overwrite"`) +- `observability_enabled`: Optional observability toggle +- `log_config`: Optional SDK logging config - `policy_refresh_interval_seconds`: Background cache refresh interval in seconds. Default `60`; set to `0` to disable background refresh. - `**kwargs`: Additional metadata @@ -184,7 +191,7 @@ Initialize Agent Control with your agent's information. **Returns:** `Agent` instance When background refresh is enabled, the SDK refreshes cache snapshots via -`GET /agents/{agent_id}/controls`. On refresh failures, it keeps the previous +`GET /agents/{agent_name}/controls`. On refresh failures, it keeps the previous snapshot (fail-open behavior). ### Decorator @@ -198,7 +205,7 @@ def control(policy: Optional[str] = None): Decorator to protect a function with server-defined controls. **Parameters:** -- `policy`: Optional policy name to use (defaults to agent's assigned policy) +- `policy`: Optional policy label for code readability when multiple policies exist **Example:** ```python @@ -290,8 +297,8 @@ import agent_control from agent_control import control, ControlViolationError agent_control.init( - agent_name="...", agent_name="550e8400-e29b-41d4-a716-446655440000", + agent_description="...", ) @control() diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index 621c9a08..ce5286b2 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -18,7 +18,8 @@ async def chat(message: str) -> str: return await assistant.respond(message) - # Apply all controls for this agent + # Optional policy label for grouping/readability; control selection still follows + # the server's active controls (policy + direct associations). @agent_control.control(policy="safety-policy") async def process(input: str) -> str: return await pipeline.run(input) @@ -375,10 +376,10 @@ def init( ] ) - # Now use @control decorator to apply the agent's policy + # Now use @control decorator to apply agent-associated controls from agent_control import control - @control() # Applies agent's assigned policy + @control() # Applies controls associated with the agent async def handle(message: str): return message @@ -627,7 +628,7 @@ async def list_agents( Returns: Dictionary containing: - agents: List of agent summaries with agent_name, - policy_id, created_at, step_count, evaluator_count + policy_ids, created_at, step_count, evaluator_count - pagination: Object with limit, total, next_cursor, has_more Raises: @@ -656,6 +657,87 @@ async def main(): return await agents.list_agents(client, cursor=cursor, limit=limit) +# ============================================================================ +# Agent Association Convenience Functions +# ============================================================================ + + +async def get_agent_policies( + agent_name: str, + server_url: str | None = None, + api_key: str | None = None, +) -> dict[str, Any]: + """List policy IDs associated with an agent.""" + _final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' + + async with AgentControlClient(base_url=_final_server_url, api_key=api_key) as client: + return await agents.get_agent_policies(client, agent_name) + + +async def add_agent_policy( + agent_name: str, + policy_id: int, + server_url: str | None = None, + api_key: str | None = None, +) -> dict[str, Any]: + """Associate a policy with an agent (additive, idempotent).""" + _final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' + + async with AgentControlClient(base_url=_final_server_url, api_key=api_key) as client: + return await agents.add_agent_policy(client, agent_name, policy_id) + + +async def remove_agent_policy_association( + agent_name: str, + policy_id: int, + server_url: str | None = None, + api_key: str | None = None, +) -> dict[str, Any]: + """Remove one policy association from an agent (idempotent).""" + _final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' + + async with AgentControlClient(base_url=_final_server_url, api_key=api_key) as client: + return await agents.remove_agent_policy_association(client, agent_name, policy_id) + + +async def remove_all_agent_policies( + agent_name: str, + server_url: str | None = None, + api_key: str | None = None, +) -> dict[str, Any]: + """Remove all policy associations from an agent.""" + _final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' + + async with AgentControlClient(base_url=_final_server_url, api_key=api_key) as client: + return await agents.remove_all_agent_policies(client, agent_name) + + +async def add_agent_control( + agent_name: str, + control_id: int, + server_url: str | None = None, + api_key: str | None = None, +) -> dict[str, Any]: + """Associate a control directly with an agent (idempotent).""" + _final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' + + async with AgentControlClient(base_url=_final_server_url, api_key=api_key) as client: + return await agents.add_agent_control(client, agent_name, control_id) + + +async def remove_agent_control( + agent_name: str, + control_id: int, + server_url: str | None = None, + api_key: str | None = None, +) -> dict[str, Any]: + """Remove a direct control association from an agent (idempotent).""" + _final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' + + async with AgentControlClient(base_url=_final_server_url, api_key=api_key) as client: + return await agents.remove_agent_control(client, agent_name, control_id) + + # ============================================================================ # Control Management Convenience Functions # ============================================================================ @@ -836,7 +918,7 @@ async def delete_control( """ Delete a control from the server. - By default, deletion fails if the control is associated with any policy. + By default, deletion fails if the control is associated with any policy or agent. Use force=True to automatically dissociate and delete. Args: @@ -848,7 +930,8 @@ async def delete_control( Returns: Dictionary containing: - success: True if control was deleted - - dissociated_from: List of policy IDs the control was removed from + - dissociated_from_policies: List of policy IDs the control was removed from + - dissociated_from_agents: List of agent names the control was removed from Raises: httpx.HTTPError: If request fails @@ -862,7 +945,11 @@ async def delete_control( async def main(): # Force delete result = await agent_control.delete_control(5, force=True) - print(f"Deleted, removed from {len(result['dissociated_from'])} policies") + print( + "Deleted, removed from " + f"{len(result['dissociated_from_policies'])} policies and " + f"{len(result['dissociated_from_agents'])} agents" + ) asyncio.run(main()) """ @@ -1066,6 +1153,12 @@ async def main(): # Agent management "get_agent", "list_agents", + "get_agent_policies", + "add_agent_policy", + "remove_agent_policy_association", + "remove_all_agent_policies", + "add_agent_control", + "remove_agent_control", # Control management "create_control", "list_controls", diff --git a/sdks/python/src/agent_control/agents.py b/sdks/python/src/agent_control/agents.py index b14e69a8..3d63cc5e 100644 --- a/sdks/python/src/agent_control/agents.py +++ b/sdks/python/src/agent_control/agents.py @@ -54,11 +54,61 @@ async def list_agents( return cast(dict[str, Any], response.json()) +async def get_agent_policies( + client: AgentControlClient, + agent_name: str, +) -> dict[str, Any]: + """List policy IDs associated with an agent.""" + normalized_name = ensure_agent_name(agent_name) + response = await client.http_client.get(f"/api/v1/agents/{normalized_name}/policies") + response.raise_for_status() + return cast(dict[str, Any], response.json()) + + +async def add_agent_policy( + client: AgentControlClient, + agent_name: str, + policy_id: int, +) -> dict[str, Any]: + """Associate a policy with an agent (additive, idempotent).""" + normalized_name = ensure_agent_name(agent_name) + response = await client.http_client.post( + f"/api/v1/agents/{normalized_name}/policies/{policy_id}" + ) + response.raise_for_status() + return cast(dict[str, Any], response.json()) + + +async def remove_agent_policy_association( + client: AgentControlClient, + agent_name: str, + policy_id: int, +) -> dict[str, Any]: + """Remove one policy association from an agent (idempotent).""" + normalized_name = ensure_agent_name(agent_name) + response = await client.http_client.delete( + f"/api/v1/agents/{normalized_name}/policies/{policy_id}" + ) + response.raise_for_status() + return cast(dict[str, Any], response.json()) + + +async def remove_all_agent_policies( + client: AgentControlClient, + agent_name: str, +) -> dict[str, Any]: + """Remove all policy associations from an agent.""" + normalized_name = ensure_agent_name(agent_name) + response = await client.http_client.delete(f"/api/v1/agents/{normalized_name}/policies") + response.raise_for_status() + return cast(dict[str, Any], response.json()) + + async def get_agent_policy( client: AgentControlClient, agent_name: str, ) -> dict[str, Any]: - """Get the policy assigned to an agent.""" + """Get the primary policy assigned to an agent (compatibility endpoint).""" normalized_name = ensure_agent_name(agent_name) response = await client.http_client.get(f"/api/v1/agents/{normalized_name}/policy") response.raise_for_status() @@ -69,13 +119,41 @@ async def remove_agent_policy( client: AgentControlClient, agent_name: str, ) -> dict[str, Any]: - """Remove the policy assignment from an agent.""" + """Remove all policy associations via singular compatibility endpoint.""" normalized_name = ensure_agent_name(agent_name) response = await client.http_client.delete(f"/api/v1/agents/{normalized_name}/policy") response.raise_for_status() return cast(dict[str, Any], response.json()) +async def add_agent_control( + client: AgentControlClient, + agent_name: str, + control_id: int, +) -> dict[str, Any]: + """Associate a control directly with an agent (idempotent).""" + normalized_name = ensure_agent_name(agent_name) + response = await client.http_client.post( + f"/api/v1/agents/{normalized_name}/controls/{control_id}" + ) + response.raise_for_status() + return cast(dict[str, Any], response.json()) + + +async def remove_agent_control( + client: AgentControlClient, + agent_name: str, + control_id: int, +) -> dict[str, Any]: + """Remove a direct control association from an agent (idempotent).""" + normalized_name = ensure_agent_name(agent_name) + response = await client.http_client.delete( + f"/api/v1/agents/{normalized_name}/controls/{control_id}" + ) + response.raise_for_status() + return cast(dict[str, Any], response.json()) + + async def list_agent_controls( client: AgentControlClient, agent_name: str, diff --git a/sdks/python/src/agent_control/control_decorators.py b/sdks/python/src/agent_control/control_decorators.py index 24789513..3370be1d 100644 --- a/sdks/python/src/agent_control/control_decorators.py +++ b/sdks/python/src/agent_control/control_decorators.py @@ -1,12 +1,12 @@ """ Control decorator for server-side protection of agent functions. -This module provides a decorator that applies server-defined policies to agent functions. -Policies contain multiple controls (regex, list, Luna2, etc.) that are managed server-side. +This module provides a decorator that applies server-defined controls to agent functions. +Controls can be associated via policies and direct agent-control links. Architecture: SERVER defines: Policies -> Controls (stage, selector, evaluator, action) - SDK decorator: just marks WHERE the policy applies + SDK decorator: just marks WHERE controls are evaluated Usage: import agent_control @@ -15,12 +15,12 @@ agent_name="my-agent-identity", ) - # Apply the agent's assigned policy + # Apply controls associated with the agent @agent_control.control() async def chat(message: str) -> str: return await assistant.respond(message) - # The server's policy contains controls that define: + # Server-side controls define: # - stage: "pre" or "post" # - selector.path: "input" or "output" # - evaluator: regex, list, Luna2 evaluator, etc. @@ -726,15 +726,14 @@ async def _execute_with_control( def control(policy: str | None = None, step_name: str | None = None) -> Callable[[F], F]: """ - Decorator to apply server-defined policy at this code location. + Decorator to apply server-defined controls at this code location. - The policy's controls (stage, selector, evaluator, action) are defined - on the SERVER. This decorator just marks WHERE to apply the policy. + Controls (stage, selector, evaluator, action) are defined on the SERVER. + This decorator marks WHERE to evaluate controls for the current agent. Args: - policy: Optional policy name for documentation. The agent's assigned - policy is automatically used. This parameter is for clarity - in code when multiple policies exist. + policy: Optional policy name for documentation. This parameter is for + clarity in code when multiple policies exist. step_name: Optional custom name for this step. If not provided, uses the function name. @@ -746,20 +745,20 @@ def control(policy: str | None = None, step_name: str | None = None) -> Callable How it works: 1. Before function execution: Calls server with stage="pre" - - Server evaluates all "pre" controls in the agent's policy + - Server evaluates all matching "pre" controls for the agent 2. Function executes 3. After function execution: Calls server with stage="post" - - Server evaluates all "post" controls in the agent's policy + - Server evaluates all matching "post" controls for the agent Example: import agent_control - # Initialize agent (connects to server, loads policy) + # Initialize agent (connects to server, loads control associations) agent_control.init( agent_name="my-bot-identity", ) - # Apply the agent's policy (all controls) + # Apply controls associated with the agent @agent_control.control() async def chat(message: str) -> str: return await assistant.respond(message) @@ -783,11 +782,12 @@ async def handle_user_input(user_message: str) -> str: PUT /api/v1/policies {"name": "safety-policy"} POST /api/v1/policies/{policy_id}/controls/{control_id} - 3. Assign policy to agent: - POST /api/v1/agents/{agent_name}/policy/{policy_id} + 3. Associate policy and/or controls with an agent: + POST /api/v1/agents/{agent_name}/policies/{policy_id} + POST /api/v1/agents/{agent_name}/controls/{control_id} """ - # The policy parameter is for documentation only - the server uses - # the agent's assigned policy automatically + # The policy parameter is for documentation only - the server evaluates + # controls associated with the agent via policy and direct links. _ = policy def decorator(func: F) -> F: diff --git a/sdks/python/src/agent_control/policies.py b/sdks/python/src/agent_control/policies.py index aef75bbf..99420b19 100644 --- a/sdks/python/src/agent_control/policies.py +++ b/sdks/python/src/agent_control/policies.py @@ -155,9 +155,11 @@ async def assign_policy_to_agent( policy_id: int ) -> dict[str, Any]: """ - Assign a policy to an agent. + Assign a single policy to an agent via compatibility endpoint. - This makes the policy active for the agent. Any existing policy assignment is replaced. + This call uses replace semantics for backward compatibility: + existing policy associations are removed and replaced with ``policy_id``. + For additive behavior, use ``agents.add_agent_policy(...)``. Args: client: AgentControlClient instance @@ -171,9 +173,9 @@ async def assign_policy_to_agent( httpx.HTTPError: If request fails HTTPException 404: Agent or policy not found """ - agent_name_str = ensure_agent_name(agent_name) + normalized_name = ensure_agent_name(agent_name) response = await client.http_client.post( - f"/api/v1/agents/{agent_name_str}/policy/{policy_id}" + f"/api/v1/agents/{normalized_name}/policy/{policy_id}" ) response.raise_for_status() return cast(dict[str, Any], response.json()) diff --git a/sdks/python/tests/README.md b/sdks/python/tests/README.md index b4677039..b1dada52 100644 --- a/sdks/python/tests/README.md +++ b/sdks/python/tests/README.md @@ -152,7 +152,7 @@ PASSED ### Function-Scoped Fixtures - `client`: Authenticated AgentProtectClient instance - `unique_name`: Unique name generator for test resources -- `test_agent_id`: Unique agent UUID for testing +- `test_agent_name`: Unique agent UUID for testing - `test_agent`: Registered test agent - `test_policy`: Created test policy - `test_control`: Created test control diff --git a/sdks/python/tests/conftest.py b/sdks/python/tests/conftest.py index 6860c8e9..eabd7ec9 100644 --- a/sdks/python/tests/conftest.py +++ b/sdks/python/tests/conftest.py @@ -100,7 +100,7 @@ def unique_name() -> str: @pytest.fixture -def test_agent_id() -> str: +def test_agent_name() -> str: """Generate a unique agent name for testing.""" return f"agent-{uuid.uuid4().hex[:12]}" @@ -108,7 +108,7 @@ def test_agent_id() -> str: @pytest_asyncio.fixture async def test_agent( client: agent_control.AgentControlClient, - test_agent_id: str, + test_agent_name: str, server_url: str ) -> AsyncGenerator[dict[str, Any], None]: """ @@ -122,7 +122,7 @@ async def test_agent( from agent_control_models import Agent agent = Agent( - agent_name=test_agent_id, + agent_name=test_agent_name, agent_description="Integration test agent", agent_created_at=datetime.now(UTC).isoformat(), agent_updated_at=None, @@ -146,8 +146,8 @@ async def test_agent( yield { "agent": agent, - "agent_name": test_agent_id, - "agent_name": test_agent_id, + "agent_name": test_agent_name, + "agent_name": test_agent_name, "response": response } diff --git a/sdks/python/tests/test_agent_id_validation.py b/sdks/python/tests/test_agent_id_validation.py index 20277e22..37fdc265 100644 --- a/sdks/python/tests/test_agent_id_validation.py +++ b/sdks/python/tests/test_agent_id_validation.py @@ -1,8 +1,8 @@ """SDK agent name validation behavior tests.""" from unittest.mock import AsyncMock, MagicMock -import pytest +import pytest from agent_control import agents, policies @@ -38,6 +38,18 @@ async def test_get_agent_policy_rejects_invalid_agent_name() -> None: client.http_client.get.assert_not_called() +@pytest.mark.asyncio +async def test_get_agent_policies_rejects_invalid_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.get = AsyncMock() + + with pytest.raises(ValueError, match="at least 10 characters"): + await agents.get_agent_policies(client, "short") + + client.http_client.get.assert_not_called() + + @pytest.mark.asyncio async def test_remove_agent_policy_rejects_invalid_agent_name() -> None: client = MagicMock() @@ -50,6 +62,66 @@ async def test_remove_agent_policy_rejects_invalid_agent_name() -> None: client.http_client.delete.assert_not_called() +@pytest.mark.asyncio +async def test_remove_all_agent_policies_rejects_invalid_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.delete = AsyncMock() + + with pytest.raises(ValueError, match="at least 10 characters"): + await agents.remove_all_agent_policies(client, "short") + + client.http_client.delete.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_agent_policy_rejects_invalid_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.post = AsyncMock() + + with pytest.raises(ValueError, match="at least 10 characters"): + await agents.add_agent_policy(client, "short", policy_id=1) + + client.http_client.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_remove_agent_policy_association_rejects_invalid_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.delete = AsyncMock() + + with pytest.raises(ValueError, match="at least 10 characters"): + await agents.remove_agent_policy_association(client, "short", policy_id=1) + + client.http_client.delete.assert_not_called() + + +@pytest.mark.asyncio +async def test_add_agent_control_rejects_invalid_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.post = AsyncMock() + + with pytest.raises(ValueError, match="at least 10 characters"): + await agents.add_agent_control(client, "short", control_id=1) + + client.http_client.post.assert_not_called() + + +@pytest.mark.asyncio +async def test_remove_agent_control_rejects_invalid_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.delete = AsyncMock() + + with pytest.raises(ValueError, match="at least 10 characters"): + await agents.remove_agent_control(client, "short", control_id=1) + + client.http_client.delete.assert_not_called() + + @pytest.mark.asyncio async def test_list_agents_normalizes_cursor() -> None: client = MagicMock() @@ -82,8 +154,7 @@ async def test_get_agent_normalizes_agent_name() -> None: client.http_client = MagicMock() client.http_client.get = AsyncMock(return_value=DummyResponse()) - agent_name = "Agent-Example_01" - await agents.get_agent(client, agent_name) + await agents.get_agent(client, "Agent-Example_01") client.http_client.get.assert_awaited_once_with("/api/v1/agents/agent-example_01") @@ -99,6 +170,17 @@ async def test_get_agent_policy_normalizes_agent_name() -> None: client.http_client.get.assert_awaited_once_with("/api/v1/agents/agent-example_01/policy") +@pytest.mark.asyncio +async def test_get_agent_policies_normalizes_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.get = AsyncMock(return_value=DummyResponse()) + + await agents.get_agent_policies(client, "Agent-Example_01") + + client.http_client.get.assert_awaited_once_with("/api/v1/agents/agent-example_01/policies") + + @pytest.mark.asyncio async def test_remove_agent_policy_normalizes_agent_name() -> None: client = MagicMock() @@ -108,3 +190,75 @@ async def test_remove_agent_policy_normalizes_agent_name() -> None: await agents.remove_agent_policy(client, "Agent-Example_01") client.http_client.delete.assert_awaited_once_with("/api/v1/agents/agent-example_01/policy") + + +@pytest.mark.asyncio +async def test_remove_all_agent_policies_normalizes_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.delete = AsyncMock(return_value=DummyResponse()) + + await agents.remove_all_agent_policies(client, "Agent-Example_01") + + client.http_client.delete.assert_awaited_once_with( + "/api/v1/agents/agent-example_01/policies" + ) + + +@pytest.mark.asyncio +async def test_add_agent_policy_normalizes_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.post = AsyncMock(return_value=DummyResponse()) + + await agents.add_agent_policy(client, "Agent-Example_01", policy_id=3) + + client.http_client.post.assert_awaited_once_with("/api/v1/agents/agent-example_01/policies/3") + + +@pytest.mark.asyncio +async def test_remove_agent_policy_association_normalizes_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.delete = AsyncMock(return_value=DummyResponse()) + + await agents.remove_agent_policy_association(client, "Agent-Example_01", policy_id=3) + + client.http_client.delete.assert_awaited_once_with( + "/api/v1/agents/agent-example_01/policies/3" + ) + + +@pytest.mark.asyncio +async def test_add_agent_control_normalizes_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.post = AsyncMock(return_value=DummyResponse()) + + await agents.add_agent_control(client, "Agent-Example_01", control_id=9) + + client.http_client.post.assert_awaited_once_with("/api/v1/agents/agent-example_01/controls/9") + + +@pytest.mark.asyncio +async def test_remove_agent_control_normalizes_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.delete = AsyncMock(return_value=DummyResponse()) + + await agents.remove_agent_control(client, "Agent-Example_01", control_id=9) + + client.http_client.delete.assert_awaited_once_with( + "/api/v1/agents/agent-example_01/controls/9" + ) + + +@pytest.mark.asyncio +async def test_assign_policy_normalizes_agent_name() -> None: + client = MagicMock() + client.http_client = MagicMock() + client.http_client.post = AsyncMock(return_value=DummyResponse()) + + await policies.assign_policy_to_agent(client, "Agent-Example_01", policy_id=11) + + client.http_client.post.assert_awaited_once_with("/api/v1/agents/agent-example_01/policy/11") diff --git a/sdks/python/tests/test_agents_api.py b/sdks/python/tests/test_agents_api.py index 271271e0..01e59b36 100644 --- a/sdks/python/tests/test_agents_api.py +++ b/sdks/python/tests/test_agents_api.py @@ -20,13 +20,13 @@ async def test_list_agent_controls_typed_calls_controls_endpoint() -> None: response.raise_for_status = Mock() response.json = Mock(return_value=response_payload) client = SimpleNamespace(http_client=SimpleNamespace(get=AsyncMock(return_value=response))) - agent_id = str(uuid4()) + agent_name = str(uuid4()) # WHEN: typed controls are requested. - result = await agent_control.agents.list_agent_controls_typed(client, agent_id) + result = await agent_control.agents.list_agent_controls_typed(client, agent_name) # THEN: wrapper calls the expected endpoint and returns a typed result. - client.http_client.get.assert_awaited_once_with(f"/api/v1/agents/{agent_id}/controls") + client.http_client.get.assert_awaited_once_with(f"/api/v1/agents/{agent_name}/controls") assert isinstance(result.controls, list) diff --git a/sdks/python/tests/test_integration_agents.py b/sdks/python/tests/test_integration_agents.py index 144081d5..66a22c43 100644 --- a/sdks/python/tests/test_integration_agents.py +++ b/sdks/python/tests/test_integration_agents.py @@ -9,16 +9,15 @@ import uuid -import pytest - import agent_control +import pytest from agent_control_models.server import AgentControlsResponse @pytest.mark.asyncio async def test_agent_registration_workflow( client: agent_control.AgentControlClient, - test_agent_id: str, + test_agent_name: str, sample_steps: list ) -> None: """ @@ -197,9 +196,107 @@ async def test_convenience_get_agent_function( print("✓ Convenience function works") +@pytest.mark.asyncio +async def test_convenience_agent_association_functions( + test_agent: dict, + test_policy: dict, + test_control: dict, + server_url: str, + api_key: str | None, +) -> None: + """Top-level convenience helpers support multi-policy and direct controls.""" + agent_name = test_agent["agent_name"] + policy_id = test_policy["policy_id"] + control_id = test_control["control_id"] + + control_data = { + "description": "Integration test control", + "enabled": True, + "execution": "server", + "scope": {"step_types": ["tool"], "stages": ["pre"]}, + "selector": {"path": "input"}, + "evaluator": { + "name": "regex", + "config": {"pattern": ".*"}, + }, + "action": {"decision": "allow"}, + "tags": ["test"], + } + + add_policy_result = await agent_control.add_agent_policy( + agent_name, + policy_id, + server_url=server_url, + api_key=api_key, + ) + assert add_policy_result["success"] is True + + policies_result = await agent_control.get_agent_policies( + agent_name, + server_url=server_url, + api_key=api_key, + ) + assert policy_id in policies_result["policy_ids"] + + add_control_result = await agent_control.add_agent_control( + agent_name, + control_id, + server_url=server_url, + api_key=api_key, + ) + assert add_control_result["success"] is True + + async with agent_control.AgentControlClient( + base_url=server_url, + api_key=api_key, + ) as client: + await agent_control.controls.set_control_data(client, control_id, control_data) + controls_result = await agent_control.agents.list_agent_controls(client, agent_name) + control_ids = {item["id"] for item in controls_result["controls"]} + assert control_id in control_ids + + remove_control_result = await agent_control.remove_agent_control( + agent_name, + control_id, + server_url=server_url, + api_key=api_key, + ) + assert remove_control_result["success"] is True + assert remove_control_result["removed_direct_association"] is True + + remove_policy_result = await agent_control.remove_agent_policy_association( + agent_name, + policy_id, + server_url=server_url, + api_key=api_key, + ) + assert remove_policy_result["success"] is True + + # Re-associate then clear all to verify remove_all convenience path. + await agent_control.add_agent_policy( + agent_name, + policy_id, + server_url=server_url, + api_key=api_key, + ) + clear_result = await agent_control.remove_all_agent_policies( + agent_name, + server_url=server_url, + api_key=api_key, + ) + assert clear_result["success"] is True + + policies_after_clear = await agent_control.get_agent_policies( + agent_name, + server_url=server_url, + api_key=api_key, + ) + assert policies_after_clear["policy_ids"] == [] + + @pytest.mark.asyncio async def test_init_function_workflow( - test_agent_id: str, + test_agent_name: str, server_url: str, api_key: str | None, sample_steps: list, @@ -214,7 +311,7 @@ async def test_init_function_workflow( """ # Initialize agent agent = agent_control.init( - agent_name=test_agent_id, + agent_name=test_agent_name, agent_description="Testing init function", agent_version="1.0.0", server_url=server_url, @@ -226,7 +323,7 @@ async def test_init_function_workflow( # Verify agent instance assert agent is not None - assert agent.agent_name == test_agent_id + assert agent.agent_name == test_agent_name assert hasattr(agent, "agent_name") # Verify current_agent() diff --git a/sdks/typescript/Makefile b/sdks/typescript/Makefile index 351e37b1..70b6708c 100644 --- a/sdks/typescript/Makefile +++ b/sdks/typescript/Makefile @@ -34,7 +34,12 @@ overlay-test: cd ../.. && uv run --package agent-control-server pytest sdks/typescript/tests/test_generate_method_names_overlay.py name-check: - @if rg -n --no-heading 'async\\s+[A-Za-z0-9]+(ApiV[0-9]+|Get|Post|Put|Patch|Delete)\\b' src/generated/sdk; then \ + @if command -v rg >/dev/null 2>&1; then \ + MATCH_CMD="rg -n --no-heading"; \ + else \ + MATCH_CMD="grep -REn"; \ + fi; \ + if eval "$$MATCH_CMD 'async\\s+[A-Za-z0-9]+(ApiV[0-9]+|Get|Post|Put|Patch|Delete)\\b' src/generated/sdk"; then \ echo "Found verbose generated method names. Update generate-method-names-overlay.py rules."; \ exit 1; \ fi diff --git a/sdks/typescript/overlays/method-names.overlay.yaml b/sdks/typescript/overlays/method-names.overlay.yaml index fd4927f7..3093d3f8 100644 --- a/sdks/typescript/overlays/method-names.overlay.yaml +++ b/sdks/typescript/overlays/method-names.overlay.yaml @@ -30,6 +30,16 @@ actions: x-speakeasy-group: agents x-speakeasy-name-override: listControls + - target: $["paths"]["/api/v1/agents/{agent_name}/controls/{control_id}"]["post"] + update: + x-speakeasy-group: agents + x-speakeasy-name-override: addControl + + - target: $["paths"]["/api/v1/agents/{agent_name}/controls/{control_id}"]["delete"] + update: + x-speakeasy-group: agents + x-speakeasy-name-override: removeControl + - target: $["paths"]["/api/v1/agents/{agent_name}/evaluators"]["get"] update: x-speakeasy-group: agents @@ -40,6 +50,26 @@ actions: x-speakeasy-group: agents x-speakeasy-name-override: getEvaluator + - target: $["paths"]["/api/v1/agents/{agent_name}/policies"]["get"] + update: + x-speakeasy-group: agents + x-speakeasy-name-override: getPolicies + + - target: $["paths"]["/api/v1/agents/{agent_name}/policies"]["delete"] + update: + x-speakeasy-group: agents + x-speakeasy-name-override: removeAllAgentPolicies + + - target: $["paths"]["/api/v1/agents/{agent_name}/policies/{policy_id}"]["post"] + update: + x-speakeasy-group: agents + x-speakeasy-name-override: addPolicy + + - target: $["paths"]["/api/v1/agents/{agent_name}/policies/{policy_id}"]["delete"] + update: + x-speakeasy-group: agents + x-speakeasy-name-override: removePolicy + - target: $["paths"]["/api/v1/agents/{agent_name}/policy"]["get"] update: x-speakeasy-group: agents diff --git a/sdks/typescript/src/client.ts b/sdks/typescript/src/client.ts index f9cd94d5..2836f23f 100644 --- a/sdks/typescript/src/client.ts +++ b/sdks/typescript/src/client.ts @@ -9,6 +9,7 @@ export type APIKeyProvider = string | (() => Promise); export interface AgentControlInitOptions { agentName: string; + agentId?: string; serverUrl: string; apiKey?: APIKeyProvider; steps?: StepSchema[]; diff --git a/sdks/typescript/src/generated/funcs/agents-add-control.ts b/sdks/typescript/src/generated/funcs/agents-add-control.ts new file mode 100644 index 00000000..08a7553f --- /dev/null +++ b/sdks/typescript/src/generated/funcs/agents-add-control.ts @@ -0,0 +1,191 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Associate control directly with agent + * + * @remarks + * Associate a control directly with an agent (idempotent). + */ +export function agentsAddControl( + client: AgentControlSDKCore, + request: + operations.AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: + operations.AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = null; + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + control_id: encodeSimple("control_id", payload.control_id, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/agents/{agent_name}/controls/{control_id}")( + pathParams, + ); + + const headers = new Headers(compactMap({ + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: + "add_agent_control_api_v1_agents__agent_name__controls__control_id__post", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "POST", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.AssocResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/agents-add-policy.ts b/sdks/typescript/src/generated/funcs/agents-add-policy.ts new file mode 100644 index 00000000..76e7dfb5 --- /dev/null +++ b/sdks/typescript/src/generated/funcs/agents-add-policy.ts @@ -0,0 +1,191 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Associate policy with agent + * + * @remarks + * Associate a policy with an agent (idempotent). + */ +export function agentsAddPolicy( + client: AgentControlSDKCore, + request: + operations.AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: + operations.AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = null; + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + policy_id: encodeSimple("policy_id", payload.policy_id, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/agents/{agent_name}/policies/{policy_id}")( + pathParams, + ); + + const headers = new Headers(compactMap({ + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: + "add_agent_policy_api_v1_agents__agent_name__policies__policy_id__post", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "POST", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.AssocResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/agents-delete-policy.ts b/sdks/typescript/src/generated/funcs/agents-delete-policy.ts index 2e8a18b6..a87b1627 100644 --- a/sdks/typescript/src/generated/funcs/agents-delete-policy.ts +++ b/sdks/typescript/src/generated/funcs/agents-delete-policy.ts @@ -28,23 +28,10 @@ import { APICall, APIPromise } from "../types/async.js"; import { Result } from "../types/fp.js"; /** - * Remove agent's policy assignment + * Remove agent's policy assignment (compatibility) * * @remarks - * Remove the policy assignment from an agent. - * - * The agent will no longer have any protection controls active. - * - * Args: - * agent_name: Agent identifier - * db: Database session (injected) - * - * Returns: - * DeletePolicyResponse with success flag - * - * Raises: - * HTTPException 404: Agent not found or agent has no policy assigned - * HTTPException 500: Database error during removal + * Compatibility endpoint that removes all policy associations. */ export function agentsDeletePolicy( client: AgentControlSDKCore, diff --git a/sdks/typescript/src/generated/funcs/agents-get-policies.ts b/sdks/typescript/src/generated/funcs/agents-get-policies.ts new file mode 100644 index 00000000..6c1c9168 --- /dev/null +++ b/sdks/typescript/src/generated/funcs/agents-get-policies.ts @@ -0,0 +1,182 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * List policies associated with agent + * + * @remarks + * List policy IDs associated with an agent. + */ +export function agentsGetPolicies( + client: AgentControlSDKCore, + request: operations.GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.GetAgentPoliciesResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: operations.GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.GetAgentPoliciesResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = null; + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/agents/{agent_name}/policies")(pathParams); + + const headers = new Headers(compactMap({ + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: "get_agent_policies_api_v1_agents__agent_name__policies_get", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "GET", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.GetAgentPoliciesResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.GetAgentPoliciesResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/agents-get-policy.ts b/sdks/typescript/src/generated/funcs/agents-get-policy.ts index cbe953f3..df3755d8 100644 --- a/sdks/typescript/src/generated/funcs/agents-get-policy.ts +++ b/sdks/typescript/src/generated/funcs/agents-get-policy.ts @@ -28,20 +28,10 @@ import { APICall, APIPromise } from "../types/async.js"; import { Result } from "../types/fp.js"; /** - * Get agent's assigned policy + * Get agent's assigned policy (compatibility) * * @remarks - * Retrieve the policy currently assigned to an agent. - * - * Args: - * agent_name: Agent identifier - * db: Database session (injected) - * - * Returns: - * GetPolicyResponse with policy ID - * - * Raises: - * HTTPException 404: Agent not found or agent has no policy assigned + * Compatibility endpoint that returns the first associated policy. */ export function agentsGetPolicy( client: AgentControlSDKCore, diff --git a/sdks/typescript/src/generated/funcs/agents-init.ts b/sdks/typescript/src/generated/funcs/agents-init.ts index 1cb02298..508b00ed 100644 --- a/sdks/typescript/src/generated/funcs/agents-init.ts +++ b/sdks/typescript/src/generated/funcs/agents-init.ts @@ -45,7 +45,7 @@ import { Result } from "../types/fp.js"; * db: Database session (injected) * * Returns: - * InitAgentResponse with created flag and active controls (if policy assigned) + * InitAgentResponse with created flag and active controls (policy-derived + direct) */ export function agentsInit( client: AgentControlSDKCore, diff --git a/sdks/typescript/src/generated/funcs/agents-list-controls.ts b/sdks/typescript/src/generated/funcs/agents-list-controls.ts index b44c3402..d0776636 100644 --- a/sdks/typescript/src/generated/funcs/agents-list-controls.ts +++ b/sdks/typescript/src/generated/funcs/agents-list-controls.ts @@ -33,15 +33,14 @@ import { Result } from "../types/fp.js"; * @remarks * List all protection controls active for an agent. * - * Controls are inherited from the agent's assigned policy. - * Returns an empty list if the agent has no policy. + * Controls include the union of policy-derived and directly associated controls. * * Args: * agent_name: Agent identifier * db: Database session (injected) * * Returns: - * AgentControlsResponse with list of controls (empty if no policy) + * AgentControlsResponse with list of active controls * * Raises: * HTTPException 404: Agent not found diff --git a/sdks/typescript/src/generated/funcs/agents-list.ts b/sdks/typescript/src/generated/funcs/agents-list.ts index 82e21470..5dcbb7c6 100644 --- a/sdks/typescript/src/generated/funcs/agents-list.ts +++ b/sdks/typescript/src/generated/funcs/agents-list.ts @@ -33,7 +33,7 @@ import { Result } from "../types/fp.js"; * @remarks * List all registered agents with cursor-based pagination. * - * Returns a summary of each agent including identifier, policy assignment, + * Returns a summary of each agent including identifier, policy associations, * and counts of registered steps and evaluators. * * Args: diff --git a/sdks/typescript/src/generated/funcs/agents-remove-all-agent-policies.ts b/sdks/typescript/src/generated/funcs/agents-remove-all-agent-policies.ts new file mode 100644 index 00000000..01234e18 --- /dev/null +++ b/sdks/typescript/src/generated/funcs/agents-remove-all-agent-policies.ts @@ -0,0 +1,185 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Remove all policy associations from agent + * + * @remarks + * Remove all policy associations from an agent. + */ +export function agentsRemoveAllAgentPolicies( + client: AgentControlSDKCore, + request: + operations.RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: + operations.RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = null; + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/agents/{agent_name}/policies")(pathParams); + + const headers = new Headers(compactMap({ + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: + "remove_all_agent_policies_api_v1_agents__agent_name__policies_delete", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "DELETE", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.AssocResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/agents-remove-control.ts b/sdks/typescript/src/generated/funcs/agents-remove-control.ts new file mode 100644 index 00000000..c23e3bec --- /dev/null +++ b/sdks/typescript/src/generated/funcs/agents-remove-control.ts @@ -0,0 +1,191 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Remove direct control association from agent + * + * @remarks + * Remove a direct control association from an agent (idempotent). + */ +export function agentsRemoveControl( + client: AgentControlSDKCore, + request: + operations.RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.RemoveAgentControlResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: + operations.RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.RemoveAgentControlResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = null; + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + control_id: encodeSimple("control_id", payload.control_id, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/agents/{agent_name}/controls/{control_id}")( + pathParams, + ); + + const headers = new Headers(compactMap({ + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: + "remove_agent_control_api_v1_agents__agent_name__controls__control_id__delete", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "DELETE", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.RemoveAgentControlResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.RemoveAgentControlResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/agents-remove-policy.ts b/sdks/typescript/src/generated/funcs/agents-remove-policy.ts new file mode 100644 index 00000000..b275052a --- /dev/null +++ b/sdks/typescript/src/generated/funcs/agents-remove-policy.ts @@ -0,0 +1,194 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { AgentControlSDKCore } from "../core.js"; +import { encodeSimple } from "../lib/encodings.js"; +import * as M from "../lib/matchers.js"; +import { compactMap } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { RequestOptions } from "../lib/sdks.js"; +import { extractSecurity, resolveGlobalSecurity } from "../lib/security.js"; +import { pathToFunc } from "../lib/url.js"; +import { AgentControlSDKError } from "../models/errors/agent-control-sdk-error.js"; +import { + ConnectionError, + InvalidRequestError, + RequestAbortedError, + RequestTimeoutError, + UnexpectedClientError, +} from "../models/errors/http-client-errors.js"; +import * as errors from "../models/errors/index.js"; +import { ResponseValidationError } from "../models/errors/response-validation-error.js"; +import { SDKValidationError } from "../models/errors/sdk-validation-error.js"; +import * as models from "../models/index.js"; +import * as operations from "../models/operations/index.js"; +import { APICall, APIPromise } from "../types/async.js"; +import { Result } from "../types/fp.js"; + +/** + * Remove policy association from agent + * + * @remarks + * Remove a policy association from an agent. + * + * Idempotent for existing resources: removing a non-associated link is a no-op. + * Missing agent/policy resources still return 404. + */ +export function agentsRemovePolicy( + client: AgentControlSDKCore, + request: + operations.RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest, + options?: RequestOptions, +): APIPromise< + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + > +> { + return new APIPromise($do( + client, + request, + options, + )); +} + +async function $do( + client: AgentControlSDKCore, + request: + operations.RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest, + options?: RequestOptions, +): Promise< + [ + Result< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >, + APICall, + ] +> { + const parsed = safeParse( + request, + (value) => + z.parse( + operations + .RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest$outboundSchema, + value, + ), + "Input validation failed", + ); + if (!parsed.ok) { + return [parsed, { status: "invalid" }]; + } + const payload = parsed.value; + const body = null; + + const pathParams = { + agent_name: encodeSimple("agent_name", payload.agent_name, { + explode: false, + charEncoding: "percent", + }), + policy_id: encodeSimple("policy_id", payload.policy_id, { + explode: false, + charEncoding: "percent", + }), + }; + + const path = pathToFunc("/api/v1/agents/{agent_name}/policies/{policy_id}")( + pathParams, + ); + + const headers = new Headers(compactMap({ + Accept: "application/json", + })); + + const secConfig = await extractSecurity(client._options.apiKeyHeader); + const securityInput = secConfig == null ? {} : { apiKeyHeader: secConfig }; + const requestSecurity = resolveGlobalSecurity(securityInput); + + const context = { + options: client._options, + baseURL: options?.serverURL ?? client._baseURL ?? "", + operationID: + "remove_agent_policy_api_v1_agents__agent_name__policies__policy_id__delete", + oAuth2Scopes: null, + + resolvedSecurity: requestSecurity, + + securitySource: client._options.apiKeyHeader, + retryConfig: options?.retries + || client._options.retryConfig + || { strategy: "none" }, + retryCodes: options?.retryCodes || ["429", "500", "502", "503", "504"], + }; + + const requestRes = client._createRequest(context, { + security: requestSecurity, + method: "DELETE", + baseURL: options?.serverURL, + path: path, + headers: headers, + body: body, + userAgent: client._options.userAgent, + timeoutMs: options?.timeoutMs || client._options.timeoutMs || -1, + }, options); + if (!requestRes.ok) { + return [requestRes, { status: "invalid" }]; + } + const req = requestRes.value; + + const doResult = await client._do(req, { + context, + errorCodes: ["422", "4XX", "5XX"], + retryConfig: context.retryConfig, + retryCodes: context.retryCodes, + }); + if (!doResult.ok) { + return [doResult, { status: "request-error", request: req }]; + } + const response = doResult.value; + + const responseFields = { + HttpMeta: { Response: response, Request: req }, + }; + + const [result] = await M.match< + models.AssocResponse, + | errors.HTTPValidationError + | AgentControlSDKError + | ResponseValidationError + | ConnectionError + | RequestAbortedError + | RequestTimeoutError + | InvalidRequestError + | UnexpectedClientError + | SDKValidationError + >( + M.json(200, models.AssocResponse$inboundSchema), + M.jsonErr(422, errors.HTTPValidationError$inboundSchema), + M.fail("4XX"), + M.fail("5XX"), + )(response, req, { extraFields: responseFields }); + if (!result.ok) { + return [result, { status: "complete", request: req, response }]; + } + + return [result, { status: "complete", request: req, response }]; +} diff --git a/sdks/typescript/src/generated/funcs/agents-update-policy.ts b/sdks/typescript/src/generated/funcs/agents-update-policy.ts index 7626a798..f120912b 100644 --- a/sdks/typescript/src/generated/funcs/agents-update-policy.ts +++ b/sdks/typescript/src/generated/funcs/agents-update-policy.ts @@ -28,24 +28,10 @@ import { APICall, APIPromise } from "../types/async.js"; import { Result } from "../types/fp.js"; /** - * Assign policy to agent + * Assign policy to agent (compatibility) * * @remarks - * Assign a policy to an agent, replacing any existing policy assignment. - * - * The agent will immediately inherit all controls from the assigned policy. - * - * Args: - * agent_name: Agent identifier - * policy_id: ID of the policy to assign - * db: Database session (injected) - * - * Returns: - * SetPolicyResponse with success flag and previous policy ID (if any) - * - * Raises: - * HTTPException 404: Agent or policy not found - * HTTPException 500: Database error during assignment + * Compatibility endpoint that replaces all policy associations with one policy. */ export function agentsUpdatePolicy( client: AgentControlSDKCore, diff --git a/sdks/typescript/src/generated/funcs/controls-delete.ts b/sdks/typescript/src/generated/funcs/controls-delete.ts index c3b65d39..ea973a24 100644 --- a/sdks/typescript/src/generated/funcs/controls-delete.ts +++ b/sdks/typescript/src/generated/funcs/controls-delete.ts @@ -33,7 +33,7 @@ import { Result } from "../types/fp.js"; * @remarks * Delete a control by ID. * - * By default, deletion fails if the control is associated with any policy. + * By default, deletion fails if the control is associated with any policy or agent. * Use force=true to automatically dissociate and delete. * * Args: @@ -42,7 +42,7 @@ import { Result } from "../types/fp.js"; * db: Database session (injected) * * Returns: - * DeleteControlResponse with success flag and list of dissociated policies + * DeleteControlResponse with success flag and dissociation details * * Raises: * HTTPException 404: Control not found diff --git a/sdks/typescript/src/generated/hooks/hooks.ts b/sdks/typescript/src/generated/hooks/hooks.ts index 7ed9e044..94ad0a32 100644 --- a/sdks/typescript/src/generated/hooks/hooks.ts +++ b/sdks/typescript/src/generated/hooks/hooks.ts @@ -18,8 +18,6 @@ import { SDKInitHook, } from "./types.js"; -import { initHooks } from "./registration.js"; - export class SDKHooks implements Hooks { sdkInitHooks: SDKInitHook[] = []; beforeCreateRequestHooks: BeforeCreateRequestHook[] = []; @@ -47,7 +45,6 @@ export class SDKHooks implements Hooks { this.registerAfterErrorHook(hook); } } - initHooks(this); } registerSDKInitHook(hook: SDKInitHook) { diff --git a/sdks/typescript/src/generated/hooks/registration.ts b/sdks/typescript/src/generated/hooks/registration.ts deleted file mode 100644 index 3a8fdeec..00000000 --- a/sdks/typescript/src/generated/hooks/registration.ts +++ /dev/null @@ -1,14 +0,0 @@ -import { Hooks } from "./types.js"; - -/* - * This file is only ever generated once on the first generation and then is free to be modified. - * Any hooks you wish to add should be registered in the initHooks function. Feel free to define them - * in this file or in separate files in the hooks folder. - */ - -// @ts-ignore remove this line when you add your first hook and hooks is used -export function initHooks(hooks: Hooks) { - // Add hooks by calling hooks.register{ClientInit/BeforeCreateRequest/BeforeRequest/AfterSuccess/AfterError}Hook - // with an instance of a hook that implements that specific Hook interface - // Hooks are registered per SDK instance, and are valid for the lifetime of the SDK instance -} diff --git a/sdks/typescript/src/generated/models/agent-controls-response.ts b/sdks/typescript/src/generated/models/agent-controls-response.ts index 0668702c..22a591ee 100644 --- a/sdks/typescript/src/generated/models/agent-controls-response.ts +++ b/sdks/typescript/src/generated/models/agent-controls-response.ts @@ -10,7 +10,7 @@ import { SDKValidationError } from "./errors/sdk-validation-error.js"; export type AgentControlsResponse = { /** - * List of controls associated with the agent via its policy + * List of active controls associated with the agent */ controls: Array; }; diff --git a/sdks/typescript/src/generated/models/agent-ref.ts b/sdks/typescript/src/generated/models/agent-ref.ts index e24952df..676ac11a 100644 --- a/sdks/typescript/src/generated/models/agent-ref.ts +++ b/sdks/typescript/src/generated/models/agent-ref.ts @@ -14,7 +14,7 @@ import { SDKValidationError } from "./errors/sdk-validation-error.js"; */ export type AgentRef = { /** - * Agent identifier + * Agent name */ agentName: string; }; diff --git a/sdks/typescript/src/generated/models/agent-summary.ts b/sdks/typescript/src/generated/models/agent-summary.ts index a2cccd70..8f9732c4 100644 --- a/sdks/typescript/src/generated/models/agent-summary.ts +++ b/sdks/typescript/src/generated/models/agent-summary.ts @@ -14,7 +14,7 @@ import { SDKValidationError } from "./errors/sdk-validation-error.js"; */ export type AgentSummary = { /** - * Number of active controls from agent's policy + * Number of active controls for this agent */ activeControlsCount: number; /** @@ -30,9 +30,9 @@ export type AgentSummary = { */ evaluatorCount: number; /** - * ID of assigned policy, if any + * IDs of policies associated with the agent */ - policyId?: number | null | undefined; + policyIds?: Array | undefined; /** * Number of steps registered with the agent */ @@ -47,7 +47,7 @@ export const AgentSummary$inboundSchema: z.ZodMiniType = agent_name: types.string(), created_at: z.optional(z.nullable(types.string())), evaluator_count: z._default(types.number(), 0), - policy_id: z.optional(z.nullable(types.number())), + policy_ids: types.optional(z.array(types.number())), step_count: z._default(types.number(), 0), }), z.transform((v) => { @@ -56,7 +56,7 @@ export const AgentSummary$inboundSchema: z.ZodMiniType = "agent_name": "agentName", "created_at": "createdAt", "evaluator_count": "evaluatorCount", - "policy_id": "policyId", + "policy_ids": "policyIds", "step_count": "stepCount", }); }), diff --git a/sdks/typescript/src/generated/models/control-summary.ts b/sdks/typescript/src/generated/models/control-summary.ts index 9e0356a5..7e89a07f 100644 --- a/sdks/typescript/src/generated/models/control-summary.ts +++ b/sdks/typescript/src/generated/models/control-summary.ts @@ -50,6 +50,10 @@ export type ControlSummary = { * Agent using this control */ usedByAgent?: AgentRef | null | undefined; + /** + * Number of unique agents using this control + */ + usedByAgentsCount: number; }; /** @internal */ @@ -67,11 +71,13 @@ export const ControlSummary$inboundSchema: z.ZodMiniType< step_types: z.optional(z.nullable(z.array(types.string()))), tags: types.optional(z.array(types.string())), used_by_agent: z.optional(z.nullable(AgentRef$inboundSchema)), + used_by_agents_count: z._default(types.number(), 0), }), z.transform((v) => { return remap$(v, { "step_types": "stepTypes", "used_by_agent": "usedByAgent", + "used_by_agents_count": "usedByAgentsCount", }); }), ); diff --git a/sdks/typescript/src/generated/models/delete-control-response.ts b/sdks/typescript/src/generated/models/delete-control-response.ts index 623abd95..40f6c252 100644 --- a/sdks/typescript/src/generated/models/delete-control-response.ts +++ b/sdks/typescript/src/generated/models/delete-control-response.ts @@ -14,9 +14,17 @@ import { SDKValidationError } from "./errors/sdk-validation-error.js"; */ export type DeleteControlResponse = { /** - * Policy IDs the control was removed from before deletion + * Deprecated: policy IDs the control was removed from before deletion */ dissociatedFrom?: Array | undefined; + /** + * Agent names the control was removed from before deletion + */ + dissociatedFromAgents?: Array | undefined; + /** + * Policy IDs the control was removed from before deletion + */ + dissociatedFromPolicies?: Array | undefined; /** * Whether the control was deleted */ @@ -30,11 +38,15 @@ export const DeleteControlResponse$inboundSchema: z.ZodMiniType< > = z.pipe( z.object({ dissociated_from: types.optional(z.array(types.number())), + dissociated_from_agents: types.optional(z.array(types.string())), + dissociated_from_policies: types.optional(z.array(types.number())), success: types.boolean(), }), z.transform((v) => { return remap$(v, { "dissociated_from": "dissociatedFrom", + "dissociated_from_agents": "dissociatedFromAgents", + "dissociated_from_policies": "dissociatedFromPolicies", }); }), ); diff --git a/sdks/typescript/src/generated/models/delete-policy-response.ts b/sdks/typescript/src/generated/models/delete-policy-response.ts index d37ef18e..760045ca 100644 --- a/sdks/typescript/src/generated/models/delete-policy-response.ts +++ b/sdks/typescript/src/generated/models/delete-policy-response.ts @@ -8,9 +8,12 @@ import { Result as SafeParseResult } from "../types/fp.js"; import * as types from "../types/primitives.js"; import { SDKValidationError } from "./errors/sdk-validation-error.js"; +/** + * Compatibility response for singular policy deletion endpoint. + */ export type DeletePolicyResponse = { /** - * Whether the policy was successfully removed + * Whether the request succeeded */ success: boolean; }; diff --git a/sdks/typescript/src/generated/models/get-agent-policies-response.ts b/sdks/typescript/src/generated/models/get-agent-policies-response.ts new file mode 100644 index 00000000..3e9413e0 --- /dev/null +++ b/sdks/typescript/src/generated/models/get-agent-policies-response.ts @@ -0,0 +1,42 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { Result as SafeParseResult } from "../types/fp.js"; +import * as types from "../types/primitives.js"; +import { SDKValidationError } from "./errors/sdk-validation-error.js"; + +export type GetAgentPoliciesResponse = { + /** + * IDs of policies associated with the agent + */ + policyIds?: Array | undefined; +}; + +/** @internal */ +export const GetAgentPoliciesResponse$inboundSchema: z.ZodMiniType< + GetAgentPoliciesResponse, + unknown +> = z.pipe( + z.object({ + policy_ids: types.optional(z.array(types.number())), + }), + z.transform((v) => { + return remap$(v, { + "policy_ids": "policyIds", + }); + }), +); + +export function getAgentPoliciesResponseFromJSON( + jsonString: string, +): SafeParseResult { + return safeParse( + jsonString, + (x) => GetAgentPoliciesResponse$inboundSchema.parse(JSON.parse(x)), + `Failed to parse 'GetAgentPoliciesResponse' from JSON`, + ); +} diff --git a/sdks/typescript/src/generated/models/get-policy-response.ts b/sdks/typescript/src/generated/models/get-policy-response.ts index fd8437df..83e05a8a 100644 --- a/sdks/typescript/src/generated/models/get-policy-response.ts +++ b/sdks/typescript/src/generated/models/get-policy-response.ts @@ -9,9 +9,12 @@ import { Result as SafeParseResult } from "../types/fp.js"; import * as types from "../types/primitives.js"; import { SDKValidationError } from "./errors/sdk-validation-error.js"; +/** + * Compatibility response for singular policy retrieval endpoint. + */ export type GetPolicyResponse = { /** - * Identifier of the policy assigned to the agent + * Associated policy ID */ policyId: number; }; diff --git a/sdks/typescript/src/generated/models/index.ts b/sdks/typescript/src/generated/models/index.ts index af193ed3..29a27023 100644 --- a/sdks/typescript/src/generated/models/index.ts +++ b/sdks/typescript/src/generated/models/index.ts @@ -39,6 +39,7 @@ export * from "./evaluator-schema.js"; export * from "./evaluator-spec.js"; export * from "./event-query-request.js"; export * from "./event-query-response.js"; +export * from "./get-agent-policies-response.js"; export * from "./get-agent-response.js"; export * from "./get-control-data-response.js"; export * from "./get-control-response.js"; @@ -58,6 +59,7 @@ export * from "./patch-agent-request.js"; export * from "./patch-agent-response.js"; export * from "./patch-control-request.js"; export * from "./patch-control-response.js"; +export * from "./remove-agent-control-response.js"; export * from "./security.js"; export * from "./set-control-data-request.js"; export * from "./set-control-data-response.js"; diff --git a/sdks/typescript/src/generated/models/init-agent-response.ts b/sdks/typescript/src/generated/models/init-agent-response.ts index c2e00b02..76ade4de 100644 --- a/sdks/typescript/src/generated/models/init-agent-response.ts +++ b/sdks/typescript/src/generated/models/init-agent-response.ts @@ -19,7 +19,7 @@ import { */ export type InitAgentResponse = { /** - * Active protection controls for the agent (if policy assigned) + * Active protection controls for the agent */ controls?: Array | undefined; /** diff --git a/sdks/typescript/src/generated/models/operations/add-agent-control-api-v1-agents-agent-name-controls-control-id-post.ts b/sdks/typescript/src/generated/models/operations/add-agent-control-api-v1-agents-agent-name-controls-control-id-post.ts new file mode 100644 index 00000000..dec32e6e --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/add-agent-control-api-v1-agents-agent-name-controls-control-id-post.ts @@ -0,0 +1,46 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; + +export type AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest = { + agentName: string; + controlId: number; +}; + +/** @internal */ +export type AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest$Outbound = + { + agent_name: string; + control_id: number; + }; + +/** @internal */ +export const AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest$outboundSchema: + z.ZodMiniType< + AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest$Outbound, + AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest + > = z.pipe( + z.object({ + agentName: z.string(), + controlId: z.int(), + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + controlId: "control_id", + }); + }), + ); + +export function addAgentControlApiV1AgentsAgentNameControlsControlIdPostRequestToJSON( + addAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest: + AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest, +): string { + return JSON.stringify( + AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest$outboundSchema + .parse(addAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest), + ); +} diff --git a/sdks/typescript/src/generated/models/operations/add-agent-policy-api-v1-agents-agent-name-policies-policy-id-post.ts b/sdks/typescript/src/generated/models/operations/add-agent-policy-api-v1-agents-agent-name-policies-policy-id-post.ts new file mode 100644 index 00000000..919d1948 --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/add-agent-policy-api-v1-agents-agent-name-policies-policy-id-post.ts @@ -0,0 +1,46 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; + +export type AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest = { + agentName: string; + policyId: number; +}; + +/** @internal */ +export type AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest$Outbound = + { + agent_name: string; + policy_id: number; + }; + +/** @internal */ +export const AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest$outboundSchema: + z.ZodMiniType< + AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest$Outbound, + AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest + > = z.pipe( + z.object({ + agentName: z.string(), + policyId: z.int(), + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + policyId: "policy_id", + }); + }), + ); + +export function addAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequestToJSON( + addAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest: + AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest, +): string { + return JSON.stringify( + AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest$outboundSchema + .parse(addAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest), + ); +} diff --git a/sdks/typescript/src/generated/models/operations/delete-control-api-v1-controls-control-id-delete.ts b/sdks/typescript/src/generated/models/operations/delete-control-api-v1-controls-control-id-delete.ts index 6d3935ba..b016ec01 100644 --- a/sdks/typescript/src/generated/models/operations/delete-control-api-v1-controls-control-id-delete.ts +++ b/sdks/typescript/src/generated/models/operations/delete-control-api-v1-controls-control-id-delete.ts @@ -8,7 +8,7 @@ import { remap as remap$ } from "../../lib/primitives.js"; export type DeleteControlApiV1ControlsControlIdDeleteRequest = { controlId: number; /** - * If true, dissociate from all policies before deleting. If false, fail if control is associated with any policy. + * If true, dissociate from all policy/agent links before deleting. If false, fail if control is associated with any policy or agent. */ force?: boolean | undefined; }; diff --git a/sdks/typescript/src/generated/models/operations/get-agent-policies-api-v1-agents-agent-name-policies-get.ts b/sdks/typescript/src/generated/models/operations/get-agent-policies-api-v1-agents-agent-name-policies-get.ts new file mode 100644 index 00000000..37010486 --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/get-agent-policies-api-v1-agents-agent-name-policies-get.ts @@ -0,0 +1,42 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; + +export type GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest = { + agentName: string; +}; + +/** @internal */ +export type GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest$Outbound = { + agent_name: string; +}; + +/** @internal */ +export const GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest$outboundSchema: + z.ZodMiniType< + GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest$Outbound, + GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest + > = z.pipe( + z.object({ + agentName: z.string(), + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + }); + }), + ); + +export function getAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequestToJSON( + getAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest: + GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest, +): string { + return JSON.stringify( + GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest$outboundSchema.parse( + getAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest, + ), + ); +} diff --git a/sdks/typescript/src/generated/models/operations/index.ts b/sdks/typescript/src/generated/models/operations/index.ts index 8e6bebc2..584ee469 100644 --- a/sdks/typescript/src/generated/models/operations/index.ts +++ b/sdks/typescript/src/generated/models/operations/index.ts @@ -2,6 +2,8 @@ * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. */ +export * from "./add-agent-control-api-v1-agents-agent-name-controls-control-id-post.js"; +export * from "./add-agent-policy-api-v1-agents-agent-name-policies-policy-id-post.js"; export * from "./add-control-to-policy-api-v1-policies-policy-id-controls-control-id-post.js"; export * from "./delete-agent-policy-api-v1-agents-agent-name-policy-delete.js"; export * from "./delete-control-api-v1-controls-control-id-delete.js"; @@ -9,6 +11,7 @@ export * from "./delete-evaluator-config-api-v1-evaluator-configs-config-id-dele export * from "./evaluate-api-v1-evaluation-post.js"; export * from "./get-agent-api-v1-agents-agent-name-get.js"; export * from "./get-agent-evaluator-api-v1-agents-agent-name-evaluators-evaluator-name-get.js"; +export * from "./get-agent-policies-api-v1-agents-agent-name-policies-get.js"; export * from "./get-agent-policy-api-v1-agents-agent-name-policy-get.js"; export * from "./get-control-api-v1-controls-control-id-get.js"; export * from "./get-control-data-api-v1-controls-control-id-data-get.js"; @@ -23,6 +26,9 @@ export * from "./list-evaluator-configs-api-v1-evaluator-configs-get.js"; export * from "./list-policy-controls-api-v1-policies-policy-id-controls-get.js"; export * from "./patch-agent-api-v1-agents-agent-name-patch.js"; export * from "./patch-control-api-v1-controls-control-id-patch.js"; +export * from "./remove-agent-control-api-v1-agents-agent-name-controls-control-id-delete.js"; +export * from "./remove-agent-policy-api-v1-agents-agent-name-policies-policy-id-delete.js"; +export * from "./remove-all-agent-policies-api-v1-agents-agent-name-policies-delete.js"; export * from "./remove-control-from-policy-api-v1-policies-policy-id-controls-control-id-delete.js"; export * from "./set-agent-policy-api-v1-agents-agent-name-policy-policy-id-post.js"; export * from "./set-control-data-api-v1-controls-control-id-data-put.js"; diff --git a/sdks/typescript/src/generated/models/operations/remove-agent-control-api-v1-agents-agent-name-controls-control-id-delete.ts b/sdks/typescript/src/generated/models/operations/remove-agent-control-api-v1-agents-agent-name-controls-control-id-delete.ts new file mode 100644 index 00000000..d230f267 --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/remove-agent-control-api-v1-agents-agent-name-controls-control-id-delete.ts @@ -0,0 +1,49 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; + +export type RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest = + { + agentName: string; + controlId: number; + }; + +/** @internal */ +export type RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest$Outbound = + { + agent_name: string; + control_id: number; + }; + +/** @internal */ +export const RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest$outboundSchema: + z.ZodMiniType< + RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest$Outbound, + RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest + > = z.pipe( + z.object({ + agentName: z.string(), + controlId: z.int(), + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + controlId: "control_id", + }); + }), + ); + +export function removeAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequestToJSON( + removeAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest: + RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest, +): string { + return JSON.stringify( + RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest$outboundSchema + .parse( + removeAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest, + ), + ); +} diff --git a/sdks/typescript/src/generated/models/operations/remove-agent-policy-api-v1-agents-agent-name-policies-policy-id-delete.ts b/sdks/typescript/src/generated/models/operations/remove-agent-policy-api-v1-agents-agent-name-policies-policy-id-delete.ts new file mode 100644 index 00000000..5cbabf9d --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/remove-agent-policy-api-v1-agents-agent-name-policies-policy-id-delete.ts @@ -0,0 +1,49 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; + +export type RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest = + { + agentName: string; + policyId: number; + }; + +/** @internal */ +export type RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest$Outbound = + { + agent_name: string; + policy_id: number; + }; + +/** @internal */ +export const RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest$outboundSchema: + z.ZodMiniType< + RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest$Outbound, + RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest + > = z.pipe( + z.object({ + agentName: z.string(), + policyId: z.int(), + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + policyId: "policy_id", + }); + }), + ); + +export function removeAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequestToJSON( + removeAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest: + RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest, +): string { + return JSON.stringify( + RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest$outboundSchema + .parse( + removeAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest, + ), + ); +} diff --git a/sdks/typescript/src/generated/models/operations/remove-all-agent-policies-api-v1-agents-agent-name-policies-delete.ts b/sdks/typescript/src/generated/models/operations/remove-all-agent-policies-api-v1-agents-agent-name-policies-delete.ts new file mode 100644 index 00000000..0e909a21 --- /dev/null +++ b/sdks/typescript/src/generated/models/operations/remove-all-agent-policies-api-v1-agents-agent-name-policies-delete.ts @@ -0,0 +1,42 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../../lib/primitives.js"; + +export type RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest = { + agentName: string; +}; + +/** @internal */ +export type RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest$Outbound = + { + agent_name: string; + }; + +/** @internal */ +export const RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest$outboundSchema: + z.ZodMiniType< + RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest$Outbound, + RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest + > = z.pipe( + z.object({ + agentName: z.string(), + }), + z.transform((v) => { + return remap$(v, { + agentName: "agent_name", + }); + }), + ); + +export function removeAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequestToJSON( + removeAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest: + RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest, +): string { + return JSON.stringify( + RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest$outboundSchema + .parse(removeAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest), + ); +} diff --git a/sdks/typescript/src/generated/models/remove-agent-control-response.ts b/sdks/typescript/src/generated/models/remove-agent-control-response.ts new file mode 100644 index 00000000..e0c61fde --- /dev/null +++ b/sdks/typescript/src/generated/models/remove-agent-control-response.ts @@ -0,0 +1,56 @@ +/* + * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. + */ + +import * as z from "zod/v4-mini"; +import { remap as remap$ } from "../lib/primitives.js"; +import { safeParse } from "../lib/schemas.js"; +import { Result as SafeParseResult } from "../types/fp.js"; +import * as types from "../types/primitives.js"; +import { SDKValidationError } from "./errors/sdk-validation-error.js"; + +/** + * Response for removing a direct agent-control association. + */ +export type RemoveAgentControlResponse = { + /** + * True if the control remains active via policy association(s) + */ + controlStillActive: boolean; + /** + * True if a direct agent-control link was removed + */ + removedDirectAssociation: boolean; + /** + * Whether the request succeeded + */ + success: boolean; +}; + +/** @internal */ +export const RemoveAgentControlResponse$inboundSchema: z.ZodMiniType< + RemoveAgentControlResponse, + unknown +> = z.pipe( + z.object({ + control_still_active: types.boolean(), + removed_direct_association: types.boolean(), + success: types.boolean(), + }), + z.transform((v) => { + return remap$(v, { + "control_still_active": "controlStillActive", + "removed_direct_association": "removedDirectAssociation", + }); + }), +); + +export function removeAgentControlResponseFromJSON( + jsonString: string, +): SafeParseResult { + return safeParse( + jsonString, + (x) => RemoveAgentControlResponse$inboundSchema.parse(JSON.parse(x)), + `Failed to parse 'RemoveAgentControlResponse' from JSON`, + ); +} diff --git a/sdks/typescript/src/generated/models/set-policy-response.ts b/sdks/typescript/src/generated/models/set-policy-response.ts index 7fa3bfe6..f600a2a0 100644 --- a/sdks/typescript/src/generated/models/set-policy-response.ts +++ b/sdks/typescript/src/generated/models/set-policy-response.ts @@ -9,13 +9,16 @@ import { Result as SafeParseResult } from "../types/fp.js"; import * as types from "../types/primitives.js"; import { SDKValidationError } from "./errors/sdk-validation-error.js"; +/** + * Compatibility response for singular policy assignment endpoint. + */ export type SetPolicyResponse = { /** - * Previous policy id if one was replaced + * Previously associated policy ID, if any */ oldPolicyId?: number | null | undefined; /** - * Whether the policy was successfully assigned + * Whether the request succeeded */ success: boolean; }; diff --git a/sdks/typescript/src/generated/sdk/agents.ts b/sdks/typescript/src/generated/sdk/agents.ts index 98c0c874..d606b384 100644 --- a/sdks/typescript/src/generated/sdk/agents.ts +++ b/sdks/typescript/src/generated/sdk/agents.ts @@ -2,14 +2,20 @@ * Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT. */ +import { agentsAddControl } from "../funcs/agents-add-control.js"; +import { agentsAddPolicy } from "../funcs/agents-add-policy.js"; import { agentsDeletePolicy } from "../funcs/agents-delete-policy.js"; import { agentsGetEvaluator } from "../funcs/agents-get-evaluator.js"; +import { agentsGetPolicies } from "../funcs/agents-get-policies.js"; import { agentsGetPolicy } from "../funcs/agents-get-policy.js"; import { agentsGet } from "../funcs/agents-get.js"; import { agentsInit } from "../funcs/agents-init.js"; import { agentsListControls } from "../funcs/agents-list-controls.js"; import { agentsListEvaluators } from "../funcs/agents-list-evaluators.js"; import { agentsList } from "../funcs/agents-list.js"; +import { agentsRemoveAllAgentPolicies } from "../funcs/agents-remove-all-agent-policies.js"; +import { agentsRemoveControl } from "../funcs/agents-remove-control.js"; +import { agentsRemovePolicy } from "../funcs/agents-remove-policy.js"; import { agentsUpdatePolicy } from "../funcs/agents-update-policy.js"; import { agentsUpdate } from "../funcs/agents-update.js"; import { ClientSDK, RequestOptions } from "../lib/sdks.js"; @@ -24,7 +30,7 @@ export class Agents extends ClientSDK { * @remarks * List all registered agents with cursor-based pagination. * - * Returns a summary of each agent including identifier, policy assignment, + * Returns a summary of each agent including identifier, policy associations, * and counts of registered steps and evaluators. * * Args: @@ -66,7 +72,7 @@ export class Agents extends ClientSDK { * db: Database session (injected) * * Returns: - * InitAgentResponse with created flag and active controls (if policy assigned) + * InitAgentResponse with created flag and active controls (policy-derived + direct) */ async init( request: models.InitAgentRequest, @@ -147,15 +153,14 @@ export class Agents extends ClientSDK { * @remarks * List all protection controls active for an agent. * - * Controls are inherited from the agent's assigned policy. - * Returns an empty list if the agent has no policy. + * Controls include the union of policy-derived and directly associated controls. * * Args: * agent_name: Agent identifier * db: Database session (injected) * * Returns: - * AgentControlsResponse with list of controls (empty if no policy) + * AgentControlsResponse with list of active controls * * Raises: * HTTPException 404: Agent not found @@ -171,6 +176,42 @@ export class Agents extends ClientSDK { )); } + /** + * Remove direct control association from agent + * + * @remarks + * Remove a direct control association from an agent (idempotent). + */ + async removeControl( + request: + operations.RemoveAgentControlApiV1AgentsAgentNameControlsControlIdDeleteRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(agentsRemoveControl( + this, + request, + options, + )); + } + + /** + * Associate control directly with agent + * + * @remarks + * Associate a control directly with an agent (idempotent). + */ + async addControl( + request: + operations.AddAgentControlApiV1AgentsAgentNameControlsControlIdPostRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(agentsAddControl( + this, + request, + options, + )); + } + /** * List agent's registered evaluator schemas * @@ -235,23 +276,84 @@ export class Agents extends ClientSDK { } /** - * Remove agent's policy assignment + * Remove all policy associations from agent * * @remarks - * Remove the policy assignment from an agent. + * Remove all policy associations from an agent. + */ + async removeAllAgentPolicies( + request: + operations.RemoveAllAgentPoliciesApiV1AgentsAgentNamePoliciesDeleteRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(agentsRemoveAllAgentPolicies( + this, + request, + options, + )); + } + + /** + * List policies associated with agent * - * The agent will no longer have any protection controls active. + * @remarks + * List policy IDs associated with an agent. + */ + async getPolicies( + request: operations.GetAgentPoliciesApiV1AgentsAgentNamePoliciesGetRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(agentsGetPolicies( + this, + request, + options, + )); + } + + /** + * Remove policy association from agent * - * Args: - * agent_name: Agent identifier - * db: Database session (injected) + * @remarks + * Remove a policy association from an agent. * - * Returns: - * DeletePolicyResponse with success flag + * Idempotent for existing resources: removing a non-associated link is a no-op. + * Missing agent/policy resources still return 404. + */ + async removePolicy( + request: + operations.RemoveAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdDeleteRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(agentsRemovePolicy( + this, + request, + options, + )); + } + + /** + * Associate policy with agent * - * Raises: - * HTTPException 404: Agent not found or agent has no policy assigned - * HTTPException 500: Database error during removal + * @remarks + * Associate a policy with an agent (idempotent). + */ + async addPolicy( + request: + operations.AddAgentPolicyApiV1AgentsAgentNamePoliciesPolicyIdPostRequest, + options?: RequestOptions, + ): Promise { + return unwrapAsync(agentsAddPolicy( + this, + request, + options, + )); + } + + /** + * Remove agent's policy assignment (compatibility) + * + * @remarks + * Compatibility endpoint that removes all policy associations. */ async deletePolicy( request: @@ -266,20 +368,10 @@ export class Agents extends ClientSDK { } /** - * Get agent's assigned policy + * Get agent's assigned policy (compatibility) * * @remarks - * Retrieve the policy currently assigned to an agent. - * - * Args: - * agent_name: Agent identifier - * db: Database session (injected) - * - * Returns: - * GetPolicyResponse with policy ID - * - * Raises: - * HTTPException 404: Agent not found or agent has no policy assigned + * Compatibility endpoint that returns the first associated policy. */ async getPolicy( request: operations.GetAgentPolicyApiV1AgentsAgentNamePolicyGetRequest, @@ -293,24 +385,10 @@ export class Agents extends ClientSDK { } /** - * Assign policy to agent + * Assign policy to agent (compatibility) * * @remarks - * Assign a policy to an agent, replacing any existing policy assignment. - * - * The agent will immediately inherit all controls from the assigned policy. - * - * Args: - * agent_name: Agent identifier - * policy_id: ID of the policy to assign - * db: Database session (injected) - * - * Returns: - * SetPolicyResponse with success flag and previous policy ID (if any) - * - * Raises: - * HTTPException 404: Agent or policy not found - * HTTPException 500: Database error during assignment + * Compatibility endpoint that replaces all policy associations with one policy. */ async updatePolicy( request: diff --git a/sdks/typescript/src/generated/sdk/controls.ts b/sdks/typescript/src/generated/sdk/controls.ts index e4871a8a..c32be744 100644 --- a/sdks/typescript/src/generated/sdk/controls.ts +++ b/sdks/typescript/src/generated/sdk/controls.ts @@ -113,7 +113,7 @@ export class Controls extends ClientSDK { * @remarks * Delete a control by ID. * - * By default, deletion fails if the control is associated with any policy. + * By default, deletion fails if the control is associated with any policy or agent. * Use force=true to automatically dissociate and delete. * * Args: @@ -122,7 +122,7 @@ export class Controls extends ClientSDK { * db: Database session (injected) * * Returns: - * DeleteControlResponse with success flag and list of dissociated policies + * DeleteControlResponse with success flag and dissociation details * * Raises: * HTTPException 404: Control not found diff --git a/sdks/typescript/tests/client-api.test.ts b/sdks/typescript/tests/client-api.test.ts index 40899b46..c6c11864 100644 --- a/sdks/typescript/tests/client-api.test.ts +++ b/sdks/typescript/tests/client-api.test.ts @@ -106,6 +106,7 @@ describe("AgentControlClient API wiring", () => { await client.agents.init({ agent: { + agentId: "550e8400-e29b-41d4-a716-446655440000", agentName: "test-agent", }, }); diff --git a/server/README.md b/server/README.md index 39435561..05780231 100644 --- a/server/README.md +++ b/server/README.md @@ -157,13 +157,13 @@ GET /api/v1/evaluators ```bash # Register or update agent -POST /api/v1/agents/init +POST /api/v1/agents/initAgent Body: { "agent": {...}, "tools": [...], "force_replace": false } # Get agent GET /api/v1/agents/{agent_name} -# List controls for agent (based on assigned policy) +# List controls for agent (union of policy-associated + directly associated controls) GET /api/v1/agents/{agent_name}/controls ``` @@ -198,8 +198,11 @@ Body: { "name": "my-policy", "description": "..." } # List policies GET /api/v1/policies -# Assign policy to agent -POST /api/v1/policies/{policy_id}/agents/{agent_name} +# Associate policy with agent (additive; policy is optional) +POST /api/v1/agents/{agent_name}/policies/{policy_id} + +# Associate control directly with agent +POST /api/v1/agents/{agent_name}/controls/{control_id} # Add control to policy POST /api/v1/policies/{policy_id}/controls/{control_id} diff --git a/server/alembic/versions/4b8c7d4a1f31_agent_policy_m2m_and_direct_agent_controls.py b/server/alembic/versions/4b8c7d4a1f31_agent_policy_m2m_and_direct_agent_controls.py new file mode 100644 index 00000000..2a4962dc --- /dev/null +++ b/server/alembic/versions/4b8c7d4a1f31_agent_policy_m2m_and_direct_agent_controls.py @@ -0,0 +1,90 @@ +"""agent policy m2m and direct agent controls + +Revision ID: 4b8c7d4a1f31 +Revises: 58920e6807fe +Create Date: 2026-03-02 16:15:00.000000 + +""" + +from alembic import op +import sqlalchemy as sa + + +# revision identifiers, used by Alembic. +revision = "4b8c7d4a1f31" +down_revision = "58920e6807fe" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "agent_controls", + sa.Column("agent_name", sa.String(length=255), nullable=False), + sa.Column("control_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint(["agent_name"], ["agents.name"]), + sa.ForeignKeyConstraint(["control_id"], ["controls.id"]), + sa.PrimaryKeyConstraint("agent_name", "control_id"), + ) + op.create_index(op.f("ix_agent_controls_agent_name"), "agent_controls", ["agent_name"]) + op.create_index(op.f("ix_agent_controls_control_id"), "agent_controls", ["control_id"]) + + op.create_table( + "agent_policies", + sa.Column("agent_name", sa.String(length=255), nullable=False), + sa.Column("policy_id", sa.Integer(), nullable=False), + sa.ForeignKeyConstraint(["agent_name"], ["agents.name"]), + sa.ForeignKeyConstraint(["policy_id"], ["policies.id"]), + sa.PrimaryKeyConstraint("agent_name", "policy_id"), + ) + op.create_index(op.f("ix_agent_policies_agent_name"), "agent_policies", ["agent_name"]) + op.create_index(op.f("ix_agent_policies_policy_id"), "agent_policies", ["policy_id"]) + + op.execute( + sa.text( + """ + INSERT INTO agent_policies (agent_name, policy_id) + SELECT name, policy_id + FROM agents + WHERE policy_id IS NOT NULL + ON CONFLICT (agent_name, policy_id) DO NOTHING + """ + ) + ) + + op.drop_index(op.f("ix_agents_policy_id"), table_name="agents") + op.drop_constraint(op.f("agents_policy_id_fkey"), "agents", type_="foreignkey") + op.drop_column("agents", "policy_id") + + +def downgrade() -> None: + # NOTE: Downgrade can only restore one policy per agent via agents.policy_id. + # If multiple policies are associated to an agent, all but MIN(policy_id) are lost. + op.add_column("agents", sa.Column("policy_id", sa.Integer(), nullable=True)) + + op.execute( + sa.text( + """ + UPDATE agents AS a + SET policy_id = ap.policy_id + FROM ( + SELECT agent_name, MIN(policy_id) AS policy_id + FROM agent_policies + GROUP BY agent_name + ) AS ap + WHERE a.name = ap.agent_name + """ + ) + ) + + op.create_foreign_key(op.f("agents_policy_id_fkey"), "agents", "policies", ["policy_id"], ["id"]) + op.create_index(op.f("ix_agents_policy_id"), "agents", ["policy_id"]) + + op.drop_index(op.f("ix_agent_policies_policy_id"), table_name="agent_policies") + op.drop_index(op.f("ix_agent_policies_agent_name"), table_name="agent_policies") + op.drop_table("agent_policies") + + # Direct agent-control links have no representation in the downgraded schema. + op.drop_index(op.f("ix_agent_controls_control_id"), table_name="agent_controls") + op.drop_index(op.f("ix_agent_controls_agent_name"), table_name="agent_controls") + op.drop_table("agent_controls") diff --git a/server/src/agent_control_server/endpoints/agents.py b/server/src/agent_control_server/endpoints/agents.py index 82cecfc9..423d38db 100644 --- a/server/src/agent_control_server/endpoints/agents.py +++ b/server/src/agent_control_server/endpoints/agents.py @@ -7,9 +7,11 @@ from agent_control_models.server import ( AgentControlsResponse, AgentSummary, + AssocResponse, ConflictMode, DeletePolicyResponse, EvaluatorSchema, + GetAgentPoliciesResponse, GetAgentResponse, GetPolicyResponse, InitAgentEvaluatorRemoval, @@ -20,13 +22,15 @@ PaginationInfo, PatchAgentRequest, PatchAgentResponse, + RemoveAgentControlResponse, SetPolicyResponse, StepKey, ) from fastapi import APIRouter, Depends from jsonschema_rs import ValidationError as JSONSchemaValidationError from pydantic import BaseModel, ValidationError -from sqlalchemy import func, or_, select +from sqlalchemy import delete, func, or_, select, union_all +from sqlalchemy.dialects.postgresql import insert as pg_insert from sqlalchemy.ext.asyncio import AsyncSession from ..db import get_async_db @@ -43,6 +47,8 @@ AgentData, Control, Policy, + agent_controls, + agent_policies, policy_controls, ) from ..services.agent_names import normalize_agent_name_or_422 @@ -51,6 +57,7 @@ parse_evaluator_ref_full, validate_config_against_schema, ) +from ..services.query_utils import escape_like_pattern from ..services.schema_compat import ( check_schema_compatibility, format_compatibility_error, @@ -85,18 +92,8 @@ def _get_builtin_evaluator_names() -> set[str]: return _BUILTIN_EVALUATOR_NAMES -async def _validate_policy_controls_for_agent( - agent: Agent, policy_id: int, db: AsyncSession -) -> list[str]: - """Validate all controls in a policy can run on this agent. - - Checks that agent-scoped evaluators referenced by controls: - 1. Exist on the agent (registered via initAgent) - 2. Have config that validates against the evaluator's schema - - Returns: - List of error messages (empty if all valid) - """ +def _validate_controls_for_agent(agent: Agent, controls: list[Control]) -> list[str]: + """Validate controls can run on this agent.""" errors: list[str] = [] # Parse agent's registered evaluators @@ -107,9 +104,6 @@ async def _validate_policy_controls_for_agent( agent_evaluators = {e.name: e for e in (agent_data.evaluators or [])} - # Get all controls for this policy - controls = await list_controls_for_policy(policy_id, db) - for control in controls: if not control.data: continue @@ -155,6 +149,14 @@ async def _validate_policy_controls_for_agent( return errors +async def _validate_policy_controls_for_agent( + agent: Agent, policy_id: int, db: AsyncSession +) -> list[str]: + """Validate all controls in a policy can run on this agent.""" + controls = await list_controls_for_policy(policy_id, db) + return _validate_controls_for_agent(agent, controls) + + def _step_registration_changed(existing_step: StepSchema, incoming_step: StepSchema) -> bool: """Return True when non-key step registration fields differ.""" return ( @@ -177,8 +179,8 @@ async def _build_overwrite_evaluator_removals( db: AsyncSession, ) -> list[InitAgentEvaluatorRemoval]: """Build evaluator removal details, including active-control references.""" - if not removed_evaluators or agent.policy_id is None: - return [InitAgentEvaluatorRemoval(name=name) for name in sorted(removed_evaluators)] + if not removed_evaluators: + return [] try: controls = await list_controls_for_agent( @@ -239,7 +241,7 @@ async def list_agents( """ List all registered agents with cursor-based pagination. - Returns a summary of each agent including identifier, policy assignment, + Returns a summary of each agent including identifier, policy associations, and counts of registered steps and evaluators. Args: @@ -255,7 +257,9 @@ async def list_agents( limit = min(max(1, limit), _MAX_PAGINATION_LIMIT) # Build base filter for name search - name_filter = Agent.name.ilike(f"%{name}%") if name else None + name_filter = ( + Agent.name.ilike(f"%{escape_like_pattern(name)}%", escape="\\") if name else None + ) # Get total count (with name filter if provided) count_query = select(func.count()).select_from(Agent) @@ -303,32 +307,54 @@ async def list_agents( if has_more and agents: next_cursor = agents[-1].name - # Batch query: Get control counts for all agents at once - # Join: Agent -> Policy -> policy_controls (junction table) -> Control - # Count distinct enabled control IDs per agent - # Performance: Filter NULL controls explicitly to avoid JSONB parsing on NULL rows - # This allows the query planner to optimize better control_counts_map: dict[str, int] = {} + policy_ids_map: dict[str, list[int]] = {} if agents: + agent_names = [agent.name for agent in agents] + + policy_ids_query = ( + select( + agent_policies.c.agent_name, + agent_policies.c.policy_id, + ) + .where(agent_policies.c.agent_name.in_(agent_names)) + .order_by(agent_policies.c.agent_name, agent_policies.c.policy_id) + ) + policy_ids_result = await db.execute(policy_ids_query) + for assoc_agent_name, policy_id in policy_ids_result.all(): + policy_ids_map.setdefault(assoc_agent_name, []).append(policy_id) + + policy_associations = ( + select( + agent_policies.c.agent_name.label("agent_name"), + policy_controls.c.control_id.label("control_id"), + ) + .select_from( + agent_policies.join( + policy_controls, agent_policies.c.policy_id == policy_controls.c.policy_id + ) + ) + .where(agent_policies.c.agent_name.in_(agent_names)) + ) + direct_associations = select( + agent_controls.c.agent_name.label("agent_name"), + agent_controls.c.control_id.label("control_id"), + ).where(agent_controls.c.agent_name.in_(agent_names)) + all_associations = union_all(policy_associations, direct_associations).subquery() + control_counts_query = ( select( - Agent.name, - func.count(func.distinct(policy_controls.c.control_id)).label("count"), + all_associations.c.agent_name, + func.count(func.distinct(all_associations.c.control_id)).label("count"), ) - .outerjoin(Policy, Agent.policy_id == Policy.id) - .outerjoin(policy_controls, Policy.id == policy_controls.c.policy_id) - .outerjoin(Control, policy_controls.c.control_id == Control.id) + .join(Control, all_associations.c.control_id == Control.id) .where( - Agent.name.in_([agent.name for agent in agents]), - # Only count enabled controls: Control must exist AND be enabled - # (enabled=true OR enabled key missing, default is True) - Control.id.is_not(None), # Exclude NULL controls (agents without policies) or_( Control.data["enabled"].astext == "true", ~Control.data.has_key("enabled"), ), ) - .group_by(Agent.name) + .group_by(all_associations.c.agent_name) ) control_counts_result = await db.execute(control_counts_query) control_counts_map = {row[0]: row[1] for row in control_counts_result.all()} @@ -351,10 +377,12 @@ async def list_agents( # Get active controls count from batched query result active_controls = control_counts_map.get(agent.name, 0) + policy_ids = policy_ids_map.get(agent.name, []) + summaries.append( AgentSummary( agent_name=agent.name, - policy_id=agent.policy_id, + policy_ids=policy_ids, created_at=agent.created_at.isoformat() if agent.created_at else None, step_count=step_count, evaluator_count=evaluator_count, @@ -398,7 +426,7 @@ async def init_agent( db: Database session (injected) Returns: - InitAgentResponse with created flag and active controls (if policy assigned) + InitAgentResponse with created flag and active controls (policy-derived + direct) """ # Check for evaluator name collisions with built-in evaluators builtin_names = _get_builtin_evaluator_names() @@ -729,10 +757,7 @@ async def init_agent( operation="update", ) - # If the existing agent has a policy, include its controls; otherwise empty list - controls = [] - if existing.policy_id is not None: - controls = await list_controls_for_agent(existing.name, db) + controls = await list_controls_for_agent(existing.name, db) return InitAgentResponse( created=created, @@ -818,46 +843,99 @@ async def get_agent(agent_name: str, db: AsyncSession = Depends(get_async_db)) - ) +async def _get_agent_or_404(agent_name: str, db: AsyncSession) -> Agent: + """Get an agent or raise AGENT_NOT_FOUND.""" + normalized_agent_name = normalize_agent_name_or_422(agent_name) + result = await db.execute(select(Agent).where(Agent.name == normalized_agent_name)) + agent: Agent | None = result.scalars().first() + if agent is None: + raise NotFoundError( + error_code=ErrorCode.AGENT_NOT_FOUND, + detail=f"Agent with name '{normalized_agent_name}' not found", + resource="Agent", + resource_id=normalized_agent_name, + hint="Verify the agent name is correct and the agent has been registered.", + ) + return agent + + @router.post( - "/{agent_name}/policy/{policy_id}", - response_model=SetPolicyResponse, - summary="Assign policy to agent", - response_description="Success status with previous policy ID", + "/{agent_name}/policies/{policy_id}", + response_model=AssocResponse, + summary="Associate policy with agent", + response_description="Success confirmation", ) -async def set_agent_policy( +async def add_agent_policy( agent_name: str, policy_id: int, db: AsyncSession = Depends(get_async_db) -) -> SetPolicyResponse: - """ - Assign a policy to an agent, replacing any existing policy assignment. - - The agent will immediately inherit all controls from the assigned policy. +) -> AssocResponse: + """Associate a policy with an agent (idempotent).""" + agent = await _get_agent_or_404(agent_name, db) - Args: - agent_name: Agent identifier - policy_id: ID of the policy to assign - db: Database session (injected) + policy_result = await db.execute(select(Policy).where(Policy.id == policy_id)) + policy: Policy | None = policy_result.scalars().first() + if policy is None: + raise NotFoundError( + error_code=ErrorCode.POLICY_NOT_FOUND, + detail=f"Policy with ID '{policy_id}' not found", + resource="Policy", + resource_id=str(policy_id), + hint="Verify the policy ID is correct and the policy has been created.", + ) - Returns: - SetPolicyResponse with success flag and previous policy ID (if any) + validation_errors = await _validate_policy_controls_for_agent(agent, policy_id, db) + if validation_errors: + raise BadRequestError( + error_code=ErrorCode.POLICY_CONTROL_INCOMPATIBLE, + detail="Policy contains controls incompatible with this agent", + hint="Ensure all controls in the policy are compatible with this agent's evaluators.", + errors=[ + ValidationErrorItem( + resource="Control", + field="evaluator", + code="incompatible", + message=err, + ) + for err in validation_errors + ], + ) - Raises: - HTTPException 404: Agent or policy not found - HTTPException 500: Database error during assignment - """ - agent_name = normalize_agent_name_or_422(agent_name) - # Find agent - result = await db.execute(select(Agent).where(Agent.name == agent_name)) - agent: Agent | None = result.scalars().first() - if agent is None: - raise NotFoundError( - error_code=ErrorCode.AGENT_NOT_FOUND, - detail=f"Agent with name '{agent_name}' not found", + try: + stmt = ( + pg_insert(agent_policies) + .values(agent_name=agent.name, policy_id=policy_id) + .on_conflict_do_nothing() + ) + await db.execute(stmt) + await db.commit() + except Exception: + await db.rollback() + _logger.error( + "Failed to associate policy '%s' with agent '%s'", + policy_id, + agent.name, + exc_info=True, + ) + raise DatabaseError( + detail=f"Failed to associate policy with agent '{agent.name}': database error", resource="Agent", - resource_id=str(agent_name), - hint="Verify the agent name is correct and the agent has been registered.", + operation="add policy association", ) - # Find policy by id + return AssocResponse(success=True) + + +@router.post( + "/{agent_name}/policy/{policy_id}", + response_model=SetPolicyResponse, + summary="Assign policy to agent (compatibility)", + response_description="Success status with previous policy ID", +) +async def set_agent_policy( + agent_name: str, policy_id: int, db: AsyncSession = Depends(get_async_db) +) -> SetPolicyResponse: + """Compatibility endpoint that replaces all policy associations with one policy.""" + agent = await _get_agent_or_404(agent_name, db) + policy_result = await db.execute(select(Policy).where(Policy.id == policy_id)) policy: Policy | None = policy_result.scalars().first() if policy is None: @@ -869,7 +947,6 @@ async def set_agent_policy( hint="Verify the policy ID is correct and the policy has been created.", ) - # Validate controls can run on this agent validation_errors = await _validate_policy_controls_for_agent(agent, policy_id, db) if validation_errors: raise BadRequestError( @@ -887,19 +964,28 @@ async def set_agent_policy( ], ) - # Store old policy ID if exists - old_policy_id: int | None = None - if agent.policy_id is not None: - old_policy_id = agent.policy_id + existing_policies_result = await db.execute( + select(agent_policies.c.policy_id) + .where(agent_policies.c.agent_name == agent.name) + .order_by(agent_policies.c.policy_id) + ) + existing_policy_ids = [row[0] for row in existing_policies_result.all()] + old_policy_id = existing_policy_ids[0] if existing_policy_ids else None - # Assign new policy - agent.policy_id = policy.id try: + await db.execute(delete(agent_policies).where(agent_policies.c.agent_name == agent.name)) + await db.execute( + pg_insert(agent_policies) + .values(agent_name=agent.name, policy_id=policy_id) + .on_conflict_do_nothing() + ) await db.commit() except Exception: await db.rollback() _logger.error( - f"Failed to assign policy '{policy_id}' to agent '{agent.name}' ({agent_name})", + "Failed to assign policy '%s' to agent '%s'", + policy_id, + agent.name, exc_info=True, ) raise DatabaseError( @@ -911,139 +997,332 @@ async def set_agent_policy( return SetPolicyResponse(success=True, old_policy_id=old_policy_id) +@router.get( + "/{agent_name}/policies", + response_model=GetAgentPoliciesResponse, + summary="List policies associated with agent", + response_description="List of policy IDs", +) +async def get_agent_policies( + agent_name: str, db: AsyncSession = Depends(get_async_db) +) -> GetAgentPoliciesResponse: + """List policy IDs associated with an agent.""" + agent = await _get_agent_or_404(agent_name, db) + result = await db.execute( + select(agent_policies.c.policy_id) + .where(agent_policies.c.agent_name == agent.name) + .order_by(agent_policies.c.policy_id) + ) + return GetAgentPoliciesResponse(policy_ids=[row[0] for row in result.all()]) + + @router.get( "/{agent_name}/policy", response_model=GetPolicyResponse, - summary="Get agent's assigned policy", + summary="Get agent's assigned policy (compatibility)", response_description="Policy ID", ) async def get_agent_policy( agent_name: str, db: AsyncSession = Depends(get_async_db) ) -> GetPolicyResponse: - """ - Retrieve the policy currently assigned to an agent. + """Compatibility endpoint that returns the first associated policy.""" + agent = await _get_agent_or_404(agent_name, db) + policy_result = await db.execute( + select(Policy.id) + .join(agent_policies, agent_policies.c.policy_id == Policy.id) + .where(agent_policies.c.agent_name == agent.name) + .order_by(Policy.id) + .limit(1) + ) + policy_id = policy_result.scalars().first() + if policy_id is None: + raise NotFoundError( + error_code=ErrorCode.POLICY_NOT_FOUND, + detail=f"Agent '{agent.name}' has no policy assigned", + resource="Policy", + hint="Assign a policy to the agent using POST /{agent_name}/policies/{policy_id}.", + ) + return GetPolicyResponse(policy_id=policy_id) - Args: - agent_name: Agent identifier - db: Database session (injected) - Returns: - GetPolicyResponse with policy ID +@router.delete( + "/{agent_name}/policies/{policy_id}", + response_model=AssocResponse, + summary="Remove policy association from agent", + response_description="Success confirmation", +) +async def remove_agent_policy( + agent_name: str, policy_id: int, db: AsyncSession = Depends(get_async_db) +) -> AssocResponse: + """Remove a policy association from an agent. - Raises: - HTTPException 404: Agent not found or agent has no policy assigned + Idempotent for existing resources: removing a non-associated link is a no-op. + Missing agent/policy resources still return 404. """ - agent_name = normalize_agent_name_or_422(agent_name) - # Find agent - result = await db.execute(select(Agent).where(Agent.name == agent_name)) - agent: Agent | None = result.scalars().first() - if agent is None: - raise NotFoundError( - error_code=ErrorCode.AGENT_NOT_FOUND, - detail=f"Agent with name '{agent_name}' not found", - resource="Agent", - resource_id=str(agent_name), - hint="Verify the agent name is correct and the agent has been registered.", - ) + agent = await _get_agent_or_404(agent_name, db) - # Check if agent has a policy - if agent.policy_id is None: + policy_result = await db.execute(select(Policy.id).where(Policy.id == policy_id)) + if policy_result.first() is None: raise NotFoundError( error_code=ErrorCode.POLICY_NOT_FOUND, - detail=f"Agent '{agent.name}' has no policy assigned", + detail=f"Policy with ID '{policy_id}' not found", resource="Policy", - hint="Assign a policy to the agent using POST /{agent_name}/policy/{policy_id}.", + resource_id=str(policy_id), + hint="Verify the policy ID is correct and the policy has been created.", ) - # Find policy - policy_result = await db.execute(select(Policy).where(Policy.id == agent.policy_id)) - policy: Policy | None = policy_result.scalars().first() - if policy is None: - raise NotFoundError( - error_code=ErrorCode.POLICY_NOT_FOUND, + try: + await db.execute( + delete(agent_policies).where( + (agent_policies.c.agent_name == agent.name) + & (agent_policies.c.policy_id == policy_id) + ) + ) + await db.commit() + except Exception: + await db.rollback() + _logger.error( + "Failed to remove policy '%s' from agent '%s'", + policy_id, + agent.name, + exc_info=True, + ) + raise DatabaseError( + detail=f"Failed to remove policy association from agent '{agent.name}': database error", + resource="Agent", + operation="remove policy association", + ) + + return AssocResponse(success=True) + + +@router.delete( + "/{agent_name}/policies", + response_model=AssocResponse, + summary="Remove all policy associations from agent", + response_description="Success confirmation", +) +async def remove_all_agent_policies( + agent_name: str, db: AsyncSession = Depends(get_async_db) +) -> AssocResponse: + """Remove all policy associations from an agent.""" + agent = await _get_agent_or_404(agent_name, db) + + try: + await db.execute(delete(agent_policies).where(agent_policies.c.agent_name == agent.name)) + await db.commit() + except Exception: + await db.rollback() + _logger.error( + "Failed to remove all policies from agent '%s'", + agent.name, + exc_info=True, + ) + raise DatabaseError( detail=( - f"Policy with ID '{agent.policy_id}' not found " - f"(referenced by agent '{agent.name}')" + f"Failed to remove policy associations from agent '{agent.name}': " + "database error" ), - resource="Policy", - resource_id=str(agent.policy_id), - hint="The referenced policy may have been deleted. Assign a new policy to the agent.", + resource="Agent", + operation="remove all policy associations", ) - return GetPolicyResponse(policy_id=policy.id) + return AssocResponse(success=True) @router.delete( "/{agent_name}/policy", response_model=DeletePolicyResponse, - summary="Remove agent's policy assignment", + summary="Remove agent's policy assignment (compatibility)", response_description="Success confirmation", ) async def delete_agent_policy( agent_name: str, db: AsyncSession = Depends(get_async_db) ) -> DeletePolicyResponse: - """ - Remove the policy assignment from an agent. + """Compatibility endpoint that removes all policy associations.""" + agent = await _get_agent_or_404(agent_name, db) - The agent will no longer have any protection controls active. + existing_policy_result = await db.execute( + select(agent_policies.c.policy_id) + .where(agent_policies.c.agent_name == agent.name) + .limit(1) + ) + if existing_policy_result.first() is None: + raise NotFoundError( + error_code=ErrorCode.POLICY_NOT_FOUND, + detail=f"Agent '{agent.name}' has no policy assigned", + resource="Policy", + hint="The agent does not have a policy to remove.", + ) - Args: - agent_name: Agent identifier - db: Database session (injected) + try: + await db.execute(delete(agent_policies).where(agent_policies.c.agent_name == agent.name)) + await db.commit() + except Exception: + await db.rollback() + _logger.error( + "Failed to remove policy associations from agent '%s'", + agent.name, + exc_info=True, + ) + raise DatabaseError( + detail=( + f"Failed to remove policy associations from agent '{agent.name}': " + "database error" + ), + resource="Agent", + operation="remove policy", + ) - Returns: - DeletePolicyResponse with success flag + return DeletePolicyResponse(success=True) - Raises: - HTTPException 404: Agent not found or agent has no policy assigned - HTTPException 500: Database error during removal - """ - agent_name = normalize_agent_name_or_422(agent_name) - # Find agent - result = await db.execute(select(Agent).where(Agent.name == agent_name)) - agent: Agent | None = result.scalars().first() - if agent is None: + +@router.post( + "/{agent_name}/controls/{control_id}", + response_model=AssocResponse, + summary="Associate control directly with agent", + response_description="Success confirmation", +) +async def add_agent_control( + agent_name: str, control_id: int, db: AsyncSession = Depends(get_async_db) +) -> AssocResponse: + """Associate a control directly with an agent (idempotent).""" + agent = await _get_agent_or_404(agent_name, db) + + control_result = await db.execute(select(Control).where(Control.id == control_id)) + control: Control | None = control_result.scalars().first() + if control is None: raise NotFoundError( - error_code=ErrorCode.AGENT_NOT_FOUND, - detail=f"Agent with name '{agent_name}' not found", + error_code=ErrorCode.CONTROL_NOT_FOUND, + detail=f"Control with ID '{control_id}' not found", + resource="Control", + resource_id=str(control_id), + hint="Verify the control ID is correct and the control has been created.", + ) + + validation_errors = _validate_controls_for_agent(agent, [control]) + if validation_errors: + raise BadRequestError( + error_code=ErrorCode.POLICY_CONTROL_INCOMPATIBLE, + detail="Control is incompatible with this agent", + hint="Ensure the control is compatible with this agent's evaluators.", + errors=[ + ValidationErrorItem( + resource="Control", + field="evaluator", + code="incompatible", + message=err, + ) + for err in validation_errors + ], + ) + + try: + stmt = ( + pg_insert(agent_controls) + .values(agent_name=agent.name, control_id=control_id) + .on_conflict_do_nothing() + ) + await db.execute(stmt) + await db.commit() + except Exception: + await db.rollback() + _logger.error( + "Failed to associate control '%s' with agent '%s'", + control_id, + agent.name, + exc_info=True, + ) + raise DatabaseError( + detail=f"Failed to associate control with agent '{agent.name}': database error", resource="Agent", - resource_id=str(agent_name), - hint="Verify the agent name is correct and the agent has been registered.", + operation="add control association", ) - # Check if agent has a policy - if agent.policy_id is None: + return AssocResponse(success=True) + + +@router.delete( + "/{agent_name}/controls/{control_id}", + response_model=RemoveAgentControlResponse, + summary="Remove direct control association from agent", + response_description="Success confirmation", +) +async def remove_agent_control( + agent_name: str, control_id: int, db: AsyncSession = Depends(get_async_db) +) -> RemoveAgentControlResponse: + """Remove a direct control association from an agent (idempotent).""" + agent = await _get_agent_or_404(agent_name, db) + + control_result = await db.execute(select(Control.id).where(Control.id == control_id)) + if control_result.first() is None: raise NotFoundError( - error_code=ErrorCode.POLICY_NOT_FOUND, - detail=f"Agent '{agent.name}' has no policy assigned", - resource="Policy", - hint="The agent does not have a policy to remove.", + error_code=ErrorCode.CONTROL_NOT_FOUND, + detail=f"Control with ID '{control_id}' not found", + resource="Control", + resource_id=str(control_id), + hint="Verify the control ID is correct and the control has been created.", ) - # Remove policy assignment - agent.policy_id = None try: + remove_direct_stmt = ( + delete(agent_controls) + .where( + (agent_controls.c.agent_name == agent.name) + & (agent_controls.c.control_id == control_id) + ) + .returning(agent_controls.c.control_id) + ) + remove_direct_result = await db.execute(remove_direct_stmt) + removed_direct_association = remove_direct_result.first() is not None + + # The control may still be active for this agent if inherited from policy association(s). + policy_inheritance_result = await db.execute( + select(policy_controls.c.control_id) + .select_from( + agent_policies.join( + policy_controls, + agent_policies.c.policy_id == policy_controls.c.policy_id, + ) + ) + .where( + (agent_policies.c.agent_name == agent.name) + & (policy_controls.c.control_id == control_id) + ) + .limit(1) + ) + control_still_active = policy_inheritance_result.first() is not None + await db.commit() except Exception: await db.rollback() _logger.error( - f"Failed to remove policy from agent '{agent.name}' ({agent_name})", + "Failed to remove control '%s' from agent '%s'", + control_id, + agent.name, exc_info=True, ) raise DatabaseError( - detail=f"Failed to remove policy from agent '{agent.name}': database error", + detail=( + f"Failed to remove control association from agent '{agent.name}': " + "database error" + ), resource="Agent", - operation="remove policy", + operation="remove control association", ) - return DeletePolicyResponse(success=True) + return RemoveAgentControlResponse( + success=True, + removed_direct_association=removed_direct_association, + control_still_active=control_still_active, + ) @router.get( "/{agent_name}/controls", response_model=AgentControlsResponse, summary="List agent's active controls", - response_description="List of controls from agent's policy", + response_description="List of controls from agent policy and direct associations", ) async def list_agent_controls( agent_name: str, db: AsyncSession = Depends(get_async_db) @@ -1051,35 +1330,20 @@ async def list_agent_controls( """ List all protection controls active for an agent. - Controls are inherited from the agent's assigned policy. - Returns an empty list if the agent has no policy. + Controls include the union of policy-derived and directly associated controls. Args: agent_name: Agent identifier db: Database session (injected) Returns: - AgentControlsResponse with list of controls (empty if no policy) + AgentControlsResponse with list of active controls Raises: HTTPException 404: Agent not found """ - agent_name = normalize_agent_name_or_422(agent_name) - result = await db.execute(select(Agent).where(Agent.name == agent_name)) - agent: Agent | None = result.scalars().first() - if agent is None: - raise NotFoundError( - error_code=ErrorCode.AGENT_NOT_FOUND, - detail=f"Agent with name '{agent_name}' not found", - resource="Agent", - resource_id=str(agent_name), - hint="Verify the agent name is correct and the agent has been registered.", - ) - - if agent.policy_id is None: - return AgentControlsResponse(controls=[]) - - controls = await list_controls_for_agent(agent_name, db) + agent = await _get_agent_or_404(agent_name, db) + controls = await list_controls_for_agent(agent.name, db) return AgentControlsResponse(controls=controls) @@ -1338,37 +1602,35 @@ async def patch_agent( if request.remove_evaluators: remove_evaluator_set = set(request.remove_evaluators) - # Check if any controls reference evaluators being removed - if agent.policy_id is not None: - # Get all controls for this agent's policy - controls = await list_controls_for_agent(agent.name, db) - referencing_controls: list[tuple[str, str]] = [] # (control_name, evaluator) - - for ctrl in controls: - evaluator_ref = ctrl.control.evaluator.name - if ":" in evaluator_ref: - ref_agent, ref_eval = evaluator_ref.split(":", 1) - # Check if this control references an evaluator we're removing - # AND it's scoped to this agent (by name match) - if ref_agent == agent.name and ref_eval in remove_evaluator_set: - referencing_controls.append((ctrl.name, ref_eval)) - - if referencing_controls: - raise ConflictError( - error_code=ErrorCode.EVALUATOR_IN_USE, - detail="Cannot remove evaluators: active controls reference them", - resource="Evaluator", - hint="Remove or update the controls that reference these evaluators first.", - errors=[ - ValidationErrorItem( - resource="Control", - field="evaluator.name", - code="in_use", - message=f"Control '{ctrl}' uses evaluator '{ev}'", - ) - for ctrl, ev in referencing_controls - ], - ) + # Check if any active controls reference evaluators being removed. + controls = await list_controls_for_agent(agent.name, db) + referencing_controls: list[tuple[str, str]] = [] # (control_name, evaluator) + + for ctrl in controls: + evaluator_ref = ctrl.control.evaluator.name + if ":" in evaluator_ref: + ref_agent, ref_eval = evaluator_ref.split(":", 1) + # Check if this control references an evaluator we're removing + # AND it's scoped to this agent (by name match) + if ref_agent == agent.name and ref_eval in remove_evaluator_set: + referencing_controls.append((ctrl.name, ref_eval)) + + if referencing_controls: + raise ConflictError( + error_code=ErrorCode.EVALUATOR_IN_USE, + detail="Cannot remove evaluators: active controls reference them", + resource="Evaluator", + hint="Remove or update the controls that reference these evaluators first.", + errors=[ + ValidationErrorItem( + resource="Control", + field="evaluator.name", + code="in_use", + message=f"Control '{ctrl}' uses evaluator '{ev}'", + ) + for ctrl, ev in referencing_controls + ], + ) new_evaluators = [] for ev in data_model.evaluators or []: diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index 97a85f17..e0d5b805 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -21,7 +21,8 @@ from fastapi import APIRouter, Depends, Query from jsonschema_rs import ValidationError as JSONSchemaValidationError from pydantic import ValidationError -from sqlalchemy import delete, func, or_, select +from sqlalchemy import Integer, String, delete, func, literal, or_, select, union_all +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from ..db import get_async_db @@ -32,11 +33,12 @@ NotFoundError, ) from ..logging_utils import get_logger -from ..models import Agent, AgentData, Control, Policy, policy_controls +from ..models import Agent, AgentData, Control, agent_controls, agent_policies, policy_controls from ..services.evaluator_utils import ( parse_evaluator_ref_full, validate_config_against_schema, ) +from ..services.query_utils import escape_like_pattern # Pagination constants _DEFAULT_PAGINATION_LIMIT = 20 @@ -235,6 +237,15 @@ async def create_control( try: await db.commit() await db.refresh(control) + except IntegrityError: + await db.rollback() + raise ConflictError( + error_code=ErrorCode.CONTROL_NAME_CONFLICT, + detail=f"Control with name '{request.name}' already exists", + resource="Control", + resource_id=request.name, + hint="Choose a different name or update the existing control.", + ) except Exception: await db.rollback() _logger.error( @@ -496,18 +507,15 @@ async def list_controls( Example: GET /controls?limit=10&enabled=true&step_type=tool """ - # Get total count (with filters applied) - count_query = select(func.count()).select_from(Control) query = select(Control).order_by(Control.id.desc()) # Apply cursor if cursor is not None: query = query.where(Control.id < cursor) - count_query = count_query.where(Control.id < cursor) # Apply name filter (case-insensitive partial match) if name is not None: - query = query.where(Control.name.ilike(f"%{name}%")) + query = query.where(Control.name.ilike(f"%{escape_like_pattern(name)}%", escape="\\")) # Don't apply to count_query - total should be pre-filter # Apply JSONB filters at database level @@ -554,7 +562,9 @@ async def list_controls( # Get total count (with same filters, but without cursor/limit) total_query = select(func.count()).select_from(Control) if name is not None: - total_query = total_query.where(Control.name.ilike(f"%{name}%")) + total_query = total_query.where( + Control.name.ilike(f"%{escape_like_pattern(name)}%", escape="\\") + ) if enabled is not None: if enabled: total_query = total_query.where( @@ -593,27 +603,45 @@ async def list_controls( if has_more: controls = controls[:-1] - # Build mapping of control_id -> agent that uses it - # Traversal: Control -> policy_controls -> Policy -> Agent + # Build mapping of control_id -> usage attribution + # Traversal includes both: + # - Control -> policy_controls -> agent_policies -> Agent + # - Control -> agent_controls -> Agent control_agent_map: dict[int, AgentRef | None] = {ctrl.id: None for ctrl in controls} + control_agent_names_map: dict[int, set[str]] = {ctrl.id: set() for ctrl in controls} + control_agent_repr_map: dict[int, str | None] = {ctrl.id: None for ctrl in controls} if controls: control_ids = [ctrl.id for ctrl in controls] - agents_query = ( + policy_agents_query = ( select( policy_controls.c.control_id, - Agent.name, + agent_policies.c.agent_name, ) .select_from(policy_controls) - .join(Policy, policy_controls.c.policy_id == Policy.id) - .join(Agent, Agent.policy_id == Policy.id) + .join(agent_policies, policy_controls.c.policy_id == agent_policies.c.policy_id) .where(policy_controls.c.control_id.in_(control_ids)) ) + direct_agents_query = ( + select( + agent_controls.c.control_id, + agent_controls.c.agent_name, + ) + .select_from(agent_controls) + .where(agent_controls.c.control_id.in_(control_ids)) + ) + agents_query = union_all(policy_agents_query, direct_agents_query) agents_result = await db.execute(agents_query) for row in agents_result.all(): control_id, agent_name = row - # Take the first agent found (1 control = 1 agent) - if control_agent_map[control_id] is None: - control_agent_map[control_id] = AgentRef(agent_name=agent_name) + control_agent_names_map[control_id].add(agent_name) + + # Keep a deterministic representative agent for backwards compatibility. + current_repr = control_agent_repr_map[control_id] + if current_repr is None or agent_name < current_repr: + control_agent_repr_map[control_id] = agent_name + control_agent_map[control_id] = AgentRef( + agent_name=agent_name + ) # Build summaries (filtering already done at DB level) summaries: list[ControlSummary] = [] @@ -632,6 +660,7 @@ async def list_controls( stages=scope.get("stages"), tags=data.get("tags", []), used_by_agent=control_agent_map.get(ctrl.id), + used_by_agents_count=len(control_agent_names_map.get(ctrl.id, set())), ) ) @@ -661,15 +690,15 @@ async def delete_control( control_id: int, force: bool = Query( False, - description="If true, dissociate from all policies before deleting. " - "If false, fail if control is associated with any policy.", + description="If true, dissociate from all policy/agent links before deleting. " + "If false, fail if control is associated with any policy or agent.", ), db: AsyncSession = Depends(get_async_db), ) -> DeleteControlResponse: """ Delete a control by ID. - By default, deletion fails if the control is associated with any policy. + By default, deletion fails if the control is associated with any policy or agent. Use force=true to automatically dissociate and delete. Args: @@ -678,7 +707,7 @@ async def delete_control( db: Database session (injected) Returns: - DeleteControlResponse with success flag and list of dissociated policies + DeleteControlResponse with success flag and dissociation details Raises: HTTPException 404: Control not found @@ -697,48 +726,74 @@ async def delete_control( hint="Verify the control ID is correct and the control has been created.", ) - # Check for associations with policies - assoc_result = await db.execute( - select(policy_controls.c.policy_id).where( - policy_controls.c.control_id == control_id - ) - ) - associated_policy_ids = [row[0] for row in assoc_result.all()] - - if associated_policy_ids and not force: + # Check for associations with policies and direct agent links. + policy_assoc_query = select( + policy_controls.c.policy_id.label("policy_id"), + literal(None, type_=String).label("agent_name"), + ).where(policy_controls.c.control_id == control_id) + agent_assoc_query = select( + literal(None, type_=Integer).label("policy_id"), + agent_controls.c.agent_name.label("agent_name"), + ).where(agent_controls.c.control_id == control_id) + assoc_result = await db.execute(union_all(policy_assoc_query, agent_assoc_query)) + + associated_policy_ids: list[int] = [] + associated_agent_names: list[str] = [] + for policy_id, agent_name in assoc_result.all(): + if policy_id is not None: + associated_policy_ids.append(policy_id) + if agent_name is not None: + associated_agent_names.append(agent_name) + + if (associated_policy_ids or associated_agent_names) and not force: + errors = [ + ValidationErrorItem( + resource="Policy", + field="controls", + code="control_in_use", + message=f"Control is associated with policy ID {pid}", + value=pid, + ) + for pid in associated_policy_ids + ] + [ + ValidationErrorItem( + resource="Agent", + field="controls", + code="control_in_use", + message=f"Control is directly associated with agent '{agent_name}'", + value=agent_name, + ) + for agent_name in associated_agent_names + ] raise ConflictError( error_code=ErrorCode.CONTROL_IN_USE, detail=( f"Control '{control.name}' is associated with " - f"{len(associated_policy_ids)} policy/policies" + f"{len(associated_policy_ids)} policy/policies and " + f"{len(associated_agent_names)} agent(s)" ), resource="Control", resource_id=control.name, hint="Use force=true to dissociate and delete, or remove associations manually first.", - errors=[ - ValidationErrorItem( - resource="Policy", - field="controls", - code="control_in_use", - message=f"Control is associated with policy ID {pid}", - value=pid, - ) - for pid in associated_policy_ids - ], + errors=errors, ) - # Remove associations if force=true - dissociated_from: list[int] = [] + # Remove associations if force=true. + dissociated_from_policies: list[int] = [] + dissociated_from_agents: list[str] = [] if associated_policy_ids: - await db.execute( - delete(policy_controls).where( - policy_controls.c.control_id == control_id - ) - ) - dissociated_from = associated_policy_ids + await db.execute(delete(policy_controls).where(policy_controls.c.control_id == control_id)) + dissociated_from_policies = associated_policy_ids + if associated_agent_names: + await db.execute(delete(agent_controls).where(agent_controls.c.control_id == control_id)) + dissociated_from_agents = associated_agent_names + if dissociated_from_policies or dissociated_from_agents: _logger.info( - f"Dissociated control '{control.name}' ({control_id}) " - f"from {len(dissociated_from)} policy/policies" + "Dissociated control '%s' (%s) from %s policy/policies and %s agent(s)", + control.name, + control_id, + len(dissociated_from_policies), + len(dissociated_from_agents), ) # Delete the control @@ -758,7 +813,12 @@ async def delete_control( operation="delete", ) - return DeleteControlResponse(success=True, dissociated_from=dissociated_from) + return DeleteControlResponse( + success=True, + dissociated_from=dissociated_from_policies, + dissociated_from_policies=dissociated_from_policies, + dissociated_from_agents=dissociated_from_agents, + ) @router.patch( @@ -890,6 +950,16 @@ async def patch_control( try: await db.commit() _logger.info(f"Updated control '{control.name}' ({control_id})") + except IntegrityError: + await db.rollback() + conflicting_name = request.name or control.name + raise ConflictError( + error_code=ErrorCode.CONTROL_NAME_CONFLICT, + detail=f"Control with name '{conflicting_name}' already exists", + resource="Control", + resource_id=conflicting_name, + hint="Choose a different name or update the existing control.", + ) except Exception: await db.rollback() _logger.error( diff --git a/server/src/agent_control_server/models.py b/server/src/agent_control_server/models.py index 383c7681..1b148d1b 100644 --- a/server/src/agent_control_server/models.py +++ b/server/src/agent_control_server/models.py @@ -1,5 +1,5 @@ import datetime as dt -from typing import Any, Optional +from typing import Any from agent_control_models.agent import StepSchema, normalize_agent_name from agent_control_models.base import BaseModel @@ -38,13 +38,31 @@ class AgentData(BaseModel): Column("control_id", ForeignKey("controls.id"), primary_key=True, index=True), ) +# Association table for Agent <> Policy many-to-many relationship +agent_policies: Table = Table( + "agent_policies", + Base.metadata, + Column("agent_name", ForeignKey("agents.name"), primary_key=True, index=True), + Column("policy_id", ForeignKey("policies.id"), primary_key=True, index=True), +) + +# Association table for Agent <> Control many-to-many direct relationship +agent_controls: Table = Table( + "agent_controls", + Base.metadata, + Column("agent_name", ForeignKey("agents.name"), primary_key=True, index=True), + Column("control_id", ForeignKey("controls.id"), primary_key=True, index=True), +) + class Policy(Base): __tablename__ = "policies" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) name: Mapped[str] = mapped_column(String(255), nullable=False, unique=True) - agents: Mapped[list["Agent"]] = relationship("Agent", back_populates="policy") + agents: Mapped[list["Agent"]] = relationship( + "Agent", secondary=lambda: agent_policies, back_populates="policies" + ) # Many-to-many: Policy <> Control (direct relationship, no ControlSet layer) controls: Mapped[list["Control"]] = relationship( "Control", secondary=lambda: policy_controls, back_populates="policies" @@ -64,6 +82,10 @@ class Control(Base): policies: Mapped[list["Policy"]] = relationship( "Policy", secondary=lambda: policy_controls, back_populates="controls" ) + # Many-to-many backref: Control <> Agent (direct relationship) + agents: Mapped[list["Agent"]] = relationship( + "Agent", secondary=lambda: agent_controls, back_populates="controls" + ) class EvaluatorConfigDB(Base): @@ -97,10 +119,12 @@ class Agent(Base): data: Mapped[dict[str, Any]] = mapped_column( JSONB, server_default=text("'{}'::jsonb"), nullable=False ) - policy_id: Mapped[int | None] = mapped_column( - ForeignKey("policies.id"), nullable=True, index=True + policies: Mapped[list["Policy"]] = relationship( + "Policy", secondary=lambda: agent_policies, back_populates="agents" + ) + controls: Mapped[list["Control"]] = relationship( + "Control", secondary=lambda: agent_controls, back_populates="agents" ) - policy: Mapped[Optional["Policy"]] = relationship("Policy", back_populates="agents") created_at: Mapped[dt.datetime] = mapped_column( DateTime(), server_default=text("CURRENT_TIMESTAMP"), nullable=False, index=True ) diff --git a/server/src/agent_control_server/services/controls.py b/server/src/agent_control_server/services/controls.py index 46f9987d..15a05e9b 100644 --- a/server/src/agent_control_server/services/controls.py +++ b/server/src/agent_control_server/services/controls.py @@ -7,11 +7,11 @@ from agent_control_models.errors import ErrorCode, ValidationErrorItem from agent_control_models.policy import Control as APIControl from pydantic import ValidationError -from sqlalchemy import select +from sqlalchemy import select, union from sqlalchemy.ext.asyncio import AsyncSession from ..errors import APIValidationError -from ..models import Agent, Control, Policy, policy_controls +from ..models import Control, agent_controls, agent_policies, policy_controls _logger = logging.getLogger(__name__) @@ -33,19 +33,30 @@ async def list_controls_for_agent( *, allow_invalid_step_name_regex: bool = False, ) -> list[APIControl]: - """Return API Control models for all configured controls associated with the agent's policy. + """Return API Control models for controls associated with the agent. - Traversal: Agent -> Policy -> Controls (direct relationship). - Uses explicit joins over association table to avoid async relationship loading. + Active controls are the de-duplicated union of: + - controls inherited from all assigned policies + - controls directly associated with the agent Note: Invalid ControlDefinition data triggers an APIValidationError. """ - stmt = ( - select(Control) - .join(policy_controls, Control.id == policy_controls.c.control_id) - .join(Policy, policy_controls.c.policy_id == Policy.id) - .join(Agent, Policy.id == Agent.policy_id) - .where(Agent.name == agent_name) + policy_control_ids = ( + select(policy_controls.c.control_id.label("control_id")) + .select_from( + policy_controls.join( + agent_policies, policy_controls.c.policy_id == agent_policies.c.policy_id + ) + ) + .where(agent_policies.c.agent_name == agent_name) + ) + direct_control_ids = select(agent_controls.c.control_id.label("control_id")).where( + agent_controls.c.agent_name == agent_name + ) + control_ids_subquery = union(policy_control_ids, direct_control_ids).subquery() + + stmt = select(Control).join( + control_ids_subquery, Control.id == control_ids_subquery.c.control_id ) result = await db.execute(stmt) diff --git a/server/src/agent_control_server/services/query_utils.py b/server/src/agent_control_server/services/query_utils.py new file mode 100644 index 00000000..7e8e5376 --- /dev/null +++ b/server/src/agent_control_server/services/query_utils.py @@ -0,0 +1,3 @@ +def escape_like_pattern(value: str) -> str: + """Escape special SQL LIKE pattern characters in user-provided input.""" + return value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_") diff --git a/server/tests/test_agents_additional.py b/server/tests/test_agents_additional.py index 2297e9a5..8c8085f8 100644 --- a/server/tests/test_agents_additional.py +++ b/server/tests/test_agents_additional.py @@ -14,12 +14,11 @@ def _init_agent( client: TestClient, *, - agent_id: str | None = None, agent_name: str | None = None, steps: list[dict] | None = None, evaluators: list[dict] | None = None, ) -> tuple[str, str]: - name = (agent_name or agent_id or f"agent-{uuid.uuid4().hex[:12]}").lower() + name = (agent_name or f"agent-{uuid.uuid4().hex[:12]}").lower() if len(name) < 10: name = f"{name}-agent".replace("--", "-") payload = { @@ -58,10 +57,10 @@ def test_list_agent_evaluators_pagination_and_get(client: TestClient) -> None: {"name": "eval-b", "description": "b", "config_schema": {"type": "object"}}, {"name": "eval-c", "description": "c", "config_schema": {}}, ] - agent_id, _ = _init_agent(client, evaluators=evaluators) + agent_name, _ = _init_agent(client, evaluators=evaluators) # When: listing with pagination - resp = client.get(f"/api/v1/agents/{agent_id}/evaluators", params={"limit": 2}) + resp = client.get(f"/api/v1/agents/{agent_name}/evaluators", params={"limit": 2}) assert resp.status_code == 200 body = resp.json() # Then: first page returns two items and a next cursor @@ -71,7 +70,7 @@ def test_list_agent_evaluators_pagination_and_get(client: TestClient) -> None: # When: fetching next page using cursor resp2 = client.get( - f"/api/v1/agents/{agent_id}/evaluators", + f"/api/v1/agents/{agent_name}/evaluators", params={"limit": 2, "cursor": body["pagination"]["next_cursor"]}, ) assert resp2.status_code == 200 @@ -81,7 +80,7 @@ def test_list_agent_evaluators_pagination_and_get(client: TestClient) -> None: assert [e["name"] for e in body2["evaluators"]] == ["eval-c"] # When: getting a specific evaluator - get_resp = client.get(f"/api/v1/agents/{agent_id}/evaluators/eval-b") + get_resp = client.get(f"/api/v1/agents/{agent_name}/evaluators/eval-b") assert get_resp.status_code == 200 evaluator = get_resp.json() # Then: evaluator details are returned @@ -95,16 +94,16 @@ def test_list_agent_evaluators_invalid_cursor_returns_first_page(client: TestCli {"name": "eval-a", "description": "a", "config_schema": {}}, {"name": "eval-b", "description": "b", "config_schema": {}}, ] - agent_id, _ = _init_agent(client, evaluators=evaluators) + agent_name, _ = _init_agent(client, evaluators=evaluators) # When: listing without cursor - resp = client.get(f"/api/v1/agents/{agent_id}/evaluators", params={"limit": 1}) + resp = client.get(f"/api/v1/agents/{agent_name}/evaluators", params={"limit": 1}) assert resp.status_code == 200 base = resp.json() # When: listing with an invalid cursor resp2 = client.get( - f"/api/v1/agents/{agent_id}/evaluators", + f"/api/v1/agents/{agent_name}/evaluators", params={"limit": 1, "cursor": "does-not-exist"}, ) assert resp2.status_code == 200 @@ -123,12 +122,12 @@ def test_init_agent_preserves_existing_steps_when_missing_from_payload( {"type": "tool", "name": "tool-a", "input_schema": {}, "output_schema": {}}, {"type": "tool", "name": "tool-b", "input_schema": {}, "output_schema": {}}, ] - agent_id, agent_name = _init_agent(client, steps=steps) + agent_name, agent_name = _init_agent(client, steps=steps) # When: re-initializing with only one of the steps payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -140,7 +139,7 @@ def test_init_agent_preserves_existing_steps_when_missing_from_payload( assert resp.status_code == 200 # Then: the missing step is preserved (initAgent only adds) - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") assert get_resp.status_code == 200 step_names = {step["name"] for step in get_resp.json()["steps"]} assert step_names == {"tool-a", "tool-b"} @@ -148,10 +147,10 @@ def test_init_agent_preserves_existing_steps_when_missing_from_payload( def test_get_agent_evaluator_not_found(client: TestClient) -> None: # Given: an existing agent with no matching evaluator - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) # When: requesting a missing evaluator - resp = client.get(f"/api/v1/agents/{agent_id}/evaluators/missing") + resp = client.get(f"/api/v1/agents/{agent_name}/evaluators/missing") # Then: 404 not found assert resp.status_code == 404 @@ -180,11 +179,11 @@ def test_patch_agent_remove_steps_and_evaluators(client: TestClient) -> None: {"name": "eval-a", "description": "a", "config_schema": {}}, {"name": "eval-b", "description": "b", "config_schema": {}}, ] - agent_id, _ = _init_agent(client, steps=steps, evaluators=evaluators) + agent_name, _ = _init_agent(client, steps=steps, evaluators=evaluators) # When: removing one step and one evaluator resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={ "remove_steps": [{"type": "tool", "name": "tool-a"}], "remove_evaluators": ["eval-b"], @@ -198,7 +197,7 @@ def test_patch_agent_remove_steps_and_evaluators(client: TestClient) -> None: assert body["evaluators_removed"] == ["eval-b"] # Then: agent data reflects removal - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") assert get_resp.status_code == 200 data = get_resp.json() assert {s["name"] for s in data["steps"]} == {"tool-b"} @@ -218,7 +217,7 @@ def test_patch_agent_remove_evaluator_in_use_conflict(client: TestClient) -> Non }, } ] - agent_id, agent_name = _init_agent(client, evaluators=evaluators) + agent_name, agent_name = _init_agent(client, evaluators=evaluators) control_payload = deepcopy(VALID_CONTROL_PAYLOAD) control_payload["evaluator"] = { @@ -230,12 +229,12 @@ def test_patch_agent_remove_evaluator_in_use_conflict(client: TestClient) -> Non policy_id = _create_policy(client) assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") assert assoc.status_code == 200 - assign = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign.status_code == 200 # When: attempting to remove evaluator in use resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_evaluators": ["custom"]}, ) @@ -281,7 +280,7 @@ def test_init_agent_rejects_builtin_evaluator_name(client: TestClient) -> None: # Given: a payload that registers an evaluator matching a built-in name payload = { "agent": { - "agent_id": str(uuid.uuid4()), + "agent_name": str(uuid.uuid4()), "agent_name": f"agent-{uuid.uuid4().hex[:12]}", "agent_description": "desc", "agent_version": "1.0", @@ -346,14 +345,14 @@ def test_list_agent_controls_corrupted_control_data_returns_422( client: TestClient, ) -> None: # Given: an agent with a policy that includes a control - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) control_payload = deepcopy(VALID_CONTROL_PAYLOAD) control_payload["evaluator"] = {"name": "regex", "config": {"pattern": "x"}} control_id = _create_control_with_data(client, control_payload) policy_id = _create_policy(client) assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") assert assoc.status_code == 200 - assign = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign.status_code == 200 # And: the control data is corrupted in the DB @@ -364,7 +363,7 @@ def test_list_agent_controls_corrupted_control_data_returns_422( ) # When: listing agent controls - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") # Then: corrupted data error is returned assert resp.status_code == 422 @@ -393,15 +392,15 @@ def test_list_agents_invalid_cursor_returns_first_page(client: TestClient) -> No def test_list_agent_evaluators_corrupted_data_returns_empty(client: TestClient) -> None: # Given: an agent with corrupted stored data - agent_id, _ = _init_agent(client, evaluators=[{"name": "eval-a", "config_schema": {}}]) + agent_name, _ = _init_agent(client, evaluators=[{"name": "eval-a", "config_schema": {}}]) with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": "{\"bad\": \"data\"}", "id": agent_id}, + {"data": "{\"bad\": \"data\"}", "id": agent_name}, ) # When: listing evaluator schemas - resp = client.get(f"/api/v1/agents/{agent_id}/evaluators") + resp = client.get(f"/api/v1/agents/{agent_name}/evaluators") # Then: empty list is returned assert resp.status_code == 200 @@ -412,7 +411,7 @@ def test_list_agent_evaluators_corrupted_data_returns_empty(client: TestClient) def test_set_agent_policy_rejects_corrupted_agent_data(client: TestClient) -> None: # Given: an agent with corrupted stored data and a policy with a control - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) policy_id = _create_policy(client) control_id = _create_control_with_data(client, VALID_CONTROL_PAYLOAD) assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") @@ -421,11 +420,11 @@ def test_set_agent_policy_rejects_corrupted_agent_data(client: TestClient) -> No with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": json.dumps({"bad": "data"}), "id": agent_id}, + {"data": json.dumps({"bad": "data"}), "id": agent_name}, ) # When: assigning policy to the agent - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: incompatible controls error is returned assert resp.status_code == 400 @@ -436,7 +435,7 @@ def test_set_agent_policy_rejects_corrupted_agent_data(client: TestClient) -> No def test_set_agent_policy_rejects_missing_agent_evaluator(client: TestClient) -> None: # Given: an agent with no evaluators and a control referencing a missing evaluator - agent_id, agent_name = _init_agent(client) + agent_name, agent_name = _init_agent(client) policy_id = _create_policy(client) control_id = _create_control_with_data(client, VALID_CONTROL_PAYLOAD) assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") @@ -454,7 +453,7 @@ def test_set_agent_policy_rejects_missing_agent_evaluator(client: TestClient) -> ) # When: assigning policy to the agent - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: incompatible controls error is returned assert resp.status_code == 400 @@ -465,7 +464,7 @@ def test_set_agent_policy_rejects_missing_agent_evaluator(client: TestClient) -> def test_set_agent_policy_rejects_invalid_agent_evaluator_config(client: TestClient) -> None: # Given: an agent with an evaluator schema requiring \"pattern\" - agent_id, agent_name = _init_agent( + agent_name, agent_name = _init_agent( client, evaluators=[ { @@ -496,7 +495,7 @@ def test_set_agent_policy_rejects_invalid_agent_evaluator_config(client: TestCli ) # When: assigning policy to the agent - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: incompatible controls error is returned assert resp.status_code == 400 @@ -531,10 +530,10 @@ def test_delete_agent_policy_agent_not_found(client: TestClient) -> None: def test_delete_agent_policy_no_policy_assigned_returns_404(client: TestClient) -> None: # Given: an agent with no policy assigned - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) # When: deleting policy - resp = client.delete(f"/api/v1/agents/{agent_id}/policy") + resp = client.delete(f"/api/v1/agents/{agent_name}/policy") # Then: policy not found error is returned assert resp.status_code == 404 @@ -543,7 +542,7 @@ def test_delete_agent_policy_no_policy_assigned_returns_404(client: TestClient) def test_list_agents_corrupted_data_sets_zero_counts(client: TestClient) -> None: # Given: an agent with corrupted data stored in the DB - agent_id, _ = _init_agent( + agent_name, _ = _init_agent( client, steps=[{"type": "tool", "name": "tool-a", "input_schema": {}, "output_schema": {}}], evaluators=[{"name": "eval-a", "config_schema": {}}], @@ -551,7 +550,7 @@ def test_list_agents_corrupted_data_sets_zero_counts(client: TestClient) -> None with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": json.dumps({"bad": "data"}), "id": agent_id}, + {"data": json.dumps({"bad": "data"}), "id": agent_name}, ) # When: listing agents @@ -560,22 +559,22 @@ def test_list_agents_corrupted_data_sets_zero_counts(client: TestClient) -> None # Then: step/evaluator counts are zeroed for corrupted data assert resp.status_code == 200 agents = {a["agent_name"]: a for a in resp.json()["agents"]} - agent = agents[agent_id] + agent = agents[agent_name] assert agent["step_count"] == 0 assert agent["evaluator_count"] == 0 def test_get_agent_corrupted_data_returns_422(client: TestClient) -> None: # Given: an agent with corrupted stored data - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": json.dumps({"bad": "data"}), "id": agent_id}, + {"data": json.dumps({"bad": "data"}), "id": agent_name}, ) # When: fetching the agent - resp = client.get(f"/api/v1/agents/{agent_id}") + resp = client.get(f"/api/v1/agents/{agent_name}") # Then: corrupted data error is returned assert resp.status_code == 422 @@ -584,16 +583,16 @@ def test_get_agent_corrupted_data_returns_422(client: TestClient) -> None: def test_get_agent_corrupted_metadata_returns_422(client: TestClient) -> None: # Given: an agent with invalid agent_metadata payload - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) corrupted = {"agent_metadata": {}, "steps": [], "evaluators": []} with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": json.dumps(corrupted), "id": agent_id}, + {"data": json.dumps(corrupted), "id": agent_name}, ) # When: fetching the agent - resp = client.get(f"/api/v1/agents/{agent_id}") + resp = client.get(f"/api/v1/agents/{agent_name}") # Then: corrupted metadata error is returned assert resp.status_code == 422 @@ -602,9 +601,9 @@ def test_get_agent_corrupted_metadata_returns_422(client: TestClient) -> None: def test_get_agent_policy_missing_policy_returns_404(client: TestClient) -> None: # Given: an agent assigned to a policy that cannot be found - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) policy_id = _create_policy(client) - assign = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign.status_code == 200 from agent_control_server.db import get_async_db @@ -619,7 +618,7 @@ def test_get_agent_policy_missing_policy_returns_404(client: TestClient) -> None with Session(engine) as session: agent_row = ( session.execute( - select(AgentModel).where(AgentModel.name == agent_id) + select(AgentModel).where(AgentModel.name == agent_name) ) .scalars() .first() @@ -640,7 +639,7 @@ async def mock_db_missing_policy() -> AsyncGenerator[AsyncSession, None]: # When: retrieving the agent policy and policy lookup returns None app.dependency_overrides[get_async_db] = mock_db_missing_policy try: - resp = client.get(f"/api/v1/agents/{agent_id}/policy") + resp = client.get(f"/api/v1/agents/{agent_name}/policy") finally: app.dependency_overrides.clear() @@ -651,7 +650,7 @@ async def mock_db_missing_policy() -> AsyncGenerator[AsyncSession, None]: def test_set_agent_policy_skips_controls_without_data(client: TestClient) -> None: # Given: an agent and a policy with a control that has no data configured - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) policy_id = _create_policy(client) control_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) assert control_resp.status_code == 200 @@ -660,7 +659,7 @@ def test_set_agent_policy_skips_controls_without_data(client: TestClient) -> Non assert assoc.status_code == 200 # When: assigning the policy to the agent - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: assignment succeeds because empty data is ignored during validation assert resp.status_code == 200 @@ -669,7 +668,7 @@ def test_set_agent_policy_skips_controls_without_data(client: TestClient) -> Non def test_set_agent_policy_skips_controls_without_evaluator_name(client: TestClient) -> None: # Given: an agent and a policy with a control missing evaluator name - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) policy_id = _create_policy(client) control_id = _create_control_with_data(client, VALID_CONTROL_PAYLOAD) assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") @@ -682,7 +681,7 @@ def test_set_agent_policy_skips_controls_without_evaluator_name(client: TestClie ) # When: assigning the policy to the agent - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: assignment succeeds because evaluator name is missing assert resp.status_code == 200 @@ -691,7 +690,7 @@ def test_set_agent_policy_skips_controls_without_evaluator_name(client: TestClie def test_list_agents_includes_active_controls_count(client: TestClient) -> None: # Given: an agent assigned to a policy with two controls - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) policy_id = _create_policy(client) control_ids = [ _create_control_with_data(client, VALID_CONTROL_PAYLOAD), @@ -700,7 +699,7 @@ def test_list_agents_includes_active_controls_count(client: TestClient) -> None: for control_id in control_ids: assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") assert assoc.status_code == 200 - assign = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign.status_code == 200 # When: listing agents @@ -735,10 +734,9 @@ def test_list_agents_valid_cursor_not_found_returns_first_page(client: TestClien def test_init_agent_adds_new_evaluator(client: TestClient) -> None: # Given: an existing agent with one evaluator agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -754,7 +752,7 @@ def test_init_agent_adds_new_evaluator(client: TestClient) -> None: "/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -766,7 +764,7 @@ def test_init_agent_adds_new_evaluator(client: TestClient) -> None: # Then: both evaluators are present assert resp2.status_code == 200 - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") names = {e["name"] for e in get_resp.json()["evaluators"]} assert names == {"eval-a", "eval-b"} @@ -774,12 +772,11 @@ def test_init_agent_adds_new_evaluator(client: TestClient) -> None: def test_init_agent_returns_controls_when_policy_assigned(client: TestClient) -> None: # Given: an agent assigned to a policy with a control agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name init_resp = client.post( "/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -794,7 +791,7 @@ def test_init_agent_returns_controls_when_policy_assigned(client: TestClient) -> control_id = _create_control_with_data(client, VALID_CONTROL_PAYLOAD) assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") assert assoc.status_code == 200 - assign = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign.status_code == 200 # When: re-initializing the agent with the same UUID @@ -802,7 +799,7 @@ def test_init_agent_returns_controls_when_policy_assigned(client: TestClient) -> "/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -821,16 +818,16 @@ def test_init_agent_returns_controls_when_policy_assigned(client: TestClient) -> def test_patch_agent_corrupted_data_returns_422(client: TestClient) -> None: # Given: an agent with corrupted stored data - agent_id, _ = _init_agent(client) + agent_name, _ = _init_agent(client) with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": json.dumps({"bad": "data"}), "id": agent_id}, + {"data": json.dumps({"bad": "data"}), "id": agent_name}, ) # When: patching the agent resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_steps": [{"type": "tool", "name": "tool-a"}]}, ) @@ -841,15 +838,15 @@ def test_patch_agent_corrupted_data_returns_422(client: TestClient) -> None: def test_get_agent_evaluator_corrupted_data_returns_404(client: TestClient) -> None: # Given: an agent with evaluator data that becomes corrupted - agent_id, _ = _init_agent(client, evaluators=[{"name": "eval-a", "config_schema": {}}]) + agent_name, _ = _init_agent(client, evaluators=[{"name": "eval-a", "config_schema": {}}]) with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": json.dumps({"bad": "data"}), "id": agent_id}, + {"data": json.dumps({"bad": "data"}), "id": agent_name}, ) # When: fetching a specific evaluator - resp = client.get(f"/api/v1/agents/{agent_id}/evaluators/eval-a") + resp = client.get(f"/api/v1/agents/{agent_name}/evaluators/eval-a") # Then: evaluator not found is returned due to corrupted data assert resp.status_code == 404 @@ -862,7 +859,7 @@ def test_init_agent_rejects_duplicate_step_names_in_single_request( # Given: a payload with duplicate step names payload = { "agent": { - "agent_id": str(uuid.uuid4()), + "agent_name": str(uuid.uuid4()), "agent_name": f"agent-{uuid.uuid4().hex[:12]}", "agent_description": "desc", "agent_version": "1.0", @@ -892,10 +889,9 @@ def test_init_agent_rejects_step_schema_conflict_across_registrations( ) -> None: # Given: an agent registered with a step agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name original_payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -916,7 +912,7 @@ def test_init_agent_rejects_step_schema_conflict_across_registrations( # When: re-registering with same step name but different schema conflicting_payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -946,10 +942,9 @@ def test_init_agent_accepts_identical_step_schema_across_registrations( ) -> None: # Given: an agent registered with a step agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "desc", "agent_version": "1.0", @@ -975,7 +970,7 @@ def test_init_agent_accepts_identical_step_schema_across_registrations( assert resp2.json()["created"] is False # Agent already exists # And: step is preserved - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") assert get_resp.status_code == 200 steps = get_resp.json()["steps"] assert len(steps) == 1 diff --git a/server/tests/test_controls_additional.py b/server/tests/test_controls_additional.py index 4be88a31..21a58f5f 100644 --- a/server/tests/test_controls_additional.py +++ b/server/tests/test_controls_additional.py @@ -2,19 +2,24 @@ import json import uuid +from collections.abc import AsyncGenerator from copy import deepcopy from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock import pytest from fastapi.testclient import TestClient from sqlalchemy import text +from sqlalchemy.exc import IntegrityError +from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session +from agent_control_server.db import get_async_db from agent_control_server.models import Control from agent_control_evaluators import RegexEvaluatorConfig from agent_control_server.endpoints import controls as controls_module -from agent_control_server.models import Control +from agent_control_server.main import app from .conftest import engine from .utils import VALID_CONTROL_PAYLOAD @@ -32,6 +37,71 @@ def _set_control_data(client: TestClient, control_id: int, data: dict) -> None: assert resp.status_code == 200, resp.text +def test_create_control_integrity_error_returns_conflict(client: TestClient) -> None: + """DB uniqueness violations during create should be surfaced as 409 conflicts.""" + + async def mock_db_integrity_error() -> AsyncGenerator[AsyncSession, None]: + mock_session = AsyncMock(spec=AsyncSession) + existing_result = MagicMock() + existing_result.first.return_value = None + + mock_session.execute = AsyncMock(return_value=existing_result) + mock_session.add = MagicMock() + mock_session.refresh = AsyncMock() + mock_session.commit = AsyncMock( + side_effect=IntegrityError( + "INSERT INTO controls ...", + {"name": "duplicate-control"}, + Exception("duplicate key value violates unique constraint"), + ) + ) + yield mock_session + + app.dependency_overrides[get_async_db] = mock_db_integrity_error + try: + resp = client.put("/api/v1/controls", json={"name": "duplicate-control"}) + finally: + app.dependency_overrides.clear() + + assert resp.status_code == 409 + assert resp.json()["error_code"] == "CONTROL_NAME_CONFLICT" + + +def test_patch_control_rename_integrity_error_returns_conflict(client: TestClient) -> None: + """DB uniqueness violations during rename should be surfaced as 409 conflicts.""" + control_obj = SimpleNamespace(id=1, name="old-control", data={}) + + async def mock_db_integrity_error() -> AsyncGenerator[AsyncSession, None]: + mock_session = AsyncMock(spec=AsyncSession) + + control_lookup_result = MagicMock() + control_lookup_result.scalars.return_value.first.return_value = control_obj + + name_lookup_result = MagicMock() + name_lookup_result.first.return_value = None + + mock_session.execute = AsyncMock( + side_effect=[control_lookup_result, name_lookup_result] + ) + mock_session.commit = AsyncMock( + side_effect=IntegrityError( + "UPDATE controls ...", + {"name": "existing-control"}, + Exception("duplicate key value violates unique constraint"), + ) + ) + yield mock_session + + app.dependency_overrides[get_async_db] = mock_db_integrity_error + try: + resp = client.patch("/api/v1/controls/1", json={"name": "existing-control"}) + finally: + app.dependency_overrides.clear() + + assert resp.status_code == 409 + assert resp.json()["error_code"] == "CONTROL_NAME_CONFLICT" + + def test_list_controls_filters_and_pagination(client: TestClient) -> None: # Given: three controls with varying data control1_id, control1_name = _create_control(client, name=f"AlphaControl-{uuid.uuid4()}") @@ -425,6 +495,32 @@ def test_delete_control_force_dissociates(client: TestClient) -> None: assert control_id not in list_resp.json()["control_ids"] +def test_delete_control_force_dissociates_direct_agent_links(client: TestClient) -> None: + # Given: a control directly associated with an agent + control_id, _ = _create_control(client) + _set_control_data(client, control_id, deepcopy(VALID_CONTROL_PAYLOAD)) + + agent_name = f"agent-{uuid.uuid4().hex[:12]}" + init_resp = client.post( + "/api/v1/agents/initAgent", + json={"agent": {"agent_name": agent_name}, "steps": []}, + ) + assert init_resp.status_code == 200 + + assoc_resp = client.post(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert assoc_resp.status_code == 200 + + # When: force-deleting the control + resp = client.delete(f"/api/v1/controls/{control_id}?force=true") + assert resp.status_code == 200 + body = resp.json() + + # Then: direct agent dissociation details are returned + assert body["success"] is True + assert body.get("dissociated_from_policies", []) == [] + assert body.get("dissociated_from_agents", []) == [agent_name] + + def test_get_control_corrupted_data_returns_none(client: TestClient) -> None: # Given: a control with corrupted data in DB control_id, control_name = _create_control(client) @@ -499,11 +595,11 @@ def test_set_control_data_agent_scoped_agent_not_found(client: TestClient) -> No def test_set_control_data_agent_scoped_evaluator_missing(client: TestClient) -> None: # Given: an agent without the referenced evaluator agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name resp = client.post( "/api/v1/agents/initAgent", json={ - "agent": {"agent_id": agent_id, "agent_name": agent_name}, + "agent": {"agent_name": agent_name, "agent_name": agent_name}, "steps": [], "evaluators": [], }, @@ -527,11 +623,11 @@ def test_set_control_data_agent_scoped_evaluator_missing(client: TestClient) -> def test_set_control_data_agent_scoped_invalid_schema(client: TestClient) -> None: # Given: an agent with evaluator schema requiring "pattern" agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name resp = client.post( "/api/v1/agents/initAgent", json={ - "agent": {"agent_id": agent_id, "agent_name": agent_name}, + "agent": {"agent_name": agent_name, "agent_name": agent_name}, "steps": [], "evaluators": [ { @@ -617,11 +713,11 @@ def test_set_control_data_agent_scoped_corrupted_agent_data_returns_422( ) -> None: # Given: an agent whose stored data is corrupted agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name resp = client.post( "/api/v1/agents/initAgent", json={ - "agent": {"agent_id": agent_id, "agent_name": agent_name}, + "agent": {"agent_name": agent_name, "agent_name": agent_name}, "steps": [], "evaluators": [{"name": "custom", "config_schema": {"type": "object"}}], }, @@ -631,7 +727,7 @@ def test_set_control_data_agent_scoped_corrupted_agent_data_returns_422( with engine.begin() as conn: conn.execute( text("UPDATE agents SET data = CAST(:data AS JSONB) WHERE name = :id"), - {"data": json.dumps({"bad": "data"}), "id": agent_id}, + {"data": json.dumps({"bad": "data"}), "id": agent_name}, ) control_id, _ = _create_control(client) diff --git a/server/tests/test_error_handling.py b/server/tests/test_error_handling.py index 756827c8..5f139431 100644 --- a/server/tests/test_error_handling.py +++ b/server/tests/test_error_handling.py @@ -32,10 +32,10 @@ def test_init_agent_rollback_on_create_failure( """Test that init_agent rolls back transaction when commit fails on create.""" # Given: a valid agent init payload agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "test", "agent_version": "1.0", @@ -63,7 +63,7 @@ def test_delete_agent_policy_rollback_on_failure( # Given: an agent with an assigned policy agent_payload = { "agent": { - "agent_id": str(uuid.uuid4()), + "agent_name": str(uuid.uuid4()), "agent_name": f"test-agent-{uuid.uuid4()}", "agent_description": "test", "agent_version": "1.0", @@ -73,14 +73,14 @@ def test_delete_agent_policy_rollback_on_failure( } r1 = client.post("/api/v1/agents/initAgent", json=agent_payload) assert r1.status_code == 200 - agent_id = agent_payload["agent"]["agent_name"] + agent_name = agent_payload["agent"]["agent_name"] policy_name = f"test-policy-{uuid.uuid4()}" r2 = client.put("/api/v1/policies", json={"name": policy_name}) assert r2.status_code == 200 policy_id = r2.json()["policy_id"] - assign_resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign_resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign_resp.status_code == 200 # And: a database session that fails on commit @@ -89,7 +89,7 @@ def test_delete_agent_policy_rollback_on_failure( with Session(db_engine) as session: existing_agent = ( - session.query(Agent).filter(Agent.name == agent_id).first() + session.query(Agent).filter(Agent.name == agent_name).first() ) assert existing_agent is not None @@ -106,7 +106,7 @@ async def mock_db_for_delete_policy() -> AsyncGenerator[AsyncSession, None]: # When: deleting policy and commit fails app.dependency_overrides[get_async_db] = mock_db_for_delete_policy try: - resp = client.delete(f"/api/v1/agents/{agent_id}/policy") + resp = client.delete(f"/api/v1/agents/{agent_name}/policy") finally: app.dependency_overrides.clear() @@ -121,10 +121,10 @@ def test_init_agent_rollback_on_update_failure( """Test that init_agent rolls back transaction when commit fails on update.""" # Given: an existing agent agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "test", "agent_version": "1.0", @@ -221,10 +221,10 @@ def test_patch_agent_rollback_on_failure( """Test that patch_agent rolls back when commit fails.""" # Given: an existing agent with a step to remove agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name payload = { "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "test", "agent_version": "1.0", @@ -248,7 +248,7 @@ def test_patch_agent_rollback_on_failure( with Session(db_engine) as session: existing_agent = ( - session.query(Agent).filter(Agent.name == agent_id).first() + session.query(Agent).filter(Agent.name == agent_name).first() ) assert existing_agent is not None @@ -266,7 +266,7 @@ async def mock_db_for_patch_agent() -> AsyncGenerator[AsyncSession, None]: app.dependency_overrides[get_async_db] = mock_db_for_patch_agent try: resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_steps": [{"type": "tool", "name": "tool_a"}]}, ) finally: @@ -349,7 +349,7 @@ def test_set_agent_policy_rollback_on_failure( # Given: an existing agent and policy agent_payload = { "agent": { - "agent_id": str(uuid.uuid4()), + "agent_name": str(uuid.uuid4()), "agent_name": f"test-agent-{uuid.uuid4()}", "agent_description": "test", "agent_version": "1.0", @@ -359,7 +359,7 @@ def test_set_agent_policy_rollback_on_failure( } r1 = client.post("/api/v1/agents/initAgent", json=agent_payload) assert r1.status_code == 200 - agent_id = agent_payload["agent"]["agent_name"] + agent_name = agent_payload["agent"]["agent_name"] policy_name = f"test-policy-{uuid.uuid4()}" r2 = client.put("/api/v1/policies", json={"name": policy_name}) @@ -373,7 +373,7 @@ def test_set_agent_policy_rollback_on_failure( with Session(db_engine) as session: existing_agent = ( session.query(Agent) - .filter(Agent.name == agent_id) + .filter(Agent.name == agent_name) .first() ) existing_policy = ( @@ -406,17 +406,27 @@ async def mock_db_for_policy_assignment() -> AsyncGenerator[AsyncSession, None]: mock_controls_result = MagicMock() mock_controls_result.scalars.return_value.unique.return_value.all.return_value = [] + # Mock existing policy associations query + mock_existing_policy_assoc_result = MagicMock() + mock_existing_policy_assoc_result.all.return_value = [] + + # Mock delete/insert association statements + mock_mutation_result = MagicMock() + # Return different results for different queries mock_session.execute = AsyncMock(side_effect=[ mock_agent_result, mock_policy_result, mock_controls_result, + mock_existing_policy_assoc_result, + mock_mutation_result, + mock_mutation_result, ]) yield mock_session app.dependency_overrides[get_async_db] = mock_db_for_policy_assignment try: - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: rollback is called and 500 error is returned assert resp.status_code == 500 diff --git a/server/tests/test_evaluator_schemas.py b/server/tests/test_evaluator_schemas.py index 785f41bf..bb952c06 100644 --- a/server/tests/test_evaluator_schemas.py +++ b/server/tests/test_evaluator_schemas.py @@ -9,13 +9,13 @@ def make_agent_payload( - agent_id: str | None = None, + agent_name: str | None = None, name: str | None = None, evaluators: list | None = None, ): """Helper to create agent payload with evaluators.""" - if agent_id is not None: - name = agent_id + if agent_name is not None: + name = agent_name elif name is None: name = f"agent-{uuid.uuid4().hex[:12]}" canonical_name = name.lower().replace(" ", "-") @@ -23,7 +23,7 @@ def make_agent_payload( canonical_name = f"{canonical_name}-agent".replace("--", "-") return { "agent": { - "agent_id": canonical_name, + "agent_name": canonical_name, "agent_name": canonical_name, "agent_description": "desc", "agent_version": "1.0", @@ -95,10 +95,10 @@ def test_init_agent_evaluator_name_collision_list(client: TestClient) -> None: def test_init_agent_update_evaluator_compatible_schema(client: TestClient) -> None: """Test updating evaluator with compatible schema change (add optional field).""" # Given: Agent with evaluator - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload1 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -116,7 +116,7 @@ def test_init_agent_update_evaluator_compatible_schema(client: TestClient) -> No # When: Updating with compatible schema (add optional field) payload2 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -142,10 +142,10 @@ def test_init_agent_update_evaluator_incompatible_schema_rejected( ) -> None: """Test that incompatible schema change is rejected.""" # Given: Agent with evaluator - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload1 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -162,7 +162,7 @@ def test_init_agent_update_evaluator_incompatible_schema_rejected( # When: Updating with incompatible schema (remove property) payload2 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -183,10 +183,10 @@ def test_init_agent_update_evaluator_incompatible_schema_rejected( def test_init_agent_update_evaluator_type_change_rejected(client: TestClient) -> None: """Test that changing property type is rejected.""" # Given: Agent with evaluator - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload1 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -202,7 +202,7 @@ def test_init_agent_update_evaluator_type_change_rejected(client: TestClient) -> # When: Changing property type payload2 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -223,10 +223,10 @@ def test_init_agent_update_evaluator_type_change_rejected(client: TestClient) -> def test_init_agent_add_required_property_rejected(client: TestClient) -> None: """Test that adding a new required property is rejected.""" # Given: Agent with evaluator - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload1 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -242,7 +242,7 @@ def test_init_agent_add_required_property_rejected(client: TestClient) -> None: # When: Adding new required property payload2 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -280,10 +280,10 @@ def test_list_agent_evaluators(client: TestClient) -> None: ) resp = client.post("/api/v1/agents/initAgent", json=payload) assert resp.status_code == 200 - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] # When: Listing evaluators - list_resp = client.get(f"/api/v1/agents/{agent_id}/evaluators") + list_resp = client.get(f"/api/v1/agents/{agent_name}/evaluators") # Then: Should return both evaluators assert list_resp.status_code == 200 data = list_resp.json() @@ -301,10 +301,10 @@ def test_list_agent_evaluators_pagination(client: TestClient) -> None: ) resp = client.post("/api/v1/agents/initAgent", json=payload) assert resp.status_code == 200 - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] # When: Fetching first page - resp1 = client.get(f"/api/v1/agents/{agent_id}/evaluators?offset=0&limit=2") + resp1 = client.get(f"/api/v1/agents/{agent_name}/evaluators?offset=0&limit=2") # Then: Should return 2 items with total=5 assert resp1.status_code == 200 data1 = resp1.json() @@ -312,7 +312,7 @@ def test_list_agent_evaluators_pagination(client: TestClient) -> None: assert data1["pagination"]["total"] == 5 # When: Fetching second page - resp2 = client.get(f"/api/v1/agents/{agent_id}/evaluators?offset=2&limit=2") + resp2 = client.get(f"/api/v1/agents/{agent_name}/evaluators?offset=2&limit=2") # Then: Should return 2 more items assert resp2.status_code == 200 data2 = resp2.json() @@ -333,10 +333,10 @@ def test_get_agent_evaluator_by_name(client: TestClient) -> None: ) resp = client.post("/api/v1/agents/initAgent", json=payload) assert resp.status_code == 200 - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] # When: Getting evaluator by name - get_resp = client.get(f"/api/v1/agents/{agent_id}/evaluators/my-eval") + get_resp = client.get(f"/api/v1/agents/{agent_name}/evaluators/my-eval") # Then: Should return evaluator details assert get_resp.status_code == 200 data = get_resp.json() @@ -350,10 +350,10 @@ def test_get_agent_evaluator_not_found(client: TestClient) -> None: payload = make_agent_payload(evaluators=[]) resp = client.post("/api/v1/agents/initAgent", json=payload) assert resp.status_code == 200 - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] # When: Getting nonexistent evaluator - get_resp = client.get(f"/api/v1/agents/{agent_id}/evaluators/nonexistent") + get_resp = client.get(f"/api/v1/agents/{agent_name}/evaluators/nonexistent") # Then: Should return 404 assert get_resp.status_code == 404 diff --git a/server/tests/test_init_agent.py b/server/tests/test_init_agent.py index 2bc46e5e..7140e736 100644 --- a/server/tests/test_init_agent.py +++ b/server/tests/test_init_agent.py @@ -17,11 +17,11 @@ def make_agent_payload( - agent_id: str | None = None, + agent_name: str | None = None, name: str = "testagent0001", steps: list[dict[str, Any]] | None = None, ) -> dict[str, Any]: - resolved_name = name if name != "testagent0001" else (agent_id or f"agent-{uuid.uuid4().hex[:12]}") + resolved_name = name if name != "testagent0001" else (agent_name or f"agent-{uuid.uuid4().hex[:12]}") canonical_name = resolved_name.lower().replace(" ", "-") if len(canonical_name) < 10: canonical_name = f"{canonical_name}-agent".replace("--", "-") @@ -36,7 +36,7 @@ def make_agent_payload( ] return { "agent": { - "agent_id": canonical_name, + "agent_name": canonical_name, "agent_name": canonical_name, "agent_description": "desc", "agent_version": "1.0", @@ -120,7 +120,7 @@ def test_agent_endpoints_normalize_mixed_case_agent_name(client: TestClient) -> assert delete_policy_resp.status_code == 200 -def test_init_agent_idempotent_same_steps(client: TestClient) -> None: +def test_init_agent_nameempotent_same_steps(client: TestClient) -> None: # Given: an init payload payload = make_agent_payload() # When: initializing the agent the first time @@ -183,8 +183,8 @@ def test_init_agent_updates_metadata_on_reinit(client: TestClient) -> None: def test_init_agent_adds_new_step(client: TestClient) -> None: # Given: an agent id and base payload - agent_id = str(uuid.uuid4()) - base = make_agent_payload(agent_id=agent_id) + agent_name = str(uuid.uuid4()) + base = make_agent_payload(agent_name=agent_name) # When: initializing the agent r1 = client.post("/api/v1/agents/initAgent", json=base) assert r1.status_code == 200 @@ -200,14 +200,14 @@ def test_init_agent_adds_new_step(client: TestClient) -> None: ] r2 = client.post( "/api/v1/agents/initAgent", - json=make_agent_payload(agent_id=agent_id, steps=steps), + json=make_agent_payload(agent_name=agent_name, steps=steps), ) assert r2.status_code == 200 # Then: the agent is not newly created assert r2.json()["created"] is False # When: fetching the agent - g = client.get(f"/api/v1/agents/{agent_id}") + g = client.get(f"/api/v1/agents/{agent_name}") assert g.status_code == 200 names = {s["name"] for s in g.json()["steps"]} # Then: both steps are present @@ -216,15 +216,15 @@ def test_init_agent_adds_new_step(client: TestClient) -> None: def test_init_agent_overwrites_step_on_signature_change(client: TestClient) -> None: # Given: a base payload for an agent - agent_id = str(uuid.uuid4()) - base = make_agent_payload(agent_id=agent_id) + agent_name = str(uuid.uuid4()) + base = make_agent_payload(agent_name=agent_name) # When: initializing the agent r1 = client.post("/api/v1/agents/initAgent", json=base) assert r1.status_code == 200 # When: updating tool_a schema changed = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, steps=[ { "type": "tool", @@ -245,8 +245,8 @@ def test_init_agent_overwrites_step_on_signature_change(client: TestClient) -> N def test_get_agent_returns_evaluators(client: TestClient) -> None: """Test that GET /agents/{id} returns evaluators.""" # Given: an agent with evaluators - agent_id = str(uuid.uuid4()) - payload = make_agent_payload(agent_id=agent_id) + agent_name = str(uuid.uuid4()) + payload = make_agent_payload(agent_name=agent_name) payload["evaluators"] = [ {"name": "eval-a", "description": "First", "config_schema": {}}, {"name": "eval-b", "description": "Second", "config_schema": {"type": "object"}}, @@ -255,7 +255,7 @@ def test_get_agent_returns_evaluators(client: TestClient) -> None: assert resp.status_code == 200 # When: fetching the agent - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") # Then: evaluators are included in the response assert get_resp.status_code == 200 data = get_resp.json() @@ -325,10 +325,10 @@ def test_set_agent_policy_first_time(client: TestClient) -> None: payload = make_agent_payload() r = client.post("/api/v1/agents/initAgent", json=payload) assert r.status_code == 200 - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] # When: assigning policy the first time - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: success and no old policy assert resp.status_code == 200 body = resp.json() @@ -341,11 +341,11 @@ def test_get_agent_policy_after_assignment(client: TestClient) -> None: policy_id = _create_policy(client) payload = make_agent_payload() client.post("/api/v1/agents/initAgent", json=payload) - agent_id = payload["agent"]["agent_id"] - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + agent_name = payload["agent"]["agent_name"] + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # When: retrieving policy - resp = client.get(f"/api/v1/agents/{agent_id}/policy") + resp = client.get(f"/api/v1/agents/{agent_name}/policy") # Then: we see the assigned policy id assert resp.status_code == 200 assert resp.json()["policy_id"] == policy_id @@ -357,11 +357,11 @@ def test_reassign_agent_policy_returns_old_id(client: TestClient) -> None: second = _create_policy(client) payload = make_agent_payload() client.post("/api/v1/agents/initAgent", json=payload) - agent_id = payload["agent"]["agent_id"] - client.post(f"/api/v1/agents/{agent_id}/policy/{first}") + agent_name = payload["agent"]["agent_name"] + client.post(f"/api/v1/agents/{agent_name}/policy/{first}") # When: reassigning to another policy - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{second}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{second}") # Then: success and old_policy_id equals the first policy id assert resp.status_code == 200 assert resp.json()["success"] is True @@ -373,17 +373,17 @@ def test_delete_agent_policy_then_get_404(client: TestClient) -> None: policy_id = _create_policy(client) payload = make_agent_payload() client.post("/api/v1/agents/initAgent", json=payload) - agent_id = payload["agent"]["agent_id"] - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + agent_name = payload["agent"]["agent_name"] + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # When: removing the policy association - del_resp = client.delete(f"/api/v1/agents/{agent_id}/policy") + del_resp = client.delete(f"/api/v1/agents/{agent_name}/policy") # Then: deletion success assert del_resp.status_code == 200 assert del_resp.json()["success"] is True # When: fetching policy after deletion - get_resp = client.get(f"/api/v1/agents/{agent_id}/policy") + get_resp = client.get(f"/api/v1/agents/{agent_name}/policy") # Then: not found assert get_resp.status_code == 404 @@ -403,11 +403,11 @@ def test_set_policy_not_found_returns_404(client: TestClient) -> None: # Given: an agent and a bogus policy id payload = make_agent_payload() client.post("/api/v1/agents/initAgent", json=payload) - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] bogus_policy = "999999999" # When: assigning a non-existent policy - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{bogus_policy}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{bogus_policy}") # Then: 404 assert resp.status_code == 404 @@ -416,10 +416,10 @@ def test_list_agent_controls_no_policy_returns_empty(client: TestClient) -> None # Given: an agent without a policy payload = make_agent_payload() client.post("/api/v1/agents/initAgent", json=payload) - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] # When: listing controls - r = client.get(f"/api/v1/agents/{agent_id}/controls") + r = client.get(f"/api/v1/agents/{agent_name}/controls") # Then: empty list assert r.status_code == 200 assert r.json()["controls"] == [] @@ -429,7 +429,7 @@ def test_list_agent_controls_with_policy(client: TestClient) -> None: # Given: an agent with a policy containing one control set and one control payload = make_agent_payload() client.post("/api/v1/agents/initAgent", json=payload) - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] # Create policy, control, and wire them pol_name = f"pol-{uuid.uuid4()}" @@ -447,10 +447,10 @@ def test_list_agent_controls_with_policy(client: TestClient) -> None: # Associate control -> policy; assign policy to agent client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # When: listing controls - r = client.get(f"/api/v1/agents/{agent_id}/controls") + r = client.get(f"/api/v1/agents/{agent_name}/controls") # Then: contains our control serialized via API model assert r.status_code == 200 body = r.json() @@ -541,43 +541,43 @@ def test_list_agents_returns_created_agents(client: TestClient) -> None: assert agent1["agent_name"] == "agent-one-01" assert agent1["step_count"] == 1 # from make_agent_payload assert agent1["evaluator_count"] == 1 - assert agent1["policy_id"] is None + assert agent1["policy_ids"] == [] assert "agent-two-02" in agent_map agent2 = agent_map["agent-two-02"] assert agent2["agent_name"] == "agent-two-02" assert agent2["step_count"] == 2 assert agent2["evaluator_count"] == 0 - assert agent2["policy_id"] is None + assert agent2["policy_ids"] == [] def test_list_agents_with_policy(client: TestClient) -> None: - """Test that list agents shows policy_id when assigned.""" + """Test that list agents shows policy_ids when assigned.""" # Given: an agent with a policy assigned payload = make_agent_payload() client.post("/api/v1/agents/initAgent", json=payload) - agent_id = payload["agent"]["agent_id"] + agent_name = payload["agent"]["agent_name"] policy_id = _create_policy(client) - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # When: listing agents resp = client.get("/api/v1/agents") - # Then: the agent shows the policy_id + # Then: the agent shows policy_ids assert resp.status_code == 200 body = resp.json() assert len(body["agents"]) == 1 - assert body["agents"][0]["policy_id"] == policy_id + assert body["agents"][0]["policy_ids"] == [policy_id] def test_list_agents_pagination(client: TestClient) -> None: """Test cursor-based pagination works correctly.""" # Given: 5 agents - agent_ids = [] + agent_names = [] for i in range(5): - agent_id = str(uuid.uuid4()) - agent_ids.append(agent_id) - payload = make_agent_payload(agent_id=agent_id, name=f"Agent {i}") + agent_name = str(uuid.uuid4()) + agent_names.append(agent_name) + payload = make_agent_payload(agent_name=agent_name, name=f"Agent {i}") r = client.post("/api/v1/agents/initAgent", json=payload) assert r.status_code == 200 diff --git a/server/tests/test_init_agent_conflict_mode.py b/server/tests/test_init_agent_conflict_mode.py index 0cd1b714..e39f1c02 100644 --- a/server/tests/test_init_agent_conflict_mode.py +++ b/server/tests/test_init_agent_conflict_mode.py @@ -13,7 +13,6 @@ def _init_payload( *, - agent_id: str, agent_name: str, agent_description: str = "desc", agent_version: str = "1.0", @@ -24,7 +23,6 @@ def _init_payload( canonical_name = agent_name.lower() payload: dict[str, Any] = { "agent": { - "agent_id": canonical_name, "agent_name": canonical_name, "agent_description": agent_description, "agent_version": agent_version, @@ -73,10 +71,8 @@ def _create_policy_with_agent_evaluator_control( def test_init_agent_overwrite_replaces_steps_and_evaluators(client: TestClient) -> None: # Given: an existing agent registration with baseline steps and evaluators. agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name create_payload = _init_payload( - agent_id=agent_id, agent_name=agent_name, steps=[ { @@ -104,7 +100,6 @@ def test_init_agent_overwrite_replaces_steps_and_evaluators(client: TestClient) # When: initAgent is called in overwrite mode with an updated registration payload. overwrite_payload = _init_payload( - agent_id=agent_id, agent_name=agent_name, agent_description="updated desc", agent_version="2.0", @@ -154,7 +149,7 @@ def test_init_agent_overwrite_replaces_steps_and_evaluators(client: TestClient) } ] - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") assert get_resp.status_code == 200 get_data = get_resp.json() assert get_data["agent"]["agent_description"] == "updated desc" @@ -165,13 +160,11 @@ def test_init_agent_overwrite_replaces_steps_and_evaluators(client: TestClient) def test_init_agent_overwrite_warns_on_removed_referenced_evaluator(client: TestClient) -> None: # Given: an agent whose assigned policy contains a control referencing an agent evaluator. agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name evaluator_name = "custom-eval" init_resp = client.post( "/api/v1/agents/initAgent", json=_init_payload( - agent_id=agent_id, agent_name=agent_name, evaluators=[{"name": evaluator_name, "config_schema": {"type": "object"}}], ), @@ -181,14 +174,13 @@ def test_init_agent_overwrite_warns_on_removed_referenced_evaluator(client: Test policy_id, control_id, control_name = _create_policy_with_agent_evaluator_control( client, agent_name=agent_name, evaluator_name=evaluator_name ) - assign_resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign_resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign_resp.status_code == 200 # When: overwrite mode removes the evaluator from the incoming registration payload. overwrite_resp = client.post( "/api/v1/agents/initAgent", json=_init_payload( - agent_id=agent_id, agent_name=agent_name, evaluators=[], conflict_mode="overwrite", @@ -209,7 +201,7 @@ def test_init_agent_overwrite_warns_on_removed_referenced_evaluator(client: Test } ] - get_resp = client.get(f"/api/v1/agents/{agent_id}/evaluators") + get_resp = client.get(f"/api/v1/agents/{agent_name}/evaluators") assert get_resp.status_code == 200 assert get_resp.json()["evaluators"] == [] @@ -217,9 +209,7 @@ def test_init_agent_overwrite_warns_on_removed_referenced_evaluator(client: Test def test_init_agent_overwrite_noop_reports_not_applied(client: TestClient) -> None: # Given: an existing agent registration and an equivalent overwrite payload. agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name payload = _init_payload( - agent_id=agent_id, agent_name=agent_name, steps=[{"type": "tool", "name": "tool-a", "input_schema": {}, "output_schema": {}}], evaluators=[{"name": "eval-a", "description": "x", "config_schema": {"type": "object"}}], diff --git a/server/tests/test_init_agent_force_replace.py b/server/tests/test_init_agent_force_replace.py index f324a641..9a3e75a4 100644 --- a/server/tests/test_init_agent_force_replace.py +++ b/server/tests/test_init_agent_force_replace.py @@ -19,12 +19,12 @@ def test_init_agent_force_replace_default_false_works_normally(client: TestClien """ # Given: New agent agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name # When: Create without force_replace (default) resp = client.post("/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "Test", "agent_version": "1.0" @@ -46,12 +46,12 @@ def test_init_agent_force_replace_false_explicit_works_normally(client: TestClie """ # Given: New agent agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name # When: Create with force_replace=false resp = client.post("/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "Test", "agent_version": "1.0" @@ -74,11 +74,11 @@ def test_init_agent_force_replace_true_on_valid_data_works_normally(client: Test """ # Given: Create agent with steps agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name resp = client.post("/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "Test", "agent_version": "1.0" @@ -93,7 +93,7 @@ def test_init_agent_force_replace_true_on_valid_data_works_normally(client: Test # When: Update with force_replace=true and add a new step resp = client.post("/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "Updated", "agent_version": "2.0" @@ -108,7 +108,7 @@ def test_init_agent_force_replace_true_on_valid_data_works_normally(client: Test # Then: Should succeed and all steps should be present assert resp.status_code == 200 - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") steps = [s["name"] for s in get_resp.json()["steps"]] assert set(steps) == {"tool1", "tool2", "tool3"} @@ -131,12 +131,12 @@ def test_init_agent_force_replace_recovers_from_corrupted_data(client: TestClien """Test that force_replace=true replaces corrupted stored data.""" # Given: an existing agent with corrupted data in the DB agent_name = f"agent-{uuid.uuid4().hex[:12]}" - agent_id = agent_name + agent_name = agent_name resp = client.post( "/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "Test", "agent_version": "1.0", @@ -156,7 +156,7 @@ def test_init_agent_force_replace_recovers_from_corrupted_data(client: TestClien "/api/v1/agents/initAgent", json={ "agent": { - "agent_id": agent_id, + "agent_name": agent_name, "agent_name": agent_name, "agent_description": "Replaced", "agent_version": "2.0", @@ -170,7 +170,7 @@ def test_init_agent_force_replace_recovers_from_corrupted_data(client: TestClien # Then: request succeeds and stored data is replaced assert resp.status_code == 200 - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") assert get_resp.status_code == 200 data = get_resp.json() assert data["agent"]["agent_description"] == "Replaced" diff --git a/server/tests/test_new_features.py b/server/tests/test_new_features.py index 01bc96ec..ad5a91c8 100644 --- a/server/tests/test_new_features.py +++ b/server/tests/test_new_features.py @@ -6,14 +6,14 @@ def make_agent_payload( - agent_id: str | None = None, + agent_name: str | None = None, name: str | None = None, steps: list | None = None, evaluators: list | None = None, ): """Helper to create agent payload.""" - if agent_id is not None: - name = agent_id + if agent_name is not None: + name = agent_name elif name is None: name = f"agent-{uuid.uuid4().hex[:12]}" canonical_name = name.lower().replace(" ", "-") @@ -21,7 +21,7 @@ def make_agent_payload( canonical_name = f"{canonical_name}-agent".replace("--", "-") return { "agent": { - "agent_id": canonical_name, + "agent_name": canonical_name, "agent_name": canonical_name, "agent_description": "desc", "agent_version": "1.0", @@ -79,10 +79,10 @@ def test_get_evaluators_schema_has_properties(client: TestClient) -> None: def test_patch_agent_remove_step(client: TestClient) -> None: """Given an agent with multiple steps, when removing one step, then only that step is removed.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, steps=[ {"type": "tool", "name": "tool1", "input_schema": {}, "output_schema": {}}, @@ -93,7 +93,7 @@ def test_patch_agent_remove_step(client: TestClient) -> None: # When: patch_resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_steps": [{"type": "tool", "name": "tool1"}]}, ) @@ -103,7 +103,7 @@ def test_patch_agent_remove_step(client: TestClient) -> None: assert data["steps_removed"] == [{"type": "tool", "name": "tool1"}] assert data["evaluators_removed"] == [] - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") steps = [s["name"] for s in get_resp.json()["steps"]] assert "tool1" not in steps assert "tool2" in steps @@ -112,10 +112,10 @@ def test_patch_agent_remove_step(client: TestClient) -> None: def test_patch_agent_remove_evaluator(client: TestClient) -> None: """Given an agent with multiple evaluators, when removing one, then only that evaluator is removed.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ {"name": "eval1", "config_schema": {}}, @@ -126,7 +126,7 @@ def test_patch_agent_remove_evaluator(client: TestClient) -> None: # When: patch_resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_evaluators": ["eval1"]}, ) @@ -135,7 +135,7 @@ def test_patch_agent_remove_evaluator(client: TestClient) -> None: data = patch_resp.json() assert data["evaluators_removed"] == ["eval1"] - get_resp = client.get(f"/api/v1/agents/{agent_id}/evaluators") + get_resp = client.get(f"/api/v1/agents/{agent_name}/evaluators") evals = [e["name"] for e in get_resp.json()["evaluators"]] assert "eval1" not in evals assert "eval2" in evals @@ -144,14 +144,14 @@ def test_patch_agent_remove_evaluator(client: TestClient) -> None: def test_patch_agent_remove_nonexistent_is_idempotent(client: TestClient) -> None: """Given an agent, when removing nonexistent items, then succeeds with empty removed lists.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" - payload = make_agent_payload(agent_id=agent_id, name=name) + payload = make_agent_payload(agent_name=agent_name, name=name) client.post("/api/v1/agents/initAgent", json=payload) # When: patch_resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={ "remove_steps": [{"type": "tool", "name": "nonexistent"}], "remove_evaluators": ["also_nonexistent"], @@ -183,10 +183,10 @@ def test_patch_agent_not_found(client: TestClient) -> None: def test_patch_agent_remove_both(client: TestClient) -> None: """Given an agent with steps and evaluators, when removing both, then both are removed.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, steps=[{"type": "tool", "name": "my_tool", "input_schema": {}, "output_schema": {}}], evaluators=[{"name": "my_eval", "config_schema": {}}], @@ -195,7 +195,7 @@ def test_patch_agent_remove_both(client: TestClient) -> None: # When: patch_resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={ "remove_steps": [{"type": "tool", "name": "my_tool"}], "remove_evaluators": ["my_eval"], @@ -212,10 +212,10 @@ def test_patch_agent_remove_both(client: TestClient) -> None: def test_patch_agent_empty_request_is_noop(client: TestClient) -> None: """Given an agent, when patching with empty lists, then nothing changes and succeeds.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, steps=[{"type": "tool", "name": "keep_me", "input_schema": {}, "output_schema": {}}], evaluators=[{"name": "keep_me_too", "config_schema": {}}], @@ -224,7 +224,7 @@ def test_patch_agent_empty_request_is_noop(client: TestClient) -> None: # When: patch_resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_steps": [], "remove_evaluators": []}, ) @@ -235,11 +235,11 @@ def test_patch_agent_empty_request_is_noop(client: TestClient) -> None: assert data["evaluators_removed"] == [] # Verify nothing was removed - get_resp = client.get(f"/api/v1/agents/{agent_id}") + get_resp = client.get(f"/api/v1/agents/{agent_name}") steps = [s["name"] for s in get_resp.json()["steps"]] assert "keep_me" in steps - get_evals = client.get(f"/api/v1/agents/{agent_id}/evaluators") + get_evals = client.get(f"/api/v1/agents/{agent_name}/evaluators") evals = [e["name"] for e in get_evals.json()["evaluators"]] assert "keep_me_too" in evals @@ -282,9 +282,9 @@ def _create_policy_with_control( def test_policy_assignment_with_builtin_evaluator(client: TestClient) -> None: """Given an agent and a policy with built-in evaluator control, when assigning policy, then succeeds.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" - payload = make_agent_payload(agent_id=agent_id, name=name) + payload = make_agent_payload(agent_name=agent_name, name=name) client.post("/api/v1/agents/initAgent", json=payload) policy_id, _ = _create_policy_with_control( @@ -301,7 +301,7 @@ def test_policy_assignment_with_builtin_evaluator(client: TestClient) -> None: ) # When: - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: assert resp.status_code == 200 @@ -310,10 +310,10 @@ def test_policy_assignment_with_builtin_evaluator(client: TestClient) -> None: def test_policy_assignment_with_registered_agent_evaluator(client: TestClient) -> None: """Given an agent with custom evaluator and matching policy, when assigning policy, then succeeds.""" # Given: - agent_id = f"agent-{uuid.uuid4().hex[:12]}" - agent_name = agent_id + agent_name = f"agent-{uuid.uuid4().hex[:12]}" + agent_name = agent_name payload = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=agent_name, evaluators=[{"name": "custom-eval", "config_schema": {"type": "object"}}], ) @@ -333,7 +333,7 @@ def test_policy_assignment_with_registered_agent_evaluator(client: TestClient) - ) # When: - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Then: assert resp.status_code == 200 @@ -342,9 +342,9 @@ def test_policy_assignment_with_registered_agent_evaluator(client: TestClient) - def test_control_creation_with_unregistered_evaluator_fails(client: TestClient) -> None: """Given an agent without evaluator, when setting control to use that evaluator, then fails.""" # Given: - agent_id = f"agent-{uuid.uuid4().hex[:12]}" - agent_name = agent_id - payload = make_agent_payload(agent_id=agent_id, name=agent_name) + agent_name = f"agent-{uuid.uuid4().hex[:12]}" + agent_name = agent_name + payload = make_agent_payload(agent_name=agent_name, name=agent_name) client.post("/api/v1/agents/initAgent", json=payload) ctl_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4().hex[:8]}"}) @@ -377,7 +377,7 @@ def test_policy_assignment_cross_agent_evaluator_fails(client: TestClient) -> No agent_a_id = f"agent-a-{uuid.uuid4().hex[:12]}" agent_a_name = agent_a_id payload_a = make_agent_payload( - agent_id=agent_a_id, + agent_name=agent_a_id, name=agent_a_name, evaluators=[{"name": "shared-eval", "config_schema": {"type": "object"}}], ) @@ -385,7 +385,7 @@ def test_policy_assignment_cross_agent_evaluator_fails(client: TestClient) -> No agent_b_id = f"agent-b-{uuid.uuid4().hex[:12]}" agent_b_name = agent_b_id - payload_b = make_agent_payload(agent_id=agent_b_id, name=agent_b_name) + payload_b = make_agent_payload(agent_name=agent_b_id, name=agent_b_name) client.post("/api/v1/agents/initAgent", json=payload_b) policy_id, _ = _create_policy_with_control( @@ -426,10 +426,10 @@ def test_policy_assignment_cross_agent_evaluator_fails(client: TestClient) -> No def test_schema_compat_nested_additional_properties_compatible(client: TestClient) -> None: """Given a nested schema, when adding optional property in nested object, then compatible.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload1 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -450,7 +450,7 @@ def test_schema_compat_nested_additional_properties_compatible(client: TestClien # When: add optional property in nested object payload2 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -479,10 +479,10 @@ def test_schema_compat_nested_additional_properties_compatible(client: TestClien def test_schema_compat_nested_type_change_incompatible(client: TestClient) -> None: """Given a nested schema, when changing nested property type, then rejected as incompatible.""" # Given: - agent_id = str(uuid.uuid4()) + agent_name = str(uuid.uuid4()) name = f"Test Agent {uuid.uuid4().hex[:8]}" payload1 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -503,7 +503,7 @@ def test_schema_compat_nested_type_change_incompatible(client: TestClient) -> No # When: change nested type from integer to string payload2 = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=name, evaluators=[ { @@ -540,10 +540,10 @@ def test_patch_agent_remove_evaluator_blocked_by_control(client: TestClient) -> Then: Returns 409 with error message about referencing control """ # Given: Agent with custom evaluator - agent_id = f"agent-{uuid.uuid4().hex[:12]}" - agent_name = agent_id + agent_name = f"agent-{uuid.uuid4().hex[:12]}" + agent_name = agent_name payload = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=agent_name, evaluators=[{"name": "my-eval", "config_schema": {"type": "object"}}], ) @@ -564,12 +564,12 @@ def test_patch_agent_remove_evaluator_blocked_by_control(client: TestClient) -> ) # And: Policy assigned to agent - assign_resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + assign_resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert assign_resp.status_code == 200 # When: Trying to remove the evaluator patch_resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_evaluators": ["my-eval"]}, ) @@ -591,10 +591,10 @@ def test_patch_agent_remove_evaluator_allowed_without_policy(client: TestClient) Then: Succeeds since no controls can reference it """ # Given: Agent with custom evaluator but no policy - agent_id = f"agent-{uuid.uuid4().hex[:12]}" - agent_name = agent_id + agent_name = f"agent-{uuid.uuid4().hex[:12]}" + agent_name = agent_name payload = make_agent_payload( - agent_id=agent_id, + agent_name=agent_name, name=agent_name, evaluators=[{"name": "my-eval", "config_schema": {"type": "object"}}], ) @@ -602,7 +602,7 @@ def test_patch_agent_remove_evaluator_allowed_without_policy(client: TestClient) # When: Removing the evaluator (no policy = no controls can reference it) patch_resp = client.patch( - f"/api/v1/agents/{agent_id}", + f"/api/v1/agents/{agent_name}", json={"remove_evaluators": ["my-eval"]}, ) diff --git a/server/tests/test_policy_integration.py b/server/tests/test_policy_integration.py index 446c520a..bd3ba15d 100644 --- a/server/tests/test_policy_integration.py +++ b/server/tests/test_policy_integration.py @@ -55,7 +55,7 @@ def _create_control(client: TestClient, name: str | None = None, data: dict | No def test_agent_gets_controls_from_policy(client: TestClient) -> None: """Agent should see all controls from its policy.""" # Given: Agent with policy containing 6 controls - agent_id, _ = _create_agent(client) + agent_name, _ = _create_agent(client) policy_id = _create_policy(client) # Create 6 controls @@ -69,11 +69,11 @@ def test_agent_gets_controls_from_policy(client: TestClient) -> None: assert resp.status_code == 200 # Assign policy to agent - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") assert resp.status_code == 200 # When: Get agent's controls - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") assert resp.status_code == 200 controls = resp.json()["controls"] @@ -88,7 +88,7 @@ def test_agent_gets_controls_from_policy(client: TestClient) -> None: def test_agent_controls_update_when_control_added_to_policy(client: TestClient) -> None: """Adding a control to policy should make it visible to agents.""" # Given: Agent → Policy → 2 controls - agent_id, _ = _create_agent(client) + agent_name, _ = _create_agent(client) policy_id = _create_policy(client) control_1_id = _create_control(client, "control-1", {"id": 1}) @@ -96,10 +96,10 @@ def test_agent_controls_update_when_control_added_to_policy(client: TestClient) client.post(f"/api/v1/policies/{policy_id}/controls/{control_1_id}") client.post(f"/api/v1/policies/{policy_id}/controls/{control_2_id}") - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Verify initial state: 2 controls - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") assert len(resp.json()["controls"]) == 2 # When: Add 3 more controls to the policy @@ -115,7 +115,7 @@ def test_agent_controls_update_when_control_added_to_policy(client: TestClient) assert resp.status_code == 200 # Then: Agent now sees 5 controls total - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") controls = resp.json()["controls"] assert len(controls) == 5 @@ -126,7 +126,7 @@ def test_agent_controls_update_when_control_added_to_policy(client: TestClient) def test_switching_agent_policy_changes_controls(client: TestClient) -> None: """Switching agent's policy should completely change its controls.""" # Given: Two policies with different controls - agent_id, _ = _create_agent(client) + agent_name, _ = _create_agent(client) # Policy A with controls {1, 2} policy_a_id = _create_policy(client, "policy-a") @@ -143,18 +143,18 @@ def test_switching_agent_policy_changes_controls(client: TestClient) -> None: client.post(f"/api/v1/policies/{policy_b_id}/controls/{control_4_id}") # Assign policy A to agent - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_a_id}") - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_a_id}") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") controls_a = resp.json()["controls"] assert len(controls_a) == 2 assert {r["id"] for r in controls_a} == {control_1_id, control_2_id} # When: Switch to policy B - resp = client.post(f"/api/v1/agents/{agent_id}/policy/{policy_b_id}") + resp = client.post(f"/api/v1/agents/{agent_name}/policy/{policy_b_id}") assert resp.status_code == 200 # Then: Agent's controls change completely - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") controls_b = resp.json()["controls"] assert len(controls_b) == 2 assert {r["id"] for r in controls_b} == {control_3_id, control_4_id} @@ -163,23 +163,23 @@ def test_switching_agent_policy_changes_controls(client: TestClient) -> None: def test_removing_agent_policy_clears_controls(client: TestClient) -> None: """Removing policy from agent should result in empty controls list.""" # Given: Agent with policy that has controls - agent_id, _ = _create_agent(client) + agent_name, _ = _create_agent(client) policy_id = _create_policy(client) control_id = _create_control(client, "control-1", {"id": 1}) client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Verify agent has controls - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") assert len(resp.json()["controls"]) > 0 # When: Remove policy from agent - resp = client.delete(f"/api/v1/agents/{agent_id}/policy") + resp = client.delete(f"/api/v1/agents/{agent_name}/policy") assert resp.status_code == 200 # Then: Agent returns empty controls list - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") assert resp.status_code == 200 assert resp.json()["controls"] == [] @@ -187,7 +187,7 @@ def test_removing_agent_policy_clears_controls(client: TestClient) -> None: def test_removing_control_from_policy_removes_from_agent(client: TestClient) -> None: """Removing control from policy should remove it from agent.""" # Given: Agent → Policy → 4 controls - agent_id, _ = _create_agent(client) + agent_name, _ = _create_agent(client) policy_id = _create_policy(client) control_1_id = _create_control(client, "control-1", {"id": 1}) @@ -199,10 +199,10 @@ def test_removing_control_from_policy_removes_from_agent(client: TestClient) -> client.post(f"/api/v1/policies/{policy_id}/controls/{control_2_id}") client.post(f"/api/v1/policies/{policy_id}/controls/{control_3_id}") client.post(f"/api/v1/policies/{policy_id}/controls/{control_4_id}") - client.post(f"/api/v1/agents/{agent_id}/policy/{policy_id}") + client.post(f"/api/v1/agents/{agent_name}/policy/{policy_id}") # Verify initial state: 4 controls - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") assert len(resp.json()["controls"]) == 4 # When: Remove 2 controls from policy @@ -212,7 +212,7 @@ def test_removing_control_from_policy_removes_from_agent(client: TestClient) -> assert resp.status_code == 200 # Then: Agent sees 2 remaining controls - resp = client.get(f"/api/v1/agents/{agent_id}/controls") + resp = client.get(f"/api/v1/agents/{agent_name}/controls") controls = resp.json()["controls"] assert len(controls) == 2 assert {r["id"] for r in controls} == {control_3_id, control_4_id} @@ -292,3 +292,330 @@ def test_control_shared_between_policies(client: TestClient) -> None: assert control_id in {c["id"] for c in resp_a.json()["controls"]} assert control_id in {c["id"] for c in resp_b.json()["controls"]} + + +def test_agent_controls_union_across_multiple_policies_with_dedupe(client: TestClient) -> None: + """Agent controls should be additive across policies and deduplicated by control id.""" + agent_name, _ = _create_agent(client) + policy_a_id = _create_policy(client, "policy-a") + policy_b_id = _create_policy(client, "policy-b") + + shared_control_id = _create_control(client) + policy_a_only_control_id = _create_control(client) + policy_b_only_control_id = _create_control(client) + + # Policy A: shared + A-only + resp = client.post(f"/api/v1/policies/{policy_a_id}/controls/{shared_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/policies/{policy_a_id}/controls/{policy_a_only_control_id}") + assert resp.status_code == 200 + + # Policy B: shared + B-only + resp = client.post(f"/api/v1/policies/{policy_b_id}/controls/{shared_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/policies/{policy_b_id}/controls/{policy_b_only_control_id}") + assert resp.status_code == 200 + + # Associate both policies with the same agent (plural additive endpoint). + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_a_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_b_id}") + assert resp.status_code == 200 + + # Agent should have both policy associations. + policies_resp = client.get(f"/api/v1/agents/{agent_name}/policies") + assert policies_resp.status_code == 200 + assert policies_resp.json()["policy_ids"] == [policy_a_id, policy_b_id] + + # Active controls should be union(policy A, policy B) with shared deduped. + controls_resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert controls_resp.status_code == 200 + controls = controls_resp.json()["controls"] + received_control_ids = {control["id"] for control in controls} + assert received_control_ids == { + shared_control_id, + policy_a_only_control_id, + policy_b_only_control_id, + } + assert len(controls) == 3 + + # list_agents count should match the deduplicated union. + agents_resp = client.get("/api/v1/agents", params={"name": agent_name}) + assert agents_resp.status_code == 200 + matching_agents = [ + agent for agent in agents_resp.json()["agents"] if agent["agent_name"] == agent_name + ] + assert len(matching_agents) == 1 + assert matching_agents[0]["active_controls_count"] == 3 + + +def test_remove_one_policy_keeps_controls_from_remaining_policies(client: TestClient) -> None: + """Removing one policy should preserve controls inherited from other policies.""" + agent_name, _ = _create_agent(client) + policy_a_id = _create_policy(client, "policy-a") + policy_b_id = _create_policy(client, "policy-b") + + shared_control_id = _create_control(client) + policy_a_only_control_id = _create_control(client) + policy_b_only_control_id = _create_control(client) + + # Policy A: shared + A-only + resp = client.post(f"/api/v1/policies/{policy_a_id}/controls/{shared_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/policies/{policy_a_id}/controls/{policy_a_only_control_id}") + assert resp.status_code == 200 + + # Policy B: shared + B-only + resp = client.post(f"/api/v1/policies/{policy_b_id}/controls/{shared_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/policies/{policy_b_id}/controls/{policy_b_only_control_id}") + assert resp.status_code == 200 + + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_a_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_b_id}") + assert resp.status_code == 200 + + # Remove only policy A. + resp = client.delete(f"/api/v1/agents/{agent_name}/policies/{policy_a_id}") + assert resp.status_code == 200 + + policies_resp = client.get(f"/api/v1/agents/{agent_name}/policies") + assert policies_resp.status_code == 200 + assert policies_resp.json()["policy_ids"] == [policy_b_id] + + controls_resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert controls_resp.status_code == 200 + received_control_ids = {control["id"] for control in controls_resp.json()["controls"]} + assert received_control_ids == {shared_control_id, policy_b_only_control_id} + + +def test_remove_all_policies_preserves_direct_controls(client: TestClient) -> None: + """Removing all policy links should keep direct agent-control associations active.""" + agent_name, _ = _create_agent(client) + policy_id = _create_policy(client) + policy_control_id = _create_control(client) + direct_control_id = _create_control(client) + + resp = client.post(f"/api/v1/policies/{policy_id}/controls/{policy_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{direct_control_id}") + assert resp.status_code == 200 + + resp = client.delete(f"/api/v1/agents/{agent_name}/policies") + assert resp.status_code == 200 + assert resp.json()["success"] is True + + policies_resp = client.get(f"/api/v1/agents/{agent_name}/policies") + assert policies_resp.status_code == 200 + assert policies_resp.json()["policy_ids"] == [] + + controls_resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert controls_resp.status_code == 200 + assert {control["id"] for control in controls_resp.json()["controls"]} == {direct_control_id} + + +def test_add_agent_policy_is_idempotent(client: TestClient) -> None: + """Adding the same policy association twice should not duplicate links.""" + agent_name, _ = _create_agent(client) + policy_id = _create_policy(client) + + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_id}") + assert resp.status_code == 200 + + policies_resp = client.get(f"/api/v1/agents/{agent_name}/policies") + assert policies_resp.status_code == 200 + assert policies_resp.json()["policy_ids"] == [policy_id] + + +def test_add_agent_control_is_idempotent(client: TestClient) -> None: + """Adding the same direct control twice should not duplicate active controls.""" + agent_name, _ = _create_agent(client) + control_id = _create_control(client) + + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert resp.status_code == 200 + + controls_resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert controls_resp.status_code == 200 + controls = controls_resp.json()["controls"] + assert {control["id"] for control in controls} == {control_id} + assert len(controls) == 1 + + +def test_agent_policy_endpoints_return_404_for_missing_resources(client: TestClient) -> None: + """Plural policy endpoints should return consistent 404s for missing agent/policy.""" + existing_agent_name, _ = _create_agent(client) + existing_policy_id = _create_policy(client) + + missing_agent_name = "missing-agent-1234" + missing_policy_id = 999999 + + # Missing agent on add/list/remove-one/remove-all. + resp = client.post(f"/api/v1/agents/{missing_agent_name}/policies/{existing_policy_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "AGENT_NOT_FOUND" + + resp = client.get(f"/api/v1/agents/{missing_agent_name}/policies") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "AGENT_NOT_FOUND" + + resp = client.delete(f"/api/v1/agents/{missing_agent_name}/policies/{existing_policy_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "AGENT_NOT_FOUND" + + resp = client.delete(f"/api/v1/agents/{missing_agent_name}/policies") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "AGENT_NOT_FOUND" + + # Missing policy on add/remove-one. + resp = client.post(f"/api/v1/agents/{existing_agent_name}/policies/{missing_policy_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "POLICY_NOT_FOUND" + + resp = client.delete(f"/api/v1/agents/{existing_agent_name}/policies/{missing_policy_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "POLICY_NOT_FOUND" + + +def test_agent_control_endpoints_return_404_for_missing_resources(client: TestClient) -> None: + """Direct control association endpoints should return 404s for missing agent/control.""" + existing_agent_name, _ = _create_agent(client) + existing_control_id = _create_control(client) + + missing_agent_name = "missing-agent-1234" + missing_control_id = 999999 + + # Missing agent on add/remove. + resp = client.post(f"/api/v1/agents/{missing_agent_name}/controls/{existing_control_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "AGENT_NOT_FOUND" + + resp = client.delete(f"/api/v1/agents/{missing_agent_name}/controls/{existing_control_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "AGENT_NOT_FOUND" + + # Missing control on add/remove. + resp = client.post(f"/api/v1/agents/{existing_agent_name}/controls/{missing_control_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "CONTROL_NOT_FOUND" + + resp = client.delete(f"/api/v1/agents/{existing_agent_name}/controls/{missing_control_id}") + assert resp.status_code == 404 + assert resp.json()["error_code"] == "CONTROL_NOT_FOUND" + + +def test_agent_gets_controls_from_direct_associations(client: TestClient) -> None: + """Agent should see controls directly associated with it.""" + agent_name, _ = _create_agent(client) + control_1_id = _create_control(client) + control_2_id = _create_control(client) + + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{control_1_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{control_2_id}") + assert resp.status_code == 200 + + resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert resp.status_code == 200 + controls = resp.json()["controls"] + assert {control["id"] for control in controls} == {control_1_id, control_2_id} + + +def test_agent_controls_are_union_of_policy_and_direct_with_dedupe(client: TestClient) -> None: + """Agent control list should union policy + direct controls and de-duplicate by control id.""" + agent_name, _ = _create_agent(client) + policy_id = _create_policy(client) + + shared_control_id = _create_control(client) + policy_only_control_id = _create_control(client) + direct_only_control_id = _create_control(client) + + # Associate shared + policy-only controls via policy. + resp = client.post(f"/api/v1/policies/{policy_id}/controls/{shared_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/policies/{policy_id}/controls/{policy_only_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_id}") + assert resp.status_code == 200 + + # Associate shared + direct-only controls directly with agent. + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{shared_control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{direct_only_control_id}") + assert resp.status_code == 200 + + # Shared control should appear only once in active controls. + resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert resp.status_code == 200 + controls = resp.json()["controls"] + received_control_ids = {control["id"] for control in controls} + assert received_control_ids == {shared_control_id, policy_only_control_id, direct_only_control_id} + assert len(controls) == 3 + + # list_agents active_controls_count should reflect deduplicated union as well. + agents_resp = client.get("/api/v1/agents", params={"name": agent_name}) + assert agents_resp.status_code == 200 + matching_agents = [ + agent for agent in agents_resp.json()["agents"] if agent["agent_name"] == agent_name + ] + assert len(matching_agents) == 1 + assert matching_agents[0]["active_controls_count"] == 3 + + +def test_remove_direct_control_keeps_policy_inherited_control_active(client: TestClient) -> None: + """Removing a direct association should keep control active when policy still provides it.""" + agent_name, _ = _create_agent(client) + policy_id = _create_policy(client) + control_id = _create_control(client) + + resp = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/policies/{policy_id}") + assert resp.status_code == 200 + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert resp.status_code == 200 + + resp = client.delete(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["removed_direct_association"] is True + assert body["control_still_active"] is True + + # Idempotent behavior when no direct link remains but policy inheritance still exists. + resp = client.delete(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert resp.status_code == 200 + body = resp.json() + assert body["removed_direct_association"] is False + assert body["control_still_active"] is True + + resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert resp.status_code == 200 + assert control_id in {control["id"] for control in resp.json()["controls"]} + + +def test_remove_direct_control_deactivates_when_not_inherited(client: TestClient) -> None: + """Removing a direct-only control should make it inactive for the agent.""" + agent_name, _ = _create_agent(client) + control_id = _create_control(client) + + resp = client.post(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert resp.status_code == 200 + + resp = client.delete(f"/api/v1/agents/{agent_name}/controls/{control_id}") + assert resp.status_code == 200 + body = resp.json() + assert body["success"] is True + assert body["removed_direct_association"] is True + assert body["control_still_active"] is False + + resp = client.get(f"/api/v1/agents/{agent_name}/controls") + assert resp.status_code == 200 + assert control_id not in {control["id"] for control in resp.json()["controls"]} diff --git a/server/tests/test_services_controls.py b/server/tests/test_services_controls.py index 797835d2..72cda230 100644 --- a/server/tests/test_services_controls.py +++ b/server/tests/test_services_controls.py @@ -7,7 +7,14 @@ from agent_control_models.errors import ErrorCode from agent_control_server.errors import APIValidationError -from agent_control_server.models import Agent, Control, Policy, policy_controls +from agent_control_server.models import ( + Agent, + Control, + Policy, + agent_controls, + agent_policies, + policy_controls, +) from agent_control_server.services.controls import list_controls_for_agent, list_controls_for_policy from .utils import VALID_CONTROL_PAYLOAD @@ -42,44 +49,52 @@ async def test_list_controls_for_policy_returns_controls(async_db) -> None: @pytest.mark.asyncio async def test_list_controls_for_agent_returns_controls(async_db) -> None: - # Given: an agent assigned to a policy with one control + # Given: an agent associated with one policy control and one direct control policy = Policy(name=f"policy-{uuid.uuid4()}") - control = Control(name=f"control-{uuid.uuid4()}", data=VALID_CONTROL_PAYLOAD) + policy_control = Control(name=f"policy-control-{uuid.uuid4()}", data=VALID_CONTROL_PAYLOAD) + direct_control = Control(name=f"direct-control-{uuid.uuid4()}", data=VALID_CONTROL_PAYLOAD) agent = Agent( name=f"agent-{uuid.uuid4()}", data={}, - policy=policy, ) - async_db.add_all([policy, control, agent]) + async_db.add_all([policy, policy_control, direct_control, agent]) await async_db.flush() await async_db.execute( - insert(policy_controls).values({"policy_id": policy.id, "control_id": control.id}) + insert(agent_policies).values({"agent_name": agent.name, "policy_id": policy.id}) + ) + await async_db.execute( + insert(policy_controls).values({"policy_id": policy.id, "control_id": policy_control.id}) + ) + await async_db.execute( + insert(agent_controls).values({"agent_name": agent.name, "control_id": direct_control.id}) ) await async_db.commit() # When: listing controls for the agent controls = await list_controls_for_agent(agent.name, async_db) - # Then: the API control is returned with expected fields - assert len(controls) == 1 - assert controls[0].name == control.name - assert controls[0].control.evaluator.name == VALID_CONTROL_PAYLOAD["evaluator"]["name"] + # Then: both policy-derived and direct controls are returned + assert len(controls) == 2 + names = {control.name for control in controls} + assert names == {policy_control.name, direct_control.name} @pytest.mark.asyncio async def test_list_controls_for_agent_corrupted_data_raises(async_db) -> None: - # Given: an agent assigned to a policy with corrupted control data + # Given: an agent associated with a policy containing corrupted control data policy = Policy(name=f"policy-{uuid.uuid4()}") control = Control(name=f"control-{uuid.uuid4()}", data={"bad": "data"}) agent = Agent( name=f"agent-{uuid.uuid4()}", data={}, - policy=policy, ) async_db.add_all([policy, control, agent]) await async_db.flush() + await async_db.execute( + insert(agent_policies).values({"agent_name": agent.name, "policy_id": policy.id}) + ) await async_db.execute( insert(policy_controls).values({"policy_id": policy.id, "control_id": control.id}) ) diff --git a/ui/README.md b/ui/README.md index a6395b07..aaabc41d 100644 --- a/ui/README.md +++ b/ui/README.md @@ -1,6 +1,6 @@ # Agent Control UI -Next.js dashboard for managing [Agent Control](https://github.com/agentcontrol/agent-control) agents, controls, and policies. Runs on port 4000. +Next.js dashboard for managing [Agent Control](https://github.com/agentcontrol/agent-control) agents and controls. Runs on port 4000. ## Prerequisites diff --git a/ui/src/core/api/client.ts b/ui/src/core/api/client.ts index 3facbcb0..337fd05c 100644 --- a/ui/src/core/api/client.ts +++ b/ui/src/core/api/client.ts @@ -65,21 +65,42 @@ export const api = { apiClient.GET('/api/v1/agents/{agent_name}/controls', { params: { path: { agent_name: agentName } }, }), - setPolicy: ( + getPolicies: (agentName: GetAgentPathParams['agent_name']) => + apiClient.GET('/api/v1/agents/{agent_name}/policies', { + params: { path: { agent_name: agentName } }, + }), + addPolicy: ( agentName: GetAgentPathParams['agent_name'], policyId: number ) => - apiClient.POST('/api/v1/agents/{agent_name}/policy/{policy_id}', { + apiClient.POST('/api/v1/agents/{agent_name}/policies/{policy_id}', { params: { path: { agent_name: agentName, policy_id: policyId } }, }), - getPolicy: (agentName: GetAgentPathParams['agent_name']) => - apiClient.GET('/api/v1/agents/{agent_name}/policy', { - params: { path: { agent_name: agentName } }, + removePolicy: ( + agentName: GetAgentPathParams['agent_name'], + policyId: number + ) => + apiClient.DELETE('/api/v1/agents/{agent_name}/policies/{policy_id}', { + params: { path: { agent_name: agentName, policy_id: policyId } }, }), - deletePolicy: (agentName: GetAgentPathParams['agent_name']) => - apiClient.DELETE('/api/v1/agents/{agent_name}/policy', { + clearPolicies: (agentName: GetAgentPathParams['agent_name']) => + apiClient.DELETE('/api/v1/agents/{agent_name}/policies', { params: { path: { agent_name: agentName } }, }), + addControl: ( + agentName: GetAgentPathParams['agent_name'], + controlId: number + ) => + apiClient.POST('/api/v1/agents/{agent_name}/controls/{control_id}', { + params: { path: { agent_name: agentName, control_id: controlId } }, + }), + removeControl: ( + agentName: GetAgentPathParams['agent_name'], + controlId: number + ) => + apiClient.DELETE('/api/v1/agents/{agent_name}/controls/{control_id}', { + params: { path: { agent_name: agentName, control_id: controlId } }, + }), }, evaluators: { list: () => apiClient.GET('/api/v1/evaluators'), @@ -143,6 +164,10 @@ export const api = { apiClient.POST('/api/v1/policies/{policy_id}/controls/{control_id}', { params: { path: { policy_id: policyId, control_id: controlId } }, }), + removeControl: (policyId: number, controlId: number) => + apiClient.DELETE('/api/v1/policies/{policy_id}/controls/{control_id}', { + params: { path: { policy_id: policyId, control_id: controlId } }, + }), }, observability: { getStats: (params: { diff --git a/ui/src/core/api/generated/api-types.ts b/ui/src/core/api/generated/api-types.ts index 339c1417..d54bbae1 100644 --- a/ui/src/core/api/generated/api-types.ts +++ b/ui/src/core/api/generated/api-types.ts @@ -15,7 +15,7 @@ export interface paths { * List all agents * @description List all registered agents with cursor-based pagination. * - * Returns a summary of each agent including identifier, policy assignment, + * Returns a summary of each agent including identifier, policy associations, * and counts of registered steps and evaluators. * * Args: @@ -134,15 +134,14 @@ export interface paths { * List agent's active controls * @description List all protection controls active for an agent. * - * Controls are inherited from the agent's assigned policy. - * Returns an empty list if the agent has no policy. + * Controls include the union of policy-derived and directly associated controls. * * Args: * agent_name: Agent identifier * db: Database session (injected) * * Returns: - * AgentControlsResponse with list of controls (empty if no policy) + * AgentControlsResponse with list of active controls * * Raises: * HTTPException 404: Agent not found @@ -156,6 +155,30 @@ export interface paths { patch?: never; trace?: never; }; + '/api/v1/agents/{agent_name}/controls/{control_id}': { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Associate control directly with agent + * @description Associate a control directly with an agent (idempotent). + */ + post: operations['add_agent_control_api_v1_agents__agent_name__controls__control_id__post']; + /** + * Remove direct control association from agent + * @description Remove a direct control association from an agent (idempotent). + */ + delete: operations['remove_agent_control_api_v1_agents__agent_name__controls__control_id__delete']; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; '/api/v1/agents/{agent_name}/evaluators': { parameters: { query?: never; @@ -223,7 +246,7 @@ export interface paths { patch?: never; trace?: never; }; - '/api/v1/agents/{agent_name}/policy': { + '/api/v1/agents/{agent_name}/policies': { parameters: { query?: never; header?: never; @@ -231,38 +254,66 @@ export interface paths { cookie?: never; }; /** - * Get agent's assigned policy - * @description Retrieve the policy currently assigned to an agent. - * - * Args: - * agent_name: Agent identifier - * db: Database session (injected) - * - * Returns: - * GetPolicyResponse with policy ID + * List policies associated with agent + * @description List policy IDs associated with an agent. + */ + get: operations['get_agent_policies_api_v1_agents__agent_name__policies_get']; + put?: never; + post?: never; + /** + * Remove all policy associations from agent + * @description Remove all policy associations from an agent. + */ + delete: operations['remove_all_agent_policies_api_v1_agents__agent_name__policies_delete']; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + '/api/v1/agents/{agent_name}/policies/{policy_id}': { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Associate policy with agent + * @description Associate a policy with an agent (idempotent). + */ + post: operations['add_agent_policy_api_v1_agents__agent_name__policies__policy_id__post']; + /** + * Remove policy association from agent + * @description Remove a policy association from an agent. * - * Raises: - * HTTPException 404: Agent not found or agent has no policy assigned + * Idempotent for existing resources: removing a non-associated link is a no-op. + * Missing agent/policy resources still return 404. + */ + delete: operations['remove_agent_policy_api_v1_agents__agent_name__policies__policy_id__delete']; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; + '/api/v1/agents/{agent_name}/policy': { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Get agent's assigned policy (compatibility) + * @description Compatibility endpoint that returns the first associated policy. */ get: operations['get_agent_policy_api_v1_agents__agent_name__policy_get']; put?: never; post?: never; /** - * Remove agent's policy assignment - * @description Remove the policy assignment from an agent. - * - * The agent will no longer have any protection controls active. - * - * Args: - * agent_name: Agent identifier - * db: Database session (injected) - * - * Returns: - * DeletePolicyResponse with success flag - * - * Raises: - * HTTPException 404: Agent not found or agent has no policy assigned - * HTTPException 500: Database error during removal + * Remove agent's policy assignment (compatibility) + * @description Compatibility endpoint that removes all policy associations. */ delete: operations['delete_agent_policy_api_v1_agents__agent_name__policy_delete']; options?: never; @@ -280,22 +331,8 @@ export interface paths { get?: never; put?: never; /** - * Assign policy to agent - * @description Assign a policy to an agent, replacing any existing policy assignment. - * - * The agent will immediately inherit all controls from the assigned policy. - * - * Args: - * agent_name: Agent identifier - * policy_id: ID of the policy to assign - * db: Database session (injected) - * - * Returns: - * SetPolicyResponse with success flag and previous policy ID (if any) - * - * Raises: - * HTTPException 404: Agent or policy not found - * HTTPException 500: Database error during assignment + * Assign policy to agent (compatibility) + * @description Compatibility endpoint that replaces all policy associations with one policy. */ post: operations['set_agent_policy_api_v1_agents__agent_name__policy__policy_id__post']; delete?: never; @@ -416,7 +453,7 @@ export interface paths { * Delete a control * @description Delete a control by ID. * - * By default, deletion fails if the control is associated with any policy. + * By default, deletion fails if the control is associated with any policy or agent. * Use force=true to automatically dissociate and delete. * * Args: @@ -425,7 +462,7 @@ export interface paths { * db: Database session (injected) * * Returns: - * DeleteControlResponse with success flag and list of dissociated policies + * DeleteControlResponse with success flag and dissociation details * * Raises: * HTTPException 404: Control not found @@ -658,7 +695,7 @@ export interface paths { * - control_execution_id: Get a specific event * - agent_name: Filter by agent * - control_ids: Filter by controls - * - actions: Filter by actions (allow, deny, steer, warn, log) + * - actions: Filter by actions (allow, deny, warn, log) * - matched: Filter by matched status * - check_stages: Filter by check stage (pre, post) * - applies_to: Filter by call type (llm_call, tool_call) @@ -967,7 +1004,7 @@ export interface components { AgentControlsResponse: { /** * Controls - * @description List of controls associated with the agent via its policy + * @description List of active controls associated with the agent */ controls: components['schemas']['Control'][]; }; @@ -978,7 +1015,7 @@ export interface components { AgentRef: { /** * Agent Name - * @description Agent identifier + * @description Agent name */ agent_name: string; }; @@ -989,7 +1026,7 @@ export interface components { AgentSummary: { /** * Active Controls Count - * @description Number of active controls from agent's policy + * @description Number of active controls for this agent * @default 0 */ active_controls_count: number; @@ -1010,10 +1047,10 @@ export interface components { */ evaluator_count: number; /** - * Policy Id - * @description ID of assigned policy, if any + * Policy Ids + * @description IDs of policies associated with the agent */ - policy_id?: number | null; + policy_ids?: number[]; /** * Step Count * @description Number of steps registered with the agent @@ -1112,29 +1149,12 @@ export interface components { * are returned from API endpoints. Unconfigured controls are filtered out. */ Control: { - control: components['schemas']['ControlDefinition']; + control: components['schemas']['ControlDefinition-Output']; /** Id */ id: number; /** Name */ name: string; }; - /** - * SteeringContext - * @description Steering context for steer actions. - * - * This model provides an extensible structure for steering guidance. - * Future fields could include severity, categories, suggested_actions, etc. - * @example { - * "message": "This large transfer requires user verification. Request 2FA code from user, verify it, then retry the transaction with verified_2fa=True." - * } - */ - SteeringContext: { - /** - * Message - * @description Guidance message explaining what needs to be corrected and how - */ - message: string; - }; /** * ControlAction * @description What to do when control matches. @@ -1146,11 +1166,76 @@ export interface components { * @enum {string} */ decision: 'allow' | 'deny' | 'steer' | 'warn' | 'log'; + /** @description Steering context object for steer actions. Strongly recommended when decision='steer' to provide correction suggestions. If not provided, the evaluator result message will be used as fallback. */ + steering_context?: components['schemas']['SteeringContext'] | null; + }; + /** + * ControlDefinition + * @description A control definition to evaluate agent interactions. + * + * This model contains only the logic and configuration. + * Identity fields (id, name) are managed by the database. + * @example { + * "action": { + * "decision": "deny" + * }, + * "description": "Block outputs containing US Social Security Numbers", + * "enabled": true, + * "evaluator": { + * "config": { + * "pattern": "\\b\\d{3}-\\d{2}-\\d{4}\\b" + * }, + * "name": "regex" + * }, + * "execution": "server", + * "scope": { + * "stages": [ + * "post" + * ], + * "step_types": [ + * "llm" + * ] + * }, + * "selector": { + * "path": "output" + * }, + * "tags": [ + * "pii", + * "compliance" + * ] + * } + */ + 'ControlDefinition-Input': { + /** @description What action to take when control matches */ + action: components['schemas']['ControlAction']; /** - * Steering Context - * @description Steering context object for steer actions. Strongly recommended when decision='steer' to provide correction suggestions. If not provided, the evaluator result message will be used as fallback. + * Description + * @description Detailed description of the control */ - steering_context?: components['schemas']['SteeringContext'] | null; + description?: string | null; + /** + * Enabled + * @description Whether this control is active + * @default true + */ + enabled: boolean; + /** @description How to evaluate the selected data */ + evaluator: components['schemas']['EvaluatorSpec']; + /** + * Execution + * @description Where this control executes + * @enum {string} + */ + execution: 'server' | 'sdk'; + /** @description Which steps and stages this control applies to */ + scope?: components['schemas']['ControlScope']; + /** @description What data to select from the payload */ + selector: components['schemas']['ControlSelector']; + /** + * Tags + * @description Tags for categorization + */ + tags?: string[]; }; /** * ControlDefinition @@ -1188,7 +1273,7 @@ export interface components { * ] * } */ - ControlDefinition: { + 'ControlDefinition-Output': { /** @description What action to take when control matches */ action: components['schemas']['ControlAction']; /** @@ -1240,7 +1325,7 @@ export interface components { * control_name: Name of the control (denormalized for queries) * check_stage: "pre" (before execution) or "post" (after execution) * applies_to: "llm_call" or "tool_call" - * action: The action taken (allow, deny, steer, warn, log) + * action: The action taken (allow, deny, warn, log) * matched: Whether the control evaluator matched * confidence: Confidence score from the evaluator (0.0-1.0) * timestamp: When the control was executed (UTC) @@ -1388,6 +1473,8 @@ export interface components { control_name: string; /** @description Evaluator result (confidence, message, metadata) */ result: components['schemas']['EvaluatorResult']; + /** @description Steering context for steer actions if configured */ + steering_context?: components['schemas']['SteeringContext'] | null; }; /** * ControlScope @@ -1485,6 +1572,7 @@ export interface components { * non_match_count: Number of times the control did not match * allow_count: Number of allow actions * deny_count: Number of deny actions + * steer_count: Number of steer actions * warn_count: Number of warn actions * log_count: Number of log actions * error_count: Number of errors during evaluation @@ -1643,6 +1731,12 @@ export interface components { tags?: string[]; /** @description Agent using this control */ used_by_agent?: components['schemas']['AgentRef'] | null; + /** + * Used By Agents Count + * @description Number of unique agents using this control + * @default 0 + */ + used_by_agents_count: number; }; /** CreateControlRequest */ CreateControlRequest: { @@ -1711,9 +1805,19 @@ export interface components { DeleteControlResponse: { /** * Dissociated From - * @description Policy IDs the control was removed from before deletion + * @description Deprecated: policy IDs the control was removed from before deletion */ dissociated_from?: number[]; + /** + * Dissociated From Agents + * @description Agent names the control was removed from before deletion + */ + dissociated_from_agents?: string[]; + /** + * Dissociated From Policies + * @description Policy IDs the control was removed from before deletion + */ + dissociated_from_policies?: number[]; /** * Success * @description Whether the control was deleted @@ -1731,11 +1835,14 @@ export interface components { */ success: boolean; }; - /** DeletePolicyResponse */ + /** + * DeletePolicyResponse + * @description Compatibility response for singular policy deletion endpoint. + */ DeletePolicyResponse: { /** * Success - * @description Whether the policy was successfully removed + * @description Whether the request succeeded */ success: boolean; }; @@ -2205,6 +2312,14 @@ export interface components { */ total: number; }; + /** GetAgentPoliciesResponse */ + GetAgentPoliciesResponse: { + /** + * Policy Ids + * @description IDs of policies associated with the agent + */ + policy_ids?: number[]; + }; /** * GetAgentResponse * @description Response containing agent details and registered steps. @@ -2226,7 +2341,7 @@ export interface components { /** GetControlDataResponse */ GetControlDataResponse: { /** @description Control data payload */ - data: components['schemas']['ControlDefinition']; + data: components['schemas']['ControlDefinition-Output']; }; /** * GetControlResponse @@ -2234,7 +2349,7 @@ export interface components { */ GetControlResponse: { /** @description Control configuration data (None if not yet configured) */ - data?: components['schemas']['ControlDefinition'] | null; + data?: components['schemas']['ControlDefinition-Output'] | null; /** * Id * @description Control ID @@ -2257,11 +2372,14 @@ export interface components { */ control_ids: number[]; }; - /** GetPolicyResponse */ + /** + * GetPolicyResponse + * @description Compatibility response for singular policy retrieval endpoint. + */ GetPolicyResponse: { /** * Policy Id - * @description Identifier of the policy assigned to the agent + * @description Associated policy ID */ policy_id: number; }; @@ -2431,7 +2549,7 @@ export interface components { InitAgentResponse: { /** * Controls - * @description Active protection controls for the agent (if policy assigned) + * @description Active protection controls for the agent */ controls?: components['schemas']['Control'][]; /** @@ -2596,13 +2714,34 @@ export interface components { */ success: boolean; }; + /** + * RemoveAgentControlResponse + * @description Response for removing a direct agent-control association. + */ + RemoveAgentControlResponse: { + /** + * Control Still Active + * @description True if the control remains active via policy association(s) + */ + control_still_active: boolean; + /** + * Removed Direct Association + * @description True if a direct agent-control link was removed + */ + removed_direct_association: boolean; + /** + * Success + * @description Whether the request succeeded + */ + success: boolean; + }; /** * SetControlDataRequest * @description Request to update control configuration data. */ SetControlDataRequest: { /** @description Control configuration data (replaces existing) */ - data: components['schemas']['ControlDefinition']; + data: components['schemas']['ControlDefinition-Input']; }; /** SetControlDataResponse */ SetControlDataResponse: { @@ -2612,16 +2751,19 @@ export interface components { */ success: boolean; }; - /** SetPolicyResponse */ + /** + * SetPolicyResponse + * @description Compatibility response for singular policy assignment endpoint. + */ SetPolicyResponse: { /** * Old Policy Id - * @description Previous policy id if one was replaced + * @description Previously associated policy ID, if any */ old_policy_id?: number | null; /** * Success - * @description Whether the policy was successfully assigned + * @description Whether the request succeeded */ success: boolean; }; @@ -2710,6 +2852,26 @@ export interface components { */ timeseries?: components['schemas']['TimeseriesBucket'][] | null; }; + /** + * SteeringContext + * @description Steering context for steer actions. + * + * This model provides an extensible structure for steering guidance. + * Future fields could include severity, categories, suggested_actions, etc. + * @example { + * "message": "This large transfer requires user verification. Request 2FA code from user, verify it, then retry the transaction with verified_2fa=True." + * } + * @example { + * "message": "Transfer exceeds daily limit. Steps: 1) Ask user for business justification, 2) Request manager approval with amount and justification, 3) If approved, retry with manager_approved=True and justification filled in." + * } + */ + SteeringContext: { + /** + * Message + * @description Guidance message explaining what needs to be corrected and how + */ + message: string; + }; /** * Step * @description Runtime payload for an agent step invocation. @@ -2922,7 +3084,7 @@ export interface components { */ ValidateControlDataRequest: { /** @description Control configuration data to validate */ - data: components['schemas']['ControlDefinition']; + data: components['schemas']['ControlDefinition-Input']; }; /** ValidateControlDataResponse */ ValidateControlDataResponse: { @@ -3097,7 +3259,7 @@ export interface operations { }; requestBody?: never; responses: { - /** @description List of controls from agent's policy */ + /** @description List of controls from agent policy and direct associations */ 200: { headers: { [name: string]: unknown; @@ -3117,6 +3279,70 @@ export interface operations { }; }; }; + add_agent_control_api_v1_agents__agent_name__controls__control_id__post: { + parameters: { + query?: never; + header?: never; + path: { + agent_name: string; + control_id: number; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Success confirmation */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['AssocResponse']; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['HTTPValidationError']; + }; + }; + }; + }; + remove_agent_control_api_v1_agents__agent_name__controls__control_id__delete: { + parameters: { + query?: never; + header?: never; + path: { + agent_name: string; + control_id: number; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Success confirmation */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['RemoveAgentControlResponse']; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['HTTPValidationError']; + }; + }; + }; + }; list_agent_evaluators_api_v1_agents__agent_name__evaluators_get: { parameters: { query?: { @@ -3183,6 +3409,132 @@ export interface operations { }; }; }; + get_agent_policies_api_v1_agents__agent_name__policies_get: { + parameters: { + query?: never; + header?: never; + path: { + agent_name: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description List of policy IDs */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['GetAgentPoliciesResponse']; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['HTTPValidationError']; + }; + }; + }; + }; + remove_all_agent_policies_api_v1_agents__agent_name__policies_delete: { + parameters: { + query?: never; + header?: never; + path: { + agent_name: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Success confirmation */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['AssocResponse']; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['HTTPValidationError']; + }; + }; + }; + }; + add_agent_policy_api_v1_agents__agent_name__policies__policy_id__post: { + parameters: { + query?: never; + header?: never; + path: { + agent_name: string; + policy_id: number; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Success confirmation */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['AssocResponse']; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['HTTPValidationError']; + }; + }; + }; + }; + remove_agent_policy_api_v1_agents__agent_name__policies__policy_id__delete: { + parameters: { + query?: never; + header?: never; + path: { + agent_name: string; + policy_id: number; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Success confirmation */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['AssocResponse']; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + 'application/json': components['schemas']['HTTPValidationError']; + }; + }; + }; + }; get_agent_policy_api_v1_agents__agent_name__policy_get: { parameters: { query?: never; @@ -3422,7 +3774,7 @@ export interface operations { delete_control_api_v1_controls__control_id__delete: { parameters: { query?: { - /** @description If true, dissociate from all policies before deleting. If false, fail if control is associated with any policy. */ + /** @description If true, dissociate from all policy/agent links before deleting. If false, fail if control is associated with any policy or agent. */ force?: boolean; }; header?: never; diff --git a/ui/src/core/api/types.ts b/ui/src/core/api/types.ts index ca2f944e..5ac120a6 100644 --- a/ui/src/core/api/types.ts +++ b/ui/src/core/api/types.ts @@ -69,14 +69,20 @@ export type GetAgentResponse = components['schemas']['GetAgentResponse']; export type ControlActionDecision = components['schemas']['ControlAction']['decision']; export type ControlExecution = - components['schemas']['ControlDefinition']['execution']; + components['schemas']['ControlDefinition-Input']['execution']; export type ControlStage = NonNullable< components['schemas']['ControlScope']['stages'] >[number]; export type ControlScope = components['schemas']['ControlScope']; export type ControlSelector = components['schemas']['ControlSelector']; export type ControlAction = components['schemas']['ControlAction']; -export type ControlDefinition = components['schemas']['ControlDefinition']; +export type ControlDefinitionInput = + components['schemas']['ControlDefinition-Input']; +export type ControlDefinitionOutput = + components['schemas']['ControlDefinition-Output']; +export type ControlDefinition = + | ControlDefinitionInput + | ControlDefinitionOutput; export type Control = components['schemas']['Control']; export type AgentControlsResponse = components['schemas']['AgentControlsResponse']; @@ -93,19 +99,14 @@ export type SetControlDataResponse = export type GetControlDataResponse = components['schemas']['GetControlDataResponse']; -// Validate control data types (not yet in generated schemas) -// TODO: replace these with generated types after running pnpm fetch-api-types -export type ValidateControlDataRequest = { - data: ControlDefinition; -}; -export type ValidateControlDataResponse = { - success: boolean; -}; +export type ValidateControlDataRequest = + components['schemas']['ValidateControlDataRequest']; +export type ValidateControlDataResponse = + components['schemas']['ValidateControlDataResponse']; export type ControlSummary = components['schemas']['ControlSummary']; export type ListControlsResponse = components['schemas']['ListControlsResponse']; -// AgentRef - reference to an agent (for used_by_agents) export type AgentRef = components['schemas']['AgentRef']; // Helper type to extract query parameters from operations diff --git a/ui/src/core/hooks/query-hooks/use-add-control-to-agent.ts b/ui/src/core/hooks/query-hooks/use-add-control-to-agent.ts index 2a4e173d..40c9ce2c 100644 --- a/ui/src/core/hooks/query-hooks/use-add-control-to-agent.ts +++ b/ui/src/core/hooks/query-hooks/use-add-control-to-agent.ts @@ -13,11 +13,9 @@ type AddControlToAgentParams = { /** * Mutation hook to add a control to an agent * Flow: - * 1. Check if agent has a policy - * 2. If no policy, create one and assign it to the agent - * 3. Create the control - * 4. Set control data (definition) - * 5. Add control to policy + * 1. Create the control + * 2. Set control data (definition) + * 3. Associate the control directly with the agent */ export function useAddControlToAgent() { const queryClient = useQueryClient(); @@ -28,96 +26,73 @@ export function useAddControlToAgent() { controlName, definition, }: AddControlToAgentParams) => { - // Step 1: Check if agent has a policy - let policyId: number; - const { data: policyData, error: policyError } = - await api.agents.getPolicy(agentId); + let createdControlId: number | null = null; - if (policyError || !policyData?.policy_id) { - // Step 2: Create a new policy and assign it to the agent - const policyName = `policy-${agentId}`; + try { + // Step 1: Create the control const { - data: createPolicyResult, - error: createPolicyError, - response: createPolicyResponse, - } = await api.policies.create(policyName); + data: createControlResult, + error: createControlError, + response: createControlResponse, + } = await api.controls.create({ name: controlName }); - if (createPolicyError || !createPolicyResult) { + if (createControlError || !createControlResult) { throw parseApiError( - createPolicyError, - 'Failed to create policy', - createPolicyResponse?.status + createControlError, + 'Failed to create control', + createControlResponse?.status ); } - policyId = createPolicyResult.policy_id; + createdControlId = createControlResult.control_id; - // Assign policy to agent - const { error: assignError, response: assignResponse } = - await api.agents.setPolicy(agentId, policyId); + // Step 2: Set control data (definition) + const { error: setDataError, response: setDataResponse } = + await api.controls.setData(createdControlId, { + data: definition, + }); - if (assignError) { + if (setDataError) { throw parseApiError( - assignError, - 'Failed to assign policy to agent', - assignResponse?.status + setDataError, + 'Failed to set control data', + setDataResponse?.status ); } - } else { - policyId = policyData.policy_id; - } - - // Step 3: Create the control - const { - data: createControlResult, - error: createControlError, - response: createControlResponse, - } = await api.controls.create({ name: controlName }); - - if (createControlError || !createControlResult) { - throw parseApiError( - createControlError, - 'Failed to create control', - createControlResponse?.status - ); - } - const controlId = createControlResult.control_id; - - // Step 4: Set control data (definition) - const { error: setDataError, response: setDataResponse } = - await api.controls.setData(controlId, { - data: definition, - }); - - if (setDataError) { - throw parseApiError( - setDataError, - 'Failed to set control data', - setDataResponse?.status - ); - } + // Step 3: Associate control directly with the agent. + const { error: associateError, response: associateResponse } = + await api.agents.addControl(agentId, createdControlId); - // Step 5: Add control to policy - const { error: addControlError, response: addControlResponse } = - await api.policies.addControl(policyId, controlId); + if (associateError) { + throw parseApiError( + associateError, + 'Failed to add control to agent', + associateResponse?.status + ); + } - if (addControlError) { - throw parseApiError( - addControlError, - 'Failed to add control to policy', - addControlResponse?.status - ); + return { controlId: createdControlId }; + } catch (error) { + // Best effort cleanup: avoid orphan controls if a later step fails. + if (createdControlId !== null) { + try { + await api.controls.delete(createdControlId, { force: true }); + } catch { + // Preserve the original error from the primary flow. + } + } + throw error; } - - return { controlId, policyId }; }, onSuccess: (_data, variables) => { // Invalidate relevant queries to refetch data - queryClient.invalidateQueries({ queryKey: ['controls'] }); queryClient.invalidateQueries({ queryKey: ['agent', variables.agentId] }); queryClient.invalidateQueries({ - queryKey: ['agentControls', variables.agentId], + queryKey: ['agent', variables.agentId, 'controls'], + }); + queryClient.invalidateQueries({ + queryKey: ['controls', 'infinite'], }); // Invalidate agents list query to refresh active controls count queryClient.invalidateQueries({ diff --git a/ui/src/core/hooks/query-hooks/use-agent-controls.ts b/ui/src/core/hooks/query-hooks/use-agent-controls.ts index 0cb18df3..069387ea 100644 --- a/ui/src/core/hooks/query-hooks/use-agent-controls.ts +++ b/ui/src/core/hooks/query-hooks/use-agent-controls.ts @@ -9,15 +9,15 @@ import type { /** * Query hook to fetch active controls for an agent * - * @param agentName - Immutable agent name (required) + * @param agentId - UUID of the agent (required) */ export function useAgentControls( - agentName: GetAgentControlsPathParams['agent_name'] + agentId: GetAgentControlsPathParams['agent_name'] ) { return useQuery({ - queryKey: ['agent', agentName, 'controls'], + queryKey: ['agent', agentId, 'controls'], queryFn: async () => { - const { data, error } = await api.agents.getControls(agentName); + const { data, error } = await api.agents.getControls(agentId); if (error) throw error; return data; }, diff --git a/ui/src/core/hooks/query-hooks/use-agent.ts b/ui/src/core/hooks/query-hooks/use-agent.ts index e01f8a30..ee5bb656 100644 --- a/ui/src/core/hooks/query-hooks/use-agent.ts +++ b/ui/src/core/hooks/query-hooks/use-agent.ts @@ -5,13 +5,13 @@ import { api } from '@/core/api/client'; import type { GetAgentPathParams, GetAgentResponse } from '@/core/api/types'; /** - * Query hook to fetch a single agent by identifier. + * Query hook to fetch a single agent by ID * - * @param agentName - Immutable agent name (required) + * @param agentId - UUID of the agent (required) * */ export function useAgent( - agentName: GetAgentPathParams['agent_name'], + agentId: GetAgentPathParams['agent_name'], options?: Omit< UseQueryOptions< GetAgentResponse, @@ -23,12 +23,12 @@ export function useAgent( > ): UseQueryResult { const { enabled, ...rest } = options ?? {}; - const isEnabled = enabled ?? Boolean(agentName); + const isEnabled = enabled ?? Boolean(agentId); return useQuery({ - queryKey: ['agent', agentName], + queryKey: ['agent', agentId], queryFn: async () => { - const { data, error } = await api.agents.get(agentName); + const { data, error } = await api.agents.get(agentId); if (error) throw error; return data; }, diff --git a/ui/src/core/hooks/query-hooks/use-delete-control.ts b/ui/src/core/hooks/query-hooks/use-delete-control.ts deleted file mode 100644 index 719efab4..00000000 --- a/ui/src/core/hooks/query-hooks/use-delete-control.ts +++ /dev/null @@ -1,45 +0,0 @@ -import { useMutation, useQueryClient } from '@tanstack/react-query'; - -import { api } from '@/core/api/client'; -import { parseApiError } from '@/core/api/errors'; - -type DeleteControlParams = { - agentId: string; - controlId: number; - /** If true, dissociate from all policies before deleting. Default true so control can be removed from agent. */ - force?: boolean; -}; - -/** - * Mutation hook to delete a control. - * Use force: true when deleting from agent detail so the control is removed from the policy and then deleted. - */ -export function useDeleteControl() { - const queryClient = useQueryClient(); - - return useMutation({ - mutationFn: async ({ controlId, force = true }: DeleteControlParams) => { - const { data, error, response } = await api.controls.delete(controlId, { - force, - }); - - if (error) { - throw parseApiError( - error, - 'Failed to delete control', - response?.status - ); - } - - return data; - }, - onSuccess: (_data, variables) => { - queryClient.invalidateQueries({ - queryKey: ['agent', variables.agentId, 'controls'], - }); - queryClient.invalidateQueries({ - queryKey: ['agents', 'infinite'], - }); - }, - }); -} diff --git a/ui/src/core/hooks/query-hooks/use-remove-control-from-agent.ts b/ui/src/core/hooks/query-hooks/use-remove-control-from-agent.ts new file mode 100644 index 00000000..93663216 --- /dev/null +++ b/ui/src/core/hooks/query-hooks/use-remove-control-from-agent.ts @@ -0,0 +1,59 @@ +import { useMutation, useQueryClient } from '@tanstack/react-query'; + +import { api } from '@/core/api/client'; +import { parseApiError } from '@/core/api/errors'; + +type RemoveControlFromAgentParams = { + agentId: string; + controlId: number; +}; + +export type RemoveControlFromAgentResult = { + success: boolean; + removed_direct_association: boolean; + control_still_active: boolean; +}; + +/** + * Mutation hook to remove a control from a specific agent. + */ +export function useRemoveControlFromAgent() { + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async ({ + agentId, + controlId, + }: RemoveControlFromAgentParams) => { + const { data, error, response } = await api.agents.removeControl( + agentId, + controlId + ); + + if (error || !data) { + throw parseApiError( + error, + 'Failed to remove control from agent', + response?.status + ); + } + + return { + success: data.success, + removed_direct_association: data.removed_direct_association, + control_still_active: data.control_still_active, + } satisfies RemoveControlFromAgentResult; + }, + onSuccess: (_data, variables) => { + queryClient.invalidateQueries({ + queryKey: ['agent', variables.agentId, 'controls'], + }); + queryClient.invalidateQueries({ + queryKey: ['controls', 'infinite'], + }); + queryClient.invalidateQueries({ + queryKey: ['agents', 'infinite'], + }); + }, + }); +} diff --git a/ui/src/core/page-components/agent-detail/agent-detail.tsx b/ui/src/core/page-components/agent-detail/agent-detail.tsx index aa251f59..ecbee7e3 100644 --- a/ui/src/core/page-components/agent-detail/agent-detail.tsx +++ b/ui/src/core/page-components/agent-detail/agent-detail.tsx @@ -74,7 +74,7 @@ const AgentDetailPage = ({ agentId, defaultTab }: AgentDetailPageProps) => { // the close animation; the effect syncs selectedControl when the modal opens. }; - const { handleDeleteControl, deleteControl } = useDeleteControlFlow({ + const { handleDeleteControl, removeControlFromAgent } = useDeleteControlFlow({ agentId, selectedControl, onCloseEditModal: handleCloseEditModal, @@ -133,7 +133,7 @@ const AgentDetailPage = ({ agentId, defaultTab }: AgentDetailPageProps) => { const columns = useControlsTableColumns({ agentId, updateControl, - deleteControl, + removeControlFromAgent, onEditControl: handleEditControl, onDeleteControl: handleDeleteControl, }); diff --git a/ui/src/core/page-components/agent-detail/controls/table-columns.tsx b/ui/src/core/page-components/agent-detail/controls/table-columns.tsx index 3eb438d7..745358a3 100644 --- a/ui/src/core/page-components/agent-detail/controls/table-columns.tsx +++ b/ui/src/core/page-components/agent-detail/controls/table-columns.tsx @@ -6,7 +6,7 @@ import { type ColumnDef } from '@tanstack/react-table'; import { useMemo } from 'react'; import type { Control } from '@/core/api/types'; -import type { useDeleteControl } from '@/core/hooks/query-hooks/use-delete-control'; +import type { useRemoveControlFromAgent } from '@/core/hooks/query-hooks/use-remove-control-from-agent'; import type { useUpdateControl } from '@/core/hooks/query-hooks/use-update-control'; import { getStepTypeLabelAndColor } from './utils'; @@ -14,7 +14,7 @@ import { getStepTypeLabelAndColor } from './utils'; type UseControlsTableColumnsParams = { agentId: string; updateControl: ReturnType; - deleteControl: ReturnType; + removeControlFromAgent: ReturnType; onEditControl: (control: Control) => void; onDeleteControl: (control: Control) => void; }; @@ -22,7 +22,7 @@ type UseControlsTableColumnsParams = { export function useControlsTableColumns({ agentId, updateControl, - deleteControl, + removeControlFromAgent, onEditControl, onDeleteControl, }: UseControlsTableColumnsParams): ColumnDef[] { @@ -190,8 +190,8 @@ export function useControlsTableColumns({ cell: ({ row }: { row: { original: Control } }) => { const control = row.original; const isDeleting = - deleteControl.isPending && - deleteControl.variables?.controlId === control.id; + removeControlFromAgent.isPending && + removeControlFromAgent.variables?.controlId === control.id; return ( onDeleteControl(control)} - aria-label="Delete control" + aria-label="Remove control from agent" disabled={isDeleting} > @@ -212,8 +212,8 @@ export function useControlsTableColumns({ [ agentId, updateControl, - deleteControl.isPending, - deleteControl.variables?.controlId, + removeControlFromAgent.isPending, + removeControlFromAgent.variables?.controlId, onEditControl, onDeleteControl, ] diff --git a/ui/src/core/page-components/agent-detail/controls/use-delete-control-flow.tsx b/ui/src/core/page-components/agent-detail/controls/use-delete-control-flow.tsx index d9123740..5dcfef73 100644 --- a/ui/src/core/page-components/agent-detail/controls/use-delete-control-flow.tsx +++ b/ui/src/core/page-components/agent-detail/controls/use-delete-control-flow.tsx @@ -3,7 +3,10 @@ import { modals } from '@mantine/modals'; import { notifications } from '@mantine/notifications'; import type { Control } from '@/core/api/types'; -import { useDeleteControl } from '@/core/hooks/query-hooks/use-delete-control'; +import { + type RemoveControlFromAgentResult, + useRemoveControlFromAgent, +} from '@/core/hooks/query-hooks/use-remove-control-from-agent'; type UseDeleteControlFlowParams = { agentId: string; @@ -16,18 +19,19 @@ export function useDeleteControlFlow({ selectedControl, onCloseEditModal, }: UseDeleteControlFlowParams) { - const deleteControl = useDeleteControl(); + const removeControlFromAgent = useRemoveControlFromAgent(); const handleDeleteControl = (control: Control) => { modals.openConfirmModal({ - title: 'Delete control?', + title: 'Remove control from agent?', children: ( - Delete "{control.name}"? This will remove the control from - this agent and delete it. This action cannot be undone. + Remove "{control.name}" from this agent? This only removes + the direct association from this agent and does not delete the control + globally. ), - labels: { confirm: 'Delete', cancel: 'Cancel' }, + labels: { confirm: 'Remove', cancel: 'Cancel' }, confirmProps: { variant: 'filled', color: 'red.7', @@ -35,17 +39,31 @@ export function useDeleteControlFlow({ }, cancelProps: { variant: 'default', size: 'sm' }, onConfirm: () => - deleteControl.mutate( + removeControlFromAgent.mutate( { agentId, controlId: control.id, - force: true, }, { - onSuccess: () => { + onSuccess: (result: RemoveControlFromAgentResult) => { + // The controls table shows active controls (direct + policy-derived), so + // remove-direct can legitimately no-op for policy-derived entries. + if (!result.removed_direct_association) { + notifications.show({ + title: 'No direct association found', + message: result.control_still_active + ? `"${control.name}" is still active for this agent through policy associations.` + : `"${control.name}" was not directly associated with this agent.`, + color: result.control_still_active ? 'yellow' : 'blue', + }); + return; + } + notifications.show({ - title: 'Control deleted', - message: `"${control.name}" has been removed.`, + title: 'Control removed', + message: result.control_still_active + ? `"${control.name}" was removed from direct associations, but remains active via policy associations.` + : `"${control.name}" has been removed from this agent.`, color: 'green', }); if (selectedControl?.id === control.id) { @@ -54,7 +72,7 @@ export function useDeleteControlFlow({ }, onError: (error) => { notifications.show({ - title: 'Failed to delete control', + title: 'Failed to remove control', message: error instanceof Error ? error.message @@ -67,5 +85,5 @@ export function useDeleteControlFlow({ }); }; - return { handleDeleteControl, deleteControl }; + return { handleDeleteControl, removeControlFromAgent }; } diff --git a/ui/src/core/page-components/agent-detail/modals/control-store/index.tsx b/ui/src/core/page-components/agent-detail/modals/control-store/index.tsx index de09c575..b6b6a79d 100644 --- a/ui/src/core/page-components/agent-detail/modals/control-store/index.tsx +++ b/ui/src/core/page-components/agent-detail/modals/control-store/index.tsx @@ -21,11 +21,7 @@ import { useEffect, useMemo, useRef, useState } from 'react'; import { ErrorBoundary } from '@/components/error-boundary'; import { api } from '@/core/api/client'; -import type { - AgentRef, - ControlDefinition, - ControlSummary, -} from '@/core/api/types'; +import type { ControlDefinition, ControlSummary } from '@/core/api/types'; import { SearchInput } from '@/core/components/search-input'; import { MODAL_NAMES, SUBMODAL_NAMES } from '@/core/constants/modal-routes'; import { useControlsInfinite } from '@/core/hooks/query-hooks/use-controls-infinite'; @@ -37,11 +33,6 @@ import { AddNewControlModal } from '../add-new-control'; import { EditControlContent } from '../edit-control/edit-control-content'; import { sanitizeControlNamePart } from '../edit-control/utils'; -// Extended ControlSummary with optional used_by_agent -type ControlSummaryWithAgent = ControlSummary & { - used_by_agent?: AgentRef | null; -}; - type ControlStoreModalProps = { opened: boolean; onClose: () => void; @@ -176,8 +167,7 @@ export function ControlStoreModal({ }; loadControl(); } - // eslint-disable-next-line react-hooks/exhaustive-deps - }, [editModalOpened, controlId, selectedControl]); + }, [editModalOpened, controlId, selectedControl, controls]); // Clear selectedControl when edit modal closes useEffect(() => { @@ -276,32 +266,25 @@ export function ControlStoreModal({ }, { id: 'agent', - header: 'Agent', + header: 'Used by', size: 150, cell: ({ row }) => { - const agent = (row.original as ControlSummaryWithAgent).used_by_agent; - const control = row.original; - if (!agent) { + const usedByAgent = row.original.used_by_agent; + if (!usedByAgent) { return ( — ); } - // Link to agent controls tab with control name filter - const href = `/agents/${agent.agent_name}/controls?q=${encodeURIComponent(control.name)}`; return ( { - e.stopPropagation(); - // Close modal when navigating to agent page - onClose(); - }} + underline="hover" > - {agent.agent_name} + {usedByAgent.agent_name} ); }, diff --git a/ui/src/core/page-components/home/home.tsx b/ui/src/core/page-components/home/home.tsx index aa810740..ff363cb0 100644 --- a/ui/src/core/page-components/home/home.tsx +++ b/ui/src/core/page-components/home/home.tsx @@ -45,6 +45,7 @@ from agent_control import control, ControlViolationError agent_control.init( agent_name="support-agent-v1", + agent_description="Customer Support Agent", server_url="http://localhost:8000", ) diff --git a/ui/tests/agent-detail.spec.ts b/ui/tests/agent-detail.spec.ts index 938dec8a..d75fe57c 100644 --- a/ui/tests/agent-detail.spec.ts +++ b/ui/tests/agent-detail.spec.ts @@ -3,7 +3,7 @@ import type { AgentControlsResponse, GetAgentResponse } from '@/core/api/types'; import { expect, mockData, mockRoutes, test } from './fixtures'; test.describe('Agent Detail Page', () => { - const agentId = 'customer-support-bot'; + const agentId = 'agent-1'; const agentUrl = `/agents/${agentId}/controls`; // Type-safe access to mock agent data @@ -412,17 +412,24 @@ test.describe('Agent Detail Page', () => { } ); - await mockedPage.route('**/api/v1/controls/*', async (route, request) => { - if (request.method() === 'DELETE') { - await route.fulfill({ - status: 200, - contentType: 'application/json', - body: JSON.stringify({ success: true }), - }); - } else { - await route.continue(); + await mockedPage.route( + '**/api/v1/agents/*/controls/*', + async (route, request) => { + if (request.method() === 'DELETE') { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + removed_direct_association: true, + control_still_active: false, + }), + }); + } else { + await route.continue(); + } } - }); + ); await mockedPage.goto(agentUrl); @@ -439,32 +446,33 @@ test.describe('Agent Detail Page', () => { await targetRow.scrollIntoViewIfNeeded(); const deleteButton = targetRow.getByRole('button', { - name: 'Delete control', + name: 'Remove control from agent', }); await deleteButton.click(); const confirmModal = mockedPage.getByRole('dialog', { - name: /Delete control\?/i, + name: /Remove control from agent\?/i, }); await expect(confirmModal).toBeVisible(); await expect( - confirmModal.getByText(/This action cannot be undone/) + confirmModal.getByText(/does not delete the control globally/i) ).toBeVisible(); const deleteRequest = mockedPage.waitForRequest( (request) => request.method() === 'DELETE' && - new RegExp(`/api/v1/controls/${deletedControlId}(\\?|$)`).test( - request.url() - ), + new RegExp( + `/api/v1/agents/[^/]+/controls/${deletedControlId}(\\?|$)` + ).test(request.url()), { timeout: 5000 } ); - await confirmModal.getByRole('button', { name: 'Delete' }).click(); + await confirmModal.getByRole('button', { name: 'Remove' }).click(); const request = await deleteRequest; - expect(request.url()).toContain(`/api/v1/controls/${deletedControlId}`); - expect(new URL(request.url()).searchParams.get('force')).toBe('true'); + expect(request.url()).toMatch( + new RegExp(`/api/v1/agents/[^/]+/controls/${deletedControlId}(\\?|$)`) + ); expect(request.method()).toBe('DELETE'); await expect(confirmModal).not.toBeVisible({ timeout: 5000 }); @@ -638,7 +646,7 @@ test.describe('Agent Detail - Empty State', () => { }); }); - await page.goto('/agents/customer-support-bot/controls'); + await page.goto('/agents/agent-1/controls'); // Check for empty state message await expect(page.getByText('No controls configured')).toBeVisible(); diff --git a/ui/tests/agent-stats.spec.ts b/ui/tests/agent-stats.spec.ts index 3cbe183d..872a251d 100644 --- a/ui/tests/agent-stats.spec.ts +++ b/ui/tests/agent-stats.spec.ts @@ -3,7 +3,7 @@ import { expect, mockData, mockRoutes, test } from './fixtures'; test.describe('Agent Monitor Tab', () => { test.beforeEach(async ({ mockedPage }) => { // Navigate to agent detail page - await mockedPage.goto('/agents/customer-support-bot/monitor'); + await mockedPage.goto('/agents/agent-1/monitor'); // Wait for the page to load await expect( mockedPage.getByRole('heading', { name: 'customer-support-bot' }) @@ -148,7 +148,7 @@ test.describe('Agent Monitor Tab - Empty State', () => { await mockRoutes.stats(page, { data: mockData.emptyStats }); // Navigate to agent detail page - await page.goto('/agents/customer-support-bot/monitor'); + await page.goto('/agents/agent-1/monitor'); await expect( page.getByRole('heading', { name: 'customer-support-bot' }) ).toBeVisible(); @@ -237,7 +237,7 @@ test.describe('Agent Monitor Tab - Refetch Flow', () => { }); // Navigate to agent detail page - await page.goto('/agents/customer-support-bot/monitor'); + await page.goto('/agents/agent-1/monitor'); await expect( page.getByRole('heading', { name: 'customer-support-bot' }) ).toBeVisible(); @@ -274,7 +274,7 @@ test.describe('Agent Monitor Tab - Error State', () => { }); // Navigate to agent detail page - await page.goto('/agents/customer-support-bot/monitor'); + await page.goto('/agents/agent-1/monitor'); await expect( page.getByRole('heading', { name: 'customer-support-bot' }) ).toBeVisible(); diff --git a/ui/tests/control-store.spec.ts b/ui/tests/control-store.spec.ts index ed984ed4..6f56b695 100644 --- a/ui/tests/control-store.spec.ts +++ b/ui/tests/control-store.spec.ts @@ -2,7 +2,7 @@ import type { Page } from '@playwright/test'; import { expect, mockData, test } from './fixtures'; -const agentUrl = '/agents/customer-support-bot/controls'; +const agentUrl = '/agents/agent-1/controls'; async function openControlStoreModal(page: Page) { await page.goto(agentUrl); @@ -48,7 +48,7 @@ test.describe('Control Store Modal', () => { ).toBeVisible(); // Status dot column has no header text, so we skip checking for it await expect( - modal.getByRole('columnheader', { name: 'Agent' }) + modal.getByRole('columnheader', { name: 'Used by' }) ).toBeVisible(); for (const control of mockData.listControls.controls) { @@ -58,18 +58,24 @@ test.describe('Control Store Modal', () => { } }); - test('displays agent links in Agent column', async ({ mockedPage }) => { + test('displays usage attribution in Used by column', async ({ + mockedPage, + }) => { const modal = await openControlStoreModal(mockedPage); - // PII Detection is used by customer-support-bot - const agentLink = modal - .getByRole('link', { name: 'customer-support-bot' }) - .first(); - await expect(agentLink).toBeVisible(); - // Link includes query param to filter by control name - await expect(agentLink).toHaveAttribute( + // Two controls show direct agent-name attribution in fixture data. + await expect(modal.getByText('customer-support-bot')).toBeVisible(); + await expect(modal.getByText('data-analysis-agent')).toBeVisible(); + // One control has no usage and renders as an em dash + await expect(modal.getByText('—')).toBeVisible(); + // Agent attribution renders as navigation links. + const customerSupportLink = modal.getByRole('link', { + name: 'customer-support-bot', + }); + await expect(customerSupportLink).toHaveCount(1); + await expect(customerSupportLink).toHaveAttribute( 'href', - '/agents/customer-support-bot/controls?q=PII%20Detection' + '/agents/customer-support-bot' ); }); @@ -478,22 +484,15 @@ test.describe('Modal Routing', () => { }); await expect(createModal).toBeVisible(); - // Mock successful API response for control creation - // Agent already has a policy (return 200 with policy_id) + // Mock successful API response for direct agent-control association await mockedPage.route( - '**/api/v1/agents/*/policy', + '**/api/v1/agents/*/controls/*', async (route, request) => { - if (request.method() === 'GET') { - await route.fulfill({ - status: 200, - contentType: 'application/json', - body: JSON.stringify({ policy_id: 1 }), - }); - } else if (request.method() === 'POST') { + if (request.method() === 'POST') { await route.fulfill({ status: 200, contentType: 'application/json', - body: JSON.stringify({}), + body: JSON.stringify({ success: true }), }); } else { await route.continue(); @@ -528,21 +527,6 @@ test.describe('Modal Routing', () => { } ); - await mockedPage.route( - '**/api/v1/policies/*/controls/*', - async (route, request) => { - if (request.method() === 'POST') { - await route.fulfill({ - status: 200, - contentType: 'application/json', - body: JSON.stringify({}), - }); - } else { - await route.continue(); - } - } - ); - // Fill out the form and submit const controlNameInput = createModal.getByPlaceholder('Enter control name'); await controlNameInput.fill('Test Control'); @@ -568,7 +552,7 @@ test.describe('Modal Routing', () => { // Start waiting for API response before clicking (must be set up before the action) const responsePromise = mockedPage.waitForResponse( - '**/api/v1/policies/*/controls/*', + '**/api/v1/agents/*/controls/*', { timeout: 10000 } ); await confirmButton.click(); @@ -611,13 +595,20 @@ test.describe('Modal Routing', () => { await expect(controlNameInput).toHaveValue(/.*-copy$/, { timeout: 5000 }); // Set up mock routes for control creation flow (copying creates a new control) - await mockedPage.route('**/api/v1/agents/*/policy', async (route) => { - await route.fulfill({ - status: 200, - contentType: 'application/json', - body: JSON.stringify({ policy_id: 1 }), - }); - }); + await mockedPage.route( + '**/api/v1/agents/*/controls/*', + async (route, request) => { + if (request.method() === 'POST') { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ success: true }), + }); + } else { + await route.continue(); + } + } + ); await mockedPage.route('**/api/v1/controls', async (route, request) => { if (request.method() === 'PUT') { @@ -646,21 +637,6 @@ test.describe('Modal Routing', () => { } ); - await mockedPage.route( - '**/api/v1/policies/*/controls/*', - async (route, request) => { - if (request.method() === 'POST') { - await route.fulfill({ - status: 200, - contentType: 'application/json', - body: JSON.stringify({}), - }); - } else { - await route.continue(); - } - } - ); - // Submit the form const saveButton = editModal.getByRole('button', { name: /Save|Create/i }); await saveButton.click(); @@ -674,7 +650,7 @@ test.describe('Modal Routing', () => { // Start waiting for API response before clicking (must be set up before the action) const responsePromise = mockedPage.waitForResponse( - '**/api/v1/policies/*/controls/*', + '**/api/v1/agents/*/controls/*', { timeout: 10000 } ); await confirmButton.click(); @@ -689,6 +665,119 @@ test.describe('Modal Routing', () => { // URL should not contain any modal parameters await expect(mockedPage).not.toHaveURL(/.*\?modal=/); }); + + test('cleans up created control when agent association fails', async ({ + mockedPage, + }) => { + await mockedPage.goto( + `${agentUrl}?modal=control-store&submodal=create&evaluator=list` + ); + + const createModal = mockedPage.getByRole('dialog', { + name: 'Create Control', + }); + await expect(createModal).toBeVisible(); + + let cleanupDeleteCalls = 0; + + await mockedPage.route( + '**/api/v1/agents/*/controls/*', + async (route, request) => { + if (request.method() === 'POST') { + await route.fulfill({ + status: 400, + contentType: 'application/json', + body: JSON.stringify({ + detail: 'Control is incompatible with this agent', + error_code: 'POLICY_CONTROL_INCOMPATIBLE', + }), + }); + } else { + await route.continue(); + } + } + ); + + await mockedPage.route('**/api/v1/controls', async (route, request) => { + if (request.method() === 'PUT') { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ control_id: 100 }), + }); + } else { + await route.continue(); + } + }); + + await mockedPage.route( + '**/api/v1/controls/*/data', + async (route, request) => { + if (request.method() === 'PUT') { + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ success: true }), + }); + } else { + await route.continue(); + } + } + ); + + await mockedPage.route('**/api/v1/controls/*', async (route, request) => { + if (request.method() === 'DELETE') { + cleanupDeleteCalls += 1; + await route.fulfill({ + status: 200, + contentType: 'application/json', + body: JSON.stringify({ + success: true, + dissociated_from_policies: [], + dissociated_from_agents: [], + }), + }); + } else { + await route.continue(); + } + }); + + const controlNameInput = createModal.getByPlaceholder('Enter control name'); + await controlNameInput.fill('Cleanup Test Control'); + const valuesTextarea = createModal.getByPlaceholder( + 'Enter values (one per line)' + ); + await valuesTextarea.fill('value-1'); + + const saveButton = createModal.getByRole('button', { + name: /Save|Create/i, + }); + await saveButton.click(); + + const confirmButton = mockedPage.locator("button:has-text('Confirm')"); + await expect(confirmButton).toBeVisible({ timeout: 5000 }); + + const failedAssociationPromise = mockedPage.waitForResponse( + (response) => + response.request().method() === 'POST' && + /\/api\/v1\/agents\/[^/]+\/controls\/\d+/.test(response.url()) && + response.status() === 400, + { timeout: 10000 } + ); + const cleanupDeletePromise = mockedPage.waitForResponse( + (response) => + response.request().method() === 'DELETE' && + /\/api\/v1\/controls\/\d+\?force=true/.test(response.url()) && + response.status() === 200, + { timeout: 10000 } + ); + + await confirmButton.click(); + await failedAssociationPromise; + await cleanupDeletePromise; + + expect(cleanupDeleteCalls).toBe(1); + }); }); test.describe('Control Store - Loading States', () => { @@ -720,7 +809,7 @@ test.describe('Control Store - Loading States', () => { }); }); - await page.goto('/agents/customer-support-bot/controls'); + await page.goto('/agents/agent-1/controls'); // Open the control store modal await page.getByTestId('add-control-button').first().click(); diff --git a/ui/tests/evaluators/helpers.ts b/ui/tests/evaluators/helpers.ts index 0ddd940a..40e38037 100644 --- a/ui/tests/evaluators/helpers.ts +++ b/ui/tests/evaluators/helpers.ts @@ -4,7 +4,7 @@ import { expect, type Page } from '@playwright/test'; -const AGENT_URL = '/agents/customer-support-bot/controls'; +const AGENT_URL = '/agents/agent-1/controls'; /** * Opens the control store and selects an evaluator to create a new control diff --git a/ui/tests/fixtures.ts b/ui/tests/fixtures.ts index 1eb14a23..04263ed5 100644 --- a/ui/tests/fixtures.ts +++ b/ui/tests/fixtures.ts @@ -21,7 +21,7 @@ import type { StatsResponse } from '@/core/hooks/query-hooks/use-agent-monitor'; const agentsList: AgentSummary[] = [ { agent_name: 'customer-support-bot', - policy_id: 1, + policy_ids: [1], created_at: '2024-01-01T00:00:00Z', step_count: 5, evaluator_count: 2, @@ -29,7 +29,7 @@ const agentsList: AgentSummary[] = [ }, { agent_name: 'data-analysis-agent', - policy_id: 2, + policy_ids: [2], created_at: '2024-01-02T00:00:00Z', step_count: 3, evaluator_count: 1, @@ -37,7 +37,7 @@ const agentsList: AgentSummary[] = [ }, { agent_name: 'code-review-assistant', - policy_id: 3, + policy_ids: [3], created_at: '2024-01-03T00:00:00Z', step_count: 8, evaluator_count: 4, @@ -156,6 +156,7 @@ const controlSummariesList: (ControlSummary & { stages: ['post'], tags: ['pii', 'compliance'], used_by_agent: { agent_name: 'customer-support-bot' }, + used_by_agents_count: 1, }, { id: 2, @@ -167,6 +168,7 @@ const controlSummariesList: (ControlSummary & { stages: ['pre'], tags: ['security'], used_by_agent: { agent_name: 'data-analysis-agent' }, + used_by_agents_count: 1, }, { id: 3, @@ -178,6 +180,7 @@ const controlSummariesList: (ControlSummary & { stages: ['pre'], tags: [], used_by_agent: null, + used_by_agents_count: 0, }, ]; diff --git a/ui/tests/search-input.spec.ts b/ui/tests/search-input.spec.ts index 4ec1f8d2..92b48b40 100644 --- a/ui/tests/search-input.spec.ts +++ b/ui/tests/search-input.spec.ts @@ -88,7 +88,9 @@ test.describe('SearchInput - Query Param Syncing', () => { // Navigate away await mockedPage.getByText('customer-support-bot').click(); - await expect(mockedPage).toHaveURL(/\/agents\/customer-support-bot/); + await expect(mockedPage).toHaveURL( + /\/agents\/customer-support-bot\/monitor/ + ); // Go back await mockedPage.goBack(); @@ -111,7 +113,7 @@ test.describe('SearchInput - Query Param Syncing', () => { test.describe('SearchInput - Agent Detail Page', () => { test('syncs search value to URL query param (q)', async ({ mockedPage }) => { - await mockedPage.goto('/agents/customer-support-bot/controls'); + await mockedPage.goto('/agents/agent-1/controls'); const searchInput = mockedPage.getByPlaceholder('Search controls...'); await searchInput.fill('PII'); @@ -124,7 +126,7 @@ test.describe('SearchInput - Agent Detail Page', () => { }); test('reads search value from URL on page load', async ({ mockedPage }) => { - await mockedPage.goto('/agents/customer-support-bot/controls?q=PII'); + await mockedPage.goto('/agents/agent-1/controls?q=PII'); // Wait for page to load await expect(mockedPage.getByRole('table')).toBeVisible();