Vectrum
Sign In
337 lines · 11.0 KB · Go
package tui

import (
	"fmt"
	"path/filepath"
	"strings"
	"time"

	"github.com/charmbracelet/bubbles/textarea"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"

	"github.com/johannSo/codey/internal/agent"
	"github.com/johannSo/codey/internal/llm"
)

// The session chooser is deliberately scoped to the working root. Sessions are
// already filed under <root>/.codey/sessions, and every row is checked against
// the root a second time on the way out of the agent package — a conversation
// about one checkout has no business being reopened in another, where none of
// the paths it discusses exist.

// openSessionPicker reads the sessions belonging to this root and puts the
// chooser up. The list comes off local disk rather than out of the network, so
// it is read here and not in a command: there is no wait to render.
func (m *model) openSessionPicker() tea.Cmd {
	if why, ok := m.sessionSwitchBlocked(); !ok {
		m.tr.addError(why)
		return m.sync()
	}
	list, err := agent.ListSessions(m.opts.SessionDir, m.opts.Root)
	if err != nil {
		m.tr.addError("could not read saved sessions from " + m.opts.SessionDir + ": " + err.Error())
		return m.sync()
	}
	if len(list) == 0 {
		m.tr.addNotice("no saved sessions for " + shortenHome(m.opts.Root) +
			" — one is written after the first turn, and only this directory's are offered")
		return m.sync()
	}

	m.sessions = list
	m.sessionIdx = sessionIndex(list, m.opts.SessionID)
	m.phase = phaseSessions
	m.closeMenu()
	return m.sync()
}

// sessionSwitchBlocked says why the conversation cannot be swapped out right
// now. Switching replaces the history the agent is building on, so anything
// mid-flight has to finish or be interrupted first.
func (m *model) sessionSwitchBlocked() (string, bool) {
	switch {
	case m.opts.Agent == nil:
		return "no agent to resume into", false
	case m.phase == phaseRunning:
		return "/resume would replace the history the turn is building on — interrupt first", false
	case m.compacting:
		return "/resume would replace the history being compacted — interrupt first", false
	case m.worktreeing:
		return "/resume would land in the middle of moving to a worktree — interrupt first", false
	}
	return "", true
}

// sessionIndex opens the chooser on the session in force, so the list says
// where you already are rather than starting at the top.
func sessionIndex(list []agent.SessionInfo, current string) int {
	for i, s := range list {
		if s.ID == current {
			return i
		}
	}
	return 0
}

// onSessionKey drives the chooser.
func (m *model) onSessionKey(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	key := msg.String()
	if key == "esc" || key == "ctrl+c" {
		m.phase, m.sessions = phaseIdle, nil
		return m, tea.Batch(m.sync(), textarea.Blink)
	}
	if len(m.sessions) == 0 {
		return m, nil
	}

	last := len(m.sessions) - 1
	switch key {
	case "up", "ctrl+p":
		m.sessionIdx = maxInt(0, m.sessionIdx-1)
	case "down", "ctrl+n":
		m.sessionIdx = minInt(last, m.sessionIdx+1)
	case "home":
		m.sessionIdx = 0
	case "end":
		m.sessionIdx = last
	case "enter", "tab":
		id := m.sessions[m.sessionIdx].ID
		m.phase, m.sessions = phaseIdle, nil
		m.resumeSession(id)
		return m, tea.Batch(m.sync(), textarea.Blink)
	default:
		if len(key) == 1 && key[0] >= '1' && key[0] <= '9' {
			if i := int(key[0] - '1'); i <= last {
				m.sessionIdx = i
			}
		}
	}
	return m, m.sync()
}

// resumeSession loads a session by id and switches to it, refusing anything
// that was not recorded in this working root.
func (m *model) resumeSession(id string) {
	if why, ok := m.sessionSwitchBlocked(); !ok {
		m.tr.addError(why)
		return
	}
	if id == m.opts.SessionID {
		m.tr.addNotice("already in session " + id)
		return
	}
	sess, err := agent.LoadSession(filepath.Join(m.opts.SessionDir, id+".json"))
	if err != nil {
		m.tr.addError("no session " + id + " in " + shortenHome(m.opts.SessionDir) + " — /resume lists this directory's")
		return
	}
	if !sess.BelongsTo(m.opts.Root) {
		m.tr.addError("session " + id + " belongs to " + shortenHome(sess.Root) +
			", not " + shortenHome(m.opts.Root) + " — resume it from there")
		return
	}
	m.applySession(sess)
}

// applySession swaps the conversation over: the agent's history, the session
// the next turn is saved into, and the scrollback, which is replayed from the
// messages so the transcript shows the conversation being resumed rather than
// an empty screen above it.
func (m *model) applySession(s agent.Session) {
	// The session being left is written out first. A turn already saves on the
	// way out, but anything a subagent filed since then, or a compaction, has
	// not been — and this is the last moment it can be.
	if m.opts.SessionID != "" && m.turns > 0 {
		if err := m.saveSession(); err != nil {
			m.tr.addError("could not save session " + m.opts.SessionID + ": " + err.Error())
		}
	}

	m.opts.Agent.Restore(s)
	m.opts.SessionID = s.ID
	// The tab is named after the conversation in it, so resuming one renames
	// the tab it lands in.
	m.title = firstLine(s.Title())

	// The restored totals are what the session cost, including what earlier
	// runs already paid for; priced is set to match so the next turn is billed
	// on its own delta rather than on the whole history a second time.
	m.usage, m.turns, m.priced = s.Usage, s.Turns, s.Usage
	// What the model last read is unknown until it reads it, so the context
	// gauge waits for the next response rather than reporting the old one's.
	m.ctxTokens = 0

	m.clearSelection()
	m.tr.clear()
	m.replay(s.Messages)
	title := s.Title()
	if title == "" {
		title = "no prompt yet"
	}
	m.tr.addNotice(fmt.Sprintf("resumed %s — %s (%s, %s)",
		s.ID, title, plural(s.Turns, "turn"), plural(len(s.Messages), "message")))
	if s.Model != "" && s.Model != m.opts.Client.Model() {
		m.tr.addNotice("it ran on " + prettyModel(s.Model) + "; this session continues on " +
			prettyModel(m.opts.Client.Model()) + " — /model switches")
	}
}

// sessionTitle is what a saved session was about, for naming the tab it is
// opened in. A session that will not load has no title and no complaint: this
// is a label, and the conversation itself is where a broken file is reported.
func sessionTitle(dir, id string) string {
	if dir == "" || id == "" {
		return ""
	}
	s, err := agent.LoadSession(filepath.Join(dir, id+".json"))
	if err != nil {
		return ""
	}
	return firstLine(s.Title())
}

// replay rebuilds the scrollback from a saved conversation. It is a reading of
// the history, not a recording of what was on screen: thinking blocks and the
// exact streaming are gone, but what was said and what was run is all there.
func (m *model) replay(msgs []llm.Message) {
	// Tool results arrive in the following user message, so they are collected
	// first and then shown with the call they answer.
	results := map[string]llm.ContentBlock{}
	for _, msg := range msgs {
		for _, b := range msg.Content {
			if b.Type == "tool_result" {
				results[b.ToolUseID] = b
			}
		}
	}

	for _, msg := range msgs {
		for _, b := range msg.Content {
			switch b.Type {
			case "text", "":
				if strings.TrimSpace(b.Text) == "" {
					continue
				}
				if msg.Role == "user" {
					m.tr.addUser(b.Text)
				} else {
					m.tr.add(entry{kind: entryAssistant, text: b.Text})
				}
			case "tool_use":
				m.tr.startTool(b.Name, string(b.Input))
				res, ok := results[b.ID]
				if !ok {
					// A call whose result never came back: the session was
					// interrupted between the two.
					m.tr.finishTool(b.Name, "(no result recorded)", false)
					continue
				}
				m.tr.finishTool(b.Name, res.Content, res.IsError)
			}
		}
	}
	m.tr.closeOpen()
}

// sessionView is the chooser: a heading, the directory the list is scoped to,
// and a row per session with when it was last touched and what it was about.
func (m *model) sessionView() string {
	width := maxInt(1, m.mainWidth())
	head := " " + styHeading.Render("Resume session")
	from := styDim.Render("in " + shortenHome(m.opts.Root))

	start, end := m.sessionWindow()

	// A column for the ages and a column for the titles, so what each session
	// was and what it cost line up under each other rather than starting
	// wherever the row before them happened to end.
	ages := make([]string, len(m.sessions))
	titles := make([]string, len(m.sessions))
	ageCol, titleCol := 0, 0
	for i, s := range m.sessions {
		ages[i] = ago(time.Since(s.Updated))
		titles[i] = s.Title
		if titles[i] == "" {
			titles[i] = styFaint.Render("(no prompt)")
		}
		if i >= start && i < end {
			ageCol = maxInt(ageCol, lipgloss.Width(ages[i]))
			titleCol = maxInt(titleCol, lipgloss.Width(titles[i]))
		}
	}
	titleCol = minInt(titleCol, maxInt(20, width/2))

	var b strings.Builder
	b.WriteString(head + "\n")
	b.WriteString(truncateStyled(" "+from+styFaint.Render(" · ↑↓ move · enter resume · esc cancel"), width) + "\n")
	for i := start; i < end; i++ {
		s := m.sessions[i]
		marker, age := "  ", styDim.Render(pad(ages[i], ageCol))
		if i == m.sessionIdx {
			marker, age = styKey.Render("❯ "), styKey.Render(pad(ages[i], ageCol))
		}
		detail := plural(s.Turns, "turn")
		if s.Model != "" {
			detail += " · " + prettyModel(s.Model)
		}
		title := pad(truncateStyled(titles[i], titleCol), titleCol)
		line := marker + fmt.Sprintf("%d. ", i+1) + age + " " + title + " " + styFaint.Render(detail)
		if s.ID == m.opts.SessionID {
			line += styOK.Render(" ✓")
		}
		b.WriteString(truncateStyled(" "+line, width) + "\n")
	}
	if n := len(m.sessions) - (end - start); n > 0 {
		b.WriteString(truncateStyled("   "+styFaint.Render(fmt.Sprintf("… +%d more", n)), width) + "\n")
	}
	return strings.TrimRight(b.String(), "\n")
}

// sessionHeadRows is how many rows come before the list: the heading, and the
// line saying which directory the sessions are from.
const sessionHeadRows = 2

// sessionWindow is the slice of the list that is on screen, shared by the view
// and by the hit test, so a click lands on the session it looks like it is on.
func (m *model) sessionWindow() (start, end int) {
	rows := m.sessionRows()
	shown := rows
	if len(m.sessions) > rows {
		shown = maxInt(1, rows-1)
	}
	if m.sessionIdx >= shown {
		start = m.sessionIdx - shown + 1
	}
	return start, minInt(len(m.sessions), start+shown)
}

// sessionRows is how many session rows fit, by the same arithmetic the model
// picker does.
func (m *model) sessionRows() int {
	room := m.height - 6
	return maxInt(1, minInt(sessionMaxRows, room))
}

// sessionMaxRows caps the list. Beyond a dozen, scanning stops working and the
// id is the faster way in — which /resume <id> is for.
const sessionMaxRows = 12

// ago says how long ago something was in the shape a list wants: short, and
// coarser the further back it goes, because nobody picks a session by the
// minute it was saved three weeks later.
func ago(d time.Duration) string {
	switch {
	case d < 0:
		return "just now"
	case d < time.Minute:
		return "just now"
	case d < time.Hour:
		return fmt.Sprintf("%dm ago", int(d.Minutes()))
	case d < 24*time.Hour:
		return fmt.Sprintf("%dh ago", int(d.Hours()))
	case d < 7*24*time.Hour:
		return fmt.Sprintf("%dd ago", int(d.Hours()/24))
	}
	return fmt.Sprintf("%dw ago", int(d.Hours()/24/7))
}