@@ -112,9 +112,9 @@ again. Whether a capability's work is replayed depends on where it runs:
         and capabilities from an agent spec file
       - Tools run again. Pass tools you need replayed in ``toolsets=`` instead.
     * - pydantic-ai-harness ``CodeMode``
-      - Not allowed: the operator raises ``ValueError``, as it does for ``code_mode=True``. This
-        includes a ``CodeMode`` inside a ``CombinedCapability`` or a wrapper such as
-        ``PrefixTools``, but not one a capability function builds when the run starts.
+      - Not allowed: the operator raises ``ValueError``. This includes a ``CodeMode`` inside a
+        ``CombinedCapability`` or a wrapper such as ``PrefixTools``, but not one a capability
+        function builds when the run starts.
 
 See :doc:`durable_execution` for how the cache works.
 
 
@@ -51,6 +51,16 @@ Changelog
   result or ``result["sample_columns"]`` on a summary. A summary carries no ``columns`` key, so
   ``result["columns"]`` raises ``KeyError`` on any table wide enough to be summarized.
 
+.. note::
+  ``AgentOperator`` no longer accepts ``code_mode``. Pass the pydantic-ai-harness capability
+  instead: replace ``code_mode=True`` with ``capabilities=[CodeMode()]``, using
+  ``from pydantic_ai_harness import CodeMode``. The capability takes its own arguments, such
+  as ``max_tool_calls``. If the task also passes ``agent_params={"capabilities": [...]}``,
+  move those into ``capabilities=`` too, since the two cannot be combined. The import runs
+  when the Dag file is parsed, so the ``code-mode`` extra is now needed by the Dag processor
+  as well as the workers. The extra's floor is now ``pydantic-ai-harness>=0.24.0``, the first
+  release with ``max_tool_calls``. See :doc:`code_mode`.
+
 0.10.0
 ......
 
 
@@ -25,9 +25,10 @@ Code mode
     Experimental: this can change or be removed in a minor release of this provider.
     See :ref:`howto/stability`.
 
-Set ``code_mode=True`` to collapse the agent's tools into a single ``run_code``
-tool powered by the `Monty <https://github.com/pydantic/monty>`__ sandbox (via
-pydantic-ai-harness). Instead of one model round-trip per tool call, the model
+Pass pydantic-ai-harness's ``CodeMode`` capability in ``capabilities=`` (see
+:ref:`capabilities`) to collapse the
+agent's tools into a single ``run_code`` tool powered by the
+`Monty <https://github.com/pydantic/monty>`__ sandbox. Instead of one model round-trip per tool call, the model
 writes a single Python snippet that calls the tools as functions -- with loops,
 conditionals, and ``asyncio.gather`` -- in one turn. For multi-tool workflows
 this cuts round-trips and token use.
@@ -68,24 +69,39 @@ mode:
   genuinely required rather than by default. See :doc:`sandbox/index` for the
   backends and their limitations.
 
-The two are not exclusive: ``code_mode=True`` and a ``SandboxToolset`` can be
+The two are not exclusive: ``CodeMode`` and a ``SandboxToolset`` can be
 enabled together, and the file tools fold into ``run_code`` while ``run_command``
 stays a tool of its own.
 
 Requires the ``code-mode`` extra::
 
     pip install "apache-airflow-providers-common-ai[code-mode]"
 
+and the import ``from pydantic_ai_harness import CodeMode`` in the Dag file.
+
 .. exampleinclude:: /../../ai/src/airflow/providers/common/ai/example_dags/example_agent.py
     :language: python
     :start-after: [START howto_operator_agent_code_mode]
     :end-before: [END howto_operator_agent_code_mode]
 
-``code_mode=True`` is the same as adding ``CodeMode()`` at the end of ``capabilities=`` (see
-:ref:`capabilities`), except that the operator builds the capability when the task runs, so the
-Dag file does not import ``pydantic_ai_harness``.
+``CodeMode`` takes its own arguments; see the
+`pydantic-ai-harness code mode docs <https://ai.pydantic.dev/harness/code-mode/>`__. The one
+you are most likely to need is ``max_tool_calls``: each ``run_code`` snippet may make at most
+that many nested tool calls, 100 by default. A snippet that asks for more fails, the calls it
+already made are not undone, and the model has to split the work across several snippets. If the
+agent fans out over a list, for example one tool call per row or per file, set it above the
+largest list you expect.
+
+Code mode cannot be combined with ``durable=True`` (see :doc:`capabilities` for which
+capabilities durable replay covers), and it turns off :doc:`tool_approval`. A tool that needs
+approval and is called from ``run_code`` does not run: the model gets an error back and may
+carry on without it, so the task can still succeed. To keep such a tool out of ``run_code``,
+list the other tools in ``CodeMode(tools=[...])``; called directly, it fails the task.
+
+Importing ``CodeMode`` in the Dag file loads pydantic-ai-harness and Monty each time the file is
+parsed. Keep code mode agents in their own Dag file if the rest of the file does not need them.
 
 .. note::
 
-    Monty is pre-1.0. The ``code-mode`` extra is opt-in so its dependency churn
+    pydantic-ai-harness is pre-1.0. The ``code-mode`` extra is opt-in so its dependency churn
     never affects the base provider install.
@@ -28,8 +28,8 @@ you use. Each is a parameter on the operator or decorator.
 - :doc:`message_history`: ``message_history`` carries a conversation across agent runs.
 - :doc:`capabilities`: ``capabilities=`` adds pydantic-ai capabilities such as ``Thinking`` and
   ``WebSearch``, and ``pydantic-ai-shields`` guardrails, to an agent.
-- :doc:`code_mode`: ``code_mode=True`` lets the model call several tools from one Python
-  snippet instead of one round trip per call.
+- :doc:`code_mode`: the ``CodeMode`` capability lets the model call several tools from one
+  Python snippet instead of one round trip per call.
 - :doc:`approval_gates`: ``require_approval=True`` pauses an LLM operator until a person
   approves, edits or rejects the output.
 - :doc:`hitl_review`: ``enable_hitl_review=True`` opens an iterative review loop on an agent,
 
@@ -227,7 +227,7 @@ Extra            Dependencies
 ``mcp``          ``pydantic-ai-slim[mcp]>=2.33.0``
 ``modal``        ``apache-airflow-providers-modal``, ``modal>=1.5.2``
 ``opensandbox``  ``opensandbox>=1.1.0``
-``code-mode``    ``pydantic-ai-harness[codemode]>=0.3.0``
+``code-mode``    ``pydantic-ai-harness[codemode]>=0.24.0``
 ``shields``      ``pydantic-ai-shields>=0.3.4``
 ``skills``       ``apache-airflow-providers-git>=0.4.0``, ``pydantic-ai-skills>=1.2.0``
 ``avro``         ``fastavro>=1.10.0; python_version < "3.14"``, ``fastavro>=1.12.1; python_version >= "3.14"``
 
@@ -251,8 +251,8 @@ Five features have pages of their own:
   on retry instead of paying for them again.
 - :doc:`../capabilities`: pass pydantic-ai capabilities and ``pydantic-ai-shields`` guardrails
   with ``capabilities=``.
-- :doc:`../code_mode`: set ``code_mode=True`` to collapse the agent's tools into a single
-  ``run_code`` tool the model drives by writing Python.
+- :doc:`../code_mode`: pass the ``CodeMode`` capability to collapse the agent's tools into a
+  single ``run_code`` tool the model drives by writing Python.
 - :doc:`../tool_approval`: mark tools that need a person's approval, and the task pauses before
   a marked call runs.
 
@@ -435,9 +435,6 @@ Parameters
   when the steps it needs are cached. Clearing a failed task instance starts
   a fresh budget but keeps the durable cache its attempts left behind, so what
   the rerun replays from that cache is free there too.
-- ``code_mode``: When ``True``, wraps the agent's tools in a single ``run_code``
-  tool that the model drives by writing Python, executed in the Monty sandbox.
-  Requires the ``code-mode`` extra. Default ``False``. See :ref:`code-mode`.
 - ``cache_prompt``: Ask the provider to cache the tool definitions, system prompt and
   conversation so later requests read them back at a discount. Default ``True``; a no-op for
   providers that cache on their own. See :ref:`agent-prompt-caching`.
 
@@ -359,7 +359,7 @@ somewhere to work". This page has the worked scenarios above and
 the limitations below to read before designing a Dag around it.
 
 Before reaching for it, check whether the actual need is narrower than that.
-``code_mode=True`` is a flag on ``AgentOperator``. It changes how the
+:ref:`Code mode <code-mode>` is a capability on ``AgentOperator``. It changes how the
 model invokes the tools it already has, letting it write code that calls several
 of them instead of emitting one call per step. It does not give the agent somewhere
 to run arbitrary code of its own. Because it needs no backend, it avoids the
 
@@ -51,7 +51,7 @@ Pydantic AI toolsets, but the Pydantic AI class they inherit from can change.
    * - ``@task.agent`` and :class:`~airflow.providers.common.ai.operators.agent.AgentOperator`
      - Runs an agent with the model from ``llm_conn_id`` and the given toolsets, and
        returns its output under the same rules as ``@task.llm``. ``durable``,
-       ``code_mode`` and per-tool approval are experimental; see below.
+       code mode and per-tool approval are experimental; see below.
    * - ``message_history`` on ``AgentOperator``
      - Seeds the run with the given conversation, as a list of messages or its JSON form,
        and publishes the finished conversation under the ``message_history`` XCom key so a
@@ -131,7 +131,7 @@ Everything this provider ships that is not in the table above is experimental.
        ``on_tool_approval_timeout`` and ``tool_approval_assigned_users``;
        :doc:`tool_approval`)
      - New, and needs Airflow 3.3. How a paused run resumes may change.
-   * - ``code_mode`` (:doc:`code_mode`), the Agent Skills toolset
+   * - Code mode (:doc:`code_mode`), the Agent Skills toolset
        (:doc:`toolsets/skills`) and the ``shields`` extra (used in :doc:`capabilities`)
      - Thin integrations of packages outside this provider whose APIs are still
        changing: ``pydantic-ai-harness``, ``pydantic-ai-skills`` and
 
@@ -127,11 +127,13 @@ Requirements and limits
 
 - Airflow 3.3 or later. On older versions a tool marked for approval fails the
   task, as it did before.
-- Not together with ``durable=True``, ``enable_hitl_review=True``, code mode
-  (``code_mode=True`` or a ``CodeMode`` capability), or a ``SandboxToolset`` that
+- Not together with ``durable=True``, ``enable_hitl_review=True``, a ``CodeMode``
+  capability, or a ``SandboxToolset`` that
   provisions its own sandbox. Each
   assumes the run finishes in one go; that sandbox, for one, is destroyed when the
-  run pauses. With any of them, a marked tool fails the task. A ``SandboxToolset``
+  run pauses. With any of them, a marked tool fails the task, except under code mode,
+  where a marked tool called from ``run_code`` does not run and the model may carry on
+  without it (see :doc:`code_mode`). A ``SandboxToolset``
   attached to a sandbox another task owns keeps its files through the pause, so it
   is allowed; the wait spends that sandbox's lifetime
   (:ref:`sandbox-attach`).
 
@@ -44,8 +44,9 @@ Model and connection errors
     version suffixes, for example), use the matching vendor connection type instead of
     the generic one.
 
-An ``ImportError`` for ``pydantic_ai.models.<vendor>`` or the vendor SDK
-    The provider is installed without the extra for that vendor. Install it, quoting the
+An ``ImportError`` for ``pydantic_ai.models.<vendor>``, the vendor SDK, or ``pydantic_ai_harness``
+    The provider is installed without the extra for that vendor, or without ``code-mode``
+    for ``pydantic_ai_harness``. Install it, quoting the
     package name so the brackets survive the shell:
 
     .. code-block:: bash
@@ -102,17 +103,14 @@ as a task failure.
     or later. Upgrade Airflow, or use ``on_uncertain="fail"`` and drop the review flags
     on an older Airflow version. See :doc:`approval_gates` and :doc:`hitl_review`.
 
-``durable=True and enable_hitl_review=True cannot be used together`` / ``durable=True and code_mode=True cannot be used together``
+``durable=True and enable_hitl_review=True cannot be used together`` / ``durable=True cannot be used with a CodeMode capability``
     Durable replay assumes a stable step order across attempts, which neither a human
     review loop nor code mode provides. Pick one. See :doc:`durable_execution`.
 
 ``message_history and enable_hitl_review=True cannot be used together``
     The post-review transcript is not recoverable today, so the operator refuses rather
     than silently dropping the reviewed turns. See :doc:`message_history`.
 
-``code_mode=True requires the 'code-mode' extra``
-    Install ``apache-airflow-providers-common-ai[code-mode]``. See :doc:`code_mode`.
-
 ``... does not support decision_policy yet``
     Only ``LLMOperator`` and ``LLMBranchOperator`` honor a ``DecisionPolicy`` with a
     confidence bar. The SQL, schema-compare and file-analysis operators run their own