From 64596bf0c1fd65b2ddf0135256fca853e9774aa8 Mon Sep 17 00:00:00 2001 From: nguyenha935 Date: Mon, 29 Jun 2026 01:24:51 +0700 Subject: [PATCH] fix: preflight compact v3 context before model calls (#1297) Run pruning before ThinkStage so oversized histories compact before the provider request is built. Persist V3 last_prompt_tokens with the post-flush persisted message count so post-run compaction calibration has current usage data. Tests: go test ./internal/pipeline -run "TestNewDefaultPipeline_PrunesBeforeThink|TestFinalizeStage_UpdateMetadataReceivesPersistedMessageCount" -count=1; go test ./internal/agent -run "^$" -count=1 Co-authored-by: GoClaw Operator --- internal/agent/loop_pipeline_callbacks.go | 9 ++- internal/pipeline/deps.go | 2 +- internal/pipeline/finalize_stage.go | 7 ++- .../pipeline/finalize_stage_metadata_test.go | 46 ++++++++++++++ internal/pipeline/pipeline.go | 4 +- internal/pipeline/pipeline_test.go | 63 +++++++++++++++++++ 6 files changed, 123 insertions(+), 8 deletions(-) create mode 100644 internal/pipeline/finalize_stage_metadata_test.go diff --git a/internal/agent/loop_pipeline_callbacks.go b/internal/agent/loop_pipeline_callbacks.go index e2d2accc..35eb36f6 100644 --- a/internal/agent/loop_pipeline_callbacks.go +++ b/internal/agent/loop_pipeline_callbacks.go @@ -90,7 +90,7 @@ type pipelineCallbackSet struct { checkReadOnly func(state *pipeline.RunState) (*providers.Message, bool) sanitizeContent func(string) string flushMessages func(ctx context.Context, sessionKey string, msgs []providers.Message) error - updateMetadata func(ctx context.Context, sessionKey string, usage providers.Usage) error + updateMetadata func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error bootstrapCleanup func(ctx context.Context, state *pipeline.RunState) error maybeSummarize func(ctx context.Context, sessionKey string) } @@ -633,12 +633,15 @@ func (l *Loop) makeFlushMessages(req *RunRequest) func(ctx context.Context, sess } } -func (l *Loop) makeUpdateMetadata(req *RunRequest) func(ctx context.Context, sessionKey string, usage providers.Usage) error { - return func(ctx context.Context, sessionKey string, usage providers.Usage) error { +func (l *Loop) makeUpdateMetadata(req *RunRequest) func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error { + return func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error { l.sessions.UpdateMetadata(ctx, sessionKey, l.model, l.provider.Name(), req.Channel) l.sessions.AccumulateTokens(ctx, sessionKey, int64(usage.PromptTokens), int64(usage.CompletionTokens)) // Persist session to DB (matching v2 finalizeRun behavior). // FlushMessages already ran, so all pending messages are in the cache. + if usage.PromptTokens > 0 { + l.sessions.SetLastPromptTokens(ctx, sessionKey, usage.PromptTokens, msgCount) + } l.sessions.Save(ctx, sessionKey) return nil } diff --git a/internal/pipeline/deps.go b/internal/pipeline/deps.go index 9eaa0616..1d75ce1e 100644 --- a/internal/pipeline/deps.go +++ b/internal/pipeline/deps.go @@ -121,7 +121,7 @@ type PipelineDeps struct { DeduplicateMediaSuffix func(content, suffix string) string IsSilentReply func(content string) bool EmitSessionCompleted func(ctx context.Context, sessionKey string, msgCount, tokensUsed, compactionCount int) - UpdateMetadata func(ctx context.Context, sessionKey string, usage providers.Usage) error + UpdateMetadata func(ctx context.Context, sessionKey string, usage providers.Usage, msgCount int) error BootstrapCleanup func(ctx context.Context, state *RunState) error MaybeSummarize func(ctx context.Context, sessionKey string) } diff --git a/internal/pipeline/finalize_stage.go b/internal/pipeline/finalize_stage.go index 16a0dd6d..f58975bd 100644 --- a/internal/pipeline/finalize_stage.go +++ b/internal/pipeline/finalize_stage.go @@ -110,16 +110,19 @@ func (s *FinalizeStage) Execute(ctx context.Context, state *RunState) error { state.Messages.AppendPending(assistantMsg) // 4. Flush remaining pending messages to session store + historyCountBeforeFlush := len(state.Messages.History()) pending := state.Messages.FlushPending() + persistablePending := persistableMessages(pending) if len(pending) > 0 && s.deps.FlushMessages != nil { - if err := s.deps.FlushMessages(ctx, state.Input.SessionKey, persistableMessages(pending)); err != nil { + if err := s.deps.FlushMessages(ctx, state.Input.SessionKey, persistablePending); err != nil { slog.Warn("finalize flush failed", "err", err) } } // 5. Update session metadata (token usage) if s.deps.UpdateMetadata != nil { - if err := s.deps.UpdateMetadata(ctx, state.Input.SessionKey, state.Think.TotalUsage); err != nil { + msgCount := historyCountBeforeFlush + len(persistablePending) + if err := s.deps.UpdateMetadata(ctx, state.Input.SessionKey, state.Think.TotalUsage, msgCount); err != nil { slog.Warn("finalize metadata update failed", "err", err) } } diff --git a/internal/pipeline/finalize_stage_metadata_test.go b/internal/pipeline/finalize_stage_metadata_test.go new file mode 100644 index 00000000..09df55c5 --- /dev/null +++ b/internal/pipeline/finalize_stage_metadata_test.go @@ -0,0 +1,46 @@ +package pipeline + +import ( + "context" + "testing" + + "github.com/nextlevelbuilder/goclaw/internal/providers" +) + +func TestFinalizeStage_UpdateMetadataReceivesPersistedMessageCount(t *testing.T) { + t.Parallel() + + var gotUsage providers.Usage + var gotMsgCount int + deps := &PipelineDeps{ + FlushMessages: func(_ context.Context, _ string, _ []providers.Message) error { + return nil + }, + UpdateMetadata: func(_ context.Context, _ string, usage providers.Usage, msgCount int) error { + gotUsage = usage + gotMsgCount = msgCount + return nil + }, + } + + stage := NewFinalizeStage(deps) + state := defaultState() + state.Messages.SetHistory([]providers.Message{ + {Role: "user", Content: "old question"}, + {Role: "assistant", Content: "old answer"}, + }) + state.Messages.AppendPending(providers.Message{Role: "assistant", Content: "tool result"}) + state.Messages.AppendPending(providers.Message{Role: "assistant", Content: "transient", Transient: true}) + state.Observe.FinalContent = "final answer" + state.Think.TotalUsage = providers.Usage{PromptTokens: 1234, CompletionTokens: 56} + + if err := stage.Execute(context.Background(), state); err != nil { + t.Fatalf("Execute() error: %v", err) + } + if gotUsage.PromptTokens != 1234 || gotUsage.CompletionTokens != 56 { + t.Fatalf("UpdateMetadata usage = %+v, want prompt=1234 completion=56", gotUsage) + } + if gotMsgCount != 4 { + t.Fatalf("UpdateMetadata msgCount = %d, want 4 persisted messages", gotMsgCount) + } +} diff --git a/internal/pipeline/pipeline.go b/internal/pipeline/pipeline.go index dce26cc1..ae99716e 100644 --- a/internal/pipeline/pipeline.go +++ b/internal/pipeline/pipeline.go @@ -27,7 +27,7 @@ func NewPipeline(setup, iteration, finalize []Stage, deps PipelineDeps) *Pipelin } // NewDefaultPipeline creates the standard 8-stage pipeline. -// Setup: [ContextStage]. Iteration: [ThinkStage, PruneStage, ToolStage, ObserveStage, CheckpointStage]. +// Setup: [ContextStage]. Iteration: [PruneStage, ThinkStage, ToolStage, ObserveStage, CheckpointStage]. // Finalize: [FinalizeStage]. func NewDefaultPipeline(deps PipelineDeps) *Pipeline { d := &deps @@ -37,8 +37,8 @@ func NewDefaultPipeline(deps PipelineDeps) *Pipeline { NewContextStage(d), } iteration := []Stage{ - NewThinkStage(d), NewPruneStage(d, memFlush), + NewThinkStage(d), NewToolStage(d), NewObserveStage(d), NewCheckpointStage(d), diff --git a/internal/pipeline/pipeline_test.go b/internal/pipeline/pipeline_test.go index ab8ea8e3..bb0b7e4d 100644 --- a/internal/pipeline/pipeline_test.go +++ b/internal/pipeline/pipeline_test.go @@ -80,6 +80,69 @@ func TestPipeline_SetupRunsOnce(t *testing.T) { } } +func TestNewDefaultPipeline_PrunesBeforeThink(t *testing.T) { + t.Parallel() + + history := []providers.Message{ + {Role: "user", Content: "old 1"}, + {Role: "assistant", Content: "old 2"}, + {Role: "user", Content: "old 3"}, + {Role: "assistant", Content: "old 4"}, + {Role: "user", Content: "old 5"}, + {Role: "assistant", Content: "old 6"}, + {Role: "user", Content: "old 7"}, + {Role: "assistant", Content: "old 8"}, + {Role: "user", Content: "old 9"}, + {Role: "assistant", Content: "old 10"}, + } + compacted := []providers.Message{{Role: "assistant", Content: "compacted history"}} + var compactCalled bool + var llmMessages []providers.Message + + deps := PipelineDeps{ + Config: PipelineConfig{ + MaxIterations: 1, + ContextWindow: 1000, + MaxTokens: 100, + }, + TokenCounter: &mockTokenCounter{countPerMessage: 100}, + LoadSessionHistory: func(_ context.Context, _ string) ([]providers.Message, string) { + return history, "" + }, + BuildMessages: func(_ context.Context, _ *RunInput, loaded []providers.Message, _ string) ([]providers.Message, error) { + msgs := []providers.Message{{Role: "system", Content: "system"}} + msgs = append(msgs, loaded...) + return msgs, nil + }, + PruneMessages: func(msgs []providers.Message, _ int) ([]providers.Message, PruneStats) { + return msgs, PruneStats{} + }, + CompactMessages: func(_ context.Context, _ []providers.Message, _ string) ([]providers.Message, error) { + compactCalled = true + return compacted, nil + }, + CallLLM: func(_ context.Context, _ *RunState, req providers.ChatRequest) (*providers.ChatResponse, error) { + llmMessages = append([]providers.Message(nil), req.Messages...) + return &providers.ChatResponse{Content: "ok", FinishReason: "stop"}, nil + }, + } + + p := NewDefaultPipeline(deps) + _, err := p.Run(context.Background(), buildMinimalRunState()) + if err != nil { + t.Fatalf("Run() error: %v", err) + } + if !compactCalled { + t.Fatal("CompactMessages was not called") + } + if len(llmMessages) != 2 { + t.Fatalf("CallLLM received %d messages, want system + compacted history: %#v", len(llmMessages), llmMessages) + } + if llmMessages[1].Content != "compacted history" { + t.Fatalf("CallLLM second message = %#v, want compacted history", llmMessages[1]) + } +} + func TestPipeline_FinalizeRunsOnce(t *testing.T) { t.Parallel() finalize := newMockStageNoResult("finalize")