agentapprovalprotocol 0.1.0__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.
- agentapprovalprotocol/__init__.py +33 -0
- agentapprovalprotocol/_client.py +201 -0
- agentapprovalprotocol/_gate.py +310 -0
- agentapprovalprotocol/_wire.py +229 -0
- agentapprovalprotocol/langgraph.py +255 -0
- agentapprovalprotocol/py.typed +0 -0
- agentapprovalprotocol/pydantic_ai.py +255 -0
- agentapprovalprotocol/receiver.py +168 -0
- agentapprovalprotocol/testing.py +173 -0
- agentapprovalprotocol-0.1.0.dist-info/METADATA +96 -0
- agentapprovalprotocol-0.1.0.dist-info/RECORD +12 -0
- agentapprovalprotocol-0.1.0.dist-info/WHEEL +4 -0
|
@@ -0,0 +1,255 @@
|
|
|
1
|
+
"""Approvals for Pydantic AI agents.
|
|
2
|
+
|
|
3
|
+
``ApprovalCapability`` gates every function tool call an agent makes,
|
|
4
|
+
including MCP and other toolset tools, with the arguments the tool will
|
|
5
|
+
run with. In sync mode the call waits for the decision. In async mode the
|
|
6
|
+
run ends with ``DeferredToolRequests``, you store its ``conversation``, and
|
|
7
|
+
a ``Receiver`` continues it when the decision arrives. On resume the
|
|
8
|
+
capability asks the provider again and ignores the approval passed in, so a
|
|
9
|
+
conversation continued by anything else cannot approve a call.
|
|
10
|
+
|
|
11
|
+
Install with ``pip install "agentapprovalprotocol[pydantic-ai]"``.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import asyncio
|
|
17
|
+
import contextvars
|
|
18
|
+
import inspect
|
|
19
|
+
import logging
|
|
20
|
+
import types
|
|
21
|
+
import typing
|
|
22
|
+
from dataclasses import dataclass, field
|
|
23
|
+
from typing import Any, Awaitable, Callable, Literal, Mapping, Sequence
|
|
24
|
+
|
|
25
|
+
from pydantic_core import to_jsonable_python
|
|
26
|
+
from pydantic_ai import DeferredToolRequests, DeferredToolResults
|
|
27
|
+
from pydantic_ai.capabilities import AbstractCapability, CapabilityOrdering
|
|
28
|
+
from pydantic_ai.conversation import Conversation
|
|
29
|
+
from pydantic_ai.exceptions import ApprovalRequired, ToolFailed, UserError
|
|
30
|
+
from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart
|
|
31
|
+
from pydantic_ai.models import ModelRequestContext
|
|
32
|
+
from pydantic_ai.tools import RunContext, ToolDefinition
|
|
33
|
+
from pydantic_ai.toolsets import AbstractToolset, WrapperToolset
|
|
34
|
+
from pydantic_ai.toolsets.abstract import ToolsetTool
|
|
35
|
+
|
|
36
|
+
from ._client import AsyncClient
|
|
37
|
+
from ._gate import FAIL_CLOSED, AsyncGate, Call, Deny, Pending
|
|
38
|
+
from ._wire import Request
|
|
39
|
+
|
|
40
|
+
log = logging.getLogger("agentapprovalprotocol")
|
|
41
|
+
|
|
42
|
+
FRAMEWORK = "pydantic-ai"
|
|
43
|
+
METADATA_KEY = "aap_request_id"
|
|
44
|
+
EDITED = "An approved call cannot run with edited arguments. The call was not run."
|
|
45
|
+
# Native tool search only finds tools among the agent's own, which are gated.
|
|
46
|
+
SEARCH_ONLY = frozenset({"tool_search"})
|
|
47
|
+
|
|
48
|
+
# Set while a gated tool runs. An agent run started inside that tool (a
|
|
49
|
+
# sub-agent) cannot pause: its caller would get the pause as a tool result.
|
|
50
|
+
_in_tool: contextvars.ContextVar[bool] = contextvars.ContextVar("aap_pydantic_ai_in_tool", default=False)
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def _allows_deferred(spec: Any) -> bool:
|
|
54
|
+
"""Whether an agent's output type includes ``DeferredToolRequests``."""
|
|
55
|
+
if spec is DeferredToolRequests:
|
|
56
|
+
return True
|
|
57
|
+
if isinstance(spec, (list, tuple)):
|
|
58
|
+
return any(_allows_deferred(item) for item in spec)
|
|
59
|
+
if typing.get_origin(spec) in (typing.Union, types.UnionType):
|
|
60
|
+
return any(_allows_deferred(item) for item in typing.get_args(spec))
|
|
61
|
+
return False
|
|
62
|
+
|
|
63
|
+
|
|
64
|
+
def _response_with(messages: Sequence[ModelMessage], tool_call_id: str) -> tuple[ModelResponse, ToolCallPart] | None:
|
|
65
|
+
for message in reversed(messages):
|
|
66
|
+
if isinstance(message, ModelResponse):
|
|
67
|
+
for part in message.parts:
|
|
68
|
+
if isinstance(part, ToolCallPart) and part.tool_call_id == tool_call_id:
|
|
69
|
+
return message, part
|
|
70
|
+
return None
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def _reasoning(response: ModelResponse | None) -> str | None:
|
|
74
|
+
if response is None:
|
|
75
|
+
return None
|
|
76
|
+
text = "\n".join(part.content for part in response.parts if isinstance(part, TextPart)).strip()
|
|
77
|
+
return text or None
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
@dataclass(init=False)
|
|
81
|
+
class ApprovalCapability(AbstractCapability[Any]):
|
|
82
|
+
"""Send every tool call to an AAP provider before it runs.
|
|
83
|
+
|
|
84
|
+
Add it to the agent's ``capabilities`` and to every agent that runs as
|
|
85
|
+
a sub-agent. Provider-run (native) tools cannot be gated, so a model
|
|
86
|
+
request that includes one is refused.
|
|
87
|
+
|
|
88
|
+
``mode="async"`` pauses only when the agent's ``output_type`` includes
|
|
89
|
+
``DeferredToolRequests``; otherwise, and inside a sub-agent, calls wait.
|
|
90
|
+
"""
|
|
91
|
+
|
|
92
|
+
gate: AsyncGate = field(repr=False)
|
|
93
|
+
context: dict[str, Any] = field(default_factory=dict)
|
|
94
|
+
|
|
95
|
+
def __init__(
|
|
96
|
+
self,
|
|
97
|
+
base_url: str | None = None,
|
|
98
|
+
token: str | None = None,
|
|
99
|
+
*,
|
|
100
|
+
mode: Literal["sync", "async"] = "sync",
|
|
101
|
+
gate: AsyncGate | None = None,
|
|
102
|
+
context: Mapping[str, Any] | None = None,
|
|
103
|
+
**gate_options: Any,
|
|
104
|
+
):
|
|
105
|
+
if gate is None:
|
|
106
|
+
if base_url is None or token is None:
|
|
107
|
+
raise ValueError("pass base_url and token, or a gate")
|
|
108
|
+
gate = AsyncGate(AsyncClient(base_url, token), mode=mode, **gate_options)
|
|
109
|
+
self.gate = gate
|
|
110
|
+
self.context = {"framework": FRAMEWORK, **(context or {})}
|
|
111
|
+
|
|
112
|
+
@classmethod
|
|
113
|
+
def get_serialization_name(cls) -> str | None:
|
|
114
|
+
return None
|
|
115
|
+
|
|
116
|
+
def get_ordering(self) -> CapabilityOrdering:
|
|
117
|
+
# Outermost, so tools other capabilities add around the agent's own
|
|
118
|
+
# (a search or code tool, say) are gated too.
|
|
119
|
+
return CapabilityOrdering(position="outermost")
|
|
120
|
+
|
|
121
|
+
def get_wrapper_toolset(self, toolset: AbstractToolset[Any]) -> AbstractToolset[Any]:
|
|
122
|
+
return _GatedToolset(toolset, approvals=self)
|
|
123
|
+
|
|
124
|
+
async def before_model_request(self, ctx: RunContext[Any], request_context: ModelRequestContext) -> ModelRequestContext:
|
|
125
|
+
kinds = {getattr(tool, "kind", type(tool).__name__) for tool in request_context.model_request_parameters.native_tools} - SEARCH_ONLY
|
|
126
|
+
if kinds:
|
|
127
|
+
raise UserError(f"Native tools run at the model provider and cannot be sent for approval: {', '.join(sorted(kinds))}. Remove them, or use a local tool instead.")
|
|
128
|
+
return request_context
|
|
129
|
+
|
|
130
|
+
async def after_tool_validate(self, ctx: RunContext[Any], *, call: ToolCallPart, tool_def: ToolDefinition, args: dict[str, Any]) -> dict[str, Any]:
|
|
131
|
+
# ToolApproved(override_args=...) replaces the call's arguments. An
|
|
132
|
+
# approval covers the arguments it was given, so refuse the edit.
|
|
133
|
+
if ctx.tool_call_approved:
|
|
134
|
+
found = _response_with(ctx.messages, call.tool_call_id)
|
|
135
|
+
if found is not None and found[1].args_as_dict() != call.args_as_dict():
|
|
136
|
+
raise ToolFailed(EDITED)
|
|
137
|
+
return args
|
|
138
|
+
|
|
139
|
+
def _call(self, name: str, args: dict[str, Any], ctx: RunContext[Any]) -> tuple[Call, bool]:
|
|
140
|
+
conversation = ctx.conversation_id
|
|
141
|
+
tool_call_id = ctx.tool_call_id or "-"
|
|
142
|
+
found = _response_with(ctx.messages, tool_call_id)
|
|
143
|
+
resume = None
|
|
144
|
+
if conversation is not None:
|
|
145
|
+
resume = {"framework": FRAMEWORK, "conversation_id": conversation}
|
|
146
|
+
call = Call(
|
|
147
|
+
key=f"{conversation or '-'}:{tool_call_id}",
|
|
148
|
+
tool=name,
|
|
149
|
+
arguments=to_jsonable_python(args),
|
|
150
|
+
reasoning=_reasoning(found[0] if found else None),
|
|
151
|
+
context=self.context,
|
|
152
|
+
resume=resume,
|
|
153
|
+
)
|
|
154
|
+
can_pause = (
|
|
155
|
+
conversation is not None
|
|
156
|
+
and ctx.tool_call_id is not None
|
|
157
|
+
and not _in_tool.get()
|
|
158
|
+
and ctx.agent is not None
|
|
159
|
+
and _allows_deferred(ctx.agent.output_type)
|
|
160
|
+
)
|
|
161
|
+
return call, can_pause
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
@dataclass
|
|
165
|
+
class _GatedToolset(WrapperToolset[Any]):
|
|
166
|
+
"""Checks each call with the gate, at the point the tool runs."""
|
|
167
|
+
|
|
168
|
+
approvals: ApprovalCapability = field(kw_only=True)
|
|
169
|
+
|
|
170
|
+
async def call_tool(self, name: str, tool_args: dict[str, Any], ctx: RunContext[Any], tool: ToolsetTool[Any]) -> Any:
|
|
171
|
+
gate = self.approvals.gate
|
|
172
|
+
try:
|
|
173
|
+
call, can_pause = self.approvals._call(name, tool_args, ctx)
|
|
174
|
+
except Exception: # noqa: BLE001 - arguments that cannot be sent are not run
|
|
175
|
+
log.exception("could not build the approval request for %s", name)
|
|
176
|
+
raise ToolFailed(FAIL_CLOSED) from None
|
|
177
|
+
verdict = await gate.check(call, can_pause=can_pause)
|
|
178
|
+
if isinstance(verdict, Pending):
|
|
179
|
+
# The run ends with DeferredToolRequests. Any approval passed back
|
|
180
|
+
# on resume is ignored: this check runs again and asks the provider.
|
|
181
|
+
raise ApprovalRequired(metadata={METADATA_KEY: verdict.request_id})
|
|
182
|
+
if isinstance(verdict, Deny):
|
|
183
|
+
raise ToolFailed(verdict.reason)
|
|
184
|
+
token = _in_tool.set(True)
|
|
185
|
+
try:
|
|
186
|
+
result = await self.wrapped.call_tool(name, tool_args, ctx, tool)
|
|
187
|
+
except BaseException:
|
|
188
|
+
if gate.report_outcomes:
|
|
189
|
+
await verdict.report("failed")
|
|
190
|
+
raise
|
|
191
|
+
finally:
|
|
192
|
+
_in_tool.reset(token)
|
|
193
|
+
if gate.report_outcomes:
|
|
194
|
+
await verdict.report("succeeded")
|
|
195
|
+
return result
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
def pending_request_ids(requests: DeferredToolRequests | None) -> dict[str, str]:
|
|
199
|
+
"""The AAP request each paused call waits on, keyed by tool call ID."""
|
|
200
|
+
if requests is None:
|
|
201
|
+
return {}
|
|
202
|
+
out = {}
|
|
203
|
+
for call in requests.approvals:
|
|
204
|
+
request_id = (requests.metadata.get(call.tool_call_id) or {}).get(METADATA_KEY)
|
|
205
|
+
if isinstance(request_id, str):
|
|
206
|
+
out[call.tool_call_id] = request_id
|
|
207
|
+
return out
|
|
208
|
+
|
|
209
|
+
|
|
210
|
+
async def _maybe(value: Any) -> Any:
|
|
211
|
+
return await value if inspect.isawaitable(value) else value
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def resumer(
|
|
215
|
+
agent: Any,
|
|
216
|
+
*,
|
|
217
|
+
load: Callable[[str], Conversation | None | Awaitable[Conversation | None]],
|
|
218
|
+
save: Callable[[Conversation], Any],
|
|
219
|
+
**run_options: Any,
|
|
220
|
+
) -> Callable[[Mapping[str, Any], Request], Awaitable[None]]:
|
|
221
|
+
"""A ``Receiver`` resumer that continues a stored conversation.
|
|
222
|
+
|
|
223
|
+
``load(conversation_id)`` returns the ``Conversation`` you saved when the
|
|
224
|
+
run paused, or ``None``. The resumed run's conversation is passed to
|
|
225
|
+
``save``. Either may be async. ``run_options`` (``deps``, for example)
|
|
226
|
+
are passed to ``agent.run``.
|
|
227
|
+
|
|
228
|
+
A conversation is continued only while it waits on the decided request,
|
|
229
|
+
one resume per conversation at a time in this process. Every paused
|
|
230
|
+
approval in it is answered, and each is checked with the provider again:
|
|
231
|
+
calls still undecided pause again.
|
|
232
|
+
"""
|
|
233
|
+
locks: dict[str, asyncio.Lock] = {}
|
|
234
|
+
|
|
235
|
+
async def resume(handle: Mapping[str, Any], request: Request) -> None:
|
|
236
|
+
conversation_id = handle.get("conversation_id")
|
|
237
|
+
if not isinstance(conversation_id, str):
|
|
238
|
+
return
|
|
239
|
+
async with locks.setdefault(conversation_id, asyncio.Lock()):
|
|
240
|
+
conversation = await _maybe(load(conversation_id))
|
|
241
|
+
if conversation is None:
|
|
242
|
+
return
|
|
243
|
+
waiting = conversation.deferred_tool_requests
|
|
244
|
+
ids = pending_request_ids(waiting)
|
|
245
|
+
if request.id not in ids.values():
|
|
246
|
+
return
|
|
247
|
+
assert waiting is not None
|
|
248
|
+
if waiting.calls or len(ids) != len(waiting.approvals):
|
|
249
|
+
log.warning("conversation %s also waits on calls that are not AAP approvals; resume it yourself", conversation_id)
|
|
250
|
+
return
|
|
251
|
+
results = DeferredToolResults(approvals={tool_call_id: True for tool_call_id in ids})
|
|
252
|
+
result = await agent.run(conversation=conversation, deferred_tool_results=results, **run_options)
|
|
253
|
+
await _maybe(save(result.conversation))
|
|
254
|
+
|
|
255
|
+
return resume
|
|
@@ -0,0 +1,168 @@
|
|
|
1
|
+
"""Webhook receiver for asynchronous mode.
|
|
2
|
+
|
|
3
|
+
The receiver verifies AAP webhooks, answers receiver verification during
|
|
4
|
+
provisioning, and wakes paused runs. A notification only wakes a run: the
|
|
5
|
+
gate fetches and consumes the decision itself, so the receiver keeps no
|
|
6
|
+
state and a repeated or forged event at most wakes a run that pauses again.
|
|
7
|
+
|
|
8
|
+
The resume handle is read from the request's ``context.resume``, which the
|
|
9
|
+
gate wrote when it created the request. Its ``framework`` names the resumer
|
|
10
|
+
to call.
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
from __future__ import annotations
|
|
14
|
+
|
|
15
|
+
import asyncio
|
|
16
|
+
import base64
|
|
17
|
+
import hashlib
|
|
18
|
+
import hmac
|
|
19
|
+
import inspect
|
|
20
|
+
import json
|
|
21
|
+
import logging
|
|
22
|
+
import time
|
|
23
|
+
from typing import Any, Awaitable, Callable, Mapping
|
|
24
|
+
|
|
25
|
+
from ._client import AAPError, AsyncClient
|
|
26
|
+
from ._wire import ProtocolError, Request, loads
|
|
27
|
+
|
|
28
|
+
log = logging.getLogger("agentapprovalprotocol")
|
|
29
|
+
|
|
30
|
+
Resumer = Callable[[Mapping[str, Any], Request], Awaitable[None] | None]
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def signing_key(secret: str) -> bytes:
|
|
34
|
+
if not secret.startswith("whsec_"):
|
|
35
|
+
raise ValueError("signing secret must start with whsec_")
|
|
36
|
+
return base64.b64decode(secret[len("whsec_") :], validate=True)
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def sign(secret: str, message_id: str, timestamp: int, body: bytes) -> str:
|
|
40
|
+
"""The ``webhook-signature`` header value for one delivery attempt."""
|
|
41
|
+
digest = hmac.new(signing_key(secret), f"{message_id}.{timestamp}.".encode() + body, hashlib.sha256).digest()
|
|
42
|
+
return "v1," + base64.b64encode(digest).decode()
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def verify(secret: str, headers: Mapping[str, str], body: bytes, *, tolerance: int = 300, now: float | None = None) -> dict[str, Any]:
|
|
46
|
+
"""Verify a delivery and return its parsed body. Raises ``ValueError``."""
|
|
47
|
+
lower = {k.lower(): v for k, v in headers.items()}
|
|
48
|
+
message_id = lower.get("webhook-id", "")
|
|
49
|
+
stamp = lower.get("webhook-timestamp", "")
|
|
50
|
+
signatures = lower.get("webhook-signature", "").split()
|
|
51
|
+
if not message_id or not stamp.isdigit() or not signatures:
|
|
52
|
+
raise ValueError("missing webhook headers")
|
|
53
|
+
if abs((now if now is not None else time.time()) - int(stamp)) > tolerance:
|
|
54
|
+
raise ValueError("webhook timestamp outside tolerance")
|
|
55
|
+
expected = sign(secret, message_id, int(stamp), body)
|
|
56
|
+
if not any(hmac.compare_digest(expected, candidate) for candidate in signatures if candidate.startswith("v1,")):
|
|
57
|
+
raise ValueError("invalid webhook signature")
|
|
58
|
+
try:
|
|
59
|
+
message = loads(body)
|
|
60
|
+
except ProtocolError as err:
|
|
61
|
+
raise ValueError("invalid webhook body") from err
|
|
62
|
+
if not isinstance(message, dict) or message.get("id") != message_id:
|
|
63
|
+
raise ValueError("webhook ID does not match the body")
|
|
64
|
+
return message
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
class Receiver:
|
|
68
|
+
"""Verifies AAP webhooks and wakes paused runs.
|
|
69
|
+
|
|
70
|
+
``registration_id`` is the ID the provisioner generated for this
|
|
71
|
+
receiver; verification messages for any other registration are refused.
|
|
72
|
+
``instance_id``, once known, restricts notifications to that instance.
|
|
73
|
+
``resumers`` maps a framework name to a function that resumes a run from
|
|
74
|
+
its resume handle.
|
|
75
|
+
"""
|
|
76
|
+
|
|
77
|
+
def __init__(
|
|
78
|
+
self,
|
|
79
|
+
*,
|
|
80
|
+
secret: str,
|
|
81
|
+
client: AsyncClient,
|
|
82
|
+
resumers: Mapping[str, Resumer],
|
|
83
|
+
registration_id: str | None = None,
|
|
84
|
+
instance_id: str | None = None,
|
|
85
|
+
tolerance: int = 300,
|
|
86
|
+
):
|
|
87
|
+
signing_key(secret)
|
|
88
|
+
self.secret = secret
|
|
89
|
+
self.client = client
|
|
90
|
+
self.resumers = dict(resumers)
|
|
91
|
+
self.registration_id = registration_id
|
|
92
|
+
self.instance_id = instance_id
|
|
93
|
+
self.tolerance = tolerance
|
|
94
|
+
self._tasks: set[asyncio.Task[None]] = set()
|
|
95
|
+
|
|
96
|
+
async def handle(self, headers: Mapping[str, str], body: bytes) -> tuple[int, bytes]:
|
|
97
|
+
"""Process one delivery. Returns the HTTP status and body to send.
|
|
98
|
+
|
|
99
|
+
A notification is acknowledged at once and the run is woken in the
|
|
100
|
+
background, because a delivery attempt has a ten second limit.
|
|
101
|
+
"""
|
|
102
|
+
try:
|
|
103
|
+
message = verify(self.secret, headers, body, tolerance=self.tolerance)
|
|
104
|
+
except ValueError:
|
|
105
|
+
return 401, b""
|
|
106
|
+
kind = message.get("type")
|
|
107
|
+
if kind == "webhook.verification":
|
|
108
|
+
if self.registration_id is None or message.get("registration_id") != self.registration_id or not isinstance(message.get("challenge"), str):
|
|
109
|
+
return 403, b""
|
|
110
|
+
return 200, json.dumps({"challenge": message["challenge"]}).encode()
|
|
111
|
+
if kind == "request.resolved":
|
|
112
|
+
if self.instance_id is not None and message.get("instance_id") != self.instance_id:
|
|
113
|
+
return 403, b""
|
|
114
|
+
request_id = message.get("request_id")
|
|
115
|
+
if isinstance(request_id, str):
|
|
116
|
+
task = asyncio.ensure_future(self.wake(request_id))
|
|
117
|
+
self._tasks.add(task)
|
|
118
|
+
task.add_done_callback(self._tasks.discard)
|
|
119
|
+
return 204, b""
|
|
120
|
+
return 400, b""
|
|
121
|
+
|
|
122
|
+
async def wake(self, request_id: str) -> bool:
|
|
123
|
+
"""Resume the run waiting on a request. Also the recovery sweep's entry point."""
|
|
124
|
+
try:
|
|
125
|
+
request = await self.client.get(request_id)
|
|
126
|
+
except (AAPError, ProtocolError):
|
|
127
|
+
log.warning("could not read request %s to resume it", request_id, exc_info=True)
|
|
128
|
+
return False
|
|
129
|
+
handle = (request.context or {}).get("resume")
|
|
130
|
+
resumer = self.resumers.get(handle.get("framework", "")) if isinstance(handle, dict) else None
|
|
131
|
+
if resumer is None:
|
|
132
|
+
log.warning("request %s has no resume handle this receiver can use", request_id)
|
|
133
|
+
return False
|
|
134
|
+
try:
|
|
135
|
+
result = resumer(handle, request)
|
|
136
|
+
if inspect.isawaitable(result):
|
|
137
|
+
await result
|
|
138
|
+
return True
|
|
139
|
+
except Exception: # noqa: BLE001 - a failed resume must not crash the receiver
|
|
140
|
+
log.exception("resuming the run for request %s failed", request_id)
|
|
141
|
+
return False
|
|
142
|
+
|
|
143
|
+
def asgi(self) -> Callable[..., Awaitable[None]]:
|
|
144
|
+
"""An ASGI app that accepts POSTed webhooks at any path it is mounted on."""
|
|
145
|
+
|
|
146
|
+
async def app(scope: dict[str, Any], receive: Callable[[], Awaitable[dict[str, Any]]], send: Callable[[dict[str, Any]], Awaitable[None]]) -> None:
|
|
147
|
+
if scope["type"] != "http":
|
|
148
|
+
return
|
|
149
|
+
if scope["method"] != "POST":
|
|
150
|
+
await _respond(send, 405, b"")
|
|
151
|
+
return
|
|
152
|
+
chunks = []
|
|
153
|
+
more = True
|
|
154
|
+
while more:
|
|
155
|
+
event = await receive()
|
|
156
|
+
chunks.append(event.get("body", b""))
|
|
157
|
+
more = event.get("more_body", False)
|
|
158
|
+
headers = {k.decode("latin-1"): v.decode("latin-1") for k, v in scope.get("headers", [])}
|
|
159
|
+
status, body = await self.handle(headers, b"".join(chunks))
|
|
160
|
+
await _respond(send, status, body)
|
|
161
|
+
|
|
162
|
+
return app
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
async def _respond(send: Callable[[dict[str, Any]], Awaitable[None]], status: int, body: bytes) -> None:
|
|
166
|
+
headers = [(b"content-type", b"application/json")] if body else []
|
|
167
|
+
await send({"type": "http.response.start", "status": status, "headers": headers})
|
|
168
|
+
await send({"type": "http.response.body", "body": body})
|
|
@@ -0,0 +1,173 @@
|
|
|
1
|
+
"""An in-memory AAP provider for tests.
|
|
2
|
+
|
|
3
|
+
``FakeProvider`` implements the AAP request operations, including consume
|
|
4
|
+
and outcome, as an httpx transport, so ``Client`` and ``AsyncClient`` talk
|
|
5
|
+
to it without a network. Framework shims run the conformance scenarios in
|
|
6
|
+
their own tests against it.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import json
|
|
12
|
+
import threading
|
|
13
|
+
import time
|
|
14
|
+
import uuid
|
|
15
|
+
from dataclasses import dataclass, field
|
|
16
|
+
from datetime import datetime, timedelta, timezone
|
|
17
|
+
from typing import Any, Callable
|
|
18
|
+
|
|
19
|
+
import httpx
|
|
20
|
+
|
|
21
|
+
from ._client import AsyncClient, Client
|
|
22
|
+
from .receiver import sign
|
|
23
|
+
|
|
24
|
+
BASE_URL = "https://provider.test/aap"
|
|
25
|
+
TOKEN = "test-token"
|
|
26
|
+
SECRET = "whsec_AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _stamp(at: datetime) -> str:
|
|
30
|
+
return at.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@dataclass
|
|
34
|
+
class Stored:
|
|
35
|
+
body: dict[str, Any]
|
|
36
|
+
id: str
|
|
37
|
+
created_at: datetime
|
|
38
|
+
deadline_at: datetime
|
|
39
|
+
status: str = "pending"
|
|
40
|
+
note: str | None = None
|
|
41
|
+
decided_at: datetime | None = None
|
|
42
|
+
expires_at: datetime | None = None
|
|
43
|
+
consumed_key: str | None = None
|
|
44
|
+
consumed_at: datetime | None = None
|
|
45
|
+
outcome: dict[str, Any] | None = None
|
|
46
|
+
|
|
47
|
+
def wire(self) -> dict[str, Any]:
|
|
48
|
+
out: dict[str, Any] = {**self.body, "id": self.id, "status": self.status, "created_at": _stamp(self.created_at), "deadline_at": _stamp(self.deadline_at)}
|
|
49
|
+
if self.decided_at:
|
|
50
|
+
decision: dict[str, Any] = {"status": self.status, "decided_at": _stamp(self.decided_at)}
|
|
51
|
+
if self.note:
|
|
52
|
+
decision["note"] = self.note
|
|
53
|
+
if self.expires_at:
|
|
54
|
+
decision["expires_at"] = _stamp(self.expires_at)
|
|
55
|
+
out["decision"] = decision
|
|
56
|
+
if self.consumed_at:
|
|
57
|
+
out["consumption"] = {"consumed_at": _stamp(self.consumed_at), **({"outcome": self.outcome} if self.outcome else {})}
|
|
58
|
+
return out
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
@dataclass
|
|
62
|
+
class FakeProvider:
|
|
63
|
+
"""In-memory provider. ``auto`` decides each new request at once."""
|
|
64
|
+
|
|
65
|
+
auto: str | None = "approved"
|
|
66
|
+
validity: timedelta = timedelta(minutes=5)
|
|
67
|
+
requests: dict[str, Stored] = field(default_factory=dict)
|
|
68
|
+
keys: dict[str, str] = field(default_factory=dict)
|
|
69
|
+
calls: list[tuple[str, str]] = field(default_factory=list)
|
|
70
|
+
#: Optional hook to replace a response: return an httpx.Response to use it.
|
|
71
|
+
intercept: Callable[[httpx.Request], httpx.Response | None] | None = None
|
|
72
|
+
|
|
73
|
+
def __post_init__(self) -> None:
|
|
74
|
+
self._lock = threading.Lock()
|
|
75
|
+
|
|
76
|
+
# Test controls
|
|
77
|
+
|
|
78
|
+
def decide(self, request_id: str, status: str = "approved", note: str | None = None) -> None:
|
|
79
|
+
with self._lock:
|
|
80
|
+
self._decide(self.requests[request_id], status, note)
|
|
81
|
+
|
|
82
|
+
def only(self) -> Stored:
|
|
83
|
+
(stored,) = self.requests.values()
|
|
84
|
+
return stored
|
|
85
|
+
|
|
86
|
+
def client(self, **kwargs: Any) -> Client:
|
|
87
|
+
return Client(BASE_URL, TOKEN, http=httpx.Client(transport=httpx.MockTransport(self.handle)), **kwargs)
|
|
88
|
+
|
|
89
|
+
def async_client(self, **kwargs: Any) -> AsyncClient:
|
|
90
|
+
return AsyncClient(BASE_URL, TOKEN, http=httpx.AsyncClient(transport=httpx.MockTransport(self.handle)), **kwargs)
|
|
91
|
+
|
|
92
|
+
def webhook(self, request_id: str, *, instance_id: str = "b29a43a7-1949-4b9f-9d61-b7ae66f85d31", secret: str = SECRET) -> tuple[dict[str, str], bytes]:
|
|
93
|
+
"""Headers and body of a signed request.resolved notification."""
|
|
94
|
+
event_id = str(uuid.uuid4())
|
|
95
|
+
body = json.dumps({"id": event_id, "type": "request.resolved", "instance_id": instance_id, "request_id": request_id, "created_at": _stamp(datetime.now(timezone.utc))}).encode()
|
|
96
|
+
stamp = int(time.time())
|
|
97
|
+
return {"webhook-id": event_id, "webhook-timestamp": str(stamp), "webhook-signature": sign(secret, event_id, stamp, body)}, body
|
|
98
|
+
|
|
99
|
+
# Transport
|
|
100
|
+
|
|
101
|
+
def _decide(self, stored: Stored, status: str, note: str | None) -> None:
|
|
102
|
+
if stored.status != "pending":
|
|
103
|
+
return
|
|
104
|
+
now = datetime.now(timezone.utc)
|
|
105
|
+
stored.status, stored.note, stored.decided_at = status, note, now
|
|
106
|
+
stored.expires_at = now + self.validity if status == "approved" else None
|
|
107
|
+
|
|
108
|
+
def handle(self, request: httpx.Request) -> httpx.Response:
|
|
109
|
+
if self.intercept and (replaced := self.intercept(request)) is not None:
|
|
110
|
+
return replaced
|
|
111
|
+
with self._lock:
|
|
112
|
+
return self._handle(request)
|
|
113
|
+
|
|
114
|
+
def _handle(self, request: httpx.Request) -> httpx.Response:
|
|
115
|
+
if request.headers.get("authorization") != f"Bearer {TOKEN}":
|
|
116
|
+
return httpx.Response(401)
|
|
117
|
+
path = request.url.path.removeprefix("/aap/v1/requests")
|
|
118
|
+
self.calls.append((request.method, path or "/"))
|
|
119
|
+
key = request.headers.get("idempotency-key")
|
|
120
|
+
if request.method == "POST" and path == "":
|
|
121
|
+
if not key:
|
|
122
|
+
return _error(400, "invalid_request", "missing_idempotency_key")
|
|
123
|
+
body = json.loads(request.content)
|
|
124
|
+
if key in self.keys:
|
|
125
|
+
stored = self.requests[self.keys[key]]
|
|
126
|
+
if json.dumps(stored.body, sort_keys=True, default=str) != json.dumps(body, sort_keys=True, default=str):
|
|
127
|
+
return _error(409, "conflict", "idempotency_conflict")
|
|
128
|
+
return httpx.Response(200, json=stored.wire())
|
|
129
|
+
now = datetime.now(timezone.utc)
|
|
130
|
+
stored = Stored(body=body, id=str(uuid.uuid4()), created_at=now, deadline_at=now + timedelta(minutes=30))
|
|
131
|
+
self.requests[stored.id] = stored
|
|
132
|
+
self.keys[key] = stored.id
|
|
133
|
+
if self.auto:
|
|
134
|
+
self._decide(stored, self.auto, None)
|
|
135
|
+
return httpx.Response(201, json=stored.wire())
|
|
136
|
+
parts = path.strip("/").split("/")
|
|
137
|
+
stored = self.requests.get(parts[0])
|
|
138
|
+
if stored is None:
|
|
139
|
+
return _error(404, "not_found", "not_found")
|
|
140
|
+
if request.method == "GET" and len(parts) == 1:
|
|
141
|
+
return httpx.Response(200, json=stored.wire())
|
|
142
|
+
if request.method == "DELETE" and len(parts) == 1:
|
|
143
|
+
if stored.status != "pending" and stored.status != "cancelled":
|
|
144
|
+
return _error(409, "conflict", "already_terminal")
|
|
145
|
+
self._decide(stored, "cancelled", None)
|
|
146
|
+
return httpx.Response(200, json=stored.wire())
|
|
147
|
+
if request.method == "POST" and parts[1:] == ["consume"]:
|
|
148
|
+
if not key:
|
|
149
|
+
return _error(400, "invalid_request", "missing_idempotency_key")
|
|
150
|
+
if stored.consumed_key == key:
|
|
151
|
+
return httpx.Response(200, json=stored.wire())
|
|
152
|
+
if stored.consumed_key is not None:
|
|
153
|
+
return _error(409, "conflict", "already_consumed")
|
|
154
|
+
if stored.status != "approved":
|
|
155
|
+
return _error(409, "conflict", "not_approved")
|
|
156
|
+
if stored.expires_at is None or datetime.now(timezone.utc) >= stored.expires_at:
|
|
157
|
+
return _error(409, "conflict", "approval_expired")
|
|
158
|
+
stored.consumed_key, stored.consumed_at = key, datetime.now(timezone.utc)
|
|
159
|
+
return httpx.Response(200, json=stored.wire())
|
|
160
|
+
if request.method == "POST" and parts[1:] == ["outcome"]:
|
|
161
|
+
report = json.loads(request.content)
|
|
162
|
+
if stored.consumed_at is None:
|
|
163
|
+
return _error(409, "conflict", "not_consumed")
|
|
164
|
+
if stored.outcome is not None:
|
|
165
|
+
same = stored.outcome["status"] == report.get("status") and stored.outcome.get("note") == report.get("note")
|
|
166
|
+
return httpx.Response(200, json=stored.wire()) if same else _error(409, "conflict", "outcome_conflict")
|
|
167
|
+
stored.outcome = {**report, "reported_at": _stamp(datetime.now(timezone.utc))}
|
|
168
|
+
return httpx.Response(200, json=stored.wire())
|
|
169
|
+
return httpx.Response(405)
|
|
170
|
+
|
|
171
|
+
|
|
172
|
+
def _error(status: int, kind: str, code: str) -> httpx.Response:
|
|
173
|
+
return httpx.Response(status, json={"error": {"type": kind, "code": code, "message": code, "request_id": str(uuid.uuid4())}})
|