Skip to content

Commit d1ee234

Browse files
committed
refactor: harden template control runtime helpers
1 parent 8bc76a0 commit d1ee234

8 files changed

Lines changed: 218 additions & 151 deletions

File tree

‎engine/src/agent_control_engine/core.py‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,9 @@
1515
from agent_control_evaluators import get_evaluator_instance
1616
from agent_control_models import (
1717
ConditionNode,
18-
ControlDefinition,
19-
ControlDefinitionRuntime,
18+
ControlAction,
2019
ControlMatch,
20+
ControlScope,
2121
EvaluationRequest,
2222
EvaluationResponse,
2323
EvaluatorResult,
@@ -43,12 +43,30 @@ def _compile_regex(pattern: str) -> Any:
4343
return re2.compile(pattern)
4444

4545

46+
class ControlDefinitionLike(Protocol):
47+
"""Runtime control shape required by the engine."""
48+
49+
enabled: bool
50+
execution: Literal["server", "sdk"]
51+
scope: ControlScope
52+
condition: ConditionNode
53+
action: ControlAction
54+
55+
4656
class ControlWithIdentity(Protocol):
4757
"""Protocol for a control with identity information."""
4858

49-
id: int
50-
name: str
51-
control: ControlDefinition | ControlDefinitionRuntime
59+
@property
60+
def id(self) -> int:
61+
"""Database identity for the control."""
62+
63+
@property
64+
def name(self) -> str:
65+
"""Human-readable name for the control."""
66+
67+
@property
68+
def control(self) -> ControlDefinitionLike:
69+
"""Runtime control payload used during evaluation."""
5270

5371

5472
@dataclass

‎models/src/agent_control_models/controls.py‎

Lines changed: 88 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -268,6 +268,8 @@ class SteeringContext(BaseModel):
268268
type JsonValue = JSONValue
269269

270270
_TEMPLATE_PARAMETER_NAME_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$")
271+
MAX_TEMPLATE_DEFINITION_DEPTH = 12
272+
MAX_TEMPLATE_DEFINITION_NODES = 1000
271273

272274

273275
def _validate_re2_value(value: str, *, field_name: str) -> str:
@@ -279,6 +281,35 @@ def _validate_re2_value(value: str, *, field_name: str) -> str:
279281
return value
280282

281283

284+
def _validate_template_definition_structure(value: JsonValue) -> JsonValue:
285+
"""Validate template definition nesting and overall size."""
286+
stack: list[tuple[JsonValue, int]] = [(value, 1)]
287+
node_count = 0
288+
289+
while stack:
290+
node, depth = stack.pop()
291+
node_count += 1
292+
293+
if depth > MAX_TEMPLATE_DEFINITION_DEPTH:
294+
raise ValueError(
295+
"definition_template nesting depth exceeds maximum of "
296+
f"{MAX_TEMPLATE_DEFINITION_DEPTH}"
297+
)
298+
299+
if node_count > MAX_TEMPLATE_DEFINITION_NODES:
300+
raise ValueError(
301+
"definition_template size exceeds maximum of "
302+
f"{MAX_TEMPLATE_DEFINITION_NODES} JSON nodes"
303+
)
304+
305+
if isinstance(node, dict):
306+
stack.extend((nested_value, depth + 1) for nested_value in node.values())
307+
elif isinstance(node, list):
308+
stack.extend((nested_value, depth + 1) for nested_value in node)
309+
310+
return value
311+
312+
282313
class TemplateParameterBase(BaseModel):
283314
"""Base definition for a template parameter."""
284315

@@ -416,6 +447,12 @@ def validate_parameter_names(
416447
)
417448
return value
418449

450+
@field_validator("definition_template")
451+
@classmethod
452+
def validate_definition_template_structure(cls, value: JsonValue) -> JsonValue:
453+
"""Limit template nesting and size to keep rendering bounded."""
454+
return _validate_template_definition_structure(value)
455+
419456

420457
class TemplateControlInput(BaseModel):
421458
"""Template-backed input payload for control create/update requests."""
@@ -626,7 +663,56 @@ def leaf_parts(self) -> ConditionLeafParts | None:
626663
ConditionNode.model_rebuild()
627664

628665

629-
class ControlDefinition(BaseModel):
666+
def _build_observability_identity(
667+
condition: ConditionNode,
668+
) -> ControlObservabilityIdentity:
669+
"""Build a stable selector/evaluator identity for a condition tree."""
670+
all_evaluators: list[str] = []
671+
all_selector_paths: list[str] = []
672+
seen_evaluators: set[str] = set()
673+
seen_selector_paths: set[str] = set()
674+
leaf_count = 0
675+
676+
for selector, evaluator in condition.iter_leaf_parts():
677+
leaf_count += 1
678+
selector_path = selector.path or "*"
679+
680+
if evaluator.name not in seen_evaluators:
681+
seen_evaluators.add(evaluator.name)
682+
all_evaluators.append(evaluator.name)
683+
684+
if selector_path not in seen_selector_paths:
685+
seen_selector_paths.add(selector_path)
686+
all_selector_paths.append(selector_path)
687+
688+
return ControlObservabilityIdentity(
689+
selector_path=all_selector_paths[0] if all_selector_paths else None,
690+
evaluator_name=all_evaluators[0] if all_evaluators else None,
691+
leaf_count=leaf_count,
692+
all_evaluators=all_evaluators,
693+
all_selector_paths=all_selector_paths,
694+
)
695+
696+
697+
class _ConditionBackedControlMixin:
698+
"""Shared helpers for control models backed by a condition tree."""
699+
700+
condition: ConditionNode
701+
702+
def iter_condition_leaves(self) -> Iterator[ConditionNode]:
703+
"""Yield leaf conditions in evaluation order."""
704+
yield from self.condition.iter_leaves()
705+
706+
def iter_condition_leaf_parts(self) -> Iterator[ConditionLeafParts]:
707+
"""Yield leaf selector/evaluator pairs in evaluation order."""
708+
yield from self.condition.iter_leaf_parts()
709+
710+
def observability_identity(self) -> ControlObservabilityIdentity:
711+
"""Return a deterministic representative identity for observability."""
712+
return _build_observability_identity(self.condition)
713+
714+
715+
class ControlDefinition(_ConditionBackedControlMixin, BaseModel):
630716
"""A control definition to evaluate agent interactions.
631717
632718
This model contains only the logic and configuration.
@@ -725,53 +811,12 @@ def to_template_control_input(self) -> TemplateControlInput:
725811
template_values=dict(self.template_values),
726812
)
727813

728-
def iter_condition_leaves(self) -> Iterator[ConditionNode]:
729-
"""Yield leaf conditions in evaluation order."""
730-
yield from self.condition.iter_leaves()
731-
732-
def iter_condition_leaf_parts(self) -> Iterator[ConditionLeafParts]:
733-
"""Yield leaf selector/evaluator pairs in evaluation order."""
734-
yield from self.condition.iter_leaf_parts()
735-
736814
def primary_leaf(self) -> ConditionNode | None:
737815
"""Return the single leaf node when the whole condition is just one leaf."""
738816
if self.condition.is_leaf():
739817
return self.condition
740818
return None
741819

742-
def observability_identity(self) -> ControlObservabilityIdentity:
743-
"""Return a deterministic representative identity for observability.
744-
745-
The representative selector/evaluator comes from the first leaf in
746-
evaluation order so composite trees still populate top-level event
747-
dimensions. The full ordered, deduped leaf context is also returned.
748-
"""
749-
all_evaluators: list[str] = []
750-
all_selector_paths: list[str] = []
751-
seen_evaluators: set[str] = set()
752-
seen_selector_paths: set[str] = set()
753-
leaf_count = 0
754-
755-
for selector, evaluator in self.iter_condition_leaf_parts():
756-
leaf_count += 1
757-
selector_path = selector.path or "*"
758-
759-
if evaluator.name not in seen_evaluators:
760-
seen_evaluators.add(evaluator.name)
761-
all_evaluators.append(evaluator.name)
762-
763-
if selector_path not in seen_selector_paths:
764-
seen_selector_paths.add(selector_path)
765-
all_selector_paths.append(selector_path)
766-
767-
return ControlObservabilityIdentity(
768-
selector_path=all_selector_paths[0] if all_selector_paths else None,
769-
evaluator_name=all_evaluators[0] if all_evaluators else None,
770-
leaf_count=leaf_count,
771-
all_evaluators=all_evaluators,
772-
all_selector_paths=all_selector_paths,
773-
)
774-
775820
model_config = {
776821
"json_schema_extra": {
777822
"examples": [
@@ -799,7 +844,7 @@ def observability_identity(self) -> ControlObservabilityIdentity:
799844
}
800845

801846

802-
class ControlDefinitionRuntime(BaseModel):
847+
class ControlDefinitionRuntime(_ConditionBackedControlMixin, BaseModel):
803848
"""Slim runtime control model that ignores template authoring metadata."""
804849

805850
model_config = ConfigDict(extra="ignore")
@@ -835,41 +880,6 @@ def validate_condition_constraints(self) -> Self:
835880
_validate_common_control_constraints(self.condition, self.action)
836881
return self
837882

838-
def iter_condition_leaves(self) -> Iterator[ConditionNode]:
839-
"""Yield leaf conditions in evaluation order."""
840-
yield from self.condition.iter_leaves()
841-
842-
def iter_condition_leaf_parts(self) -> Iterator[ConditionLeafParts]:
843-
"""Yield leaf selector/evaluator pairs in evaluation order."""
844-
yield from self.condition.iter_leaf_parts()
845-
846-
def observability_identity(self) -> ControlObservabilityIdentity:
847-
"""Return a deterministic representative identity for observability."""
848-
all_evaluators: list[str] = []
849-
all_selector_paths: list[str] = []
850-
seen_evaluators: set[str] = set()
851-
seen_selector_paths: set[str] = set()
852-
leaf_count = 0
853-
854-
for selector, evaluator in self.iter_condition_leaf_parts():
855-
leaf_count += 1
856-
selector_path = selector.path or "*"
857-
858-
if evaluator.name not in seen_evaluators:
859-
seen_evaluators.add(evaluator.name)
860-
all_evaluators.append(evaluator.name)
861-
862-
if selector_path not in seen_selector_paths:
863-
seen_selector_paths.add(selector_path)
864-
all_selector_paths.append(selector_path)
865-
866-
return ControlObservabilityIdentity(
867-
selector_path=all_selector_paths[0] if all_selector_paths else None,
868-
evaluator_name=all_evaluators[0] if all_evaluators else None,
869-
leaf_count=leaf_count,
870-
all_evaluators=all_evaluators,
871-
all_selector_paths=all_selector_paths,
872-
)
873883

874884
class EvaluatorResult(BaseModel):
875885
"""Result from a control evaluator.

‎models/tests/test_control_templates.py‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,13 @@
3636
}
3737

3838

39+
def _nested_template_value(depth: int) -> object:
40+
value: object = "leaf"
41+
for _ in range(depth):
42+
value = {"nested": value}
43+
return value
44+
45+
3946
def test_control_definition_requires_template_fields_together() -> None:
4047
with pytest.raises(
4148
ValidationError,
@@ -98,6 +105,32 @@ def test_template_definition_rejects_invalid_parameter_name() -> None:
98105
)
99106

100107

108+
def test_template_definition_rejects_excessive_nesting() -> None:
109+
with pytest.raises(
110+
ValidationError,
111+
match="definition_template nesting depth exceeds maximum",
112+
):
113+
TemplateDefinition.model_validate(
114+
{
115+
"parameters": {},
116+
"definition_template": _nested_template_value(12),
117+
}
118+
)
119+
120+
121+
def test_template_definition_rejects_excessive_size() -> None:
122+
with pytest.raises(
123+
ValidationError,
124+
match="definition_template size exceeds maximum",
125+
):
126+
TemplateDefinition.model_validate(
127+
{
128+
"parameters": {},
129+
"definition_template": list(range(1001)),
130+
}
131+
)
132+
133+
101134
def test_create_control_request_parses_template_payload_as_template_input() -> None:
102135
request = CreateControlRequest.model_validate(
103136
{

‎server/src/agent_control_server/endpoints/controls.py‎

Lines changed: 6 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
1-
from collections.abc import Iterator
2-
31
from agent_control_engine import list_evaluators
4-
from agent_control_models import ConditionNode, ControlDefinition, TemplateControlInput
2+
from agent_control_models import ControlDefinition, TemplateControlInput
53
from agent_control_models.errors import ErrorCode, ValidationErrorItem
64
from agent_control_models.server import (
75
AgentRef,
@@ -39,6 +37,7 @@
3937
)
4038
from ..logging_utils import get_logger
4139
from ..models import Agent, AgentData, Control, agent_controls, agent_policies, policy_controls
40+
from ..services.condition_traversal import iter_condition_leaves_with_paths
4241
from ..services.control_definitions import parse_control_definition_or_api_error
4342
from ..services.control_templates import remap_template_api_error, render_template_control_input
4443
from ..services.evaluator_utils import (
@@ -59,32 +58,6 @@
5958
template_router = APIRouter(prefix="/control-templates", tags=["controls"])
6059

6160
_logger = get_logger(__name__)
62-
63-
64-
def _iter_condition_leaves(
65-
node: ConditionNode,
66-
*,
67-
path: str = "data.condition",
68-
) -> Iterator[tuple[str, ConditionNode]]:
69-
"""Yield each leaf condition with its dot/bracket field path."""
70-
if node.is_leaf():
71-
yield path, node
72-
return
73-
74-
if node.and_ is not None:
75-
for index, child in enumerate(node.and_):
76-
yield from _iter_condition_leaves(child, path=f"{path}.and[{index}]")
77-
return
78-
79-
if node.or_ is not None:
80-
for index, child in enumerate(node.or_):
81-
yield from _iter_condition_leaves(child, path=f"{path}.or[{index}]")
82-
return
83-
84-
if node.not_ is not None:
85-
yield from _iter_condition_leaves(node.not_, path=f"{path}.not")
86-
87-
8861
def _serialize_control_definition(control_def: ControlDefinition) -> dict[str, object]:
8962
"""Serialize control data for storage while omitting null scope fields."""
9063
data_json = control_def.model_dump(
@@ -130,7 +103,10 @@ async def _validate_control_definition(
130103
"""Validate evaluator config for a control definition."""
131104
available_evaluators = list_evaluators()
132105
agent_data_by_name: dict[str, AgentData] = {}
133-
for field_prefix, leaf in _iter_condition_leaves(control_def.condition):
106+
for field_prefix, leaf in iter_condition_leaves_with_paths(
107+
control_def.condition,
108+
path="data.condition",
109+
):
134110
leaf_parts = leaf.leaf_parts()
135111
if leaf_parts is None:
136112
continue

0 commit comments

Comments
 (0)