kaxil opened a new pull request, #73559:
URL: https://github.com/apache/airflow/pull/73559
`SandboxToolset` provisions its own sandbox on the model's first tool call
and destroys it when the run ends. Three things cannot be done inside that
shape, and they have come up separately: a credential cannot come from a
connection (the spec is fixed at parse time), a file the agent built cannot
leave (only the model's text-only, capped context crosses the boundary), and
HITL regeneration starts a second run whose history describes files that no
longer exist, which is why #73529 refuses `enable_hitl_review` with a sandbox.
This PR gives them one answer. A `@task` creates the sandbox and hands the
agent only its handle:
```python
@task
def provision(**context) -> str:
backend = ModalSandboxBackend()
sandbox = backend.create(spec=SandboxSpec(owner=dag_run_owner(context)))
backend.write_file(sandbox, "/workspace/orders.csv", staged_bytes) #
the task's credential, not the sandbox's
return sandbox
analyse = AgentOperator(..., toolsets=[SandboxToolset(ModalSandboxBackend(),
attach_to="{{ ti.xcom_pull(task_ids='provision') }}")])
@task(trigger_rule=TriggerRule.ALL_DONE)
def collect(sandbox: str):
try:
report = backend.read_file(sandbox, "/workspace/report.md",
max_bytes=16 * MIB)
finally:
backend.destroy(sandbox)
```
The toolset uses that sandbox for the run, leaves it standing at the end,
and the task that created it reads the artifact out through the backend and
destroys it. `enable_hitl_review` is now accepted with an attached toolset,
since the regenerated run finds the first run's files. `durable=True` stays
refused: a replayed tool result is not re-executed, so the workspace would not
move with the transcript.
Live run on Modal with a real model (Opus 4.8 through a gateway).
`provision` stages the input, `analyse` attaches and writes the report,
`collect` reads it out after the agent run ended and destroys the sandbox. The
two red tasks are agents that were refused before running a single tool: one
presented the wrong owner, one a handle that does not exist.
## Design rationale
**A handle alone is never enough.** An upstream XCom can be written by an
agent, and a valid id for someone else's sandbox in the same Modal workspace is
a valid id. So the provisioning task stamps an owner (`SandboxSpec.owner`), and
the toolset refuses to attach unless the sandbox carries the owner it presents.
By default that is the Dag run, so two tasks in one run need no shared secret;
`owner=` on the toolset covers a sandbox provisioned under another name. The
docs are explicit about what the check is: the tags are written with the same
vendor credential the attaching task holds, so it stops a run reaching the
wrong sandbox by mistake and gives attribution. It is not a boundary between
authors.
**One agent run holds a sandbox at a time.** Attaching marks the sandbox
with the task instance (Dag, run, task, map index, without the try number), and
the run's end clears it. A different task is refused while the mark is there;
the same task attaching again is allowed, so a retry finds the files of an
attempt that died without releasing. Modal's tags have no conditional write, so
the claim is a read-then-write that is read back once. That catches the
sequential mistakes; two runs starting in the same instant can both pass, and
the docs say so rather than promising a lock.
**The lifetime and the network policy travel with the sandbox.** `create`
stamps the expiry and the network policy as tags. On attach the toolset
shortens commands to what is left, whatever backend the sandbox is on, and the
`run_command` description tells the model whose sandbox it is, what it can
reach, and how much time remained when the run began. The note is fixed at
attach rather than recomputed each step, so tool definitions stay stable for
provider prompt caching. A backend that provisions a sandbox with an owner
refuses an `idle_timeout`, since the first gap between tasks would reclaim it.
**Why the ownership rules live on a base class.** `AttachableSandboxBackend`
adds two vendor primitives, `read_tags` and `write_tags`, and the owner, holder
and expiry rules are written once on top of them. Modal implements the two
primitives; `sbx` runs a microVM on the worker that created it and cannot be
reached from another task, so it stays a plain `SandboxBackend`, refuses
`SandboxSpec.owner`, and the toolset refuses `attach_to` for it at construction.
**Why `toolsets` is not a template field.** The handle has to be templated
to come from XCom, and the natural move is to add `toolsets` to
`AgentOperator.template_fields`. That would write the toolset list into the
serialized Dag by repr, and a pydantic-ai wrapper's repr carries the wrapped
object's memory address, so any Dag composing toolsets with `.prefixed()` or a
combined toolset would get a new Dag version on every parse. The operator
instead overrides `render_template_fields` and renders the template fields a
toolset declares, looking inside wrappers and combinations. A handle that
renders to nothing (the provisioning task pushed no XCom, or the agent ran
first) fails the run instead of falling back to a sandbox of the toolset's own.
## Gotchas
- The reviewer's wait in HITL spends the provisioning backend's
`sandbox_timeout`; size it for the review, or set `hitl_timeout` below what
will be left.
- A collecting task on `ALL_DONE` also runs when provisioning itself failed
and the handle is `None`; the example checks that first, and the Modal backend
now refuses a non-string handle with a message naming the cause instead of
failing inside the SDK.
- The `modal` extra's floor moves from 1.5.0 to 1.5.2: reading tags back is
gated to V1 sandboxes in 1.5.0 and 1.5.1, so attaching would fail on a V2
sandbox there.
- A run that cannot clear its mark retries three times, then logs the holder
it left behind; only that holder, or the task that created the sandbox, can use
it until the lifetime reclaims it.
The system test gains a task that provisions, attaches two agent runs in
turn, refuses a run with the wrong owner, reads the file out through the
backend and destroys the sandbox; it passed live against Modal.
--
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]
