kaxil opened a new pull request, #73993: URL: https://github.com/apache/airflow/pull/73993
Stacked on #73984 (first-class `capabilities=`); the first commit here is that PR's. Review the second commit. Common AI passes what a tool returns, and the error text it hands the model, through Airflow's secret masker, so a connection password in a database error or a hook's return value reaches the model, the provider and traces as `***`. That masking was a `WrapperToolset` around the toolsets `AgentOperator` could find: `toolsets=`, `agent_params["toolsets"]`, and a `Toolset` capability holding a toolset. Everything else reached the model unmasked: - tools from any other capability: `MCP`, `PrefixTools`, `CombinedCapability`, a `Toolset` built per run, or a Dag author's own capability exposing tools - function tools passed in `agent_params["tools"]` - every tool given to `LLMOperator` and the other LLM operators, which pass `agent_params` straight to `PydanticAIHook.create_agent` `create_agent` now adds a `MaskingCapability` that masks around every tool call with pydantic-ai's `wrap_tool_execute` hook, so it applies to every agent the provider builds, including one a `@task` builds with the hook itself. Two related leaks are fixed along the way: - **Retry and failure messages.** By the time a capability sees a tool's `ModelRetry` or `ToolFailed`, pydantic-ai has wrapped it in a `ToolRetryError` or `ToolFailedError`, and the model is sent the `RetryPromptPart` or failed `ToolReturnPart` inside, not the exception's text. Those parts are masked, and an ordinary retry is not logged as a tool failure. - **`ToolReturn.metadata`.** Not sent to the model, but kept in message history and recorded in traces; it was never masked. Extending masking to these tools surfaced a bug in the masker itself: it rebuilt a dataclass result with `dataclasses.replace`, which raises on an `InitVar` without a default and resets `init=False` fields. It now copies the object and masks the copy's fields. ## Design rationale **Why in the hook, not the operator.** Every Common AI operator builds its Pydantic AI agent with `PydanticAIHook.create_agent`, so it is the one place that covers `AgentOperator`, the LLM operators and hand-built agents alike. `AgentOperator` no longer adds anything itself. **Why toolset-level masking stays.** The durable cache (`durable=True`) sits inside the toolset call, below the capability hooks, and must only ever store masked results. So `toolsets=` results are masked twice; the second pass is small next to the model call that consumes the result. **Why innermost, and appended last.** Other capabilities should only ever see a masked result, so `MaskingCapability` declares the innermost position. pydantic-ai breaks ties between capabilities that all ask for innermost by list order, so `create_agent` appends it after the caller's capabilities. A test pins both. **Why not mask at the model-request boundary as well.** A `before_model_request` hook could mask every tool-return and retry part just before it goes to the model, which would also catch the paths below. pydantic-ai keeps the unmasked originals in message history, so that layer would re-mask the whole conversation on every request, and the paths it adds need Dag-author code that puts a secret into a validator's message or a capability's own result. They are listed as not masked in the security docs instead. ## Gotchas These are still not masked, and the security page now says so: errors raised while validating a tool's arguments or the agent's output, a result a capability returns in place of running the tool, and tools the model provider runs itself (native `WebSearch`, `WebFetch` or `ImageGeneration`, and `MCP(native=True)`). The observability page no longer says tool output reaches traces unredacted; tool arguments still do. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
