hexgate 0.1.1__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (98) hide show
  1. fortify/__init__.py +78 -0
  2. fortify/adapters/__init__.py +0 -0
  3. fortify/adapters/google/__init__.py +3 -0
  4. fortify/adapters/google/runner.py +106 -0
  5. fortify/adapters/google/tools.py +57 -0
  6. fortify/adapters/google/wrapper.py +47 -0
  7. fortify/adapters/langchain/__init__.py +7 -0
  8. fortify/adapters/langchain/agent.py +133 -0
  9. fortify/adapters/langchain/tools.py +248 -0
  10. fortify/adapters/langchain/wrapper.py +70 -0
  11. fortify/adapters/openai/__init__.py +3 -0
  12. fortify/adapters/openai/runner.py +125 -0
  13. fortify/adapters/openai/tools.py +57 -0
  14. fortify/adapters/openai/wrapper.py +47 -0
  15. fortify/adapters/pydantic_ai/__init__.py +7 -0
  16. fortify/adapters/pydantic_ai/agent.py +111 -0
  17. fortify/adapters/pydantic_ai/tools.py +40 -0
  18. fortify/adapters/pydantic_ai/wrapper.py +94 -0
  19. fortify/agents/__init__.py +57 -0
  20. fortify/agents/builtin/__init__.py +1 -0
  21. fortify/agents/builtin/researcher/agent.yaml +7 -0
  22. fortify/agents/builtin/researcher/policy.yaml +10 -0
  23. fortify/agents/builtin/researcher/system.md +5 -0
  24. fortify/agents/factory.py +683 -0
  25. fortify/agents/loader.py +790 -0
  26. fortify/agents/models.py +17 -0
  27. fortify/agents/prompts/agent_system.md +14 -0
  28. fortify/audit.py +292 -0
  29. fortify/bootstrap.py +32 -0
  30. fortify/cli/__init__.py +42 -0
  31. fortify/cli/_common.py +306 -0
  32. fortify/cli/chat.py +282 -0
  33. fortify/cli/policy/__init__.py +17 -0
  34. fortify/cli/policy/main.py +527 -0
  35. fortify/cli/register/__init__.py +8 -0
  36. fortify/cli/register/fortify.py +139 -0
  37. fortify/cli/register/google.py +104 -0
  38. fortify/cli/register/langchain.py +76 -0
  39. fortify/cli/register/main.py +113 -0
  40. fortify/cli/register/manifest.py +65 -0
  41. fortify/cli/register/models.py +107 -0
  42. fortify/cli/register/openai.py +78 -0
  43. fortify/cli/register/pydantic_ai.py +114 -0
  44. fortify/cli/register/register.py +71 -0
  45. fortify/cli/serve.py +337 -0
  46. fortify/cli/state.py +156 -0
  47. fortify/cloud/__init__.py +32 -0
  48. fortify/cloud/attenuate.py +105 -0
  49. fortify/cloud/biscuit.py +172 -0
  50. fortify/cloud/client.py +327 -0
  51. fortify/config/__init__.py +1 -0
  52. fortify/config/settings.py +50 -0
  53. fortify/runtime/__init__.py +53 -0
  54. fortify/runtime/command_policy.py +289 -0
  55. fortify/runtime/context.py +129 -0
  56. fortify/runtime/sandbox_runtime.py +115 -0
  57. fortify/runtime/srt.py +88 -0
  58. fortify/runtime/workspace.py +330 -0
  59. fortify/security/__init__.py +141 -0
  60. fortify/security/bundle.py +399 -0
  61. fortify/security/constraints.py +252 -0
  62. fortify/security/decision.py +145 -0
  63. fortify/security/enforcer.py +83 -0
  64. fortify/security/errors.py +11 -0
  65. fortify/security/file_scope.py +78 -0
  66. fortify/security/models.py +59 -0
  67. fortify/security/policy.py +182 -0
  68. fortify/security/policy_set.py +266 -0
  69. fortify/security/rego.py +334 -0
  70. fortify/security/rego_wasm.py +213 -0
  71. fortify/security/signing.py +122 -0
  72. fortify/security/source.py +372 -0
  73. fortify/security/wasm_engine.py +486 -0
  74. fortify/streaming/__init__.py +57 -0
  75. fortify/streaming/events.py +206 -0
  76. fortify/streaming/normalize.py +429 -0
  77. fortify/tools/__init__.py +21 -0
  78. fortify/tools/bash.py +49 -0
  79. fortify/tools/decorators.py +158 -0
  80. fortify/tools/fetch.py +72 -0
  81. fortify/tools/files/__init__.py +9 -0
  82. fortify/tools/files/_common.py +35 -0
  83. fortify/tools/files/edit_file.py +57 -0
  84. fortify/tools/files/glob.py +48 -0
  85. fortify/tools/files/grep.py +91 -0
  86. fortify/tools/files/read_file.py +39 -0
  87. fortify/tools/files/write_file.py +36 -0
  88. fortify/tools/refund.py +53 -0
  89. fortify/tools/websearch.py +72 -0
  90. fortify/tracing/__init__.py +1 -0
  91. fortify/tracing/langfuse.py +68 -0
  92. fortify/utils/__init__.py +1 -0
  93. fortify/utils/retry.py +58 -0
  94. hexgate-0.1.1.dist-info/METADATA +1440 -0
  95. hexgate-0.1.1.dist-info/RECORD +98 -0
  96. hexgate-0.1.1.dist-info/WHEEL +5 -0
  97. hexgate-0.1.1.dist-info/entry_points.txt +2 -0
  98. hexgate-0.1.1.dist-info/top_level.txt +1 -0
fortify/__init__.py ADDED
@@ -0,0 +1,78 @@
1
+ """Public package surface for fortify."""
2
+
3
+ from fortify.agents.factory import (
4
+ create_agent,
5
+ enforce_policy,
6
+ invoke_agent,
7
+ stream_agent,
8
+ stream_agent_raw,
9
+ )
10
+ from fortify.agents.loader import (
11
+ clear_registered_agents,
12
+ list_available_agents,
13
+ list_builtin_agents,
14
+ list_local_agents,
15
+ list_registered_agents,
16
+ load_agent,
17
+ load_builtin_agent,
18
+ load_fortify_agent,
19
+ load_local_agent,
20
+ load_registered_agent,
21
+ register_agent,
22
+ unregister_agent,
23
+ )
24
+ from fortify.cli.register import AgentManifest, create_manifest
25
+ from fortify.cloud import FortifyClient, FortifyConfig
26
+ from fortify.runtime import LocalWorkspace, ToolUseContext, User, Workspace
27
+ from fortify.security import AgentPolicy
28
+ from fortify.tools import (
29
+ agent_tool,
30
+ bash,
31
+ edit_file,
32
+ glob,
33
+ grep,
34
+ read_file,
35
+ refund_order,
36
+ write_file,
37
+ )
38
+ from fortify.tools.fetch import fetch
39
+ from fortify.tools.websearch import web_search
40
+
41
+ __all__ = [
42
+ "AgentManifest",
43
+ "AgentPolicy",
44
+ "FortifyClient",
45
+ "FortifyConfig",
46
+ "LocalWorkspace",
47
+ "ToolUseContext",
48
+ "User",
49
+ "Workspace",
50
+ "agent_tool",
51
+ "bash",
52
+ "edit_file",
53
+ "clear_registered_agents",
54
+ "create_agent",
55
+ "create_manifest",
56
+ "enforce_policy",
57
+ "fetch",
58
+ "glob",
59
+ "grep",
60
+ "invoke_agent",
61
+ "list_available_agents",
62
+ "list_builtin_agents",
63
+ "list_local_agents",
64
+ "list_registered_agents",
65
+ "load_agent",
66
+ "load_builtin_agent",
67
+ "load_fortify_agent",
68
+ "load_local_agent",
69
+ "load_registered_agent",
70
+ "register_agent",
71
+ "read_file",
72
+ "refund_order",
73
+ "stream_agent",
74
+ "stream_agent_raw",
75
+ "unregister_agent",
76
+ "web_search",
77
+ "write_file",
78
+ ]
File without changes
@@ -0,0 +1,3 @@
1
+ from fortify.adapters.google.runner import FortifyRunner
2
+
3
+ __all__ = ["FortifyRunner"]
@@ -0,0 +1,106 @@
1
+ """Google ADK ``Runner`` wrapper: opens a :class:`User` scope around each
2
+ ``Runner.run*`` call so the wrapped tools' enforcers can resolve the
3
+ active role. Langfuse propagation mirrors User identity into spans.
4
+ """
5
+
6
+ import asyncio
7
+ import os
8
+ from contextlib import contextmanager
9
+ from typing import Any, AsyncGenerator, Generator
10
+
11
+ import nest_asyncio
12
+ from google.adk.agents import BaseAgent
13
+ from google.adk.runners import Runner
14
+ from google.adk.sessions import BaseSessionService
15
+ from google.genai import types
16
+ from langfuse import get_client, propagate_attributes
17
+ from openinference.instrumentation.google_adk import GoogleADKInstrumentor
18
+
19
+ from fortify.adapters.google.wrapper import wrap_google_agent
20
+ from fortify.runtime import User
21
+
22
+
23
+ class FortifyRunner:
24
+ """Runner for Google ADK agents with Fortify tool policy and observability."""
25
+
26
+ def __init__(
27
+ self,
28
+ *,
29
+ agent: BaseAgent,
30
+ app_name: str,
31
+ session_service: BaseSessionService,
32
+ api_key: str | None = None,
33
+ **runner_kwargs: Any,
34
+ ):
35
+ self.api_key = api_key or os.getenv("FORTIFY_KEY")
36
+ if self.api_key is None:
37
+ raise ValueError(
38
+ "FORTIFY_KEY is not set. Pass api_key= explicitly or set FORTIFY_KEY environment variable."
39
+ )
40
+ # Policy is baked at construction; the Runner is built once and reused
41
+ # since role resolution happens at call time via the User contextvar.
42
+ self._wrapped_agent = wrap_google_agent(agent, api_key=self.api_key)
43
+ self._runner = Runner(
44
+ agent=self._wrapped_agent,
45
+ app_name=app_name,
46
+ session_service=session_service,
47
+ **runner_kwargs,
48
+ )
49
+ self._agent_name = getattr(agent, "name", "default")
50
+
51
+ def _setup_observability(self):
52
+ """Install Langfuse + GoogleADKInstrumentor (idempotent)."""
53
+ try:
54
+ asyncio.get_running_loop()
55
+ except RuntimeError:
56
+ # No running loop: safe to patch (and only useful for sync entry points).
57
+ # Patching a live loop breaks asyncio.current_task() on Python 3.12+.
58
+ nest_asyncio.apply()
59
+ get_client()
60
+ GoogleADKInstrumentor().instrument()
61
+
62
+ @contextmanager
63
+ def _propagate(self, user: User):
64
+ """Propagate User identity into Langfuse spans for the block."""
65
+ kwargs: dict[str, Any] = {"tags": [f"google.runner.run.{self._agent_name}"]}
66
+ kwargs["user_id"] = user.user_id
67
+ kwargs["session_id"] = user.session_id
68
+ kwargs["metadata"] = {"user_role": user.role}
69
+ with propagate_attributes(**kwargs):
70
+ yield
71
+
72
+ def run(
73
+ self,
74
+ *,
75
+ new_message: types.Content,
76
+ user: User,
77
+ **kwargs: Any,
78
+ ) -> Generator[Any, None, None]:
79
+ """Run the Google ADK agent synchronously, yielding events."""
80
+ self._setup_observability()
81
+ with user.sync_scope(), self._propagate(user):
82
+ yield from self._runner.run(
83
+ user_id=user.user_id,
84
+ session_id=user.session_id,
85
+ new_message=new_message,
86
+ **kwargs,
87
+ )
88
+
89
+ async def run_async(
90
+ self,
91
+ *,
92
+ new_message: types.Content | None = None,
93
+ user: User,
94
+ **kwargs: Any,
95
+ ) -> AsyncGenerator[Any, None]:
96
+ """Run the Google ADK agent asynchronously, yielding events."""
97
+ self._setup_observability()
98
+ async with user:
99
+ with self._propagate(user):
100
+ async for event in self._runner.run_async(
101
+ user_id=user.user_id,
102
+ session_id=user.session_id,
103
+ new_message=new_message,
104
+ **kwargs,
105
+ ):
106
+ yield event
@@ -0,0 +1,57 @@
1
+ """Google ADK adapter: wrap ``BaseTool`` so ``run_async`` consults a
2
+ :class:`PolicyEnforcer` first. Non-allow outcomes render as markered
3
+ strings the model sees as tool output.
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ import copy
9
+ import functools
10
+ from collections.abc import Callable
11
+ from typing import Any, Union
12
+
13
+ from google.adk.tools.base_tool import BaseTool
14
+ from google.adk.tools.function_tool import FunctionTool
15
+ from google.adk.tools.tool_context import ToolContext
16
+
17
+ from fortify.security.enforcer import PolicyEnforcer
18
+
19
+
20
+ ToolEntry = Union[BaseTool, Callable[..., Any]]
21
+
22
+
23
+ def _normalize(tool: ToolEntry) -> BaseTool:
24
+ """Coerce a tool entry into a ``BaseTool`` (plain callables → FunctionTool)."""
25
+ if isinstance(tool, BaseTool):
26
+ return tool
27
+ if callable(tool):
28
+ return FunctionTool(func=tool)
29
+ raise TypeError(
30
+ f"Cannot install policy on tool {tool!r}: expected google.adk BaseTool "
31
+ f"or callable, got {type(tool).__name__}."
32
+ )
33
+
34
+
35
+ def wrap_tool(tool: ToolEntry, enforcer: PolicyEnforcer) -> BaseTool:
36
+ """Return a copy of ``tool`` with ``run_async`` gated by ``enforcer``."""
37
+ base = _normalize(tool)
38
+ name = base.name
39
+ original_run_async = base.run_async
40
+
41
+ @functools.wraps(original_run_async, updated=())
42
+ async def guarded_run_async(
43
+ *, args: dict[str, Any], tool_context: ToolContext
44
+ ) -> Any:
45
+ decision = enforcer.decide(name, args or {})
46
+ if decision.allowed:
47
+ return await original_run_async(args=args, tool_context=tool_context)
48
+ return decision.as_error_message()
49
+
50
+ wrapped = copy.copy(base)
51
+ wrapped.run_async = guarded_run_async
52
+ return wrapped
53
+
54
+
55
+ def wrap_tools(tools: list[ToolEntry], enforcer: PolicyEnforcer) -> list[BaseTool]:
56
+ """Return a fresh list of policy-gated copies."""
57
+ return [wrap_tool(t, enforcer) for t in tools]
@@ -0,0 +1,47 @@
1
+ """Google ADK adapter: build a :class:`PolicySet`, construct one
2
+ :class:`PolicyEnforcer`, and return a clone of the agent whose tools
3
+ are policy-gated. User-agnostic at wrap time — role resolution happens
4
+ inside the enforcer via the :class:`User` contextvar.
5
+ """
6
+
7
+ from __future__ import annotations
8
+
9
+ from google.adk.agents import BaseAgent
10
+
11
+ from fortify import audit
12
+ from fortify.adapters.google.tools import wrap_tools
13
+ from fortify.security import AgentPolicy, BaseToolPolicy, PolicySet
14
+ from fortify.security.enforcer import PolicyEnforcer
15
+ from fortify.security.policy_set import DEFAULT_ROLE_NAME
16
+
17
+
18
+ def build_policy_set(
19
+ api_key: str, # noqa: ARG001 — reserved for the future Fortify-cloud fetch
20
+ agent_name: str, # noqa: ARG001 — same
21
+ tool_names: list[str],
22
+ ) -> PolicySet:
23
+ """Placeholder allow-all one-role bundle. TODO: cloud-fetch via FortifyClient."""
24
+ default_policy = AgentPolicy(
25
+ tools={name: BaseToolPolicy(mode="allow") for name in tool_names}
26
+ )
27
+ return PolicySet({DEFAULT_ROLE_NAME: default_policy})
28
+
29
+
30
+ def wrap_google_agent(agent: BaseAgent, *, api_key: str) -> BaseAgent:
31
+ """Return a clone of ``agent`` with policy-gated tools.
32
+
33
+ Caller must open a :class:`User` scope around the run.
34
+ ``NEEDS_APPROVAL`` outcomes surface as ``[approval_required]``-prefixed
35
+ strings in tool results; ``[policy_denied]`` for denials.
36
+ """
37
+ audit_sender = audit.configure(api_key)
38
+
39
+ agent_name = getattr(agent, "name", "default")
40
+ tools = list(getattr(agent, "tools", []) or [])
41
+ tool_names = [getattr(t, "name", getattr(t, "__name__", "tool")) for t in tools]
42
+ policy_set = build_policy_set(api_key, agent_name, tool_names)
43
+ enforcer = PolicyEnforcer(
44
+ policy_set, agent_name=agent_name, audit_sender=audit_sender
45
+ )
46
+ guarded_tools = wrap_tools(tools, enforcer)
47
+ return agent.model_copy(update={"tools": guarded_tools})
@@ -0,0 +1,7 @@
1
+ from fortify.adapters.langchain.agent import FortifyLangchainAgent
2
+ from fortify.adapters.langchain.wrapper import wrap_langchain_agent
3
+
4
+ __all__ = [
5
+ "FortifyLangchainAgent",
6
+ "wrap_langchain_agent",
7
+ ]
@@ -0,0 +1,133 @@
1
+ """Proxy around a pre-built ``CompiledStateGraph`` for Fortify-aware calls."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any, AsyncIterator, Iterator
6
+
7
+ from langchain_core.runnables import RunnableConfig
8
+ from langfuse import get_client, propagate_attributes
9
+ from langfuse.langchain import CallbackHandler
10
+ from langgraph.graph.state import CompiledStateGraph
11
+
12
+ from fortify.runtime import User
13
+
14
+
15
+ class FortifyLangchainAgent:
16
+ """Proxy around a ``CompiledStateGraph`` that opens a User scope per call.
17
+
18
+ Tools are already enforcer-installed at construction (by
19
+ :func:`wrap_langchain_agent`). This proxy pushes the active
20
+ :class:`User` onto the contextvar and propagates identity into
21
+ Langfuse spans. ``user`` is per-call, so one proxy serves many
22
+ users concurrently.
23
+ """
24
+
25
+ def __init__(
26
+ self,
27
+ *,
28
+ agent: CompiledStateGraph,
29
+ api_key: str,
30
+ tool_names: list[str],
31
+ ) -> None:
32
+ self._agent = agent
33
+ self._api_key = api_key
34
+ self._tool_names = tool_names
35
+ self._langfuse = get_client()
36
+ self._callback_handler = CallbackHandler()
37
+
38
+ def _propagate_kwargs(self, user: User, method: str) -> dict[str, Any]:
39
+ return {
40
+ "tags": [f"langchain.agent.{method}"],
41
+ "user_id": user.user_id,
42
+ "session_id": user.session_id,
43
+ "metadata": {"user_role": user.role},
44
+ }
45
+
46
+ def _with_callbacks(self, config: RunnableConfig | None) -> RunnableConfig:
47
+ """Append the Fortify callback handler to ``config['callbacks']``."""
48
+ merged: RunnableConfig = dict(config) if config else {}
49
+ callbacks = list(merged.get("callbacks") or [])
50
+ if self._callback_handler not in callbacks:
51
+ callbacks.append(self._callback_handler)
52
+ merged["callbacks"] = callbacks
53
+ return merged
54
+
55
+ async def ainvoke(
56
+ self,
57
+ input: dict[str, Any],
58
+ *,
59
+ user: User,
60
+ config: RunnableConfig | None = None,
61
+ **kwargs: Any,
62
+ ) -> dict[str, Any]:
63
+ """Invoke the agent asynchronously inside a User scope."""
64
+ async with user:
65
+ with propagate_attributes(**self._propagate_kwargs(user, "ainvoke")):
66
+ return await self._agent.ainvoke(
67
+ input, self._with_callbacks(config), **kwargs
68
+ )
69
+
70
+ def invoke(
71
+ self,
72
+ input: dict[str, Any],
73
+ *,
74
+ user: User,
75
+ config: RunnableConfig | None = None,
76
+ **kwargs: Any,
77
+ ) -> dict[str, Any]:
78
+ """Invoke the agent synchronously inside a User scope."""
79
+ with user.sync_scope():
80
+ with propagate_attributes(**self._propagate_kwargs(user, "invoke")):
81
+ return self._agent.invoke(input, self._with_callbacks(config), **kwargs)
82
+
83
+ async def astream(
84
+ self,
85
+ input: dict[str, Any],
86
+ *,
87
+ user: User,
88
+ config: RunnableConfig | None = None,
89
+ **kwargs: Any,
90
+ ) -> AsyncIterator[dict[str, Any]]:
91
+ """Stream the agent asynchronously inside a User scope."""
92
+ async with user:
93
+ with propagate_attributes(**self._propagate_kwargs(user, "astream")):
94
+ async for chunk in self._agent.astream(
95
+ input, self._with_callbacks(config), **kwargs
96
+ ):
97
+ yield chunk
98
+
99
+ def stream(
100
+ self,
101
+ input: dict[str, Any],
102
+ *,
103
+ user: User,
104
+ config: RunnableConfig | None = None,
105
+ **kwargs: Any,
106
+ ) -> Iterator[dict[str, Any]]:
107
+ """Stream the agent synchronously inside a User scope."""
108
+ with user.sync_scope():
109
+ with propagate_attributes(**self._propagate_kwargs(user, "stream")):
110
+ yield from self._agent.stream(
111
+ input, self._with_callbacks(config), **kwargs
112
+ )
113
+
114
+ async def astream_events(
115
+ self,
116
+ input: dict[str, Any],
117
+ version: str,
118
+ *,
119
+ user: User,
120
+ config: RunnableConfig | None = None,
121
+ **kwargs: Any,
122
+ ) -> AsyncIterator[dict[str, Any]]:
123
+ """Stream the agent events asynchronously inside a User scope."""
124
+ async with user:
125
+ with propagate_attributes(**self._propagate_kwargs(user, "astream_events")):
126
+ async for event in self._agent.astream_events(
127
+ input, version, config=self._with_callbacks(config), **kwargs
128
+ ):
129
+ yield event
130
+
131
+ def __getattr__(self, name: str) -> Any:
132
+ """Delegate unknown attributes to the wrapped agent."""
133
+ return getattr(self._agent, name)
@@ -0,0 +1,248 @@
1
+ """LangChain adapter for :class:`PolicyEnforcer`.
2
+
3
+ :class:`GuardedTool` wraps a ``BaseTool`` (used by
4
+ :meth:`FortifyAgent.enforce_policy`, which rebuilds the graph) and
5
+ carries an optional ``approval_handler`` for inline ``NEEDS_APPROVAL``
6
+ resolution.
7
+ :func:`install_enforcer_on_tool` mutates ``StructuredTool``'s ``func``/
8
+ ``coroutine`` in place (used by :func:`wrap_langchain_agent` for
9
+ pre-built ``CompiledStateGraph``s) and always renders non-allow as a
10
+ structured error — approval flows wire in on the host side.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ import functools
16
+ from collections.abc import Awaitable, Callable
17
+ from inspect import isawaitable
18
+ from typing import Any
19
+
20
+ from langchain_core.tools import BaseTool
21
+ from langchain_core.tools.structured import StructuredTool
22
+ from pydantic import ConfigDict
23
+
24
+ from fortify.agents.factory import ApprovalHandler
25
+ from fortify.security.decision import Decision, DecisionOutcome
26
+ from fortify.security.enforcer import PolicyEnforcer
27
+ from fortify.tools.decorators import TOOL_METADATA_ATTR
28
+
29
+
30
+ def _copy_tool_metadata(source: Any, target: Any) -> Any:
31
+ """Copy fortify tool metadata (tracing labels, etc.) onto a wrapper."""
32
+ metadata = getattr(source, TOOL_METADATA_ATTR, None)
33
+ if metadata is not None:
34
+ setattr(target, TOOL_METADATA_ATTR, metadata)
35
+ return target
36
+
37
+
38
+ def _resolve_approval_sync(handler: ApprovalHandler, decision: Decision) -> bool:
39
+ """Resolve a NEEDS_APPROVAL decision in a sync caller (rejects coroutines)."""
40
+ if isinstance(handler, bool):
41
+ return handler
42
+ result = handler(decision)
43
+ if isawaitable(result):
44
+ raise RuntimeError(
45
+ "approval_handler returned a coroutine; sync tool invocation cannot "
46
+ "await it — use ainvoke/astream/astream_events"
47
+ )
48
+ return bool(result)
49
+
50
+
51
+ async def _resolve_approval_async(handler: ApprovalHandler, decision: Decision) -> bool:
52
+ """Resolve a NEEDS_APPROVAL decision in an async caller."""
53
+ if isinstance(handler, bool):
54
+ return handler
55
+ result = handler(decision)
56
+ if isawaitable(result):
57
+ result = await result
58
+ return bool(result)
59
+
60
+
61
+ class GuardedTool(BaseTool):
62
+ """LangChain tool wrapper that consults a :class:`PolicyEnforcer`.
63
+
64
+ ALLOW delegates to the wrapped tool; non-ALLOW renders
65
+ ``Decision.as_error_payload()`` so the LLM sees governance failures
66
+ as tool output. NEEDS_APPROVAL is treated as denial unless
67
+ ``approval_handler`` (callable taking the :class:`Decision`, or a
68
+ ``bool`` shorthand) returns truthy.
69
+ """
70
+
71
+ model_config = ConfigDict(arbitrary_types_allowed=True)
72
+
73
+ wrapped_tool: BaseTool
74
+ enforcer: PolicyEnforcer | None = None
75
+ approval_handler: ApprovalHandler | None = None
76
+
77
+ @classmethod
78
+ def wrap(
79
+ cls,
80
+ tool: BaseTool,
81
+ *,
82
+ enforcer: PolicyEnforcer | None = None,
83
+ approval_handler: ApprovalHandler | None = None,
84
+ ) -> "GuardedTool":
85
+ """Return a GuardedTool delegating to ``tool`` after policy check.
86
+
87
+ Idempotent re-wrap: an existing ``GuardedTool`` is unwrapped once
88
+ so enforcers don't stack; fields fall through unless explicitly
89
+ overridden.
90
+ """
91
+ if isinstance(tool, cls):
92
+ inner = tool.wrapped_tool
93
+ resolved_enforcer = enforcer if enforcer is not None else tool.enforcer
94
+ resolved_approval = (
95
+ approval_handler
96
+ if approval_handler is not None
97
+ else tool.approval_handler
98
+ )
99
+ else:
100
+ inner = tool
101
+ resolved_enforcer = enforcer
102
+ resolved_approval = approval_handler
103
+
104
+ guarded = cls(
105
+ name=inner.name,
106
+ description=inner.description,
107
+ args_schema=inner.args_schema,
108
+ return_direct=inner.return_direct,
109
+ verbose=inner.verbose,
110
+ callbacks=inner.callbacks,
111
+ tags=inner.tags,
112
+ metadata=inner.metadata,
113
+ handle_tool_error=inner.handle_tool_error,
114
+ handle_validation_error=inner.handle_validation_error,
115
+ response_format=inner.response_format,
116
+ extras=inner.extras,
117
+ wrapped_tool=inner,
118
+ enforcer=resolved_enforcer,
119
+ approval_handler=resolved_approval,
120
+ )
121
+ return _copy_tool_metadata(inner, guarded)
122
+
123
+ async def _arun(self, *args: Any, **kwargs: Any) -> Any:
124
+ if self.enforcer is not None:
125
+ decision = self.enforcer.decide(self.name, kwargs)
126
+ if not decision.allowed:
127
+ if (
128
+ decision.outcome is DecisionOutcome.NEEDS_APPROVAL
129
+ and self.approval_handler is not None
130
+ and await _resolve_approval_async(self.approval_handler, decision)
131
+ ):
132
+ pass # approved → fall through and invoke
133
+ else:
134
+ return {"ok": False, "error": decision.as_error_payload()}
135
+ return await self._invoke_wrapped_async(*args, **kwargs)
136
+
137
+ def _run(self, *args: Any, **kwargs: Any) -> Any:
138
+ if self.enforcer is not None:
139
+ decision = self.enforcer.decide(self.name, kwargs)
140
+ if not decision.allowed:
141
+ if (
142
+ decision.outcome is DecisionOutcome.NEEDS_APPROVAL
143
+ and self.approval_handler is not None
144
+ and _resolve_approval_sync(self.approval_handler, decision)
145
+ ):
146
+ pass
147
+ else:
148
+ return {"ok": False, "error": decision.as_error_payload()}
149
+ return self._invoke_wrapped_sync(*args, **kwargs)
150
+
151
+ async def _invoke_wrapped_async(self, *args: Any, **kwargs: Any) -> Any:
152
+ """Call the wrapped tool without re-entering LangChain instrumentation."""
153
+ if isinstance(self.wrapped_tool, StructuredTool):
154
+ if self.wrapped_tool.coroutine is not None:
155
+ return await self.wrapped_tool.coroutine(*args, **kwargs)
156
+ if self.wrapped_tool.func is not None:
157
+ return self.wrapped_tool.func(*args, **kwargs)
158
+ return await self.wrapped_tool._arun(*args, **kwargs)
159
+
160
+ def _invoke_wrapped_sync(self, *args: Any, **kwargs: Any) -> Any:
161
+ if (
162
+ isinstance(self.wrapped_tool, StructuredTool)
163
+ and self.wrapped_tool.func is not None
164
+ ):
165
+ return self.wrapped_tool.func(*args, **kwargs)
166
+ return self.wrapped_tool._run(*args, **kwargs)
167
+
168
+
169
+ # ---------------------------------------------------------------------------
170
+ # In-place installer for retrofitting existing CompiledStateGraph tools.
171
+ # ---------------------------------------------------------------------------
172
+
173
+ _ORIGINAL_FUNC_ATTR = "_fortify_original_func"
174
+ _ORIGINAL_COROUTINE_ATTR = "_fortify_original_coroutine"
175
+ _INSTALLED_ATTR = "_fortify_enforcer_installed"
176
+
177
+
178
+ def install_enforcer_on_tool(
179
+ tool: BaseTool,
180
+ *,
181
+ enforcer: PolicyEnforcer,
182
+ ) -> BaseTool:
183
+ """Install :class:`PolicyEnforcer` gating on ``tool`` in place.
184
+
185
+ Same semantics as :class:`GuardedTool` but mutates ``StructuredTool``'s
186
+ ``func``/``coroutine`` instead of constructing a wrapper — use when
187
+ the tool is already bound to a ``CompiledStateGraph``. Idempotent:
188
+ re-install restores captured originals first so gates don't stack.
189
+ Non-allow outcomes render as the structured error dict; approval
190
+ flows belong on the host side, not on this in-place installer.
191
+ """
192
+ name = tool.name
193
+ original_func: Callable[..., Any] | None = getattr(tool, _ORIGINAL_FUNC_ATTR, None)
194
+ if original_func is None:
195
+ original_func = getattr(tool, "func", None)
196
+ original_coroutine: Callable[..., Awaitable[Any]] | None = getattr(
197
+ tool, _ORIGINAL_COROUTINE_ATTR, None
198
+ )
199
+ if original_coroutine is None:
200
+ original_coroutine = getattr(tool, "coroutine", None)
201
+
202
+ if original_func is None and original_coroutine is None:
203
+ raise TypeError(
204
+ f"Cannot install policy on tool {name!r}: it is a "
205
+ f"{type(tool).__name__} without `func`/`coroutine` attributes. "
206
+ "In-place wrapping only supports StructuredTool-style tools."
207
+ )
208
+
209
+ if original_func is not None:
210
+ captured_func = original_func
211
+
212
+ @functools.wraps(captured_func)
213
+ def guarded_func(*args: Any, **kwargs: Any) -> Any:
214
+ decision = enforcer.decide(name, kwargs)
215
+ if decision.allowed:
216
+ return captured_func(*args, **kwargs)
217
+ return {"ok": False, "error": decision.as_error_payload()}
218
+
219
+ setattr(tool, _ORIGINAL_FUNC_ATTR, captured_func)
220
+ tool.func = guarded_func
221
+
222
+ if original_coroutine is not None:
223
+ captured_coroutine = original_coroutine
224
+
225
+ @functools.wraps(captured_coroutine)
226
+ async def guarded_coroutine(*args: Any, **kwargs: Any) -> Any:
227
+ decision = enforcer.decide(name, kwargs)
228
+ if decision.allowed:
229
+ return await captured_coroutine(*args, **kwargs)
230
+ return {"ok": False, "error": decision.as_error_payload()}
231
+
232
+ setattr(tool, _ORIGINAL_COROUTINE_ATTR, captured_coroutine)
233
+ tool.coroutine = guarded_coroutine
234
+
235
+ tool.handle_tool_error = True
236
+ setattr(tool, _INSTALLED_ATTR, True)
237
+ return tool
238
+
239
+
240
+ def install_enforcer_on_tools(
241
+ tools: list[BaseTool],
242
+ *,
243
+ enforcer: PolicyEnforcer,
244
+ ) -> list[BaseTool]:
245
+ """Install enforcement on every StructuredTool-style tool in place."""
246
+ for t in tools:
247
+ install_enforcer_on_tool(t, enforcer=enforcer)
248
+ return tools