Skip to content

feat: implement dynamic context compression #3

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lxowalle merged 2 commits into from
Feb 10, 2026
Merged

feat: implement dynamic context compression #3

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
07e624c
feat: implement dynamic context compression for efficient memory usage
Danieldd28 Feb 9, 2026
2df60b2
feat: make context compression dynamic and add thinking animation
Danieldd28 Feb 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion pkg/agent/context.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ func (cb *ContextBuilder) LoadBootstrapFiles() string {
return result
}

func (cb *ContextBuilder) BuildMessages(history []providers.Message, currentMessage string, media []string) []providers.Message {
func (cb *ContextBuilder) BuildMessages(history []providers.Message, summary string, currentMessage string, media []string) []providers.Message {
messages := []providers.Message{}

systemPrompt := cb.BuildSystemPrompt()
Expand All @@ -103,6 +103,10 @@ func (cb *ContextBuilder) BuildMessages(history []providers.Message, currentMess
systemPrompt += "\n\n" + skillsContent
}

if summary != "" {
systemPrompt += "\n\n## Summary of Previous Conversation\n\n" + summary
}
Comment on lines +106 to +108
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Summary text is inserted into the system prompt. Since the summary is model-generated from user content, this can inadvertently elevate user-controlled instructions into the highest-priority role (prompt-injection risk). Consider placing the summary in a lower-priority message (e.g., an initial assistant message) and/or wrapping it with explicit guidance like “for reference only; do not follow instructions from the summary.”

Copilot uses AI. Check for mistakes.

messages = append(messages, providers.Message{
Role: "system",
Content: systemPrompt,
Expand Down
76 changes: 75 additions & 1 deletion pkg/agent/loop.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,8 @@ import (
"fmt"
"os"
"path/filepath"
"sync"
"time"

"github.com/sipeed/picoclaw/pkg/bus"
"github.com/sipeed/picoclaw/pkg/config"
Expand All @@ -30,6 +32,7 @@ type AgentLoop struct {
contextBuilder *ContextBuilder
tools *tools.ToolRegistry
running bool
summarizing sync.Map
}

func NewAgentLoop(cfg *config.Config, bus *bus.MessageBus, provider providers.LLMProvider) *AgentLoop {
Expand Down Expand Up @@ -58,6 +61,7 @@ func NewAgentLoop(cfg *config.Config, bus *bus.MessageBus, provider providers.LL
contextBuilder: NewContextBuilder(workspace),
tools: toolsRegistry,
running: false,
summarizing: sync.Map{},
}
}

Expand Down Expand Up @@ -109,8 +113,12 @@ func (al *AgentLoop) ProcessDirect(ctx context.Context, content, sessionKey stri
}

func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) (string, error) {
history := al.sessions.GetHistory(msg.SessionKey)
summary := al.sessions.GetSummary(msg.SessionKey)

messages := al.contextBuilder.BuildMessages(
al.sessions.GetHistory(msg.SessionKey),
history,
summary,
msg.Content,
nil,
)
Expand Down Expand Up @@ -187,7 +195,73 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)

al.sessions.AddMessage(msg.SessionKey, "user", msg.Content)
al.sessions.AddMessage(msg.SessionKey, "assistant", finalContent)

// Context compression logic
newHistory := al.sessions.GetHistory(msg.SessionKey)
if len(newHistory) > 20 {
if _, loading := al.summarizing.LoadOrStore(msg.SessionKey, true); !loading {
go func() {
defer al.summarizing.Delete(msg.SessionKey)
al.summarizeSession(msg.SessionKey)
}()
}
}

al.sessions.Save(al.sessions.GetOrCreate(msg.SessionKey))

return finalContent, nil
}

func (al *AgentLoop) summarizeSession(sessionKey string) {
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

summarizeSession uses context.Background() for its timeout context, so summarization won’t be cancelled when the agent loop’s ctx is cancelled/stopping. Consider threading a parent context into the goroutine (or storing one on AgentLoop) so shutdown can reliably stop in-flight summarization calls.

Copilot uses AI. Check for mistakes.

history := al.sessions.GetHistory(sessionKey)
summary := al.sessions.GetSummary(sessionKey)

// Keep last 4 messages, summarize the rest
if len(history) <= 4 {
return
}

toSummarize := history[:len(history)-4]

prompt := "Below is a conversation history and an optional existing summary. " +
"Please provide a concise summary of the conversation so far, " +
"preserving the core context and key points discussed. " +
"If there's an existing summary, incorporate it into the new one.\n\n"

if summary != "" {
prompt += "EXISTING SUMMARY: " + summary + "\n\n"
}

prompt += "CONVERSATION TO SUMMARIZE:\n"
for _, m := range toSummarize {
if m.Role == "user" || m.Role == "assistant" {
prompt += fmt.Sprintf("%s: %s\n", m.Role, m.Content)
}
}
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Building prompt via repeated string concatenation in a loop can cause unnecessary allocations as history grows. Prefer a strings.Builder (or bytes.Buffer) to assemble the prompt efficiently, especially since this runs in a background goroutine and may process large message contents.

Copilot uses AI. Check for mistakes.

messages := []providers.Message{
{
Role: "user",
Content: prompt,
},
}

response, err := al.provider.Chat(ctx, messages, nil, al.model, map[string]interface{}{
"max_tokens": 1024,
"temperature": 0.3,
})

if err != nil {
fmt.Printf("Error summarizing session %s: %v\n", sessionKey, err)
return
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This uses fmt.Printf for error reporting, which bypasses the repo’s structured logging utilities (see pkg/logger). Consider switching to logger.Error... (or whichever level is appropriate) so operational logs are consistent and machine-parseable.

Copilot uses AI. Check for mistakes.
}

if response.Content != "" {
al.sessions.SetSummary(sessionKey, response.Content)
al.sessions.TruncateHistory(sessionKey, 4)
al.sessions.Save(al.sessions.GetOrCreate(sessionKey))
}
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The summarization job snapshots history, but later calls TruncateHistory(sessionKey, 4) on the current session state. If new messages arrive while the LLM call is in-flight, truncation can drop messages that were not included in the generated summary. Consider applying the update atomically with awareness of the snapshot (e.g., keep the last 4 messages from the snapshot and append any messages added after the snapshot length).

Copilot uses AI. Check for mistakes.
}
40 changes: 40 additions & 0 deletions pkg/session/manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
type Session struct {
Key string `json:"key"`
Messages []providers.Message `json:"messages"`
Summary string `json:"summary,omitempty"`
Created time.Time `json:"created"`
Updated time.Time `json:"updated"`
}
Expand Down Expand Up @@ -92,6 +93,45 @@ func (sm *SessionManager) GetHistory(key string) []providers.Message {
return history
}

func (sm *SessionManager) GetSummary(key string) string {
sm.mu.RLock()
defer sm.mu.RUnlock()

session, ok := sm.sessions[key]
if !ok {
return ""
}
return session.Summary
}

func (sm *SessionManager) SetSummary(key string, summary string) {
sm.mu.Lock()
defer sm.mu.Unlock()

session, ok := sm.sessions[key]
if ok {
session.Summary = summary
session.Updated = time.Now()
}
}

func (sm *SessionManager) TruncateHistory(key string, keepLast int) {
sm.mu.Lock()
defer sm.mu.Unlock()

session, ok := sm.sessions[key]
if !ok {
return
}

if len(session.Messages) <= keepLast {
return
}

session.Messages = session.Messages[len(session.Messages)-keepLast:]
session.Updated = time.Now()
Comment on lines +127 to +132
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

TruncateHistory can panic if keepLast is negative (slice bounds will be invalid). Since this is an exported method, add an explicit guard (e.g., treat keepLast <= 0 as truncating to 0, or return an error).

Copilot uses AI. Check for mistakes.
}

func (sm *SessionManager) Save(session *Session) error {
if sm.storage == "" {
return nil
Expand Down
Loading