Source code for tools.aws.helpers

"""Generic boto3 invoke helper for tools/aws handlers."""

from __future__ import annotations

import asyncio
import jsonutil as json
from typing import Any

from tools.aws.base import _get_session, _dumps, _err, strip_rm
from tools.manage_api_keys import missing_api_key_error


[docs] async def aws_method( ctx, service: str, method: str, *, region: str | None = None, request_json: str | None = None, ) -> str: """Invoke an arbitrary boto3 ``client(service).method(**kwargs)`` call. Generic escape hatch that lets any AWS tool reach an arbitrary boto3 operation: it resolves the caller's credentials, decodes the optional ``request_json`` payload into keyword arguments, runs the blocking boto3 call off the event loop, and serialises the result. This is why every per-service JSON tool (ec2, s3, lambda, iam, etc.) can be defined as a thin wrapper rather than hand-coding each API method. Resolves a boto3 ``Session`` via :func:`tools.aws.base._get_session` (which reads the user's stored AWS key from Redis or falls back to ``AWS_*`` env vars), parses ``request_json`` with :func:`jsonutil.loads`, and executes the nested ``_run`` closure inside :func:`asyncio.to_thread` so the network call never stalls the loop. Results are serialised with :func:`tools.aws.base._dumps` and failures wrapped by :func:`tools.aws.base._err`; missing credentials surface :func:`tools.manage_api_keys.missing_api_key_error`. Any AWS-side mutation depends entirely on which ``method`` is requested -- it has no Redis, event-bus, KG, or LLM interaction of its own. Called directly by the per-service ``sts`` handlers in :mod:`tools.aws.sts` and by the ``_handler`` closure that :func:`wrap_aws_method` returns for every other AWS service module; also exercised directly in ``tests/test_aws_tools.py``. Args: ctx: The tool ``ToolContext`` (or ``None``) used to resolve AWS credentials. service: The boto3 service name to build a client for (e.g. ``ec2``). method: The client method to invoke on that service (e.g. ``describe_instances``). region: Optional AWS region name; the session default is used when ``None``. request_json: Optional JSON string decoded into the keyword arguments passed to the boto3 method. Returns: str: The JSON-serialized boto3 response on success, or a ``{"error": ...}`` JSON string when credentials are missing, the JSON is invalid, or the API call raises. """ sess = await _get_session(ctx) if not sess: return _err(missing_api_key_error("aws")) try: kwargs: dict[str, Any] = {} if request_json: kwargs = json.loads(request_json) def _run(): """Build the boto3 client and run the requested API method synchronously. Closure executed on a worker thread by the enclosing ``aws_method`` via ``asyncio.to_thread`` so the blocking boto3 network call never stalls the event loop. It instantiates a client for ``service`` on the parent's ``sess`` boto3 ``Session`` (passing ``region_name`` only when ``region`` is set), looks up ``method`` on that client with ``getattr``, and invokes it with the JSON-derived ``kwargs`` from the enclosing scope. Dict responses are passed through :func:`tools.aws.base.strip_rm` to drop the boto3 ``ResponseMetadata`` envelope before returning. Performs a live AWS API call (any AWS-side mutation depends on which ``method`` was requested); it has no Redis, event-bus, KG, or LLM interaction. Called only by the enclosing ``aws_method`` coroutine via ``asyncio.to_thread``; no other internal callers exist. Returns: The boto3 response with ``ResponseMetadata`` stripped when it is a dict, otherwise the raw response object unchanged. """ c = ( sess.client(service, region_name=region) if region else sess.client(service) ) fn = getattr(c, method) resp = fn(**kwargs) return strip_rm(resp) if isinstance(resp, dict) else resp return _dumps(await asyncio.to_thread(_run)) except json.JSONDecodeError as e: return _err(f"Invalid JSON: {e}") except Exception as e: return _err(str(e))
[docs] def wrap_aws_method(service: str, method: str): """Build an async tool handler bound to one boto3 ``service`` and ``method``. Factory that captures a service/method pair and returns an async ``_handler`` closure suitable for use as the ``handler`` of an AWS tool definition. This is what lets each ``tools/aws/*`` module declare dozens of JSON-driven operations declaratively (a list of ``(name, method, description)`` tuples) instead of writing a coroutine per API call; the returned closure simply forwards to :func:`aws_method` with the captured names. Pure closure construction with no side effects of its own -- it touches no Redis, event bus, KG, or LLM; all credential resolution, JSON parsing, the threaded boto3 call, and result/error serialisation happen later inside :func:`aws_method` when the returned handler runs. Called by the per-service tool builders across the ``tools/aws/*`` modules (ec2, s3, lambda, iam, rds, ecs, route53, cloudwatch, sqs, sns, ecr, elbv2, dynamodb, logs, secretsmanager) to populate their tool lists -- e.g. ``wrap_aws_method("ec2", meth)`` inside :func:`tools.aws.ec2._ec2_json_tools`. Args: service: The boto3 service name the produced handler will target (e.g. ``ec2``). method: The client method the produced handler will invoke (e.g. ``run_instances``). Returns: An async ``_handler(request_json=None, region=None, ctx=None) -> str`` coroutine function that delegates to :func:`aws_method`. """ async def _handler( request_json: str | None = None, region: str | None = None, ctx=None, ) -> str: """Tool handler that forwards a JSON request to a boto3 ``service.method``. Closure returned by ``wrap_aws_method`` and stored as the ``handler`` for a per-service AWS tool definition. It simply delegates to the module-level :func:`aws_method` coroutine, binding the ``service`` and ``method`` names captured from the enclosing scope and passing through the caller-supplied ``request_json``, ``region``, and tool ``ctx``. All credential resolution (via :func:`tools.aws.base._get_session`), JSON parsing, the threaded boto3 call, and result/error serialization happen inside ``aws_method``; this wrapper adds no side effects of its own and touches no Redis, event bus, KG, or LLM. Instances of this closure are invoked by the tool runtime by the ``handler`` key of the AWS tool definitions assembled across the ``tools/aws/*`` modules (elbv2, lambda, logs, s3, ec2, etc.) and folded into ``tools/aws_tools.py``\\ 's aggregated, privilege-gated ``TOOLS``. No direct internal callers were found. Args: request_json: Optional JSON string of keyword arguments forwarded verbatim to the boto3 ``service.method`` call. region: Optional AWS region name; the session default is used when ``None``. ctx: The tool ``ToolContext`` used to resolve AWS credentials. Returns: str: The JSON-serialized boto3 response on success, or a ``{"error": ...}`` JSON string when credentials are missing, the JSON is invalid, or the API call raises. """ return await aws_method( ctx, service, method, region=region, request_json=request_json ) return _handler