Validate task tool arguments against their declared types#4355
Validate task tool arguments against their declared types#4355Michael-WhiteCapData wants to merge 1 commit into
Conversation
FunctionTool.register_with_docket registered the raw function and let Docket bind arguments by signature. Docket does not coerce argument values, so a parameter typed as a Pydantic model reached the function as a raw dict under a task call, while a synchronous call validates it through FunctionTool.run -> TypeAdapter.validate_python. Accessing a model attribute then raised AttributeError under task=True only. Register a wrapper that keeps the raw function's signature (so Docket still resolves FastMCP and Docket-native dependencies) and coerces the client-facing arguments via cached TypeAdapters before invoking it, mirroring the synchronous path. Injected dependency parameters are excluded from coercion. Task tools are always async, so the wrapper preserves the coroutine / async-generator nature of the function. Fixes PrefectHQ#4349
|
This PR has been automatically closed because you are not assigned to the issue it references. Per CONTRIBUTING.md, an external PR must reference an issue that is assigned to its author. To proceed:
Maintainers: reopen this PR or remove the |
marvin-context-protocol
Bot
added
bug
|
Thanks for the report. This issue goes beyond what our contributor guidelines ask for — we just need a short problem description and an MRE. Please see our contributing guidelines and condense this issue. We'll triage it once it's trimmed down. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: eaaf5f788e
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| for name in inspect.signature(fn).parameters | ||
| if name not in exclude and name in hints |
There was a problem hiding this comment.
Existing exclude_args support intentionally allows hidden parameters with non-serializable types (for example ServerSession) because the schema path removes them before creating Pydantic adapters. This new task registration iterates the original signature and builds adapters for every annotated non-DI parameter, so a task-enabled tool with exclude_args=["session"] and another typed client arg will construct a TypeAdapter for the hidden session parameter during server startup and fail instead of registering. The excluded names need to be carried here or the same pruned wrapper used by schema generation should be reused.
Useful? React with 👍 / 👎.
|
|
||
| hints = get_type_hints(fn, include_extras=True) | ||
| adapters = { | ||
| name: get_cached_typeadapter(hints[name]) |
There was a problem hiding this comment.
FastMCP supports constraints supplied as a Pydantic Field default (x: int = Field(ge=1)) and the synchronous path enforces them by validating the whole wrapper function. Here the adapter is built from get_type_hints, which returns only int for that parameter, so task calls validate/coerce the type but drop ge, aliases, and other Field metadata; for example, task=True with {"x": 0} runs while the normal call rejects it. Build from the function signature/field metadata rather than the bare hint.
Useful? React with 👍 / 👎.
1 participant
Fixes #4349.
Problem
A tool parameter typed as a Pydantic model is validated into a model instance on a synchronous call, but arrives as a raw
dictwhen the same tool runs as a task. Accessing a model attribute (e.g.items[0].value) therefore raisesAttributeErrorundertask=Trueonly.Cause
FunctionTool.register_with_docketregisters the raw function, and Docket binds arguments by signature without coercing their values. The synchronous path validates arguments throughFunctionTool.run→TypeAdapter.validate_python; the task path skips that step, so model-typed parameters stay dicts.Fix
Register a thin wrapper that:
functools.wraps, so Docket still resolves all dependencies — FastMCP's (Context, Progress) and Docket-native (Retry, Timeout, ConcurrencyLimit); andTypeAdapters before invoking the function, mirroring the synchronous path. Injected dependency parameters are excluded from coercion.Task tools are always async (sync functions are rejected at registration), so the wrapper preserves the coroutine / async-generator nature of the wrapped function and leaves Docket's execution path unchanged.
Tests
Added a regression test in
tests/server/tasks/test_task_tools.pycovering both the sync and task paths. The fulltests/server/tasks/andtests/tools/suites pass (663 passed, 1 skipped);ruff checkandruff formatare clean on the changed files.