Wrapping

wrapping

Tool-call wrapping, tracing, and per-key locking runtime.

Builds the synchronous and async wrapper closures handed to MCPServer for each discovered tool: kwarg unwrapping, implicit-path warnings, external session tracing, ToolResult flattening, and per-key concurrency locking (active only in HTTP mode).

This module is a leaf — it imports only axm_mcp.concurrency at runtime. Shared structural protocols (ToolEntry, ToolLike, PlainTool, ToolResultLike) live in axm_mcp.discovery and are referenced here under TYPE_CHECKING (annotations are strings via from __future__ import annotations) plus string-literal cast targets, keeping the runtime import edge one-directional (discovery -> wrapping).

build_wrappers(name, tool, *, shared_mode=False, write_contract_resolver=None, preserve_result=False)

Build the (sync, async) wrapper pair for one tool.

The single construction seam shared by the direct MCP registration path (:func:axm_mcp.discovery.register_one) and the facade path (:class:axm_mcp.facade.catalog.ToolCatalog). Both callers pass the same shared-mode flag and per-request write-contract resolver, so write-scope enforcement cannot diverge between exposure routes. In dedicated mode, omitting the resolver preserves the environment-backed fallback. The same wrappers also keep kwarg unwrapping, tracing, flattening and locking aligned.

Returns:

Type	Description
_SyncWrapper	(sync_wrapper, async_wrapper) where the sync wrapper carries the
_AnyWrapper	trace/flatten/exception contract and the async wrapper adds the HTTP
tuple[_SyncWrapper, _AnyWrapper]	to_thread offload plus the optional per-key lock.

Source code in packages/axm-mcp/src/axm_mcp/wrapping.py

def build_wrappers(
    name: str,
    tool: ToolEntry,
    *,
    shared_mode: bool = False,
    write_contract_resolver: Callable[[], WriteContract | None] | None = None,
    preserve_result: bool = False,
) -> tuple[_SyncWrapper, _AnyWrapper]:
    """Build the ``(sync, async)`` wrapper pair for one tool.

    The single construction seam shared by the direct MCP registration path
    (:func:`axm_mcp.discovery.register_one`) and the facade path
    (:class:`axm_mcp.facade.catalog.ToolCatalog`). Both callers pass the same
    shared-mode flag and per-request write-contract resolver, so write-scope
    enforcement cannot diverge between exposure routes. In dedicated mode,
    omitting the resolver preserves the environment-backed fallback. The same
    wrappers also keep kwarg unwrapping, tracing, flattening and locking aligned.

    Returns:
        ``(sync_wrapper, async_wrapper)`` where the sync wrapper carries the
        trace/flatten/exception contract and the async wrapper adds the HTTP
        ``to_thread`` offload plus the optional per-key lock.
    """
    resolver: Callable[[], WriteContract | None]
    if write_contract_resolver is None:
        from axm.tools.write_scope import write_contract_from_env

        write_contract_from_env()
        resolver = write_contract_from_env
    else:
        resolver = write_contract_resolver
    is_plain = callable(tool) and not hasattr(tool, "execute")
    # Protocol tools already trace via orchestrator.run_tool()
    ctx = _WrapperCtx(
        name=name,
        should_trace=not name.startswith("protocol_"),
        write_contract_resolver=resolver,
        shared_mode=shared_mode,
        path_defaults=path_defaults(getattr(tool, "execute", tool)),
        preserve_result=preserve_result,
    )
    base_sync_wrapper = (
        _build_plain_wrapper(ctx, tool) if is_plain else _build_tool_wrapper(ctx, tool)
    )

    sync_wrapper = base_sync_wrapper

    exec_doc = getattr(getattr(tool, "execute", tool), "__doc__", None)
    sync_wrapper.__doc__ = exec_doc or f"Execute {name} tool."
    # Resolve before lock selection too: "." and the equivalent absolute path
    # must serialize on the same repository key. Execution consumes the prepared
    # payload without looking up the session or resolving it again.
    execution = (
        _build_plain_wrapper(ctx, tool, prepared=True)
        if is_plain
        else _build_tool_wrapper(ctx, tool, prepared=True)
    )
    locked = _wrap_with_lock(execution, name)

    async def async_wrapper(**kwargs: object) -> dict[str, object] | str:
        refusal = _write_refusal(ctx, kwargs)
        if refusal is not None:
            return refusal
        return await locked(**kwargs)

    async_wrapper.__doc__ = sync_wrapper.__doc__
    return sync_wrapper, async_wrapper

flatten_result(result)

Flatten a ToolResult into a JSON-friendly dict.

Spreads result.data first, then sets the envelope keys (success/error/hint) deterministically. Any reserved key already present in result.data is relocated to data_{key} (with a warning) so the envelope is never clobbered and the data value is never silently lost.

Source code in packages/axm-mcp/src/axm_mcp/wrapping.py

def flatten_result(result: ToolResultLike) -> dict[str, object]:
    """Flatten a ToolResult into a JSON-friendly dict.

    Spreads ``result.data`` first, then sets the envelope keys
    (``success``/``error``/``hint``) deterministically. Any reserved key
    already present in ``result.data`` is relocated to ``data_{key}`` (with a
    warning) so the envelope is never clobbered and the data value is never
    silently lost.
    """
    output: dict[str, object] = dict(getattr(result, "data", None) or {})
    for key in _RESERVED_KEYS:
        if key in output:
            namespaced = f"data_{key}"
            logger.warning(
                "ToolResult.data key %r collides with the envelope; relocating to %r",
                key,
                namespaced,
            )
            output[namespaced] = output.pop(key)
    # ``success`` missing → False (never silently promoted to a passing result).
    success = getattr(result, "success", False)
    output["success"] = success is True
    if getattr(result, "error", None):
        output["error"] = result.error
    if not isinstance(success, bool):
        output["error"] = "malformed ToolResult: success must be a boolean" + (
            f"; {output['error']}" if output.get("error") else ""
        )
    hint = getattr(result, "hint", None)
    if hint:
        output["hint"] = hint
    return output

log_external_step(tool_name, tool_args, success, result_str, duration_ms)

Instrumentation seam for non-protocol tool calls.

Currently a no-op. This is the hook point where an execution engine can observe each MCP tool call (name, args, outcome, duration). The legacy axm-engine tracing wiring was removed when engine was deprecated; a future axm-loom-based tracer should re-attach here. Any implementation MUST swallow its own errors — tracing must never break tool execution.

Source code in packages/axm-mcp/src/axm_mcp/wrapping.py

def log_external_step(
    tool_name: str,
    tool_args: dict[str, object],
    success: bool,
    result_str: str,
    duration_ms: int,
) -> None:
    """Instrumentation seam for non-protocol tool calls.

    Currently a no-op. This is the hook point where an execution engine
    can observe each MCP tool call (name, args, outcome, duration). The
    legacy ``axm-engine`` tracing wiring was removed when engine was
    deprecated; a future ``axm-loom``-based tracer should re-attach here.
    Any implementation MUST swallow its own errors — tracing must never
    break tool execution.
    """