Decode base64 for bytes-typed HookToolset parameters

[ColtenOuO]

Summary

JSON has no binary type, so HookToolset advertises a hook method's bytes-typed
parameters to the model as strings — and nothing converted them back. A method such
as S3Hook.load_bytes(bytes_data: bytes, ...) received a plain str and failed deep
inside BytesIO with an opaque TypeError: a bytes-like object is required, not 'str',
which gave no hint about where the bad value came from.

What that string was meant to contain was never stated either. Models routinely emit
base64 for binary payloads, so reading it as UTF-8 would have written the base64 text
to storage: a successful upload of the wrong bytes. The contract is now explicit.

Changes

• bytes-typed parameters are advertised with contentEncoding: "base64", and the
  same instruction is appended to the parameter description — most function-calling
  APIs drop the keyword, and the description is the field the model is guaranteed to see.
• call_tool decodes those parameters with base64.b64decode(value, validate=True).
  A value that is not valid base64 is raised as ModelRetry, the only exception
  pydantic-ai feeds back to the model, so it can correct the call instead of failing
  the run. Decoding stays in the innermost toolset so bytes never reach
  CachingToolset, which fingerprints tool arguments with a plain json.dumps.
• Annotations are resolved per parameter. get_type_hints is all-or-nothing, so one
  TYPE_CHECKING-only import elsewhere in a signature — CloudKMSHook.encrypt has
  exactly that — would otherwise turn the conversion off for the whole method silently.
  As a side effect that method's parameters stop being advertised as untyped.
• The JSON Schema and the set of bytes parameters come from a single introspection
  pass at build time, driven by one condition, so a parameter can never be advertised
  as base64 without also being decoded.

Only bytes and bytes | None are affected. A parameter that already accepts text,
such as LambdaHook.invoke_lambda's payload: bytes | str | None, is passed through
untouched.

Still open

The return direction has the same problem — _serialize_for_llm hands bytes to
json.dumps(..., default=str), so the model receives a Python repr. A proposal is
under discussion in the review thread and will follow once the shape is agreed.

---

Was generative AI tooling used to co-author this PR?

• Yes — Claude Code (Sonnet 5)

[ColtenOuO]

It looks like the CI needs to be restarted.

[potiuk]

Thanks — the bug is real and well diagnosed. _TYPE_MAP maps bytes to {"type": "string"} for the tool schema, nothing converted back, so a hook method like S3Hook.load_bytes failed deep inside BytesIO with a TypeError that gave no hint about where the bad value came from.

The implementation handles the cases I'd want: Optional[bytes] and bytes | None both resolve, real bytes values pass through untouched, and the tests fail without the fix since the fake hook reports the received type.

One thing to settle before this goes in: what the string actually contains.

The schema advertises {"type": "string"} with no contentEncoding, so the contract with the model is undefined. Models routinely emit base64 for binary payloads. If one does that here, the base64 text gets UTF-8 encoded and written to storage — wrong bytes, no error, upload reported as successful. That swaps a loud crash for silent data corruption, which is the worse of the two.

Worth making the contract explicit rather than leaving it implied. Two directions, and I don't think this should be decided in a review thread alone — it sets the behaviour for every bytes parameter the toolset will ever expose:

1. Declare the encoding in the schema and follow it:

   bytes: {"type": "string", "contentEncoding": "base64"},

   with a base64 decode in call_tool. This is the standard JSON Schema mechanism, it supports genuinely binary payloads, and plain text sent by mistake fails loudly on decode instead of silently storing the wrong thing.

2. Keep UTF-8 and say so in the schema, via a description on bytes-typed parameters telling the model to send plain text. Smaller change, but it caps these tools at text-only payloads.

Either is defensible; leaving it unstated is the option I'd avoid. Input from whoever owns common.ai would be useful here, since this is a provider-wide behaviour decision rather than a detail of this PR.

Two smaller points:

except Exception around get_type_hints is broader than the project prefers — NameError (unresolvable forward reference) and TypeError are the realistic failures. It matters more than usual here: the fallback to param.annotation under from __future__ import annotations yields the string "bytes", which _resolves_to_bytes returns False for, so coercion silently stops happening rather than failing. Worth narrowing the except and leaving a comment about that fallback.

_bytes_param_names re-runs inspect.signature and get_type_hints on every call_tool. _build_json_schema_from_signature already does both when the tools are built, so the parameter names could be resolved once there and reused per call.

---

Drafted-by: Claude Code (Opus 5); reviewed by @potiuk before posting

[ColtenOuO]

Thanks! agreed on base64, since it's the only option that can carry real binary payloads without ever needing revisiting.

bytes: {"type": "string", "contentEncoding": "base64"},

Decoding with base64.b64decode(value, validate=True), validate=True matters because the default silently discards invalid characters instead of erroring, which is the same silent-corruption problem this PR is meant to fix. I'll re-raise as a ValueError naming the parameter so the model gets a usable retry message. No UTF-8 fallback on decode failure — that would put us back to an undefined contract.

On the two smaller points: I'll narrow to except (NameError, TypeError) with a comment on the from __future__ import annotations fallback issue, and have _build_json_schema_from_signature return the bytes param names once, cached per tool, instead of re-introspecting on every call_tool.

Happy to hold for common.ai maintainer input given this is provider-wide.

[kaxil]

+1 on the base64 contract agreed above. One constraint for the rework: the decode has to stay here in call_tool rather than move into build_args_validator, even though #70096 makes the validator look like the natural home. Validated args flow through the whole toolset chain, and CachingToolset fingerprints tool_args before delegating (_digest is json.dumps with no default=), so bytes values there would turn every bytes-param call into an unverifiable fingerprint under durable=True. Building a new dict inside the innermost toolset, as this PR already does, keeps bytes out of that path, so worth a short code comment to make the placement survive the rework. Since contentEncoding is ignored by most function-calling APIs, the base64 instruction should also be appended to each parameter's schema description after the docstring enrichment loop in get_tools (that loop overwrites descriptions, and e.g. S3Hook.load_bytes's ":param bytes_data: bytes to set as content for the key" would otherwise be all the model sees), plus a sentence in toolsets.rst documenting the contract.

[kaxil]

wdyt @gopidesupavan , what's your opinion?

[gopidesupavan]

Agree with all your points. I checked _digest, it is json.dumps without default=, so bytes in tool_args will fingerprint as None and those calls lose replay verification with durable=True. So keeping the decode in call_tool makes sense.

one more thing from my side.

Output side has the same problem. _serialize_for_llm passes bytes to json.dumps(..., default=str), so the model gets the Python repr:

json.dumps(b'\x89PNG', default=str)  # -> "b'\\x89PNG'"

GCSHook.download returns bytes and our docs recommend exposing such read-only methods. So download -> load_bytes through an agent already writes wrong data, and base64 only on the input side does not fix it. I think we should define both directions in this PR. @ColtenOuO FYI

[ColtenOuO]

Confirmed, and agreed both directions belong here. Proposal below.

Same wire format both ways, so a round trip is a pass-through and the model never transforms anything:

def _json_default(value: Any) -> str:
    if isinstance(value, (bytes, bytearray, memoryview)):
        return _b64(value)
    return str(value)


def _serialize_for_llm(value: Any) -> str:
    if value is None:
        return "null"
    if isinstance(value, str):
        return value
    if isinstance(value, (bytes, bytearray, memoryview)):
        return _b64(value)
    try:
        return json.dumps(value, default=_json_default)
    except (TypeError, ValueError):
        return str(value)

Top-level bytes bypass json.dumps, mirroring the existing str early return — no surrounding quotes, so the value goes straight back into a bytes parameter. Nested bytes are handled by swapping default=, leaving str() as the fallback for datetimes and friends.

Correctness comes from the runtime isinstance, not from annotations, and deliberately so: download is overloaded on filename, CloudKMSHook.decrypt yields no hints at all, and nested dict[str, bytes] gives no signal either. Annotations only decide whether the model gets told — " Binary output is returned base64-encoded." appended to the description when the return annotation mentions bytes. That check is looser than the argument side's _resolves_to_bytes, which is safe here because it governs a sentence of documentation rather than how data is handled; the worst case is a missing hint, never a wrong payload.

Not adding contentEncoding to return_schema: download returns a path string when filename is passed, so the declaration would be wrong half the time. The argument side can declare it because parameter types are static.

toolsets.rst gets both directions plus the limitation that bytes nested inside a JSON structure are indistinguishable from ordinary strings once encoded.

[ColtenOuO]

Thanks all — reworked and rebased onto current main (the branch predated #70096, which had rewritten the same region of hook.py). Title and description updated to match, since the change is no longer a UTF-8 coercion.

Point by point:

Review point	How it was handled
Undefined string contract	_TYPE_MAP[bytes] now carries contentEncoding: "base64"; call_tool decodes with b64decode(value, validate=True)
contentEncoding ignored by most APIs	The instruction is appended to each parameter's description, after the docstring enrichment loop that would otherwise overwrite it
Document the contract	New "Binary parameters" section in toolsets.rst
Keep decoding out of the args validator	Unchanged placement, with a comment at the site recording the CachingToolset fingerprint reason so it survives the next refactor
except Exception too broad	Gone entirely, replaced by the per-parameter resolution below
Per-parameter annotation resolution	_resolve_annotations substitutes unresolvable names with Any and retries, so one bad name no longer discards the whole signature
Re-introspecting on every call	Resolved once in get_tools and stored per tool; no lru_cache on bound methods
VAR_POSITIONAL / VAR_KEYWORD drift	The two passes are now one, and the skip happens before the bytes check, so the drift is structurally impossible
Test duplication, untested union guard	The three duplicated cases are gone; a single parametrised test covers bytes, bytes | None, bytes | str, list[bytes], str and an unannotated parameter

One correction to what I said earlier in this thread. I wrote that decode failures would be re-raised as ValueError so the model could retry. That was wrong, and @kaxil is right: pydantic-ai only turns ModelRetry (and ValidationError from the args-validation stage) into a retry prompt — a ValueError raised inside call_tool fails the run without the model ever seeing it. The code raises ModelRetry. The UnicodeEncodeError case in the same comment disappears with the UTF-8 encode itself.

Two things that surfaced while implementing this:

A leak I introduced and caught before pushing. Putting contentEncoding in _TYPE_MAP made it recurse into every union and list branch, so bytes | str and list[bytes] | None advertised base64 while _resolves_to_bytes — correctly — excluded them from decoding. LambdaHook.invoke_lambda and LevelDBHook.run both hit that shape, and a model trusting the schema would have handed the hook the encoded text: exactly the corruption this PR exists to prevent. contentEncoding is now applied in the same if that populates the decode set, and the parametrised test asserts both sides agree for every case.

The annotation fix reaches further than the bytes parameters. _build_json_schema_from_signature shared the all-or-nothing get_type_hints call, so CloudKMSHook.encrypt was advertising every parameter as untyped, not just failing to decode plaintext. That is fixed by the same change.

Deliberately left alone for now:

• The return direction. Agreed it belongs here; I've replied with a concrete proposal in that thread rather than implementing ahead of the discussion, since it turns on a question the argument side doesn't have (return types can't be resolved statically — GCSHook.download is overloaded on filename).
• list[bytes]. LevelDBHook.run decodes key/value but not keys/values in the same signature. Not a regression, and the parametrised test pins the current behaviour so it can't drift silently, but worth folding in if you'd rather the coercion recursed into containers.

---

Drafted-by: Claude Code (Opus 5); reviewed by @ColtenOuO before posting

[ColtenOuO]

Apologies for the slow reply :(

I spent quite a bit of time reviewing the feedback and looking into the related issues today. Thank you all for the thoughtful discussion and feedback!