kenn-io--agentsview/internal/db/messages_diff.go

// ABOUTME: Diff-based session message replacement: updates only the
// ABOUTME: rows that changed instead of delete+reinserting the session.
package db

import (
	"context"
	"database/sql"
	"fmt"
	"strings"
)

// diffDeleteChunkSize bounds IN(...) parameter lists when clearing
// tool rows for updated messages, mirroring the insert chunk limits.
const diffDeleteChunkSize = 500

// messageInsertArgs returns the values for insertMessageCols in
// declaration order. insertMessagesTx and the diff UPDATE share it
// so the persisted tuple has exactly one definition.
func messageInsertArgs(m Message) []any {
	return []any{
		m.SessionID, m.Ordinal, m.Role, m.Content,
		m.ThinkingText,
		m.Timestamp, m.HasThinking, m.HasToolUse,
		m.ContentLength, m.IsSystem,
		m.Model, string(m.TokenUsage),
		m.ContextTokens, m.OutputTokens,
		m.HasContextTokens, m.HasOutputTokens,
		m.ClaudeMessageID, m.ClaudeRequestID,
		m.SourceType, m.SourceSubtype, m.SourceUUID,
		m.SourceParentUUID, m.IsSidechain, m.IsCompactBoundary,
	}
}

// messageUpdateSetClause is the UPDATE assignment list derived from
// insertMessageCols, so the two can never drift apart.
var messageUpdateSetClause = func() string {
	cols := strings.Split(insertMessageCols, ",")
	parts := make([]string, 0, len(cols))
	for _, c := range cols {
		parts = append(parts, strings.TrimSpace(c)+" = ?")
	}
	return strings.Join(parts, ", ")
}()

// messageRowEqual reports whether two messages would persist the
// same message row, tool_calls rows, and tool_result_events rows.
// Both sides go through the same resolve helpers the insert path
// uses, so equality is defined on exactly the persisted tuples.
func messageRowEqual(a, b Message) bool {
	aArgs, bArgs := messageInsertArgs(a), messageInsertArgs(b)
	for i := range aArgs {
		if aArgs[i] != bArgs[i] {
			return false
		}
	}

	aCalls := resolveToolCalls([]Message{a}, []int64{0})
	bCalls := resolveToolCalls([]Message{b}, []int64{0})
	if len(aCalls) != len(bCalls) {
		return false
	}
	for i := range aCalls {
		if !toolCallRowEqual(aCalls[i], bCalls[i]) {
			return false
		}
	}

	aEvents := resolveToolResultEvents([]Message{a})
	bEvents := resolveToolResultEvents([]Message{b})
	if len(aEvents) != len(bEvents) {
		return false
	}
	for i := range aEvents {
		if aEvents[i] != bEvents[i] {
			return false
		}
	}
	return true
}

// toolCallRowEqual compares the persisted tool_calls columns except
// message_id (rowid-derived on both sides of a diff).
func toolCallRowEqual(a, b ToolCall) bool {
	return a.SessionID == b.SessionID &&
		a.ToolName == b.ToolName &&
		a.Category == b.Category &&
		a.ToolUseID == b.ToolUseID &&
		a.InputJSON == b.InputJSON &&
		a.SkillName == b.SkillName &&
		a.ResultContentLength == b.ResultContentLength &&
		a.ResultContent == b.ResultContent &&
		a.SubagentSessionID == b.SubagentSessionID &&
		a.FilePath == b.FilePath &&
		a.CallIndex == b.CallIndex
}

type messageDiffUpdate struct {
	id  int64
	msg Message
}

type messageDiffPlan struct {
	updates []messageDiffUpdate
	inserts []Message
}

// planStoredMessageDiff loads the session's stored messages and
// plans an in-place diff against msgs. ok=false means the caller
// must use the full replace path; a load failure also degrades to
// full replace rather than failing the write.
func (db *DB) planStoredMessageDiff(
	sessionID string, msgs []Message,
) (messageDiffPlan, []Message, bool, bool) {
	stored, err := db.GetAllMessages(
		context.Background(), sessionID,
	)
	if err != nil {
		return messageDiffPlan{}, nil, false, false
	}
	plan, useDiff := planSessionMessageDiff(stored, msgs)
	return plan, stored, useDiff, true
}

func transcriptMessagesEqual(stored, incoming []Message) bool {
	if len(stored) != len(incoming) {
		return false
	}
	byOrdinal := make(map[int]Message, len(stored))
	for _, msg := range stored {
		if _, exists := byOrdinal[msg.Ordinal]; exists {
			return false
		}
		byOrdinal[msg.Ordinal] = msg
	}
	seen := make(map[int]bool, len(incoming))
	for _, msg := range incoming {
		if seen[msg.Ordinal] {
			return false
		}
		seen[msg.Ordinal] = true
		old, exists := byOrdinal[msg.Ordinal]
		if !exists || !transcriptMessageEqual(old, msg) {
			return false
		}
	}
	return true
}

func transcriptMessageEqual(a, b Message) bool {
	comparableTranscriptMessage := func(msg Message) Message {
		msg.ID = 0
		msg.ContentLength = 0
		msg.TokenUsage = nil
		msg.ClaudeMessageID = ""
		msg.ClaudeRequestID = ""
		msg.SourceType = ""
		msg.SourceUUID = ""
		msg.SourceParentUUID = ""
		msg.IsSidechain = false
		msg.ToolCalls = append([]ToolCall(nil), msg.ToolCalls...)
		for i := range msg.ToolCalls {
			msg.ToolCalls[i].ResultContentLength = 0
			msg.ToolCalls[i].ResultEvents = append(
				[]ToolResultEvent(nil),
				msg.ToolCalls[i].ResultEvents...,
			)
			for j := range msg.ToolCalls[i].ResultEvents {
				msg.ToolCalls[i].ResultEvents[j].ContentLength = 0
			}
		}
		return msg
	}
	return messageRowEqual(
		comparableTranscriptMessage(a),
		comparableTranscriptMessage(b),
	)
}

// planSessionMessageDiff classifies incoming messages against stored
// rows by ordinal. It refuses (ok=false) when the diff cannot be
// applied safely or profitably:
//   - a stored ordinal is absent from the incoming set (truncation
//     or reordering needs delete handling and pin re-matching);
//   - duplicate ordinals on either side;
//   - a changed row's source_uuid differs from the stored row's
//     (the ordinal now holds a different message, so pins must be
//     re-matched by source_uuid — full replace's job);
//   - more than half the stored rows changed (an ordinal-shifting
//     rewrite, where full replace's source_uuid pin re-matching and
//     bulk FTS handling are the right tool).
func planSessionMessageDiff(
	stored, incoming []Message,
) (messageDiffPlan, bool) {
	if len(stored) == 0 {
		return messageDiffPlan{}, false
	}
	byOrdinal := make(map[int]Message, len(stored))
	for _, m := range stored {
		if _, dup := byOrdinal[m.Ordinal]; dup {
			return messageDiffPlan{}, false
		}
		byOrdinal[m.Ordinal] = m
	}

	var plan messageDiffPlan
	seen := make(map[int]bool, len(incoming))
	for _, m := range incoming {
		if seen[m.Ordinal] {
			return messageDiffPlan{}, false
		}
		seen[m.Ordinal] = true
		old, exists := byOrdinal[m.Ordinal]
		switch {
		case !exists:
			plan.inserts = append(plan.inserts, m)
		case !messageRowEqual(old, m):
			if old.SourceUUID != m.SourceUUID {
				return messageDiffPlan{}, false
			}
			plan.updates = append(plan.updates, messageDiffUpdate{
				id:  old.ID,
				msg: m,
			})
		}
	}
	for ord := range byOrdinal {
		if !seen[ord] {
			return messageDiffPlan{}, false
		}
	}
	if 2*len(plan.updates) > len(stored) {
		return messageDiffPlan{}, false
	}
	return plan, true
}

// applySessionMessageDiffTx persists a planned diff: changed rows
// are updated in place (keeping rowids, so pins survive and the FTS
// triggers reindex only those rows), their tool rows are rebuilt,
// and new ordinals are inserted through the normal insert path.
func applySessionMessageDiffTx(
	tx *sql.Tx, sessionID string, plan messageDiffPlan,
) error {
	if len(plan.updates) > 0 {
		updateSQL := "UPDATE messages SET " +
			messageUpdateSetClause + " WHERE id = ?"
		ids := make([]int64, 0, len(plan.updates))
		ordinals := make([]int, 0, len(plan.updates))
		msgs := make([]Message, 0, len(plan.updates))
		for _, u := range plan.updates {
			args := append(messageInsertArgs(u.msg), u.id)
			if _, err := tx.Exec(updateSQL, args...); err != nil {
				return fmt.Errorf(
					"updating message ord=%d: %w",
					u.msg.Ordinal, err,
				)
			}
			ids = append(ids, u.id)
			ordinals = append(ordinals, u.msg.Ordinal)
			msgs = append(msgs, u.msg)
		}
		if err := deleteToolRowsForMessagesTx(
			tx, sessionID, ids, ordinals,
		); err != nil {
			return err
		}
		if err := insertToolCallsTx(
			tx, resolveToolCalls(msgs, ids),
		); err != nil {
			return err
		}
		if err := insertToolResultEventsTx(
			tx, resolveToolResultEvents(msgs),
		); err != nil {
			return err
		}
	}

	if len(plan.inserts) > 0 {
		ids, err := insertMessagesTx(tx, plan.inserts)
		if err != nil {
			return err
		}
		if err := insertToolCallsTx(
			tx, resolveToolCalls(plan.inserts, ids),
		); err != nil {
			return err
		}
		if err := insertToolResultEventsTx(
			tx, resolveToolResultEvents(plan.inserts),
		); err != nil {
			return err
		}
	}
	return nil
}

// deleteToolRowsForMessagesTx clears tool_calls and
// tool_result_events for the updated messages so their rebuilt rows
// cannot duplicate the stale ones.
func deleteToolRowsForMessagesTx(
	tx *sql.Tx, sessionID string, ids []int64, ordinals []int,
) error {
	for start := 0; start < len(ids); start += diffDeleteChunkSize {
		end := min(start+diffDeleteChunkSize, len(ids))

		idArgs := make([]any, 0, end-start)
		for _, id := range ids[start:end] {
			idArgs = append(idArgs, id)
		}
		if _, err := tx.Exec(
			"DELETE FROM tool_calls WHERE message_id IN ("+
				placeholderList(len(idArgs))+")",
			idArgs...,
		); err != nil {
			return fmt.Errorf("deleting stale tool_calls: %w", err)
		}

		ordArgs := make([]any, 0, 1+end-start)
		ordArgs = append(ordArgs, sessionID)
		for _, ord := range ordinals[start:end] {
			ordArgs = append(ordArgs, ord)
		}
		if _, err := tx.Exec(
			"DELETE FROM tool_result_events WHERE session_id = ?"+
				" AND tool_call_message_ordinal IN ("+
				placeholderList(len(ordArgs)-1)+")",
			ordArgs...,
		); err != nil {
			return fmt.Errorf(
				"deleting stale tool_result_events: %w", err,
			)
		}
	}
	return nil
}

func placeholderList(n int) string {
	return strings.TrimSuffix(
		strings.Repeat("?,", n), ",",
	)
}