diff --git a/engine/src/agent_control_engine/core.py b/engine/src/agent_control_engine/core.py index 7dbba3c7..ab170e82 100644 --- a/engine/src/agent_control_engine/core.py +++ b/engine/src/agent_control_engine/core.py @@ -21,6 +21,7 @@ EvaluationRequest, EvaluationResponse, EvaluatorResult, + JSONObject, ) from .selectors import select_data @@ -295,6 +296,7 @@ async def _evaluate_leaf( node: ConditionNode, request: EvaluationRequest, semaphore: asyncio.Semaphore, + extensions: JSONObject | None, ) -> _ConditionEvaluation: """Evaluate a leaf selector/evaluator pair. @@ -320,7 +322,7 @@ async def _evaluate_leaf( timeout = DEFAULT_EVALUATOR_TIMEOUT result = await asyncio.wait_for( - evaluator.evaluate_with_context(data, request.step), + evaluator.evaluate_with_extensions(data, request.step, extensions), timeout=timeout, ) except TimeoutError: @@ -440,17 +442,20 @@ async def _evaluate_condition( node: ConditionNode, request: EvaluationRequest, semaphore: asyncio.Semaphore, + extensions: JSONObject | None, ) -> _ConditionEvaluation: """Evaluate a recursive condition tree.""" if node.is_leaf(): - return await self._evaluate_leaf(item, node, request, semaphore) + return await self._evaluate_leaf(item, node, request, semaphore, extensions) kind = node.kind() children = node.children_in_order() child_evaluations: list[_ConditionEvaluation] = [] if kind == "not": - child_eval = await self._evaluate_condition(item, children[0], request, semaphore) + child_eval = await self._evaluate_condition( + item, children[0], request, semaphore, extensions + ) trace = { "type": "not", "evaluated": True, @@ -478,7 +483,9 @@ async def _evaluate_condition( return _ConditionEvaluation(result=result, trace=trace) for index, child in enumerate(children): - child_eval = await self._evaluate_condition(item, child, request, semaphore) + child_eval = await self._evaluate_condition( + item, child, request, semaphore, extensions + ) child_evaluations.append(child_eval) if child_eval.result.error: @@ -624,7 +631,11 @@ def get_applicable_controls( return applicable - async def process(self, request: EvaluationRequest) -> EvaluationResponse: + async def process( + self, + request: EvaluationRequest, + extensions: JSONObject | None = None, + ) -> EvaluationResponse: """Process controls in parallel with cancel-on-deny. All applicable controls are evaluated concurrently. If any control @@ -632,6 +643,7 @@ async def process(self, request: EvaluationRequest) -> EvaluationResponse: Args: request: The evaluation request containing step and context + extensions: Opaque authenticated metadata supplied by the caller. Returns: EvaluationResponse with is_safe status and any matches @@ -667,6 +679,7 @@ async def evaluate_control(eval_task: _EvalTask) -> None: eval_task.item.control.condition, request, semaphore, + extensions, ) eval_task.result = evaluation.result diff --git a/engine/tests/test_core.py b/engine/tests/test_core.py index e16d437f..34c620d5 100644 --- a/engine/tests/test_core.py +++ b/engine/tests/test_core.py @@ -20,6 +20,7 @@ EvaluationRequest, EvaluatorResult, EvaluatorSpec, + JSONObject, SteeringContext, Step, ) @@ -40,14 +41,16 @@ class SimpleConfig(BaseModel): _execution_log: list[str] = [] _blocker_event: asyncio.Event | None = None _context_calls: list[tuple[Any, Step]] = [] +_extension_calls: list[JSONObject | None] = [] def reset_test_state() -> None: """Reset shared test state.""" - global _execution_log, _blocker_event, _context_calls + global _execution_log, _blocker_event, _context_calls, _extension_calls _execution_log = [] _blocker_event = asyncio.Event() _context_calls = [] + _extension_calls = [] class AllowEvaluator(Evaluator[SimpleConfig]): @@ -182,6 +185,12 @@ async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult: _context_calls.append((data, step)) return EvaluatorResult(matched=False, confidence=1.0, message="context received") + async def evaluate_with_extensions( + self, data: Any, step: Step, extensions: JSONObject | None + ) -> EvaluatorResult: + _extension_calls.append(extensions) + return await self.evaluate_with_context(data, step) + @dataclass class MockControlWithIdentity: @@ -303,6 +312,49 @@ async def test_context_evaluator_receives_selected_data_and_complete_step() -> N assert _context_calls == [("answer", step)] +@pytest.mark.asyncio +async def test_engine_passes_opaque_extensions_to_evaluator() -> None: + # Given: an authenticated extension mapping and a context-aware evaluator + engine = ControlEngine( + [make_control(1, "context", "test-context", action="observe", path="input")] + ) + extensions: JSONObject = { + "namespace_key": "namespace-1", + "target_type": "custom_target", + "target_id": "target-9", + "metadata": {"provider_field": "opaque_value"}, + } + + # When: evaluating a request with the trusted extensions + await engine.process( + EvaluationRequest( + agent_name="00000000-0000-0000-0000-000000000001", + step=Step(type="llm", name="test-step", input="question"), + stage="pre", + ), + extensions=extensions, + ) + + # Then: the engine passes the mapping unchanged without interpreting it + assert _extension_calls == [extensions] + + +@pytest.mark.asyncio +async def test_default_extension_hook_preserves_existing_evaluator_behavior() -> None: + # Given: an evaluator that implements only the original evaluate method + evaluator = AllowEvaluator(SimpleConfig()) + step = Step(type="llm", name="test-step", input="question") + + # When: the new hook is called with opaque extensions + result = await evaluator.evaluate_with_extensions( + "question", step, {"provider_field": "opaque_value"} + ) + + # Then: the existing evaluation method still handles the request + assert result.message == "Allowed" + assert _execution_log == ["allow:default:start", "allow:default:end"] + + @pytest.mark.asyncio async def test_cached_context_evaluator_handles_concurrent_steps_without_retaining_state() -> None: # Given: one cached evaluator configuration and two independent requests diff --git a/evaluators/builtin/src/agent_control_evaluators/_base.py b/evaluators/builtin/src/agent_control_evaluators/_base.py index d83c0c3a..100ebd80 100644 --- a/evaluators/builtin/src/agent_control_evaluators/_base.py +++ b/evaluators/builtin/src/agent_control_evaluators/_base.py @@ -7,7 +7,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar -from agent_control_models import EvaluatorResult, Step +from agent_control_models import EvaluatorResult, JSONObject, Step from agent_control_models.base import BaseModel if TYPE_CHECKING: @@ -178,6 +178,19 @@ async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult: """ return await self.evaluate(data) + async def evaluate_with_extensions( + self, + data: Any, + step: Step, + extensions: JSONObject | None, + ) -> EvaluatorResult: + """Evaluate with opaque trusted request extensions. + + Existing evaluators retain their behavior because the default + implementation delegates to :meth:`evaluate_with_context`. + """ + return await self.evaluate_with_context(data, step) + def get_timeout_seconds(self) -> float: """Get timeout in seconds from config or metadata default.""" timeout_ms: int = getattr(self.config, "timeout_ms", self.metadata.timeout_ms) diff --git a/evaluators/builtin/tests/test_base.py b/evaluators/builtin/tests/test_base.py index d81cc3f5..5dfe6915 100644 --- a/evaluators/builtin/tests/test_base.py +++ b/evaluators/builtin/tests/test_base.py @@ -121,6 +121,19 @@ async def test_contextual_evaluation_delegates_to_existing_evaluate(self): assert result.matched is True assert result.metadata == {"data": "selected data"} + @pytest.mark.asyncio + async def test_extension_evaluation_delegates_to_existing_context_hook(self): + """The extension hook preserves behavior for legacy evaluators.""" + evaluator = MockEvaluator.from_dict({"should_match": True}) + step = Step(type="llm", name="answer", input="full input") + + result = await evaluator.evaluate_with_extensions( + "selected data", step, {"provider_field": "opaque_value"} + ) + + assert result.matched is True + assert result.metadata == {"data": "selected data"} + def test_evaluator_config_stored(self): """Test that evaluator stores config.""" evaluator = MockEvaluator.from_dict({"should_match": True}) diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py index d00d1457..cca27c66 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/client.py @@ -228,6 +228,15 @@ class ScorerInvokeRecord(BaseModel): dataset_output: JSONValue = None +class GalileoExecutionContext(BaseModel): + """Authenticated identity fields required by Galileo scorer invocation.""" + + organization_id: str = Field(min_length=1) + user_id: str = Field(min_length=1) + project_id: str = Field(min_length=1) + run_id: str = Field(min_length=1) + + class ScorerInvokeRequest(BaseModel): """Request payload for Luna scorer invocation. @@ -246,6 +255,7 @@ class ScorerInvokeRequest(BaseModel): scorer_label: str | None = Field(default=None, min_length=1) inputs: ScorerInvokeInputs record: ScorerInvokeRecord | None = None + execution_context: GalileoExecutionContext | None = None config: ScorerInvokeConfig = Field(default_factory=ScorerInvokeConfig) @model_validator(mode="after") @@ -474,6 +484,7 @@ async def invoke( input: JSONValue = None, output: JSONValue = None, step: Step | None = None, + execution_context: GalileoExecutionContext | None = None, config: ScorerInvokeConfig | JSONObject | None = None, timeout: float = DEFAULT_TIMEOUT_SECS, headers: dict[str, str] | None = None, @@ -488,6 +499,7 @@ async def invoke( input: Optional user/system prompt text. output: Optional model response text. step: Optional complete runtime step used for structured dual-write. + execution_context: Optional authenticated Galileo scorer context. config: Optional Orbit-supported scorer invocation configuration. timeout: Request timeout in seconds. headers: Additional request headers. @@ -531,6 +543,7 @@ async def invoke( selected_input=input, selected_output=output, ), + execution_context=execution_context, config=invoke_config, ).to_dict() diff --git a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py index e4147a9a..99085e26 100644 --- a/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py +++ b/evaluators/contrib/galileo/src/agent_control_evaluator_galileo/luna/evaluator.py @@ -10,9 +10,9 @@ import httpx from agent_control_evaluators import Evaluator, EvaluatorMetadata, register_evaluator -from agent_control_models import EvaluatorResult, JSONValue, Step +from agent_control_models import EvaluatorResult, JSONObject, JSONValue, Step -from .client import GalileoLunaClient, ScorerInvokeResponse +from .client import GalileoExecutionContext, GalileoLunaClient, ScorerInvokeResponse from .config import LunaEvaluatorConfig, coerce_number logger = logging.getLogger(__name__) @@ -212,7 +212,72 @@ async def evaluate_with_context(self, data: Any, step: Step) -> EvaluatorResult: """ return await self._evaluate(data, step=step) - async def _evaluate(self, data: Any, *, step: Step | None) -> EvaluatorResult: + async def evaluate_with_extensions( + self, + data: Any, + step: Step, + extensions: JSONObject | None, + ) -> EvaluatorResult: + """Evaluate using opaque authenticated metadata supplied by Agent Control.""" + return await self._evaluate( + data, + step=step, + extensions=extensions, + ) + + @staticmethod + def _execution_context_from_extensions( + extensions: JSONObject | None, + ) -> GalileoExecutionContext | None: + """Translate trusted opaque auth metadata into Galileo's request context.""" + if extensions is None: + return None + + metadata = extensions.get("metadata") + metadata_obj = metadata if isinstance(metadata, dict) else {} + organization_id = extensions.get("namespace_key") + # Orbit's runtime-auth contract defines caller_id as the authenticated + # user ID for this flow. Keep that contract mapping inside Galileo. + user_id = extensions.get("caller_id") + project_id = metadata_obj.get("project_id") + target_type = extensions.get("target_type") + target_id = extensions.get("target_id") + run_id = target_id if target_type == "log_stream" else None + + missing = [ + field + for field, value in ( + ("organization_id", organization_id), + ("user_id", user_id), + ("project_id", project_id), + ("run_id", run_id), + ) + if not isinstance(value, str) or not value + ] + if missing: + raise ValueError( + "Authenticated execution metadata is missing required fields: " + + ", ".join(missing) + ) + + assert isinstance(organization_id, str) + assert isinstance(user_id, str) + assert isinstance(project_id, str) + assert isinstance(run_id, str) + return GalileoExecutionContext( + organization_id=organization_id, + user_id=user_id, + project_id=project_id, + run_id=run_id, + ) + + async def _evaluate( + self, + data: Any, + *, + step: Step | None, + extensions: JSONObject | None = None, + ) -> EvaluatorResult: """Run a Luna evaluation with optional structured runtime context.""" input_text, output_text = self._prepare_payload(data) 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: ) try: + execution_context = self._execution_context_from_extensions(extensions) scorer_kwargs = self._scorer_kwargs() if step is not None: scorer_kwargs["step"] = step + if execution_context is not None: + scorer_kwargs["execution_context"] = execution_context response = await self._get_client().invoke( **scorer_kwargs, input=input_text if _has_text(input_text) else None, diff --git a/evaluators/contrib/galileo/tests/test_luna_evaluator.py b/evaluators/contrib/galileo/tests/test_luna_evaluator.py index 4334ac5e..b368f24d 100644 --- a/evaluators/contrib/galileo/tests/test_luna_evaluator.py +++ b/evaluators/contrib/galileo/tests/test_luna_evaluator.py @@ -627,6 +627,7 @@ def test_client_reports_invalid_connection_tuning_env( @pytest.mark.asyncio async def test_client_posts_to_luna_invoke_scorer_invoke(self) -> None: from agent_control_evaluator_galileo.luna import GalileoLunaClient + from agent_control_evaluator_galileo.luna.client import GalileoExecutionContext captured: dict[str, object] = {} @@ -654,6 +655,12 @@ def handler(request: httpx.Request) -> httpx.Response: scorer_id=SCORER_ID, input="user prompt", output="model answer", + execution_context=GalileoExecutionContext( + organization_id="org-1", + user_id="user-2", + project_id="project-3", + run_id="run-4", + ), config={"request_timeout_seconds": 7}, headers={"Galileo-API-Key": "blocked", "X-Request-ID": "safe-id"}, ) @@ -665,6 +672,12 @@ def handler(request: httpx.Request) -> httpx.Response: assert captured["url"] == "http://luna-invoke:8090/api/v1/scorers/invoke" expected_body = { "scorer_id": SCORER_ID, + "execution_context": { + "organization_id": "org-1", + "user_id": "user-2", + "project_id": "project-3", + "run_id": "run-4", + }, "inputs": {"query": "user prompt", "response": "model answer"}, "config": {"request_timeout_seconds": 7.0}, } @@ -1049,6 +1062,71 @@ async def test_evaluator_contextual_hook_forwards_complete_step(self) -> None: timeout=10.0, ) + @patch.dict(os.environ, LUNA_ENV) + @pytest.mark.asyncio + async def test_evaluator_consumes_verified_opaque_extensions(self) -> None: + from agent_control_evaluator_galileo.luna import LunaEvaluator, ScorerInvokeResponse + from agent_control_evaluator_galileo.luna.client import ( + GalileoExecutionContext, + GalileoLunaClient, + ) + + evaluator = LunaEvaluator.from_dict( + {"scorer_id": "scorer-123", "threshold": 0.5, "operator": "gte"} + ) + extensions = { + "namespace_key": "org-1", + "caller_id": "verified-user-2", + "target_type": "log_stream", + "target_id": "run-4", + "metadata": {"project_id": "project-3"}, + } + + with patch.object(GalileoLunaClient, "invoke", new_callable=AsyncMock) as mock_invoke: + mock_invoke.return_value = ScorerInvokeResponse(score=0.8, status="success") + result = await evaluator.evaluate_with_extensions( + "selected input", Step(type="llm", name="answer", input="prompt"), extensions + ) + + assert result.matched is True + mock_invoke.assert_awaited_once_with( + scorer_id="scorer-123", + step=Step(type="llm", name="answer", input="prompt"), + execution_context=GalileoExecutionContext( + organization_id="org-1", + user_id="verified-user-2", + project_id="project-3", + run_id="run-4", + ), + input="selected input", + output=None, + config=None, + timeout=10.0, + ) + + @patch.dict(os.environ, LUNA_ENV) + @pytest.mark.asyncio + async def test_evaluator_requires_authenticated_caller_id(self) -> None: + from agent_control_evaluator_galileo.luna import LunaEvaluator + from agent_control_evaluator_galileo.luna.client import GalileoLunaClient + + evaluator = LunaEvaluator.from_dict({"scorer_id": "scorer-123"}) + extensions = { + "namespace_key": "org-1", + "target_type": "log_stream", + "target_id": "run-4", + "metadata": {"project_id": "project-3"}, + } + + with patch.object(GalileoLunaClient, "invoke", new_callable=AsyncMock) as mock_invoke: + result = await evaluator.evaluate_with_extensions( + "selected input", Step(type="llm", name="answer", input="prompt"), extensions + ) + + assert result.error is not None + assert "user_id" in result.error + mock_invoke.assert_not_called() + @patch.dict(os.environ, LUNA_ENV) @pytest.mark.asyncio async def test_evaluator_labels_forwarded_scorer_version_id_as_requested(self) -> None: diff --git a/server/src/agent_control_server/auth_framework/core.py b/server/src/agent_control_server/auth_framework/core.py index c0c0a59f..dd32c070 100644 --- a/server/src/agent_control_server/auth_framework/core.py +++ b/server/src/agent_control_server/auth_framework/core.py @@ -28,6 +28,7 @@ from enum import StrEnum from typing import Any, Protocol, runtime_checkable +from agent_control_models import JSONObject from fastapi import Request @@ -83,6 +84,8 @@ class Principal: grant_expires_at: When the upstream grant expires. Used by the runtime-token exchange endpoint to bound the local token's lifetime. + extensions: Opaque trusted metadata returned by the authorizer. + Generic authorization code preserves this without interpreting it. """ namespace_key: str @@ -92,6 +95,7 @@ class Principal: target_id: str | None = None scopes: tuple[str, ...] = () grant_expires_at: datetime | None = None + extensions: JSONObject | None = None ContextBuilder = Callable[[Request], dict[str, Any] | Awaitable[dict[str, Any]]] diff --git a/server/src/agent_control_server/auth_framework/providers/http_upstream.py b/server/src/agent_control_server/auth_framework/providers/http_upstream.py index c2a18b4f..464ed2cb 100644 --- a/server/src/agent_control_server/auth_framework/providers/http_upstream.py +++ b/server/src/agent_control_server/auth_framework/providers/http_upstream.py @@ -48,6 +48,7 @@ from typing import Any import httpx +from agent_control_models import JSONObject from agent_control_models.errors import ErrorCode, ErrorReason from fastapi import Request from prometheus_client import Counter, Histogram @@ -55,6 +56,7 @@ BaseModel, ConfigDict, Field, + TypeAdapter, ValidationError, field_validator, model_validator, @@ -80,6 +82,7 @@ "Duration of auth upstream HTTP attempts made by Agent Control.", ("operation", "outcome"), ) +_JSON_OBJECT_ADAPTER: TypeAdapter[JSONObject] = TypeAdapter(JSONObject) class _UpstreamGrant(BaseModel): @@ -91,7 +94,7 @@ class _UpstreamGrant(BaseModel): with a 502. """ - model_config = ConfigDict(extra="ignore", strict=True) + model_config = ConfigDict(extra="allow", strict=True) namespace_key: str = Field(min_length=1) is_admin: bool = False @@ -449,6 +452,21 @@ def _parse_principal(self, response: httpx.Response) -> Principal: hint="Contact the operator.", ) from exc + try: + extensions = _JSON_OBJECT_ADAPTER.validate_python(grant.model_extra or {}) + except ValidationError as exc: + _logger.error( + "Auth upstream returned malformed extension metadata: %s", + exc.errors(), + ) + raise APIError( + status_code=502, + error_code=ErrorCode.AUTH_MISCONFIGURED, + reason=ErrorReason.INTERNAL_ERROR, + detail="Authorization service returned malformed extension metadata.", + hint="Contact the operator.", + ) from exc + return Principal( namespace_key=grant.namespace_key, is_admin=grant.is_admin, @@ -457,6 +475,7 @@ def _parse_principal(self, response: httpx.Response) -> Principal: target_id=grant.target_id, scopes=grant.scopes, grant_expires_at=grant.expires_at, + extensions=extensions or None, ) diff --git a/server/src/agent_control_server/auth_framework/providers/local_jwt.py b/server/src/agent_control_server/auth_framework/providers/local_jwt.py index dd208739..0b5b10bb 100644 --- a/server/src/agent_control_server/auth_framework/providers/local_jwt.py +++ b/server/src/agent_control_server/auth_framework/providers/local_jwt.py @@ -101,6 +101,7 @@ async def resolve_identity(self, request: Request, operation: Operation) -> Prin target_id=claims.target_id, scopes=claims.scopes, grant_expires_at=claims.expires_at, + extensions=claims.extensions, ) async def authorize( diff --git a/server/src/agent_control_server/auth_framework/runtime_token.py b/server/src/agent_control_server/auth_framework/runtime_token.py index 54c59fbb..e40e008e 100644 --- a/server/src/agent_control_server/auth_framework/runtime_token.py +++ b/server/src/agent_control_server/auth_framework/runtime_token.py @@ -33,10 +33,13 @@ from typing import Any import jwt +from agent_control_models import JSONObject +from pydantic import TypeAdapter, ValidationError _ALGORITHM = "HS256" _ISSUER = "agent-control/server" _DOMAIN = "runtime" +_JSON_OBJECT_ADAPTER: TypeAdapter[JSONObject] = TypeAdapter(JSONObject) class RuntimeTokenError(Exception): @@ -65,6 +68,7 @@ class RuntimeTokenClaims: expires_at: datetime issued_at: datetime jti: str + extensions: JSONObject | None = None def mint_runtime_token( @@ -77,6 +81,7 @@ def mint_runtime_token( secret: str, ttl_seconds: int, upstream_expires_at: datetime | None = None, + extensions: JSONObject | None = None, now: datetime | None = None, ) -> tuple[str, RuntimeTokenClaims]: """Mint a runtime token. Returns ``(token, claims)``. @@ -116,6 +121,12 @@ def mint_runtime_token( "(e.g., tz=UTC); naive datetimes are not supported." ) + if extensions is not None: + try: + extensions = _JSON_OBJECT_ADAPTER.validate_python(extensions) + except ValidationError as exc: + raise RuntimeTokenError("Runtime token extensions must be a JSON object.") from exc + issued_at = now or datetime.now(UTC) if upstream_expires_at is not None and upstream_expires_at <= issued_at: # Minting with an already-expired ``exp`` would return a 200 with @@ -142,6 +153,8 @@ def mint_runtime_token( "exp": int(expires_at.timestamp()), "jti": jti, } + if extensions: + payload["extensions"] = extensions token = jwt.encode(payload, secret, algorithm=_ALGORITHM) claims = RuntimeTokenClaims( namespace_key=namespace_key, @@ -152,6 +165,7 @@ def mint_runtime_token( expires_at=expires_at, issued_at=issued_at, jti=jti, + extensions=extensions or None, ) return token, claims @@ -187,6 +201,7 @@ def verify_runtime_token(token: str, secret: str) -> RuntimeTokenClaims: actor_id = payload.get("actor_id") target_type = payload.get("target_type") target_id = payload.get("target_id") + raw_extensions = payload.get("extensions") if not isinstance(namespace_key, str) or not namespace_key: raise RuntimeTokenError("Runtime token missing namespace_key.") if not isinstance(actor_id, str) or not actor_id: @@ -196,6 +211,13 @@ def verify_runtime_token(token: str, secret: str) -> RuntimeTokenClaims: if not isinstance(target_id, str) or not target_id: raise RuntimeTokenError("Runtime token missing target_id.") + extensions: JSONObject | None = None + if raw_extensions is not None: + try: + extensions = _JSON_OBJECT_ADAPTER.validate_python(raw_extensions) + except ValidationError as exc: + raise RuntimeTokenError("Runtime token has malformed extensions.") from exc + raw_scopes = payload.get("scopes", []) if not isinstance(raw_scopes, list) or not all(isinstance(s, str) for s in raw_scopes): raise RuntimeTokenError("Runtime token has malformed scopes.") @@ -214,4 +236,5 @@ def verify_runtime_token(token: str, secret: str) -> RuntimeTokenClaims: expires_at=datetime.fromtimestamp(payload["exp"], tz=UTC), issued_at=datetime.fromtimestamp(payload["iat"], tz=UTC), jti=jti, + extensions=extensions, ) diff --git a/server/src/agent_control_server/endpoints/auth.py b/server/src/agent_control_server/endpoints/auth.py index 2372eb6f..f2797bee 100644 --- a/server/src/agent_control_server/endpoints/auth.py +++ b/server/src/agent_control_server/endpoints/auth.py @@ -211,6 +211,7 @@ async def runtime_token_exchange( secret=config.secret, ttl_seconds=config.ttl_seconds, upstream_expires_at=principal.grant_expires_at, + extensions=principal.extensions, ) except UpstreamGrantExpiredError as exc: # Upstream returned a grant whose ``expires_at`` is already in diff --git a/server/src/agent_control_server/endpoints/evaluation.py b/server/src/agent_control_server/endpoints/evaluation.py index a31d757d..7a46829e 100644 --- a/server/src/agent_control_server/endpoints/evaluation.py +++ b/server/src/agent_control_server/endpoints/evaluation.py @@ -9,6 +9,7 @@ ControlMatch, EvaluationRequest, EvaluationResponse, + JSONObject, ) from agent_control_models.errors import ErrorCode, ValidationErrorItem from fastapi import APIRouter, Depends, Request @@ -117,6 +118,22 @@ def _sanitize_evaluation_response(response: EvaluationResponse) -> EvaluationRes ) +def _principal_extensions(principal: Principal) -> JSONObject | None: + """Build the opaque evaluator extension from the authenticated principal.""" + if principal.target_type is None or principal.target_id is None: + return None + + extensions: JSONObject = { + "namespace_key": principal.namespace_key, + "target_type": principal.target_type, + "target_id": principal.target_id, + "metadata": principal.extensions or {}, + } + if principal.caller_id is not None: + extensions["caller_id"] = principal.caller_id + return extensions + + async def _evaluation_context(request: Request) -> dict[str, object]: """Surface target identifiers to the runtime authorizer.""" try: @@ -199,7 +216,7 @@ async def evaluate( engine_controls = await _load_engine_controls(request, principal) engine = ControlEngine(engine_controls) try: - raw_response = await engine.process(request) + raw_response = await engine.process(request, extensions=_principal_extensions(principal)) except ValueError: _logger.exception("Evaluation failed due to invalid configuration or input") raise APIValidationError( diff --git a/server/tests/test_auth_framework.py b/server/tests/test_auth_framework.py index eee6789a..80400c25 100644 --- a/server/tests/test_auth_framework.py +++ b/server/tests/test_auth_framework.py @@ -842,6 +842,51 @@ def test_runtime_token_round_trips(): assert decoded.scopes == ("runtime.use",) +def test_runtime_token_rejects_non_json_extensions_when_minting(): + from agent_control_server.auth_framework.runtime_token import ( + RuntimeTokenError, + mint_runtime_token, + ) + + with pytest.raises(RuntimeTokenError, match="extensions must be a JSON object"): + mint_runtime_token( + namespace_key="default", + actor_id="actor", + target_type="target", + target_id="target-id", + scopes=("runtime.use",), + secret=_TEST_SECRET, + ttl_seconds=60, + extensions={"provider_field": object()}, + ) + + +def test_runtime_token_rejects_malformed_signed_extensions(): + import jwt + + from agent_control_server.auth_framework.runtime_token import ( + RuntimeTokenError, + mint_runtime_token, + verify_runtime_token, + ) + + token, _ = mint_runtime_token( + namespace_key="default", + actor_id="actor", + target_type="target", + target_id="target-id", + scopes=("runtime.use",), + secret=_TEST_SECRET, + ttl_seconds=60, + ) + payload = jwt.decode(token, _TEST_SECRET, algorithms=["HS256"]) + payload["extensions"] = [] + malformed_token = jwt.encode(payload, _TEST_SECRET, algorithm="HS256") + + with pytest.raises(RuntimeTokenError, match="malformed extensions"): + verify_runtime_token(malformed_token, _TEST_SECRET) + + def test_runtime_token_rejects_wrong_secret(): from agent_control_server.auth_framework.runtime_token import ( RuntimeTokenError, @@ -1333,11 +1378,13 @@ async def test_http_upstream_accepts_iso_datetime_and_array_scopes(): lambda req: httpx.Response( 200, json={ - "namespace_key": "org-1", + "namespace_key": "namespace-1", "is_admin": False, "scopes": ["runtime.use", "runtime.read_only"], - "target_type": "log_stream", - "target_id": "ls-1", + "target_type": "custom_target", + "target_id": "target-1", + "provider_field": "opaque_value", + "nested": {"opaque": "metadata"}, "expires_at": iso_expiry, }, ) @@ -1345,16 +1392,42 @@ async def test_http_upstream_accepts_iso_datetime_and_array_scopes(): principal = await provider.authorize( _build_request(), Operation.RUNTIME_TOKEN_EXCHANGE, - context={"target_type": "log_stream", "target_id": "ls-1"}, + context={"target_type": "custom_target", "target_id": "target-1"}, ) - assert principal.namespace_key == "org-1" + assert principal.namespace_key == "namespace-1" assert principal.scopes == ("runtime.use", "runtime.read_only") - assert principal.target_type == "log_stream" - assert principal.target_id == "ls-1" + assert principal.target_type == "custom_target" + assert principal.target_id == "target-1" + assert principal.extensions == { + "provider_field": "opaque_value", + "nested": {"opaque": "metadata"}, + } assert principal.grant_expires_at is not None assert principal.grant_expires_at.isoformat() == iso_expiry +@pytest.mark.asyncio +async def test_http_upstream_rejects_invalid_extension_metadata(monkeypatch): + from pydantic import TypeAdapter + + from agent_control_server.auth_framework.providers import http_upstream + + provider = _build_upstream( + lambda req: httpx.Response( + 200, + json={"namespace_key": "namespace-1", "provider_field": "opaque_value"}, + ) + ) + # Force the defensive conversion error path for invalid extension payloads. + monkeypatch.setattr(http_upstream, "_JSON_OBJECT_ADAPTER", TypeAdapter(str)) + + with pytest.raises(APIError) as exc_info: + await provider.authorize(_build_request(), Operation.RUNTIME_TOKEN_EXCHANGE) + + assert exc_info.value.status_code == 502 + assert "malformed extension metadata" in exc_info.value.detail + + @pytest.mark.asyncio async def test_http_upstream_rejects_target_grant_mismatch(): provider = _build_upstream( diff --git a/server/tests/test_runtime_token_exchange_endpoint.py b/server/tests/test_runtime_token_exchange_endpoint.py index a4c5d440..6a3a8645 100644 --- a/server/tests/test_runtime_token_exchange_endpoint.py +++ b/server/tests/test_runtime_token_exchange_endpoint.py @@ -14,8 +14,11 @@ from datetime import UTC, datetime, timedelta from unittest.mock import patch +import httpx +import jwt import pytest from fastapi.testclient import TestClient +from starlette.requests import Request from agent_control_server.auth_framework import Operation, Principal from agent_control_server.auth_framework.config import ( @@ -26,9 +29,8 @@ clear_authorizers, set_authorizer, ) -from agent_control_server.auth_framework.providers import ( - LocalJwtVerifyProvider, -) +from agent_control_server.auth_framework.providers import HttpUpstreamAuthProvider, LocalJwtVerifyProvider +from agent_control_server.auth_framework.providers.http_upstream import HttpUpstreamConfig from agent_control_server.auth_framework.runtime_token import RuntimeTokenError _TEST_SECRET = "test-runtime-secret-12345678901234567890" @@ -125,6 +127,73 @@ def test_exchange_endpoint_mints_token_when_configured( assert body["expires_at"] +@pytest.mark.asyncio +async def test_upstream_extensions_round_trip_through_signed_runtime_token( + runtime_config_enabled, +): + from agent_control_server.endpoints.auth import ( + RuntimeTokenExchangeRequest, + runtime_token_exchange, + ) + + trusted_metadata = {"provider_field": "opaque_value", "nested": {"opaque": "value"}} + + def authz_response(_request: httpx.Request) -> httpx.Response: + return httpx.Response( + 200, + json={ + "namespace_key": "namespace-1", + "caller_id": "opaque-caller", + "target_type": "custom_target", + "target_id": "target-9", + "scopes": ["runtime.use"], + **trusted_metadata, + }, + ) + + async with httpx.AsyncClient(transport=httpx.MockTransport(authz_response)) as client: + provider = HttpUpstreamAuthProvider( + HttpUpstreamConfig(url="https://authz.test/check"), client=client + ) + upstream_request = Request( + {"type": "http", "headers": [], "method": "POST", "path": "/"} + ) + principal = await provider.authorize( + upstream_request, + Operation.RUNTIME_TOKEN_EXCHANGE, + context={"target_type": "custom_target", "target_id": "target-9"}, + ) + + assert principal.extensions == trusted_metadata + response = await runtime_token_exchange( + RuntimeTokenExchangeRequest(target_type="custom_target", target_id="target-9"), + principal=principal, + ) + token_payload = jwt.decode( + response.token, + _TEST_SECRET, + algorithms=["HS256"], + issuer="agent-control/server", + ) + assert "provider_field" not in token_payload + assert token_payload["extensions"] == trusted_metadata + + verify_request = Request( + { + "type": "http", + "headers": [(b"authorization", f"Bearer {response.token}".encode())], + "method": "POST", + "path": "/", + } + ) + verified_principal = await LocalJwtVerifyProvider(secret=_TEST_SECRET).authorize( + verify_request, + Operation.RUNTIME_USE, + context={"target_type": "custom_target", "target_id": "target-9"}, + ) + assert verified_principal.extensions == trusted_metadata + + @pytest.mark.parametrize( ("actor_id", "expected_log_actor_id"), [