refactor(hooks): make the unload on_agent_switch builtin pure by dgageot · Pull Request #2706 · docker/docker-agent

dgageot · 2026-05-07T15:41:54Z

Refactors the existing unload on_agent_switch builtin (added in #2684) so the runtime no longer knows about the unload concept. The hook is now a pure, runtime-agnostic HTTP POST driven entirely by data the runtime ships on hooks.Input.

Before

BuiltinUnload was registered as a method on *LocalRuntime.
The runtime imported provider.Unloader and called Unload(ctx) on every model implementing it.
DMR's Client.Unload lived next to its other HTTP plumbing.

After

The runtime ships a snapshot of the previous agent's model endpoints on every on_agent_switch dispatch via a new Input.FromAgentModels []ModelEndpoint field ({provider, model, base_url, unload_api}).
The unload builtin moves to pkg/hooks/builtins/unload.go, depends only on net/http + pkg/hooks, and POSTs {"model": "<id>"} to the resolved _unload URL of every DMR endpoint in the snapshot.
provider.Unloader, dmr.Client.Unload, pkg/runtime/unload.go, and the runtime-side registration are deleted.
DMR's resolved base URL is exposed via a new base.Config.BaseURL field; the private dmr.Client.baseURL is folded into the embedded Config.BaseURL to remove the drift hazard.

What this buys us

The runtime no longer knows the word "unload" — any future "act on the previous agent's models" hook reuses the same FromAgentModels snapshot.
The builtin is testable with just httptest (no fake team / no fake Unloader).
External command hooks on on_agent_switch see the same data, so users can write their own curl-based unload (or warm, health-check, …) without us shipping a builtin.
task lint clean; task test green; net −55 lines vs. the first cut after the simplification commit.

Trade-offs worth flagging

This is a meaningful architectural shift, not a free win. Reviewers should weigh:

The DMR-style URL convention (/v1 → /_unload) now lives in pkg/hooks/builtins, not next to DMR's sibling _configure math.
Provider-type dispatch is stringly-typed (m.Provider == "dmr") instead of an interface assertion.
ModelEndpoint becomes an external JSON wire contract for any command hook on on_agent_switch.
The runtime always builds the snapshot, even when no hook listens; the old design was lazy.

Two small follow-ups would close most of these gaps without reverting:

Move unloadURL into pkg/model/provider/dmr as an exported pure helper; the builtin imports it (and a dmr.ProviderType constant) instead of hard-coding the convention and the literal.
Build FromAgentModels only when executor.Has(EventOnAgentSwitch) is true.

Happy to fold either into this PR if reviewers prefer.

Commits

refactor(hooks): make the unload builtin pure, decouple from runtime
refactor(hooks/builtins): tighten the unload builtin (URL helpers collapsed into one unloadURL, unloadOne extracted with defer cancel(), dead constant + parameter removed, tests deduplicated)

Validation

task lint — golangci-lint: 0 issues; custom 12-cop analyzer: 987 files, no offenses; go mod tidy clean.
task test — every package green, including the affected pkg/hooks/..., pkg/runtime/..., pkg/model/provider/..., pkg/config/latest.
New tests pin the runtime → hook handoff (TestExecuteOnAgentSwitchHooks_PopulatesFromAgentModels, …_FromAgentModelsNilWhenFromEmpty) and the builtin's behaviour end-to-end (TestUnload_* against httptest).

Builds on / supersedes the runtime-coupled implementation from #2684.

aheritier

LGTM. Solid architectural cleanup — the runtime no longer knows about "unload" as a concept, and the ModelEndpoint snapshot generalises cleanly to any future "act on the previous agent" hook.

Strong points:

Clean separation: hooks layer no longer reaches into runtime types; runtime layer no longer imports provider.Unloader.
Net −55 LOC after the simplification commit, despite gaining a public wire contract.
Tests are well-organised — the unloadURL table covers every branch of the resolution algorithm in one place, and the no-op contract is pinned with a single tabular test.
Good test additions on the runtime side (TestExecuteOnAgentSwitchHooks_PopulatesFromAgentModels + the nil-on-empty-from variant).

Trade-offs flagged inline (all non-blocking):

m.Provider != "dmr" is stringly-typed — worth a dmr.ProviderType constant.
The DMR /v1 → /_unload convention now lives in pkg/hooks/builtins, away from its DMR siblings — your follow-up #1 would fix this and pairs well with #1 above.
Snapshot is built unconditionally — your follow-up #2 (gating on executor.Has(...)).
Asymmetry with http_post (http.DefaultClient vs. NewSafeClient) is correct but worth a one-liner.

I'd merge with #1+#2 from the diff folded in if quick, otherwise approve as-is and treat them as follow-ups. CI's already green and the contract is stable either way.

aheritier · 2026-05-07T16:06:42Z

+		return nil, nil
+	}
+	for _, m := range in.FromAgentModels {
+		if m.Provider != "dmr" {


Non-blocking: stringly-typed provider dispatch (m.Provider != "dmr") is fragile. A typo in a future caller, or a future DMR-compatible provider with a different name, will silently no-op.

Worth either:

exporting a dmr.ProviderType = "dmr" constant from pkg/model/provider/dmr and importing it here (your follow-up Telemetry #1), or

documenting the convention in the ModelEndpoint.Provider doc comment so the contract is discoverable.

I'd lean toward the constant — pkg/hooks/builtins already imports pkg/hooks so the dependency direction stays clean.

aheritier · 2026-05-07T16:06:44Z

+	if strings.HasPrefix(m.UnloadAPI, "http://") || strings.HasPrefix(m.UnloadAPI, "https://") {
+		return m.UnloadAPI, nil
+	}
+	if m.BaseURL == "" && m.UnloadAPI == "" {
+		return "", nil
+	}
+	u, err := url.Parse(m.BaseURL)
+	if err != nil || u.Scheme == "" || u.Host == "" {
+		return "", fmt.Errorf("base_url %q is not absolute; cannot resolve unload endpoint", m.BaseURL)
+	}
+	switch {
+	case m.UnloadAPI == "":
+		u.Path = strings.TrimSuffix(strings.TrimSuffix(u.Path, "/"), "/v1") + "/_unload"
+	case strings.HasPrefix(m.UnloadAPI, "/"):
+		u.Path = m.UnloadAPI
+	default:
+		u.Path = "/" + m.UnloadAPI
+	}
+	return u.String(), nil
+}


Non-blocking, but worth folding in: this function knows DMR's URL convention (/v1 → /_unload) and lives next to neither dmr.buildConfigureURL nor any other DMR plumbing. If DMR's convention drifts (e.g. a /v2 engines path, or a different unload route), there's no obvious place to find this code.

Your own follow-up #1 — exporting unloadURL from pkg/model/provider/dmr and calling it from here — would restore the cohesion. Combined with the typed dmr.ProviderType from the comment above, the builtin becomes a dumb dispatcher and DMR owns its conventions.

Doesn't have to land here, but the diff is small and you've already done most of the thinking. I'd merge with this folded in if it's quick.

aheritier · 2026-05-07T16:06:46Z

 		FromAgent:       fromAgent,
 		ToAgent:         toAgent,
 		AgentSwitchKind: kind,
+		FromAgentModels: r.fromAgentModels(fromAgent),


Non-blocking: the snapshot is built on every on_agent_switch regardless of whether anything reads it. For a chain with no on_agent_switch hook this is a wasted alloc + team lookup per transition.

Your follow-up #2 (gating on executor.Has(EventOnAgentSwitch)) is the right call — agent switches aren't hot, but the cheap path stays cheap. Fine to defer to a separate PR.

aheritier · 2026-05-07T16:06:48Z

+// ModelEndpoint identifies one of an agent's configured models plus
+// the HTTP endpoint that hosts it, when one is known. It is the wire
+// format used by [Input.FromAgentModels] so hooks can reach a
+// model-serving endpoint without depending on runtime-only types.
+type ModelEndpoint struct {
+	// Provider is the provider type ("openai", "anthropic", "dmr", ...).
+	Provider string `json:"provider,omitempty"`
+	// Model is the resolved model identifier.
+	Model string `json:"model,omitempty"`
+	// BaseURL is the resolved HTTP base URL of the provider, when known
+	// (set by providers that talk to a configurable HTTP endpoint, e.g.
+	// Docker Model Runner). Empty for cloud providers that don't expose
+	// a stable per-instance base URL on the runtime side.
+	BaseURL string `json:"base_url,omitempty"`
+	// UnloadAPI is the per-model unload path or absolute URL copied
+	// verbatim from the model's `unload_api` provider option. Empty
+	// when the user hasn't configured an override; the unload builtin
+	// falls back to a provider-specific default in that case.
+	UnloadAPI string `json:"unload_api,omitempty"`
+}
+


Praise: nice shape. ModelEndpoint as a self-contained, runtime-free wire format makes the on-agent-switch contract genuinely portable — command hooks now receive the same data as builtins via JSON, and any future "act on the previous agent" hook reuses this slice without touching runtime internals. The decoupling is the real win of this PR.

aheritier · 2026-05-07T16:06:50Z

+	_, err := unload(t.Context(), dmrInput(server, "/custom/unload"), nil)
+	require.NoError(t, err)
+	assert.Equal(t, "/custom/unload", gotPath)
+}


Question / nit: a test with FromAgentModels containing a mix of DMR + non-DMR endpoints would pin the per-element filter (currently only implicitly covered via the all-cloud no-op case). Cheap to add as another dmrInput-style helper input; ignore if you think the existing coverage is enough.

aheritier · 2026-05-07T16:06:54Z

+		return fmt.Errorf("building unload request: %w", err)
+	}
+	req.Header.Set("Content-Type", "application/json")
+


Question: http.DefaultClient here vs. the httpclient.NewSafeClient used by the new http_post builtin (#2705). The asymmetry is justified — DMR runs on loopback, which the SSRF dialer would block — but it's not obvious to a future reader. Worth a one-liner comment noting why the safe client isn't used here (operator-supplied URL, expected localhost target).

docker-agent · 2026-05-07T16:13:21Z

❌ PR Review Failed — The review agent encountered an error and could not complete the review. View logs.

Move the on_agent_switch `unload` builtin out of pkg/runtime and pkg/model/provider/dmr into pkg/hooks/builtins. The builtin no longer depends on the runtime team or the provider.Unloader interface; it reads a snapshot of the previous agent's model endpoints from hooks.Input.FromAgentModels and POSTs to the resolved `_unload` URL over plain HTTP. The runtime ships generic data on every on_agent_switch dispatch (provider, model id, resolved base URL, optional unload_api override) and no longer knows the word 'unload'. Cross-provider chains stay safe: non-DMR endpoints are silently skipped.

Same behaviour, smaller surface: - Collapse the three URL helpers (resolveUnloadURL, rebaseURL, defaultUnloadURL) into one unloadURL(ModelEndpoint) function. The three branches (absolute override, relative override, default derivation) are now visible side-by-side in a single switch. - Extract per-model work into unloadOne so the per-call timeout uses defer cancel() and the call site collapses to a single warn-log on failure. - Drop the unused unloadProviderDMR constant (one literal, one site). - Drop the unused *http.Client parameter from the POST helper; it was always http.DefaultClient, which is reachable from httptest servers anyway. - Drop the redundant len(configured)==0 early return in the runtime's fromAgentModels: the loop handles the empty case, and an empty vs nil slice are wire-equivalent under `omitempty`. - Tests: replace the two URL-helper tables with one TestUnloadURL table covering every branch; collapse the three no-op tests (empty FromAgent, equal From/To, non-DMR providers, no endpoint) into one table-driven TestUnload_NoOpInputs; add a small dmrInput builder so each test highlights only the field it cares about. Net 55 lines removed; `task lint` clean; `task test` green.

Address review feedback on docker#2706: - Export pkg/model/provider/dmr.ProviderType and UnloadURL so the unload builtin becomes a dumb dispatcher and DMR owns the provider literal + the /v1 -> /_unload URL convention (sibling to the existing /_configure helper). Move the URL-resolution table test into the dmr package where the helper now lives. - Gate the FromAgentModels snapshot in executeOnAgentSwitchHooks on executor.Has(EventOnAgentSwitch) so audit-free deployments don't pay the team-lookup + per-model allocation on every agent switch. - Add a mixed-providers test pinning the per-element DMR filter: cloud entries (openai, anthropic) in the snapshot must be silently skipped, only DMR endpoints get POSTed. - Add a comment on http.DefaultClient explaining why the SSRF-safe client used by http_post is wrong here (DMR is loopback).

dgageot requested a review from a team as a code owner May 7, 2026 15:41

aheritier previously approved these changes May 7, 2026

View reviewed changes

dgageot marked this pull request as draft May 7, 2026 16:14

dgageot added 3 commits May 7, 2026 18:18

dgageot dismissed aheritier’s stale review via 33378c9 May 7, 2026 16:25

dgageot force-pushed the board/50f4b45c23b1511a branch from 274b578 to 33378c9 Compare May 7, 2026 16:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

refactor(hooks): make the unload on_agent_switch builtin pure#2706

refactor(hooks): make the unload on_agent_switch builtin pure#2706
dgageot wants to merge 3 commits intodocker:mainfrom
dgageot:board/50f4b45c23b1511a

dgageot commented May 7, 2026

Uh oh!

aheritier left a comment

Uh oh!

aheritier May 7, 2026

Uh oh!

aheritier May 7, 2026

Uh oh!

aheritier May 7, 2026

Uh oh!

aheritier May 7, 2026

Uh oh!

aheritier May 7, 2026

Uh oh!

aheritier May 7, 2026

Uh oh!

docker-agent commented May 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

dgageot commented May 7, 2026

Before

After

What this buys us

Trade-offs worth flagging

Commits

Validation

Uh oh!

aheritier left a comment

Choose a reason for hiding this comment

Uh oh!

aheritier May 7, 2026

Choose a reason for hiding this comment

Uh oh!

aheritier May 7, 2026

Choose a reason for hiding this comment

Uh oh!

aheritier May 7, 2026

Choose a reason for hiding this comment

Uh oh!

aheritier May 7, 2026

Choose a reason for hiding this comment

Uh oh!

aheritier May 7, 2026

Choose a reason for hiding this comment

Uh oh!

aheritier May 7, 2026

Choose a reason for hiding this comment

Uh oh!

docker-agent commented May 7, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants