From 8c2d9a8c09ceefbbc16201df31d4057630373032 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Thu, 4 Jun 2026 04:39:33 +0530 Subject: [PATCH] fix(adapters): propagate budget halts out of the LangChain callback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The callback handler raised BudgetExceeded / ApprovalDenied from its hooks, but LangChain swallows exceptions thrown inside a callback (it logs them and keeps running) unless the handler sets raise_error=True. So on a real LangChain agent the halt was silently dropped and the chain kept spending past budget — enforcement was a no-op. Set raise_error=True so the deterministic halt actually stops the run. Tests: a unit test drives the hooks against the stub daemon and asserts the loop budget halts at the cap (and that raise_error stays set, so this can't silently regress); a tool-gating test asserts a denied tool raises ApprovalDenied; and a skip-unless-langchain integration test runs a real FakeListLLM in a runaway loop and confirms it stops after exactly the budgeted number of calls. --- sdks/python/riskkernel/adapters/langchain.py | 7 +++ sdks/python/tests/test_sdk.py | 54 ++++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/sdks/python/riskkernel/adapters/langchain.py b/sdks/python/riskkernel/adapters/langchain.py index 462c4be..9b7dc6e 100644 --- a/sdks/python/riskkernel/adapters/langchain.py +++ b/sdks/python/riskkernel/adapters/langchain.py @@ -42,6 +42,13 @@ class RiskKernelCallbackHandler(_base_handler()): # type: ignore[misc] tool_side_effect: side-effect label reported for gated tools. """ + # LangChain swallows exceptions raised inside a callback — it logs them and + # keeps running — UNLESS the handler sets raise_error=True. Without this, a + # BudgetExceeded (or ApprovalDenied) raised in a hook below would be silently + # dropped and the chain would keep spending past its budget. This single flag + # is what makes the deterministic halt actually stop the LangChain run. + raise_error = True + def __init__(self, run: Run, gate_tools: bool = False, tool_side_effect: str = "tool"): self.run = run diff --git a/sdks/python/tests/test_sdk.py b/sdks/python/tests/test_sdk.py index e49a8ee..e968193 100644 --- a/sdks/python/tests/test_sdk.py +++ b/sdks/python/tests/test_sdk.py @@ -9,6 +9,15 @@ from riskkernel.errors import ApprovalDenied, BudgetExceeded +def _has_langchain() -> bool: + try: + import langchain_core # noqa: F401 + + return True + except Exception: + return False + + class _State: def __init__(self): self.steps = 0 @@ -170,6 +179,51 @@ def test_proxy_config(self): self.assertTrue(cfg["base_url"].endswith("/v1")) self.assertEqual(cfg["headers"]["X-RiskKernel-Run-Id"], "run-1") + # --- LangChain adapter --- + + def test_langchain_handler_enforces_loop_budget(self): + from riskkernel.adapters.langchain import RiskKernelCallbackHandler + + with self.rt.governed_run(name="t", budget=self.rt.budget(loops=2)) as run: + h = RiskKernelCallbackHandler(run) + # Must stay True, or LangChain silently swallows the halt and the chain + # keeps spending past budget (see the handler comment + integration test). + self.assertTrue(h.raise_error) + h.on_chat_model_start({}, []) # step 1 + h.on_chat_model_start({}, []) # step 2 + with self.assertRaises(BudgetExceeded) as cm: + h.on_chat_model_start({}, []) # step 3 → over the loop budget + self.assertEqual(cm.exception.reason, "loop_budget_exceeded") + + def test_langchain_handler_gates_denied_tool(self): + from riskkernel.adapters.langchain import RiskKernelCallbackHandler + + with self.rt.governed_run(name="t") as run: + run._client.get_approval = lambda _id: { + "id": "ap-1", "status": "denied", "reason": "no"} + h = RiskKernelCallbackHandler(run, gate_tools=True) + with self.assertRaises(ApprovalDenied): + h.on_tool_start({"name": "deploy"}, "ship it") + + @unittest.skipUnless(_has_langchain(), "langchain-core not installed") + def test_langchain_integration_stops_runaway_loop(self): + # The real proof: a runaway LangChain loop must actually stop. on_llm_start + # raises BudgetExceeded at the cap, and raise_error=True lets it propagate + # out of llm.invoke() — without that flag LangChain would swallow it. + from langchain_core.language_models.fake import FakeListLLM + + from riskkernel.adapters.langchain import RiskKernelCallbackHandler + + with self.rt.governed_run(name="t", budget=self.rt.budget(loops=2)) as run: + h = RiskKernelCallbackHandler(run) + llm = FakeListLLM(responses=["keep going"] * 10) + calls = 0 + with self.assertRaises(BudgetExceeded): + while True: # a deliberately runaway loop + llm.invoke("step", config={"callbacks": [h]}) + calls += 1 + self.assertEqual(calls, 2) # 2 calls allowed; the 3rd halted + if __name__ == "__main__": unittest.main()