Skip to content

Commit 4e6b025

Browse files
feat(evaluators): add extensions/metadata for additional fields that we want to pass as context to execution of evaluator (#276)
## Summary Pass the authenticated execution context through Agent Control to Galileo evaluators so scorer invocation can receive the Orbit-authorized organization, caller, project, and target information. - Preserve upstream authorization response extras as opaque `Principal.extensions` metadata. - Add an optional opaque `extensions` claim to newly minted runtime JWTs. Existing claims, expiration behavior, and tokens without extensions remain compatible. - Pass verified principal extensions through a provider-neutral evaluator hook. Existing evaluators continue using their current evaluation path. - In the Galileo evaluator only, translate the trusted extension envelope into Galileo's `execution_context`. For the current Orbit contract, `caller_id` is the authenticated user ID. A `log_stream` target's authenticated target ID becomes `run_id`. - Return a clear evaluator error when required authenticated fields are missing. ## Scope - API/runtime contract: additive opaque `extensions` claim; no named Galileo fields are added to generic principal or JWT models. - Internal: carry verified extensions from auth exchange through runtime-token verification and engine evaluation. - Galileo: serialize the local execution context with scorer invocation requests. - Out of scope: scorer grants, scorer authorization/execution, API-key forwarding to Runners, record factory changes. ## Trust and Compatibility Execution metadata comes only from the verified runtime principal. Evaluator step payloads are not used to construct it, and user-supplied tenant identity is not trusted. No scorer grant is acquired or cached, and the application `GALILEO_API_KEY` is not sent to Runners. Legacy runtime tokens without extensions and evaluators that do not override the extension hook remain supported. ## Testing and Coverage - [x] Engine tests: 120 passed, including opaque extension pass-through and legacy evaluator fallback. - [x] Built-in evaluator tests: 311 passed, including default extension-hook delegation. - [x] Galileo tests: 156 passed, including context serialization, missing caller handling, and API-key non-forwarding. - [x] Focused server auth tests: 6 passed, including malformed extension handling and signed-token round-trip. - [x] `make check` passed in PR CI; Codecov patch check passed. - [x] `make lint` and `make typecheck` passed locally.
1 parent 0b2a9d8 commit 4e6b025

15 files changed

Lines changed: 479 additions & 22 deletions

File tree

‎engine/src/agent_control_engine/core.py‎

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
EvaluationRequest,
2222
EvaluationResponse,
2323
EvaluatorResult,
24+
JSONObject,
2425
)
2526

2627
from .selectors import select_data
@@ -295,6 +296,7 @@ async def _evaluate_leaf(
295296
node: ConditionNode,
296297
request: EvaluationRequest,
297298
semaphore: asyncio.Semaphore,
299+
extensions: JSONObject | None,
298300
) -> _ConditionEvaluation:
299301
"""Evaluate a leaf selector/evaluator pair.
300302
@@ -320,7 +322,7 @@ async def _evaluate_leaf(
320322
timeout = DEFAULT_EVALUATOR_TIMEOUT
321323

322324
result = await asyncio.wait_for(
323-
evaluator.evaluate_with_context(data, request.step),
325+
evaluator.evaluate_with_extensions(data, request.step, extensions),
324326
timeout=timeout,
325327
)
326328
except TimeoutError:
@@ -440,17 +442,20 @@ async def _evaluate_condition(
440442
node: ConditionNode,
441443
request: EvaluationRequest,
442444
semaphore: asyncio.Semaphore,
445+
extensions: JSONObject | None,
443446
) -> _ConditionEvaluation:
444447
"""Evaluate a recursive condition tree."""
445448
if node.is_leaf():
446-
return await self._evaluate_leaf(item, node, request, semaphore)
449+
return await self._evaluate_leaf(item, node, request, semaphore, extensions)
447450

448451
kind = node.kind()
449452
children = node.children_in_order()
450453
child_evaluations: list[_ConditionEvaluation] = []
451454

452455
if kind == "not":
453-
child_eval = await self._evaluate_condition(item, children[0], request, semaphore)
456+
child_eval = await self._evaluate_condition(
457+
item, children[0], request, semaphore, extensions
458+
)
454459
trace = {
455460
"type": "not",
456461
"evaluated": True,
@@ -478,7 +483,9 @@ async def _evaluate_condition(
478483
return _ConditionEvaluation(result=result, trace=trace)
479484

480485
for index, child in enumerate(children):
481-
child_eval = await self._evaluate_condition(item, child, request, semaphore)
486+
child_eval = await self._evaluate_condition(
487+
item, child, request, semaphore, extensions
488+
)
482489
child_evaluations.append(child_eval)
483490

484491
if child_eval.result.error:
@@ -624,14 +631,19 @@ def get_applicable_controls(
624631

625632
return applicable
626633

627-
async def process(self, request: EvaluationRequest) -> EvaluationResponse:
634+
async def process(
635+
self,
636+
request: EvaluationRequest,
637+
extensions: JSONObject | None = None,
638+
) -> EvaluationResponse:
628639
"""Process controls in parallel with cancel-on-deny.
629640
630641
All applicable controls are evaluated concurrently. If any control
631642
matches with action=deny, remaining evaluations are cancelled.
632643
633644
Args:
634645
request: The evaluation request containing step and context
646+
extensions: Opaque authenticated metadata supplied by the caller.
635647
636648
Returns:
637649
EvaluationResponse with is_safe status and any matches
@@ -667,6 +679,7 @@ async def evaluate_control(eval_task: _EvalTask) -> None:
667679
eval_task.item.control.condition,
668680
request,
669681
semaphore,
682+
extensions,
670683
)
671684
eval_task.result = evaluation.result
672685

‎engine/tests/test_core.py‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
EvaluationRequest,
2121
EvaluatorResult,
2222
EvaluatorSpec,
23+
JSONObject,
2324
SteeringContext,
2425
Step,
2526
)
@@ -40,14 +41,16 @@ class SimpleConfig(BaseModel):
4041
_execution_log: list[str] = []
4142
_blocker_event: asyncio.Event | None = None
4243
_context_calls: list[tuple[Any, Step]] = []
44+
_extension_calls: list[JSONObject | None] = []
4345

4446

4547
def reset_test_state() -> None:
4648
"""Reset shared test state."""
47-
global _execution_log, _blocker_event, _context_calls
49+
global _execution_log, _blocker_event, _context_calls, _extension_calls
4850
_execution_log = []
4951
_blocker_event = asyncio.Event()
5052
_context_calls = []
53+
_extension_calls = []
5154

5255

5356
class AllowEvaluator(Evaluator[SimpleConfig]):
@@ -182,6 +185,12 @@ async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult:
182185
_context_calls.append((data, step))
183186
return EvaluatorResult(matched=False, confidence=1.0, message="context received")
184187

188+
async def evaluate_with_extensions(
189+
self, data: Any, step: Step, extensions: JSONObject | None
190+
) -> EvaluatorResult:
191+
_extension_calls.append(extensions)
192+
return await self.evaluate_with_context(data, step)
193+
185194

186195
@dataclass
187196
class MockControlWithIdentity:
@@ -303,6 +312,49 @@ async def test_context_evaluator_receives_selected_data_and_complete_step() -> N
303312
assert _context_calls == [("answer", step)]
304313

305314

315+
@pytest.mark.asyncio
316+
async def test_engine_passes_opaque_extensions_to_evaluator() -> None:
317+
# Given: an authenticated extension mapping and a context-aware evaluator
318+
engine = ControlEngine(
319+
[make_control(1, "context", "test-context", action="observe", path="input")]
320+
)
321+
extensions: JSONObject = {
322+
"namespace_key": "namespace-1",
323+
"target_type": "custom_target",
324+
"target_id": "target-9",
325+
"metadata": {"provider_field": "opaque_value"},
326+
}
327+
328+
# When: evaluating a request with the trusted extensions
329+
await engine.process(
330+
EvaluationRequest(
331+
agent_name="00000000-0000-0000-0000-000000000001",
332+
step=Step(type="llm", name="test-step", input="question"),
333+
stage="pre",
334+
),
335+
extensions=extensions,
336+
)
337+
338+
# Then: the engine passes the mapping unchanged without interpreting it
339+
assert _extension_calls == [extensions]
340+
341+
342+
@pytest.mark.asyncio
343+
async def test_default_extension_hook_preserves_existing_evaluator_behavior() -> None:
344+
# Given: an evaluator that implements only the original evaluate method
345+
evaluator = AllowEvaluator(SimpleConfig())
346+
step = Step(type="llm", name="test-step", input="question")
347+
348+
# When: the new hook is called with opaque extensions
349+
result = await evaluator.evaluate_with_extensions(
350+
"question", step, {"provider_field": "opaque_value"}
351+
)
352+
353+
# Then: the existing evaluation method still handles the request
354+
assert result.message == "Allowed"
355+
assert _execution_log == ["allow:default:start", "allow:default:end"]
356+
357+
306358
@pytest.mark.asyncio
307359
async def test_cached_context_evaluator_handles_concurrent_steps_without_retaining_state() -> None:
308360
# Given: one cached evaluator configuration and two independent requests

‎evaluators/builtin/src/agent_control_evaluators/_base.py‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
from dataclasses import dataclass
88
from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar
99

10-
from agent_control_models import EvaluatorResult, Step
10+
from agent_control_models import EvaluatorResult, JSONObject, Step
1111
from agent_control_models.base import BaseModel
1212

1313
if TYPE_CHECKING:
@@ -178,6 +178,19 @@ async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult:
178178
"""
179179
return await self.evaluate(data)
180180

181+
async def evaluate_with_extensions(
182+
self,
183+
data: Any,
184+
step: Step,
185+
extensions: JSONObject | None,
186+
) -> EvaluatorResult:
187+
"""Evaluate with opaque trusted request extensions.
188+
189+
Existing evaluators retain their behavior because the default
190+
implementation delegates to :meth:`evaluate_with_context`.
191+
"""
192+
return await self.evaluate_with_context(data, step)
193+
181194
def get_timeout_seconds(self) -> float:
182195
"""Get timeout in seconds from config or metadata default."""
183196
timeout_ms: int = getattr(self.config, "timeout_ms", self.metadata.timeout_ms)

‎evaluators/builtin/tests/test_base.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,19 @@ async def test_contextual_evaluation_delegates_to_existing_evaluate(self):
121121
assert result.matched is True
122122
assert result.metadata == {"data": "selected data"}
123123

124+
@pytest.mark.asyncio
125+
async def test_extension_evaluation_delegates_to_existing_context_hook(self):
126+
"""The extension hook preserves behavior for legacy evaluators."""
127+
evaluator = MockEvaluator.from_dict({"should_match": True})
128+
step = Step(type="llm", name="answer", input="full input")
129+
130+
result = await evaluator.evaluate_with_extensions(
131+
"selected data", step, {"provider_field": "opaque_value"}
132+
)
133+
134+
assert result.matched is True
135+
assert result.metadata == {"data": "selected data"}
136+
124137
def test_evaluator_config_stored(self):
125138
"""Test that evaluator stores config."""
126139
evaluator = MockEvaluator.from_dict({"should_match": True})

‎evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -228,6 +228,15 @@ class ScorerInvokeRecord(BaseModel):
228228
dataset_output: JSONValue = None
229229

230230

231+
class GalileoExecutionContext(BaseModel):
232+
"""Authenticated identity fields required by Galileo scorer invocation."""
233+
234+
organization_id: str = Field(min_length=1)
235+
user_id: str = Field(min_length=1)
236+
project_id: str = Field(min_length=1)
237+
run_id: str = Field(min_length=1)
238+
239+
231240
class ScorerInvokeRequest(BaseModel):
232241
"""Request payload for Luna scorer invocation.
233242
@@ -246,6 +255,7 @@ class ScorerInvokeRequest(BaseModel):
246255
scorer_label: str | None = Field(default=None, min_length=1)
247256
inputs: ScorerInvokeInputs
248257
record: ScorerInvokeRecord | None = None
258+
execution_context: GalileoExecutionContext | None = None
249259
config: ScorerInvokeConfig = Field(default_factory=ScorerInvokeConfig)
250260

251261
@model_validator(mode="after")
@@ -474,6 +484,7 @@ async def invoke(
474484
input: JSONValue = None,
475485
output: JSONValue = None,
476486
step: Step | None = None,
487+
execution_context: GalileoExecutionContext | None = None,
477488
config: ScorerInvokeConfig | JSONObject | None = None,
478489
timeout: float = DEFAULT_TIMEOUT_SECS,
479490
headers: dict[str, str] | None = None,
@@ -488,6 +499,7 @@ async def invoke(
488499
input: Optional user/system prompt text.
489500
output: Optional model response text.
490501
step: Optional complete runtime step used for structured dual-write.
502+
execution_context: Optional authenticated Galileo scorer context.
491503
config: Optional Orbit-supported scorer invocation configuration.
492504
timeout: Request timeout in seconds.
493505
headers: Additional request headers.
@@ -531,6 +543,7 @@ async def invoke(
531543
selected_input=input,
532544
selected_output=output,
533545
),
546+
execution_context=execution_context,
534547
config=invoke_config,
535548
).to_dict()
536549

‎evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py‎

Lines changed: 71 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,9 @@
1010

1111
import httpx
1212
from agent_control_evaluators import Evaluator, EvaluatorMetadata, register_evaluator
13-
from agent_control_models import EvaluatorResult, JSONValue, Step
13+
from agent_control_models import EvaluatorResult, JSONObject, JSONValue, Step
1414

15-
from .client import GalileoLunaClient, ScorerInvokeResponse
15+
from .client import GalileoExecutionContext, GalileoLunaClient, ScorerInvokeResponse
1616
from .config import LunaEvaluatorConfig, coerce_number
1717

1818
logger = logging.getLogger(__name__)
@@ -212,7 +212,72 @@ async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult:
212212
"""
213213
return await self._evaluate(data, step=step)
214214

215-
async def _evaluate(self, data: Any, *, step: Step | None) -> EvaluatorResult:
215+
async def evaluate_with_extensions(
216+
self,
217+
data: Any,
218+
step: Step,
219+
extensions: JSONObject | None,
220+
) -> EvaluatorResult:
221+
"""Evaluate using opaque authenticated metadata supplied by Agent Control."""
222+
return await self._evaluate(
223+
data,
224+
step=step,
225+
extensions=extensions,
226+
)
227+
228+
@staticmethod
229+
def _execution_context_from_extensions(
230+
extensions: JSONObject | None,
231+
) -> GalileoExecutionContext | None:
232+
"""Translate trusted opaque auth metadata into Galileo's request context."""
233+
if extensions is None:
234+
return None
235+
236+
metadata = extensions.get("metadata")
237+
metadata_obj = metadata if isinstance(metadata, dict) else {}
238+
organization_id = extensions.get("namespace_key")
239+
# Orbit's runtime-auth contract defines caller_id as the authenticated
240+
# user ID for this flow. Keep that contract mapping inside Galileo.
241+
user_id = extensions.get("caller_id")
242+
project_id = metadata_obj.get("project_id")
243+
target_type = extensions.get("target_type")
244+
target_id = extensions.get("target_id")
245+
run_id = target_id if target_type == "log_stream" else None
246+
247+
missing = [
248+
field
249+
for field, value in (
250+
("organization_id", organization_id),
251+
("user_id", user_id),
252+
("project_id", project_id),
253+
("run_id", run_id),
254+
)
255+
if not isinstance(value, str) or not value
256+
]
257+
if missing:
258+
raise ValueError(
259+
"Authenticated execution metadata is missing required fields: "
260+
+ ", ".join(missing)
261+
)
262+
263+
assert isinstance(organization_id, str)
264+
assert isinstance(user_id, str)
265+
assert isinstance(project_id, str)
266+
assert isinstance(run_id, str)
267+
return GalileoExecutionContext(
268+
organization_id=organization_id,
269+
user_id=user_id,
270+
project_id=project_id,
271+
run_id=run_id,
272+
)
273+
274+
async def _evaluate(
275+
self,
276+
data: Any,
277+
*,
278+
step: Step | None,
279+
extensions: JSONObject | None = None,
280+
) -> EvaluatorResult:
216281
"""Run a Luna evaluation with optional structured runtime context."""
217282
input_text, output_text = self._prepare_payload(data)
218283
if not (_has_text(input_text) or _has_text(output_text)):
@@ -224,9 +289,12 @@ async def _evaluate(self, data: Any, *, step: Step | None) -> EvaluatorResult:
224289
)
225290

226291
try:
292+
execution_context = self._execution_context_from_extensions(extensions)
227293
scorer_kwargs = self._scorer_kwargs()
228294
if step is not None:
229295
scorer_kwargs["step"] = step
296+
if execution_context is not None:
297+
scorer_kwargs["execution_context"] = execution_context
230298
response = await self._get_client().invoke(
231299
**scorer_kwargs,
232300
input=input_text if _has_text(input_text) else None,

0 commit comments

Comments
 (0)