fix(agent): prevent model switch deadlock

commit: 63ffe334b76eab6f0dfea7dfce116cd4f6b48ac7
author: brice · 2mo ago · 2026-08-05 03:10 UTC
committer: brice · 2mo ago · 2026-08-05 03:22 UTC
parent: bce908a
signature: GPG VerifiedCD88F9E564E0102646FB6077A1A07C7C7C3E349D · arch-gpg Verified for the committer email

2 files · +91 -46

M packages/agent/cli.go +55 -46


@@ -275,6 +275,52 @@ func liveInteractiveAgent(iv modes.Interactive, fallback core.Agent) *core.Age
275	275		return fallback
276	276		}
277	277
	278	+	func newSessionTransition(mu *sync.RWMutex) func(func()) {
	279	+	return func(fn func()) {
	280	+	mu.Lock()
	281	+	defer mu.Unlock()
	282	+	fn()
	283	+	}
	284	+	}
	285	+
	286	+	// newPersistModelCallback returns the callback used by the TUI while its
	287	+	// session transition write lock is held. It must not acquire that lock again
	288	+	// because sync.RWMutex is not reentrant.
	289	+	func newPersistModelCallback(persistMu sync.Mutex, sess core.Session, activeProvider, activeModel string, notify func(string)) func(string, string) {
	290	+	return func(providerName, model string) {
	291	+	providerName = strings.TrimSpace(providerName)
	292	+	model = strings.TrimSpace(model)
	293	+
	294	+	persistMu.Lock()
	295	+	var sessionErr error
	296	+	if *sess != nil {
	297	+	sessionErr = (*sess).UpdateModel(providerName, model)
	298	+	}
	299	+	*activeProvider = providerName
	300	+	*activeModel = model
	301	+	persistMu.Unlock()
	302	+	if sessionErr != nil {
	303	+	if notify != nil {
	304	+	notify("model change was not persisted: " + sessionErr.Error())
	305	+	}
	306	+	return
	307	+	}
	308	+
	309	+	cfg, err := LoadConfig()
	310	+	if err != nil {
	311	+	if notify != nil {
	312	+	notify("model changed for this run, but config could not be read: " + err.Error())
	313	+	}
	314	+	return
	315	+	}
	316	+	cfg.Provider = providerName
	317	+	cfg.Model = model
	318	+	if err := SaveConfig(cfg); err != nil && notify != nil {
	319	+	notify("model changed for this run, but config could not be saved: " + err.Error())
	320	+	}
	321	+	}
	322	+	}
	323	+
278	324		func closeAgentLSP(ag *core.Agent) {
279	325		if ag != nil {
280	326		_ = tools.CloseLSPManagers(ag.Tools)

@@ -2053,15 +2099,11 @@ func runInteractive(ctx context.Context, args Args, version string) error {
2053	2099		writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
2054	2100		sessBaselineMsgs = len(currentAg.Messages())
2055	2101		},
2056		-	SessionTransition: func(fn func()) {
2057		-	sessionTransitionMu.Lock()
2058		-	defer sessionTransitionMu.Unlock()
2059		-	fn()
2060		-	},
2061		-	Extensions: extMgr,
2062		-	Swarm: swarmMgr,
2063		-	ResolveSubagent: resolveSubagent,
2064		-	ChangelogChan: changelogCh,
	2102	+	SessionTransition: newSessionTransition(&sessionTransitionMu),
	2103	+	Extensions: extMgr,
	2104	+	Swarm: swarmMgr,
	2105	+	ResolveSubagent: resolveSubagent,
	2106	+	ChangelogChan: changelogCh,
2065	2107		OnChangelogDismiss: func() {
2066	2108		// For dev builds (0.0.0) store the actual release version
2067	2109		// so the same changelog doesn't show again next launch.

@@ -2092,44 +2134,11 @@ func runInteractive(ctx context.Context, args Args, version string) error {
2092	2134		},
2093	2135		NoYolo: args.NoYolo,
2094	2136		ConfirmGate: confirmGate,
2095		-	PersistModel: func(providerName, model string) {
2096		-	providerName = strings.TrimSpace(providerName)
2097		-	model = strings.TrimSpace(model)
2098		-	// Persist the active session metadata before publishing the
2099		-	// selection used by future resumes. The live agent has already
2100		-	// switched pairs, so keep the in-memory selection aligned even
2101		-	// when the append fails and report that persistence error below.
2102		-	sessionTransitionMu.RLock()
2103		-	defer sessionTransitionMu.RUnlock()
2104		-	persistMu.Lock()
2105		-	var sessionErr error
2106		-	if sess != nil {
2107		-	sessionErr = sess.UpdateModel(providerName, model)
2108		-	}
2109		-	activeProvider = providerName
2110		-	activeModel = model
2111		-	persistMu.Unlock()
2112		-	if sessionErr != nil {
2113		-	if iv != nil {
2114		-	iv.Notify("session", "error", "model change was not persisted: "+sessionErr.Error())
2115		-	}
2116		-	return
2117		-	}
2118		-
2119		-	// Update config.json so next launch uses the same pick.
2120		-	cfg, cfgErr := LoadConfig()
2121		-	if cfgErr != nil {
2122		-	if iv != nil {
2123		-	iv.Notify("session", "error", "model changed for this run, but config could not be read: "+cfgErr.Error())
2124		-	}
2125		-	return
2126		-	}
2127		-	cfg.Provider = providerName
2128		-	cfg.Model = model
2129		-	if cfgErr = SaveConfig(cfg); cfgErr != nil && iv != nil {
2130		-	iv.Notify("session", "error", "model changed for this run, but config could not be saved: "+cfgErr.Error())
	2137	+	PersistModel: newPersistModelCallback(&persistMu, &sess, &activeProvider, &activeModel, func(message string) {
	2138	+	if iv != nil {
	2139	+	iv.Notify("session", "error", message)
2131	2140		}
2132		-	},
	2141	+	}),
2133	2142		CurrentSessionTitle: func() string {
2134	2143		sessionTransitionMu.RLock()
2135	2144		defer sessionTransitionMu.RUnlock()

 	return fallback
 }
+func newSessionTransition(mu *sync.RWMutex) func(func()) {
+	return func(fn func()) {
+		mu.Lock()
+		defer mu.Unlock()
+		fn()
+	}
+}
+// newPersistModelCallback returns the callback used by the TUI while its
+// session transition write lock is held. It must not acquire that lock again
+// because sync.RWMutex is not reentrant.
+func newPersistModelCallback(persistMu *sync.Mutex, sess **core.Session, activeProvider, activeModel *string, notify func(string)) func(string, string) {
+	return func(providerName, model string) {
+		providerName = strings.TrimSpace(providerName)
+		model = strings.TrimSpace(model)
+		persistMu.Lock()
+		var sessionErr error
+		if *sess != nil {
+			sessionErr = (*sess).UpdateModel(providerName, model)
+		}
+		*activeProvider = providerName
+		*activeModel = model
+		persistMu.Unlock()
+		if sessionErr != nil {
+			if notify != nil {
+				notify("model change was not persisted: " + sessionErr.Error())
+			}
+			return
+		}
+		cfg, err := LoadConfig()
+		if err != nil {
+			if notify != nil {
+				notify("model changed for this run, but config could not be read: " + err.Error())
+			}
+			return
+		}
+		cfg.Provider = providerName
+		cfg.Model = model
+		if err := SaveConfig(cfg); err != nil && notify != nil {
+			notify("model changed for this run, but config could not be saved: " + err.Error())
+		}
+	}
+}
 func closeAgentLSP(ag *core.Agent) {
 	if ag != nil {
 		_ = tools.CloseLSPManagers(ag.Tools)
 			writeNewTranscriptLocked(currentAg, sess, sessBaselineMsgs)
 			sessBaselineMsgs = len(currentAg.Messages())
 		},
-		SessionTransition: func(fn func()) {
+		SessionTransition: newSessionTransition(&sessionTransitionMu),
-			sessionTransitionMu.Lock()
+		Extensions:        extMgr,
-			defer sessionTransitionMu.Unlock()
+		Swarm:             swarmMgr,
-			fn()
+		ResolveSubagent:   resolveSubagent,
-		},
+		ChangelogChan:     changelogCh,
-		Extensions:      extMgr,
-		Swarm:           swarmMgr,
-		ResolveSubagent: resolveSubagent,
-		ChangelogChan:   changelogCh,
 		OnChangelogDismiss: func() {
 			// For dev builds (0.0.0) store the actual release version
 			// so the same changelog doesn't show again next launch.
 		},
 		NoYolo:      args.NoYolo,
 		ConfirmGate: confirmGate,
-		PersistModel: func(providerName, model string) {
+		PersistModel: newPersistModelCallback(&persistMu, &sess, &activeProvider, &activeModel, func(message string) {
-			providerName = strings.TrimSpace(providerName)
+			if iv != nil {
-			model = strings.TrimSpace(model)
+				iv.Notify("session", "error", message)
-			// Persist the active session metadata before publishing the
-			// selection used by future resumes. The live agent has already
-			// switched pairs, so keep the in-memory selection aligned even
-			// when the append fails and report that persistence error below.
-			sessionTransitionMu.RLock()
-			defer sessionTransitionMu.RUnlock()
-			persistMu.Lock()
-			var sessionErr error
-			if sess != nil {
-				sessionErr = sess.UpdateModel(providerName, model)
-			}
-			activeProvider = providerName
-			activeModel = model
-			persistMu.Unlock()
-			if sessionErr != nil {
-				if iv != nil {
-					iv.Notify("session", "error", "model change was not persisted: "+sessionErr.Error())
-				}
-				return
-			}
-			// Update config.json so next launch uses the same pick.
-			cfg, cfgErr := LoadConfig()
-			if cfgErr != nil {
-				if iv != nil {
-					iv.Notify("session", "error", "model changed for this run, but config could not be read: "+cfgErr.Error())
-				}
-				return
-			}
-			cfg.Provider = providerName
-			cfg.Model = model
-			if cfgErr = SaveConfig(cfg); cfgErr != nil && iv != nil {
-				iv.Notify("session", "error", "model changed for this run, but config could not be saved: "+cfgErr.Error())
 			}
-		},
+		}),
 		CurrentSessionTitle: func() string {
 			sessionTransitionMu.RLock()
 			defer sessionTransitionMu.RUnlock()

M packages/agent/cli_session_test.go +36 -0


@@ -5,7 +5,9 @@ import (
5	5		"errors"
6	6		"fmt"
7	7		"os"
	8	+	"sync"
8	9		"testing"
	10	+	"time"
9	11
10	12		"github.com/patriceckhart/zot/packages/agent/modes"
11	13		"github.com/patriceckhart/zot/packages/core"

@@ -94,6 +96,40 @@ func TestLiveInteractiveAgentFallsBackBeforeInteractiveConstruction(t *testing.T
94	96		}
95	97		}
96	98
	99	+	func TestPersistModelCallbackDoesNotReenterSessionTransition(t *testing.T) {
	100	+	t.Setenv("ZOT_HOME", t.TempDir())
	101	+	sess, err := core.NewSession(t.TempDir(), t.TempDir(), "old-provider", "old-model", "test")
	102	+	if err != nil {
	103	+	t.Fatal(err)
	104	+	}
	105	+	defer sess.Close()
	106	+
	107	+	var persistMu sync.Mutex
	108	+	activeProvider, activeModel := "old-provider", "old-model"
	109	+	persistModel := newPersistModelCallback(&persistMu, &sess, &activeProvider, &activeModel, nil)
	110	+
	111	+	var transitionMu sync.RWMutex
	112	+	sessionTransition := newSessionTransition(&transitionMu)
	113	+	done := make(chan struct{})
	114	+	go func() {
	115	+	sessionTransition(func() {
	116	+	persistModel("new-provider", "new-model")
	117	+	})
	118	+	close(done)
	119	+	}()
	120	+	select {
	121	+	case <-done:
	122	+	case <-time.After(time.Second):
	123	+	t.Fatal("model persistence deadlocked inside the session transition")
	124	+	}
	125	+	if sess.Meta.Provider != "new-provider" \|\| sess.Meta.Model != "new-model" {
	126	+	t.Fatalf("session model = %q/%q, want new-provider/new-model", sess.Meta.Provider, sess.Meta.Model)
	127	+	}
	128	+	if activeProvider != "new-provider" \|\| activeModel != "new-model" {
	129	+	t.Fatalf("active model = %q/%q, want new-provider/new-model", activeProvider, activeModel)
	130	+	}
	131	+	}
	132	+
97	133		func firstMessageText(messages []provider.Message) string {
98	134		if len(messages) == 0 \|\| len(messages[0].Content) == 0 {
99	135		return ""


@@ -5,7 +5,9 @@ import (
5	"errors"	5		"errors"
6	"fmt"	6		"fmt"
7	"os"	7		"os"
		8	+	"sync"
8	"testing"	9		"testing"
		10	+	"time"
9		11
10	"github.com/patriceckhart/zot/packages/agent/modes"	12		"github.com/patriceckhart/zot/packages/agent/modes"
11	"github.com/patriceckhart/zot/packages/core"	13		"github.com/patriceckhart/zot/packages/core"

@@ -94,6 +96,40 @@ func TestLiveInteractiveAgentFallsBackBeforeInteractiveConstruction(t *testing.T
94	}	96		}
95	}	97		}
96		98
		99	+	func TestPersistModelCallbackDoesNotReenterSessionTransition(t *testing.T) {
		100	+	t.Setenv("ZOT_HOME", t.TempDir())
		101	+	sess, err := core.NewSession(t.TempDir(), t.TempDir(), "old-provider", "old-model", "test")
		102	+	if err != nil {
		103	+	t.Fatal(err)
		104	+	}
		105	+	defer sess.Close()
		106	+
		107	+	var persistMu sync.Mutex
		108	+	activeProvider, activeModel := "old-provider", "old-model"
		109	+	persistModel := newPersistModelCallback(&persistMu, &sess, &activeProvider, &activeModel, nil)
		110	+
		111	+	var transitionMu sync.RWMutex
		112	+	sessionTransition := newSessionTransition(&transitionMu)
		113	+	done := make(chan struct{})
		114	+	go func() {
		115	+	sessionTransition(func() {
		116	+	persistModel("new-provider", "new-model")
		117	+	})
		118	+	close(done)
		119	+	}()
		120	+	select {
		121	+	case <-done:
		122	+	case <-time.After(time.Second):
		123	+	t.Fatal("model persistence deadlocked inside the session transition")
		124	+	}
		125	+	if sess.Meta.Provider != "new-provider" \|\| sess.Meta.Model != "new-model" {
		126	+	t.Fatalf("session model = %q/%q, want new-provider/new-model", sess.Meta.Provider, sess.Meta.Model)
		127	+	}
		128	+	if activeProvider != "new-provider" \|\| activeModel != "new-model" {
		129	+	t.Fatalf("active model = %q/%q, want new-provider/new-model", activeProvider, activeModel)
		130	+	}
		131	+	}
		132	+
97	func firstMessageText(messages []provider.Message) string {	133		func firstMessageText(messages []provider.Message) string {
98	if len(messages) == 0 \|\| len(messages[0].Content) == 0 {	134		if len(messages) == 0 \|\| len(messages[0].Content) == 0 {
99	return ""	135		return ""