feat(api): accept model overrides on session creation and add runtime model switching endpoints#2791
feat(api): accept model overrides on session creation and add runtime model switching endpoints#2791dgageot wants to merge 10 commits into
Conversation
Allow callers to specify agent_model_overrides and custom_models_used when creating sessions via POST /api/sessions, and wire the runtime with model switcher configuration to apply those overrides when the session is first run.
Add GET /api/sessions/:id/models to list available models and the current override (if any), and PATCH/POST /api/sessions/:id/model to apply a model override for the session's current agent. Extend ModelChoice with JSON tags for wire format and add SessionModelsResponse type. Both endpoints return 422 if the runtime doesn't support model switching.
PATCH is the canonical verb for updating the agent's model. POST is also accepted because pkg/runtime Client.SetAgentModel (used by RemoteRuntime) was historically a POST; keep both so a client upgrade is not required.
Fix a data race in AvailableSessionModels where rs.session is read without holding sm.mux (the lock protecting reads/writes of AgentModelOverrides and CustomModelsUsed by SetSessionAgentModel). In SetSessionAgentModel, snapshot the in-memory override state before mutating the runtime and session. If sessionStore.UpdateSession fails after the runtime mutation, roll back both the in-memory session state and the runtime override so callers don't observe a runtime/store mismatch on the next request. Fixes issues identified in review and validated with go test -race.
Extract a startAttachedServer helper that wires up a SessionManager + HTTP server with t.Cleanup(ln.Close) so the Serve goroutine exits cleanly when the test finishes. Replaces boilerplate in all model-switching tests and prevents goroutine leaks. Add tests for key picker scenarios: - TestAttachedServer_GetSessionModels_DefaultMarkedCurrent: default marked current when no override is active - TestAttachedServer_GetSessionModels_AppendsCustomModels: session's custom model history appears in the picker - TestSessionManager_SetSessionAgentModel_RollsBackOnStoreFailure: store failure rolls back both in-memory session and runtime state - TestDecorateModelChoices (table): corner cases for the decorate helper Update TestAttachedServer_GetSessionModels to verify the IsCurrent decoration flow end-to-end and to remove the manual IsCurrent: true from the fixture (now correctly set by the manager).
Map ErrSessionNotRunning to 404 so callers can tell apart a stale session ID from a server-side problem. Introduce the sentinel error and replace four ad-hoc errors.New() calls with it: SteerSession, FollowUpSession, AvailableSessionModels, SetSessionAgentModel. Update the HTTP handlers to check for this sentinel explicitly and return 404 accordingly. Also add TestAttachedServer_ModelEndpoints_404WhenNotRunning to verify both GET /sessions/:id/models and PATCH /sessions/:id/model return 404 when no runtime is attached.
docker-agent
left a comment
docker-agent
left a comment
There was a problem hiding this comment.
Assessment: 🟡 NEEDS ATTENTION
This PR introduces runtime model switching (GET/PATCH/POST endpoints) and session-creation model overrides. The implementation is well-structured with a rollback path and a new error sentinel. Three findings were identified in the changed code.
| customRefs = rs.session.CustomModelsUsed | ||
| } | ||
|
|
||
| choices := runtime.DecorateModelChoices(rs.runtime.AvailableModels(ctx), current, customRefs) |
There was a problem hiding this comment.
[HIGH] Lock held across blocking runtime I/O in AvailableSessionModels
sm.mux.Lock() / defer sm.mux.Unlock() wraps the entire function body, including the call to rs.runtime.AvailableModels(ctx) — a context-accepting call that can perform network I/O (e.g. models.dev catalog lookups, provider discovery). Any concurrent caller of RunSession, DeleteSession, SetSessionAgentModel, UpdateSessionTitle, or any other method that acquires sm.mux will block for the full duration of that network call.
The same pattern occurs in SetSessionAgentModel, which calls rs.runtime.SetAgentModel(ctx, agentName, modelRef) under sm.mux.
Fix: snapshot the rs pointer (and any session fields needed) under the lock, release the lock, then invoke the runtime method without holding sm.mux. The runtime is already concurrency-safe via its own internal locking.
| case errors.Is(err, ErrSessionNotRunning): | ||
| return echo.NewHTTPError(http.StatusNotFound, err.Error()) | ||
| default: | ||
| return echo.NewHTTPError(http.StatusBadRequest, err.Error()) |
There was a problem hiding this comment.
[MEDIUM] Unknown errors from SetSessionAgentModel mapped to HTTP 400 instead of 500
The default branch here returns http.StatusBadRequest (400) for any error that is not ErrModelSwitchingNotSupported or ErrSessionNotRunning. HTTP 400 means the client sent an invalid request — but unexpected errors from the session store or from the runtime (e.g. a network failure, an internal validation error wrapped in fmt.Errorf) are server-side faults, not client mistakes.
By contrast, the companion getSessionModels handler (a few lines above) correctly uses http.StatusInternalServerError (500) in its default branch.
Fix: Change the default case to return http.StatusInternalServerError.
default:
return echo.NewHTTPError(http.StatusInternalServerError, err.Error())| for i := range models { | ||
| if currentRef != "" { | ||
| if models[i].Ref == currentRef { | ||
| models[i].IsCurrent = true |
There was a problem hiding this comment.
[MEDIUM] DecorateModelChoices mutates the caller-supplied models slice in-place
models[i].IsCurrent = true and models[i].IsCurrent = models[i].IsDefault modify elements of the slice that was passed in by the caller. In Go, a slice header is copied at the call boundary but the underlying array is shared. If AvailableModels() returns a slice backed by an internal cache (rather than a freshly allocated copy), these index-assignment mutations will silently update the cached structs.
A subsequent call to AvailableModels() could then return entries already marked IsCurrent = true (or false) from a previous decoration call, producing incorrect picker state for a different session or agent.
Fix: Make a defensive copy of the input slice at the top of DecorateModelChoices:
result := make([]ModelChoice, len(models))
copy(result, models)
// ... operate on result, not modelsand return result instead of the mutated input.
|
Thanks for the review! Addressed all three findings in 3 separate commits: [MEDIUM]
|
| group.GET("/sessions/:id/models", s.getSessionModels) | ||
| // PATCH is the canonical verb for updating the agent's model. POST is | ||
| // also accepted because pkg/runtime Client.SetAgentModel (used by | ||
| // RemoteRuntime) was historically a POST; keep both so a client | ||
| // upgrade is not required. | ||
| group.PATCH("/sessions/:id/model", s.setSessionModel) |
There was a problem hiding this comment.
Blocking — docs: The three new endpoints (GET /api/sessions/:id/models, PATCH /api/sessions/:id/model, POST /api/sessions/:id/model) aren't reflected in docs/features/api-server/index.md's endpoint table (around line 44). Worth adding rows so the API server doc stays in sync with the routes registered here. Also worth a short note that POST is accepted for backward compat with RemoteRuntime.
(Could be a follow-up doc PR if you'd rather — but since CHANGELOG is the only doc updated in this branch and that's auto-generated, the API page is the only place a user can actually discover these.)
| switch { | ||
| case errors.Is(err, ErrModelSwitchingNotSupported): | ||
| return echo.NewHTTPError(http.StatusUnprocessableEntity, err.Error()) |
There was a problem hiding this comment.
Question: Mapping ErrModelSwitchingNotSupported to 422 is a defensible choice (the URL/method is fine, the runtime just can't honour it), but 501 Not Implemented or 409 Conflict are both more conventional for "this server doesn't support this capability" / "the resource is in a state that disallows this action". 422 traditionally signals "the body parsed but is semantically wrong", which doesn't really match here — there's no body for GET. Curious whether you considered 501 and rejected it. Not blocking, just worth a sentence in the doc-comment justifying the choice.
| return c.JSON(http.StatusOK, map[string]string{"status": "updated"}) | ||
| } | ||
|
|
||
| // getSessionModels lists the models the user can pick from for the |
There was a problem hiding this comment.
Question (about the API surface, not this handler specifically): the existing createSession binds straight into session.Session and SessionManager.CreateSession cherry-picks the fields it honours (now extended to AgentModelOverrides / CustomModelsUsed). A client that POSTs {"input_tokens": 999, "starred": true, "id": "..."} will silently have those ignored — correct behaviour, but the contract is implicit.
Worth either (a) introducing an explicit api.CreateSessionRequest DTO so the wire surface is type-level visible, or (b) at minimum a doc-comment on CreateSession listing exactly which template fields are honoured. Suggestion only; not a blocker for this PR.
| // session and store never observe interleaved updates. The manager-wide | ||
| // sm.mux is only held briefly while reading or mutating session fields, | ||
| // never while calling into the runtime or the store. | ||
| func (sm *SessionManager) SetSessionAgentModel(ctx context.Context, sessionID, modelRef string) (string, string, error) { |
There was a problem hiding this comment.
Suggestion: SetSessionAgentModel returns (string, string, error) — readable inline but easy to misuse at call sites (which arg is agent, which is model?). The handler in server.go already destructures into agentName, modelRef. Consider returning a small named struct (e.g. type ModelSwitchResult struct { Agent, Model string }) to make the API self-documenting, or rename to (agentName, modelRef, err) in a doc-comment line. Same applies to AvailableSessionModels which returns four positional values. Not blocking.
| sm.mux.Unlock() | ||
|
|
||
| rollback := prevOverride | ||
| if !hadOverride { | ||
| rollback = "" | ||
| } | ||
| if rbErr := rs.runtime.SetAgentModel(ctx, agentName, rollback); rbErr != nil { |
There was a problem hiding this comment.
Suggestion: the rollback block uses two passes (in-memory under sm.mux, then runtime without it) which is fine, but the rollback ref calculation is duplicated:
rollback := prevOverride
if !hadOverride {
rollback = ""
}This same (prevOverride, hadOverride) → ref projection is what the in-memory rollback computes too. Worth extracting a tiny helper, e.g. prevRef := prevOverrideRef(hadOverride, prevOverride) and using it in both rollback branches and as the SetAgentModel arg, so a future reader doesn't have to re-derive that "missing key means clear" rule twice. Non-blocking.
| // GET /api/sessions/:id/models; renaming a tag is a breaking change. | ||
| type ModelChoice struct { | ||
| // Name is the display name (config key) | ||
| Name string | ||
| Name string `json:"name"` | ||
| // Ref is the model reference used internally (e.g., "my_model" or "openai/gpt-4o") | ||
| Ref string | ||
| Ref string `json:"ref"` | ||
| // Provider is the provider name (e.g., "openai", "anthropic") | ||
| Provider string | ||
| Provider string `json:"provider,omitempty"` | ||
| // Model is the specific model name (e.g., "gpt-4o", "claude-sonnet-4-0") | ||
| Model string | ||
| Model string `json:"model,omitempty"` | ||
| // IsDefault indicates this is the agent's configured default model | ||
| IsDefault bool | ||
| IsDefault bool `json:"is_default,omitempty"` | ||
| // IsCurrent indicates this is the currently active model for the agent | ||
| IsCurrent bool | ||
| IsCurrent bool `json:"is_current,omitempty"` | ||
| // IsCustom indicates this is a custom model from the session history (not from config) | ||
| IsCustom bool | ||
| IsCustom bool `json:"is_custom,omitempty"` | ||
| // IsCatalog indicates this is a model from the models.dev catalog | ||
| IsCatalog bool | ||
| IsCatalog bool `json:"is_catalog,omitempty"` | ||
|
|
||
| // The fields below are populated (best-effort) from the models.dev | ||
| // catalog. They are optional and may all be zero/empty when no | ||
| // catalog entry is found for the model. | ||
|
|
||
| // Family is the model family (e.g., "claude", "gpt"). | ||
| Family string | ||
| Family string `json:"family,omitempty"` | ||
| // InputCost is the price (in USD) per 1M input tokens. | ||
| InputCost float64 | ||
| InputCost float64 `json:"input_cost,omitempty"` | ||
| // OutputCost is the price (in USD) per 1M output tokens. | ||
| OutputCost float64 | ||
| OutputCost float64 `json:"output_cost,omitempty"` | ||
| // CacheReadCost is the price (in USD) per 1M cached input tokens. | ||
| CacheReadCost float64 | ||
| CacheReadCost float64 `json:"cache_read_cost,omitempty"` | ||
| // CacheWriteCost is the price (in USD) per 1M cache-write tokens. | ||
| CacheWriteCost float64 | ||
| CacheWriteCost float64 `json:"cache_write_cost,omitempty"` | ||
| // ContextLimit is the maximum context window size in tokens. | ||
| ContextLimit int | ||
| ContextLimit int `json:"context_limit,omitempty"` | ||
| // OutputLimit is the maximum number of tokens the model can produce | ||
| // in a single response. | ||
| OutputLimit int64 | ||
| OutputLimit int64 `json:"output_limit,omitempty"` | ||
| // InputModalities lists the input modalities supported by the model | ||
| // (e.g., "text", "image", "audio"). | ||
| InputModalities []string | ||
| InputModalities []string `json:"input_modalities,omitempty"` | ||
| // OutputModalities lists the output modalities the model can produce. | ||
| OutputModalities []string | ||
| OutputModalities []string `json:"output_modalities,omitempty"` | ||
| } | ||
|
|
There was a problem hiding this comment.
Suggestion: ModelChoice is now part of the wire format (per the new doc-comment) — every public field has a JSON tag. The catalog-only fields below (Family, InputCost, OutputCost, CacheReadCost, CacheWriteCost, ContextLimit, OutputLimit, InputModalities, OutputModalities) are now serialised on every response; that's a lot of optional surface to commit to. If any of those fields might rename or restructure (e.g. consolidating costs into a sub-object), now is the time to consider it before clients depend on them.
Also: Ref and Name carry no omitempty while everything else does — intentional? They'll always be present, so the asymmetry is fine, just calling it out.
| // If the override points at an inline provider/model not in the | ||
| // runtime's list nor in the session's history, fabricate a synthetic | ||
| // choice so the picker can still highlight the active selection. | ||
| if !currentFound && strings.Contains(currentRef, "/") { | ||
| prov, name, _ := strings.Cut(currentRef, "/") | ||
| result = append(result, ModelChoice{ | ||
| Name: currentRef, | ||
| Ref: currentRef, | ||
| Provider: prov, | ||
| Model: name, | ||
| IsCurrent: true, | ||
| IsCustom: true, |
There was a problem hiding this comment.
Question: the synthesized choice for an inline override only fires when the ref contains /. If the override is a config key (my_model) that has been removed from the agent config since the session was created, the picker silently won't highlight anything and IsCurrent is false everywhere. The unit test "non-provider override does not synthesize a fabricated choice" enshrines this as intentional — but is that really what callers want? An entry like {Name: currentRef, Ref: currentRef, IsCurrent: true, IsCustom: true} (with empty Provider/Model) would at least let the UI surface "you're on a model that no longer exists in the config". Non-blocking design call.
| // SessionModelsResponse is the response returned by | ||
| // GET /api/sessions/:id/models. CurrentModelRef is the active override for | ||
| // the named agent (empty when the agent is using its configured default). | ||
| type SessionModelsResponse struct { | ||
| Agent string `json:"agent"` | ||
| CurrentModelRef string `json:"current_model_ref,omitempty"` | ||
| Models []ModelChoice `json:"models"` |
There was a problem hiding this comment.
Question: the response embeds ref for the agent default but no other notion of "default" beyond the per-choice IsDefault flag. Should SessionModelsResponse include a DefaultModelRef string too, so a client showing "Reset to default" doesn't have to scan Models[].IsDefault? Minor; ignore if you'd rather keep it lean.
| // SetSessionModelRequest is the body of PATCH /api/sessions/:id/model. | ||
| // An empty Model clears the override and reverts to the agent's default. | ||
| type SetSessionModelRequest struct { | ||
| Model string `json:"model"` |
There was a problem hiding this comment.
Nit: missing the Assisted-By: docker-agent trailer on a couple of commits in this branch (feat: support model overrides when creating sessions uses Assisted-By: Docker Agent — capitalised — and a couple of the later commits don't have a body at all). Project convention per git-commit-conventions is the lowercase docker-agent form. Cosmetic, but worth fixing if you squash-merge with a clean trailer; harmless if you merge as-is.
| models := a.runtime.AvailableModels(ctx) | ||
|
|
||
| // Determine the currently active model for this agent | ||
| agentName := a.runtime.CurrentAgentName() |
There was a problem hiding this comment.
Praise: clean refactor — 66 lines deleted, 5 added, identical behaviour, and the extracted helper is now exhaustively unit-tested. Good catch on factoring this out.
4 participants
Fixes #2780 and #2781.
What changes
POST /api/sessionshonours model selection (#2780)The handler binds the request body to a
session.Sessionso the existingagent_model_overridesandcustom_models_usedJSON fields are now pickedup and copied onto the new session in
SessionManager.CreateSession. Theruntime applies them the first time it is built for the session (see
runtimeForSession), so a freshly created session can be pinned to aspecific model without an extra round-trip.
Runtime model switching over HTTP (#2781)
Two new endpoints:
GET /api/sessions/:id/models→ list ofruntime.ModelChoicefor thesession's current agent, with
IsCurrentset on the active selection(override or default), and any inline
provider/modelref from thesession's history (or override) appended as a synthetic choice.
PATCH /api/sessions/:id/model→ applies a per-agent override andpersists it. Empty
modelclears the override and reverts to theagent's configured default.
POST /api/sessions/:id/modelis also accepted because the existingpkg/runtimeClient.SetAgentModel(used byRemoteRuntime) washistorically a POST — so no client upgrade is required.
Status code mapping:
Implementation notes
runtime.ModelChoicedirectly as the wire type (added JSON tags).No duplicate
pkg/api.ModelChoice.App.AvailableModelsnow delegates to a sharedruntime.DecorateModelChoiceshelper, which is also used bySessionManager.AvailableSessionModels. HTTP and TUI clients see thesame picker state.
SessionManager.SetSessionAgentModelsnapshots state before mutationand rolls back both the in-memory session and the runtime override
if
sessionStore.UpdateSessionfails.ErrSessionNotRunningdistinguishes "session not foundor not running" from generic errors so the HTTP layer can map it to
404 deterministically.
runtimeForSessionre-applies persisted overrides via a smallapplyStoredOverrideshelper. Failures are logged at WARN — a storedoverride that no longer resolves must not prevent session resumption.
Tests
pkg/server/session_models_test.go(new) — end-to-end coverage forGET/PATCH/POST, IsCurrent decoration, custom models appended, empty
body clears override, store-write rollback, runtime-failure rollback,
404 when no runtime attached, 422 when runtime doesn't support
switching, POST verb compat.
pkg/runtime/model_switcher_test.go— table test forDecorateModelChoicescorner cases.pkg/server/session_manager_test.go—SupportsModelSwitchingonfakeRuntime.Validation
task lint— clean.go test -race -count=3 ./pkg/server/... ./pkg/runtime/... ./pkg/api/...— green.go test— green except pre-existing environmental failuresin
pkg/teamloader(require Docker Model Runner running locally,unrelated to this PR).
Commits