toolbound 0.2.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.
@@ -0,0 +1,1461 @@
1
+ """The Gateway: one execution pipeline for every tool.
2
+
3
+ REST operations and Python tools pass through the same steps::
4
+
5
+ resolve tool (visibility) -> validate arguments -> resolve trusted bindings
6
+ -> open cursor -> required scopes -> authorize() -> rate limit
7
+ -> approval (if required) -> replay check -> execute -> shape output -> result
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import asyncio
13
+ import base64
14
+ import concurrent.futures
15
+ import functools
16
+ import inspect
17
+ import json
18
+ import logging
19
+ import random
20
+ import re
21
+ import secrets
22
+ import threading
23
+ import time
24
+ import warnings
25
+ import weakref
26
+ from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
27
+ from dataclasses import dataclass, replace
28
+ from pathlib import Path as FsPath
29
+ from typing import Any, TypeVar, cast
30
+ from urllib.parse import urlencode, urljoin, urlsplit
31
+
32
+ from .approval import ApprovalManager, PendingApproval
33
+ from .budget import ByteEstimator
34
+ from .context import ExecutionContext
35
+ from .contracts import (
36
+ SAFE_METHODS,
37
+ ApiKeyAuth,
38
+ BasicAuth,
39
+ BearerAuth,
40
+ Effect,
41
+ NoAuth,
42
+ Operation,
43
+ Output,
44
+ PythonTool,
45
+ Tool,
46
+ resolve_binding,
47
+ )
48
+ from .errors import (
49
+ ApprovalInvalid,
50
+ ConfigurationError,
51
+ ErrorCode,
52
+ NetworkBlocked,
53
+ ToolboundWarning,
54
+ ToolFailure,
55
+ )
56
+ from .exposure import Strategy, ToolDefinition, group_definitions
57
+ from .network import build_url, check_url, same_origin
58
+ from .observe import Event, NullObserver
59
+ from .paths import MISSING, parse_path
60
+ from .ports import (
61
+ Approver,
62
+ AuthorizationRequest,
63
+ Authorizer,
64
+ CredentialProvider,
65
+ Decision,
66
+ HttpRequest,
67
+ HttpResponse,
68
+ Observer,
69
+ StateStore,
70
+ TokenEstimator,
71
+ Transport,
72
+ )
73
+ from .pyschema import handler_parameters, infer_input_schema
74
+ from .ratelimit import RateLimiter
75
+ from .result import ResultMeta, ToolError, ToolResult
76
+ from .schema import RESERVED_ARGUMENTS, ArgumentValidator, normalize_input_schema
77
+ from .shaping import CompiledOutput, ShapeResult, compile_output, process, shape_error_body
78
+ from .signing import Signer, canonical_digest
79
+ from .state import MemoryStateStore, claim_or_fail
80
+
81
+ logger = logging.getLogger("toolbound")
82
+
83
+ _CURSOR_PREFIX = "tbc2"
84
+ _MAX_CURSOR = 4096
85
+ #: Above this many bytes, output shaping runs in a worker thread so a large
86
+ #: upstream response cannot stall the event loop.
87
+ _OFFLOAD_BYTES = 256 * 1024
88
+ _USER_AGENT = "toolbound/0.2"
89
+ _REDIRECTS = frozenset({301, 302, 303, 307, 308})
90
+ _REQUEST_ID_HEADERS = ("x-request-id", "x-github-request-id", "request-id", "x-amzn-requestid")
91
+ _SAFE_ID = re.compile(r"^[A-Za-z0-9._:/-]{1,80}$")
92
+ _JSON_TYPE = re.compile(r"^application/(?:[\w.+-]+\+)?json$")
93
+
94
+ F = TypeVar("F", bound=Callable[..., Any])
95
+
96
+
97
+ @dataclass(slots=True)
98
+ class _Registered:
99
+ spec: Tool
100
+ schema: dict[str, Any]
101
+ validator: ArgumentValidator
102
+ output: CompiledOutput
103
+ fingerprint: str
104
+ selectable: bool
105
+ has_cursor: bool
106
+ accepts_ctx: bool = False
107
+
108
+
109
+ class _ApprovalNeeded(Exception):
110
+ def __init__(self, pending: PendingApproval) -> None:
111
+ self.pending = pending
112
+
113
+
114
+ def allow_all(request: AuthorizationRequest) -> bool:
115
+ """Authorizer that allows every visible tool. Use only for demos and tests."""
116
+ return True
117
+
118
+
119
+ def deny_all(request: AuthorizationRequest) -> bool:
120
+ """Authorizer that denies everything."""
121
+ return False
122
+
123
+
124
+ class Gateway:
125
+ """Registers tools and executes model tool calls under trusted policy.
126
+
127
+ Args:
128
+ authorize: **Required.** Your policy hook ``(AuthorizationRequest) -> bool |
129
+ Decision`` (sync or async). Pass :func:`allow_all` explicitly to
130
+ disable application-level authorization.
131
+ credentials: Resolves credential names to secret values.
132
+ approver: Optional inline approval hook. Without one, calls that need
133
+ approval return ``approval_required`` with a :class:`PendingApproval`
134
+ for your UI, and you finish them with :meth:`resume`.
135
+ transport: HTTP transport. Defaults to :class:`toolbound.SecureTransport`.
136
+ estimator: Token estimator. Defaults to :class:`ByteEstimator`.
137
+ observer: Receives structured events. Defaults to no-op.
138
+ state: Store for single-use approvals and call-id replay protection.
139
+ secret: HMAC key (>= 32 bytes) for cursors and approval grants. If
140
+ omitted a random per-process key is used, so cursors and grants do
141
+ not survive restarts or work across processes.
142
+ visibility: Optional ``(ctx, tool) -> bool`` deciding which tools a
143
+ user can even see. Hidden tools behave exactly like unknown tools.
144
+ transforms: Named trusted transforms usable from declarative shapes.
145
+ max_argument_bytes: Upper bound for serialized tool arguments.
146
+ """
147
+
148
+ def __init__(
149
+ self,
150
+ *,
151
+ authorize: Authorizer,
152
+ credentials: CredentialProvider | None = None,
153
+ approver: Approver | None = None,
154
+ transport: Transport | None = None,
155
+ estimator: TokenEstimator | None = None,
156
+ observer: Observer | None = None,
157
+ state: StateStore | None = None,
158
+ secret: bytes | str | None = None,
159
+ visibility: Callable[[ExecutionContext, Tool], bool] | None = None,
160
+ transforms: Mapping[str, Callable[[Any], Any]] | None = None,
161
+ approval_ttl: float = 600,
162
+ cursor_ttl: float = 3600,
163
+ call_id_ttl: float = 86_400,
164
+ max_argument_bytes: int = 64 * 1024,
165
+ sleep: Callable[[float], Awaitable[Any]] = asyncio.sleep,
166
+ rng: random.Random | None = None,
167
+ ) -> None:
168
+ if authorize is None or not callable(authorize):
169
+ raise ConfigurationError(
170
+ "Gateway(authorize=...) is required. Pass your policy function, or "
171
+ "toolbound.allow_all to explicitly disable application authorization."
172
+ )
173
+ self._authorize_hook = authorize
174
+ self._credentials = credentials
175
+ self._approver = approver
176
+ self._transport = transport
177
+ self._owns_transport = transport is None
178
+ self._estimator: TokenEstimator = estimator if estimator is not None else ByteEstimator()
179
+ self._observer: Observer = observer if observer is not None else NullObserver()
180
+ self._state: StateStore = state if state is not None else MemoryStateStore()
181
+ self._signer = Signer(secret)
182
+ self._approvals = ApprovalManager(self._signer, self._state, ttl_seconds=approval_ttl)
183
+ self._visibility = visibility
184
+ self._transforms = dict(transforms or {})
185
+ self._cursor_ttl = cursor_ttl
186
+ self._call_id_ttl = call_id_ttl
187
+ self._max_argument_bytes = max_argument_bytes
188
+ self._sleep = sleep
189
+ self._rng = rng or random.Random() # noqa: S311 - jitter, not crypto
190
+ self._tools: dict[str, _Registered] = {}
191
+ self._groups: dict[str, list[str]] = {}
192
+ # asyncio primitives belong to one event loop; keep them per loop.
193
+ self._semaphores: weakref.WeakKeyDictionary[
194
+ asyncio.AbstractEventLoop, dict[str, asyncio.Semaphore]
195
+ ] = weakref.WeakKeyDictionary()
196
+ self._sem_lock = threading.Lock()
197
+ self._rate = RateLimiter()
198
+ self._warned_ephemeral = False
199
+ self._executor: concurrent.futures.ThreadPoolExecutor | None = None
200
+
201
+ # ------------------------------------------------------------------ lifecycle
202
+
203
+ async def __aenter__(self) -> Gateway:
204
+ return self
205
+
206
+ async def __aexit__(self, *exc: object) -> None:
207
+ await self.aclose()
208
+
209
+ async def aclose(self) -> None:
210
+ """Close the transport if the gateway created it."""
211
+ if self._transport is not None and self._owns_transport:
212
+ await self._transport.aclose()
213
+ self._transport = None
214
+ if self._executor is not None:
215
+ self._executor.shutdown(wait=False, cancel_futures=True)
216
+ self._executor = None
217
+
218
+ def _get_transport(self) -> Transport:
219
+ if self._transport is None:
220
+ from toolbound.transport.http import SecureTransport
221
+
222
+ self._transport = SecureTransport()
223
+ self._owns_transport = True
224
+ return self._transport
225
+
226
+ # ------------------------------------------------------------------ registration
227
+
228
+ def register(self, *tools: Tool) -> Gateway:
229
+ """Register REST operations and/or Python tools. Returns ``self``."""
230
+ for spec in tools:
231
+ self._register_one(spec)
232
+ return self
233
+
234
+ def tool(
235
+ self,
236
+ func: Callable[..., Any] | None = None,
237
+ /,
238
+ *,
239
+ name: str | None = None,
240
+ description: str | None = None,
241
+ input: Mapping[str, Any] | None = None,
242
+ effect: Effect | str | None = None,
243
+ output: Output | None = None,
244
+ bind: Mapping[str, Any] | None = None,
245
+ required_scopes: Sequence[str] = (),
246
+ tags: Sequence[str] = (),
247
+ group: str | None = None,
248
+ requires_approval: bool | None = None,
249
+ timeout: float = 30.0,
250
+ ) -> Any:
251
+ """Decorator registering a trusted Python function as a tool.
252
+
253
+ The input schema is inferred from type hints unless ``input`` is given;
254
+ the description defaults to the first paragraph of the docstring.
255
+ """
256
+
257
+ def decorate(fn: F) -> F:
258
+ doc = inspect.getdoc(fn) or ""
259
+ spec = PythonTool(
260
+ name=name or fn.__name__,
261
+ description=description or doc.split("\n\n")[0].strip(),
262
+ input=dict(input) if input is not None else {},
263
+ effect=cast(Effect, None if effect is None else Effect(effect)),
264
+ output=output if output is not None else Output(),
265
+ bind=dict(bind or {}),
266
+ required_scopes=tuple(required_scopes),
267
+ tags=tuple(tags),
268
+ group=group,
269
+ requires_approval=requires_approval,
270
+ handler=fn,
271
+ timeout=timeout,
272
+ )
273
+ self._register_one(spec, infer=input is None)
274
+ return fn
275
+
276
+ return decorate(func) if func is not None else decorate
277
+
278
+ def load(self, path: str | FsPath) -> Gateway:
279
+ """Load declarative contracts from a ``.yaml``/``.json`` file or directory."""
280
+ from toolbound.loaders import load_contracts
281
+
282
+ return self.register(*load_contracts(path))
283
+
284
+ @property
285
+ def tool_names(self) -> list[str]:
286
+ return list(self._tools)
287
+
288
+ def get_tool(self, name: str) -> Tool:
289
+ return self._tools[name].spec
290
+
291
+ def _register_one(self, spec: Tool, *, infer: bool = False) -> None:
292
+ if not isinstance(spec, Operation | PythonTool):
293
+ raise ConfigurationError(f"cannot register {type(spec).__name__}; expected a Tool")
294
+ if spec.name in self._tools or spec.name in self._groups:
295
+ raise ConfigurationError(f"a tool named {spec.name!r} is already registered")
296
+ if spec.group is not None and spec.group in self._tools:
297
+ raise ConfigurationError(f"group {spec.group!r} clashes with a tool name")
298
+
299
+ bound_names = set(spec.bind)
300
+ if isinstance(spec, PythonTool) and (infer or not spec.input):
301
+ raw_schema: Mapping[str, Any] = infer_input_schema(spec.handler, exclude=bound_names)
302
+ else:
303
+ raw_schema = spec.input
304
+ schema = normalize_input_schema(spec.name, raw_schema)
305
+ props: dict[str, Any] = schema["properties"]
306
+ required = set(schema.get("required", []))
307
+
308
+ clash = bound_names & set(props)
309
+ if clash:
310
+ raise ConfigurationError(
311
+ f"tool {spec.name}: bound parameter(s) {sorted(clash)} must not appear in the "
312
+ "model-visible input"
313
+ )
314
+ reserved = RESERVED_ARGUMENTS & set(props)
315
+ if reserved:
316
+ raise ConfigurationError(
317
+ f"tool {spec.name}: argument names {sorted(reserved)} are reserved"
318
+ )
319
+
320
+ accepts_ctx = False
321
+ if isinstance(spec, Operation):
322
+ for param in spec.path_params:
323
+ if param in bound_names:
324
+ continue
325
+ if param not in props:
326
+ raise ConfigurationError(
327
+ f"tool {spec.name}: path parameter {{{param}}} is neither an input "
328
+ "property nor bound"
329
+ )
330
+ if param not in required:
331
+ raise ConfigurationError(
332
+ f"tool {spec.name}: path parameter {param!r} must be required"
333
+ )
334
+ known = set(props) | bound_names
335
+ for listed in (spec.query or ()) + (spec.body or ()):
336
+ if listed not in known:
337
+ raise ConfigurationError(
338
+ f"tool {spec.name}: {listed!r} in query/body is not an input or binding"
339
+ )
340
+ pagination = spec.pagination
341
+ if pagination is not None and pagination.param in known:
342
+ raise ConfigurationError(
343
+ f"tool {spec.name}: pagination param {pagination.param!r} must not be an input"
344
+ )
345
+ if pagination is not None and pagination.next_path:
346
+ parse_path(pagination.next_path)
347
+ else:
348
+ names, var_kw, accepts_ctx = handler_parameters(spec.handler)
349
+ missing = (set(props) | bound_names) - names
350
+ if missing and not var_kw:
351
+ raise ConfigurationError(
352
+ f"tool {spec.name}: handler does not accept argument(s) {sorted(missing)}"
353
+ )
354
+
355
+ compiled = compile_output(spec.output, self._transforms)
356
+ selectable = spec.output.selectable
357
+ has_cursor = spec.effect is Effect.READ and (
358
+ spec.output.continuation
359
+ or (isinstance(spec, Operation) and spec.pagination is not None)
360
+ )
361
+ if selectable:
362
+ keys = compiled.selectable_fields
363
+ always = sorted(compiled.required)
364
+ desc = "Optional subset of fields to return; omit for all."
365
+ if always:
366
+ desc += f" Always included: {', '.join(always)}."
367
+ props["fields"] = {
368
+ "type": "array",
369
+ "items": {"type": "string", "enum": keys},
370
+ "uniqueItems": True,
371
+ "description": desc,
372
+ }
373
+ if has_cursor:
374
+ props["cursor"] = {
375
+ "type": "string",
376
+ "maxLength": 4096,
377
+ "description": "Opaque next_cursor from a previous result, to get more.",
378
+ }
379
+
380
+ validator = ArgumentValidator(schema, max_argument_bytes=self._max_argument_bytes)
381
+ fingerprint = canonical_digest(_fingerprint_source(spec, schema))
382
+ self._tools[spec.name] = _Registered(
383
+ spec=spec,
384
+ schema=schema,
385
+ validator=validator,
386
+ output=compiled,
387
+ fingerprint=fingerprint,
388
+ selectable=selectable,
389
+ has_cursor=has_cursor,
390
+ accepts_ctx=accepts_ctx,
391
+ )
392
+ if spec.group is not None:
393
+ self._groups.setdefault(spec.group, []).append(spec.name)
394
+
395
+ # ------------------------------------------------------------------ exposure
396
+
397
+ def definitions(
398
+ self,
399
+ ctx: ExecutionContext,
400
+ *,
401
+ tags: Iterable[str] | None = None,
402
+ names: Iterable[str] | None = None,
403
+ strategy: Strategy = "direct",
404
+ ) -> list[ToolDefinition]:
405
+ """Tool definitions visible to ``ctx``, optionally filtered by tags/names.
406
+
407
+ ``strategy="grouped"`` merges tools sharing a ``group`` into one tool with
408
+ an ``action`` argument (fewer definitions for larger catalogs).
409
+ """
410
+ wanted_tags = set(tags) if tags is not None else None
411
+ wanted_names = set(names) if names is not None else None
412
+ visible: list[_Registered] = []
413
+ for reg in self._tools.values():
414
+ if wanted_names is not None and reg.spec.name not in wanted_names:
415
+ continue
416
+ if wanted_tags is not None and not wanted_tags & set(reg.spec.tags):
417
+ continue
418
+ if self._visible(reg.spec, ctx):
419
+ visible.append(reg)
420
+ direct = [
421
+ ToolDefinition(r.spec.name, r.spec.description, json.loads(json.dumps(r.schema)))
422
+ for r in visible
423
+ ]
424
+ if strategy == "direct":
425
+ return direct
426
+ if strategy != "grouped":
427
+ raise ValueError("strategy must be 'direct' or 'grouped'")
428
+ out: list[ToolDefinition] = []
429
+ grouped: dict[str, list[ToolDefinition]] = {}
430
+ for reg, definition in zip(visible, direct, strict=True):
431
+ if reg.spec.group is None:
432
+ out.append(definition)
433
+ else:
434
+ grouped.setdefault(reg.spec.group, []).append(definition)
435
+ for group, members in grouped.items():
436
+ out.append(members[0] if len(members) == 1 else group_definitions(group, members))
437
+ return out
438
+
439
+ def _visible(self, spec: Tool, ctx: ExecutionContext) -> bool:
440
+ if not set(spec.required_scopes) <= ctx.scopes:
441
+ return False
442
+ return self._visibility is None or bool(self._visibility(ctx, spec))
443
+
444
+ # ------------------------------------------------------------------ approvals
445
+
446
+ def pending_approval(self, request: PendingApproval | str) -> PendingApproval:
447
+ """Rebuild and verify a :class:`PendingApproval` from its ``token``.
448
+
449
+ Use it when an approval request comes back from storage, a browser or a
450
+ workflow engine. Raises :class:`~toolbound.ApprovalInvalid` if the token
451
+ is forged, altered, expired, or from another gateway secret.
452
+ """
453
+ return self._approvals.open(request)
454
+
455
+ def grant(self, request: PendingApproval | str, *, approved_by: str) -> str:
456
+ """Issue a single-use grant for a gateway-issued approval request (host code only).
457
+
458
+ ``request`` is a :class:`PendingApproval` or its ``token``. Only the
459
+ token is trusted: the other fields of a ``PendingApproval`` object are
460
+ ignored. Raises :class:`~toolbound.ApprovalInvalid` for a forged or
461
+ expired request.
462
+ """
463
+ return self._approvals.grant(request, approved_by=approved_by)
464
+
465
+ async def deny(self, request: PendingApproval | str) -> PendingApproval:
466
+ """Record that a human rejected a request; it can never be granted afterwards.
467
+
468
+ Raises :class:`~toolbound.ApprovalInvalid` for a forged or expired request.
469
+ """
470
+ return await self._approvals.deny(request)
471
+
472
+ async def resume(
473
+ self,
474
+ request: PendingApproval | str,
475
+ *,
476
+ ctx: ExecutionContext,
477
+ approved_by: str,
478
+ call_id: str | None = None,
479
+ ) -> ToolResult:
480
+ """Approve a request and execute it in one step (host code only).
481
+
482
+ The tool and arguments come from the verified request token, and the
483
+ call runs under the live ``ctx``, which must match the requesting actor
484
+ and tenant. Never raises for an invalid request: it returns an
485
+ ``approval_invalid`` result.
486
+ """
487
+ if not isinstance(ctx, ExecutionContext):
488
+ raise TypeError("ctx must be an ExecutionContext created by trusted code")
489
+ try:
490
+ pending = self._approvals.open(request)
491
+ token = self._approvals.grant(pending, approved_by=approved_by)
492
+ except ApprovalInvalid as failure:
493
+ return ToolResult(
494
+ ok=False,
495
+ tool=request.tool if isinstance(request, PendingApproval) else "<approval>",
496
+ error=ToolError(failure.code, failure.message),
497
+ call_id=call_id,
498
+ )
499
+ return await self.call(
500
+ pending.tool, dict(pending.arguments), ctx=ctx, approval=token, call_id=call_id
501
+ )
502
+
503
+ # ------------------------------------------------------------------ execution
504
+
505
+ async def call(
506
+ self,
507
+ name: str,
508
+ arguments: Mapping[str, Any] | None,
509
+ *,
510
+ ctx: ExecutionContext,
511
+ call_id: str | None = None,
512
+ approval: str | None = None,
513
+ ) -> ToolResult:
514
+ """Execute one tool call. Never raises for expected failures.
515
+
516
+ Args:
517
+ name: Tool (or group) name chosen by the model.
518
+ arguments: Arguments chosen by the model (untrusted).
519
+ ctx: Trusted execution context from your application.
520
+ call_id: The provider's tool-call id. For write/destructive tools a
521
+ repeated id is rejected, preventing accidental double execution.
522
+ approval: A grant token from :meth:`grant` (host code only).
523
+ """
524
+ if not isinstance(ctx, ExecutionContext):
525
+ raise TypeError("ctx must be an ExecutionContext created by trusted code")
526
+ started = time.perf_counter()
527
+ display = name if isinstance(name, str) and len(name) <= 64 else "<invalid>"
528
+ self._emit("call.start", display, ctx)
529
+ http_status: list[int] = []
530
+ try:
531
+ name, arguments = self._resolve_group(name, arguments, ctx)
532
+ reg = self._tools.get(name) if isinstance(name, str) else None
533
+ if reg is None or not self._visible(reg.spec, ctx):
534
+ raise ToolFailure(ErrorCode.UNKNOWN_TOOL, f"unknown tool {display!r}")
535
+ display = reg.spec.name
536
+ result = await self._run(reg, arguments, ctx, call_id, approval, http_status)
537
+ except _ApprovalNeeded as needed:
538
+ self._emit("approval.required", display, ctx, outcome="approval_required")
539
+ result = ToolResult(
540
+ ok=False,
541
+ tool=display,
542
+ error=ToolError(
543
+ ErrorCode.APPROVAL_REQUIRED,
544
+ "This action needs human approval. The request was sent to the user; "
545
+ "wait for their decision instead of retrying.",
546
+ ),
547
+ approval=needed.pending,
548
+ call_id=call_id,
549
+ )
550
+ except ToolFailure as failure:
551
+ if failure.code in (ErrorCode.FORBIDDEN,):
552
+ self._emit("policy.denied", display, ctx, outcome=failure.code.value)
553
+ elif failure.code is ErrorCode.NETWORK_BLOCKED:
554
+ self._emit("network.blocked", display, ctx, outcome=failure.code.value)
555
+ result = ToolResult(
556
+ ok=False,
557
+ tool=display,
558
+ error=ToolError(
559
+ failure.code,
560
+ failure.message,
561
+ retryable=failure.retryable,
562
+ details=dict(failure.details),
563
+ retry_after=failure.retry_after,
564
+ ),
565
+ http_status=failure.http_status,
566
+ call_id=call_id,
567
+ )
568
+ except asyncio.CancelledError:
569
+ self._emit("call.end", display, ctx, outcome="cancelled")
570
+ raise
571
+ except Exception:
572
+ logger.exception("tool %s failed with an unexpected error", display)
573
+ result = ToolResult(
574
+ ok=False,
575
+ tool=display,
576
+ error=ToolError(ErrorCode.INTERNAL_ERROR, "the tool failed with an internal error"),
577
+ call_id=call_id,
578
+ )
579
+ duration = (time.perf_counter() - started) * 1000
580
+ result = replace(result, duration_ms=duration)
581
+ self._emit(
582
+ "call.end",
583
+ display,
584
+ ctx,
585
+ outcome="ok" if result.ok else result.error.code.value if result.error else None,
586
+ duration_ms=duration,
587
+ http_status=result.http_status,
588
+ )
589
+ return result
590
+
591
+ def _resolve_group(
592
+ self, name: str, arguments: Mapping[str, Any] | None, ctx: ExecutionContext
593
+ ) -> tuple[str, Mapping[str, Any] | None]:
594
+ if not isinstance(name, str) or name not in self._groups:
595
+ return name, arguments
596
+ if not isinstance(arguments, Mapping):
597
+ raise ToolFailure(ErrorCode.INVALID_ARGUMENTS, "arguments must be a JSON object")
598
+ extra = set(arguments) - {"action", "arguments"}
599
+ if extra:
600
+ raise ToolFailure(ErrorCode.INVALID_ARGUMENTS, f"unknown argument(s) {sorted(extra)}")
601
+ action = arguments.get("action")
602
+ visible = [n for n in self._groups[name] if self._visible(self._tools[n].spec, ctx)]
603
+ if action not in visible: # hidden actions look exactly like missing ones
604
+ raise ToolFailure(ErrorCode.INVALID_ARGUMENTS, f"action must be one of {visible}")
605
+ inner = arguments.get("arguments", {})
606
+ return str(action), inner
607
+
608
+ async def _run(
609
+ self,
610
+ reg: _Registered,
611
+ arguments: Mapping[str, Any] | None,
612
+ ctx: ExecutionContext,
613
+ call_id: str | None,
614
+ approval: str | None,
615
+ http_status: list[int],
616
+ ) -> ToolResult:
617
+ spec = reg.spec
618
+ raw = {} if arguments is None else arguments
619
+ args = reg.validator.validate(dict(raw) if isinstance(raw, Mapping) else raw)
620
+ fields = args.pop("fields", None) if reg.selectable else None
621
+ cursor_token = args.pop("cursor", None) if reg.has_cursor else None
622
+ bound = self._resolve_bound(spec, ctx)
623
+ cursor_state = (
624
+ self._open_cursor(cursor_token, reg, ctx, args, bound) if cursor_token else None
625
+ )
626
+
627
+ missing_scopes = set(spec.required_scopes) - ctx.scopes
628
+ if missing_scopes:
629
+ raise ToolFailure(ErrorCode.FORBIDDEN, "missing required permission for this tool")
630
+ decision = await self._authorize(AuthorizationRequest(ctx, spec, dict(args), dict(bound)))
631
+ if not decision.allowed:
632
+ raise ToolFailure(ErrorCode.FORBIDDEN, decision.reason or "not permitted")
633
+ if spec.rate_limit is not None:
634
+ self._rate.check(spec.name, spec.rate_limit, ctx)
635
+
636
+ grant: dict[str, Any] | None = None
637
+ if spec.needs_approval:
638
+ grant = await self._ensure_approval(reg, ctx, args, bound, approval)
639
+
640
+ async def commit() -> None:
641
+ """Claim the call id, then use up the approval: right before execution,
642
+ after every check that can fail without side effects."""
643
+ if call_id is not None and spec.effect is not Effect.READ:
644
+ key = (
645
+ f"call:{canonical_digest([ctx.actor, ctx.tenant])[:16]}:"
646
+ f"{canonical_digest([spec.name, str(call_id)])}"
647
+ )
648
+ if not await claim_or_fail(self._state, key, self._call_id_ttl):
649
+ raise ToolFailure(
650
+ ErrorCode.DUPLICATE_CALL,
651
+ "this tool call was already executed; it will not run twice",
652
+ )
653
+ if grant is not None:
654
+ await self._approvals.claim(grant, ctx)
655
+ self._emit(
656
+ "approval.granted",
657
+ spec.name,
658
+ ctx,
659
+ outcome="ok",
660
+ attributes={"approved_by": grant.get("by")},
661
+ )
662
+
663
+ idempotency_key = (str(grant["i"]) if grant is not None else None) or (
664
+ canonical_digest([ctx.actor, ctx.tenant, spec.name, str(call_id)])[:32]
665
+ if call_id
666
+ else None
667
+ )
668
+ if isinstance(spec, Operation):
669
+ shaped, meta_source, next_upstream, status = await self._execute_rest(
670
+ reg, spec, ctx, args, bound, cursor_state, fields, idempotency_key, commit
671
+ )
672
+ http_status.append(status)
673
+ else:
674
+ await commit()
675
+ shaped = await self._execute_python(reg, spec, ctx, args, bound, cursor_state, fields)
676
+ meta_source, next_upstream, status = {"tool": spec.name}, None, None
677
+
678
+ next_cursor = None
679
+ if reg.has_cursor:
680
+ current_upstream = cursor_state.get("u") if cursor_state else None
681
+ if shaped.next_offset is not None:
682
+ next_cursor = self._issue_cursor(
683
+ reg, ctx, args, bound, current_upstream, shaped.next_offset
684
+ )
685
+ elif next_upstream is not None:
686
+ next_cursor = self._issue_cursor(reg, ctx, args, bound, next_upstream, 0)
687
+
688
+ report = shaped.report
689
+ meta = ResultMeta(
690
+ truncated=report.truncated,
691
+ omitted=tuple(report.omissions()),
692
+ redacted=report.redacted,
693
+ warnings=tuple(report.warnings),
694
+ next_cursor=next_cursor,
695
+ items_returned=len(shaped.data) if isinstance(shaped.data, list) else None,
696
+ items_total=shaped.total_items,
697
+ estimated_tokens=shaped.estimated_tokens,
698
+ source=meta_source,
699
+ )
700
+ for warning in report.warnings:
701
+ self._emit(
702
+ "schema.warning", spec.name, ctx, outcome="warning", attributes={"warning": warning}
703
+ )
704
+ return ToolResult(
705
+ ok=True,
706
+ tool=spec.name,
707
+ data=shaped.data,
708
+ meta=meta,
709
+ http_status=status,
710
+ call_id=call_id,
711
+ )
712
+
713
+ def _resolve_bound(self, spec: Tool, ctx: ExecutionContext) -> dict[str, Any]:
714
+ bound: dict[str, Any] = {}
715
+ for name, binding in spec.bind.items():
716
+ try:
717
+ bound[name] = resolve_binding(binding, ctx)
718
+ except LookupError as exc:
719
+ logger.error(
720
+ "tool %s: context value %s is missing for binding %r", spec.name, exc, name
721
+ )
722
+ raise ToolFailure(
723
+ ErrorCode.FORBIDDEN, "the current session lacks context required by this tool"
724
+ ) from None
725
+ return bound
726
+
727
+ async def _authorize(self, request: AuthorizationRequest) -> Decision:
728
+ outcome = self._authorize_hook(request)
729
+ if inspect.isawaitable(outcome):
730
+ outcome = await outcome
731
+ if isinstance(outcome, Decision):
732
+ return outcome
733
+ if outcome is True:
734
+ return Decision.allow()
735
+ if outcome is False:
736
+ return Decision.deny()
737
+ raise ConfigurationError("authorize() must return bool or Decision")
738
+
739
+ async def _ensure_approval(
740
+ self,
741
+ reg: _Registered,
742
+ ctx: ExecutionContext,
743
+ args: Mapping[str, Any],
744
+ bound: Mapping[str, Any],
745
+ token: str | None,
746
+ ) -> dict[str, Any]:
747
+ """Return a verified (not yet used) grant, or stop for approval."""
748
+ spec = reg.spec
749
+ if token is not None:
750
+ return self._approvals.verify(
751
+ token,
752
+ ctx=ctx,
753
+ tool=spec.name,
754
+ fingerprint=reg.fingerprint,
755
+ arguments=args,
756
+ bound=bound,
757
+ )
758
+ self._warn_if_ephemeral()
759
+ pending = self._approvals.create(
760
+ ctx=ctx,
761
+ tool=spec.name,
762
+ effect=spec.effect,
763
+ fingerprint=reg.fingerprint,
764
+ arguments=args,
765
+ bound=bound,
766
+ target=self._approval_target(reg, args, bound),
767
+ )
768
+ if self._approver is None:
769
+ raise _ApprovalNeeded(pending)
770
+ decision = self._approver(pending)
771
+ if inspect.isawaitable(decision):
772
+ decision = await decision
773
+ if decision is not True:
774
+ raise ToolFailure(ErrorCode.APPROVAL_DENIED, "the user did not approve this action")
775
+ grant = self._approvals.grant(pending, approved_by="inline-approver")
776
+ return self._approvals.verify(
777
+ grant,
778
+ ctx=ctx,
779
+ tool=spec.name,
780
+ fingerprint=reg.fingerprint,
781
+ arguments=args,
782
+ bound=bound,
783
+ )
784
+
785
+ def _approval_target(
786
+ self, reg: _Registered, args: Mapping[str, Any], bound: Mapping[str, Any]
787
+ ) -> dict[str, Any]:
788
+ spec = reg.spec
789
+ if not isinstance(spec, Operation):
790
+ return {"kind": "python", "tool": spec.name}
791
+ values = {k: args[k] for k in reg.schema["properties"] if k in args}
792
+ values.update(bound)
793
+ path_values = {p: values[p] for p in spec.path_params}
794
+ remaining = {k: v for k, v in values.items() if k not in path_values}
795
+ query_names, body_names = _placement(spec, remaining)
796
+ query = [(k, remaining[k]) for k in query_names if k in remaining]
797
+ return {
798
+ "kind": "rest",
799
+ "service": spec.service.name,
800
+ "method": spec.method,
801
+ "url": build_url(spec.service.base_url, spec.path, path_values, query),
802
+ "body_fields": [k for k in body_names if k in remaining],
803
+ }
804
+
805
+ def _warn_if_ephemeral(self) -> None:
806
+ if self._signer.ephemeral and not self._warned_ephemeral:
807
+ self._warned_ephemeral = True
808
+ warnings.warn(
809
+ "Gateway has no secret=, so it uses a random per-process key: cursors and "
810
+ "approvals stop working after a restart and across workers. Pass a 32+ byte "
811
+ "secret from your secret manager in production.",
812
+ ToolboundWarning,
813
+ stacklevel=3,
814
+ )
815
+
816
+ # ------------------------------------------------------------------ cursors
817
+
818
+ def _cursor_binding(
819
+ self,
820
+ reg: _Registered,
821
+ ctx: ExecutionContext,
822
+ args: Mapping[str, Any],
823
+ bound: Mapping[str, Any],
824
+ ) -> dict[str, Any]:
825
+ return {
826
+ "t": reg.spec.name,
827
+ "f": reg.fingerprint[:16],
828
+ "a": canonical_digest([ctx.actor, ctx.tenant])[:16],
829
+ "h": canonical_digest(dict(args))[:16],
830
+ # Trusted values too: a cursor minted under one org/repo binding must not
831
+ # replay that binding after the session's context changed.
832
+ "b": canonical_digest(dict(bound))[:16],
833
+ }
834
+
835
+ def _issue_cursor(
836
+ self,
837
+ reg: _Registered,
838
+ ctx: ExecutionContext,
839
+ args: Mapping[str, Any],
840
+ bound: Mapping[str, Any],
841
+ upstream: Any,
842
+ offset: int,
843
+ ) -> str | None:
844
+ self._warn_if_ephemeral()
845
+ payload = {**self._cursor_binding(reg, ctx, args, bound), "u": upstream, "o": offset}
846
+ token = self._signer.seal(_CURSOR_PREFIX, payload, self._cursor_ttl)
847
+ if len(token) > _MAX_CURSOR:
848
+ logger.warning("tool %s: upstream cursor too large to carry; stopping", reg.spec.name)
849
+ return None
850
+ return token
851
+
852
+ def _open_cursor(
853
+ self,
854
+ token: Any,
855
+ reg: _Registered,
856
+ ctx: ExecutionContext,
857
+ args: Mapping[str, Any],
858
+ bound: Mapping[str, Any],
859
+ ) -> dict[str, Any]:
860
+ payload = self._signer.unseal(_CURSOR_PREFIX, token, max_length=_MAX_CURSOR)
861
+ if payload is None:
862
+ raise ToolFailure(ErrorCode.INVALID_CURSOR, "cursor is invalid or has expired")
863
+ expected = self._cursor_binding(reg, ctx, args, bound)
864
+ if any(payload.get(k) != v for k, v in expected.items()):
865
+ raise ToolFailure(
866
+ ErrorCode.INVALID_CURSOR,
867
+ "cursor belongs to a different tool, user, or set of arguments",
868
+ )
869
+ if not isinstance(payload.get("o"), int) or payload["o"] < 0:
870
+ raise ToolFailure(ErrorCode.INVALID_CURSOR, "cursor is invalid")
871
+ return payload
872
+
873
+ # ------------------------------------------------------------------ python tools
874
+
875
+ async def _execute_python(
876
+ self,
877
+ reg: _Registered,
878
+ spec: PythonTool,
879
+ ctx: ExecutionContext,
880
+ args: dict[str, Any],
881
+ bound: dict[str, Any],
882
+ cursor_state: dict[str, Any] | None,
883
+ fields: list[str] | None,
884
+ ) -> ShapeResult:
885
+ kwargs = {**args, **bound}
886
+ if reg.accepts_ctx:
887
+ kwargs["ctx"] = ctx
888
+ try:
889
+ async with asyncio.timeout(spec.timeout):
890
+ if inspect.iscoroutinefunction(spec.handler):
891
+ value = await spec.handler(**kwargs)
892
+ else:
893
+ # A private pool: stuck handlers can never starve the default
894
+ # executor that DNS resolution (getaddrinfo) relies on.
895
+ value = await asyncio.get_running_loop().run_in_executor(
896
+ self._handler_pool(), functools.partial(spec.handler, **kwargs)
897
+ )
898
+ if inspect.isawaitable(value):
899
+ value = await value
900
+ except TimeoutError:
901
+ if spec.effect is not Effect.READ:
902
+ # The handler may have done its work (or still be running in a thread).
903
+ raise ToolFailure(
904
+ ErrorCode.OUTCOME_UNKNOWN,
905
+ f"tool timed out after {spec.timeout:g}s; the action may or may not have "
906
+ "been applied. Check the current state before trying again.",
907
+ retryable=False,
908
+ ) from None
909
+ raise ToolFailure(
910
+ ErrorCode.TIMEOUT, f"tool timed out after {spec.timeout:g}s"
911
+ ) from None
912
+ try:
913
+ size = len(json.dumps(value))
914
+ except (TypeError, ValueError):
915
+ raise ToolFailure(
916
+ ErrorCode.INTERNAL_ERROR, "the tool returned a value that is not JSON-serializable"
917
+ ) from None
918
+ offset = int(cursor_state["o"]) if cursor_state else 0
919
+ return await self._process(reg, value, size, offset=offset, fields=fields, secrets=())
920
+
921
+ async def _process(
922
+ self,
923
+ reg: _Registered,
924
+ value: Any,
925
+ size: int,
926
+ *,
927
+ offset: int,
928
+ fields: list[str] | None,
929
+ secrets: Sequence[str],
930
+ ) -> ShapeResult:
931
+ def run() -> ShapeResult:
932
+ return process(
933
+ reg.output,
934
+ value,
935
+ estimator=self._estimator,
936
+ offset=offset,
937
+ selected=fields,
938
+ secrets=secrets,
939
+ )
940
+
941
+ if size > _OFFLOAD_BYTES:
942
+ return await asyncio.to_thread(run)
943
+ return run()
944
+
945
+ # ------------------------------------------------------------------ REST
946
+
947
+ async def _execute_rest(
948
+ self,
949
+ reg: _Registered,
950
+ op: Operation,
951
+ ctx: ExecutionContext,
952
+ args: dict[str, Any],
953
+ bound: dict[str, Any],
954
+ cursor_state: dict[str, Any] | None,
955
+ fields: list[str] | None,
956
+ idempotency_key: str | None,
957
+ commit: Callable[[], Awaitable[None]],
958
+ ) -> tuple[ShapeResult, dict[str, Any], Any, int]:
959
+ svc = op.service
960
+ # Deterministic order (schema order, then bindings) regardless of model output order.
961
+ values = {k: args[k] for k in reg.schema["properties"] if k in args}
962
+ values.update(bound)
963
+ path_values = {p: values[p] for p in op.path_params}
964
+ remaining = {k: v for k, v in values.items() if k not in path_values}
965
+ query_names, body_names = _placement(op, remaining)
966
+ query = [(k, remaining[k]) for k in query_names if k in remaining]
967
+ body = {k: remaining[k] for k in body_names if k in remaining}
968
+
969
+ upstream = cursor_state.get("u") if cursor_state else None
970
+ pagination = op.pagination
971
+ page_number: int | None = None
972
+ if pagination is not None and pagination.type == "page":
973
+ page_number = int(upstream) if isinstance(upstream, int) else pagination.start
974
+ query.append((str(pagination.param), page_number))
975
+ elif pagination is not None and pagination.type == "cursor" and upstream is not None:
976
+ query.append((str(pagination.param), upstream))
977
+
978
+ secrets: list[str] = []
979
+ headers: dict[str, str] = {"Accept": "application/json", "User-Agent": _USER_AGENT}
980
+ headers.update(svc.headers)
981
+ headers.update(op.headers)
982
+ credential_header: str | None = None
983
+ auth = svc.auth
984
+ if not isinstance(auth, NoAuth):
985
+ secret = await self._credential(auth.credential, ctx, svc.name)
986
+ secrets.append(secret)
987
+ if isinstance(auth, BearerAuth):
988
+ credential_header = "Authorization"
989
+ headers["Authorization"] = f"Bearer {secret}"
990
+ elif isinstance(auth, BasicAuth):
991
+ credential_header = "Authorization"
992
+ headers["Authorization"] = "Basic " + base64.b64encode(secret.encode()).decode()
993
+ secrets.append(headers["Authorization"].split(" ", 1)[1])
994
+ elif isinstance(auth, ApiKeyAuth):
995
+ if auth.header is not None:
996
+ credential_header = auth.header
997
+ headers[auth.header] = secret
998
+ else:
999
+ query.append((str(auth.query), secret))
1000
+ if op.idempotency_header is not None:
1001
+ headers[op.idempotency_header] = idempotency_key or _random_key()
1002
+
1003
+ if pagination is not None and pagination.type == "link" and isinstance(upstream, str):
1004
+ url = upstream
1005
+ expected = build_url(svc.base_url, op.path, path_values, [])
1006
+ if not same_origin(url, expected) or urlsplit(url).path != urlsplit(expected).path:
1007
+ raise ToolFailure(ErrorCode.INVALID_CURSOR, "cursor does not match this operation")
1008
+ if isinstance(auth, ApiKeyAuth) and auth.query is not None:
1009
+ url += ("&" if "?" in url else "?") + urlencode({auth.query: secrets[0]})
1010
+ else:
1011
+ url = build_url(svc.base_url, op.path, path_values, query)
1012
+ check_url(url, svc.network)
1013
+
1014
+ payload: bytes | None = None
1015
+ if body or (op.method in ("POST", "PUT", "PATCH") and op.body is not None):
1016
+ if op.body_format == "json":
1017
+ payload = json.dumps(body, ensure_ascii=False).encode()
1018
+ headers["Content-Type"] = "application/json"
1019
+ else:
1020
+ payload = urlencode(list(body.items()), doseq=True).encode()
1021
+ headers["Content-Type"] = "application/x-www-form-urlencoded"
1022
+ if len(payload) > svc.limits.max_request_bytes:
1023
+ raise ToolFailure(ErrorCode.INVALID_ARGUMENTS, "request body is too large")
1024
+
1025
+ request = HttpRequest(
1026
+ method=op.method,
1027
+ url=url,
1028
+ headers=tuple(headers.items()),
1029
+ body=payload,
1030
+ network=svc.network,
1031
+ timeout=svc.limits.timeout,
1032
+ connect_timeout=svc.limits.connect_timeout,
1033
+ max_response_bytes=svc.limits.max_response_bytes,
1034
+ )
1035
+ await commit() # everything above can fail without side effects
1036
+ try:
1037
+ response = await self._send(op, request, ctx, credential_header)
1038
+ except ToolFailure as failure:
1039
+ if op.effect is not Effect.READ and failure.request_sent is not False:
1040
+ # Some attempt or redirect hop may have reached the upstream.
1041
+ raise ToolFailure(
1042
+ ErrorCode.OUTCOME_UNKNOWN,
1043
+ "the upstream did not confirm the result; the change may or may not have "
1044
+ "been applied. Check the current state before trying again.",
1045
+ retryable=False,
1046
+ ) from None
1047
+ raise
1048
+ source = {"service": svc.name, "operation": op.name, "status": response.status}
1049
+ request_id = next(
1050
+ (response.headers[h] for h in _REQUEST_ID_HEADERS if h in response.headers), None
1051
+ )
1052
+ if request_id and _SAFE_ID.match(request_id):
1053
+ source["request_id"] = request_id
1054
+
1055
+ if response.status >= 400:
1056
+ raise self._upstream_failure(reg, response, secrets, write=op.effect is not Effect.READ)
1057
+ if 300 <= response.status < 400:
1058
+ if op.effect is not Effect.READ:
1059
+ raise ToolFailure(
1060
+ ErrorCode.OUTCOME_UNKNOWN,
1061
+ f"the upstream answered with a redirect (HTTP {response.status}) after "
1062
+ "receiving the request; the change may or may not have been applied. "
1063
+ "Check the current state before trying again.",
1064
+ retryable=False,
1065
+ http_status=response.status,
1066
+ )
1067
+ raise ToolFailure(
1068
+ ErrorCode.UPSTREAM_ERROR,
1069
+ f"upstream redirected (HTTP {response.status}) and redirects are disabled",
1070
+ http_status=response.status,
1071
+ )
1072
+ if len(response.body) > _OFFLOAD_BYTES:
1073
+ decoded = await asyncio.to_thread(_decode_body, response, reg.output.spec.accept_text)
1074
+ else:
1075
+ decoded = _decode_body(response, reg.output.spec.accept_text)
1076
+ offset = int(cursor_state["o"]) if cursor_state else 0
1077
+ shaped = await self._process(
1078
+ reg, decoded, len(response.body), offset=offset, fields=fields, secrets=secrets
1079
+ )
1080
+ next_upstream = None
1081
+ if pagination is not None and shaped.next_offset is None:
1082
+ next_upstream = _next_upstream(op, response, decoded, reg, page_number, shaped)
1083
+ return shaped, source, next_upstream, response.status
1084
+
1085
+ async def _credential(self, name: str, ctx: ExecutionContext, service: str) -> str:
1086
+ if self._credentials is None:
1087
+ raise ConfigurationError(
1088
+ f"service {service} needs credential {name!r} but no credentials= provider is set"
1089
+ )
1090
+ value = self._credentials.get(name, ctx)
1091
+ if inspect.isawaitable(value):
1092
+ value = await value
1093
+ if not value or not isinstance(value, str):
1094
+ logger.error("credential %r for service %s is not available", name, service)
1095
+ raise ToolFailure(
1096
+ ErrorCode.INTERNAL_ERROR, f"credentials for service {service} are not configured"
1097
+ )
1098
+ if any(c in value for c in "\r\n\x00"):
1099
+ raise ToolFailure(ErrorCode.INTERNAL_ERROR, "configured credential is malformed")
1100
+ return value
1101
+
1102
+ async def _send(
1103
+ self,
1104
+ op: Operation,
1105
+ request: HttpRequest,
1106
+ ctx: ExecutionContext,
1107
+ credential_header: str | None,
1108
+ ) -> HttpResponse:
1109
+ svc = op.service
1110
+ retry = svc.retry
1111
+ # Retry automatically only what cannot change state twice: declared reads over
1112
+ # safe methods, or requests that carry a stable idempotency key.
1113
+ retry_allowed = (
1114
+ op.effect is Effect.READ and op.method in SAFE_METHODS
1115
+ ) or op.idempotency_header is not None
1116
+ attempts = retry.max_attempts if retry_allowed else 1
1117
+ semaphore = self._semaphore(f"{svc.name}|{svc.base_url}", svc.limits.max_concurrency)
1118
+ loop = asyncio.get_running_loop()
1119
+ deadline = loop.time() + svc.limits.total_timeout
1120
+ try:
1121
+ async with asyncio.timeout_at(deadline), semaphore:
1122
+ return await self._attempts(op, request, ctx, credential_header, attempts, deadline)
1123
+ except TimeoutError:
1124
+ raise ToolFailure(
1125
+ ErrorCode.TIMEOUT,
1126
+ f"the call exceeded its {svc.limits.total_timeout:g}s deadline",
1127
+ request_sent=None,
1128
+ ) from None
1129
+
1130
+ def _handler_pool(self) -> concurrent.futures.ThreadPoolExecutor:
1131
+ with self._sem_lock:
1132
+ if self._executor is None:
1133
+ self._executor = concurrent.futures.ThreadPoolExecutor(
1134
+ max_workers=32, thread_name_prefix="toolbound-tool"
1135
+ )
1136
+ return self._executor
1137
+
1138
+ def _semaphore(self, service: str, limit: int) -> asyncio.Semaphore:
1139
+ loop = asyncio.get_running_loop()
1140
+ with self._sem_lock:
1141
+ per_loop = self._semaphores.get(loop)
1142
+ if per_loop is None:
1143
+ per_loop = self._semaphores[loop] = {}
1144
+ semaphore = per_loop.get(service)
1145
+ if semaphore is None:
1146
+ semaphore = per_loop[service] = asyncio.Semaphore(limit)
1147
+ return semaphore
1148
+
1149
+ async def _attempts(
1150
+ self,
1151
+ op: Operation,
1152
+ request: HttpRequest,
1153
+ ctx: ExecutionContext,
1154
+ credential_header: str | None,
1155
+ attempts: int,
1156
+ deadline: float,
1157
+ ) -> HttpResponse:
1158
+ retry = op.service.retry
1159
+ loop = asyncio.get_running_loop()
1160
+ maybe_sent = False # did any earlier attempt possibly reach the upstream?
1161
+ for attempt in range(1, attempts + 1):
1162
+ try:
1163
+ response = await self._send_following_redirects(op, request, credential_header)
1164
+ except ToolFailure as failure:
1165
+ if failure.request_sent is not False:
1166
+ maybe_sent = True
1167
+ elif maybe_sent:
1168
+ failure.request_sent = None # this attempt was not sent, an earlier one was
1169
+ self._emit(
1170
+ "http.attempt", op.name, ctx, outcome=failure.code.value, attempt=attempt
1171
+ )
1172
+ if failure.code in (ErrorCode.TIMEOUT, ErrorCode.UPSTREAM_UNAVAILABLE) and (
1173
+ attempt < attempts
1174
+ ):
1175
+ delay = self._backoff(retry.backoff_base, retry.backoff_max, attempt)
1176
+ if loop.time() + delay >= deadline:
1177
+ raise
1178
+ await self._sleep(delay)
1179
+ continue
1180
+ raise
1181
+ self._emit(
1182
+ "http.attempt",
1183
+ op.name,
1184
+ ctx,
1185
+ outcome="ok",
1186
+ attempt=attempt,
1187
+ http_status=response.status,
1188
+ )
1189
+ if response.status in retry.retry_statuses and attempt < attempts:
1190
+ if response.status != 429: # a 5xx does not prove nothing was applied
1191
+ maybe_sent = True
1192
+ retry_after = _retry_after(response)
1193
+ if retry_after is not None and retry_after > retry.max_retry_after:
1194
+ return response
1195
+ delay = (
1196
+ retry_after
1197
+ if retry_after is not None
1198
+ else self._backoff(retry.backoff_base, retry.backoff_max, attempt)
1199
+ )
1200
+ if loop.time() + delay >= deadline:
1201
+ return response # no time left to wait and retry
1202
+ await self._sleep(delay)
1203
+ continue
1204
+ return response
1205
+ raise AssertionError("unreachable") # pragma: no cover
1206
+
1207
+ def _backoff(self, base: float, cap: float, attempt: int) -> float:
1208
+ return self._rng.uniform(0, min(cap, base * (2 ** (attempt - 1))))
1209
+
1210
+ async def _send_following_redirects(
1211
+ self, op: Operation, request: HttpRequest, credential_header: str | None
1212
+ ) -> HttpResponse:
1213
+ policy = op.service.network
1214
+ transport = self._get_transport()
1215
+ current = request
1216
+ hops = 0
1217
+ while True:
1218
+ try:
1219
+ response = await transport.send(current)
1220
+ except ToolFailure as failure:
1221
+ if hops > 0: # an earlier hop already delivered the request
1222
+ failure.request_sent = None
1223
+ raise
1224
+ if response.status not in _REDIRECTS or "location" not in response.headers:
1225
+ return response
1226
+ if not policy.follow_redirects:
1227
+ return response
1228
+ hops += 1
1229
+ if hops > policy.max_redirects:
1230
+ raise ToolFailure(
1231
+ ErrorCode.UPSTREAM_ERROR, "too many redirects", http_status=response.status
1232
+ )
1233
+ target = urljoin(current.url, response.headers["location"])
1234
+ try:
1235
+ check_url(target, policy)
1236
+ except ToolFailure as failure:
1237
+ failure.request_sent = None # the first hop was delivered
1238
+ raise
1239
+ headers = list(current.headers)
1240
+ if not same_origin(current.url, target):
1241
+ drop = {"authorization", "cookie"}
1242
+ if credential_header:
1243
+ drop.add(credential_header.lower())
1244
+ # Query-string API keys are not copied either: Location is used verbatim.
1245
+ headers = [(k, v) for k, v in headers if k.lower() not in drop]
1246
+ method, body = current.method, current.body
1247
+ if response.status == 303 or (response.status in (301, 302) and method == "POST"):
1248
+ method, body = "GET", None
1249
+ headers = [(k, v) for k, v in headers if k.lower() != "content-type"]
1250
+ current = replace(current, url=target, headers=tuple(headers), method=method, body=body)
1251
+
1252
+ def _upstream_failure(
1253
+ self,
1254
+ reg: _Registered,
1255
+ response: HttpResponse,
1256
+ secrets: Sequence[str],
1257
+ *,
1258
+ write: bool = False,
1259
+ ) -> ToolFailure:
1260
+ status = response.status
1261
+ details: dict[str, object] = {"http_status": status}
1262
+ try:
1263
+ body = json.loads(response.body) if response.body else None
1264
+ except (ValueError, UnicodeDecodeError):
1265
+ body = None
1266
+ shaped = shape_error_body(reg.output, body, secrets=secrets)
1267
+ if shaped:
1268
+ details["upstream"] = shaped
1269
+ if status == 429:
1270
+ return ToolFailure(
1271
+ ErrorCode.RATE_LIMITED,
1272
+ "upstream rate limit reached (HTTP 429)",
1273
+ http_status=status,
1274
+ details=details,
1275
+ retry_after=_retry_after(response),
1276
+ )
1277
+ if status in (401, 403):
1278
+ message = f"upstream refused the request (HTTP {status})"
1279
+ elif status == 404:
1280
+ message = "upstream resource not found (HTTP 404)"
1281
+ elif status >= 500:
1282
+ if write:
1283
+ # A server error after a write does not prove nothing changed.
1284
+ return ToolFailure(
1285
+ ErrorCode.UPSTREAM_ERROR,
1286
+ f"upstream service error (HTTP {status}); the change may or may not have "
1287
+ "been applied. Check the current state before trying again.",
1288
+ retryable=False,
1289
+ http_status=status,
1290
+ details=details,
1291
+ )
1292
+ return ToolFailure(
1293
+ ErrorCode.UPSTREAM_ERROR,
1294
+ f"upstream service error (HTTP {status})",
1295
+ retryable=True,
1296
+ http_status=status,
1297
+ details=details,
1298
+ )
1299
+ else:
1300
+ message = f"upstream rejected the request (HTTP {status})"
1301
+ return ToolFailure(
1302
+ ErrorCode.UPSTREAM_ERROR, message, retryable=False, http_status=status, details=details
1303
+ )
1304
+
1305
+ # ------------------------------------------------------------------ events
1306
+
1307
+ def _emit(self, name: str, tool: str, ctx: ExecutionContext, **fields: Any) -> None:
1308
+ try:
1309
+ self._observer.on_event(
1310
+ Event(
1311
+ name=name,
1312
+ tool=tool,
1313
+ actor=ctx.actor,
1314
+ request_id=ctx.request_id,
1315
+ tenant=ctx.tenant,
1316
+ **fields,
1317
+ )
1318
+ )
1319
+ except Exception:
1320
+ logger.exception("observer failed")
1321
+
1322
+
1323
+ # ---------------------------------------------------------------------- helpers
1324
+
1325
+
1326
+ def _fingerprint_source(spec: Tool, schema: Mapping[str, Any]) -> dict[str, Any]:
1327
+ source: dict[str, Any] = {
1328
+ "name": spec.name,
1329
+ "effect": spec.effect.value,
1330
+ "schema": schema,
1331
+ "bind": sorted(spec.bind),
1332
+ "scopes": sorted(spec.required_scopes),
1333
+ }
1334
+ if isinstance(spec, Operation):
1335
+ source.update(
1336
+ kind="rest",
1337
+ method=spec.method,
1338
+ path=spec.path,
1339
+ base_url=spec.service.base_url,
1340
+ service=spec.service.name,
1341
+ query=spec.query,
1342
+ body=spec.body,
1343
+ )
1344
+ else:
1345
+ source.update(kind="python", handler=getattr(spec.handler, "__qualname__", "?"))
1346
+ return source
1347
+
1348
+
1349
+ def _placement(op: Operation, remaining: Mapping[str, Any]) -> tuple[list[str], list[str]]:
1350
+ names = list(remaining)
1351
+ if op.query is not None and op.body is not None:
1352
+ return list(op.query), list(op.body)
1353
+ body_method = op.method in ("POST", "PUT", "PATCH")
1354
+ if op.query is not None:
1355
+ rest = [n for n in names if n not in op.query]
1356
+ return list(op.query), rest if body_method else []
1357
+ if op.body is not None:
1358
+ rest = [n for n in names if n not in op.body]
1359
+ return rest, list(op.body)
1360
+ return ([], names) if body_method else (names, [])
1361
+
1362
+
1363
+ def _random_key() -> str:
1364
+ return secrets.token_hex(16)
1365
+
1366
+
1367
+ def _retry_after(response: HttpResponse) -> float | None:
1368
+ value = response.headers.get("retry-after")
1369
+ if not value:
1370
+ return None
1371
+ try:
1372
+ return max(0.0, float(value))
1373
+ except ValueError:
1374
+ pass
1375
+ try:
1376
+ from email.utils import parsedate_to_datetime
1377
+
1378
+ when = parsedate_to_datetime(value)
1379
+ return max(0.0, when.timestamp() - time.time())
1380
+ except (TypeError, ValueError):
1381
+ return None
1382
+
1383
+
1384
+ def _decode_body(response: HttpResponse, accept_text: bool) -> Any:
1385
+ if response.status in (204, 205) or not response.body:
1386
+ return None
1387
+ content_type = response.headers.get("content-type", "").split(";")[0].strip().lower()
1388
+ if _JSON_TYPE.match(content_type) or (not content_type and response.body[:1] in (b"{", b"[")):
1389
+ try:
1390
+ return json.loads(response.body)
1391
+ except (ValueError, UnicodeDecodeError):
1392
+ raise ToolFailure(
1393
+ ErrorCode.UPSTREAM_SCHEMA_MISMATCH,
1394
+ "upstream returned malformed JSON",
1395
+ retryable=False,
1396
+ ) from None
1397
+ if accept_text and content_type.startswith("text/"):
1398
+ charset = "utf-8"
1399
+ for part in response.headers.get("content-type", "").split(";")[1:]:
1400
+ key, _, val = part.strip().partition("=")
1401
+ if key.lower() == "charset" and val:
1402
+ charset = val.strip('"')
1403
+ try:
1404
+ return response.body.decode(charset, errors="replace")
1405
+ except LookupError:
1406
+ return response.body.decode("utf-8", errors="replace")
1407
+ raise ToolFailure(
1408
+ ErrorCode.UNSUPPORTED_CONTENT_TYPE,
1409
+ f"upstream returned unsupported content type {content_type or 'unknown'!r}",
1410
+ retryable=False,
1411
+ )
1412
+
1413
+
1414
+ _LINK_NEXT = re.compile(r'<([^>]+)>\s*;[^,]*rel="?next"?', re.IGNORECASE)
1415
+
1416
+
1417
+ def _next_upstream(
1418
+ op: Operation,
1419
+ response: HttpResponse,
1420
+ decoded: Any,
1421
+ reg: _Registered,
1422
+ page_number: int | None,
1423
+ shaped: ShapeResult,
1424
+ ) -> Any:
1425
+ pagination = op.pagination
1426
+ assert pagination is not None
1427
+ if pagination.type == "cursor":
1428
+ value = parse_path(str(pagination.next_path)).get(decoded)
1429
+ if value is MISSING or value in (None, "", False):
1430
+ return None
1431
+ if not isinstance(value, str | int) or len(str(value)) > 2048:
1432
+ return None
1433
+ return value
1434
+ if pagination.type == "link":
1435
+ match = _LINK_NEXT.search(response.headers.get("link", ""))
1436
+ if not match:
1437
+ return None
1438
+ target = urljoin(response.url, match.group(1))
1439
+ try:
1440
+ check_url(target, op.service.network)
1441
+ except NetworkBlocked:
1442
+ return None
1443
+ if (
1444
+ not same_origin(target, response.url)
1445
+ or urlsplit(target).path != urlsplit(response.url).path
1446
+ ):
1447
+ return None
1448
+ if isinstance(op.service.auth, ApiKeyAuth) and op.service.auth.query:
1449
+ # never store a credential inside a cursor
1450
+ parts = urlsplit(target)
1451
+ from urllib.parse import parse_qsl
1452
+
1453
+ kept = [(k, v) for k, v in parse_qsl(parts.query) if k != op.service.auth.query]
1454
+ target = parts._replace(query=urlencode(kept)).geturl()
1455
+ return target
1456
+ count = shaped.total_items
1457
+ if count is None or count == 0 or page_number is None:
1458
+ return None
1459
+ if pagination.page_size is not None and count < pagination.page_size:
1460
+ return None
1461
+ return page_number + 1