feat: recover corrupt sessions

Carlos Alexandro Becker created 3 months ago

Signed-off-by: Carlos Alexandro Becker <[e-mail]>

Change summary

internal/agent/agent.go              |  8 ++++
internal/agent/coordinator.go        | 53 ++++++++++++++++++++++++++++++
internal/tui/components/chat/chat.go |  9 +++++
3 files changed, 70 insertions(+)

Detailed changes

internal/agent/agent.go

@@ -621,6 +621,14 @@ func (a *sessionAgent) Summarize(ctx context.Context, sessionID string, opts fan
 			deleteErr := a.messages.Delete(ctx, summaryMessage.ID)
 			return deleteErr
 		}
+		// Summarization failed. Mark the message as finished with an error
+		// so the UI doesn't get stuck in a "Summarizing" state forever.
+		summaryMessage.FinishThinking()
+		summaryMessage.AddFinish(message.FinishReasonError, "Summarization failed", err.Error())
+		updateErr := a.messages.Update(ctx, summaryMessage)
+		if updateErr != nil {
+			return updateErr
+		}
 		return err
 	}
 

internal/agent/coordinator.go

@@ -54,6 +54,7 @@ type Coordinator interface {
 	QueuedPromptsList(sessionID string) []string
 	ClearQueue(sessionID string)
 	Summarize(context.Context, string) error
+	RecoverSession(ctx context.Context, sessionID string) error
 	Model() Model
 	UpdateModels(ctx context.Context) error
 }
@@ -834,6 +835,58 @@ func (c *coordinator) Summarize(ctx context.Context, sessionID string) error {
 	return c.currentAgent.Summarize(ctx, sessionID, getProviderOptions(c.currentAgent.Model(), providerCfg))
 }
 
+func (c *coordinator) RecoverSession(ctx context.Context, sessionID string) error {
+	// Skip recovery if session is currently active
+	if c.currentAgent.IsSessionBusy(sessionID) {
+		return nil
+	}
+
+	msgs, err := c.messages.List(ctx, sessionID)
+	if err != nil {
+		return fmt.Errorf("failed to list messages: %w", err)
+	}
+
+	for _, msg := range msgs {
+		if msg.IsFinished() {
+			continue
+		}
+
+		// Handle incomplete summary messages
+		if msg.IsSummaryMessage {
+			msg.FinishThinking()
+			msg.AddFinish(message.FinishReasonError, "Summarization interrupted", "Session was interrupted during summarization")
+			if updateErr := c.messages.Update(ctx, msg); updateErr != nil {
+				slog.Error("Failed to recover summary message", "message_id", msg.ID, "error", updateErr)
+			}
+			continue
+		}
+
+		// Handle incomplete assistant messages with tool calls
+		if msg.Role == message.Assistant && len(msg.ToolCalls()) > 0 {
+			// Mark any unfinished tool calls as finished
+			updated := false
+			for _, tc := range msg.ToolCalls() {
+				if !tc.Finished {
+					msg.FinishToolCall(tc.ID)
+					updated = true
+				}
+			}
+			if updated {
+				if updateErr := c.messages.Update(ctx, msg); updateErr != nil {
+					slog.Error("Failed to update tool calls", "message_id", msg.ID, "error", updateErr)
+				}
+			}
+			msg.FinishThinking()
+			msg.AddFinish(message.FinishReasonError, "Response interrupted", "Session was interrupted during tool execution")
+			if updateErr := c.messages.Update(ctx, msg); updateErr != nil {
+				slog.Error("Failed to recover assistant message", "message_id", msg.ID, "error", updateErr)
+			}
+		}
+	}
+
+	return nil
+}
+
 func (c *coordinator) isUnauthorized(err error) bool {
 	var providerErr *fantasy.ProviderError
 	return errors.As(err, &providerErr) && providerErr.StatusCode == http.StatusUnauthorized

internal/tui/components/chat/chat.go

@@ -2,6 +2,7 @@ package chat
 
 import (
 	"context"
+	"log/slog"
 	"time"
 
 	"charm.land/bubbles/v2/key"
@@ -544,6 +545,14 @@ func (m *messageListCmp) SetSession(session session.Session) tea.Cmd {
 	}
 
 	m.session = session
+
+	// Recover any incomplete messages from interrupted sessions
+	if m.app.AgentCoordinator != nil {
+		if recoverErr := m.app.AgentCoordinator.RecoverSession(context.Background(), session.ID); recoverErr != nil {
+			slog.Error("Failed to recover session", "session_id", session.ID, "error", recoverErr)
+		}
+	}
+
 	sessionMessages, err := m.app.Messages.List(context.Background(), session.ID)
 	if err != nil {
 		return util.ReportError(err)