opencode/picoclaw
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge branch 'main' into feat/mutil_keys

Browse source
This commit is contained in:
stark 2026-03-04 20:04:30 +08:00
commit dfd18da13a
20 changed files with 2422 additions and 122 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.env.example
View file

@ -17,4 +17,4 @@
# BRAVE_SEARCH_API_KEY=BSA... # BRAVE_SEARCH_API_KEY=BSA...
# ── Timezone ────────────────────────────── # ── Timezone ──────────────────────────────
TZ=Asia/Tokyo TZ=Asia/Shanghai

BIN
assets/wechat.png
View file

Binary file not shown.
Side by side Swipe Overlay

Before

Width:  |  Height:  |  Size: 140 KiB

After

Width:  |  Height:  |  Size: 96 KiB

1
config/config.example.json
View file

@ -59,6 +59,7 @@
"discord": { "discord": {
"enabled": false, "enabled": false,
"token": "YOUR_DISCORD_BOT_TOKEN", "token": "YOUR_DISCORD_BOT_TOKEN",
"proxy": "",
"allow_from": [], "allow_from": [],
"group_trigger": { "group_trigger": {
"mention_only": false "mention_only": false

78
pkg/agent/instance.go
View file

@ -18,22 +18,24 @@ import (
// AgentInstance represents a fully configured agent with its own workspace, // AgentInstance represents a fully configured agent with its own workspace,
// session manager, context builder, and tool registry. // session manager, context builder, and tool registry.
type AgentInstance struct { type AgentInstance struct {
ID string ID string
Name string Name string
Model string Model string
Fallbacks []string Fallbacks []string
Workspace string Workspace string
MaxIterations int MaxIterations int
MaxTokens int MaxTokens int
Temperature float64 Temperature float64
ContextWindow int ContextWindow int
Provider providers.LLMProvider SummarizeMessageThreshold int
Sessions *session.SessionManager SummarizeTokenPercent int
ContextBuilder *ContextBuilder Provider providers.LLMProvider
Tools *tools.ToolRegistry Sessions *session.SessionManager
Subagents *config.SubagentsConfig ContextBuilder *ContextBuilder
SkillsFilter []string Tools *tools.ToolRegistry
Candidates []providers.FallbackCandidate Subagents *config.SubagentsConfig
SkillsFilter []string
Candidates []providers.FallbackCandidate
} }
// NewAgentInstance creates an agent instance from config. // NewAgentInstance creates an agent instance from config.
@ -101,6 +103,16 @@ func NewAgentInstance(
temperature = *defaults.Temperature temperature = *defaults.Temperature
} }
summarizeMessageThreshold := defaults.SummarizeMessageThreshold
if summarizeMessageThreshold == 0 {
summarizeMessageThreshold = 20
}
summarizeTokenPercent := defaults.SummarizeTokenPercent
if summarizeTokenPercent == 0 {
summarizeTokenPercent = 75
}
// Resolve fallback candidates // Resolve fallback candidates
modelCfg := providers.ModelConfig{ modelCfg := providers.ModelConfig{
Primary: model, Primary: model,
@ -149,22 +161,24 @@ func NewAgentInstance(
candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList) candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)
return &AgentInstance{ return &AgentInstance{
ID: agentID, ID: agentID,
Name: agentName, Name: agentName,
Model: model, Model: model,
Fallbacks: fallbacks, Fallbacks: fallbacks,
Workspace: workspace, Workspace: workspace,
MaxIterations: maxIter, MaxIterations: maxIter,
MaxTokens: maxTokens, MaxTokens: maxTokens,
Temperature: temperature, Temperature: temperature,
ContextWindow: maxTokens, ContextWindow: maxTokens,
Provider: provider, SummarizeMessageThreshold: summarizeMessageThreshold,
Sessions: sessionsManager, SummarizeTokenPercent: summarizeTokenPercent,
ContextBuilder: contextBuilder, Provider: provider,
Tools: toolsRegistry, Sessions: sessionsManager,
Subagents: subagents, ContextBuilder: contextBuilder,
SkillsFilter: skillsFilter, Tools: toolsRegistry,
Candidates: candidates, Subagents: subagents,
SkillsFilter: skillsFilter,
Candidates: candidates,
} }
} }

111
pkg/agent/loop.go
View file

@ -118,6 +118,11 @@ func registerSharedTools(
PerplexityAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Perplexity.APIKey, cfg.Tools.Web.Perplexity.APIKeys), PerplexityAPIKeys: config.MergeAPIKeys(cfg.Tools.Web.Perplexity.APIKey, cfg.Tools.Web.Perplexity.APIKeys),
PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults, PerplexityMaxResults: cfg.Tools.Web.Perplexity.MaxResults,
PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled, PerplexityEnabled: cfg.Tools.Web.Perplexity.Enabled,
GLMSearchAPIKey: cfg.Tools.Web.GLMSearch.APIKey,
GLMSearchBaseURL: cfg.Tools.Web.GLMSearch.BaseURL,
GLMSearchEngine: cfg.Tools.Web.GLMSearch.SearchEngine,
GLMSearchMaxResults: cfg.Tools.Web.GLMSearch.MaxResults,
GLMSearchEnabled: cfg.Tools.Web.GLMSearch.Enabled,
Proxy: cfg.Tools.Web.Proxy, Proxy: cfg.Tools.Web.Proxy,
}) })
if err != nil { if err != nil {
@ -964,62 +969,76 @@ func (al *AgentLoop) runLLMIteration(
// Save assistant message with tool calls to session // Save assistant message with tool calls to session
agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg) agent.Sessions.AddFullMessage(opts.SessionKey, assistantMsg)
// Execute tool calls // Execute tool calls in parallel
for _, tc := range normalizedToolCalls { type indexedAgentResult struct {
argsJSON, _ := json.Marshal(tc.Arguments) result *tools.ToolResult
argsPreview := utils.Truncate(string(argsJSON), 200) tc providers.ToolCall
logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), }
map[string]any{
"agent_id": agent.ID,
"tool": tc.Name,
"iteration": iteration,
})
// Create async callback for tools that implement AsyncTool agentResults := make([]indexedAgentResult, len(normalizedToolCalls))
// NOTE: Following openclaw's design, async tools do NOT send results directly to users. var wg sync.WaitGroup
// Instead, they notify the agent via PublishInbound, and the agent decides
// whether to forward the result to the user (in processSystemMessage). for i, tc := range normalizedToolCalls {
asyncCallback := func(callbackCtx context.Context, result *tools.ToolResult) { agentResults[i].tc = tc
// Log the async completion but don't send directly to user
// The agent will handle user notification via processSystemMessage wg.Add(1)
if !result.Silent && result.ForUser != "" { go func(idx int, tc providers.ToolCall) {
logger.InfoCF("agent", "Async tool completed, agent will handle notification", defer wg.Done()
map[string]any{
"tool": tc.Name, argsJSON, _ := json.Marshal(tc.Arguments)
"content_len": len(result.ForUser), argsPreview := utils.Truncate(string(argsJSON), 200)
}) logger.InfoCF("agent", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview),
map[string]any{
"agent_id": agent.ID,
"tool": tc.Name,
"iteration": iteration,
})
// Create async callback for tools that implement AsyncTool
asyncCallback := func(callbackCtx context.Context, result *tools.ToolResult) {
if !result.Silent && result.ForUser != "" {
logger.InfoCF("agent", "Async tool completed, agent will handle notification",
map[string]any{
"tool": tc.Name,
"content_len": len(result.ForUser),
})
}
} }
}
toolResult := agent.Tools.ExecuteWithContext( toolResult := agent.Tools.ExecuteWithContext(
ctx, ctx,
tc.Name, tc.Name,
tc.Arguments, tc.Arguments,
opts.Channel, opts.Channel,
opts.ChatID, opts.ChatID,
asyncCallback, asyncCallback,
) )
agentResults[idx].result = toolResult
}(i, tc)
}
wg.Wait()
// Process results in original order (send to user, save to session)
for _, r := range agentResults {
// Send ForUser content to user immediately if not Silent // Send ForUser content to user immediately if not Silent
if !toolResult.Silent && toolResult.ForUser != "" && opts.SendResponse { if !r.result.Silent && r.result.ForUser != "" && opts.SendResponse {
al.bus.PublishOutbound(ctx, bus.OutboundMessage{ al.bus.PublishOutbound(ctx, bus.OutboundMessage{
Channel: opts.Channel, Channel: opts.Channel,
ChatID: opts.ChatID, ChatID: opts.ChatID,
Content: toolResult.ForUser, Content: r.result.ForUser,
}) })
logger.DebugCF("agent", "Sent tool result to user", logger.DebugCF("agent", "Sent tool result to user",
map[string]any{ map[string]any{
"tool": tc.Name, "tool": r.tc.Name,
"content_len": len(toolResult.ForUser), "content_len": len(r.result.ForUser),
}) })
} }
// If tool returned media refs, publish them as outbound media // If tool returned media refs, publish them as outbound media
if len(toolResult.Media) > 0 && opts.SendResponse { if len(r.result.Media) > 0 && opts.SendResponse {
parts := make([]bus.MediaPart, 0, len(toolResult.Media)) parts := make([]bus.MediaPart, 0, len(r.result.Media))
for _, ref := range toolResult.Media { for _, ref := range r.result.Media {
part := bus.MediaPart{Ref: ref} part := bus.MediaPart{Ref: ref}
// Populate metadata from MediaStore when available
if al.mediaStore != nil { if al.mediaStore != nil {
if _, meta, err := al.mediaStore.ResolveWithMeta(ref); err == nil { if _, meta, err := al.mediaStore.ResolveWithMeta(ref); err == nil {
part.Filename = meta.Filename part.Filename = meta.Filename
@ -1037,15 +1056,15 @@ func (al *AgentLoop) runLLMIteration(
} }
// Determine content for LLM based on tool result // Determine content for LLM based on tool result
contentForLLM := toolResult.ForLLM contentForLLM := r.result.ForLLM
if contentForLLM == "" && toolResult.Err != nil { if contentForLLM == "" && r.result.Err != nil {
contentForLLM = toolResult.Err.Error() contentForLLM = r.result.Err.Error()
} }
toolResultMsg := providers.Message{ toolResultMsg := providers.Message{
Role: "tool", Role: "tool",
Content: contentForLLM, Content: contentForLLM,
ToolCallID: tc.ID, ToolCallID: r.tc.ID,
} }
messages = append(messages, toolResultMsg) messages = append(messages, toolResultMsg)
@ -1081,9 +1100,9 @@ func (al *AgentLoop) updateToolContexts(agent *AgentInstance, channel, chatID st
func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, chatID string) { func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, chatID string) {
newHistory := agent.Sessions.GetHistory(sessionKey) newHistory := agent.Sessions.GetHistory(sessionKey)
tokenEstimate := al.estimateTokens(newHistory) tokenEstimate := al.estimateTokens(newHistory)
threshold := agent.ContextWindow * 75 / 100 threshold := agent.ContextWindow * agent.SummarizeTokenPercent / 100
if len(newHistory) > 20 || tokenEstimate > threshold { if len(newHistory) > agent.SummarizeMessageThreshold || tokenEstimate > threshold {
summarizeKey := agent.ID + ":" + sessionKey summarizeKey := agent.ID + ":" + sessionKey
if _, loading := al.summarizing.LoadOrStore(summarizeKey, true); !loading { if _, loading := al.summarizing.LoadOrStore(summarizeKey, true); !loading {
go func() { go func() {

40
pkg/channels/discord/discord.go
View file

@ -3,12 +3,15 @@ package discord
import ( import (
"context" "context"
"fmt" "fmt"
"net/http"
"net/url"
"os" "os"
"strings" "strings"
"sync" "sync"
"time" "time"
"github.com/bwmarrin/discordgo" "github.com/bwmarrin/discordgo"
"github.com/gorilla/websocket"
"github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/bus"
"github.com/sipeed/picoclaw/pkg/channels" "github.com/sipeed/picoclaw/pkg/channels"
@ -40,6 +43,9 @@ func NewDiscordChannel(cfg config.DiscordConfig, bus *bus.MessageBus) (*DiscordC
return nil, fmt.Errorf("failed to create discord session: %w", err) return nil, fmt.Errorf("failed to create discord session: %w", err)
} }
if err := applyDiscordProxy(session, cfg.Proxy); err != nil {
return nil, err
}
base := channels.NewBaseChannel("discord", cfg, bus, cfg.AllowFrom, base := channels.NewBaseChannel("discord", cfg, bus, cfg.AllowFrom,
channels.WithMaxMessageLength(2000), channels.WithMaxMessageLength(2000),
channels.WithGroupTrigger(cfg.GroupTrigger), channels.WithGroupTrigger(cfg.GroupTrigger),
@ -465,9 +471,43 @@ func (c *DiscordChannel) StartTyping(ctx context.Context, chatID string) (func()
func (c *DiscordChannel) downloadAttachment(url, filename string) string { func (c *DiscordChannel) downloadAttachment(url, filename string) string {
return utils.DownloadFile(url, filename, utils.DownloadOptions{ return utils.DownloadFile(url, filename, utils.DownloadOptions{
LoggerPrefix: "discord", LoggerPrefix: "discord",
ProxyURL: c.config.Proxy,
}) })
} }
func applyDiscordProxy(session *discordgo.Session, proxyAddr string) error {
var proxyFunc func(*http.Request) (*url.URL, error)
if proxyAddr != "" {
proxyURL, err := url.Parse(proxyAddr)
if err != nil {
return fmt.Errorf("invalid discord proxy URL %q: %w", proxyAddr, err)
}
proxyFunc = http.ProxyURL(proxyURL)
} else if os.Getenv("HTTP_PROXY") != "" || os.Getenv("HTTPS_PROXY") != "" {
proxyFunc = http.ProxyFromEnvironment
}
if proxyFunc == nil {
return nil
}
transport := &http.Transport{Proxy: proxyFunc}
session.Client = &http.Client{
Timeout: sendTimeout,
Transport: transport,
}
if session.Dialer != nil {
dialerCopy := *session.Dialer
dialerCopy.Proxy = proxyFunc
session.Dialer = &dialerCopy
} else {
session.Dialer = &websocket.Dialer{Proxy: proxyFunc}
}
return nil
}
// stripBotMention removes the bot mention from the message content. // stripBotMention removes the bot mention from the message content.
// Discord mentions have the format <@USER_ID> or <@!USER_ID> (with nickname). // Discord mentions have the format <@USER_ID> or <@!USER_ID> (with nickname).
func (c *DiscordChannel) stripBotMention(text string) string { func (c *DiscordChannel) stripBotMention(text string) string {

91
pkg/channels/discord/discord_test.go Normal file
View file

@ -0,0 +1,91 @@
package discord
import (
"net/http"
"net/url"
"testing"
"github.com/bwmarrin/discordgo"
)
func TestApplyDiscordProxy_CustomProxy(t *testing.T) {
session, err := discordgo.New("Bot test-token")
if err != nil {
t.Fatalf("discordgo.New() error: %v", err)
}
if err = applyDiscordProxy(session, "http://127.0.0.1:7890"); err != nil {
t.Fatalf("applyDiscordProxy() error: %v", err)
}
req, err := http.NewRequest("GET", "https://discord.com/api/v10/gateway", nil)
if err != nil {
t.Fatalf("http.NewRequest() error: %v", err)
}
restProxy := session.Client.Transport.(*http.Transport).Proxy
restProxyURL, err := restProxy(req)
if err != nil {
t.Fatalf("rest proxy func error: %v", err)
}
if got, want := restProxyURL.String(), "http://127.0.0.1:7890"; got != want {
t.Fatalf("REST proxy = %q, want %q", got, want)
}
wsProxyURL, err := session.Dialer.Proxy(req)
if err != nil {
t.Fatalf("ws proxy func error: %v", err)
}
if got, want := wsProxyURL.String(), "http://127.0.0.1:7890"; got != want {
t.Fatalf("WS proxy = %q, want %q", got, want)
}
}
func TestApplyDiscordProxy_FromEnvironment(t *testing.T) {
t.Setenv("HTTP_PROXY", "http://127.0.0.1:8888")
t.Setenv("http_proxy", "http://127.0.0.1:8888")
t.Setenv("HTTPS_PROXY", "http://127.0.0.1:8888")
t.Setenv("https_proxy", "http://127.0.0.1:8888")
t.Setenv("ALL_PROXY", "")
t.Setenv("all_proxy", "")
t.Setenv("NO_PROXY", "")
t.Setenv("no_proxy", "")
session, err := discordgo.New("Bot test-token")
if err != nil {
t.Fatalf("discordgo.New() error: %v", err)
}
if err = applyDiscordProxy(session, ""); err != nil {
t.Fatalf("applyDiscordProxy() error: %v", err)
}
req, err := http.NewRequest("GET", "https://discord.com/api/v10/gateway", nil)
if err != nil {
t.Fatalf("http.NewRequest() error: %v", err)
}
gotURL, err := session.Dialer.Proxy(req)
if err != nil {
t.Fatalf("ws proxy func error: %v", err)
}
wantURL, err := url.Parse("http://127.0.0.1:8888")
if err != nil {
t.Fatalf("url.Parse() error: %v", err)
}
if gotURL.String() != wantURL.String() {
t.Fatalf("WS proxy = %q, want %q", gotURL.String(), wantURL.String())
}
}
func TestApplyDiscordProxy_InvalidProxyURL(t *testing.T) {
session, err := discordgo.New("Bot test-token")
if err != nil {
t.Fatalf("discordgo.New() error: %v", err)
}
if err = applyDiscordProxy(session, "://bad-proxy"); err == nil {
t.Fatal("applyDiscordProxy() expected error for invalid proxy URL, got nil")
}
}

14
pkg/config/config.go
View file

@ -181,6 +181,8 @@ type AgentDefaults struct {
MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"`
Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"`
MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"`
SummarizeMessageThreshold int `json:"summarize_message_threshold" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_MESSAGE_THRESHOLD"`
SummarizeTokenPercent int `json:"summarize_token_percent" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_TOKEN_PERCENT"`
MaxMediaSize int `json:"max_media_size,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_MEDIA_SIZE"` MaxMediaSize int `json:"max_media_size,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_MEDIA_SIZE"`
} }
@ -272,6 +274,7 @@ type FeishuConfig struct {
type DiscordConfig struct { type DiscordConfig struct {
Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_DISCORD_ENABLED"` Enabled bool `json:"enabled" env:"PICOCLAW_CHANNELS_DISCORD_ENABLED"`
Token string `json:"token" env:"PICOCLAW_CHANNELS_DISCORD_TOKEN"` Token string `json:"token" env:"PICOCLAW_CHANNELS_DISCORD_TOKEN"`
Proxy string `json:"proxy" env:"PICOCLAW_CHANNELS_DISCORD_PROXY"`
AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_DISCORD_ALLOW_FROM"` AllowFrom FlexibleStringSlice `json:"allow_from" env:"PICOCLAW_CHANNELS_DISCORD_ALLOW_FROM"`
MentionOnly bool `json:"mention_only" env:"PICOCLAW_CHANNELS_DISCORD_MENTION_ONLY"` MentionOnly bool `json:"mention_only" env:"PICOCLAW_CHANNELS_DISCORD_MENTION_ONLY"`
GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"` GroupTrigger GroupTriggerConfig `json:"group_trigger,omitempty"`
@ -548,11 +551,22 @@ type PerplexityConfig struct {
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_MAX_RESULTS"` MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_PERPLEXITY_MAX_RESULTS"`
} }
type GLMSearchConfig struct {
Enabled bool `json:"enabled" env:"PICOCLAW_TOOLS_WEB_GLM_ENABLED"`
APIKey string `json:"api_key" env:"PICOCLAW_TOOLS_WEB_GLM_API_KEY"`
BaseURL string `json:"base_url" env:"PICOCLAW_TOOLS_WEB_GLM_BASE_URL"`
// SearchEngine specifies the search backend: "search_std" (default),
// "search_pro", "search_pro_sogou", or "search_pro_quark".
SearchEngine string `json:"search_engine" env:"PICOCLAW_TOOLS_WEB_GLM_SEARCH_ENGINE"`
MaxResults int `json:"max_results" env:"PICOCLAW_TOOLS_WEB_GLM_MAX_RESULTS"`
}
type WebToolsConfig struct { type WebToolsConfig struct {
Brave BraveConfig `json:"brave"` Brave BraveConfig `json:"brave"`
Tavily TavilyConfig `json:"tavily"` Tavily TavilyConfig `json:"tavily"`
DuckDuckGo DuckDuckGoConfig `json:"duckduckgo"` DuckDuckGo DuckDuckGoConfig `json:"duckduckgo"`
Perplexity PerplexityConfig `json:"perplexity"` Perplexity PerplexityConfig `json:"perplexity"`
GLMSearch GLMSearchConfig `json:"glm_search"`
// Proxy is an optional proxy URL for web tools (http/https/socks5/socks5h). // Proxy is an optional proxy URL for web tools (http/https/socks5/socks5h).
// For authenticated proxies, prefer HTTP_PROXY/HTTPS_PROXY env vars instead of embedding credentials in config. // For authenticated proxies, prefer HTTP_PROXY/HTTPS_PROXY env vars instead of embedding credentials in config.
Proxy string `json:"proxy,omitempty" env:"PICOCLAW_TOOLS_WEB_PROXY"` Proxy string `json:"proxy,omitempty" env:"PICOCLAW_TOOLS_WEB_PROXY"`

12
pkg/config/config_test.go
View file

@ -435,6 +435,18 @@ func TestLoadConfig_WebToolsProxy(t *testing.T) {
} }
// TestDefaultConfig_DMScope verifies the default dm_scope value // TestDefaultConfig_DMScope verifies the default dm_scope value
// TestDefaultConfig_SummarizationThresholds verifies summarization defaults
func TestDefaultConfig_SummarizationThresholds(t *testing.T) {
cfg := DefaultConfig()
if cfg.Agents.Defaults.SummarizeMessageThreshold != 20 {
t.Errorf("SummarizeMessageThreshold = %d, want 20", cfg.Agents.Defaults.SummarizeMessageThreshold)
}
if cfg.Agents.Defaults.SummarizeTokenPercent != 75 {
t.Errorf("SummarizeTokenPercent = %d, want 75", cfg.Agents.Defaults.SummarizeTokenPercent)
}
}
func TestDefaultConfig_DMScope(t *testing.T) { func TestDefaultConfig_DMScope(t *testing.T) {
cfg := DefaultConfig() cfg := DefaultConfig()

23
pkg/config/defaults.go
View file

@ -26,13 +26,15 @@ func DefaultConfig() *Config {
return &Config{ return &Config{
Agents: AgentsConfig{ Agents: AgentsConfig{
Defaults: AgentDefaults{ Defaults: AgentDefaults{
Workspace: workspacePath, Workspace: workspacePath,
RestrictToWorkspace: true, RestrictToWorkspace: true,
Provider: "", Provider: "",
Model: "", Model: "",
MaxTokens: 32768, MaxTokens: 32768,
Temperature: nil, // nil means use provider default Temperature: nil, // nil means use provider default
MaxToolIterations: 50, MaxToolIterations: 50,
SummarizeMessageThreshold: 20,
SummarizeTokenPercent: 75,
}, },
}, },
Bindings: []AgentBinding{}, Bindings: []AgentBinding{},
@ -349,6 +351,13 @@ func DefaultConfig() *Config {
APIKeys: []string{"YOUR_PERPLEXITY_API_KEY"}, APIKeys: []string{"YOUR_PERPLEXITY_API_KEY"},
MaxResults: 5, MaxResults: 5,
}, },
GLMSearch: GLMSearchConfig{
Enabled: false,
APIKey: "",
BaseURL: "https://open.bigmodel.cn/api/paas/v4/web_search",
SearchEngine: "search_std",
MaxResults: 5,
},
}, },
Cron: CronToolsConfig{ Cron: CronToolsConfig{
ExecTimeoutMinutes: 5, ExecTimeoutMinutes: 5,

460
pkg/memory/jsonl.go Normal file
View file

@ -0,0 +1,460 @@
package memory
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"hash/fnv"
"log"
"os"
"path/filepath"
"strings"
"sync"
"time"
"github.com/sipeed/picoclaw/pkg/fileutil"
"github.com/sipeed/picoclaw/pkg/providers"
)
const (
// numLockShards is the fixed number of mutexes used to serialize
// per-session access. Using a sharded array instead of a map keeps
// memory bounded regardless of how many sessions are created over
// the lifetime of the process — important for a long-running daemon.
numLockShards = 64
// maxLineSize is the maximum size of a single JSON line in a .jsonl
// file. Tool results (read_file, web search, etc.) can be large, so
// we set a generous limit. The scanner starts at 64 KB and grows
// only as needed up to this cap.
maxLineSize = 10 * 1024 * 1024 // 10 MB
)
// sessionMeta holds per-session metadata stored in a .meta.json file.
type sessionMeta struct {
Key string `json:"key"`
Summary string `json:"summary"`
Skip int `json:"skip"`
Count int `json:"count"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
// JSONLStore implements Store using append-only JSONL files.
//
// Each session is stored as two files:
//
// {sanitized_key}.jsonl — one JSON-encoded message per line, append-only
// {sanitized_key}.meta.json — session metadata (summary, logical truncation offset)
//
// Messages are never physically deleted from the JSONL file. Instead,
// TruncateHistory records a "skip" offset in the metadata file and
// GetHistory ignores lines before that offset. This keeps all writes
// append-only, which is both fast and crash-safe.
type JSONLStore struct {
dir string
locks [numLockShards]sync.Mutex
}
// NewJSONLStore creates a new JSONL-backed store rooted at dir.
func NewJSONLStore(dir string) (*JSONLStore, error) {
err := os.MkdirAll(dir, 0o755)
if err != nil {
return nil, fmt.Errorf("memory: create directory: %w", err)
}
return &JSONLStore{dir: dir}, nil
}
// sessionLock returns a mutex for the given session key.
// Keys are mapped to a fixed pool of shards via FNV hash, so
// memory usage is O(1) regardless of total session count.
func (s *JSONLStore) sessionLock(key string) *sync.Mutex {
h := fnv.New32a()
h.Write([]byte(key))
return &s.locks[h.Sum32()%numLockShards]
}
func (s *JSONLStore) jsonlPath(key string) string {
return filepath.Join(s.dir, sanitizeKey(key)+".jsonl")
}
func (s *JSONLStore) metaPath(key string) string {
return filepath.Join(s.dir, sanitizeKey(key)+".meta.json")
}
// sanitizeKey converts a session key to a safe filename component.
// Mirrors pkg/session.sanitizeFilename so that migration paths match.
//
// Note: this is a lossy mapping — "telegram:123" and "telegram_123"
// both produce the same filename. This is an intentional tradeoff:
// keys with colons (e.g. from channels) are by far the common case,
// and a bidirectional encoding (like URL-encoding) would complicate
// file listings and debugging.
func sanitizeKey(key string) string {
return strings.ReplaceAll(key, ":", "_")
}
// readMeta loads the metadata file for a session.
// Returns a zero-value sessionMeta if the file does not exist.
func (s *JSONLStore) readMeta(key string) (sessionMeta, error) {
data, err := os.ReadFile(s.metaPath(key))
if os.IsNotExist(err) {
return sessionMeta{Key: key}, nil
}
if err != nil {
return sessionMeta{}, fmt.Errorf("memory: read meta: %w", err)
}
var meta sessionMeta
err = json.Unmarshal(data, &meta)
if err != nil {
return sessionMeta{}, fmt.Errorf("memory: decode meta: %w", err)
}
return meta, nil
}
// writeMeta atomically writes the metadata file using the project's
// standard WriteFileAtomic (temp + fsync + rename).
func (s *JSONLStore) writeMeta(key string, meta sessionMeta) error {
data, err := json.MarshalIndent(meta, "", " ")
if err != nil {
return fmt.Errorf("memory: encode meta: %w", err)
}
return fileutil.WriteFileAtomic(s.metaPath(key), data, 0o644)
}
// readMessages reads valid JSON lines from a .jsonl file, skipping
// the first `skip` lines without unmarshaling them. This avoids the
// cost of json.Unmarshal on logically truncated messages.
// Malformed trailing lines (e.g. from a crash) are silently skipped.
func readMessages(path string, skip int) ([]providers.Message, error) {
f, err := os.Open(path)
if os.IsNotExist(err) {
return []providers.Message{}, nil
}
if err != nil {
return nil, fmt.Errorf("memory: open jsonl: %w", err)
}
defer f.Close()
var msgs []providers.Message
scanner := bufio.NewScanner(f)
// Allow large lines for tool results (read_file, web search, etc.).
scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize)
lineNum := 0
for scanner.Scan() {
line := scanner.Bytes()
if len(line) == 0 {
continue
}
lineNum++
if lineNum <= skip {
continue
}
var msg providers.Message
if err := json.Unmarshal(line, &msg); err != nil {
// Corrupt line — likely a partial write from a crash.
// Log so operators know data was skipped, but don't
// fail the entire read; this is the standard JSONL
// recovery pattern.
log.Printf("memory: skipping corrupt line %d in %s: %v",
lineNum, filepath.Base(path), err)
continue
}
msgs = append(msgs, msg)
}
if scanner.Err() != nil {
return nil, fmt.Errorf("memory: scan jsonl: %w", scanner.Err())
}
if msgs == nil {
msgs = []providers.Message{}
}
return msgs, nil
}
// countLines counts the total number of non-empty lines in a .jsonl file.
// Used by TruncateHistory to reconcile a stale meta.Count without
// the overhead of unmarshaling every message.
func countLines(path string) (int, error) {
f, err := os.Open(path)
if os.IsNotExist(err) {
return 0, nil
}
if err != nil {
return 0, fmt.Errorf("memory: open jsonl: %w", err)
}
defer f.Close()
n := 0
scanner := bufio.NewScanner(f)
scanner.Buffer(make([]byte, 0, 64*1024), maxLineSize)
for scanner.Scan() {
if len(scanner.Bytes()) > 0 {
n++
}
}
return n, scanner.Err()
}
func (s *JSONLStore) AddMessage(
_ context.Context, sessionKey, role, content string,
) error {
return s.addMsg(sessionKey, providers.Message{
Role: role,
Content: content,
})
}
func (s *JSONLStore) AddFullMessage(
_ context.Context, sessionKey string, msg providers.Message,
) error {
return s.addMsg(sessionKey, msg)
}
// addMsg is the shared implementation for AddMessage and AddFullMessage.
func (s *JSONLStore) addMsg(sessionKey string, msg providers.Message) error {
l := s.sessionLock(sessionKey)
l.Lock()
defer l.Unlock()
// Append the message as a single JSON line.
line, err := json.Marshal(msg)
if err != nil {
return fmt.Errorf("memory: marshal message: %w", err)
}
line = append(line, '\n')
f, err := os.OpenFile(
s.jsonlPath(sessionKey),
os.O_CREATE|os.O_WRONLY|os.O_APPEND,
0o644,
)
if err != nil {
return fmt.Errorf("memory: open jsonl for append: %w", err)
}
_, writeErr := f.Write(line)
if writeErr != nil {
f.Close()
return fmt.Errorf("memory: append message: %w", writeErr)
}
// Flush to physical storage before closing. This matches the
// durability guarantee of writeMeta and rewriteJSONL (which use
// WriteFileAtomic with fsync). Without Sync, a power loss could
// leave the append in the kernel page cache only — lost on reboot.
if syncErr := f.Sync(); syncErr != nil {
f.Close()
return fmt.Errorf("memory: sync jsonl: %w", syncErr)
}
if closeErr := f.Close(); closeErr != nil {
return fmt.Errorf("memory: close jsonl: %w", closeErr)
}
// Update metadata.
meta, err := s.readMeta(sessionKey)
if err != nil {
return err
}
now := time.Now()
if meta.Count == 0 && meta.CreatedAt.IsZero() {
meta.CreatedAt = now
}
meta.Count++
meta.UpdatedAt = now
return s.writeMeta(sessionKey, meta)
}
func (s *JSONLStore) GetHistory(
_ context.Context, sessionKey string,
) ([]providers.Message, error) {
l := s.sessionLock(sessionKey)
l.Lock()
defer l.Unlock()
meta, err := s.readMeta(sessionKey)
if err != nil {
return nil, err
}
// Pass meta.Skip so readMessages skips those lines without
// unmarshaling them — avoids wasted CPU on truncated messages.
msgs, err := readMessages(s.jsonlPath(sessionKey), meta.Skip)
if err != nil {
return nil, err
}
return msgs, nil
}
func (s *JSONLStore) GetSummary(
_ context.Context, sessionKey string,
) (string, error) {
l := s.sessionLock(sessionKey)
l.Lock()
defer l.Unlock()
meta, err := s.readMeta(sessionKey)
if err != nil {
return "", err
}
return meta.Summary, nil
}
func (s *JSONLStore) SetSummary(
_ context.Context, sessionKey, summary string,
) error {
l := s.sessionLock(sessionKey)
l.Lock()
defer l.Unlock()
meta, err := s.readMeta(sessionKey)
if err != nil {
return err
}
now := time.Now()
if meta.CreatedAt.IsZero() {
meta.CreatedAt = now
}
meta.Summary = summary
meta.UpdatedAt = now
return s.writeMeta(sessionKey, meta)
}
func (s *JSONLStore) TruncateHistory(
_ context.Context, sessionKey string, keepLast int,
) error {
l := s.sessionLock(sessionKey)
l.Lock()
defer l.Unlock()
meta, err := s.readMeta(sessionKey)
if err != nil {
return err
}
// Always reconcile meta.Count with the actual line count on disk.
// A crash between the JSONL append and the meta update in addMsg
// leaves meta.Count stale (e.g. file has 101 lines but meta says
// 100). Counting lines is cheap — no unmarshal, just a scan — and
// TruncateHistory is not a hot path, so always re-count.
n, countErr := countLines(s.jsonlPath(sessionKey))
if countErr != nil {
return countErr
}
meta.Count = n
if keepLast <= 0 {
meta.Skip = meta.Count
} else {
effective := meta.Count - meta.Skip
if keepLast < effective {
meta.Skip = meta.Count - keepLast
}
}
meta.UpdatedAt = time.Now()
return s.writeMeta(sessionKey, meta)
}
func (s *JSONLStore) SetHistory(
_ context.Context,
sessionKey string,
history []providers.Message,
) error {
l := s.sessionLock(sessionKey)
l.Lock()
defer l.Unlock()
meta, err := s.readMeta(sessionKey)
if err != nil {
return err
}
now := time.Now()
if meta.CreatedAt.IsZero() {
meta.CreatedAt = now
}
meta.Skip = 0
meta.Count = len(history)
meta.UpdatedAt = now
// Write meta BEFORE rewriting the JSONL file. If we crash between
// the two writes, meta has Skip=0 and the old file is still intact,
// so GetHistory reads from line 1 — returning "too many" messages
// rather than losing data. The next SetHistory call corrects this.
err = s.writeMeta(sessionKey, meta)
if err != nil {
return err
}
return s.rewriteJSONL(sessionKey, history)
}
// Compact physically rewrites the JSONL file, dropping all logically
// skipped lines. This reclaims disk space that accumulates after
// repeated TruncateHistory calls.
//
// It is safe to call at any time; if there is nothing to compact
// (skip == 0) the method returns immediately.
func (s *JSONLStore) Compact(
_ context.Context, sessionKey string,
) error {
l := s.sessionLock(sessionKey)
l.Lock()
defer l.Unlock()
meta, err := s.readMeta(sessionKey)
if err != nil {
return err
}
if meta.Skip == 0 {
return nil
}
// Read only the active messages, skipping truncated lines
// without unmarshaling them.
active, err := readMessages(s.jsonlPath(sessionKey), meta.Skip)
if err != nil {
return err
}
// Write meta BEFORE rewriting the JSONL file. If the process
// crashes between the two writes, meta has Skip=0 and the old
// (uncompacted) file is still intact, so GetHistory reads from
// line 1 — returning previously-truncated messages rather than
// losing data. The next Compact or TruncateHistory corrects this.
meta.Skip = 0
meta.Count = len(active)
meta.UpdatedAt = time.Now()
err = s.writeMeta(sessionKey, meta)
if err != nil {
return err
}
return s.rewriteJSONL(sessionKey, active)
}
// rewriteJSONL atomically replaces the JSONL file with the given messages
// using the project's standard WriteFileAtomic (temp + fsync + rename).
func (s *JSONLStore) rewriteJSONL(
sessionKey string, msgs []providers.Message,
) error {
var buf bytes.Buffer
for i, msg := range msgs {
line, err := json.Marshal(msg)
if err != nil {
return fmt.Errorf("memory: marshal message %d: %w", i, err)
}
buf.Write(line)
buf.WriteByte('\n')
}
return fileutil.WriteFileAtomic(s.jsonlPath(sessionKey), buf.Bytes(), 0o644)
}
func (s *JSONLStore) Close() error {
return nil
}

835
pkg/memory/jsonl_test.go Normal file
View file

@ -0,0 +1,835 @@
package memory
import (
"context"
"os"
"path/filepath"
"sync"
"testing"
"github.com/sipeed/picoclaw/pkg/providers"
)
func newTestStore(t *testing.T) *JSONLStore {
t.Helper()
store, err := NewJSONLStore(t.TempDir())
if err != nil {
t.Fatalf("NewJSONLStore: %v", err)
}
return store
}
func TestNewJSONLStore_CreatesDirectory(t *testing.T) {
dir := filepath.Join(t.TempDir(), "nested", "sessions")
store, err := NewJSONLStore(dir)
if err != nil {
t.Fatalf("NewJSONLStore: %v", err)
}
defer store.Close()
info, err := os.Stat(dir)
if err != nil {
t.Fatalf("Stat: %v", err)
}
if !info.IsDir() {
t.Errorf("expected directory, got file")
}
}
func TestAddMessage_BasicRoundtrip(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
err := store.AddMessage(ctx, "s1", "user", "hello")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
err = store.AddMessage(ctx, "s1", "assistant", "hi there")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
history, err := store.GetHistory(ctx, "s1")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 2 {
t.Fatalf("expected 2 messages, got %d", len(history))
}
if history[0].Role != "user" || history[0].Content != "hello" {
t.Errorf("msg[0] = %+v", history[0])
}
if history[1].Role != "assistant" || history[1].Content != "hi there" {
t.Errorf("msg[1] = %+v", history[1])
}
}
func TestAddMessage_AutoCreatesSession(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
// Adding a message to a non-existent session should work.
err := store.AddMessage(ctx, "new-session", "user", "first message")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
history, err := store.GetHistory(ctx, "new-session")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Fatalf("expected 1 message, got %d", len(history))
}
}
func TestAddFullMessage_WithToolCalls(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
msg := providers.Message{
Role: "assistant",
Content: "Let me search that.",
ToolCalls: []providers.ToolCall{
{
ID: "call_abc",
Type: "function",
Function: &providers.FunctionCall{
Name: "web_search",
Arguments: `{"q":"golang jsonl"}`,
},
},
},
}
err := store.AddFullMessage(ctx, "tc", msg)
if err != nil {
t.Fatalf("AddFullMessage: %v", err)
}
history, err := store.GetHistory(ctx, "tc")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Fatalf("expected 1, got %d", len(history))
}
if len(history[0].ToolCalls) != 1 {
t.Fatalf("expected 1 tool call, got %d", len(history[0].ToolCalls))
}
tc := history[0].ToolCalls[0]
if tc.ID != "call_abc" {
t.Errorf("tool call ID = %q", tc.ID)
}
if tc.Function == nil || tc.Function.Name != "web_search" {
t.Errorf("tool call function = %+v", tc.Function)
}
}
func TestAddFullMessage_ToolCallID(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
msg := providers.Message{
Role: "tool",
Content: "search results here",
ToolCallID: "call_abc",
}
err := store.AddFullMessage(ctx, "tr", msg)
if err != nil {
t.Fatalf("AddFullMessage: %v", err)
}
history, err := store.GetHistory(ctx, "tr")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Fatalf("expected 1, got %d", len(history))
}
if history[0].ToolCallID != "call_abc" {
t.Errorf("ToolCallID = %q", history[0].ToolCallID)
}
}
func TestGetHistory_EmptySession(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
history, err := store.GetHistory(ctx, "nonexistent")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if history == nil {
t.Fatal("expected non-nil empty slice")
}
if len(history) != 0 {
t.Errorf("expected 0 messages, got %d", len(history))
}
}
func TestGetHistory_Ordering(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
for i := 0; i < 5; i++ {
err := store.AddMessage(
ctx, "order",
"user",
string(rune('a'+i)),
)
if err != nil {
t.Fatalf("AddMessage(%d): %v", i, err)
}
}
history, err := store.GetHistory(ctx, "order")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 5 {
t.Fatalf("expected 5, got %d", len(history))
}
for i := 0; i < 5; i++ {
expected := string(rune('a' + i))
if history[i].Content != expected {
t.Errorf("msg[%d].Content = %q, want %q", i, history[i].Content, expected)
}
}
}
func TestSetSummary_GetSummary(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
// No summary yet.
summary, err := store.GetSummary(ctx, "s1")
if err != nil {
t.Fatalf("GetSummary: %v", err)
}
if summary != "" {
t.Errorf("expected empty, got %q", summary)
}
// Set a summary.
err = store.SetSummary(ctx, "s1", "talked about Go")
if err != nil {
t.Fatalf("SetSummary: %v", err)
}
summary, err = store.GetSummary(ctx, "s1")
if err != nil {
t.Fatalf("GetSummary: %v", err)
}
if summary != "talked about Go" {
t.Errorf("summary = %q", summary)
}
// Update summary.
err = store.SetSummary(ctx, "s1", "updated summary")
if err != nil {
t.Fatalf("SetSummary: %v", err)
}
summary, err = store.GetSummary(ctx, "s1")
if err != nil {
t.Fatalf("GetSummary: %v", err)
}
if summary != "updated summary" {
t.Errorf("summary = %q", summary)
}
}
func TestTruncateHistory_KeepLast(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
for i := 0; i < 10; i++ {
err := store.AddMessage(
ctx, "trunc",
"user",
string(rune('a'+i)),
)
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
err := store.TruncateHistory(ctx, "trunc", 4)
if err != nil {
t.Fatalf("TruncateHistory: %v", err)
}
history, err := store.GetHistory(ctx, "trunc")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 4 {
t.Fatalf("expected 4, got %d", len(history))
}
// Should be the last 4: g, h, i, j
if history[0].Content != "g" {
t.Errorf("first kept = %q, want 'g'", history[0].Content)
}
if history[3].Content != "j" {
t.Errorf("last kept = %q, want 'j'", history[3].Content)
}
}
func TestTruncateHistory_KeepZero(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
for i := 0; i < 5; i++ {
err := store.AddMessage(ctx, "empty", "user", "msg")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
err := store.TruncateHistory(ctx, "empty", 0)
if err != nil {
t.Fatalf("TruncateHistory: %v", err)
}
history, err := store.GetHistory(ctx, "empty")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 0 {
t.Errorf("expected 0, got %d", len(history))
}
}
func TestTruncateHistory_KeepMoreThanExists(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
for i := 0; i < 3; i++ {
err := store.AddMessage(ctx, "few", "user", "msg")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
// Keep 100, but only 3 exist — should keep all.
err := store.TruncateHistory(ctx, "few", 100)
if err != nil {
t.Fatalf("TruncateHistory: %v", err)
}
history, err := store.GetHistory(ctx, "few")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 3 {
t.Errorf("expected 3, got %d", len(history))
}
}
func TestSetHistory_ReplacesAll(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
// Add some initial messages.
for i := 0; i < 5; i++ {
err := store.AddMessage(ctx, "replace", "user", "old")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
// Replace with new history.
newHistory := []providers.Message{
{Role: "user", Content: "new1"},
{Role: "assistant", Content: "new2"},
}
err := store.SetHistory(ctx, "replace", newHistory)
if err != nil {
t.Fatalf("SetHistory: %v", err)
}
history, err := store.GetHistory(ctx, "replace")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 2 {
t.Fatalf("expected 2, got %d", len(history))
}
if history[0].Content != "new1" || history[1].Content != "new2" {
t.Errorf("history = %+v", history)
}
}
func TestSetHistory_ResetsSkip(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
// Add messages and truncate.
for i := 0; i < 10; i++ {
err := store.AddMessage(ctx, "skip-reset", "user", "old")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
err := store.TruncateHistory(ctx, "skip-reset", 3)
if err != nil {
t.Fatalf("TruncateHistory: %v", err)
}
// SetHistory should reset skip to 0.
newHistory := []providers.Message{
{Role: "user", Content: "fresh"},
}
err = store.SetHistory(ctx, "skip-reset", newHistory)
if err != nil {
t.Fatalf("SetHistory: %v", err)
}
history, err := store.GetHistory(ctx, "skip-reset")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Fatalf("expected 1, got %d", len(history))
}
if history[0].Content != "fresh" {
t.Errorf("content = %q", history[0].Content)
}
}
func TestColonInKey(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
err := store.AddMessage(ctx, "telegram:123", "user", "hi")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
history, err := store.GetHistory(ctx, "telegram:123")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Fatalf("expected 1, got %d", len(history))
}
// Verify the file is named with underscore.
jsonlFile := filepath.Join(store.dir, "telegram_123.jsonl")
if _, statErr := os.Stat(jsonlFile); statErr != nil {
t.Errorf("expected file %s to exist: %v", jsonlFile, statErr)
}
}
func TestCompact_RemovesSkippedMessages(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
// Write 10 messages, then truncate to keep last 3.
for i := 0; i < 10; i++ {
err := store.AddMessage(ctx, "compact", "user", string(rune('a'+i)))
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
err := store.TruncateHistory(ctx, "compact", 3)
if err != nil {
t.Fatalf("TruncateHistory: %v", err)
}
// Before compact: file still has 10 lines.
allOnDisk, err := readMessages(store.jsonlPath("compact"), 0)
if err != nil {
t.Fatalf("readMessages: %v", err)
}
if len(allOnDisk) != 10 {
t.Fatalf("before compact: expected 10 on disk, got %d", len(allOnDisk))
}
// Compact.
err = store.Compact(ctx, "compact")
if err != nil {
t.Fatalf("Compact: %v", err)
}
// After compact: file should have only 3 lines.
allOnDisk, err = readMessages(store.jsonlPath("compact"), 0)
if err != nil {
t.Fatalf("readMessages: %v", err)
}
if len(allOnDisk) != 3 {
t.Fatalf("after compact: expected 3 on disk, got %d", len(allOnDisk))
}
// GetHistory should still return the same 3 messages.
history, err := store.GetHistory(ctx, "compact")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 3 {
t.Fatalf("expected 3, got %d", len(history))
}
if history[0].Content != "h" || history[2].Content != "j" {
t.Errorf("wrong content: %+v", history)
}
}
func TestCompact_NoOpWhenNoSkip(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
for i := 0; i < 5; i++ {
err := store.AddMessage(ctx, "noop", "user", "msg")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
// Compact without prior truncation — should be a no-op.
err := store.Compact(ctx, "noop")
if err != nil {
t.Fatalf("Compact: %v", err)
}
history, err := store.GetHistory(ctx, "noop")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 5 {
t.Errorf("expected 5, got %d", len(history))
}
}
func TestCompact_ThenAppend(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
for i := 0; i < 8; i++ {
err := store.AddMessage(ctx, "cap", "user", string(rune('a'+i)))
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
err := store.TruncateHistory(ctx, "cap", 2)
if err != nil {
t.Fatalf("TruncateHistory: %v", err)
}
err = store.Compact(ctx, "cap")
if err != nil {
t.Fatalf("Compact: %v", err)
}
// Append after compaction should work correctly.
err = store.AddMessage(ctx, "cap", "user", "new")
if err != nil {
t.Fatalf("AddMessage after compact: %v", err)
}
history, err := store.GetHistory(ctx, "cap")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 3 {
t.Fatalf("expected 3, got %d", len(history))
}
// g, h (kept from truncation), new (appended after compaction).
if history[0].Content != "g" {
t.Errorf("first = %q, want 'g'", history[0].Content)
}
if history[2].Content != "new" {
t.Errorf("last = %q, want 'new'", history[2].Content)
}
}
func TestTruncateHistory_StaleMetaCount(t *testing.T) {
// Simulates a crash between JSONL append and meta update in addMsg:
// file has N+1 lines but meta.Count is still N. TruncateHistory must
// reconcile with the real line count so that keepLast is accurate.
store := newTestStore(t)
ctx := context.Background()
// Write 10 messages normally (meta.Count = 10).
for i := 0; i < 10; i++ {
err := store.AddMessage(ctx, "stale", "user", string(rune('a'+i)))
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
// Simulate crash: append a line to JSONL but do NOT update meta.
// This leaves meta.Count = 10 while the file has 11 lines.
jsonlPath := store.jsonlPath("stale")
f, err := os.OpenFile(jsonlPath, os.O_WRONLY|os.O_APPEND, 0o644)
if err != nil {
t.Fatalf("open for append: %v", err)
}
_, err = f.WriteString(`{"role":"user","content":"orphan"}` + "\n")
if err != nil {
t.Fatalf("write orphan: %v", err)
}
f.Close()
// TruncateHistory(keepLast=4) should keep the last 4 of 11 lines,
// not the last 4 of 10.
err = store.TruncateHistory(ctx, "stale", 4)
if err != nil {
t.Fatalf("TruncateHistory: %v", err)
}
history, err := store.GetHistory(ctx, "stale")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 4 {
t.Fatalf("expected 4, got %d", len(history))
}
// Last 4 of [a,b,c,d,e,f,g,h,i,j,orphan] = [h,i,j,orphan]
if history[0].Content != "h" {
t.Errorf("first kept = %q, want 'h'", history[0].Content)
}
if history[3].Content != "orphan" {
t.Errorf("last kept = %q, want 'orphan'", history[3].Content)
}
}
func TestCrashRecovery_PartialLine(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
// Write a valid message first.
err := store.AddMessage(ctx, "crash", "user", "valid")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
// Simulate a crash by appending a partial JSON line directly.
jsonlPath := store.jsonlPath("crash")
f, err := os.OpenFile(jsonlPath, os.O_WRONLY|os.O_APPEND, 0o644)
if err != nil {
t.Fatalf("open for append: %v", err)
}
_, err = f.WriteString(`{"role":"user","content":"incomple`)
if err != nil {
t.Fatalf("write partial: %v", err)
}
f.Close()
// GetHistory should return only the valid message.
history, err := store.GetHistory(ctx, "crash")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Fatalf("expected 1 valid message, got %d", len(history))
}
if history[0].Content != "valid" {
t.Errorf("content = %q", history[0].Content)
}
}
func TestPersistence_AcrossInstances(t *testing.T) {
dir := t.TempDir()
ctx := context.Background()
// Write with first instance.
store1, err := NewJSONLStore(dir)
if err != nil {
t.Fatalf("NewJSONLStore: %v", err)
}
err = store1.AddMessage(ctx, "persist", "user", "remember me")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
err = store1.SetSummary(ctx, "persist", "a test session")
if err != nil {
t.Fatalf("SetSummary: %v", err)
}
store1.Close()
// Read with second instance.
store2, err := NewJSONLStore(dir)
if err != nil {
t.Fatalf("NewJSONLStore: %v", err)
}
defer store2.Close()
history, err := store2.GetHistory(ctx, "persist")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 || history[0].Content != "remember me" {
t.Errorf("history = %+v", history)
}
summary, err := store2.GetSummary(ctx, "persist")
if err != nil {
t.Fatalf("GetSummary: %v", err)
}
if summary != "a test session" {
t.Errorf("summary = %q", summary)
}
}
func TestConcurrent_AddAndRead(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
var wg sync.WaitGroup
const goroutines = 10
const msgsPerGoroutine = 20
// Concurrent writes.
for g := 0; g < goroutines; g++ {
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < msgsPerGoroutine; i++ {
_ = store.AddMessage(ctx, "concurrent", "user", "msg")
}
}()
}
wg.Wait()
history, err := store.GetHistory(ctx, "concurrent")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
expected := goroutines * msgsPerGoroutine
if len(history) != expected {
t.Errorf("expected %d messages, got %d", expected, len(history))
}
}
func TestConcurrent_SummarizeRace(t *testing.T) {
// Simulates the #704 race: one goroutine adds messages while
// another truncates + sets summary — like summarizeSession().
store := newTestStore(t)
ctx := context.Background()
// Seed with some messages.
for i := 0; i < 20; i++ {
err := store.AddMessage(ctx, "race", "user", "seed")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
}
var wg sync.WaitGroup
// Writer goroutine (main agent loop).
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 50; i++ {
_ = store.AddMessage(ctx, "race", "user", "new")
}
}()
// Summarizer goroutine (background task).
wg.Add(1)
go func() {
defer wg.Done()
for i := 0; i < 10; i++ {
_ = store.SetSummary(ctx, "race", "summary")
_ = store.TruncateHistory(ctx, "race", 5)
}
}()
wg.Wait()
// Verify the store is still in a consistent state.
_, err := store.GetHistory(ctx, "race")
if err != nil {
t.Fatalf("GetHistory after race: %v", err)
}
_, err = store.GetSummary(ctx, "race")
if err != nil {
t.Fatalf("GetSummary after race: %v", err)
}
}
func TestMultipleSessions_Isolation(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
err := store.AddMessage(ctx, "s1", "user", "msg for s1")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
err = store.AddMessage(ctx, "s2", "user", "msg for s2")
if err != nil {
t.Fatalf("AddMessage: %v", err)
}
h1, err := store.GetHistory(ctx, "s1")
if err != nil {
t.Fatalf("GetHistory s1: %v", err)
}
h2, err := store.GetHistory(ctx, "s2")
if err != nil {
t.Fatalf("GetHistory s2: %v", err)
}
if len(h1) != 1 || h1[0].Content != "msg for s1" {
t.Errorf("s1 history = %+v", h1)
}
if len(h2) != 1 || h2[0].Content != "msg for s2" {
t.Errorf("s2 history = %+v", h2)
}
}
func BenchmarkAddMessage(b *testing.B) {
dir := b.TempDir()
store, err := NewJSONLStore(dir)
if err != nil {
b.Fatalf("NewJSONLStore: %v", err)
}
defer store.Close()
ctx := context.Background()
b.ResetTimer()
for i := 0; i < b.N; i++ {
_ = store.AddMessage(ctx, "bench", "user", "benchmark message content")
}
}
func BenchmarkGetHistory_100(b *testing.B) {
dir := b.TempDir()
store, err := NewJSONLStore(dir)
if err != nil {
b.Fatalf("NewJSONLStore: %v", err)
}
defer store.Close()
ctx := context.Background()
for i := 0; i < 100; i++ {
_ = store.AddMessage(ctx, "bench", "user", "message content")
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _ = store.GetHistory(ctx, "bench")
}
}
func BenchmarkGetHistory_1000(b *testing.B) {
dir := b.TempDir()
store, err := NewJSONLStore(dir)
if err != nil {
b.Fatalf("NewJSONLStore: %v", err)
}
defer store.Close()
ctx := context.Background()
for i := 0; i < 1000; i++ {
_ = store.AddMessage(ctx, "bench", "user", "message content")
}
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, _ = store.GetHistory(ctx, "bench")
}
}

108
pkg/memory/migration.go Normal file
View file

@ -0,0 +1,108 @@
package memory
import (
"context"
"encoding/json"
"fmt"
"log"
"os"
"path/filepath"
"strings"
"time"
"github.com/sipeed/picoclaw/pkg/providers"
)
// jsonSession mirrors pkg/session.Session for migration purposes.
type jsonSession 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"`
}
// MigrateFromJSON reads legacy sessions/*.json files from sessionsDir,
// writes them into the Store, and renames each migrated file to
// .json.migrated as a backup. Returns the number of sessions migrated.
//
// Files that fail to parse are logged and skipped. Already-migrated
// files (.json.migrated) are ignored, making the function idempotent.
func MigrateFromJSON(
ctx context.Context, sessionsDir string, store Store,
) (int, error) {
entries, err := os.ReadDir(sessionsDir)
if os.IsNotExist(err) {
return 0, nil
}
if err != nil {
return 0, fmt.Errorf("memory: read sessions dir: %w", err)
}
migrated := 0
for _, entry := range entries {
if entry.IsDir() {
continue
}
name := entry.Name()
if !strings.HasSuffix(name, ".json") {
continue
}
// Skip already-migrated files.
if strings.HasSuffix(name, ".migrated") {
continue
}
srcPath := filepath.Join(sessionsDir, name)
data, readErr := os.ReadFile(srcPath)
if readErr != nil {
log.Printf("memory: migrate: skip %s: %v", name, readErr)
continue
}
var sess jsonSession
if parseErr := json.Unmarshal(data, &sess); parseErr != nil {
log.Printf("memory: migrate: skip %s: %v", name, parseErr)
continue
}
// Use the key from the JSON content, not the filename.
// Filenames are sanitized (":" → "_") but keys are not.
key := sess.Key
if key == "" {
key = strings.TrimSuffix(name, ".json")
}
// Use SetHistory (atomic replace) instead of per-message
// AddFullMessage. This makes migration idempotent: if the
// process crashes after writing messages but before the
// rename below, a retry replaces the partial data cleanly
// instead of duplicating messages.
if setErr := store.SetHistory(ctx, key, sess.Messages); setErr != nil {
return migrated, fmt.Errorf(
"memory: migrate %s: set history: %w",
name, setErr,
)
}
if sess.Summary != "" {
if sumErr := store.SetSummary(ctx, key, sess.Summary); sumErr != nil {
return migrated, fmt.Errorf(
"memory: migrate %s: set summary: %w",
name, sumErr,
)
}
}
// Rename to .migrated as backup (not delete).
renameErr := os.Rename(srcPath, srcPath+".migrated")
if renameErr != nil {
log.Printf("memory: migrate: rename %s: %v", name, renameErr)
}
migrated++
}
return migrated, nil
}

384
pkg/memory/migration_test.go Normal file
View file

@ -0,0 +1,384 @@
package memory
import (
"context"
"encoding/json"
"os"
"path/filepath"
"testing"
"time"
"github.com/sipeed/picoclaw/pkg/providers"
)
func writeJSONSession(
t *testing.T, dir string, filename string, sess jsonSession,
) {
t.Helper()
data, err := json.MarshalIndent(sess, "", " ")
if err != nil {
t.Fatalf("marshal session: %v", err)
}
err = os.WriteFile(filepath.Join(dir, filename), data, 0o644)
if err != nil {
t.Fatalf("write session file: %v", err)
}
}
func TestMigrateFromJSON_Basic(t *testing.T) {
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
writeJSONSession(t, sessionsDir, "test.json", jsonSession{
Key: "test",
Messages: []providers.Message{
{Role: "user", Content: "hello"},
{Role: "assistant", Content: "hi"},
},
Summary: "A greeting.",
Created: time.Now(),
Updated: time.Now(),
})
count, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("MigrateFromJSON: %v", err)
}
if count != 1 {
t.Errorf("expected 1 migrated, got %d", count)
}
history, err := store.GetHistory(ctx, "test")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 2 {
t.Fatalf("expected 2 messages, got %d", len(history))
}
if history[0].Content != "hello" || history[1].Content != "hi" {
t.Errorf("unexpected messages: %+v", history)
}
summary, err := store.GetSummary(ctx, "test")
if err != nil {
t.Fatalf("GetSummary: %v", err)
}
if summary != "A greeting." {
t.Errorf("summary = %q", summary)
}
}
func TestMigrateFromJSON_WithToolCalls(t *testing.T) {
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
writeJSONSession(t, sessionsDir, "tools.json", jsonSession{
Key: "tools",
Messages: []providers.Message{
{
Role: "assistant",
Content: "Searching...",
ToolCalls: []providers.ToolCall{
{
ID: "call_1",
Type: "function",
Function: &providers.FunctionCall{
Name: "web_search",
Arguments: `{"q":"test"}`,
},
},
},
},
{
Role: "tool",
Content: "result",
ToolCallID: "call_1",
},
},
Created: time.Now(),
Updated: time.Now(),
})
count, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("MigrateFromJSON: %v", err)
}
if count != 1 {
t.Errorf("expected 1, got %d", count)
}
history, err := store.GetHistory(ctx, "tools")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 2 {
t.Fatalf("expected 2 messages, got %d", len(history))
}
if len(history[0].ToolCalls) != 1 {
t.Fatalf("expected 1 tool call, got %d", len(history[0].ToolCalls))
}
if history[0].ToolCalls[0].Function.Name != "web_search" {
t.Errorf("function = %q", history[0].ToolCalls[0].Function.Name)
}
if history[1].ToolCallID != "call_1" {
t.Errorf("ToolCallID = %q", history[1].ToolCallID)
}
}
func TestMigrateFromJSON_MultipleFiles(t *testing.T) {
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
for i := 0; i < 3; i++ {
key := string(rune('a' + i))
writeJSONSession(t, sessionsDir, key+".json", jsonSession{
Key: key,
Messages: []providers.Message{{Role: "user", Content: "msg " + key}},
Created: time.Now(),
Updated: time.Now(),
})
}
count, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("MigrateFromJSON: %v", err)
}
if count != 3 {
t.Errorf("expected 3, got %d", count)
}
for i := 0; i < 3; i++ {
key := string(rune('a' + i))
history, histErr := store.GetHistory(ctx, key)
if histErr != nil {
t.Fatalf("GetHistory(%q): %v", key, histErr)
}
if len(history) != 1 {
t.Errorf("session %q: expected 1 msg, got %d", key, len(history))
}
}
}
func TestMigrateFromJSON_InvalidJSON(t *testing.T) {
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
// One valid, one invalid.
writeJSONSession(t, sessionsDir, "good.json", jsonSession{
Key: "good",
Messages: []providers.Message{{Role: "user", Content: "ok"}},
Created: time.Now(),
Updated: time.Now(),
})
err := os.WriteFile(
filepath.Join(sessionsDir, "bad.json"),
[]byte("{invalid json"),
0o644,
)
if err != nil {
t.Fatalf("write bad file: %v", err)
}
count, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("MigrateFromJSON: %v", err)
}
if count != 1 {
t.Errorf("expected 1 (bad file skipped), got %d", count)
}
history, err := store.GetHistory(ctx, "good")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Errorf("expected 1 message, got %d", len(history))
}
}
func TestMigrateFromJSON_RenamesFiles(t *testing.T) {
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
writeJSONSession(t, sessionsDir, "rename.json", jsonSession{
Key: "rename",
Messages: []providers.Message{{Role: "user", Content: "hi"}},
Created: time.Now(),
Updated: time.Now(),
})
_, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("MigrateFromJSON: %v", err)
}
// Original .json should not exist.
_, statErr := os.Stat(filepath.Join(sessionsDir, "rename.json"))
if !os.IsNotExist(statErr) {
t.Error("rename.json should have been renamed")
}
// .json.migrated should exist.
_, statErr = os.Stat(
filepath.Join(sessionsDir, "rename.json.migrated"),
)
if statErr != nil {
t.Errorf("rename.json.migrated should exist: %v", statErr)
}
}
func TestMigrateFromJSON_Idempotent(t *testing.T) {
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
writeJSONSession(t, sessionsDir, "idem.json", jsonSession{
Key: "idem",
Messages: []providers.Message{{Role: "user", Content: "once"}},
Created: time.Now(),
Updated: time.Now(),
})
count1, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("first migration: %v", err)
}
if count1 != 1 {
t.Errorf("first run: expected 1, got %d", count1)
}
// Second run should find only .migrated files, skip them.
count2, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("second migration: %v", err)
}
if count2 != 0 {
t.Errorf("second run: expected 0, got %d", count2)
}
history, err := store.GetHistory(ctx, "idem")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Errorf("expected 1 message, got %d", len(history))
}
}
func TestMigrateFromJSON_ColonInKey(t *testing.T) {
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
// File is named telegram_123 (sanitized), but the key inside is telegram:123.
writeJSONSession(t, sessionsDir, "telegram_123.json", jsonSession{
Key: "telegram:123",
Messages: []providers.Message{{Role: "user", Content: "from telegram"}},
Created: time.Now(),
Updated: time.Now(),
})
count, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("MigrateFromJSON: %v", err)
}
if count != 1 {
t.Errorf("expected 1, got %d", count)
}
// Accessible via the original key "telegram:123".
history, err := store.GetHistory(ctx, "telegram:123")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history) != 1 {
t.Fatalf("expected 1 message, got %d", len(history))
}
if history[0].Content != "from telegram" {
t.Errorf("content = %q", history[0].Content)
}
// In the file-based store, "telegram:123" and "telegram_123" both
// sanitize to the same filename, so they share storage. This is
// expected — the colon-to-underscore mapping is a one-way function.
history2, err := store.GetHistory(ctx, "telegram_123")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
if len(history2) != 1 {
t.Errorf("expected 1 (same file), got %d", len(history2))
}
}
func TestMigrateFromJSON_RetryAfterCrash(t *testing.T) {
// Simulates a crash during migration: first run writes messages
// but doesn't rename the .json file. Second run must replace
// (not duplicate) the messages thanks to SetHistory semantics.
sessionsDir := t.TempDir()
store := newTestStore(t)
ctx := context.Background()
writeJSONSession(t, sessionsDir, "retry.json", jsonSession{
Key: "retry",
Messages: []providers.Message{
{Role: "user", Content: "one"},
{Role: "assistant", Content: "two"},
},
Created: time.Now(),
Updated: time.Now(),
})
// First migration succeeds — writes messages and renames file.
count, err := MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("first migration: %v", err)
}
if count != 1 {
t.Fatalf("expected 1, got %d", count)
}
// Simulate "crash before rename": restore the .json file.
src := filepath.Join(sessionsDir, "retry.json.migrated")
dst := filepath.Join(sessionsDir, "retry.json")
if renameErr := os.Rename(src, dst); renameErr != nil {
t.Fatalf("restore .json: %v", renameErr)
}
// Second migration should re-import without duplicating messages.
count, err = MigrateFromJSON(ctx, sessionsDir, store)
if err != nil {
t.Fatalf("second migration: %v", err)
}
if count != 1 {
t.Fatalf("expected 1, got %d", count)
}
history, err := store.GetHistory(ctx, "retry")
if err != nil {
t.Fatalf("GetHistory: %v", err)
}
// Must be exactly 2 messages (not 4 from duplication).
if len(history) != 2 {
t.Fatalf("expected 2 messages (no duplicates), got %d", len(history))
}
if history[0].Content != "one" || history[1].Content != "two" {
t.Errorf("unexpected messages: %+v", history)
}
}
func TestMigrateFromJSON_NonexistentDir(t *testing.T) {
store := newTestStore(t)
ctx := context.Background()
count, err := MigrateFromJSON(ctx, "/nonexistent/path", store)
if err != nil {
t.Fatalf("MigrateFromJSON: %v", err)
}
if count != 0 {
t.Errorf("expected 0, got %d", count)
}
}

42
pkg/memory/store.go Normal file
View file

@ -0,0 +1,42 @@
package memory
import (
"context"
"github.com/sipeed/picoclaw/pkg/providers"
)
// Store defines an interface for persistent session storage.
// Each method is an atomic operation — there is no separate Save() call.
type Store interface {
// AddMessage appends a simple text message to a session.
AddMessage(ctx context.Context, sessionKey, role, content string) error
// AddFullMessage appends a complete message (with tool calls, etc.) to a session.
AddFullMessage(ctx context.Context, sessionKey string, msg providers.Message) error
// GetHistory returns all messages for a session in insertion order.
// Returns an empty slice (not nil) if the session does not exist.
GetHistory(ctx context.Context, sessionKey string) ([]providers.Message, error)
// GetSummary returns the conversation summary for a session.
// Returns an empty string if no summary exists.
GetSummary(ctx context.Context, sessionKey string) (string, error)
// SetSummary updates the conversation summary for a session.
SetSummary(ctx context.Context, sessionKey, summary string) error
// TruncateHistory removes all but the last keepLast messages from a session.
// If keepLast <= 0, all messages are removed.
TruncateHistory(ctx context.Context, sessionKey string, keepLast int) error
// SetHistory replaces all messages in a session with the provided history.
SetHistory(ctx context.Context, sessionKey string, history []providers.Message) error
// Compact reclaims storage by physically removing logically truncated
// data. Backends that do not accumulate dead data may return nil.
Compact(ctx context.Context, sessionKey string) error
// Close releases any resources held by the store.
Close() error
}

9
pkg/tools/message.go
View file

@ -3,6 +3,7 @@ package tools
import ( import (
"context" "context"
"fmt" "fmt"
"sync/atomic"
) )
type SendCallback func(channel, chatID, content string) error type SendCallback func(channel, chatID, content string) error
@ -11,7 +12,7 @@ type MessageTool struct {
sendCallback SendCallback sendCallback SendCallback
defaultChannel string defaultChannel string
defaultChatID string defaultChatID string
sentInRound bool // Tracks whether a message was sent in the current processing round sentInRound atomic.Bool // Tracks whether a message was sent in the current processing round
} }
func NewMessageTool() *MessageTool { func NewMessageTool() *MessageTool {
@ -50,12 +51,12 @@ func (t *MessageTool) Parameters() map[string]any {
func (t *MessageTool) SetContext(channel, chatID string) { func (t *MessageTool) SetContext(channel, chatID string) {
t.defaultChannel = channel t.defaultChannel = channel
t.defaultChatID = chatID t.defaultChatID = chatID
t.sentInRound = false // Reset send tracking for new processing round t.sentInRound.Store(false) // Reset send tracking for new processing round
} }
// HasSentInRound returns true if the message tool sent a message during the current round. // HasSentInRound returns true if the message tool sent a message during the current round.
func (t *MessageTool) HasSentInRound() bool { func (t *MessageTool) HasSentInRound() bool {
return t.sentInRound return t.sentInRound.Load()
} }
func (t *MessageTool) SetSendCallback(callback SendCallback) { func (t *MessageTool) SetSendCallback(callback SendCallback) {
@ -94,7 +95,7 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolRes
} }
} }
t.sentInRound = true t.sentInRound.Store(true)
// Silent: user already received the message directly // Silent: user already received the message directly
return &ToolResult{ return &ToolResult{
ForLLM: fmt.Sprintf("Message sent to %s:%s", channel, chatID), ForLLM: fmt.Sprintf("Message sent to %s:%s", channel, chatID),

69
pkg/tools/toolloop.go
View file

@ -10,6 +10,7 @@ import (
"context" "context"
"encoding/json" "encoding/json"
"fmt" "fmt"
"sync"
"github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/logger"
"github.com/sipeed/picoclaw/pkg/providers" "github.com/sipeed/picoclaw/pkg/providers"
@ -121,37 +122,53 @@ func RunToolLoop(
} }
messages = append(messages, assistantMsg) messages = append(messages, assistantMsg)
// 7. Execute tool calls // 7. Execute tool calls in parallel
for _, tc := range normalizedToolCalls { type indexedResult struct {
argsJSON, _ := json.Marshal(tc.Arguments) result *ToolResult
argsPreview := utils.Truncate(string(argsJSON), 200) tc providers.ToolCall
logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview), }
map[string]any{
"tool": tc.Name,
"iteration": iteration,
})
// Execute tool (no async callback for subagents - they run independently) results := make([]indexedResult, len(normalizedToolCalls))
var toolResult *ToolResult var wg sync.WaitGroup
if config.Tools != nil {
toolResult = config.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, channel, chatID, nil) for i, tc := range normalizedToolCalls {
} else { results[i].tc = tc
toolResult = ErrorResult("No tools available")
wg.Add(1)
go func(idx int, tc providers.ToolCall) {
defer wg.Done()
argsJSON, _ := json.Marshal(tc.Arguments)
argsPreview := utils.Truncate(string(argsJSON), 200)
logger.InfoCF("toolloop", fmt.Sprintf("Tool call: %s(%s)", tc.Name, argsPreview),
map[string]any{
"tool": tc.Name,
"iteration": iteration,
})
var toolResult *ToolResult
if config.Tools != nil {
toolResult = config.Tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, channel, chatID, nil)
} else {
toolResult = ErrorResult("No tools available")
}
results[idx].result = toolResult
}(i, tc)
}
wg.Wait()
// Append results in original order
for _, r := range results {
contentForLLM := r.result.ForLLM
if contentForLLM == "" && r.result.Err != nil {
contentForLLM = r.result.Err.Error()
} }
// Determine content for LLM messages = append(messages, providers.Message{
contentForLLM := toolResult.ForLLM
if contentForLLM == "" && toolResult.Err != nil {
contentForLLM = toolResult.Err.Error()
}
// Add tool result message
toolResultMsg := providers.Message{
Role: "tool", Role: "tool",
Content: contentForLLM, Content: contentForLLM,
ToolCallID: tc.ID, ToolCallID: r.tc.ID,
} })
messages = append(messages, toolResultMsg)
} }
} }

106
pkg/tools/web.go
View file

@ -486,6 +486,88 @@ func (p *PerplexitySearchProvider) Search(ctx context.Context, query string, cou
return "", fmt.Errorf("all api keys failed, last error: %w", lastErr) return "", fmt.Errorf("all api keys failed, last error: %w", lastErr)
} }
type GLMSearchProvider struct {
apiKey string
baseURL string
searchEngine string
proxy string
client *http.Client
}
func (p *GLMSearchProvider) Search(ctx context.Context, query string, count int) (string, error) {
searchURL := p.baseURL
if searchURL == "" {
searchURL = "https://open.bigmodel.cn/api/paas/v4/web_search"
}
payload := map[string]any{
"search_query": query,
"search_engine": p.searchEngine,
"search_intent": false,
"count": count,
"content_size": "medium",
}
bodyBytes, err := json.Marshal(payload)
if err != nil {
return "", fmt.Errorf("failed to marshal payload: %w", err)
}
req, err := http.NewRequestWithContext(ctx, "POST", searchURL, bytes.NewReader(bodyBytes))
if err != nil {
return "", fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+p.apiKey)
resp, err := p.client.Do(req)
if err != nil {
return "", fmt.Errorf("request failed: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
if err != nil {
return "", fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("GLM Search API error (status %d): %s", resp.StatusCode, string(body))
}
var searchResp struct {
SearchResult []struct {
Title string `json:"title"`
Content string `json:"content"`
Link string `json:"link"`
} `json:"search_result"`
}
if err := json.Unmarshal(body, &searchResp); err != nil {
return "", fmt.Errorf("failed to parse response: %w", err)
}
results := searchResp.SearchResult
if len(results) == 0 {
return fmt.Sprintf("No results for: %s", query), nil
}
var lines []string
lines = append(lines, fmt.Sprintf("Results for: %s (via GLM Search)", query))
for i, item := range results {
if i >= count {
break
}
lines = append(lines, fmt.Sprintf("%d. %s\n %s", i+1, item.Title, item.Link))
if item.Content != "" {
lines = append(lines, fmt.Sprintf(" %s", item.Content))
}
}
return strings.Join(lines, "\n"), nil
}
type WebSearchTool struct { type WebSearchTool struct {
provider SearchProvider provider SearchProvider
maxResults int maxResults int
@ -504,6 +586,11 @@ type WebSearchToolOptions struct {
PerplexityAPIKeys []string PerplexityAPIKeys []string
PerplexityMaxResults int PerplexityMaxResults int
PerplexityEnabled bool PerplexityEnabled bool
GLMSearchAPIKey string
GLMSearchBaseURL string
GLMSearchEngine string
GLMSearchMaxResults int
GLMSearchEnabled bool
Proxy string Proxy string
} }
@ -553,6 +640,25 @@ func NewWebSearchTool(opts WebSearchToolOptions) (*WebSearchTool, error) {
if opts.DuckDuckGoMaxResults > 0 { if opts.DuckDuckGoMaxResults > 0 {
maxResults = opts.DuckDuckGoMaxResults maxResults = opts.DuckDuckGoMaxResults
} }
} else if opts.GLMSearchEnabled && opts.GLMSearchAPIKey != "" {
client, err := createHTTPClient(opts.Proxy, searchTimeout)
if err != nil {
return nil, fmt.Errorf("failed to create HTTP client for GLM Search: %w", err)
}
searchEngine := opts.GLMSearchEngine
if searchEngine == "" {
searchEngine = "search_std"
}
provider = &GLMSearchProvider{
apiKey: opts.GLMSearchAPIKey,
baseURL: opts.GLMSearchBaseURL,
searchEngine: searchEngine,
proxy: opts.Proxy,
client: client,
}
if opts.GLMSearchMaxResults > 0 {
maxResults = opts.GLMSearchMaxResults
}
} else { } else {
return nil, nil return nil, nil
} }

136
pkg/tools/web_test.go
View file

@ -683,7 +683,7 @@ func TestWebTool_TavilySearch_Success(t *testing.T) {
} }
func TestAPIKeyPool(t *testing.T) { func TestAPIKeyPool(t *testing.T) {
pool := NewAPIKeyPool([]string{"key1", " key2 ", "key3"}) pool := NewAPIKeyPool([]string{"key1", "key2", "key3"})
if len(pool.keys) != 3 { if len(pool.keys) != 3 {
t.Fatalf("expected 3 keys, got %d", len(pool.keys)) t.Fatalf("expected 3 keys, got %d", len(pool.keys))
} }
@ -705,7 +705,7 @@ func TestAPIKeyPool(t *testing.T) {
t.Errorf("expected key1, got %s", k) t.Errorf("expected key1, got %s", k)
} }
emptyPool := NewAPIKeyPool([]string{" "}) emptyPool := NewAPIKeyPool([]string{})
if k := emptyPool.Get(); k != "" { if k := emptyPool.Get(); k != "" {
t.Errorf("expected empty string, got %s", k) t.Errorf("expected empty string, got %s", k)
} }
@ -779,3 +779,135 @@ func TestWebTool_TavilySearch_Failover(t *testing.T) {
t.Errorf("Expected failover to second key and success result, got: %s", result.ForUser) t.Errorf("Expected failover to second key and success result, got: %s", result.ForUser)
} }
} }
func TestWebTool_GLMSearch_Success(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method != "POST" {
t.Errorf("Expected POST request, got %s", r.Method)
}
if r.Header.Get("Content-Type") != "application/json" {
t.Errorf("Expected Content-Type application/json, got %s", r.Header.Get("Content-Type"))
}
if r.Header.Get("Authorization") != "Bearer test-glm-key" {
t.Errorf("Expected Authorization Bearer test-glm-key, got %s", r.Header.Get("Authorization"))
}
var payload map[string]any
json.NewDecoder(r.Body).Decode(&payload)
if payload["search_query"] != "test query" {
t.Errorf("Expected search_query 'test query', got %v", payload["search_query"])
}
if payload["search_engine"] != "search_std" {
t.Errorf("Expected search_engine 'search_std', got %v", payload["search_engine"])
}
response := map[string]any{
"id": "web-search-test",
"created": 1709568000,
"search_result": []map[string]any{
{
"title": "Test GLM Result",
"content": "GLM search snippet",
"link": "https://example.com/glm",
"media": "Example",
"publish_date": "2026-03-04",
},
},
}
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
json.NewEncoder(w).Encode(response)
}))
defer server.Close()
tool, err := NewWebSearchTool(WebSearchToolOptions{
GLMSearchEnabled: true,
GLMSearchAPIKey: "test-glm-key",
GLMSearchBaseURL: server.URL,
GLMSearchEngine: "search_std",
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
result := tool.Execute(context.Background(), map[string]any{
"query": "test query",
})
if result.IsError {
t.Errorf("Expected success, got IsError=true: %s", result.ForLLM)
}
if !strings.Contains(result.ForUser, "Test GLM Result") {
t.Errorf("Expected 'Test GLM Result' in output, got: %s", result.ForUser)
}
if !strings.Contains(result.ForUser, "https://example.com/glm") {
t.Errorf("Expected URL in output, got: %s", result.ForUser)
}
if !strings.Contains(result.ForUser, "via GLM Search") {
t.Errorf("Expected 'via GLM Search' in output, got: %s", result.ForUser)
}
}
func TestWebTool_GLMSearch_APIError(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusUnauthorized)
w.Write([]byte(`{"error":"invalid api key"}`))
}))
defer server.Close()
tool, err := NewWebSearchTool(WebSearchToolOptions{
GLMSearchEnabled: true,
GLMSearchAPIKey: "bad-key",
GLMSearchBaseURL: server.URL,
GLMSearchEngine: "search_std",
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
result := tool.Execute(context.Background(), map[string]any{
"query": "test query",
})
if !result.IsError {
t.Errorf("Expected IsError=true for 401 response")
}
if !strings.Contains(result.ForLLM, "status 401") {
t.Errorf("Expected status 401 in error, got: %s", result.ForLLM)
}
}
func TestWebTool_GLMSearch_Priority(t *testing.T) {
// GLM Search should only be selected when all other providers are disabled
tool, err := NewWebSearchTool(WebSearchToolOptions{
DuckDuckGoEnabled: true,
DuckDuckGoMaxResults: 5,
GLMSearchEnabled: true,
GLMSearchAPIKey: "test-key",
GLMSearchBaseURL: "https://example.com",
GLMSearchEngine: "search_std",
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
// DuckDuckGo should win over GLM Search
if _, ok := tool.provider.(*DuckDuckGoSearchProvider); !ok {
t.Errorf("Expected DuckDuckGoSearchProvider when both enabled, got %T", tool.provider)
}
// With DuckDuckGo disabled, GLM Search should be selected
tool2, err := NewWebSearchTool(WebSearchToolOptions{
DuckDuckGoEnabled: false,
GLMSearchEnabled: true,
GLMSearchAPIKey: "test-key",
GLMSearchBaseURL: "https://example.com",
GLMSearchEngine: "search_std",
})
if err != nil {
t.Fatalf("NewWebSearchTool() error: %v", err)
}
if _, ok := tool2.provider.(*GLMSearchProvider); !ok {
t.Errorf("Expected GLMSearchProvider when only GLM enabled, got %T", tool2.provider)
}
}

23
pkg/utils/media.go
View file

@ -3,6 +3,7 @@ package utils
import ( import (
"io" "io"
"net/http" "net/http"
"net/url"
"os" "os"
"path/filepath" "path/filepath"
"strings" "strings"
@ -52,11 +53,12 @@ type DownloadOptions struct {
Timeout time.Duration Timeout time.Duration
ExtraHeaders map[string]string ExtraHeaders map[string]string
LoggerPrefix string LoggerPrefix string
ProxyURL string
} }
// DownloadFile downloads a file from URL to a local temp directory. // DownloadFile downloads a file from URL to a local temp directory.
// Returns the local file path or empty string on error. // Returns the local file path or empty string on error.
func DownloadFile(url, filename string, opts DownloadOptions) string { func DownloadFile(urlStr, filename string, opts DownloadOptions) string {
// Set defaults // Set defaults
if opts.Timeout == 0 { if opts.Timeout == 0 {
opts.Timeout = 60 * time.Second opts.Timeout = 60 * time.Second
@ -78,7 +80,7 @@ func DownloadFile(url, filename string, opts DownloadOptions) string {
localPath := filepath.Join(mediaDir, uuid.New().String()[:8]+"_"+safeName) localPath := filepath.Join(mediaDir, uuid.New().String()[:8]+"_"+safeName)
// Create HTTP request // Create HTTP request
req, err := http.NewRequest("GET", url, nil) req, err := http.NewRequest("GET", urlStr, nil)
if err != nil { if err != nil {
logger.ErrorCF(opts.LoggerPrefix, "Failed to create download request", map[string]any{ logger.ErrorCF(opts.LoggerPrefix, "Failed to create download request", map[string]any{
"error": err.Error(), "error": err.Error(),
@ -92,11 +94,24 @@ func DownloadFile(url, filename string, opts DownloadOptions) string {
} }
client := &http.Client{Timeout: opts.Timeout} client := &http.Client{Timeout: opts.Timeout}
if opts.ProxyURL != "" {
proxyURL, parseErr := url.Parse(opts.ProxyURL)
if parseErr != nil {
logger.ErrorCF(opts.LoggerPrefix, "Invalid proxy URL for download", map[string]any{
"error": parseErr.Error(),
"proxy": opts.ProxyURL,
})
return ""
}
client.Transport = &http.Transport{
Proxy: http.ProxyURL(proxyURL),
}
}
resp, err := client.Do(req) resp, err := client.Do(req)
if err != nil { if err != nil {
logger.ErrorCF(opts.LoggerPrefix, "Failed to download file", map[string]any{ logger.ErrorCF(opts.LoggerPrefix, "Failed to download file", map[string]any{
"error": err.Error(), "error": err.Error(),
"url": url, "url": urlStr,
}) })
return "" return ""
} }
@ -105,7 +120,7 @@ func DownloadFile(url, filename string, opts DownloadOptions) string {
if resp.StatusCode != http.StatusOK { if resp.StatusCode != http.StatusOK {
logger.ErrorCF(opts.LoggerPrefix, "File download returned non-200 status", map[string]any{ logger.ErrorCF(opts.LoggerPrefix, "File download returned non-200 status", map[string]any{
"status": resp.StatusCode, "status": resp.StatusCode,
"url": url, "url": urlStr,
}) })
return "" return ""
} }