Update output token limit when switching models

A same-provider model switch reused the running agent and changed only
its model ID, so MaxTokens kept the previous model's output budget.
Switching from a model with a small or missing limit to a larger one,
e.g. onto Bedrock Claude Opus 5.5, kept truncating responses with
max_tokens.

Adopt the new model's MaxOutput in the interactive model swap, RPC
set_model, SDK Runtime.SetModel, and the bot's runtime reload.
Context-window checks already look up the current model on each use.

commit: c63bfd310615583691e2c8a4192de4bcd99b3e2d
author: Raymond Gasper · 23d ago · 2026-09-29 13:44 UTC
parent: 2ff82fa
signature: unsigned

8 files · +107 -3

M docs/rpc.md +1 -1


@@ -164,7 +164,7 @@ Drop the entire transcript. Equivalent to the `/clear` slash command.
164	164
165	165		### `set_model`
166	166
167		-	Switch model within the same provider.
	167	+	Switch model within the same provider. Subsequent prompts use the new model's output-token limit.
168	168
169	169		```json
170	170		{"id":"7","type":"set_model","model":"claude-sonnet-4-5"}


@@ -164,7 +164,7 @@ Drop the entire transcript. Equivalent to the `/clear` slash command.
164			164
165		### `set_model`	165		### `set_model`
166			166
167	-	Switch model within the same provider.	167	+	Switch model within the same provider. Subsequent prompts use the new model's output-token limit.
168			168
169		```json	169		```json
170		{"id":"7","type":"set_model","model":"claude-sonnet-4-5"}	170		{"id":"7","type":"set_model","model":"claude-sonnet-4-5"}

M packages/agent/botcmd.go +1 -0


@@ -272,6 +272,7 @@ func botRun(spec *botSpec, rawTail []string, version string) error {
272	272		}
273	273		agent.Client = next.NewClient()
274	274		agent.Model = next.Model
	275	+	agent.MaxTokens = next.MaxOutput
275	276		runner.UpdateRuntimeConfig(next.Provider, next.Model, next.AuthMethod, next.CWD)
276	277		return nil
277	278		},


@@ -272,6 +272,7 @@ func botRun(spec *botSpec, rawTail []string, version string) error {
272	}	272		}
273	agent.Client = next.NewClient()	273		agent.Client = next.NewClient()
274	agent.Model = next.Model	274		agent.Model = next.Model
		275	+	agent.MaxTokens = next.MaxOutput
275	runner.UpdateRuntimeConfig(next.Provider, next.Model, next.AuthMethod, next.CWD)	276		runner.UpdateRuntimeConfig(next.Provider, next.Model, next.AuthMethod, next.CWD)
276	return nil	277		return nil
277	},	278		},

M packages/agent/modes/interactive.go +3 -0


@@ -5793,6 +5793,9 @@ func (i *Interactive) swapModel(prov, model string, builder func(string, string)
5793	5793		i.mu.Lock()
5794	5794		i.cfg.Model = m.ID
5795	5795		i.agent.Model = m.ID
	5796	+	// The output budget is per model; keeping the previous model's
	5797	+	// value can truncate responses (or exceed the new model's cap).
	5798	+	i.agent.MaxTokens = m.MaxOutput
5796	5799		i.statusOK = "model: " + m.ID
5797	5800		i.statusErr = ""
5798	5801		i.mu.Unlock()


@@ -5793,6 +5793,9 @@ func (i *Interactive) swapModel(prov, model string, builder func(string, string)
5793	i.mu.Lock()	5793		i.mu.Lock()
5794	i.cfg.Model = m.ID	5794		i.cfg.Model = m.ID
5795	i.agent.Model = m.ID	5795		i.agent.Model = m.ID
		5796	+	// The output budget is per model; keeping the previous model's
		5797	+	// value can truncate responses (or exceed the new model's cap).
		5798	+	i.agent.MaxTokens = m.MaxOutput
5796	i.statusOK = "model: " + m.ID	5799		i.statusOK = "model: " + m.ID
5797	i.statusErr = ""	5800		i.statusErr = ""
5798	i.mu.Unlock()	5801		i.mu.Unlock()

A packages/agent/modes/interactive_model_swap_test.go +39 -0

@@ -0,0 +1,39 @@
	1	+	package modes
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/patriceckhart/zot/packages/core"
	7	+	"github.com/patriceckhart/zot/packages/provider"
	8	+	)
	9	+
	10	+	// A same-provider /model switch reuses the agent, so it must also adopt
	11	+	// the new model's output budget. Keeping the previous model's MaxTokens
	12	+	// truncated responses (e.g. a 4096 budget carried onto a 128k model).
	13	+	func TestSameProviderModelSwapUpdatesMaxTokens(t *testing.T) {
	14	+	const from, to = "claude-3-5-haiku-20241022", "claude-opus-4-1-20250805"
	15	+	fromModel, err := provider.FindModel("anthropic", from)
	16	+	if err != nil {
	17	+	t.Fatal(err)
	18	+	}
	19	+	toModel, err := provider.FindModel("anthropic", to)
	20	+	if err != nil {
	21	+	t.Fatal(err)
	22	+	}
	23	+	if fromModel.MaxOutput == toModel.MaxOutput \|\| toModel.MaxOutput <= 0 {
	24	+	t.Fatalf("test needs distinct positive limits, got %d and %d", fromModel.MaxOutput, toModel.MaxOutput)
	25	+	}
	26	+
	27	+	ag := core.NewAgent(nil, from, "", nil)
	28	+	ag.MaxTokens = fromModel.MaxOutput
	29	+	i := &Interactive{agent: ag, cfg: InteractiveConfig{Provider: "anthropic", Model: from}}
	30	+
	31	+	i.swapModel("anthropic", to, nil, false)
	32	+
	33	+	if ag.Model != to {
	34	+	t.Fatalf("agent model = %q, want %q", ag.Model, to)
	35	+	}
	36	+	if ag.MaxTokens != toModel.MaxOutput {
	37	+	t.Fatalf("agent MaxTokens = %d, want %d", ag.MaxTokens, toModel.MaxOutput)
	38	+	}
	39	+	}

@@ -0,0 +1,39 @@
	1	+	package modes
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/patriceckhart/zot/packages/core"
	7	+	"github.com/patriceckhart/zot/packages/provider"
	8	+	)
	9	+
	10	+	// A same-provider /model switch reuses the agent, so it must also adopt
	11	+	// the new model's output budget. Keeping the previous model's MaxTokens
	12	+	// truncated responses (e.g. a 4096 budget carried onto a 128k model).
	13	+	func TestSameProviderModelSwapUpdatesMaxTokens(t *testing.T) {
	14	+	const from, to = "claude-3-5-haiku-20241022", "claude-opus-4-1-20250805"
	15	+	fromModel, err := provider.FindModel("anthropic", from)
	16	+	if err != nil {
	17	+	t.Fatal(err)
	18	+	}
	19	+	toModel, err := provider.FindModel("anthropic", to)
	20	+	if err != nil {
	21	+	t.Fatal(err)
	22	+	}
	23	+	if fromModel.MaxOutput == toModel.MaxOutput \|\| toModel.MaxOutput <= 0 {
	24	+	t.Fatalf("test needs distinct positive limits, got %d and %d", fromModel.MaxOutput, toModel.MaxOutput)
	25	+	}
	26	+
	27	+	ag := core.NewAgent(nil, from, "", nil)
	28	+	ag.MaxTokens = fromModel.MaxOutput
	29	+	i := &Interactive{agent: ag, cfg: InteractiveConfig{Provider: "anthropic", Model: from}}
	30	+
	31	+	i.swapModel("anthropic", to, nil, false)
	32	+
	33	+	if ag.Model != to {
	34	+	t.Fatalf("agent model = %q, want %q", ag.Model, to)
	35	+	}
	36	+	if ag.MaxTokens != toModel.MaxOutput {
	37	+	t.Fatalf("agent MaxTokens = %d, want %d", ag.MaxTokens, toModel.MaxOutput)
	38	+	}
	39	+	}

M packages/agent/rpc.go +3 -1


@@ -289,11 +289,13 @@ func (s *rpcServer) dispatch(cmd, id string, raw []byte) {
289	289		s.writeError(id, cmd, err.Error())
290	290		return
291	291		}
292		-	if _, err := provider.FindModel(s.provider, req.Model); err != nil {
	292	+	m, err := provider.FindModel(s.provider, req.Model)
	293	+	if err != nil {
293	294		s.writeError(id, cmd, err.Error())
294	295		return
295	296		}
296	297		s.agent.Model = req.Model
	298	+	s.agent.MaxTokens = m.MaxOutput
297	299		s.model = req.Model
298	300		s.writeResponse(id, cmd, map[string]any{"model": req.Model})
299	301

 			s.writeError(id, cmd, err.Error())
 			return
 		}
-		if _, err := provider.FindModel(s.provider, req.Model); err != nil {
+		m, err := provider.FindModel(s.provider, req.Model)
+		if err != nil {
 			s.writeError(id, cmd, err.Error())
 			return
 		}
 		s.agent.Model = req.Model
+		s.agent.MaxTokens = m.MaxOutput
 		s.model = req.Model
 		s.writeResponse(id, cmd, map[string]any{"model": req.Model})

A packages/agent/rpc_setmodel_test.go +34 -0

@@ -0,0 +1,34 @@
	1	+	package agent
	2	+
	3	+	import (
	4	+	"bytes"
	5	+	"context"
	6	+	"encoding/json"
	7	+	"testing"
	8	+
	9	+	"github.com/patriceckhart/zot/packages/core"
	10	+	"github.com/patriceckhart/zot/packages/provider"
	11	+	)
	12	+
	13	+	func TestRPCSetModelUpdatesMaxTokens(t *testing.T) {
	14	+	const to = "claude-opus-4-1-20250805"
	15	+	want, err := provider.FindModel("anthropic", to)
	16	+	if err != nil {
	17	+	t.Fatal(err)
	18	+	}
	19	+	ag := core.NewAgent(&startCaptureClient{}, "claude-3-5-haiku-20241022", "", nil)
	20	+	ag.MaxTokens = 4096
	21	+	var out bytes.Buffer
	22	+	s := &rpcServer{ctx: context.Background(), agent: ag, provider: "anthropic", model: ag.Model, out: &out}
	23	+
	24	+	raw, err := json.Marshal(map[string]string{"model": to})
	25	+	if err != nil {
	26	+	t.Fatal(err)
	27	+	}
	28	+	s.dispatch("set_model", "1", raw)
	29	+	s.inFlight.Wait()
	30	+
	31	+	if ag.Model != to \|\| ag.MaxTokens != want.MaxOutput {
	32	+	t.Fatalf("after set_model: model=%q MaxTokens=%d, want %q %d (out=%s)", ag.Model, ag.MaxTokens, to, want.MaxOutput, out.String())
	33	+	}
	34	+	}

@@ -0,0 +1,34 @@
	1	+	package agent
	2	+
	3	+	import (
	4	+	"bytes"
	5	+	"context"
	6	+	"encoding/json"
	7	+	"testing"
	8	+
	9	+	"github.com/patriceckhart/zot/packages/core"
	10	+	"github.com/patriceckhart/zot/packages/provider"
	11	+	)
	12	+
	13	+	func TestRPCSetModelUpdatesMaxTokens(t *testing.T) {
	14	+	const to = "claude-opus-4-1-20250805"
	15	+	want, err := provider.FindModel("anthropic", to)
	16	+	if err != nil {
	17	+	t.Fatal(err)
	18	+	}
	19	+	ag := core.NewAgent(&startCaptureClient{}, "claude-3-5-haiku-20241022", "", nil)
	20	+	ag.MaxTokens = 4096
	21	+	var out bytes.Buffer
	22	+	s := &rpcServer{ctx: context.Background(), agent: ag, provider: "anthropic", model: ag.Model, out: &out}
	23	+
	24	+	raw, err := json.Marshal(map[string]string{"model": to})
	25	+	if err != nil {
	26	+	t.Fatal(err)
	27	+	}
	28	+	s.dispatch("set_model", "1", raw)
	29	+	s.inFlight.Wait()
	30	+
	31	+	if ag.Model != to \|\| ag.MaxTokens != want.MaxOutput {
	32	+	t.Fatalf("after set_model: model=%q MaxTokens=%d, want %q %d (out=%s)", ag.Model, ag.MaxTokens, to, want.MaxOutput, out.String())
	33	+	}
	34	+	}

M packages/agent/sdk/sdk.go +3 -1


@@ -299,10 +299,12 @@ func (r *Runtime) SetModel(model string) error {
299	299		if r.agent == nil {
300	300		return fmt.Errorf("sdk: no agent")
301	301		}
302		-	if _, err := provider.FindModel(r.provider, model); err != nil {
	302	+	m, err := provider.FindModel(r.provider, model)
	303	+	if err != nil {
303	304		return err
304	305		}
305	306		r.agent.Model = model
	307	+	r.agent.MaxTokens = m.MaxOutput
306	308		r.model = model
307	309		return nil
308	310		}

 	if r.agent == nil {
 		return fmt.Errorf("sdk: no agent")
 	}
-	if _, err := provider.FindModel(r.provider, model); err != nil {
+	m, err := provider.FindModel(r.provider, model)
+	if err != nil {
 		return err
 	}
 	r.agent.Model = model
+	r.agent.MaxTokens = m.MaxOutput
 	r.model = model
 	return nil
 }

A packages/agent/sdk/setmodel_test.go +23 -0

@@ -0,0 +1,23 @@
	1	+	package sdk
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/patriceckhart/zot/packages/core"
	7	+	"github.com/patriceckhart/zot/packages/provider"
	8	+	)
	9	+
	10	+	func TestRuntimeSetModelUpdatesMaxTokens(t *testing.T) {
	11	+	const to = "claude-opus-4-1-20250805"
	12	+	want, err := provider.FindModel("anthropic", to)
	13	+	if err != nil {
	14	+	t.Fatal(err)
	15	+	}
	16	+	r := &Runtime{agent: &core.Agent{MaxTokens: 4096}, provider: "anthropic"}
	17	+	if err := r.SetModel(to); err != nil {
	18	+	t.Fatal(err)
	19	+	}
	20	+	if r.agent.MaxTokens != want.MaxOutput {
	21	+	t.Fatalf("MaxTokens = %d, want %d", r.agent.MaxTokens, want.MaxOutput)
	22	+	}
	23	+	}

@@ -0,0 +1,23 @@
	1	+	package sdk
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/patriceckhart/zot/packages/core"
	7	+	"github.com/patriceckhart/zot/packages/provider"
	8	+	)
	9	+
	10	+	func TestRuntimeSetModelUpdatesMaxTokens(t *testing.T) {
	11	+	const to = "claude-opus-4-1-20250805"
	12	+	want, err := provider.FindModel("anthropic", to)
	13	+	if err != nil {
	14	+	t.Fatal(err)
	15	+	}
	16	+	r := &Runtime{agent: &core.Agent{MaxTokens: 4096}, provider: "anthropic"}
	17	+	if err := r.SetModel(to); err != nil {
	18	+	t.Fatal(err)
	19	+	}
	20	+	if r.agent.MaxTokens != want.MaxOutput {
	21	+	t.Fatalf("MaxTokens = %d, want %d", r.agent.MaxTokens, want.MaxOutput)
	22	+	}
	23	+	}