Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 2 additions & 4 deletions models/src/agent_control_models/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -311,12 +311,10 @@ class GetControlResponse(BaseModel):

id: int = Field(..., description="Control ID")
name: str = Field(..., description="Control name")
data: ControlDefinition | UnrenderedTemplateControl | None = Field(
None,
data: ControlDefinition | UnrenderedTemplateControl = Field(
description=(
"Control configuration data. A ControlDefinition for raw/rendered "
"controls, an UnrenderedTemplateControl for unrendered templates, "
"or None if not yet configured."
"controls or an UnrenderedTemplateControl for unrendered templates."
),
)

Expand Down
2 changes: 1 addition & 1 deletion sdks/python/src/agent_control/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -1055,7 +1055,7 @@ async def get_control(
Dictionary containing:
- id: Control ID
- name: Control name
- data: Control definition or None if not configured
- data: Control definition or unrendered template control data

Raises:
httpx.HTTPError: If request fails or control not found
Expand Down
5 changes: 2 additions & 3 deletions sdks/python/src/agent_control/controls.py
Original file line number Diff line number Diff line change
Expand Up @@ -107,7 +107,7 @@ async def get_control(
Dictionary containing:
- id: Control ID
- name: Control name
- data: Control definition (condition, action, scope, etc.) or None if not configured
- data: Control definition or unrendered template control data

Raises:
httpx.HTTPError: If request fails
Expand All @@ -117,8 +117,7 @@ async def get_control(
async with AgentControlClient() as client:
control = await get_control(client, control_id=5)
print(f"Control: {control['name']}")
if control['data']:
print(f"Execution: {control['data']['execution']}")
print(f"Enabled: {control['data']['enabled']}")
"""
response = await client.http_client.get(f"/api/v1/controls/{control_id}")
response.raise_for_status()
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ import { Result } from "../types/fp.js";
* Raises:
* HTTPException 404: Control not found
* HTTPException 409: New name conflicts with existing control
* HTTPException 422: Cannot update enabled status (control has no data configured)
* HTTPException 422: Cannot update metadata for corrupted control data
* HTTPException 500: Database error during update
*/
export function controlsUpdateMetadata(
Expand Down
18 changes: 7 additions & 11 deletions sdks/typescript/src/generated/models/get-control-response.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ import {
} from "./unrendered-template-control.js";

/**
* Control configuration data. A ControlDefinition for raw/rendered controls, an UnrenderedTemplateControl for unrendered templates, or None if not yet configured.
* Control configuration data. A ControlDefinition for raw/rendered controls or an UnrenderedTemplateControl for unrendered templates.
*/
export type GetControlResponseData =
| ControlDefinitionOutput
Expand All @@ -29,9 +29,9 @@ export type GetControlResponseData =
*/
export type GetControlResponse = {
/**
* Control configuration data. A ControlDefinition for raw/rendered controls, an UnrenderedTemplateControl for unrendered templates, or None if not yet configured.
* Control configuration data. A ControlDefinition for raw/rendered controls or an UnrenderedTemplateControl for unrendered templates.
*/
data?: ControlDefinitionOutput | UnrenderedTemplateControl | null | undefined;
data: ControlDefinitionOutput | UnrenderedTemplateControl;
/**
* Control ID
*/
Expand Down Expand Up @@ -66,14 +66,10 @@ export const GetControlResponse$inboundSchema: z.ZodMiniType<
GetControlResponse,
unknown
> = z.object({
data: z.optional(
z.nullable(
smartUnion([
ControlDefinitionOutput$inboundSchema,
UnrenderedTemplateControl$inboundSchema,
]),
),
),
data: smartUnion([
ControlDefinitionOutput$inboundSchema,
UnrenderedTemplateControl$inboundSchema,
]),
id: types.number(),
name: types.string(),
});
Expand Down
2 changes: 1 addition & 1 deletion sdks/typescript/src/generated/sdk/controls.ts
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ export class Controls extends ClientSDK {
* Raises:
* HTTPException 404: Control not found
* HTTPException 409: New name conflicts with existing control
* HTTPException 422: Cannot update enabled status (control has no data configured)
* HTTPException 422: Cannot update metadata for corrupted control data
* HTTPException 500: Database error during update
*/
async updateMetadata(
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,227 @@
"""add control versions and soft-delete unusable legacy controls

Revision ID: c1e9f9c4a1d2
Revises: 5f2b5f4e1a90
Create Date: 2026-04-15 12:00:00.000000

"""

from __future__ import annotations

import datetime as dt
import logging
from typing import Any

import sqlalchemy as sa
from alembic import op
from pydantic import ValidationError
from sqlalchemy import inspect
from sqlalchemy.dialects import postgresql

from agent_control_models import ControlDefinition, UnrenderedTemplateControl

# revision identifiers, used by Alembic.
revision = "c1e9f9c4a1d2"
down_revision = "5f2b5f4e1a90"
branch_labels = None
depends_on = None

_logger = logging.getLogger("alembic.runtime.migration")

_BACKFILL_NOTE = "Backfilled from existing control"


def _classify_control_payload(data: Any) -> tuple[bool, str | None]:
"""Return whether a legacy control payload is still usable."""
if data == {}:
return False, "empty payload"
if not isinstance(data, dict):
return False, "invalid control payload"

try:
UnrenderedTemplateControl.model_validate(data)
except ValidationError:
pass
else:
return True, None

try:
ControlDefinition.model_validate(data)
except ValidationError:
return False, "invalid control payload"

return True, None


def _snapshot_payload(
*,
name: str,
data: Any,
deleted_at: dt.datetime | None,
) -> dict[str, Any]:
"""Build the JSON snapshot persisted in control_versions."""
return {
"name": name,
"data": data,
"deleted_at": deleted_at.isoformat() if deleted_at is not None else None,
"cloned_control_id": None,
}


def upgrade() -> None:
op.add_column("controls", sa.Column("deleted_at", sa.DateTime(timezone=True), nullable=True))
op.drop_constraint("controls_name_key", "controls", type_="unique")
op.create_index(
"idx_controls_name_active",
"controls",
["name"],
unique=True,
postgresql_where=sa.text("deleted_at IS NULL"),
)

op.create_table(
"control_versions",
sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
sa.Column("control_id", sa.Integer(), nullable=False),
sa.Column("version_num", sa.Integer(), nullable=False),
sa.Column("event_type", sa.String(length=255), nullable=False),
sa.Column("snapshot", postgresql.JSONB(astext_type=sa.Text()), nullable=False),
sa.Column("note", sa.Text(), nullable=True),
sa.Column(
"created_at",
sa.DateTime(timezone=True),
server_default=sa.text("CURRENT_TIMESTAMP"),
nullable=False,
),
sa.ForeignKeyConstraint(["control_id"], ["controls.id"]),
sa.PrimaryKeyConstraint("id"),
sa.UniqueConstraint(
"control_id",
"version_num",
name="uq_control_versions_control_version",
),
)
op.create_index(
"idx_control_versions_control_created",
"control_versions",
["control_id", sa.literal_column("created_at DESC")],
unique=False,
)

bind = op.get_bind()
db_inspector = inspect(bind)

controls = sa.table(
"controls",
sa.column("id", sa.Integer()),
sa.column("name", sa.String()),
sa.column("data", postgresql.JSONB(astext_type=sa.Text())),
sa.column("deleted_at", sa.DateTime(timezone=True)),
)
control_versions = sa.table(
"control_versions",
sa.column("control_id", sa.Integer()),
sa.column("version_num", sa.Integer()),
sa.column("event_type", sa.String()),
sa.column("snapshot", postgresql.JSONB(astext_type=sa.Text())),
sa.column("note", sa.Text()),
)
policy_controls = sa.table(
"policy_controls",
sa.column("policy_id", sa.Integer()),
sa.column("control_id", sa.Integer()),
)
agent_controls = sa.table(
"agent_controls",
sa.column("agent_name", sa.String()),
sa.column("control_id", sa.Integer()),
)

store_publications = None
if db_inspector.has_table("control_stores_controls"):
store_publications = sa.table(
"control_stores_controls",
sa.column("store_id", sa.Integer()),
sa.column("control_id", sa.Integer()),
)

rows = bind.execute(
sa.select(
controls.c.id,
controls.c.name,
controls.c.data,
).order_by(controls.c.id)
).mappings()

auto_deleted_controls: list[str] = []
for row in rows:
control_id = int(row["id"])
control_name = str(row["name"])
control_data = row["data"]
usable, reason = _classify_control_payload(control_data)

bind.execute(
sa.insert(control_versions).values(
control_id=control_id,
version_num=1,
event_type="migration_backfill",
snapshot=_snapshot_payload(
name=control_name,
data=control_data,
deleted_at=None,
),
note=_BACKFILL_NOTE,
)
)

if usable:
continue

if store_publications is not None:
bind.execute(
sa.delete(store_publications).where(
store_publications.c.control_id == control_id
)
)
bind.execute(
sa.delete(policy_controls).where(policy_controls.c.control_id == control_id)
)
bind.execute(
sa.delete(agent_controls).where(agent_controls.c.control_id == control_id)
)

deleted_at = dt.datetime.now(dt.UTC)
bind.execute(
sa.update(controls)
.where(controls.c.id == control_id)
.values(deleted_at=deleted_at)
)
bind.execute(
sa.insert(control_versions).values(
control_id=control_id,
version_num=2,
event_type="migration_autodelete",
snapshot=_snapshot_payload(
name=control_name,
data=control_data,
deleted_at=deleted_at,
),
note=f"Auto-soft-deleted during migration: {reason}",
)
)
auto_deleted_controls.append(f"{control_id}:{control_name}")

if auto_deleted_controls:
_logger.warning(
"Auto-soft-deleted %d unusable controls during migration: %s",
len(auto_deleted_controls),
", ".join(auto_deleted_controls),
)


def downgrade() -> None:
op.drop_index("idx_control_versions_control_created", table_name="control_versions")
op.drop_table("control_versions")
op.drop_index("idx_controls_name_active", table_name="controls")
op.create_unique_constraint("controls_name_key", "controls", ["name"])
op.drop_column("controls", "deleted_at")
12 changes: 7 additions & 5 deletions server/src/agent_control_server/endpoints/agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,9 +114,6 @@ def _validate_controls_for_agent(agent: Agent, controls: list[Control]) -> list[
agent_evaluators = {e.name: e for e in (agent_data.evaluators or [])}

for control in controls:
if not control.data:
continue

# Skip unrendered template controls — they have no evaluators to validate.
if (
isinstance(control.data, dict)
Expand Down Expand Up @@ -399,6 +396,7 @@ async def list_agents(
)
.join(Control, all_associations.c.control_id == Control.id)
.where(
Control.deleted_at.is_(None),
or_(
Control.data["enabled"].astext == "true",
~Control.data.has_key("enabled"),
Expand Down Expand Up @@ -1248,7 +1246,9 @@ async def add_agent_control(
"""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_result = await db.execute(
select(Control).where(Control.id == control_id, Control.deleted_at.is_(None))
)
control: Control | None = control_result.scalars().first()
if control is None:
raise NotFoundError(
Expand Down Expand Up @@ -1314,7 +1314,9 @@ async def remove_agent_control(
"""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))
control_result = await db.execute(
select(Control.id).where(Control.id == control_id, Control.deleted_at.is_(None))
)
if control_result.first() is None:
raise NotFoundError(
error_code=ErrorCode.CONTROL_NOT_FOUND,
Expand Down
Loading
Loading