Reject SDK model switches during active turns

2 files · +32 -3

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


@@ -292,13 +292,17 @@ func (r *Runtime) Compact(ctx context.Context, customInstructions string) (Compa
292	292		}
293	293
294	294		// SetModel switches the active model. Same provider only; for cross-
295		-	// provider switches, create a new Runtime.
	295	+	// provider switches, create a new Runtime. Returns ErrBusy during a
	296	+	// Prompt or Compact so the model and output limit cannot change mid-turn.
296	297		func (r *Runtime) SetModel(model string) error {
297	298		r.mu.Lock()
298	299		defer r.mu.Unlock()
299	300		if r.agent == nil {
300	301		return fmt.Errorf("sdk: no agent")
301	302		}
	303	+	if r.activeCancel != nil {
	304	+	return ErrBusy
	305	+	}
302	306		m, err := provider.FindModel(r.provider, model)
303	307		if err != nil {
304	308		return err

@@ -377,8 +381,8 @@ func (r *Runtime) ListModels() []ModelInfo {
377	381
378	382		// ---- public errors ----
379	383
380		-	// ErrBusy is returned when a Prompt or Compact is started while
381		-	// another is in flight on the same Runtime.
	384	+	// ErrBusy is returned when a Prompt, Compact, or SetModel is called
	385	+	// during an active Prompt or Compact on the same Runtime.
382	386		var ErrBusy = errors.New("sdk: runtime is busy")
383	387
384	388		// ErrClosed is returned by methods on a Runtime after Close.

 }
 // SetModel switches the active model. Same provider only; for cross-
-// provider switches, create a new Runtime.
+// provider switches, create a new Runtime. Returns ErrBusy during a
+// Prompt or Compact so the model and output limit cannot change mid-turn.
 func (r *Runtime) SetModel(model string) error {
 	r.mu.Lock()
 	defer r.mu.Unlock()
 	if r.agent == nil {
 		return fmt.Errorf("sdk: no agent")
 	}
+	if r.activeCancel != nil {
+		return ErrBusy
+	}
 	m, err := provider.FindModel(r.provider, model)
 	if err != nil {
 		return err
 // ---- public errors ----
-// ErrBusy is returned when a Prompt or Compact is started while
+// ErrBusy is returned when a Prompt, Compact, or SetModel is called
-// another is in flight on the same Runtime.
+// during an active Prompt or Compact on the same Runtime.
 var ErrBusy = errors.New("sdk: runtime is busy")
 // ErrClosed is returned by methods on a Runtime after Close.

M packages/agent/sdk/setmodel_test.go +25 -0

@@ -1,12 +1,37 @@
1	1		package sdk
2	2
3	3		import (
	4	+	"errors"
4	5		"testing"
5	6
6	7		"github.com/patriceckhart/zot/packages/core"
7	8		"github.com/patriceckhart/zot/packages/provider"
8	9		)
9	10
	11	+	func TestRuntimeSetModelRejectsBusyAgent(t *testing.T) {
	12	+	const from, to = "claude-3-5-haiku-20241022", "claude-opus-4-1-20250805"
	13	+	r := &Runtime{
	14	+	agent: &core.Agent{Model: from, MaxTokens: 4096},
	15	+	provider: "anthropic",
	16	+	model: from,
	17	+	activeCancel: func() {}, // Prompt and Compact both hold this until complete.
	18	+	}
	19	+	if err := r.SetModel(to); !errors.Is(err, ErrBusy) {
	20	+	t.Fatalf("SetModel during active turn = %v, want ErrBusy", err)
	21	+	}
	22	+	if r.model != from \|\| r.agent.Model != from \|\| r.agent.MaxTokens != 4096 {
	23	+	t.Fatalf("busy SetModel changed model or limit: runtime=%q agent=%q max=%d", r.model, r.agent.Model, r.agent.MaxTokens)
	24	+	}
	25	+
	26	+	r.activeCancel = nil
	27	+	if err := r.SetModel(to); err != nil {
	28	+	t.Fatalf("SetModel after turn: %v", err)
	29	+	}
	30	+	if r.model != to \|\| r.agent.Model != to {
	31	+	t.Fatalf("SetModel after turn: runtime=%q agent=%q, want %q", r.model, r.agent.Model, to)
	32	+	}
	33	+	}
	34	+
10	35		func TestRuntimeSetModelUpdatesMaxTokens(t *testing.T) {
11	36		const to = "claude-opus-4-1-20250805"
12	37		want, err := provider.FindModel("anthropic", to)

@@ -1,12 +1,37 @@
1	package sdk	1		package sdk
2		2
3	import (	3		import (
		4	+	"errors"
4	"testing"	5		"testing"
5		6
6	"github.com/patriceckhart/zot/packages/core"	7		"github.com/patriceckhart/zot/packages/core"
7	"github.com/patriceckhart/zot/packages/provider"	8		"github.com/patriceckhart/zot/packages/provider"
8	)	9		)
9		10
		11	+	func TestRuntimeSetModelRejectsBusyAgent(t *testing.T) {
		12	+	const from, to = "claude-3-5-haiku-20241022", "claude-opus-4-1-20250805"
		13	+	r := &Runtime{
		14	+	agent: &core.Agent{Model: from, MaxTokens: 4096},
		15	+	provider: "anthropic",
		16	+	model: from,
		17	+	activeCancel: func() {}, // Prompt and Compact both hold this until complete.
		18	+	}
		19	+	if err := r.SetModel(to); !errors.Is(err, ErrBusy) {
		20	+	t.Fatalf("SetModel during active turn = %v, want ErrBusy", err)
		21	+	}
		22	+	if r.model != from \|\| r.agent.Model != from \|\| r.agent.MaxTokens != 4096 {
		23	+	t.Fatalf("busy SetModel changed model or limit: runtime=%q agent=%q max=%d", r.model, r.agent.Model, r.agent.MaxTokens)
		24	+	}
		25	+
		26	+	r.activeCancel = nil
		27	+	if err := r.SetModel(to); err != nil {
		28	+	t.Fatalf("SetModel after turn: %v", err)
		29	+	}
		30	+	if r.model != to \|\| r.agent.Model != to {
		31	+	t.Fatalf("SetModel after turn: runtime=%q agent=%q, want %q", r.model, r.agent.Model, to)
		32	+	}
		33	+	}
		34	+
10	func TestRuntimeSetModelUpdatesMaxTokens(t *testing.T) {	35		func TestRuntimeSetModelUpdatesMaxTokens(t *testing.T) {
11	const to = "claude-opus-4-1-20250805"	36		const to = "claude-opus-4-1-20250805"
12	want, err := provider.FindModel("anthropic", to)	37		want, err := provider.FindModel("anthropic", to)