tiennm99/goclaw
1
0 Fork 0
mirror of https://github.com/tiennm99/goclaw.git synced 2026-10-14 04:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: preflight compact v3 context before model calls (#1297)

Browse Source 
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 <operator@goclaw>
This commit is contained in:
nguyenha935andGoClaw Operator authored and GitHub committed 2026-06-29 01:24:51 +07:00
1 parent 4675aa3fd4
commit 64596bf0c1
6 files changed
+123 -8
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

internal/agent/loop_pipeline_callbacks.go
+6 -3
View File
@@ -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
}
internal/pipeline/deps.go
+1 -1
View File
@@ -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)
}
internal/pipeline/finalize_stage.go
+5 -2
View File
@@ -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)
}
}
internal/pipeline/finalize_stage_metadata_test.go
+46
View File
@@ -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)
}
}
internal/pipeline/pipeline.go
+2 -2
View File
@@ -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),
internal/pipeline/pipeline_test.go
+63
View File
@@ -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")