fix(tui): give broken streams a graceful stop
Retry temporary disconnects and allow one durable reset for protocol failures. Show terminal failures, stop polling, and keep drafts and pending work intact.
671 unmodified linesExpand all6726736746756766776786796806816826836846856846866876886896909 unmodified lines700701702703702704705706707708655 unmodified linesExpand all671 unmodified linesExpand all // a tick dropped under a modal would leave the face frozen for good sid, forward = sm.SessionID, true case ChatStatusMsg: sid, forward = sm.SessionID, true case ChatStreamResultMsg: sid, forward = sm.SessionID, true case ChatStreamClosedMsg: sid, forward = sm.SessionID, true case ChatInterruptMsg: sid, forward = sm.SessionID, true case ClipboardImagePastedMsg: sid, forward = sm.SessionID, true case ChatStreamEventMsg: sid, forward = sm.SessionID, true if m.ActiveSession != nil && sid == m.ActiveSession.ID && sm.Event.Type == daemon.EventUser && !sm.Event.Replayed { if m.ActiveSession != nil && sid == m.ActiveSession.ID && sm.Generation == m.Chat.Generation && sm.Event.Type == daemon.EventUser && !sm.Event.Replayed { m.ClearNotices() } case ChatOperationPollMsg: sid, forward = sm.SessionID, true9 unmodified lines var cmd tea.Cmd if m.ActiveSession != nil && sid == m.ActiveSession.ID { m.Chat, cmd = m.Chat.Update(msg) } if event, ok := msg.(ChatStreamEventMsg); ok && event.Event.Type == daemon.EventReset { if event, ok := msg.(ChatStreamEventMsg); ok && event.SessionID == m.Chat.SessionID && event.Generation == m.Chat.Generation && event.Event.Type == daemon.EventReset { cmd = tea.Batch(cmd, m.loadSettingsCmd(m.SettingsGen)) } return m, cmd }655 unmodified linesExpand all102 unmodified linesExpand all103104105106107108109110111112113114115116117118119120121122123124175 unmodified linesExpand all300301302303290304305306307308309142 unmodified linesExpand all45245345445544145645745845946024 unmodified lines485486487488489490491492493494495496430 unmodified linesExpand all92792892993091293193293393493511 unmodified lines9479489499509519329529539549559569579583 unmodified lines9629639649659669679689699709719729739749759769459469479489499779789799809529819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710089561009101010111012101310141015330 unmodified linesExpand all1346134713481349129613501351135213531354135513021356135713581359136015 unmodified lines137613771378137913261380138113821383138427 unmodified lines14121413141414151362141614171418141914202 unmodified lines142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446965 unmodified linesExpand all102 unmodified linesExpand all SessionID string Generation int64}
// ChatStreamResultMsg follows all accepted events from the same subscription.type ChatStreamResultMsg struct { SessionID string Generation int64 Err error Recovering bool}
type streamDelivery struct { Event *daemon.StreamEvent Err error Recovering bool}
type ChatStreamClosedMsg struct { SessionID string Generation int64}175 unmodified linesExpand all // inFlight is the last progress of the call the action row follows. It // outlives Progress, which the result clears, so the row can still say // what the call did once it ends. inFlight *daemon.ToolProgress eventChan chan daemon.StreamEvent eventChan chan streamDelivery streamStopped bool AttachedImage *daemon.ImageAttachment // window is the context window of windowModel, read once per model so // the footer can say how full the context is. window *int142 unmodified linesExpand all Styles: DefaultStyles, transcript: newTranscriptState(), streamCtx: ctx, streamCancel: cancel, eventChan: make(chan daemon.StreamEvent, 16), eventChan: make(chan streamDelivery, 16), }
return m}24 unmodified lines }}
func (m *ChatModel) ClearNotices() { // A successful mutation does not repair a stopped subscription. if m.streamStopped { return } if len(m.Notices) > 0 { m.Notices.Clear() m.syncViewportHeight() }430 unmodified linesExpand all return m.isSending || m.pendingSendCount() > 0 || m.Stopping || m.Progress != nil || m.Status.Running && !m.Status.Idle}
func (m *ChatModel) startAnimation() tea.Cmd { if m.animationActive || !m.animating() { if m.streamStopped || m.animationActive || !m.animating() { return nil } m.animationActive = true return m.progressTickCmd()11 unmodified lines case evt, ok := <-ch: if !ok { return ChatStreamClosedMsg{SessionID: sessID, Generation: gen} } if evt.Event != nil { return ChatStreamEventMsg{SessionID: sessID, Generation: gen, Event: evt} return ChatStreamEventMsg{SessionID: sessID, Generation: gen, Event: *evt.Event} } return ChatStreamResultMsg{SessionID: sessID, Generation: gen, Err: evt.Err, Recovering: evt.Recovering} case <-ctx.Done(): return ChatStreamClosedMsg{SessionID: sessID, Generation: gen} } }3 unmodified lines client, ctx, ch := m.client, m.streamCtx, m.eventChan go func() { defer close(ch) defer client.ResetStream() deliver := func(value streamDelivery) bool { select { case ch <- value: return true case <-ctx.Done(): return false } } recovered := false delay := 500 * time.Millisecond for ctx.Err() == nil { _ = client.Stream(ctx, olderPageRows, func(evt daemon.StreamEvent) error { select { case ch <- evt: return nil case <-ctx.Done(): err := client.StreamWithProgress(ctx, olderPageRows, func(evt daemon.StreamEvent) error { if !deliver(streamDelivery{Event: &evt}) { return ctx.Err() } }) return nil }, func() { delay = 500 * time.Millisecond }) if ctx.Err() != nil { return } var failure *daemon.StreamError if errors.As(err, &failure) { switch failure.Kind { case daemon.StreamProtocol: if recovered { deliver(streamDelivery{Err: err}) return } recovered = true if !deliver(streamDelivery{Err: err, Recovering: true}) { return } client.RequireStreamReset() case daemon.StreamTerminal: deliver(streamDelivery{Err: err}) return } } timer := time.NewTimer(delay) select { case <-ctx.Done(): timer.Stop() return case <-time.After(500 * time.Millisecond): case <-timer.C: } delay = min(delay*2, 5*time.Second) } }() return m.waitForNextEvent()}330 unmodified linesExpand all } return m, nil
case ChatStatusPollMsg: if msg.SessionID != m.SessionID || msg.Generation != m.Generation { if msg.SessionID != m.SessionID || msg.Generation != m.Generation || m.streamStopped { return m, nil } return m, m.statusCmd()
case ChatStatusMsg: if msg.SessionID != m.SessionID || msg.Generation != m.Generation { if msg.SessionID != m.SessionID || msg.Generation != m.Generation || m.streamStopped { return m, nil } if msg.Err == nil && msg.Status != nil && msg.Revision == m.statusRevision { m.Status = *msg.Status15 unmodified lines } return m, tea.Batch(m.startAnimation(), m.statusPollCmd())
case ChatStreamEventMsg: if msg.SessionID != m.SessionID || msg.Generation != m.Generation { if msg.SessionID != m.SessionID || msg.Generation != m.Generation || m.streamStopped { return m, nil } // only a live event outdates a status reply in flight if !msg.Event.Replayed {27 unmodified lines } return m, nil
case ChatProgressTickMsg: if msg.SessionID != m.SessionID || msg.Generation != m.Generation { if msg.SessionID != m.SessionID || msg.Generation != m.Generation || m.streamStopped { return m, nil } if !m.animating() { m.animationActive = false2 unmodified lines // the face lives in the status line, so the transcript stays as it is m.ProgressFrame++ return m, m.progressTickCmd()
case ChatStreamResultMsg: if msg.SessionID != m.SessionID || msg.Generation != m.Generation || m.streamStopped { return m, nil } m.Progress = nil m.clearAction() if msg.Recovering { m.AddNotice(fmt.Sprintf("Stream data could not be read; refreshing history: %v", msg.Err)) return m, m.waitForNextEvent() } m.streamStopped = true m.animationActive = false m.AddError(fmt.Sprintf("Session stream stopped: %v. Reopen the session to reconnect.", msg.Err)) m.refreshViewportContent() return m, nil
case ChatStreamClosedMsg: return m, nil
case ChatOperationPollMsg:965 unmodified linesExpand all123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286//go:build unix
// Stream faults are injected ahead of the shared real daemon; the TUI must// recover its durable transcript once, then stop visibly if the wire stays bad.package e2e
import ( "bufio" "context" "encoding/json" "errors" "fmt" "io" "net" "net/http" "net/http/httptest" "net/http/httputil" "net/url" "strings" "sync" "testing" "time"
"albedo/cli/internal/daemon" "albedo/cli/internal/tui" tea "charm.land/bubbletea/v2")
type streamFaultProxy struct { connection *daemon.Connection mu sync.Mutex cursors []string firstCursor int64 reconnected chan struct{}}
func faultSessionStream(t *testing.T, session string, fault string) *streamFaultProxy { t.Helper() destination, err := url.Parse(conn(t).BaseURL()) if err != nil { t.Fatal(err) } forward := httputil.NewSingleHostReverseProxy(destination) forward.FlushInterval = -1 proxy := &streamFaultProxy{reconnected: make(chan struct{}, 1)} server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/sessions/"+session+"/stream" { proxy.mu.Lock() proxy.cursors = append(proxy.cursors, r.URL.Query().Get("after_seq")) count := len(proxy.cursors) proxy.mu.Unlock() if fault == "eof" { if count == 1 { outgoing := r.Clone(r.Context()) outgoing.URL.Scheme, outgoing.URL.Host = destination.Scheme, destination.Host outgoing.RequestURI = "" response, err := conn(t).HTTPClient().Do(outgoing) if err != nil { http.Error(w, err.Error(), 502) return } defer response.Body.Close() scanner := bufio.NewScanner(response.Body) scanner.Buffer(make([]byte, 4096), 10*1024*1024) w.Header().Set("Content-Type", "text/event-stream") for scanner.Scan() { line := scanner.Text() _, _ = fmt.Fprintln(w, line) if strings.HasPrefix(line, "data: ") { var envelope struct { Cursor int64 `json:"cursor"` } if err := json.Unmarshal([]byte(strings.TrimPrefix(line, "data: ")), &envelope); err != nil { t.Error(err) } proxy.mu.Lock() proxy.firstCursor = envelope.Cursor proxy.mu.Unlock() } if line == "" { return } } if err := scanner.Err(); err != nil { t.Errorf("reading injected EOF frame: %v", err) } else { t.Error("daemon stream ended before a complete SSE frame") } return } select { case proxy.reconnected <- struct{}{}: default: } } else if count == 1 || fault == "repeated" { w.Header().Set("Content-Type", "text/event-stream") _, _ = io.WriteString(w, "data: "+`{"cursor":10,"events":[{"type":"text","text":"poison must stay hidden"},{"type":"message","role":"assistant","text":false}]}`+"\n\n") return } } forward.ServeHTTP(w, r) })) t.Cleanup(server.Close) snapshot := conn(t).Snapshot() snapshot.Port = server.Listener.Addr().(*net.TCPAddr).Port proxy.connection = daemon.NewConnection(snapshot, "") t.Cleanup(proxy.connection.HTTPClient().CloseIdleConnections) return proxy}
func (proxy *streamFaultProxy) requests() []string { proxy.mu.Lock() defer proxy.mu.Unlock() return append([]string(nil), proxy.cursors...)}
// Commands run independently, like Bubble Tea: one pending stream read must// not prevent the status result or a terminal notice from reaching the model.func streamMessages(ctx context.Context, cmd tea.Cmd, inbox chan<- tea.Msg) { if cmd == nil { return } go func() { msg := cmd() if batch, ok := msg.(tea.BatchMsg); ok { for _, next := range batch { streamMessages(ctx, next, inbox) } return } if msg != nil { select { case inbox <- msg: case <-ctx.Done(): } } }()}
func pumpStreamUntil(t *testing.T, driver *tuiDriver, condition func(tea.Msg) bool) tea.Msg { t.Helper() ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) defer cancel() inbox := make(chan tea.Msg, 32) streamMessages(ctx, driver.App.Chat.Init(), inbox) for { select { case msg := <-inbox: cmd := driver.Update(msg) if condition(msg) { return msg } streamMessages(ctx, cmd, inbox) case <-ctx.Done(): t.Fatalf("stream did not reach expected state:\n%s", driver.View()) } }}
func TestTUIStreamProtocolFailureRecoversDurableHistoryOnce(t *testing.T) { profile := providerRoute(t, echoReply) session := daemonSession(t, newSession(t, t.TempDir())) if _, err := daemon.NewChatClient(conn(t), session.ID).Send(t.Context(), "durable stream recovery", nil); err != nil { t.Fatal(err) } waitIdle(t, session.ID, profile, 1) proxy := faultSessionStream(t, session.ID, "once") driver := driveTUIWithConnection(t, &session, proxy.connection) t.Cleanup(driver.App.Chat.Close) driver.Update(tea.WindowSizeMsg{Width: 100, Height: 50}) recoverySeen := false pumpStreamUntil(t, driver, func(msg tea.Msg) bool { if result, ok := msg.(tui.ChatStreamResultMsg); ok && result.Recovering { recoverySeen = true } event, ok := msg.(tui.ChatStreamEventMsg) return ok && event.Event.Type == daemon.EventMessage && strings.Contains(event.Event.Text, "durable stream recovery") }) view := driver.View() if !recoverySeen || strings.Contains(view, "poison must stay hidden") || strings.Count(view, "durable stream recovery") != 2 { t.Fatalf("reset lost or duplicated durable transcript:\n%s", view) } requests := proxy.requests() if len(requests) != 2 || requests[1] != "-1" { t.Fatalf("protocol recovery did not request exactly one durable reset: %v", requests) }}
func TestTUIProtocolFailureStopsUntilTheSessionIsReopened(t *testing.T) { providerRoute(t, echoReply) session := daemonSession(t, newSession(t, t.TempDir())) proxy := faultSessionStream(t, session.ID, "repeated") driver := driveTUIWithConnection(t, &session, proxy.connection) t.Cleanup(func() { driver.App.Chat.Close() }) terminal := pumpStreamUntil(t, driver, func(msg tea.Msg) bool { result, ok := msg.(tui.ChatStreamResultMsg) return ok && !result.Recovering }) failure, ok := errors.AsType[*daemon.StreamError](terminal.(tui.ChatStreamResultMsg).Err) if !ok || failure.Kind != daemon.StreamProtocol || !driver.App.Chat.Notices.HasError() || strings.Contains(driver.View(), "poison must stay hidden") { t.Fatalf("protocol stop not visible or malformed text escaped:\n%s", driver.View()) } // Wait beyond both initial retry delays: a stopped attachment must stay stopped. time.Sleep(1600 * time.Millisecond) if requests := proxy.requests(); len(requests) != 2 { t.Fatalf("poison stream retried forever: %v", requests) } if cmd := driver.Update(tui.ChatStatusPollMsg{SessionID: session.ID, Generation: driver.App.Chat.Generation}); cmd != nil { t.Fatal("terminal attachment resumed status polling") } driver.Update(tui.FolderOpenSessionMsg{Session: session}) before := driver.View() if cmd := driver.Update(terminal); cmd != nil || driver.View() != before { t.Fatal("stale terminal failure affected the reopened session") }}
func TestTUIDeletedSessionStopsItsStreamAndStatusPolling(t *testing.T) { providerRoute(t, echoReply) session := daemonSession(t, newSession(t, t.TempDir())) if _, err := daemon.DeleteSession(t.Context(), conn(t), session.ID, false); err != nil { t.Fatal(err) } driver := driveTUI(t, &session) t.Cleanup(driver.App.Chat.Close) msg := pumpStreamUntil(t, driver, func(msg tea.Msg) bool { result, ok := msg.(tui.ChatStreamResultMsg) return ok && !result.Recovering }) result := msg.(tui.ChatStreamResultMsg) failure, ok := errors.AsType[*daemon.StreamError](result.Err) api, apiOK := errors.AsType[*daemon.APIError](result.Err) if !ok || failure.Kind != daemon.StreamTerminal || !apiOK || api.StatusCode != http.StatusNotFound { t.Fatalf("deleted session did not terminate: %v", result.Err) } if cmd := driver.Update(tui.ChatStatusPollMsg{SessionID: session.ID, Generation: driver.App.Chat.Generation}); cmd != nil { t.Fatal("deleted session kept polling") } if !driver.App.Chat.Notices.HasError() || !strings.Contains(driver.View(), api.Error()) { t.Fatalf("deleted session left no visible failure reason:\n%s", driver.View()) }}
func TestTUIStreamEOFReconnectsFromConsumedCursor(t *testing.T) { profile := providerRoute(t, echoReply) session := daemonSession(t, newSession(t, t.TempDir())) client := daemon.NewChatClient(conn(t), session.ID) if _, err := client.Send(t.Context(), "before EOF", nil); err != nil { t.Fatal(err) } waitIdle(t, session.ID, profile, 1) proxy := faultSessionStream(t, session.ID, "eof") driver := driveTUIWithConnection(t, &session, proxy.connection) t.Cleanup(driver.App.Chat.Close) driver.Update(tea.WindowSizeMsg{Width: 100, Height: 60}) sent := make(chan error, 1) go func() { select { case <-proxy.reconnected: _, err := client.Send(t.Context(), "after EOF", nil) sent <- err case <-t.Context().Done(): sent <- t.Context().Err() } }() pumpStreamUntil(t, driver, func(msg tea.Msg) bool { if result, ok := msg.(tui.ChatStreamResultMsg); ok { t.Errorf("ordinary EOF reported a terminal or protocol failure: %v", result.Err) } event, ok := msg.(tui.ChatStreamEventMsg) return ok && event.Event.Type == daemon.EventMessage && strings.Contains(event.Event.Text, "after EOF") }) if err := <-sent; err != nil { t.Fatal(err) } proxy.mu.Lock() requests, first := append([]string(nil), proxy.cursors...), proxy.firstCursor proxy.mu.Unlock() if len(requests) != 2 || requests[1] != fmt.Sprint(first) { t.Fatalf("EOF lost consumed cursor: requests=%v, cursor=%d", requests, first) } view := driver.View() if strings.Count(view, "before EOF") != 2 || strings.Count(view, "after EOF") != 2 { t.Fatalf("EOF replay lost or duplicated rows:\n%s", view) }}