diff --git a/examples/agent_control_demo/setup_controls.py b/examples/agent_control_demo/setup_controls.py index 09811fcb..3c3a97c1 100644 --- a/examples/agent_control_demo/setup_controls.py +++ b/examples/agent_control_demo/setup_controls.py @@ -77,10 +77,10 @@ async def create_control( ) -> int: """Create a control with the given definition.""" try: - # Step 1: Create the control (just the name) + # Create and validate the control atomically. response = await client.http_client.put( "/api/v1/controls", - json={"name": name} + json={"name": name, "data": control_definition}, ) if response.status_code == 409: @@ -93,13 +93,6 @@ async def create_control( response.raise_for_status() control_id = response.json().get("control_id") - # Step 2: Set the control data - response = await client.http_client.put( - f"/api/v1/controls/{control_id}/data", - json={"data": control_definition} - ) - response.raise_for_status() - print(f"✓ Created control '{name}' with ID: {control_id}") return control_id diff --git a/examples/deepeval/setup_controls.py b/examples/deepeval/setup_controls.py index 86caad58..faddb15e 100755 --- a/examples/deepeval/setup_controls.py +++ b/examples/deepeval/setup_controls.py @@ -193,32 +193,38 @@ async def setup_demo(quiet: bool = False): description = control_spec["description"] try: - # Create control + # Create and validate the control atomically. resp = await client.put( "/api/v1/controls", - json={"name": control_name}, + json={"name": control_name, "data": definition}, ) - if resp.status_code == 409: + control_exists = resp.status_code == 409 + if control_exists: # Control exists, get its ID resp = await client.get("/api/v1/controls", params={"name": control_name}) resp.raise_for_status() - controls = resp.json().get("controls", []) + controls = [ + control + for control in resp.json().get("controls", []) + if control.get("name") == control_name + ] if controls: control_id = controls[0]["id"] controls_updated += 1 else: + print(f" ❌ Could not find exact control match for '{control_name}'") continue else: resp.raise_for_status() control_id = resp.json()["control_id"] controls_created += 1 - # Set control definition - resp = await client.put( - f"/api/v1/controls/{control_id}/data", - json={"data": definition}, - ) - resp.raise_for_status() + if control_exists: + resp = await client.put( + f"/api/v1/controls/{control_id}/data", + json={"data": definition}, + ) + resp.raise_for_status() # Associate control directly with the agent resp = await client.post(f"/api/v1/agents/{agent_name}/controls/{control_id}") diff --git a/models/src/agent_control_models/server.py b/models/src/agent_control_models/server.py index cdc48cdc..bf388707 100644 --- a/models/src/agent_control_models/server.py +++ b/models/src/agent_control_models/server.py @@ -119,6 +119,10 @@ class CreateControlRequest(BaseModel): ..., description="Unique control name (letters, numbers, hyphens, underscores)", ) + data: ControlDefinition = Field( + ..., + description="Control definition to validate and store during creation", + ) class InitAgentRequest(BaseModel): diff --git a/sdks/python/ARCHITECTURE.md b/sdks/python/ARCHITECTURE.md index 6c1d8db3..ff67180b 100644 --- a/sdks/python/ARCHITECTURE.md +++ b/sdks/python/ARCHITECTURE.md @@ -113,7 +113,7 @@ async with agent_control.AgentControlClient() as client: - `GET /api/v1/controls/{control_id}/rules` - List control rules **Functions**: -- `async def create_control(client, name)` - Create a new control +- `async def create_control(client, name, data)` - Create a new configured control - `async def add_rule_to_control(client, control_id, rule_id)` - Associate rule - `async def remove_rule_from_control(client, control_id, rule_id)` - Dissociate rule - `async def list_control_rules(client, control_id)` - List all rules in control @@ -124,7 +124,25 @@ import agent_control async with agent_control.AgentControlClient() as client: # Create control - result = await agent_control.controls.create_control(client, "pii-protection") + result = await agent_control.controls.create_control( + client, + "pii-protection", + { + "description": "PII protection", + "enabled": True, + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["post"]}, + "condition": { + "selector": {"path": "output"}, + "evaluator": { + "name": "regex", + "config": {"pattern": "\\\\d{3}-\\\\d{2}-\\\\d{4}", "flags": []}, + }, + }, + "action": {"decision": "deny"}, + "tags": ["security"], + }, + ) control_id = result["control_id"] # Add rule to control diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index b3e41d90..33658fb4 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -935,26 +935,23 @@ async def main(): async def create_control( name: str, - data: dict[str, Any] | ControlDefinition | None = None, + data: dict[str, Any] | ControlDefinition, server_url: str | None = None, api_key: str | None = None, ) -> dict[str, Any]: """ - Create a new control, optionally with configuration. - - If `data` is provided, the control is created and configured in one call. - Otherwise, use `agent_control.controls.set_control_data()` to configure it later. + Create a new control with configuration. Args: name: Unique name for the control - data: Optional control definition with a condition tree, action, scope, etc. + data: Control definition with a condition tree, action, scope, etc. server_url: Optional server URL (defaults to AGENT_CONTROL_URL env var) api_key: Optional API key for authentication (defaults to AGENT_CONTROL_API_KEY env var) Returns: Dictionary containing: - control_id: ID of the created control - - configured: True if data was set, False if only name was created + - configured: Always True because create requires data Raises: httpx.HTTPError: If request fails @@ -984,9 +981,6 @@ async def main(): ) print(f"Created control {result['control_id']}") - # Or create without config (configure later) - result = await agent_control.create_control(name="my-control") - asyncio.run(main()) """ _final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' diff --git a/sdks/python/src/agent_control/control_decorators.py b/sdks/python/src/agent_control/control_decorators.py index 1da113b1..d1edb06e 100644 --- a/sdks/python/src/agent_control/control_decorators.py +++ b/sdks/python/src/agent_control/control_decorators.py @@ -793,8 +793,7 @@ async def handle_user_input(user_message: str) -> str: Server Setup (separate from agent code): 1. Create controls via API: - PUT /api/v1/controls {"name": "block-toxic-inputs"} - PUT /api/v1/controls/{id}/data {"data": {...}} + PUT /api/v1/controls {"name": "block-toxic-inputs", "data": {...}} 2. Create policy and add controls: PUT /api/v1/policies {"name": "safety-policy"} diff --git a/sdks/python/src/agent_control/controls.py b/sdks/python/src/agent_control/controls.py index 1305fca4..767f2c86 100644 --- a/sdks/python/src/agent_control/controls.py +++ b/sdks/python/src/agent_control/controls.py @@ -118,13 +118,10 @@ async def get_control( async def create_control( client: AgentControlClient, name: str, - data: dict[str, Any] | ControlDefinition | None = None, + data: dict[str, Any] | ControlDefinition, ) -> dict[str, Any]: """ - Create a new control with a unique name, optionally with configuration. - - If `data` is provided, the control is created and configured in one call. - Otherwise, use `set_control_data()` to configure it later. + Create a new control with a unique name and configuration. Control names are canonicalized by the API (leading/trailing whitespace is trimmed); callers may pass trimmed names for consistency. @@ -132,12 +129,12 @@ async def create_control( Args: client: AgentControlClient instance name: Unique name for the control - data: Optional control definition (condition tree, action, scope, etc.) + data: Control definition (condition tree, action, scope, etc.) Returns: Dictionary containing: - control_id: ID of the created control - - configured: True if data was set, False if only name was created + - configured: Always True because create requires data Raises: httpx.HTTPError: If request fails @@ -147,11 +144,6 @@ async def create_control( Example: async with AgentControlClient() as client: - # Create without configuration (configure later) - result = await create_control(client, "pii-protection") - control_id = result["control_id"] - - # Or create with configuration in one call result = await create_control( client, name="ssn-blocker", @@ -170,22 +162,20 @@ async def create_control( ) print(f"Created and configured control: {result['control_id']}") """ - # Step 1: Create the control with name + payload: dict[str, Any] = {"name": name} + if isinstance(data, ControlDefinition): + payload["data"] = data.model_dump(mode="json", exclude_none=True) + else: + payload["data"] = cast(dict[str, Any], data) + response = await client.http_client.put( "/api/v1/controls", - json={"name": name} + json=payload, ) response.raise_for_status() result = cast(dict[str, Any], response.json()) - # Step 2: If data provided, configure the control - if data is not None: - control_id = result["control_id"] - await set_control_data(client, control_id, data) - result["configured"] = True - else: - result["configured"] = False - + result["configured"] = True return result diff --git a/sdks/python/tests/conftest.py b/sdks/python/tests/conftest.py index eabd7ec9..2110594c 100644 --- a/sdks/python/tests/conftest.py +++ b/sdks/python/tests/conftest.py @@ -186,7 +186,22 @@ async def test_control( """ result = await agent_control.controls.create_control( client, - f"test-control-{unique_name}" + f"test-control-{unique_name}", + { + "description": "SDK integration test control", + "enabled": True, + "execution": "server", + "scope": {"step_types": ["llm"], "stages": ["pre"]}, + "condition": { + "selector": {"path": "input"}, + "evaluator": { + "name": "regex", + "config": {"pattern": "test", "flags": []}, + }, + }, + "action": {"decision": "deny"}, + "tags": ["sdk-test"], + }, ) yield result diff --git a/sdks/typescript/src/generated/funcs/controls-create.ts b/sdks/typescript/src/generated/funcs/controls-create.ts index 05cf703a..188be0c4 100644 --- a/sdks/typescript/src/generated/funcs/controls-create.ts +++ b/sdks/typescript/src/generated/funcs/controls-create.ts @@ -30,13 +30,13 @@ import { Result } from "../types/fp.js"; * Create a new control * * @remarks - * Create a new control with a unique name and empty data. + * Create a new control with a unique name. * * Controls define protection logic and can be added to policies. - * Use the PUT /{control_id}/data endpoint to set control configuration. + * Control data is required and is validated before anything is inserted. * * Args: - * request: Control creation request with unique name + * request: Control creation request with unique name and data * db: Database session (injected) * * Returns: diff --git a/sdks/typescript/src/generated/models/create-control-request.ts b/sdks/typescript/src/generated/models/create-control-request.ts index 925b2aee..83a198ba 100644 --- a/sdks/typescript/src/generated/models/create-control-request.ts +++ b/sdks/typescript/src/generated/models/create-control-request.ts @@ -3,8 +3,22 @@ */ import * as z from "zod/v4-mini"; +import { + ControlDefinitionInput, + ControlDefinitionInput$Outbound, + ControlDefinitionInput$outboundSchema, +} from "./control-definition-input.js"; export type CreateControlRequest = { + /** + * A control definition to evaluate agent interactions. + * + * @remarks + * + * This model contains only the logic and configuration. + * Identity fields (id, name) are managed by the database. + */ + data: ControlDefinitionInput; /** * Unique control name (letters, numbers, hyphens, underscores) */ @@ -13,6 +27,7 @@ export type CreateControlRequest = { /** @internal */ export type CreateControlRequest$Outbound = { + data: ControlDefinitionInput$Outbound; name: string; }; @@ -21,6 +36,7 @@ export const CreateControlRequest$outboundSchema: z.ZodMiniType< CreateControlRequest$Outbound, CreateControlRequest > = z.object({ + data: ControlDefinitionInput$outboundSchema, name: z.string(), }); diff --git a/sdks/typescript/src/generated/sdk/controls.ts b/sdks/typescript/src/generated/sdk/controls.ts index c32be744..67583335 100644 --- a/sdks/typescript/src/generated/sdk/controls.ts +++ b/sdks/typescript/src/generated/sdk/controls.ts @@ -56,13 +56,13 @@ export class Controls extends ClientSDK { * Create a new control * * @remarks - * Create a new control with a unique name and empty data. + * Create a new control with a unique name. * * Controls define protection logic and can be added to policies. - * Use the PUT /{control_id}/data endpoint to set control configuration. + * Control data is required and is validated before anything is inserted. * * Args: - * request: Control creation request with unique name + * request: Control creation request with unique name and data * db: Database session (injected) * * Returns: diff --git a/sdks/typescript/tests/client-api.test.ts b/sdks/typescript/tests/client-api.test.ts index c6c11864..fe5a98db 100644 --- a/sdks/typescript/tests/client-api.test.ts +++ b/sdks/typescript/tests/client-api.test.ts @@ -77,6 +77,25 @@ describe("AgentControlClient API wiring", () => { await client.controls.create({ name: "deny-pii", + data: { + action: { + decision: "deny", + }, + condition: { + evaluator: { + name: "regex", + config: { pattern: "pii" }, + }, + selector: { + path: "input", + }, + }, + execution: "server", + scope: { + stages: ["pre"], + stepTypes: ["llm"], + }, + }, }); expect(fetchMock).toHaveBeenCalledTimes(1); @@ -87,6 +106,26 @@ describe("AgentControlClient API wiring", () => { expect(request.headers.get("content-type")).toContain("application/json"); await expect(request.clone().json()).resolves.toEqual({ name: "deny-pii", + data: { + action: { + decision: "deny", + }, + condition: { + evaluator: { + name: "regex", + config: { pattern: "pii" }, + }, + selector: { + path: "input", + }, + }, + enabled: true, + execution: "server", + scope: { + stages: ["pre"], + step_types: ["llm"], + }, + }, }); }); diff --git a/server/src/agent_control_server/endpoints/controls.py b/server/src/agent_control_server/endpoints/controls.py index a8ab1354..4f83df96 100644 --- a/server/src/agent_control_server/endpoints/controls.py +++ b/server/src/agent_control_server/endpoints/controls.py @@ -81,6 +81,21 @@ def _iter_condition_leaves( yield from _iter_condition_leaves(node.not_, path=f"{path}.not") +def _serialize_control_definition(control_def: ControlDefinition) -> dict[str, object]: + """Serialize control data for storage while omitting null scope fields.""" + data_json = control_def.model_dump( + mode="json", + by_alias=True, + exclude_none=True, + exclude_unset=True, + ) + if "scope" in data_json and isinstance(data_json["scope"], dict): + data_json["scope"] = { + k: v for k, v in data_json["scope"].items() if v is not None + } + return data_json + + async def _validate_control_definition( control_def: ControlDefinition, db: AsyncSession ) -> None: @@ -252,13 +267,13 @@ async def create_control( request: CreateControlRequest, db: AsyncSession = Depends(get_async_db) ) -> CreateControlResponse: """ - Create a new control with a unique name and empty data. + Create a new control with a unique name. Controls define protection logic and can be added to policies. - Use the PUT /{control_id}/data endpoint to set control configuration. + Control data is required and is validated before anything is inserted. Args: - request: Control creation request with unique name + request: Control creation request with unique name and data db: Database session (injected) Returns: @@ -279,7 +294,10 @@ async def create_control( hint="Choose a different name or update the existing control.", ) - control = Control(name=request.name, data={}) + await _validate_control_definition(request.data, db) + control_data = _serialize_control_definition(request.data) + + control = Control(name=request.name, data=control_data) db.add(control) try: await db.commit() @@ -451,19 +469,7 @@ async def set_control_data( # Validate evaluator config using shared logic await _validate_control_definition(request.data, db) - data_json = request.data.model_dump( - mode="json", - by_alias=True, - exclude_none=True, - exclude_unset=True, - ) - # Ensure scope does not store null/None for step_names or other optional fields, - # so round-trip (save then load) preserves step selection in the UI. - if "scope" in data_json and isinstance(data_json["scope"], dict): - data_json["scope"] = { - k: v for k, v in data_json["scope"].items() if v is not None - } - control.data = data_json + control.data = _serialize_control_definition(request.data) try: await db.commit() except Exception: diff --git a/server/src/agent_control_server/main.py b/server/src/agent_control_server/main.py index 50429fa5..6fbeee2b 100644 --- a/server/src/agent_control_server/main.py +++ b/server/src/agent_control_server/main.py @@ -140,7 +140,7 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: ## Quick Start 1. Register your agent with `/api/v1/agents/initAgent` -2. Create controls with `/api/v1/controls` and configure them +2. Create configured controls with `/api/v1/controls` 3. Create a policy and add controls to it 4. Assign the policy to your agent 5. Query agent's active controls with `/api/v1/agents/{agent_name}/controls` diff --git a/server/tests/test_agents_additional.py b/server/tests/test_agents_additional.py index 8c69f79f..1044a89f 100644 --- a/server/tests/test_agents_additional.py +++ b/server/tests/test_agents_additional.py @@ -6,6 +6,9 @@ from fastapi.testclient import TestClient from sqlalchemy import text +from sqlalchemy.orm import Session + +from agent_control_server.models import Control from .utils import VALID_CONTROL_PAYLOAD, canonicalize_control_payload from .conftest import engine @@ -36,15 +39,24 @@ def _init_agent( def _create_control_with_data(client: TestClient, data: dict) -> int: - resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) - assert resp.status_code == 200 - control_id = resp.json()["control_id"] - set_resp = client.put( - f"/api/v1/controls/{control_id}/data", - json={"data": canonicalize_control_payload(data)}, + resp = client.put( + "/api/v1/controls", + json={ + "name": f"control-{uuid.uuid4()}", + "data": canonicalize_control_payload(data), + }, ) - assert set_resp.status_code == 200, set_resp.text - return control_id + assert resp.status_code == 200, resp.text + return resp.json()["control_id"] + + +def _insert_unconfigured_control() -> int: + control = Control(name=f"control-{uuid.uuid4()}", data={}) + with Session(engine) as session: + session.add(control) + session.commit() + session.refresh(control) + return int(control.id) def _create_policy(client: TestClient) -> int: @@ -661,9 +673,7 @@ def test_set_agent_policy_skips_controls_without_data(client: TestClient) -> Non # Given: an agent and a policy with a control that has no data configured 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 - control_id = control_resp.json()["control_id"] + control_id = _insert_unconfigured_control() assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") assert assoc.status_code == 200 diff --git a/server/tests/test_auth.py b/server/tests/test_auth.py index d31522a8..7f6d8d6f 100644 --- a/server/tests/test_auth.py +++ b/server/tests/test_auth.py @@ -8,6 +8,9 @@ from agent_control_server import __version__ as server_version from agent_control_server.config import auth_settings +from .utils import VALID_CONTROL_PAYLOAD + + class TestHealthEndpoint: """Health endpoint should always be accessible without authentication.""" @@ -163,7 +166,11 @@ class TestAdminWriteEndpointAuthorization: @pytest.mark.parametrize( ("method", "path", "json_body"), [ - ("PUT", "/api/v1/controls", {"name": "control-authz-blocked"}), + ( + "PUT", + "/api/v1/controls", + {"name": "control-authz-blocked", "data": _VALID_CONTROL_DATA}, + ), ("PUT", "/api/v1/controls/1/data", {"data": _VALID_CONTROL_DATA}), ("PATCH", "/api/v1/controls/1", {"enabled": False}), ("DELETE", "/api/v1/controls/1", None), @@ -248,7 +255,10 @@ def test_non_admin_key_can_init_agent_and_fetch_controls( def test_admin_key_allowed_on_representative_mutations(self, admin_client: TestClient) -> None: control_name = f"control-authz-{uuid.uuid4().hex[:8]}" - control_response = admin_client.put("/api/v1/controls", json={"name": control_name}) + control_response = admin_client.put( + "/api/v1/controls", + json={"name": control_name, "data": VALID_CONTROL_PAYLOAD}, + ) assert control_response.status_code == 200 control_id = control_response.json()["control_id"] diff --git a/server/tests/test_control_compatibility.py b/server/tests/test_control_compatibility.py index 6485f1c3..f6e39fd5 100644 --- a/server/tests/test_control_compatibility.py +++ b/server/tests/test_control_compatibility.py @@ -52,16 +52,13 @@ def test_set_agent_policy_accepts_legacy_stored_control_payload(client: TestClie agent_name = _init_agent(client) policy_id = _create_policy(client) - control_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) + control_resp = client.put( + "/api/v1/controls", + json={"name": f"control-{uuid.uuid4()}", "data": VALID_CONTROL_PAYLOAD}, + ) assert control_resp.status_code == 200 control_id = control_resp.json()["control_id"] - set_resp = client.put( - f"/api/v1/controls/{control_id}/data", - json={"data": VALID_CONTROL_PAYLOAD}, - ) - assert set_resp.status_code == 200 - assoc = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") assert assoc.status_code == 200 @@ -83,7 +80,10 @@ def test_get_control_data_returns_canonical_shape_for_legacy_stored_payload( client: TestClient, ) -> None: # Given: a control whose stored row has been reverted to the legacy flat shape - control_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) + control_resp = client.put( + "/api/v1/controls", + json={"name": f"control-{uuid.uuid4()}", "data": VALID_CONTROL_PAYLOAD}, + ) assert control_resp.status_code == 200 control_id = control_resp.json()["control_id"] @@ -112,16 +112,13 @@ def test_list_agent_controls_returns_canonical_shape_for_legacy_stored_payload( agent_name = _init_agent(client) policy_id = _create_policy(client) - control_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) + control_resp = client.put( + "/api/v1/controls", + json={"name": f"control-{uuid.uuid4()}", "data": VALID_CONTROL_PAYLOAD}, + ) assert control_resp.status_code == 200 control_id = control_resp.json()["control_id"] - set_resp = client.put( - f"/api/v1/controls/{control_id}/data", - json={"data": VALID_CONTROL_PAYLOAD}, - ) - assert set_resp.status_code == 200 - 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_name}/policy/{policy_id}") @@ -151,7 +148,10 @@ def test_get_control_data_rejects_partial_legacy_stored_payload( client: TestClient, ) -> None: # Given: a stored control row with only one half of the legacy flat shape - control_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) + control_resp = client.put( + "/api/v1/controls", + json={"name": f"control-{uuid.uuid4()}", "data": VALID_CONTROL_PAYLOAD}, + ) assert control_resp.status_code == 200 control_id = control_resp.json()["control_id"] diff --git a/server/tests/test_controls.py b/server/tests/test_controls.py index 51520e7a..bc20f5da 100644 --- a/server/tests/test_controls.py +++ b/server/tests/test_controls.py @@ -3,29 +3,101 @@ from typing import Any from fastapi.testclient import TestClient +from sqlalchemy.orm import Session +from agent_control_server.models import Control -def create_control(client: TestClient) -> int: +from .conftest import engine + + +def create_control(client: TestClient, data: dict[str, Any] | None = None) -> int: name = f"control-{uuid.uuid4()}" - resp = client.put("/api/v1/controls", json={"name": name}) + payload = data if data is not None else VALID_CONTROL_DATA + resp = client.put("/api/v1/controls", json={"name": name, "data": payload}) assert resp.status_code == 200 cid = resp.json()["control_id"] assert isinstance(cid, int) return cid +def create_unconfigured_control(name: str | None = None) -> int: + control = Control(name=name or f"control-{uuid.uuid4()}", data={}) + with Session(engine) as session: + session.add(control) + session.commit() + session.refresh(control) + return int(control.id) + + def test_create_control_returns_id(client: TestClient) -> None: # Given: no prior controls # When: creating a control via API - resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) + resp = client.put( + "/api/v1/controls", + json={"name": f"control-{uuid.uuid4()}", "data": VALID_CONTROL_DATA}, + ) # Then: a control_id is returned (integer) assert resp.status_code == 200 assert isinstance(resp.json()["control_id"], int) +def test_create_control_with_data_stores_configured_payload(client: TestClient) -> None: + # Given: a valid control payload included during create + name = f"control-{uuid.uuid4()}" + + # When: creating the control with data in one request + resp = client.put("/api/v1/controls", json={"name": name, "data": VALID_CONTROL_DATA}) + + # Then: the control is created successfully + assert resp.status_code == 200, resp.text + control_id = resp.json()["control_id"] + + # When: reading back its data + data_resp = client.get(f"/api/v1/controls/{control_id}/data") + + # Then: the configured payload was stored immediately + assert data_resp.status_code == 200 + data = data_resp.json()["data"] + assert data["description"] == VALID_CONTROL_DATA["description"] + assert data["execution"] == VALID_CONTROL_DATA["execution"] + assert data["condition"]["evaluator"] == VALID_CONTROL_DATA["condition"]["evaluator"] + + +def test_create_control_invalid_data_returns_422_without_persisting(client: TestClient) -> None: + # Given: a create request whose control data fails evaluator validation + name = f"control-{uuid.uuid4()}" + invalid_data = deepcopy(VALID_CONTROL_DATA) + invalid_data["condition"]["evaluator"] = { + "name": "list", + "config": { + "values": ["a", "b"], + "logic": "invalid_logic", + "match_on": "match", + }, + } + + # When: creating the control with invalid data + resp = client.put("/api/v1/controls", json={"name": name, "data": invalid_data}) + + # Then: the request is rejected + assert resp.status_code == 422 + + # And: no shell control was persisted + list_resp = client.get("/api/v1/controls", params={"name": name}) + assert list_resp.status_code == 200 + body = list_resp.json() + assert body["pagination"]["total"] == 0 + assert body["controls"] == [] + + +def test_create_control_without_data_returns_422(client: TestClient) -> None: + resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4()}"}) + assert resp.status_code == 422 + + def test_get_control_data_initially_unconfigured(client: TestClient) -> None: - # Given: a newly created control (no data set yet) - control_id = create_control(client) + # Given: a legacy control row with no data set + control_id = create_unconfigured_control() # When: fetching its data resp = client.get(f"/api/v1/controls/{control_id}/data") # Then: 422 because empty data is not a valid ControlDefinition (RFC 7807 format) @@ -51,8 +123,8 @@ def test_get_control_data_initially_unconfigured(client: TestClient) -> None: } def test_set_control_data_replaces_existing(client: TestClient) -> None: - # Given: a control with empty data - control_id = create_control(client) + # Given: a legacy control with empty data + control_id = create_unconfigured_control() # When: setting data payload = VALID_CONTROL_DATA resp_put = client.put(f"/api/v1/controls/{control_id}/data", json={"data": payload}) @@ -175,10 +247,10 @@ def test_set_control_data_requires_body_with_data_key(client: TestClient) -> Non def test_create_control_duplicate_name_409(client: TestClient) -> None: # Given: a specific control name name = f"dup-control-{uuid.uuid4()}" - r1 = client.put("/api/v1/controls", json={"name": name}) + r1 = client.put("/api/v1/controls", json={"name": name, "data": VALID_CONTROL_DATA}) assert r1.status_code == 200 # When: creating again with the same name - r2 = client.put("/api/v1/controls", json={"name": name}) + r2 = client.put("/api/v1/controls", json={"name": name, "data": VALID_CONTROL_DATA}) # Then: conflict assert r2.status_code == 409 @@ -189,17 +261,15 @@ def test_create_control_duplicate_name_409(client: TestClient) -> None: def test_get_control_returns_metadata(client: TestClient) -> None: - """Test GET /controls/{id} returns control id, name, and data.""" - # Given: a control with a specific name + """Test GET /controls/{id} returns id, name, and None data for legacy rows.""" + # Given: a legacy control with a specific name and no configured data name = f"test-control-{uuid.uuid4()}" - resp = client.put("/api/v1/controls", json={"name": name}) - assert resp.status_code == 200 - control_id = resp.json()["control_id"] + control_id = create_unconfigured_control(name) # When: fetching the control get_resp = client.get(f"/api/v1/controls/{control_id}") - # Then: returns id, name, and data (None for unconfigured) + # Then: returns id, name, and data (None for legacy unconfigured rows) assert get_resp.status_code == 200 body = get_resp.json() assert body["id"] == control_id diff --git a/server/tests/test_controls_additional.py b/server/tests/test_controls_additional.py index a5be9777..bafbab67 100644 --- a/server/tests/test_controls_additional.py +++ b/server/tests/test_controls_additional.py @@ -26,13 +26,28 @@ from .utils import VALID_CONTROL_PAYLOAD -def _create_control(client: TestClient, name: str | None = None) -> tuple[int, str]: +def _create_control( + client: TestClient, + name: str | None = None, + data: dict | None = None, +) -> tuple[int, str]: control_name = name or f"control-{uuid.uuid4()}" - resp = client.put("/api/v1/controls", json={"name": control_name}) + payload = deepcopy(data) if data is not None else deepcopy(VALID_CONTROL_PAYLOAD) + resp = client.put("/api/v1/controls", json={"name": control_name, "data": payload}) assert resp.status_code == 200 return resp.json()["control_id"], control_name +def _insert_unconfigured_control(name: str | None = None) -> tuple[int, str]: + control_name = name or f"control-{uuid.uuid4()}" + control = Control(name=control_name, data={}) + with Session(engine) as session: + session.add(control) + session.commit() + session.refresh(control) + return int(control.id), control_name + + def _set_control_data(client: TestClient, control_id: int, data: dict) -> None: resp = client.put(f"/api/v1/controls/{control_id}/data", json={"data": data}) assert resp.status_code == 200, resp.text @@ -60,7 +75,10 @@ async def mock_db_integrity_error() -> AsyncGenerator[AsyncSession, None]: app.dependency_overrides[get_async_db] = mock_db_integrity_error try: - resp = client.put("/api/v1/controls", json={"name": "duplicate-control"}) + resp = client.put( + "/api/v1/controls", + json={"name": "duplicate-control", "data": VALID_CONTROL_PAYLOAD}, + ) finally: app.dependency_overrides.clear() @@ -199,7 +217,7 @@ def test_list_controls_filters_and_pagination(client: TestClient) -> None: def test_patch_control_enabled_requires_data(client: TestClient) -> None: # Given: a control without configured data - control_id, _ = _create_control(client) + control_id, _ = _insert_unconfigured_control() # When: toggling enabled without data resp = client.patch(f"/api/v1/controls/{control_id}", json={"enabled": False}) @@ -242,7 +260,10 @@ def test_patch_control_rename_with_spaces_rejected(client: TestClient) -> None: def test_create_control_trimmed_name_stored(client: TestClient) -> None: """Control names are canonicalized at the API boundary: leading/trailing whitespace is trimmed.""" - resp = client.put("/api/v1/controls", json={"name": " trimmed-control "}) + resp = client.put( + "/api/v1/controls", + json={"name": " trimmed-control ", "data": VALID_CONTROL_PAYLOAD}, + ) assert resp.status_code == 200 control_id = resp.json()["control_id"] get_resp = client.get(f"/api/v1/controls/{control_id}") diff --git a/server/tests/test_controls_validation.py b/server/tests/test_controls_validation.py index 25a295f8..2761bdd3 100644 --- a/server/tests/test_controls_validation.py +++ b/server/tests/test_controls_validation.py @@ -10,7 +10,7 @@ def create_control(client: TestClient) -> int: name = f"control-{uuid.uuid4()}" - resp = client.put("/api/v1/controls", json={"name": name}) + resp = client.put("/api/v1/controls", json={"name": name, "data": VALID_CONTROL_PAYLOAD}) assert resp.status_code == 200 return resp.json()["control_id"] diff --git a/server/tests/test_error_handling.py b/server/tests/test_error_handling.py index f9689851..12e5b519 100644 --- a/server/tests/test_error_handling.py +++ b/server/tests/test_error_handling.py @@ -11,6 +11,8 @@ from agent_control_server.db import get_async_db +from .utils import VALID_CONTROL_PAYLOAD + async def mock_db_with_commit_failure() -> AsyncGenerator[AsyncSession, None]: """Mock database session that fails on commit.""" @@ -287,7 +289,10 @@ def test_create_control_rollback_on_failure( # When: commit fails during control creation app.dependency_overrides[get_async_db] = mock_db_with_commit_failure try: - resp = client.put("/api/v1/controls", json={"name": control_name}) + resp = client.put( + "/api/v1/controls", + json={"name": control_name, "data": VALID_CONTROL_PAYLOAD}, + ) # Then: rollback is called and 500 error is returned assert resp.status_code == 500 @@ -302,7 +307,10 @@ def test_delete_control_rollback_on_failure( """Test that delete_control rolls back when commit fails.""" # Given: an existing control control_name = f"test-control-{uuid.uuid4()}" - create_resp = client.put("/api/v1/controls", json={"name": control_name}) + create_resp = client.put( + "/api/v1/controls", + json={"name": control_name, "data": VALID_CONTROL_PAYLOAD}, + ) assert create_resp.status_code == 200 control_id = create_resp.json()["control_id"] @@ -441,7 +449,10 @@ def test_set_control_data_rollback_on_failure( """Test that set_control_data rolls back transaction when commit fails.""" # Given: an existing control control_name = f"test-control-{uuid.uuid4()}" - r1 = client.put("/api/v1/controls", json={"name": control_name}) + r1 = client.put( + "/api/v1/controls", + json={"name": control_name, "data": VALID_CONTROL_PAYLOAD}, + ) assert r1.status_code == 200 control_id = r1.json()["control_id"] @@ -497,7 +508,10 @@ def test_patch_control_rollback_on_failure( """Test that patch_control rolls back when commit fails.""" # Given: an existing control control_name = f"test-control-{uuid.uuid4()}" - create_resp = client.put("/api/v1/controls", json={"name": control_name}) + create_resp = client.put( + "/api/v1/controls", + json={"name": control_name, "data": VALID_CONTROL_PAYLOAD}, + ) assert create_resp.status_code == 200 control_id = create_resp.json()["control_id"] diff --git a/server/tests/test_evaluation_e2e.py b/server/tests/test_evaluation_e2e.py index 7ebb03b9..3ba5cde6 100644 --- a/server/tests/test_evaluation_e2e.py +++ b/server/tests/test_evaluation_e2e.py @@ -236,12 +236,14 @@ def test_evaluation_deny_precedence(client: TestClient): "evaluator": {"name": "regex", "config": {"pattern": "keyword"}}, "action": {"decision": "deny"} } - resp = client.put("/api/v1/controls", json={"name": f"deny-control-{uuid.uuid4()}"}) - deny_control_id = resp.json()["control_id"] - client.put( - f"/api/v1/controls/{deny_control_id}/data", - json={"data": canonicalize_control_payload(control_deny)}, + resp = client.put( + "/api/v1/controls", + json={ + "name": f"deny-control-{uuid.uuid4()}", + "data": canonicalize_control_payload(control_deny), + }, ) + deny_control_id = resp.json()["control_id"] # Add Control to Agent's Policy client.post(f"/api/v1/policies/{policy_id}/controls/{deny_control_id}") diff --git a/server/tests/test_evaluation_error_handling.py b/server/tests/test_evaluation_error_handling.py index 5a3db01b..942dca66 100644 --- a/server/tests/test_evaluation_error_handling.py +++ b/server/tests/test_evaluation_error_handling.py @@ -17,11 +17,11 @@ def test_evaluation_with_agent_scoped_evaluator_missing(client: TestClient): - """Test that referencing missing agent evaluator fails at policy assignment. + """Test that referencing a missing agent evaluator fails during control creation. Given: A control referencing agent:evaluator that doesn't exist - When: Attempting to assign policy - Then: Returns 400 with clear error message + When: Creating the control + Then: Returns 422 EVALUATOR_NOT_FOUND """ # Given: an agent without evaluators agent_name = f"testagent-{uuid.uuid4().hex[:12]}" @@ -47,17 +47,15 @@ def test_evaluation_with_agent_scoped_evaluator_missing(client: TestClient): "action": {"decision": "deny"} } - # When: creating the control shell - control_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4().hex[:8]}"}) - assert control_resp.status_code == 200 - control_id = control_resp.json()["control_id"] - - # When: setting control data with a missing agent-scoped evaluator - set_resp = client.put(f"/api/v1/controls/{control_id}/data", json={"data": control_data}) + # When: creating the control with a missing agent-scoped evaluator + set_resp = client.put( + "/api/v1/controls", + json={"name": f"control-{uuid.uuid4().hex[:8]}", "data": control_data}, + ) - # Then: a validation or not-found error is returned - # This will fail because the agent doesn't exist yet - assert set_resp.status_code in [404, 422] + # Then: the missing evaluator is surfaced deterministically + assert set_resp.status_code == 422 + assert set_resp.json()["error_code"] == "EVALUATOR_NOT_FOUND" def test_evaluation_control_with_invalid_config_caught_early(client: TestClient): @@ -67,12 +65,7 @@ def test_evaluation_control_with_invalid_config_caught_early(client: TestClient) When: Setting control data Then: Returns 422 with validation error """ - # Given: a control shell to configure - control_resp = client.put("/api/v1/controls", json={"name": f"control-{uuid.uuid4().hex[:8]}"}) - assert control_resp.status_code == 200 - control_id = control_resp.json()["control_id"] - - # When: setting control data with invalid regex config (missing required 'pattern') + # When: creating a control with invalid regex config (missing required 'pattern') control_data = { "description": "Test control", "enabled": True, @@ -86,7 +79,10 @@ def test_evaluation_control_with_invalid_config_caught_early(client: TestClient) "action": {"decision": "deny"} } - set_resp = client.put(f"/api/v1/controls/{control_id}/data", json={"data": control_data}) + set_resp = client.put( + "/api/v1/controls", + json={"name": f"control-{uuid.uuid4().hex[:8]}", "data": control_data}, + ) # Then: a validation error is returned assert set_resp.status_code == 422 diff --git a/server/tests/test_init_agent.py b/server/tests/test_init_agent.py index 7140e736..2dfe9eaa 100644 --- a/server/tests/test_init_agent.py +++ b/server/tests/test_init_agent.py @@ -437,13 +437,10 @@ def test_list_agent_controls_with_policy(client: TestClient) -> None: policy_id = pol.json()["policy_id"] ctl_name = f"control-{uuid.uuid4()}" - ctl = client.put("/api/v1/controls", json={"name": ctl_name}) - control_id = ctl.json()["control_id"] - - # Set control data from .utils import VALID_CONTROL_PAYLOAD data_payload = VALID_CONTROL_PAYLOAD - client.put(f"/api/v1/controls/{control_id}/data", json={"data": data_payload}) + ctl = client.put("/api/v1/controls", json={"name": ctl_name, "data": data_payload}) + control_id = ctl.json()["control_id"] # Associate control -> policy; assign policy to agent client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") diff --git a/server/tests/test_init_agent_conflict_mode.py b/server/tests/test_init_agent_conflict_mode.py index 9c29d0d1..8a9e2ce1 100644 --- a/server/tests/test_init_agent_conflict_mode.py +++ b/server/tests/test_init_agent_conflict_mode.py @@ -41,21 +41,17 @@ def _create_policy_with_agent_evaluator_control( agent_name: str, evaluator_name: str, ) -> tuple[int, int, str]: - control_name = f"control-{uuid.uuid4().hex[:8]}" - create_control_resp = client.put("/api/v1/controls", json={"name": control_name}) - assert create_control_resp.status_code == 200 - control_id = create_control_resp.json()["control_id"] - control_data = deepcopy(VALID_CONTROL_PAYLOAD) + control_name = f"control-{uuid.uuid4().hex[:8]}" control_data["condition"]["evaluator"] = { "name": f"{agent_name}:{evaluator_name}", "config": {}, } - set_data_resp = client.put( - f"/api/v1/controls/{control_id}/data", - json={"data": control_data}, + create_control_resp = client.put( + "/api/v1/controls", json={"name": control_name, "data": control_data} ) - assert set_data_resp.status_code == 200 + assert create_control_resp.status_code == 200 + control_id = create_control_resp.json()["control_id"] policy_name = f"policy-{uuid.uuid4().hex[:8]}" create_policy_resp = client.put("/api/v1/policies", json={"name": policy_name}) diff --git a/server/tests/test_new_features.py b/server/tests/test_new_features.py index f79682a2..435273a7 100644 --- a/server/tests/test_new_features.py +++ b/server/tests/test_new_features.py @@ -264,17 +264,13 @@ def _create_policy_with_control( policy_id = pol_resp.json()["policy_id"] # Create control - ctl_resp = client.put("/api/v1/controls", json={"name": control_name}) + ctl_resp = client.put( + "/api/v1/controls", + json={"name": control_name, "data": canonicalize_control_payload(control_data)}, + ) assert ctl_resp.status_code == 200 control_id = ctl_resp.json()["control_id"] - # Set control data - data_resp = client.put( - f"/api/v1/controls/{control_id}/data", - json={"data": canonicalize_control_payload(control_data)}, - ) - assert data_resp.status_code == 200 - # Add control to policy client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") @@ -342,20 +338,18 @@ 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 an agent without evaluator, when creating a control that uses it, then fails.""" # Given: 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]}"}) - control_id = ctl_resp.json()["control_id"] - # When: data_resp = client.put( - f"/api/v1/controls/{control_id}/data", + "/api/v1/controls", json={ + "name": f"control-{uuid.uuid4().hex[:8]}", "data": { "execution": "server", "scope": {"step_types": ["llm"], "stages": ["pre"]}, diff --git a/server/tests/test_policies.py b/server/tests/test_policies.py index f40e20dc..e961a989 100644 --- a/server/tests/test_policies.py +++ b/server/tests/test_policies.py @@ -9,6 +9,8 @@ from agent_control_server.db import get_async_db from agent_control_server.models import Control, Policy +from .utils import VALID_CONTROL_PAYLOAD + def _create_policy(client: TestClient) -> int: name = f"pol-{uuid.uuid4()}" @@ -19,7 +21,7 @@ def _create_policy(client: TestClient) -> int: def _create_control(client: TestClient) -> int: name = f"ctrl-{uuid.uuid4()}" - r = client.put("/api/v1/controls", json={"name": name}) + r = client.put("/api/v1/controls", json={"name": name, "data": VALID_CONTROL_PAYLOAD}) assert r.status_code == 200 return r.json()["control_id"] diff --git a/server/tests/test_policy_integration.py b/server/tests/test_policy_integration.py index bd3ba15d..f3f28f34 100644 --- a/server/tests/test_policy_integration.py +++ b/server/tests/test_policy_integration.py @@ -1,6 +1,8 @@ """Integration tests for the full policy → control chain.""" +import json import uuid +from copy import deepcopy from fastapi.testclient import TestClient @@ -38,18 +40,13 @@ def _create_policy(client: TestClient, name: str | None = None) -> int: def _create_control(client: TestClient, name: str | None = None, data: dict | None = None) -> int: """Helper: Create a control and return control_id.""" control_name = name or f"control-{uuid.uuid4()}" - resp = client.put("/api/v1/controls", json={"name": control_name}) + payload = deepcopy(VALID_CONTROL_PAYLOAD) + marker = json.dumps(data, sort_keys=True) if data is not None else control_name + payload["description"] = f"Name: {control_name}, Marker: {marker}" + payload["condition"]["evaluator"]["config"]["pattern"] = marker + resp = client.put("/api/v1/controls", json={"name": control_name, "data": payload}) assert resp.status_code == 200 - control_id = resp.json()["control_id"] - - # Always set valid data, using name/data in description for traceability - payload = VALID_CONTROL_PAYLOAD.copy() - payload["description"] = f"Name: {control_name}, Data: {data}" - - resp = client.put(f"/api/v1/controls/{control_id}/data", json={"data": payload}) - assert resp.status_code == 200 - - return control_id + return resp.json()["control_id"] def test_agent_gets_controls_from_policy(client: TestClient) -> None: diff --git a/server/tests/utils.py b/server/tests/utils.py index 64c455c5..92cf6c4d 100644 --- a/server/tests/utils.py +++ b/server/tests/utils.py @@ -49,25 +49,21 @@ def create_and_assign_policy( # 1. Create Control control_name = f"control-{uuid.uuid4()}" - resp = client.put("/api/v1/controls", json={"name": control_name}) + resp = client.put("/api/v1/controls", json={"name": control_name, "data": control_config}) assert resp.status_code == 200 control_id = resp.json()["control_id"] - # 2. Configure Control - resp = client.put(f"/api/v1/controls/{control_id}/data", json={"data": control_config}) - assert resp.status_code == 200 - - # 3. Create Policy + # 2. Create Policy policy_name = f"policy-{uuid.uuid4()}" resp = client.put("/api/v1/policies", json={"name": policy_name}) assert resp.status_code == 200 policy_id = resp.json()["policy_id"] - # 4. Add Control to Policy (direct relationship) + # 3. Add Control to Policy (direct relationship) resp = client.post(f"/api/v1/policies/{policy_id}/controls/{control_id}") assert resp.status_code == 200 - # 5. Register Agent + # 4. Register Agent normalized_agent_name = agent_name.lower() if len(normalized_agent_name) < 10: normalized_agent_name = f"{normalized_agent_name}-agent".replace("--", "-") @@ -79,7 +75,7 @@ def create_and_assign_policy( }) assert resp.status_code == 200 - # 6. Assign Policy to Agent + # 5. Assign Policy to Agent resp = client.post(f"/api/v1/agents/{normalized_agent_name}/policy/{policy_id}") assert resp.status_code == 200 diff --git a/ui/src/core/api/generated/api-types.ts b/ui/src/core/api/generated/api-types.ts index 55b9538c..01625731 100644 --- a/ui/src/core/api/generated/api-types.ts +++ b/ui/src/core/api/generated/api-types.ts @@ -441,13 +441,13 @@ export interface paths { get: operations['list_controls_api_v1_controls_get']; /** * Create a new control - * @description Create a new control with a unique name and empty data. + * @description Create a new control with a unique name. * * Controls define protection logic and can be added to policies. - * Use the PUT /{control_id}/data endpoint to set control configuration. + * Control data is required and is validated before anything is inserted. * * Args: - * request: Control creation request with unique name + * request: Control creation request with unique name and data * db: Database session (injected) * * Returns: @@ -1961,6 +1961,8 @@ export interface components { }; /** CreateControlRequest */ CreateControlRequest: { + /** @description Control definition to validate and store during creation */ + data: components['schemas']['ControlDefinition-Input']; /** * Name * @description Unique control name (letters, numbers, hyphens, underscores) @@ -2928,7 +2930,7 @@ export interface components { enabled?: boolean | null; /** * Name - * @description New name for the control + * @description New control name (letters, numbers, hyphens, underscores) */ name?: string | null; }; @@ -3335,10 +3337,6 @@ export interface components { }; /** ValidationError */ ValidationError: { - /** Context */ - ctx?: Record; - /** Input */ - input?: unknown; /** Location */ loc: (string | number)[]; /** Message */ 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 40c9ce2c..932699aa 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,9 +13,8 @@ type AddControlToAgentParams = { /** * Mutation hook to add a control to an agent * Flow: - * 1. Create the control - * 2. Set control data (definition) - * 3. Associate the control directly with the agent + * 1. Create the control with its definition + * 2. Associate the control directly with the agent */ export function useAddControlToAgent() { const queryClient = useQueryClient(); @@ -29,12 +28,12 @@ export function useAddControlToAgent() { let createdControlId: number | null = null; try { - // Step 1: Create the control + // Step 1: Create and validate the control atomically. const { data: createControlResult, error: createControlError, response: createControlResponse, - } = await api.controls.create({ name: controlName }); + } = await api.controls.create({ name: controlName, data: definition }); if (createControlError || !createControlResult) { throw parseApiError( @@ -46,21 +45,7 @@ export function useAddControlToAgent() { createdControlId = createControlResult.control_id; - // Step 2: Set control data (definition) - const { error: setDataError, response: setDataResponse } = - await api.controls.setData(createdControlId, { - data: definition, - }); - - if (setDataError) { - throw parseApiError( - setDataError, - 'Failed to set control data', - setDataResponse?.status - ); - } - - // Step 3: Associate control directly with the agent. + // Step 2: Associate control directly with the agent. const { error: associateError, response: associateResponse } = await api.agents.addControl(agentId, createdControlId); diff --git a/ui/src/core/hooks/query-hooks/use-create-control.ts b/ui/src/core/hooks/query-hooks/use-create-control.ts index 3a89d46f..521795ba 100644 --- a/ui/src/core/hooks/query-hooks/use-create-control.ts +++ b/ui/src/core/hooks/query-hooks/use-create-control.ts @@ -9,31 +9,20 @@ type CreateControlParams = { }; /** - * Mutation hook to create a new control with its definition - * 1. Creates the control with a name - * 2. Sets the control data (definition) + * Mutation hook to create a new control with its definition in one request. */ export function useCreateControl() { const queryClient = useQueryClient(); return useMutation({ mutationFn: async ({ name, definition }: CreateControlParams) => { - // Step 1: Create control with name + // Create and validate the control atomically on the server. const { data: createResult, error: createError } = - await api.controls.create({ name }); + await api.controls.create({ name, data: definition }); if (createError) throw createError; if (!createResult) throw new Error('Failed to create control'); - - const controlId = createResult.control_id; - - // Step 2: Set control data (definition) - const { data: setDataResult, error: setDataError } = - await api.controls.setData(controlId, { data: definition }); - - if (setDataError) throw setDataError; - - return { controlId, ...setDataResult }; + return { controlId: createResult.control_id, success: true }; }, onSuccess: () => { // Invalidate relevant queries to refetch data