From 9b4150475b040c9d3058be6daadd450b00cddacb Mon Sep 17 00:00:00 2001 From: NickHuo Date: Thu, 9 Jul 2026 12:34:52 -0700 Subject: [PATCH 1/4] Add config-driven agent output policies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Agents can declare `output_policies` (module:function refs) in the controller config. The global controller writes them to Redis per agent (agent::output_policies), and the local controller resolves and runs them automatically after the agent produces a result — no manual on_done registration by the caller. Policies are fire-and-forget side effects: the framework passes the output to each configured policy and ignores its return value, so policies observe/react (audit, notify, enforce externally) but do not modify the result. Framework only finds and runs; all decision logic lives inside the policy. --- examples/config/global_controller.yaml | 8 + examples/policies/finance.py | 61 ++++++ tests/run_tests.sh | 4 + tests/test_output_policy.py | 187 ++++++++++++++++++ ventis/controller/global_controller.py | 25 +++ ventis/controller/local_controller.py | 69 ++++++- .../templates/config/global_controller.yaml | 6 + 7 files changed, 358 insertions(+), 2 deletions(-) create mode 100644 examples/policies/finance.py create mode 100644 tests/test_output_policy.py diff --git a/examples/config/global_controller.yaml b/examples/config/global_controller.yaml index 8e8539d..0ba55ee 100644 --- a/examples/config/global_controller.yaml +++ b/examples/config/global_controller.yaml @@ -22,6 +22,14 @@ agents: memory: 512 # Memory in MB # Path to the agent entrypoint script entrypoint: agents/finance_agent.py + # Output policies run automatically after this agent produces a result. + # Each entry is a `module:function` reference resolved by the local + # controller. Policies are side-effect only (audit / notify / enforce): + # they react to the output but do not modify it, so their order does not + # matter. Decision logic lives inside the policy itself. + output_policies: + - policies.finance:audit_log + - policies.finance:alert_on_sensitive - name: MarketResearchAgent host: localhost diff --git a/examples/policies/finance.py b/examples/policies/finance.py new file mode 100644 index 0000000..6334b2d --- /dev/null +++ b/examples/policies/finance.py @@ -0,0 +1,61 @@ +# Example output policies for FinanceAgent. +# +# An output policy is any callable with the signature: +# +# policy(output, ctx) -> None +# +# - `output` is the agent method's return value. +# - `ctx` carries request metadata: at least `request_id`, `service`, +# `function`, plus whatever request context was propagated +# (e.g. `origin`, `tags`). +# +# Policies are SIDE-EFFECT ONLY: they observe/react to the output (audit, +# notify, enforce) but do not modify it. Any return value is ignored, so the +# policies are independent of one another and their order does not affect the +# result. The framework only finds and runs the configured policies; the +# decision of whether/how to act lives inside each policy. +# +# Bind them in the controller config: +# +# agents: +# - name: FinanceAgent +# output_policies: +# - policies.finance:audit_log +# - policies.finance:alert_on_sensitive + +import logging + +logger = logging.getLogger(__name__) + +# Toy list of tokens we never want to leak in a response. +_SENSITIVE = ("SSN", "password", "api_key") + + +def audit_log(output, ctx): + """Record every result for auditing. Does not modify the output.""" + logger.info( + "[audit] request=%s %s.%s -> %r", + ctx.get("request_id"), + ctx.get("service"), + ctx.get("function"), + output, + ) + + +def alert_on_sensitive(output, ctx): + """Warn if a result appears to contain sensitive tokens. + + Self-guards: only inspects string output; anything else is ignored. + Reacts (logs a warning) but does not change the output. + """ + if not isinstance(output, str): + return + for token in _SENSITIVE: + if token in output: + logger.warning( + "[alert] sensitive token %r in result of %s.%s (request=%s)", + token, + ctx.get("service"), + ctx.get("function"), + ctx.get("request_id"), + ) diff --git a/tests/run_tests.sh b/tests/run_tests.sh index e9e1dc7..fb4cfd3 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -47,6 +47,10 @@ ORIG_CWD=$(pwd) # Assuming the script was called from inside the ventis repo root SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" &> /dev/null && pwd)" +echo "-------------------------------------------" +echo ">> Running Output-Policy Unit Tests..." +python "$SCRIPT_DIR/test_output_policy.py" || exit 1 + echo "-------------------------------------------" echo ">> Running Integration Tests..." python "$SCRIPT_DIR/test_integration.py" || exit 1 diff --git a/tests/test_output_policy.py b/tests/test_output_policy.py new file mode 100644 index 0000000..f333f08 --- /dev/null +++ b/tests/test_output_policy.py @@ -0,0 +1,187 @@ +"""Unit tests for the agent output-policy mechanism (issues/003). + +Self-contained: no Docker / Redis / gRPC required. We stub the generated +gRPC modules so the controller module imports, inject fake policy modules +via sys.modules, and exercise the real +``LocalController._load_output_policies`` / ``_run_output_policies`` against +a fake Redis. Run: ``uv run python tests/test_output_policy.py``. + +Policies are side-effect only: each receives the same agent output, their +return value is ignored, and they never modify the result. +""" + +import json +import os +import sys +import types + +# Make the repo root importable so `import ventis...` works when run directly. +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) + +# --- Stub the generated gRPC modules the controller imports at module load. --- +_pb2 = types.ModuleType("local_controler_pb2") + + +class _JsonResponse: + def __init__(self, resonse=""): + self.resonse = resonse + + +_pb2.JsonResponse = _JsonResponse +sys.modules["local_controler_pb2"] = _pb2 + +_pb2_grpc = types.ModuleType("local_controler_pb2_grpc") + + +class _Servicer: # used as a base class at import time by the frontend + pass + + +class _Stub: + def __init__(self, channel): + pass + + +_pb2_grpc.LocalControllerServicer = _Servicer +_pb2_grpc.LocalControllerStub = _Stub +_pb2_grpc.add_LocalControllerServicer_to_server = lambda servicer, server: None +sys.modules["local_controler_pb2_grpc"] = _pb2_grpc + +from ventis.controller.local_controller import LocalController, OUTPUT_POLICY_KEY_FMT # noqa: E402 + + +# --- Fake Redis: only .get is used by _load_output_policies. --- +class FakeRedis: + def __init__(self, store=None): + self.store = store or {} + + def get(self, key): + return self.store.get(key) + + +# --- Fake policy module, resolved via importlib "module:function". --- +_seen = [] # (tag, output, request_id) recorded by side-effect policies + + +def _record_a(output, ctx): + _seen.append(("a", output, ctx.get("request_id"))) + + +def _record_b(output, ctx): + _seen.append(("b", output, ctx.get("request_id"))) + + +def _tries_to_transform(output, ctx): + # Returns a value on purpose; the framework must IGNORE it. + _seen.append(("t", output, ctx.get("request_id"))) + return "MUTATED" + + +def _boom(output, ctx): + raise RuntimeError("intentional policy failure") + + +_polmod = types.ModuleType("fake_policies") +_polmod.record_a = _record_a +_polmod.record_b = _record_b +_polmod.tries_to_transform = _tries_to_transform +_polmod.boom = _boom +sys.modules["fake_policies"] = _polmod + + +def _controller_with(refs): + """Build a bare controller wired to a fake Redis holding `refs` for 'Svc'.""" + lc = LocalController.__new__(LocalController) + lc._output_policies = {} + store = {} + if refs is not None: + store[OUTPUT_POLICY_KEY_FMT.format(service="Svc")] = json.dumps(refs) + lc.redis = FakeRedis(store) + return lc + + +def test_all_policies_run_on_same_output(): + _seen.clear() + lc = _controller_with(["fake_policies:record_a", "fake_policies:record_b"]) + ret = lc._run_output_policies("Svc", "hello", {"request_id": "r1"}) + # Side-effect only: nothing is returned... + assert ret is None, ret + # ...and every policy saw the SAME original output. + assert _seen == [("a", "hello", "r1"), ("b", "hello", "r1")], _seen + + +def test_return_value_is_ignored(): + _seen.clear() + lc = _controller_with(["fake_policies:tries_to_transform", "fake_policies:record_b"]) + lc._run_output_policies("Svc", "orig", {"request_id": "r2"}) + # The transform policy ran but its "MUTATED" return is dropped: record_b + # still receives the original output, not "MUTATED". + assert _seen == [("t", "orig", "r2"), ("b", "orig", "r2")], _seen + + +def test_no_policies_configured_is_noop(): + _seen.clear() + lc = _controller_with(None) # nothing in Redis + assert lc._run_output_policies("Svc", "x", {}) is None + assert _seen == [] + + +def test_exception_is_isolated(): + _seen.clear() + lc = _controller_with(["fake_policies:boom", "fake_policies:record_b"]) + lc._run_output_policies("Svc", "hi", {"request_id": "r3"}) + # boom raises and is skipped; record_b still runs. + assert _seen == [("b", "hi", "r3")], _seen + + +def test_unresolvable_reference_is_skipped(): + _seen.clear() + lc = _controller_with(["no_such_module:nope", "fake_policies:record_a"]) + lc._run_output_policies("Svc", "hi", {"request_id": "r4"}) + assert _seen == [("a", "hi", "r4")], _seen + + +def test_order_does_not_affect_result(): + # Both orderings run both policies; result (the passed-in output) is never + # changed regardless of order. + _seen.clear() + lc1 = _controller_with(["fake_policies:record_a", "fake_policies:record_b"]) + lc1._run_output_policies("Svc", "same", {}) + forward = list(_seen) + _seen.clear() + lc2 = _controller_with(["fake_policies:record_b", "fake_policies:record_a"]) + lc2._run_output_policies("Svc", "same", {}) + reverse = list(_seen) + # Different call order, but every policy saw the identical unchanged output. + assert {o for _, o, _ in forward} == {"same"} + assert {o for _, o, _ in reverse} == {"same"} + + +def test_resolution_is_cached(): + _seen.clear() + lc = _controller_with(["fake_policies:record_a"]) + lc._run_output_policies("Svc", "a", {}) + # Mutate the backing store; cache should mean it's not re-read. + lc.redis.store[OUTPUT_POLICY_KEY_FMT.format(service="Svc")] = json.dumps([]) + lc._run_output_policies("Svc", "b", {}) + assert [t for t, _, _ in _seen] == ["a", "a"], _seen + + +if __name__ == "__main__": + tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] + failed = 0 + for t in tests: + try: + t() + print(f"PASS {t.__name__}") + except AssertionError as e: + failed += 1 + print(f"FAIL {t.__name__}: {e}") + except Exception as e: # noqa: BLE001 + failed += 1 + print(f"ERROR {t.__name__}: {type(e).__name__}: {e}") + print("-" * 43) + if failed: + print(f"{failed}/{len(tests)} test(s) failed.") + sys.exit(1) + print(f"All {len(tests)} output-policy tests passed.") diff --git a/ventis/controller/global_controller.py b/ventis/controller/global_controller.py index 58655c1..fe9ebe2 100644 --- a/ventis/controller/global_controller.py +++ b/ventis/controller/global_controller.py @@ -81,6 +81,7 @@ def __init__(self, config_path): self._build_routing_table() self._write_resource_specs() self._load_and_write_policies() + self._write_output_policies() logger.info("Global controller initialized with %d controller(s).", len(self.controllers)) # Start background cleanup thread @@ -238,6 +239,30 @@ def _write_resource_specs(self): "replicas": str(len(placements)), }) + def _write_output_policies(self): + """Write each agent's configured output policies to Redis on every node. + + Output policies are declared per-agent in the controller config under + ``output_policies`` as a list of ``module:function`` references. They are + stored under ``agent::output_policies`` as a JSON list so the local + controller running that agent can resolve and run them after execution. + Agents without any output policies get an empty list (safe no-op). + """ + targets = list(self.node_redis.values()) if self.node_redis else [self.redis] + total = 0 + for ctrl in self.controllers: + name = ctrl["name"] + policies = ctrl.get("output_policies", []) or [] + policies_json = json.dumps(policies) + for redis_client in targets: + redis_client.set(f"agent:{name}:output_policies", policies_json) + if policies: + total += 1 + logger.info("Output policies for %s: %s", name, policies) + + logger.info("Output policies written to %d Redis instance(s) for %d agent(s).", + len(targets), total) + def _load_and_write_policies(self): """Load policy rules from config/policy.yaml and write to all Redis instances.""" config_dir = os.path.dirname(os.path.abspath(self.config_path)) diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 98e1c2b..85db53b 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -38,6 +38,7 @@ ROUTING_ENDPOINTS_KEY = "routing_table:endpoints" ROUTING_STATEFUL_KEY = "routing_table:stateful" POLICY_RULES_KEY = "policy:rules" +OUTPUT_POLICY_KEY_FMT = "agent:{service}:output_policies" class LocalController(object): @@ -72,6 +73,10 @@ def __init__(self, port=50051): # Policy rules cache (loaded lazily from Redis) self._policy_rules = None + # Output-policy cache: service name -> list of resolved callables + # (loaded lazily from Redis, resolved via importlib) + self._output_policies = {} + # Thread pool for executing agent methods concurrently. # This prevents deadlocks when an agent method creates nested Futures # that need to be routed through the same controller's request queue. @@ -159,6 +164,57 @@ def _check_policy(self, service, context): logger.warning("No policy rule matched for context=%s, denying access to %s", context, service) return False + # ------------------------------------------------------------------ # + # Output policies (run automatically after agent execution) # + # ------------------------------------------------------------------ # + + def _load_output_policies(self, service): + """Resolve the output policies configured for ``service`` into callables. + + Reads ``agent::output_policies`` (a JSON list of + ``module:function`` references) from Redis and resolves each into a + callable via importlib. Results are cached per service. References that + fail to resolve are logged and skipped so a bad reference cannot block + agent execution. + """ + if service in self._output_policies: + return self._output_policies[service] + + refs_json = self.redis.get(OUTPUT_POLICY_KEY_FMT.format(service=service)) + refs = json.loads(refs_json) if refs_json else [] + + resolved = [] + for ref in refs: + try: + module_name, _, func_name = ref.partition(":") + if not module_name or not func_name: + raise ValueError(f"expected 'module:function', got '{ref}'") + module = importlib.import_module(module_name) + fn = getattr(module, func_name) + resolved.append((ref, fn)) + logger.info("Loaded output policy '%s' for service %s", ref, service) + except Exception as e: + logger.error("Failed to load output policy '%s' for %s: %s", ref, service, e) + + self._output_policies[service] = resolved + return resolved + + def _run_output_policies(self, service, output, ctx): + """Run each configured output policy on the agent's result. + + Policies are side-effect only (audit, notify, enforce): every policy + receives the same agent output and its return value is ignored, so + policies do not modify the result and are independent of one another + (order does not affect correctness). A policy that raises is logged and + skipped so it cannot break the others or the result write-back. + Decision logic lives inside each policy — this method only dispatches. + """ + for ref, policy in self._load_output_policies(service): + try: + policy(output, ctx) + except Exception as e: + logger.error("Output policy '%s' failed for %s: %s", ref, service, e) + # ------------------------------------------------------------------ # # Endpoint resolution (affinity / load balancing) # # ------------------------------------------------------------------ # @@ -279,7 +335,7 @@ def _process_request(self, data): return if endpoint == self._my_endpoint: - self._executor.submit(self._execute_locally, service, function, args, future_id, origin, request_id) + self._executor.submit(self._execute_locally, service, function, args, future_id, origin, request_id, context) else: # Register the target as a consumer for any Future args # so results get pushed to its Redis via WriteResult. @@ -348,7 +404,7 @@ def _resolve_future_args(self, args, poll_interval=0.01, timeout=300): resolved[key] = value return resolved - def _execute_locally(self, service, function, args, future_id, origin=None, request_id=None): + def _execute_locally(self, service, function, args, future_id, origin=None, request_id=None, context=None): """Execute a request on the local agent and write the result to Redis.""" # Propagate the request_id context into this worker thread if request_id: @@ -370,6 +426,15 @@ def _execute_locally(self, service, function, args, future_id, origin=None, requ logger.info("Executing %s.%s (future=%s) locally", service, function, future_id) result = method(**args) + # Run any output policies bound to this agent, automatically. + # Policies react to the result (audit / notify / enforce); they do + # not modify it, so ordering between them doesn't affect the output. + policy_ctx = dict(context or {}) + policy_ctx.setdefault("request_id", request_id) + policy_ctx.setdefault("service", service) + policy_ctx.setdefault("function", function) + self._run_output_policies(service, result, policy_ctx) + # Serialize the result if isinstance(result, (dict, list)): serialized = json.dumps(result) diff --git a/ventis/templates/config/global_controller.yaml b/ventis/templates/config/global_controller.yaml index 74167bf..3469c76 100644 --- a/ventis/templates/config/global_controller.yaml +++ b/ventis/templates/config/global_controller.yaml @@ -11,6 +11,12 @@ agents: cpu: 1 memory: 512 entrypoint: agents/example_agent.py + # Optional: policies run automatically after the agent returns a result. + # Each entry is a `module:function` reference. Policies are side-effect + # only (audit / notify / enforce) — they react to the output but do not + # modify it, so their order does not matter. + # output_policies: + # - policies.example:audit_log - name: VllmAgent host: localhost From b4a87b9993505d3e0833d456e740e43154633c9a Mon Sep 17 00:00:00 2001 From: NickHuo Date: Thu, 9 Jul 2026 12:57:55 -0700 Subject: [PATCH 2/4] Harden output-policy loading and add real-module tests - Isolate policy loading (Redis read + JSON parse) inside _run_output_policies so a malformed stored value or Redis blip cannot propagate and let _execute_locally overwrite a good agent result. - Add tests exercising the real examples/policies/finance.py through the actual resolution + dispatch path, plus a load-failure isolation test. --- tests/test_output_policy.py | 87 +++++++++++++++++++++++++++ ventis/controller/local_controller.py | 9 ++- 2 files changed, 95 insertions(+), 1 deletion(-) diff --git a/tests/test_output_policy.py b/tests/test_output_policy.py index f333f08..f31f18c 100644 --- a/tests/test_output_policy.py +++ b/tests/test_output_policy.py @@ -141,6 +141,19 @@ def test_unresolvable_reference_is_skipped(): assert _seen == [("a", "hi", "r4")], _seen +def test_load_failure_is_isolated(): + # A malformed Redis value (or a Redis read error) during policy loading must + # NOT propagate out of _run_output_policies — otherwise _execute_locally's + # outer except would overwrite a successful agent result with an error. + _seen.clear() + lc = LocalController.__new__(LocalController) + lc._output_policies = {} + lc.redis = FakeRedis({OUTPUT_POLICY_KEY_FMT.format(service="Svc"): "{not json"}) + # Must return quietly rather than raise. + assert lc._run_output_policies("Svc", "hi", {"request_id": "r5"}) is None + assert _seen == [] + + def test_order_does_not_affect_result(): # Both orderings run both policies; result (the passed-in output) is never # changed regardless of order. @@ -167,6 +180,80 @@ def test_resolution_is_cached(): assert [t for t, _, _ in _seen] == ["a", "a"], _seen +# --- Real-module tests: exercise the actual examples/policies/finance.py --- +# through the real resolution + dispatch path (no fakes). This is what proves +# the module this PR ships is usable, not just the logic in isolation. +import logging # noqa: E402 + +_EXAMPLES_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "examples")) + + +def _capture(logger_name): + """Attach a collecting handler to `logger_name`; return (records, detach).""" + records = [] + + class _H(logging.Handler): + def emit(self, record): + records.append(record) + + lg = logging.getLogger(logger_name) + h = _H() + lg.addHandler(h) + prev_level = lg.level + lg.setLevel(logging.DEBUG) + + def detach(): + lg.removeHandler(h) + lg.setLevel(prev_level) + + return records, detach + + +def _real_policy_controller(): + # Put examples/ on sys.path so `import policies.finance` resolves — the + # same layout a local deploy runs in (project root == examples/). + if _EXAMPLES_DIR not in sys.path: + sys.path.insert(0, _EXAMPLES_DIR) + lc = LocalController.__new__(LocalController) + lc._output_policies = {} + lc.redis = FakeRedis({ + OUTPUT_POLICY_KEY_FMT.format(service="FinanceAgent"): json.dumps([ + "policies.finance:audit_log", + "policies.finance:alert_on_sensitive", + ]) + }) + return lc + + +def test_real_example_policies_resolve_and_run(): + lc = _real_policy_controller() + records, detach = _capture("policies.finance") + try: + ctx = {"request_id": "r9", "service": "FinanceAgent", "function": "run"} + ret = lc._run_output_policies("FinanceAgent", "leaked SSN here", ctx) + finally: + detach() + assert ret is None + msgs = [r.getMessage() for r in records] + # audit_log actually ran... + assert any("[audit]" in m for m in msgs), msgs + # ...and alert_on_sensitive fired on the SSN token. + assert any("sensitive token" in m and "SSN" in m for m in msgs), msgs + + +def test_real_example_alert_quiet_when_clean(): + lc = _real_policy_controller() + records, detach = _capture("policies.finance") + try: + ctx = {"request_id": "r10", "service": "FinanceAgent", "function": "run"} + lc._run_output_policies("FinanceAgent", "nothing sensitive", ctx) + finally: + detach() + msgs = [r.getMessage() for r in records] + assert any("[audit]" in m for m in msgs), msgs # audit still logs + assert not any("sensitive token" in m for m in msgs), msgs # no false alert + + if __name__ == "__main__": tests = [v for k, v in sorted(globals().items()) if k.startswith("test_")] failed = 0 diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index 85db53b..aad2f86 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -208,8 +208,15 @@ def _run_output_policies(self, service, output, ctx): (order does not affect correctness). A policy that raises is logged and skipped so it cannot break the others or the result write-back. Decision logic lives inside each policy — this method only dispatches. + Loading (Redis read + JSON parse) is isolated too, so a bad stored + value or a Redis blip cannot corrupt the agent's result write-back. """ - for ref, policy in self._load_output_policies(service): + try: + policies = self._load_output_policies(service) + except Exception as e: + logger.error("Failed to load output policies for %s: %s", service, e) + return + for ref, policy in policies: try: policy(output, ctx) except Exception as e: From 5eede1f9ec893941ca8694d1d87be60ec0ad2b4a Mon Sep 17 00:00:00 2001 From: NickHuo Date: Thu, 9 Jul 2026 13:14:04 -0700 Subject: [PATCH 3/4] Make framework policy-context fields authoritative Extract _policy_context and set request_id/service/function by direct assignment instead of setdefault, so a caller-controlled request context (from baggage) cannot spoof the service/request_id/function a policy sees. Add tests covering the override and the None-context case. --- tests/test_output_policy.py | 22 ++++++++++++++++++++++ ventis/controller/local_controller.py | 20 ++++++++++++++++---- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/tests/test_output_policy.py b/tests/test_output_policy.py index f31f18c..9df13d4 100644 --- a/tests/test_output_policy.py +++ b/tests/test_output_policy.py @@ -119,6 +119,28 @@ def test_return_value_is_ignored(): assert _seen == [("t", "orig", "r2"), ("b", "orig", "r2")], _seen +def test_policy_context_framework_fields_override_user_context(): + # The request context is caller-controlled (arrives via baggage). Framework + # fields must be authoritative, so a spoofed service/request_id/function in + # the user context cannot shadow the real ones handed to a policy. + ctx = LocalController._policy_context( + {"service": "SPOOF", "request_id": "evil", "function": "hack", "origin": "analyst"}, + request_id="real-req", + service="FinanceAgent", + function="run", + ) + assert ctx["service"] == "FinanceAgent" + assert ctx["request_id"] == "real-req" + assert ctx["function"] == "run" + # Non-framework keys from the request context are preserved. + assert ctx["origin"] == "analyst" + + +def test_policy_context_handles_none_context(): + ctx = LocalController._policy_context(None, request_id="r", service="S", function="f") + assert ctx == {"request_id": "r", "service": "S", "function": "f"} + + def test_no_policies_configured_is_noop(): _seen.clear() lc = _controller_with(None) # nothing in Redis diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index aad2f86..b088ec6 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -199,6 +199,21 @@ def _load_output_policies(self, service): self._output_policies[service] = resolved return resolved + @staticmethod + def _policy_context(context, request_id, service, function): + """Build the ctx handed to output policies. + + Starts from the request context (caller-controlled, arrives via + baggage) and then sets the framework fields authoritatively — they + OVERRIDE any same-named keys in the request context so a caller cannot + spoof the service / request_id / function a policy sees. + """ + ctx = dict(context or {}) + ctx["request_id"] = request_id + ctx["service"] = service + ctx["function"] = function + return ctx + def _run_output_policies(self, service, output, ctx): """Run each configured output policy on the agent's result. @@ -436,10 +451,7 @@ def _execute_locally(self, service, function, args, future_id, origin=None, requ # Run any output policies bound to this agent, automatically. # Policies react to the result (audit / notify / enforce); they do # not modify it, so ordering between them doesn't affect the output. - policy_ctx = dict(context or {}) - policy_ctx.setdefault("request_id", request_id) - policy_ctx.setdefault("service", service) - policy_ctx.setdefault("function", function) + policy_ctx = self._policy_context(context, request_id, service, function) self._run_output_policies(service, result, policy_ctx) # Serialize the result From 55fcbf9aaee821f7bde038e007f708d9a8b3e7f8 Mon Sep 17 00:00:00 2001 From: NickHuo Date: Thu, 9 Jul 2026 14:23:52 -0700 Subject: [PATCH 4/4] Address Copilot review: guard result write-back from policy plumbing - _policy_context coerces a non-mapping context to empty instead of letting dict() raise (a non-dict context from baggage could otherwise abort _execute_locally and clobber a successful result). - _run_output_policies hands policies a deepcopy of the output so a policy that mutates a dict/list in place cannot change what is serialized and written back; uncopyable outputs skip policies safely. - Add tests for non-dict context coercion and in-place mutation safety. --- tests/test_output_policy.py | 26 ++++++++++++++++++++++++++ ventis/controller/local_controller.py | 22 ++++++++++++++++++++-- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/tests/test_output_policy.py b/tests/test_output_policy.py index 9df13d4..055f974 100644 --- a/tests/test_output_policy.py +++ b/tests/test_output_policy.py @@ -81,11 +81,20 @@ def _boom(output, ctx): raise RuntimeError("intentional policy failure") +def _mutate_in_place(output, ctx): + # A misbehaving policy that mutates a mutable result in place. + if isinstance(output, dict): + output["injected"] = True + elif isinstance(output, list): + output.append("injected") + + _polmod = types.ModuleType("fake_policies") _polmod.record_a = _record_a _polmod.record_b = _record_b _polmod.tries_to_transform = _tries_to_transform _polmod.boom = _boom +_polmod.mutate_in_place = _mutate_in_place sys.modules["fake_policies"] = _polmod @@ -141,6 +150,23 @@ def test_policy_context_handles_none_context(): assert ctx == {"request_id": "r", "service": "S", "function": "f"} +def test_policy_context_coerces_non_dict_context(): + # A caller-controlled context that isn't a mapping (e.g., a string from + # malformed baggage) must not raise — it's coerced to just the framework + # fields so policy plumbing can never break the result write-back. + ctx = LocalController._policy_context("not-a-dict", request_id="r", service="S", function="f") + assert ctx == {"request_id": "r", "service": "S", "function": "f"} + + +def test_policy_cannot_mutate_caller_result(): + # A policy that mutates its argument in place must not affect the original + # result object the controller writes back to Redis. + lc = _controller_with(["fake_policies:mutate_in_place"]) + original = {"a": 1} + lc._run_output_policies("Svc", original, {"request_id": "r"}) + assert original == {"a": 1}, original # unchanged: policy saw a defensive copy + + def test_no_policies_configured_is_noop(): _seen.clear() lc = _controller_with(None) # nothing in Redis diff --git a/ventis/controller/local_controller.py b/ventis/controller/local_controller.py index b088ec6..b0565c6 100644 --- a/ventis/controller/local_controller.py +++ b/ventis/controller/local_controller.py @@ -2,6 +2,7 @@ # Starts the gRPC frontend server and polls the request queue for incoming requests. # Routes requests to the correct agent — either locally or by forwarding to another controller. +import copy import json import logging import os @@ -207,8 +208,12 @@ def _policy_context(context, request_id, service, function): baggage) and then sets the framework fields authoritatively — they OVERRIDE any same-named keys in the request context so a caller cannot spoof the service / request_id / function a policy sees. + + A non-mapping context (e.g., a malformed string/list from baggage) is + coerced to empty rather than raising, so policy plumbing can never + break the agent's result write-back. """ - ctx = dict(context or {}) + ctx = dict(context) if isinstance(context, dict) else {} ctx["request_id"] = request_id ctx["service"] = service ctx["function"] = function @@ -231,9 +236,22 @@ def _run_output_policies(self, service, output, ctx): except Exception as e: logger.error("Failed to load output policies for %s: %s", service, e) return + if not policies: + return + + # Hand policies a defensive copy: a policy that mutates a mutable + # result (dict/list) in place must not change what gets serialized and + # written back to the caller. If the result can't be copied, skip + # policies rather than risk corrupting the write-back. + try: + safe_output = copy.deepcopy(output) + except Exception as e: + logger.error("Could not copy output for policies of %s; skipping: %s", service, e) + return + for ref, policy in policies: try: - policy(output, ctx) + policy(safe_output, ctx) except Exception as e: logger.error("Output policy '%s' failed for %s: %s", ref, service, e)