zero/internal/tui/session.go

package tui

import (
	"encoding/json"
	"errors"
	"fmt"
	"strings"

	"github.com/Gitlawb/zero/internal/agent"
	"github.com/Gitlawb/zero/internal/sandbox"
	"github.com/Gitlawb/zero/internal/sessions"
	"github.com/Gitlawb/zero/internal/tools"
)

const tuiSessionTitleLimit = 80

type pendingSessionEvent struct {
	Type    sessions.EventType
	Payload any
}

func (m model) ensureActiveSession(prompt string) (model, error) {
	if m.activeSession.SessionID != "" {
		return m, nil
	}

	session, err := m.sessionStore.Create(sessions.CreateInput{
		Title:    tuiSessionTitle(prompt),
		Cwd:      m.cwd,
		ModelID:  m.modelName,
		Provider: m.providerName,
	})
	if err != nil {
		return m, err
	}
	m.activeSession = session
	m.sessionEvents = []sessions.Event{}
	return m, nil
}

func (m model) appendSessionEvent(eventType sessions.EventType, payload any) (model, error) {
	if m.activeSession.SessionID == "" {
		return m, nil
	}

	event, err := m.sessionStore.AppendEvent(m.activeSession.SessionID, sessions.AppendEventInput{
		Type:    eventType,
		Payload: payload,
	})
	if err != nil {
		return m, err
	}
	m.activeSession.UpdatedAt = event.CreatedAt
	m.activeSession.EventCount = event.Sequence
	m.activeSession.LastEventType = event.Type
	m.sessionEvents = append(m.sessionEvents, event)
	return m, nil
}

func (m model) appendSessionEvents(events []pendingSessionEvent) (model, []transcriptRow) {
	rows := []transcriptRow{}
	for _, event := range events {
		next, err := m.appendSessionEvent(event.Type, event.Payload)
		if err != nil {
			rows = append(rows, transcriptRow{kind: rowError, text: "session record error: " + err.Error()})
			continue
		}
		m = next
	}
	return m, rows
}

// appendSessionEventsTo persists events into a specific (non-active) session —
// the late flush of a run cancelled before a /resume switched sessions. The
// active session's in-memory metadata is deliberately untouched.
func (m model) appendSessionEventsTo(sessionID string, events []pendingSessionEvent) []transcriptRow {
	rows := []transcriptRow{}
	if m.sessionStore == nil || sessionID == "" {
		return rows
	}
	for _, event := range events {
		if _, err := m.sessionStore.AppendEvent(sessionID, sessions.AppendEventInput{
			Type:    event.Type,
			Payload: event.Payload,
		}); err != nil {
			rows = append(rows, transcriptRow{kind: rowError, text: "session record error: " + err.Error()})
		}
	}
	return rows
}

// flushableSessionEvents selects the events worth persisting from a run that was
// cancelled mid-flight. The cancel path already records a single "Run cancelled."
// error, so the goroutine's trailing EventError (the ctx-cancellation error) is
// dropped to avoid a duplicate; everything else it accumulated before the cancel
// — tool calls/results, permission events, usage, and the EventSessionCheckpoint
// blobs that /rewind depends on — is kept.
func flushableSessionEvents(events []pendingSessionEvent) []pendingSessionEvent {
	flushable := make([]pendingSessionEvent, 0, len(events))
	for _, event := range events {
		if event.Type == sessions.EventError {
			continue
		}
		flushable = append(flushable, event)
	}
	return flushable
}

func tuiSessionTitle(prompt string) string {
	// cutRunes keeps the cut on a rune boundary — a bare byte slice could split
	// a multi-byte rune and persist invalid UTF-8 into the session metadata.
	title := cutRunes(strings.Join(strings.Fields(prompt), " "), tuiSessionTitleLimit)
	if title == "" {
		return "Zero TUI session"
	}
	return title
}

func (m model) handleResumeCommand(args string) (model, string) {
	args = strings.TrimSpace(args)
	if args == "" {
		return m, m.resumeText("")
	}

	session, err := m.resolveResumeSession(args)
	if err != nil {
		return m, "Sessions\n" + err.Error()
	}
	events, err := m.sessionStore.ReadEvents(session.SessionID)
	if err != nil {
		return m, "Sessions\nerror: " + err.Error()
	}

	m.activeSession = *session
	m.sessionEvents = append([]sessions.Event{}, events...)
	if m.providerName == "" {
		m.providerName = session.Provider
	}
	if m.modelName == "" {
		m.modelName = session.ModelID
	}

	rows := initialTranscript()
	rows = appendRow(rows, rowSystem, formatResumeSummary(*session, len(events)))
	for _, row := range transcriptRowsFromSessionEvents(events) {
		rows = appendTranscriptRow(rows, row)
	}
	m.transcript = rows
	// Every rehydrated row is settled by construction, so resetting the flush
	// frontier sends the whole resumed history to native scrollback in one
	// batch — scrollable, selectable, and O(1) for every later frame.
	m.resetFlushFrontier("· resumed ·")
	return m, ""
}

func (m model) sessionPrompt(prompt string) string {
	if m.activeSession.SessionID == "" || len(m.sessionEvents) == 0 {
		return prompt
	}
	return sessions.FormatExecPrompt(prompt, sessions.PreparedExec{
		Mode:          sessions.ModeResume,
		Session:       m.activeSession,
		ContextEvents: append([]sessions.Event{}, m.sessionEvents...),
	})
}

func (m model) resolveResumeSession(args string) (*sessions.Metadata, error) {
	if strings.EqualFold(args, "latest") {
		// Latest *resumable* conversation, so "latest" never lands on a child or
		// spec sub-run. An explicit `/resume <id>` below still resolves any session.
		latest, err := m.sessionStore.LatestResumable()
		if err != nil {
			return nil, err
		}
		if latest == nil {
			return nil, errors.New("no zero sessions available to resume")
		}
		return latest, nil
	}

	session, err := m.sessionStore.Get(args)
	if err != nil {
		return nil, err
	}
	if session == nil {
		return nil, fmt.Errorf("zero session not found: %s", args)
	}
	return session, nil
}

func formatResumeSummary(session sessions.Metadata, eventCount int) string {
	return renderCommandOutput(commandOutput{
		Title:  "Resumed Zero session",
		Status: commandStatusOK,
		Sections: []commandSection{{
			Title: "Session",
			Lines: []string{
				"id: " + session.SessionID,
				"title: " + displayValue(session.Title, "untitled"),
				"model: " + displayValue(session.ModelID, "none"),
				"provider: " + displayValue(session.Provider, "none"),
				fmt.Sprintf("events: %d", eventCount),
			},
		}},
	})
}

func transcriptRowsFromSessionEvents(events []sessions.Event) []transcriptRow {
	rows := []transcriptRow{}
	// Rehydrated rows all carry runID 0, so repeated provider tool-call ids
	// (e.g. Gemini's per-turn gemini_tool_N) get the same per-occurrence
	// disambiguation the live runner applies — without it, dedup would drop
	// every tool card after the first occurrence of an id.
	callSeq := map[string]int{}
	for _, event := range events {
		payload := sessionPayload(event)
		switch event.Type {
		case sessions.EventMessage:
			role := strings.ToLower(payloadString(payload, "role"))
			switch role {
			case "ask_user":
				rows = append(rows, askUserTranscriptRow(askUserRequestFromPayload(payload)))
				continue
			case "ask_user_answers":
				if text := askUserAnswersText(payload); text != "" {
					rows = append(rows, transcriptRow{kind: rowSystem, text: text})
				}
				continue
			}
			content := payloadString(payload, "content")
			if content == "" {
				continue
			}
			switch role {
			case "user":
				rows = append(rows, transcriptRow{kind: rowUser, text: content})
			case "assistant":
				// A persisted assistant message was a turn's final answer, so it keeps
				// the final-answer rail on rehydration. Tool/timing counters were not
				// recorded; the done line omits those segments.
				rows = append(rows, transcriptRow{kind: rowAssistant, text: content, final: true})
			default:
				rows = append(rows, transcriptRow{kind: rowSystem, text: content})
			}
		case sessions.EventToolCall:
			name := payloadString(payload, "name")
			if name == "" {
				name = "unknown"
			}
			id := payloadString(payload, "id")
			callSeq[id]++
			rows = append(rows, transcriptRow{
				kind:   rowToolCall,
				id:     effectiveToolRowID(id, callSeq[id]),
				text:   "tool call: " + name,
				tool:   name,
				detail: argHint(payloadString(payload, "arguments")),
				arg:    argHintSecondary(payloadString(payload, "arguments")),
			})
		case sessions.EventPermission, sessions.EventPermissionRequest, sessions.EventPermissionDecision:
			rows = append(rows, permissionTranscriptRow(permissionEventFromPayload(payload)))
		case sessions.EventToolResult:
			name := payloadString(payload, "name")
			if name == "" {
				name = "unknown"
			}
			status := tools.Status(payloadString(payload, "status"))
			if status == "" {
				status = tools.StatusOK
			}
			output := payloadString(payload, "output")
			id := firstNonEmptyString(payloadString(payload, "toolCallId"), payloadString(payload, "id"))
			rows = append(rows, transcriptRow{
				kind:   rowToolResult,
				id:     effectiveToolRowID(id, callSeq[id]),
				text:   fmt.Sprintf("tool result: %s %s %s", name, status, truncateTUIOutput(output, tuiToolOutputLimit)),
				tool:   name,
				status: status,
				detail: output,
			})
		case sessions.EventError:
			if message := payloadString(payload, "message"); message != "" {
				rows = append(rows, transcriptRow{kind: rowError, text: message})
			}
		case sessions.EventCompaction:
			if summary := payloadString(payload, "summary"); summary != "" {
				rows = append(rows, transcriptRow{kind: rowSystem, text: summary})
			}
		case sessions.EventSessionFork:
			parentID := payloadString(payload, "parentSessionId")
			if parentID != "" {
				rows = append(rows, transcriptRow{kind: rowSystem, text: "forked from session: " + parentID})
			}
		}
	}
	return rows
}

func sessionPayload(event sessions.Event) map[string]any {
	payload := map[string]any{}
	if len(event.Payload) == 0 {
		return payload
	}
	_ = json.Unmarshal(event.Payload, &payload)
	return payload
}

func permissionEventFromPayload(payload map[string]any) agent.PermissionEvent {
	name := payloadString(payload, "name")
	if name == "" {
		name = payloadString(payload, "toolName")
	}
	event := agent.PermissionEvent{
		ToolCallID:        firstNonEmptyString(payloadString(payload, "toolCallId"), payloadString(payload, "id")),
		ToolName:          name,
		Action:            agent.PermissionAction(payloadString(payload, "action")),
		Permission:        payloadString(payload, "permission"),
		PermissionGranted: payloadBool(payload, "permissionGranted"),
		PermissionMode:    agent.PermissionMode(payloadString(payload, "permissionMode")),
		Autonomy:          payloadString(payload, "autonomy"),
		SideEffect:        payloadString(payload, "sideEffect"),
		Reason:            payloadString(payload, "reason"),
		DecisionReason:    payloadString(payload, "decisionReason"),
		GrantMatched:      payloadBool(payload, "grantMatched"),
	}
	if risk, ok := payloadMap(payload, "risk"); ok {
		event.Risk = sandbox.Risk{
			Level:  sandbox.RiskLevel(payloadString(risk, "level")),
			Reason: payloadString(risk, "reason"),
		}
	}
	if violation, ok := payloadMap(payload, "violation"); ok {
		event.Violation = &sandbox.Violation{
			Code:        sandbox.ViolationCode(payloadString(violation, "code")),
			ToolName:    payloadString(violation, "toolName"),
			Action:      sandbox.Action(payloadString(violation, "action")),
			Risk:        event.Risk,
			Path:        payloadString(violation, "path"),
			Reason:      payloadString(violation, "reason"),
			Recoverable: payloadBool(violation, "recoverable"),
		}
		if nestedRisk, ok := payloadMap(violation, "risk"); ok {
			event.Violation.Risk = sandbox.Risk{
				Level:  sandbox.RiskLevel(payloadString(nestedRisk, "level")),
				Reason: payloadString(nestedRisk, "reason"),
			}
		}
	}
	return event
}

// askUserRequestFromPayload rebuilds the request persisted by
// askUserSessionPayload, so ask_user exchanges survive /resume instead of
// silently vanishing from rehydrated history.
func askUserRequestFromPayload(payload map[string]any) agent.AskUserRequest {
	request := agent.AskUserRequest{
		ToolCallID: payloadString(payload, "toolCallId"),
		Header:     payloadString(payload, "header"),
	}
	raw, ok := payload["questions"].([]any)
	if !ok {
		return request
	}
	for _, entry := range raw {
		fields, ok := entry.(map[string]any)
		if !ok {
			continue
		}
		question := agent.AskUserQuestion{
			Question:    payloadString(fields, "question"),
			MultiSelect: payloadBool(fields, "multiSelect"),
		}
		if options, ok := fields["options"].([]any); ok {
			for _, option := range options {
				if text, ok := option.(string); ok {
					question.Options = append(question.Options, text)
				}
			}
		}
		request.Questions = append(request.Questions, question)
	}
	return request
}

// askUserAnswersText renders persisted ask_user answers for rehydration.
func askUserAnswersText(payload map[string]any) string {
	raw, ok := payload["answers"].([]any)
	if !ok {
		return ""
	}
	lines := make([]string, 0, len(raw))
	for index, entry := range raw {
		text, _ := entry.(string)
		if strings.TrimSpace(text) == "" {
			text = "(skipped)"
		}
		lines = append(lines, fmt.Sprintf("%d. %s", index+1, text))
	}
	if len(lines) == 0 {
		return ""
	}
	return "Answers\n" + strings.Join(lines, "\n")
}

func payloadString(payload map[string]any, key string) string {
	value := payload[key]
	switch typed := value.(type) {
	case string:
		return typed
	case float64, bool:
		return fmt.Sprint(typed)
	case nil:
		return ""
	default:
		data, err := json.Marshal(typed)
		if err != nil {
			return ""
		}
		return string(data)
	}
}

func payloadBool(payload map[string]any, key string) bool {
	value := payload[key]
	switch typed := value.(type) {
	case bool:
		return typed
	case string:
		return strings.EqualFold(strings.TrimSpace(typed), "true")
	default:
		return false
	}
}

func payloadMap(payload map[string]any, key string) (map[string]any, bool) {
	value, ok := payload[key].(map[string]any)
	return value, ok
}

func firstNonEmptyString(values ...string) string {
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			return strings.TrimSpace(value)
		}
	}
	return ""
}