okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.pettui: pace live reply redraws, cache settled markdown, draw only visible picker rows, drop per-event window reads6e2130b

main
6.8 KiB210 linesraw
1package tui
2
3import (
4	"albedo/cli/internal/daemon"
5	"errors"
6	"time"
7
8	tea "charm.land/bubbletea/v2"
9)
10
11func (m ChatModel) animating() bool {
12	return m.isSending || m.pendingSendCount() > 0 || m.Stopping || m.latestProgress() != nil || m.Status.Running && !m.Status.Idle || m.reaching() != "" || m.backgroundJobCount() > 0
13}
14
15func (m *ChatModel) startAnimation() tea.Cmd {
16	if m.streamStopped || m.animationActive || !m.animating() {
17		return nil
18	}
19	m.animationActive = true
20	return m.progressTickCmd()
21}
22
23func (m ChatModel) progressTickCmd() tea.Cmd {
24	id, generation := m.SessionID, m.Generation
25	return tea.Tick(faceInterval, func(time.Time) tea.Msg { return ChatProgressTickMsg{SessionID: id, Generation: generation} })
26}
27
28// onlyGrew is whether an event did nothing but lengthen the reply streaming
29// in, given the stream state and history length from before it.
30func (m ChatModel) onlyGrew(before transcriptState, entries int) bool {
31	after := m.transcript
32	return before.activeText != "" && len(after.activeText) > len(before.activeText) &&
33		after.activeKind == before.activeKind && after.activeMessageID == before.activeMessageID &&
34		m.History.Len() == entries
35}
36
37// drawLiveLater redraws the growing reply once its interval has passed.
38func (m *ChatModel) drawLiveLater() tea.Cmd {
39	if m.liveDue {
40		return nil
41	}
42	m.liveDue = true
43	id, generation := m.SessionID, m.Generation
44	return tea.Tick(time.Until(m.liveDrawn.Add(m.liveInterval)), func(time.Time) tea.Msg {
45		return ChatLiveDrawMsg{SessionID: id, Generation: generation}
46	})
47}
48
49func (m ChatModel) waitForNextEvent() tea.Cmd {
50	sessID, gen, ch, ctx := m.SessionID, m.Generation, m.eventChan, m.streamCtx
51	return func() tea.Msg {
52		select {
53		case evt, ok := <-ch:
54			if !ok {
55				return ChatStreamClosedMsg{SessionID: sessID, Generation: gen}
56			}
57			if evt.Event != nil {
58				return ChatStreamEventMsg{SessionID: sessID, Generation: gen, Event: *evt.Event}
59			}
60			return ChatStreamResultMsg{SessionID: sessID, Generation: gen, Err: evt.Err, Recovering: evt.Recovering}
61		case <-ctx.Done():
62			return ChatStreamClosedMsg{SessionID: sessID, Generation: gen}
63		}
64	}
65}
66
67func (m ChatModel) startStreamSubscription() tea.Cmd {
68	client, ctx, ch := m.client, m.streamCtx, m.eventChan
69	go func() {
70		defer close(ch)
71		defer client.ResetStream()
72		deliver := func(value streamDelivery) bool {
73			select {
74			case ch <- value:
75				return true
76			case <-ctx.Done():
77				return false
78			}
79		}
80		recovered := false
81		delay := 500 * time.Millisecond
82		for ctx.Err() == nil {
83			err := client.StreamWithProgress(ctx, olderPageRows, func(evt daemon.StreamEvent) error {
84				if !deliver(streamDelivery{Event: &evt}) {
85					return ctx.Err()
86				}
87				return nil
88			}, func() { delay = 500 * time.Millisecond })
89			if ctx.Err() != nil {
90				return
91			}
92			if failure, ok := errors.AsType[*daemon.StreamError](err); ok {
93				switch failure.Kind {
94				case daemon.StreamProtocol:
95					if recovered {
96						deliver(streamDelivery{Err: err})
97						return
98					}
99					recovered = true
100					if !deliver(streamDelivery{Err: err, Recovering: true}) {
101						return
102					}
103					client.ResetStream()
104				case daemon.StreamTerminal:
105					deliver(streamDelivery{Err: err})
106					return
107				}
108			}
109			timer := time.NewTimer(delay)
110			select {
111			case <-ctx.Done():
112				timer.Stop()
113				return
114			case <-timer.C:
115			}
116			delay = min(delay*2, 5*time.Second)
117		}
118	}()
119	return m.waitForNextEvent()
120}
121
122func (m ChatModel) clearCopyStatusCmd() tea.Cmd {
123	id, generation, revision := m.SessionID, m.Generation, m.copyStatusRevision
124	return tea.Tick(3*time.Second, func(time.Time) tea.Msg {
125		return ChatClearCopyStatusMsg{SessionID: id, Generation: generation, Revision: revision}
126	})
127}
128
129func (m ChatModel) statusCmd() tea.Cmd {
130	client, ctx, id, generation, revision := m.client, m.streamCtx, m.SessionID, m.Generation, m.statusRevision
131	return func() tea.Msg {
132		captured, err := client.GetSession(ctx)
133		return ChatStatusMsg{SessionID: id, Generation: generation, Revision: revision, Status: &captured.Status, Glances: captured.Glances, Err: err}
134	}
135}
136
137// windowCmd reads the context window when usage names a model whose window
138// has not been read yet. Only a usage event asks, so a failed read is tried
139// again at the next one rather than on every event.
140func (m ChatModel) windowCmd() tea.Cmd {
141	if m.Usage == nil || m.windowModel != nil && *m.windowModel == m.Usage.Model {
142		return nil
143	}
144	client, ctx, id, generation, model := m.client, m.streamCtx, m.SessionID, m.Generation, m.Usage.Model
145	return func() tea.Msg {
146		tokens, err := client.ContextWindow(ctx)
147		return ChatWindowMsg{SessionID: id, Generation: generation, Model: model, Tokens: tokens, Err: err}
148	}
149}
150
151// cacheFadeCmd wakes the footer when the cached count reaches its next step.
152func (m ChatModel) cacheFadeCmd() tea.Cmd {
153	if m.Usage == nil {
154		return nil
155	}
156	now := time.Now().UnixMilli()
157	for _, step := range m.Usage.CacheFade {
158		if step.At > now {
159			id, generation, usage := m.SessionID, m.Generation, m.Usage
160			return tea.Tick(time.Duration(step.At-now)*time.Millisecond, func(time.Time) tea.Msg {
161				return ChatCacheFadeMsg{SessionID: id, Generation: generation, Usage: usage}
162			})
163		}
164	}
165	return nil
166}
167
168// hostHomeCmd asks a remote workspace's host where its home is, once per
169// workspace, after the kernel attached: the probe has answered by then.
170func (m *ChatModel) hostHomeCmd() tea.Cmd {
171	host, _ := daemon.SplitLocation(m.Workspace)
172	if host == "" || m.hostAsked == m.Workspace || m.Status.KernelLink != "attached" {
173		return nil
174	}
175	m.hostAsked = m.Workspace
176	client, ctx, id, generation, workspace := m.client, m.streamCtx, m.SessionID, m.Generation, m.Workspace
177	return func() tea.Msg {
178		status, err := client.Host(ctx, host)
179		return ChatHostMsg{SessionID: id, Generation: generation, Workspace: workspace, Status: status, Err: err}
180	}
181}
182
183// statusPollCmd schedules the next status read: often while anything is live,
184// rarely while the session rests. Only the newest scheduled poll fires, so a
185// status read started elsewhere never adds a second polling loop.
186func (m *ChatModel) statusPollCmd() tea.Cmd {
187	m.statusPoll++
188	m.statusPollResting = !m.animating() && m.Status.KernelLink != "booting" && m.Status.KernelLink != "reattaching"
189	every := 750 * time.Millisecond
190	if m.statusPollResting {
191		every = 5 * time.Second
192	}
193	id, generation, poll := m.SessionID, m.Generation, m.statusPoll
194	return tea.Tick(every, func(time.Time) tea.Msg { return ChatStatusPollMsg{SessionID: id, Generation: generation, Poll: poll} })
195}
196
197func (m ChatModel) Init() tea.Cmd {
198	commands := []tea.Cmd{m.startStreamSubscription(), m.statusCmd()}
199	for _, pending := range m.pendingUsers {
200		if pending.Handle != nil && !pending.Expired {
201			commands = append(commands, m.queryOperationCmd(pending.Handle))
202		}
203	}
204	for _, pending := range m.pendingContinuations {
205		if !pending.Expired {
206			commands = append(commands, m.queryOperationCmd(pending.Handle))
207		}
208	}
209	return tea.Batch(commands...)
210}