fix(agent): harden session picker cancellation
- commit
- e7279a4c65ed4ab4b267c18a26136b808d5dd5c8
- author
- brice · · 2026-08-05 09:07 UTC
- parent
- e121559
- signature
- GPG Verified
CD88F9E564E0102646FB6077A1A07C7C7C3E349D· arch-gpg Verified for the committer email
4 files · +54 -4
M packages/agent/modes/interactive.go +7 -1
| @@ -1671,7 +1671,7 @@ func (i *Interactive) redraw() { | |||
| 1671 | 1671 | // in runSlash already handles the busy case per-command: safe | |
| 1672 | 1672 | // ones run immediately, destructive ones cancel the turn first. | |
| 1673 | 1673 | i.fileSuggest.SetCWD(i.cfg.CWD) | |
| 1674 | - | mainInputFocused := len(dialog) == 0 || (i.confirmDialog.Active() && !i.confirmDialog.Focused() && !i.confirmChildActive()) | |
| 1674 | + | mainInputFocused := len(dialog) == 0 || (i.confirmDialog.Active() && !i.confirmDialog.Focused() && !i.confirmChildActiveLocked()) | |
| 1675 | 1675 | if mainInputFocused && i.suggest.Active(currentInput) { | |
| 1676 | 1676 | suggest = i.suggest.Render(currentInput, i.cfg.Theme, cols) | |
| 1677 | 1677 | } else if mainInputFocused && i.fileSuggest.Active(currentInput) { | |
| @@ -2346,6 +2346,12 @@ func (i *Interactive) toggleBtwToolExpansion() { | |||
| 2346 | 2346 | // confirmChildActive reports whether an interaction opened from slash input | |
| 2347 | 2347 | // currently owns the keyboard while a tool confirmation remains pending. | |
| 2348 | 2348 | func (i *Interactive) confirmChildActive() bool { | |
| 2349 | + | i.mu.Lock() | |
| 2350 | + | defer i.mu.Unlock() | |
| 2351 | + | return i.confirmChildActiveLocked() | |
| 2352 | + | } | |
| 2353 | + | ||
| 2354 | + | func (i *Interactive) confirmChildActiveLocked() bool { | |
| 2349 | 2355 | return len(i.helpBlock) > 0 || | |
| 2350 | 2356 | len(i.extNotes) > 0 || | |
| 2351 | 2357 | i.dialog.Active() || | |
| @@ -1671,7 +1671,7 @@ func (i *Interactive) redraw() { | |||||
| 1671 | // in runSlash already handles the busy case per-command: safe | 1671 | // in runSlash already handles the busy case per-command: safe | ||
| 1672 | // ones run immediately, destructive ones cancel the turn first. | 1672 | // ones run immediately, destructive ones cancel the turn first. | ||
| 1673 | i.fileSuggest.SetCWD(i.cfg.CWD) | 1673 | i.fileSuggest.SetCWD(i.cfg.CWD) | ||
| 1674 | - | mainInputFocused := len(dialog) == 0 || (i.confirmDialog.Active() && !i.confirmDialog.Focused() && !i.confirmChildActive()) | 1674 | + | mainInputFocused := len(dialog) == 0 || (i.confirmDialog.Active() && !i.confirmDialog.Focused() && !i.confirmChildActiveLocked()) |
| 1675 | if mainInputFocused && i.suggest.Active(currentInput) { | 1675 | if mainInputFocused && i.suggest.Active(currentInput) { | ||
| 1676 | suggest = i.suggest.Render(currentInput, i.cfg.Theme, cols) | 1676 | suggest = i.suggest.Render(currentInput, i.cfg.Theme, cols) | ||
| 1677 | } else if mainInputFocused && i.fileSuggest.Active(currentInput) { | 1677 | } else if mainInputFocused && i.fileSuggest.Active(currentInput) { | ||
| @@ -2346,6 +2346,12 @@ func (i *Interactive) toggleBtwToolExpansion() { | |||||
| 2346 | // confirmChildActive reports whether an interaction opened from slash input | 2346 | // confirmChildActive reports whether an interaction opened from slash input | ||
| 2347 | // currently owns the keyboard while a tool confirmation remains pending. | 2347 | // currently owns the keyboard while a tool confirmation remains pending. | ||
| 2348 | func (i *Interactive) confirmChildActive() bool { | 2348 | func (i *Interactive) confirmChildActive() bool { | ||
| 2349 | + | i.mu.Lock() | |||
| 2350 | + | defer i.mu.Unlock() | |||
| 2351 | + | return i.confirmChildActiveLocked() | |||
| 2352 | + | } | |||
| 2353 | + | ||||
| 2354 | + | func (i *Interactive) confirmChildActiveLocked() bool { | |||
| 2349 | return len(i.helpBlock) > 0 || | 2355 | return len(i.helpBlock) > 0 || | ||
| 2350 | len(i.extNotes) > 0 || | 2356 | len(i.extNotes) > 0 || | ||
| 2351 | i.dialog.Active() || | 2357 | i.dialog.Active() || | ||
M packages/agent/modes/session_dialog.go +14 -3
| @@ -53,6 +53,7 @@ type sessionDialog struct { | |||
| 53 | 53 | loadingStartedAt time.Time | |
| 54 | 54 | loadGeneration uint64 | |
| 55 | 55 | loadCancel context.CancelFunc | |
| 56 | + | loadDone chan struct{} | |
| 56 | 57 | loadSlots []sessionLoadSlot | |
| 57 | 58 | ||
| 58 | 59 | // MaxRows is the maximum number of session rows the dialog | |
| @@ -89,6 +90,7 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||
| 89 | 90 | if d.loadCancel != nil { | |
| 90 | 91 | d.loadCancel() | |
| 91 | 92 | } | |
| 93 | + | previousDone := d.loadDone | |
| 92 | 94 | if parent == nil { | |
| 93 | 95 | parent = context.Background() | |
| 94 | 96 | } | |
| @@ -96,6 +98,8 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||
| 96 | 98 | d.loadCancel = cancel | |
| 97 | 99 | d.loadGeneration++ | |
| 98 | 100 | generation := d.loadGeneration | |
| 101 | + | done := make(chan struct{}) | |
| 102 | + | d.loadDone = done | |
| 99 | 103 | d.sessions = nil | |
| 100 | 104 | d.cursor = 0 | |
| 101 | 105 | d.viewTop = 0 | |
| @@ -114,6 +118,7 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||
| 114 | 118 | d.loading = false | |
| 115 | 119 | events := make(chan sessionLoadEvent) | |
| 116 | 120 | close(events) | |
| 121 | + | close(done) | |
| 117 | 122 | return events | |
| 118 | 123 | } | |
| 119 | 124 | ||
| @@ -127,7 +132,11 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||
| 127 | 132 | } | |
| 128 | 133 | } | |
| 129 | 134 | go func() { | |
| 135 | + | defer close(done) | |
| 130 | 136 | defer close(events) | |
| 137 | + | if previousDone != nil { | |
| 138 | + | <-previousDone | |
| 139 | + | } | |
| 131 | 140 | select { | |
| 132 | 141 | case <-ctx.Done(): | |
| 133 | 142 | return | |
| @@ -404,9 +413,11 @@ func formatSessionRowPlain(s core.SessionSummary, maxWidth int) string { | |||
| 404 | 413 | if summary == "" { | |
| 405 | 414 | summary = "(empty)" | |
| 406 | 415 | } | |
| 407 | - | summary = strings.ReplaceAll(summary, "\n", " ") | |
| 416 | + | provider := sanitizeSessionTreeText(s.Provider) | |
| 417 | + | model := sanitizeSessionTreeText(s.Model) | |
| 418 | + | summary = sanitizeSessionTreeText(summary) | |
| 408 | 419 | left := fmt.Sprintf("%-14s %s/%s %d msgs $%.4f ", | |
| 409 | - | when, s.Provider, s.Model, s.MessageCount, s.TotalCost) | |
| 420 | + | when, provider, model, s.MessageCount, s.TotalCost) | |
| 410 | 421 | room := maxWidth - len([]rune(left)) | |
| 411 | 422 | if room < 4 { | |
| 412 | 423 | room = 4 | |
| @@ -449,7 +460,7 @@ func (d *sessionDialog) loadingMessage(th tui.Theme) string { | |||
| 449 | 460 | if d.loadingTotal > 0 { | |
| 450 | 461 | progress = fmt.Sprintf("loading sessions (%d/%d)", d.loadingDone, d.loadingTotal) | |
| 451 | 462 | } | |
| 452 | - | return th.FG256(th.Spinner, frames[idx]) + " " + progress + " (esc cancel)" | |
| 463 | + | return th.FG256(th.Spinner, frames[idx]) + " " + th.FG256(th.Muted, progress+" (esc cancel)") | |
| 453 | 464 | } | |
| 454 | 465 | ||
| 455 | 466 | func formatRelative(t time.Time) string { | |
| @@ -53,6 +53,7 @@ type sessionDialog struct { | |||||
| 53 | loadingStartedAt time.Time | 53 | loadingStartedAt time.Time | ||
| 54 | loadGeneration uint64 | 54 | loadGeneration uint64 | ||
| 55 | loadCancel context.CancelFunc | 55 | loadCancel context.CancelFunc | ||
| 56 | + | loadDone chan struct{} | |||
| 56 | loadSlots []sessionLoadSlot | 57 | loadSlots []sessionLoadSlot | ||
| 57 | 58 | ||||
| 58 | // MaxRows is the maximum number of session rows the dialog | 59 | // MaxRows is the maximum number of session rows the dialog | ||
| @@ -89,6 +90,7 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||||
| 89 | if d.loadCancel != nil { | 90 | if d.loadCancel != nil { | ||
| 90 | d.loadCancel() | 91 | d.loadCancel() | ||
| 91 | } | 92 | } | ||
| 93 | + | previousDone := d.loadDone | |||
| 92 | if parent == nil { | 94 | if parent == nil { | ||
| 93 | parent = context.Background() | 95 | parent = context.Background() | ||
| 94 | } | 96 | } | ||
| @@ -96,6 +98,8 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||||
| 96 | d.loadCancel = cancel | 98 | d.loadCancel = cancel | ||
| 97 | d.loadGeneration++ | 99 | d.loadGeneration++ | ||
| 98 | generation := d.loadGeneration | 100 | generation := d.loadGeneration | ||
| 101 | + | done := make(chan struct{}) | |||
| 102 | + | d.loadDone = done | |||
| 99 | d.sessions = nil | 103 | d.sessions = nil | ||
| 100 | d.cursor = 0 | 104 | d.cursor = 0 | ||
| 101 | d.viewTop = 0 | 105 | d.viewTop = 0 | ||
| @@ -114,6 +118,7 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||||
| 114 | d.loading = false | 118 | d.loading = false | ||
| 115 | events := make(chan sessionLoadEvent) | 119 | events := make(chan sessionLoadEvent) | ||
| 116 | close(events) | 120 | close(events) | ||
| 121 | + | close(done) | |||
| 117 | return events | 122 | return events | ||
| 118 | } | 123 | } | ||
| 119 | 124 | ||||
| @@ -127,7 +132,11 @@ func (d *sessionDialog) Open(parent context.Context, root, cwd string) <-chan se | |||||
| 127 | } | 132 | } | ||
| 128 | } | 133 | } | ||
| 129 | go func() { | 134 | go func() { | ||
| 135 | + | defer close(done) | |||
| 130 | defer close(events) | 136 | defer close(events) | ||
| 137 | + | if previousDone != nil { | |||
| 138 | + | <-previousDone | |||
| 139 | + | } | |||
| 131 | select { | 140 | select { | ||
| 132 | case <-ctx.Done(): | 141 | case <-ctx.Done(): | ||
| 133 | return | 142 | return | ||
| @@ -404,9 +413,11 @@ func formatSessionRowPlain(s core.SessionSummary, maxWidth int) string { | |||||
| 404 | if summary == "" { | 413 | if summary == "" { | ||
| 405 | summary = "(empty)" | 414 | summary = "(empty)" | ||
| 406 | } | 415 | } | ||
| 407 | - | summary = strings.ReplaceAll(summary, "\n", " ") | 416 | + | provider := sanitizeSessionTreeText(s.Provider) |
| 417 | + | model := sanitizeSessionTreeText(s.Model) | |||
| 418 | + | summary = sanitizeSessionTreeText(summary) | |||
| 408 | left := fmt.Sprintf("%-14s %s/%s %d msgs $%.4f ", | 419 | left := fmt.Sprintf("%-14s %s/%s %d msgs $%.4f ", | ||
| 409 | - | when, s.Provider, s.Model, s.MessageCount, s.TotalCost) | 420 | + | when, provider, model, s.MessageCount, s.TotalCost) |
| 410 | room := maxWidth - len([]rune(left)) | 421 | room := maxWidth - len([]rune(left)) | ||
| 411 | if room < 4 { | 422 | if room < 4 { | ||
| 412 | room = 4 | 423 | room = 4 | ||
| @@ -449,7 +460,7 @@ func (d *sessionDialog) loadingMessage(th tui.Theme) string { | |||||
| 449 | if d.loadingTotal > 0 { | 460 | if d.loadingTotal > 0 { | ||
| 450 | progress = fmt.Sprintf("loading sessions (%d/%d)", d.loadingDone, d.loadingTotal) | 461 | progress = fmt.Sprintf("loading sessions (%d/%d)", d.loadingDone, d.loadingTotal) | ||
| 451 | } | 462 | } | ||
| 452 | - | return th.FG256(th.Spinner, frames[idx]) + " " + progress + " (esc cancel)" | 463 | + | return th.FG256(th.Spinner, frames[idx]) + " " + th.FG256(th.Muted, progress+" (esc cancel)") |
| 453 | } | 464 | } | ||
| 454 | 465 | ||||
| 455 | func formatRelative(t time.Time) string { | 466 | func formatRelative(t time.Time) string { | ||
M packages/agent/modes/session_dialog_test.go +16 -0
| @@ -4,12 +4,28 @@ import ( | |||
| 4 | 4 | "context" | |
| 5 | 5 | "strings" | |
| 6 | 6 | "testing" | |
| 7 | + | "unicode" | |
| 7 | 8 | ||
| 8 | 9 | "github.com/patriceckhart/zot/packages/core" | |
| 9 | 10 | "github.com/patriceckhart/zot/packages/provider" | |
| 10 | 11 | "github.com/patriceckhart/zot/packages/tui" | |
| 11 | 12 | ) | |
| 12 | 13 | ||
| 14 | + | func TestFormatSessionRowPlainSanitizesControlBytes(t *testing.T) { | |
| 15 | + | row := formatSessionRowPlain(core.SessionSummary{ | |
| 16 | + | Provider: "test\x1b]0;bad\a", | |
| 17 | + | Model: "model\x1b[31m", | |
| 18 | + | MessageCount: 1, | |
| 19 | + | FirstUserText: "hello\x1b[2J\nworld", | |
| 20 | + | }, 120) | |
| 21 | + | if strings.IndexFunc(row, unicode.IsControl) >= 0 { | |
| 22 | + | t.Fatalf("session row contains control characters: %q", row) | |
| 23 | + | } | |
| 24 | + | if !strings.Contains(row, "hello world") { | |
| 25 | + | t.Fatalf("sanitized session text missing: %q", row) | |
| 26 | + | } | |
| 27 | + | } | |
| 28 | + | ||
| 13 | 29 | func TestSessionDialogLoadsEntriesWithoutBlockingOpen(t *testing.T) { | |
| 14 | 30 | root := t.TempDir() | |
| 15 | 31 | cwd := t.TempDir() | |
| @@ -4,12 +4,28 @@ import ( | |||||
| 4 | "context" | 4 | "context" | ||
| 5 | "strings" | 5 | "strings" | ||
| 6 | "testing" | 6 | "testing" | ||
| 7 | + | "unicode" | |||
| 7 | 8 | ||||
| 8 | "github.com/patriceckhart/zot/packages/core" | 9 | "github.com/patriceckhart/zot/packages/core" | ||
| 9 | "github.com/patriceckhart/zot/packages/provider" | 10 | "github.com/patriceckhart/zot/packages/provider" | ||
| 10 | "github.com/patriceckhart/zot/packages/tui" | 11 | "github.com/patriceckhart/zot/packages/tui" | ||
| 11 | ) | 12 | ) | ||
| 12 | 13 | ||||
| 14 | + | func TestFormatSessionRowPlainSanitizesControlBytes(t *testing.T) { | |||
| 15 | + | row := formatSessionRowPlain(core.SessionSummary{ | |||
| 16 | + | Provider: "test\x1b]0;bad\a", | |||
| 17 | + | Model: "model\x1b[31m", | |||
| 18 | + | MessageCount: 1, | |||
| 19 | + | FirstUserText: "hello\x1b[2J\nworld", | |||
| 20 | + | }, 120) | |||
| 21 | + | if strings.IndexFunc(row, unicode.IsControl) >= 0 { | |||
| 22 | + | t.Fatalf("session row contains control characters: %q", row) | |||
| 23 | + | } | |||
| 24 | + | if !strings.Contains(row, "hello world") { | |||
| 25 | + | t.Fatalf("sanitized session text missing: %q", row) | |||
| 26 | + | } | |||
| 27 | + | } | |||
| 28 | + | ||||
| 13 | func TestSessionDialogLoadsEntriesWithoutBlockingOpen(t *testing.T) { | 29 | func TestSessionDialogLoadsEntriesWithoutBlockingOpen(t *testing.T) { | ||
| 14 | root := t.TempDir() | 30 | root := t.TempDir() | ||
| 15 | cwd := t.TempDir() | 31 | cwd := t.TempDir() | ||
M packages/core/session_snapshot_test.go +17 -0
| @@ -3,6 +3,7 @@ package core | |||
| 3 | 3 | import ( | |
| 4 | 4 | "context" | |
| 5 | 5 | "encoding/json" | |
| 6 | + | "errors" | |
| 6 | 7 | "os" | |
| 7 | 8 | "path/filepath" | |
| 8 | 9 | "strings" | |
| @@ -392,6 +393,22 @@ func TestHiddenBranchesAreTreeVisibleButFlatListHidden(t *testing.T) { | |||
| 392 | 393 | } | |
| 393 | 394 | } | |
| 394 | 395 | ||
| 396 | + | func TestForEachStrictJSONLLineContextStopsAfterCancellation(t *testing.T) { | |
| 397 | + | ctx, cancel := context.WithCancel(context.Background()) | |
| 398 | + | calls := 0 | |
| 399 | + | err := forEachStrictJSONLLineContext(ctx, strings.NewReader("first\nsecond\n"), func(line []byte, _ int) error { | |
| 400 | + | calls++ | |
| 401 | + | cancel() | |
| 402 | + | return nil | |
| 403 | + | }) | |
| 404 | + | if !errors.Is(err, context.Canceled) { | |
| 405 | + | t.Fatalf("context-aware JSONL read error = %v, want context.Canceled", err) | |
| 406 | + | } | |
| 407 | + | if calls != 1 { | |
| 408 | + | t.Fatalf("context-aware JSONL callback count = %d, want 1", calls) | |
| 409 | + | } | |
| 410 | + | } | |
| 411 | + | ||
| 395 | 412 | func TestDescribeSessionContextHonorsCancellation(t *testing.T) { | |
| 396 | 413 | root := t.TempDir() | |
| 397 | 414 | session, err := NewSession(root, "/workspace", "test", "model", "test") | |
| @@ -3,6 +3,7 @@ package core | |||||
| 3 | import ( | 3 | import ( | ||
| 4 | "context" | 4 | "context" | ||
| 5 | "encoding/json" | 5 | "encoding/json" | ||
| 6 | + | "errors" | |||
| 6 | "os" | 7 | "os" | ||
| 7 | "path/filepath" | 8 | "path/filepath" | ||
| 8 | "strings" | 9 | "strings" | ||
| @@ -392,6 +393,22 @@ func TestHiddenBranchesAreTreeVisibleButFlatListHidden(t *testing.T) { | |||||
| 392 | } | 393 | } | ||
| 393 | } | 394 | } | ||
| 394 | 395 | ||||
| 396 | + | func TestForEachStrictJSONLLineContextStopsAfterCancellation(t *testing.T) { | |||
| 397 | + | ctx, cancel := context.WithCancel(context.Background()) | |||
| 398 | + | calls := 0 | |||
| 399 | + | err := forEachStrictJSONLLineContext(ctx, strings.NewReader("first\nsecond\n"), func(line []byte, _ int) error { | |||
| 400 | + | calls++ | |||
| 401 | + | cancel() | |||
| 402 | + | return nil | |||
| 403 | + | }) | |||
| 404 | + | if !errors.Is(err, context.Canceled) { | |||
| 405 | + | t.Fatalf("context-aware JSONL read error = %v, want context.Canceled", err) | |||
| 406 | + | } | |||
| 407 | + | if calls != 1 { | |||
| 408 | + | t.Fatalf("context-aware JSONL callback count = %d, want 1", calls) | |||
| 409 | + | } | |||
| 410 | + | } | |||
| 411 | + | ||||
| 395 | func TestDescribeSessionContextHonorsCancellation(t *testing.T) { | 412 | func TestDescribeSessionContextHonorsCancellation(t *testing.T) { | ||
| 396 | root := t.TempDir() | 413 | root := t.TempDir() | ||
| 397 | session, err := NewSession(root, "/workspace", "test", "model", "test") | 414 | session, err := NewSession(root, "/workspace", "test", "model", "test") | ||