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
12 changes: 12 additions & 0 deletions models/src/agent_control_models/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,11 @@
from .server import (
AgentRef,
AgentSummary,
CloneAndBindControlRequest,
CloneAndBindControlResponse,
CloneAndBindTargetBinding,
ConflictMode,
ControlAttachments,
ControlSummary,
ControlVersionSummary,
CreateControlBindingRequest,
Expand All @@ -107,9 +111,11 @@
PatchControlBindingResponse,
PatchControlRequest,
PatchControlResponse,
PolicyRef,
RenderControlTemplateRequest,
RenderControlTemplateResponse,
StepKey,
TargetAttachmentRef,
UpsertControlBindingRequest,
UpsertControlBindingResponse,
ValidateControlDataRequest,
Expand Down Expand Up @@ -176,7 +182,11 @@
# Server models
"AgentRef",
"AgentSummary",
"CloneAndBindControlRequest",
"CloneAndBindControlResponse",
"CloneAndBindTargetBinding",
"ConflictMode",
"ControlAttachments",
"ControlVersionSummary",
"ControlSummary",
"CreateControlBindingRequest",
Expand All @@ -200,9 +210,11 @@
"PatchControlBindingResponse",
"PatchControlRequest",
"PatchControlResponse",
"PolicyRef",
"RenderControlTemplateRequest",
"RenderControlTemplateResponse",
"StepKey",
"TargetAttachmentRef",
"UpsertControlBindingRequest",
"UpsertControlBindingResponse",
"ValidateControlDataRequest",
Expand Down
2 changes: 2 additions & 0 deletions models/src/agent_control_models/errors.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ class ErrorCode(StrEnum):
AUTH_INVALID_KEY = "AUTH_INVALID_KEY"
AUTH_INSUFFICIENT_PRIVILEGES = "AUTH_INSUFFICIENT_PRIVILEGES"
AUTH_MISCONFIGURED = "AUTH_MISCONFIGURED"
AUTH_UPSTREAM_REJECTED = "AUTH_UPSTREAM_REJECTED"

# Resource Not Found (2xx pattern)
RESOURCE_NOT_FOUND = "RESOURCE_NOT_FOUND" # Generic fallback
Expand Down Expand Up @@ -363,6 +364,7 @@ def make_error_type(error_code: ErrorCode) -> str:
ErrorCode.AUTH_INVALID_KEY: "Invalid API Key",
ErrorCode.AUTH_INSUFFICIENT_PRIVILEGES: "Insufficient Privileges",
ErrorCode.AUTH_MISCONFIGURED: "Authentication Misconfigured",
ErrorCode.AUTH_UPSTREAM_REJECTED: "Authorization Upstream Rejected Request",
# Not found errors
ErrorCode.RESOURCE_NOT_FOUND: "Resource Not Found",
ErrorCode.AGENT_NOT_FOUND: "Agent Not Found",
Expand Down
119 changes: 117 additions & 2 deletions models/src/agent_control_models/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
from .agent import Agent, StepSchema
from .base import BaseModel
from .controls import (
ControlAction,
ControlDefinition,
TemplateControlInput,
TemplateDefinition,
Expand Down Expand Up @@ -347,6 +348,9 @@ class GetControlResponse(BaseModel):

id: int = Field(..., description="Control ID")
name: str = Field(..., description="Control name")
cloned_from_control_id: int | None = Field(
None, description="Source control ID when this control is a clone."
)
data: ControlDefinition | UnrenderedTemplateControl = Field(
description=(
"Control configuration data. A ControlDefinition for raw/rendered "
Expand Down Expand Up @@ -514,14 +518,60 @@ class AgentRef(BaseModel):
agent_name: str = Field(..., description="Agent name")


class PolicyRef(BaseModel):
"""Reference to a policy attached to a control."""

policy_id: int = Field(..., description="Policy ID")


class TargetAttachmentRef(BaseModel):
"""Reference to a target binding attached to a control."""

binding_id: int = Field(..., description="Control binding ID")
target_type: str = Field(..., description="Opaque target kind")
target_id: str = Field(..., description="Opaque target identifier")
enabled: bool = Field(..., description="Whether this target binding is enabled")


class ControlAttachments(BaseModel):
"""Attachments for a listed control."""

agents: list[AgentRef] = Field(
default_factory=list,
description="Direct agent associations for this control",
)
policies: list[PolicyRef] = Field(
default_factory=list,
description="Policy associations for this control",
)
targets: list[TargetAttachmentRef] = Field(
default_factory=list,
description="Target bindings for this control",
)
targets_total: int = Field(
default=0,
description="Total target bindings matching the attachment filters",
)
targets_truncated: bool = Field(
default=False,
description="Whether the target bindings list was capped",
)


class ControlSummary(BaseModel):
"""Summary of a control for list responses."""

id: int = Field(..., description="Control ID")
name: str = Field(..., description="Control name")
cloned_from_control_id: int | None = Field(
None, description="Source control ID when this control is a clone."
)
description: str | None = Field(None, description="Control description")
enabled: bool = Field(True, description="Whether control is enabled")
execution: str | None = Field(None, description="'server' or 'sdk'")
action: ControlAction | None = Field(
None, description="Action applied when the control matches."
)
step_types: list[str] | None = Field(None, description="Step types in scope")
stages: list[str] | None = Field(None, description="Evaluation stages in scope")
tags: list[str] = Field(default_factory=list, description="Control tags")
Expand All @@ -542,6 +592,13 @@ class ControlSummary(BaseModel):
used_by_agents_count: int = Field(
0, description="Number of unique agents using this control"
)
attachments: ControlAttachments | None = Field(
None,
description=(
"Expanded attachment details. Present when list controls is called "
"with include_attachments=true."
),
)


class ListControlsResponse(BaseModel):
Expand Down Expand Up @@ -580,7 +637,7 @@ class GetControlVersionResponse(BaseModel):
...,
description=(
"Raw persisted snapshot of the control state at this version, including "
"metadata such as name, deleted_at, and cloned_control_id."
"metadata such as name, deleted_at, and cloned_from_control_id."
),
)

Expand Down Expand Up @@ -635,6 +692,50 @@ class PatchControlResponse(BaseModel):
]


class CloneAndBindTargetBinding(BaseModel):
"""Target binding to create for a cloned control."""

model_config = ConfigDict(extra="forbid")

target_type: ControlBindingTargetField = Field(
...,
description="Opaque attachment kind (caller-defined; e.g. 'environment', 'session').",
)
target_id: ControlBindingTargetField = Field(
..., description="Opaque external identifier within the target_type."
)
enabled: bool = Field(
default=True,
description="Whether the created binding is active.",
)


class CloneAndBindControlRequest(BaseModel):
"""Request to clone a control and attach the clone to one target."""

model_config = ConfigDict(extra="forbid")

name: SlugName | None = Field(
None,
description=(
"Optional unique name for the cloned control. If omitted, the server "
"generates a name from the source control name."
),
)
target_binding: CloneAndBindTargetBinding = Field(
..., description="Target binding to create for the cloned control."
)


class CloneAndBindControlResponse(BaseModel):
"""Response from cloning and binding a control."""

id: int = Field(..., description="Identifier of the cloned control.")
name: str = Field(..., description="Name of the cloned control.")
cloned_from_control_id: int = Field(..., description="Source control ID.")
binding_id: int = Field(..., description="Identifier of the created binding.")


class CreateControlBindingRequest(BaseModel):
"""Request to attach a control to an opaque external target."""

Expand Down Expand Up @@ -741,6 +842,21 @@ class UpsertControlBindingResponse(BaseModel):
enabled: bool = Field(..., description="Current enabled value.")


class PatchControlBindingByKeyRequest(BaseModel):
"""Request to update an existing control binding by natural key."""

target_type: ControlBindingTargetField = Field(
..., description="Opaque attachment kind."
)
target_id: ControlBindingTargetField = Field(
..., description="Opaque external identifier within the target_type."
)
control_id: int = Field(
..., gt=0, description="ID of the bound control."
)
enabled: bool = Field(..., description="New enabled value for the binding.")


class DeleteControlBindingByKeyRequest(BaseModel):
"""Request to detach a control binding by natural key (idempotent)."""

Expand All @@ -759,4 +875,3 @@ class DeleteControlBindingByKeyResponse(BaseModel):
"binding existed."
),
)

59 changes: 58 additions & 1 deletion sdks/python/src/agent_control/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ async def handle_input(user_message: str) -> str:
set_trace_context_provider,
)

from . import agents, controls, evaluation, evaluators, policies
from . import agents, control_bindings, controls, evaluation, evaluators, policies
from ._control_registry import (
StepSchemaDict,
get_registered_steps,
Expand Down Expand Up @@ -1019,10 +1019,14 @@ async def list_controls(
name: str | None = None,
enabled: bool | None = None,
template_backed: bool | None = None,
cloned: bool | None = None,
step_type: str | None = None,
stage: Literal["pre", "post"] | None = None,
execution: Literal["server", "sdk"] | None = None,
tag: str | None = None,
include_attachments: bool = False,
attachment_target_type: str | None = None,
attachment_target_id: str | None = None,
) -> dict[str, Any]:
"""
List all controls from the server with optional filtering.
Expand All @@ -1035,10 +1039,14 @@ async def list_controls(
name: Optional filter by name (partial, case-insensitive)
enabled: Optional filter by enabled status
template_backed: Optional filter by whether the control is template-backed
cloned: Optional filter by whether the control was cloned from another control
step_type: Optional filter by step type (built-ins: 'tool', 'llm')
stage: Optional filter by stage ('pre' or 'post')
execution: Optional filter by execution ('server' or 'sdk')
tag: Optional filter by tag
include_attachments: Whether to include attachment details
attachment_target_type: Optional target binding type filter for attachments
attachment_target_id: Optional target binding ID filter for attachments

Returns:
Dictionary containing:
Expand Down Expand Up @@ -1079,10 +1087,14 @@ async def main():
name=name,
enabled=enabled,
template_backed=template_backed,
cloned=cloned,
step_type=step_type,
stage=stage,
execution=execution,
tag=tag,
include_attachments=include_attachments,
attachment_target_type=attachment_target_type,
attachment_target_id=attachment_target_id,
)


Expand Down Expand Up @@ -1147,6 +1159,49 @@ async def main():
return await controls.create_control(client, name, data=data)


async def clone_and_bind_control(
control_id: int,
*,
target_type: str,
target_id: str,
name: str | None = None,
enabled: bool = True,
server_url: str | None = None,
api_key: str | None = None,
api_key_header: str | None = None,
) -> dict[str, Any]:
"""
Clone an existing control and bind the clone to a target.

Args:
control_id: Source control ID to clone
target_type: Opaque attachment kind
target_id: Opaque external target identifier
name: Optional unique name for the cloned control
enabled: Whether the created binding is active
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 id, name, cloned_from_control_id, and binding_id.
"""
_final_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000'

async with _ad_hoc_client(
server_url=_final_server_url,
api_key=api_key,
api_key_header=api_key_header,
) as client:
return await controls.clone_and_bind_control(
client,
control_id,
target_type=target_type,
target_id=target_id,
name=name,
enabled=enabled,
)


async def validate_control_data(
data: dict[str, Any] | ControlDefinition | TemplateControlInput,
server_url: str | None = None,
Expand Down Expand Up @@ -1502,6 +1557,7 @@ async def main():
"add_agent_control",
"remove_agent_control",
# Control management
"clone_and_bind_control",
"create_control",
"list_controls",
"get_control",
Expand All @@ -1520,6 +1576,7 @@ async def main():
"agents",
"policies",
"controls",
"control_bindings",
"evaluation",
"evaluators",
# Policy-Control management
Expand Down
Loading
Loading