From cd2e407b29d249ccced5767d40bec915d6fdb3ca Mon Sep 17 00:00:00 2001 From: Viet Tran Date: Wed, 11 Mar 2026 07:57:54 +0700 Subject: [PATCH] fix: auto-persist cleaned history when orphan tool messages detected (#128) sanitizeHistory now returns dropped count so callers know when orphaned tool_use/tool_result messages were removed. When orphans are found in buildMessages, the full session history is sanitized and persisted, preventing repeated warnings on every request. - Add SetHistory() to SessionStore interface and both implementations - Adapt memoryflush caller to new two-return signature - Change sanitize log level from Warn to Debug --- internal/agent/loop_history.go | 37 ++++++++++++++++++++++------- internal/agent/loop_history_test.go | 12 +++++----- internal/agent/memoryflush.go | 3 ++- internal/sessions/manager.go | 11 +++++++++ internal/store/pg/sessions_ops.go | 9 +++++++ internal/store/session_store.go | 1 + 6 files changed, 57 insertions(+), 16 deletions(-) diff --git a/internal/agent/loop_history.go b/internal/agent/loop_history.go index e26438a6..c3a0f892 100644 --- a/internal/agent/loop_history.go +++ b/internal/agent/loop_history.go @@ -151,7 +151,18 @@ func (l *Loop) buildMessages(ctx context.Context, history []providers.Message, s // History pipeline matching TS: limitHistoryTurns → pruneContext → sanitizeHistory. trimmed := limitHistoryTurns(history, historyLimit) pruned := pruneContextMessages(trimmed, l.contextWindow, l.contextPruningCfg) - messages = append(messages, sanitizeHistory(pruned)...) + sanitized, droppedCount := sanitizeHistory(pruned) + messages = append(messages, sanitized...) + + // If orphaned messages were found and dropped, persist the cleaned history + // back to the session store so the same orphans don't trigger on every request. + if droppedCount > 0 { + slog.Info("sanitizeHistory: cleaned session history", + "session", sessionKey, "dropped", droppedCount) + cleanedHistory, _ := sanitizeHistory(history) + l.sessions.SetHistory(sessionKey, cleanedHistory) + l.sessions.Save(sessionKey) + } // Current user message messages = append(messages, providers.Message{ @@ -270,21 +281,26 @@ func limitHistoryTurns(msgs []providers.Message, limit int) []providers.Message // - Orphaned tool messages at start of history (after truncation) // - tool_result without matching tool_use in preceding assistant message // - assistant with tool_calls but missing tool_results -func sanitizeHistory(msgs []providers.Message) []providers.Message { +// sanitizeHistory repairs tool_use/tool_result pairing in session history. +// Returns the cleaned messages and the number of messages that were dropped or synthesized. +func sanitizeHistory(msgs []providers.Message) ([]providers.Message, int) { if len(msgs) == 0 { - return msgs + return msgs, 0 } + dropped := 0 + // 1. Skip leading orphaned tool messages (no preceding assistant with tool_calls). start := 0 for start < len(msgs) && msgs[start].Role == "tool" { - slog.Warn("dropping orphaned tool message at history start", + slog.Debug("sanitizeHistory: dropping orphaned tool message at history start", "tool_call_id", msgs[start].ToolCallID) + dropped++ start++ } if start >= len(msgs) { - return nil + return nil, dropped } // 2. Walk through messages ensuring tool_result follows matching tool_use. @@ -309,30 +325,33 @@ func sanitizeHistory(msgs []providers.Message) []providers.Message { result = append(result, toolMsg) delete(expectedIDs, toolMsg.ToolCallID) } else { - slog.Warn("dropping mismatched tool result", + slog.Debug("sanitizeHistory: dropping mismatched tool result", "tool_call_id", toolMsg.ToolCallID) + dropped++ } } // Synthesize missing tool results for id := range expectedIDs { - slog.Warn("synthesizing missing tool result", "tool_call_id", id) + slog.Debug("sanitizeHistory: synthesizing missing tool result", "tool_call_id", id) result = append(result, providers.Message{ Role: "tool", Content: "[Tool result missing — session was compacted]", ToolCallID: id, }) + dropped++ } } else if msg.Role == "tool" { // Orphaned tool message mid-history (no preceding assistant with matching tool_calls) - slog.Warn("dropping orphaned tool message mid-history", + slog.Debug("sanitizeHistory: dropping orphaned tool message mid-history", "tool_call_id", msg.ToolCallID) + dropped++ } else { result = append(result, msg) } } - return result + return result, dropped } func (l *Loop) maybeSummarize(ctx context.Context, sessionKey string) { diff --git a/internal/agent/loop_history_test.go b/internal/agent/loop_history_test.go index b44680f2..0bedcc13 100644 --- a/internal/agent/loop_history_test.go +++ b/internal/agent/loop_history_test.go @@ -93,7 +93,7 @@ func TestLimitHistoryTurns_LimitExceedsTotal(t *testing.T) { } func TestSanitizeHistory_Empty(t *testing.T) { - got := sanitizeHistory(nil) + got, _ := sanitizeHistory(nil) if len(got) != 0 { t.Errorf("expected empty, got %d", len(got)) } @@ -106,7 +106,7 @@ func TestSanitizeHistory_DropsLeadingOrphanedTools(t *testing.T) { {Role: "user", Content: "hello"}, {Role: "assistant", Content: "hi"}, } - got := sanitizeHistory(msgs) + got, _ := sanitizeHistory(msgs) if len(got) != 2 { t.Fatalf("expected 2 messages, got %d", len(got)) } @@ -126,7 +126,7 @@ func TestSanitizeHistory_MatchesToolResults(t *testing.T) { {Role: "tool", Content: "written", ToolCallID: "tc2"}, {Role: "assistant", Content: "done"}, } - got := sanitizeHistory(msgs) + got, _ := sanitizeHistory(msgs) if len(got) != 5 { t.Fatalf("expected 5, got %d", len(got)) } @@ -143,7 +143,7 @@ func TestSanitizeHistory_SynthesizesMissingToolResult(t *testing.T) { // tc2 is missing {Role: "user", Content: "next"}, } - got := sanitizeHistory(msgs) + got, _ := sanitizeHistory(msgs) // user + assistant + tc1 result + synthesized tc2 result + user if len(got) != 5 { @@ -175,7 +175,7 @@ func TestSanitizeHistory_DropsMismatchedToolResult(t *testing.T) { {Role: "tool", Content: "stray", ToolCallID: "unknown_id"}, {Role: "user", Content: "next"}, } - got := sanitizeHistory(msgs) + got, _ := sanitizeHistory(msgs) // The stray tool message should be dropped, tc1 result kept for _, m := range got { @@ -192,7 +192,7 @@ func TestSanitizeHistory_DropsOrphanedToolMidHistory(t *testing.T) { {Role: "tool", Content: "orphan mid", ToolCallID: "tc_orphan"}, {Role: "user", Content: "bye"}, } - got := sanitizeHistory(msgs) + got, _ := sanitizeHistory(msgs) for _, m := range got { if m.ToolCallID == "tc_orphan" { diff --git a/internal/agent/memoryflush.go b/internal/agent/memoryflush.go index 8f4caad2..a870c501 100644 --- a/internal/agent/memoryflush.go +++ b/internal/agent/memoryflush.go @@ -132,7 +132,8 @@ func (l *Loop) runMemoryFlush(ctx context.Context, sessionKey string, settings * if len(recentHistory) > 10 { recentHistory = recentHistory[len(recentHistory)-10:] } - messages = append(messages, sanitizeHistory(recentHistory)...) + sanitized, _ := sanitizeHistory(recentHistory) + messages = append(messages, sanitized...) // Flush prompt messages = append(messages, providers.Message{ diff --git a/internal/sessions/manager.go b/internal/sessions/manager.go index 24765a6a..14ef8f19 100644 --- a/internal/sessions/manager.go +++ b/internal/sessions/manager.go @@ -278,6 +278,17 @@ func (m *Manager) TruncateHistory(key string, keepLast int) { s.Updated = time.Now() } +// SetHistory replaces a session's message history with the given slice. +func (m *Manager) SetHistory(key string, msgs []providers.Message) { + m.mu.Lock() + defer m.mu.Unlock() + + if s, ok := m.sessions[key]; ok { + s.Messages = msgs + s.Updated = time.Now() + } +} + // Reset clears a session's history and summary. func (m *Manager) Reset(key string) { m.mu.Lock() diff --git a/internal/store/pg/sessions_ops.go b/internal/store/pg/sessions_ops.go index 95dee2e1..5033b6de 100644 --- a/internal/store/pg/sessions_ops.go +++ b/internal/store/pg/sessions_ops.go @@ -19,6 +19,15 @@ func (s *PGSessionStore) TruncateHistory(key string, keepLast int) { } } +func (s *PGSessionStore) SetHistory(key string, msgs []providers.Message) { + s.mu.Lock() + defer s.mu.Unlock() + if data, ok := s.cache[key]; ok { + data.Messages = msgs + data.Updated = time.Now() + } +} + func (s *PGSessionStore) Reset(key string) { s.mu.Lock() defer s.mu.Unlock() diff --git a/internal/store/session_store.go b/internal/store/session_store.go index a510d8b0..5f4d0435 100644 --- a/internal/store/session_store.go +++ b/internal/store/session_store.go @@ -103,6 +103,7 @@ type SessionStore interface { SetLastPromptTokens(key string, tokens, msgCount int) GetLastPromptTokens(key string) (tokens, msgCount int) TruncateHistory(key string, keepLast int) + SetHistory(key string, msgs []providers.Message) Reset(key string) Delete(key string) error List(agentID string) []SessionInfo