From fd382e68e281c4863493d7ef86c24137e8234010 Mon Sep 17 00:00:00 2001 From: Boris Bliznioukov Date: Fri, 6 Mar 2026 14:30:37 +0100 Subject: [PATCH] feat(session): extract summarization into SessionManager with race-safe async merging Signed-off-by: Boris Bliznioukov --- README.fr.md | 6 +- README.ja.md | 6 +- README.md | 168 +++++---- README.pt-br.md | 6 +- README.vi.md | 6 +- README.zh.md | 12 +- config/config.example.json | 6 +- docs/design/DDR-session-summarization.md | 220 +++++++++++ pkg/agent/instance.go | 95 +++-- pkg/agent/loop.go | 211 +---------- pkg/config/config.go | 146 +++++++- pkg/config/config_test.go | 42 ++- pkg/config/defaults.go | 20 +- pkg/session/manager.go | 280 +++++++++++++- pkg/session/summarizer.go | 102 +++++ pkg/session/summarizer_test.go | 454 +++++++++++++++++++++++ 16 files changed, 1401 insertions(+), 379 deletions(-) create mode 100644 docs/design/DDR-session-summarization.md create mode 100644 pkg/session/summarizer.go create mode 100644 pkg/session/summarizer_test.go diff --git a/README.fr.md b/README.fr.md index 08a1926b6..61435bf22 100644 --- a/README.fr.md +++ b/README.fr.md @@ -885,7 +885,11 @@ picoclaw agent -m "Bonjour, comment ça va ?" { "agents": { "defaults": { - "model": "anthropic/claude-opus-4-5" + "model": "anthropic/claude-opus-4-5", + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "providers": { diff --git a/README.ja.md b/README.ja.md index c4c5b27a0..d53b2bfb4 100644 --- a/README.ja.md +++ b/README.ja.md @@ -838,7 +838,11 @@ HEARTBEAT_OK 応答 ユーザーが直接結果を受け取る { "agents": { "defaults": { - "model": "anthropic/claude-opus-4-5" + "model": "anthropic/claude-opus-4-5", + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "providers": { diff --git a/README.md b/README.md index db127a85f..1812cd551 100644 --- a/README.md +++ b/README.md @@ -45,12 +45,11 @@ > [!CAUTION] > **🚨 SECURITY & OFFICIAL CHANNELS / 安全声明** > -> * **NO CRYPTO:** PicoClaw has **NO** official token/coin. All claims on `pump.fun` or other trading platforms are **SCAMS**. -> -> * **OFFICIAL DOMAIN:** The **ONLY** official website is **[picoclaw.io](https://picoclaw.io)**, and company website is **[sipeed.com](https://sipeed.com)** -> * **Warning:** Many `.ai/.org/.com/.net/...` domains are registered by third parties. -> * **Warning:** picoclaw is in early development now and may have unresolved network security issues. Do not deploy to production environments before the v1.0 release. -> * **Note:** picoclaw has recently merged a lot of PRs, which may result in a larger memory footprint (10–20MB) in the latest versions. We plan to prioritize resource optimization as soon as the current feature set reaches a stable state. +> - **NO CRYPTO:** PicoClaw has **NO** official token/coin. All claims on `pump.fun` or other trading platforms are **SCAMS**. +> - **OFFICIAL DOMAIN:** The **ONLY** official website is **[picoclaw.io](https://picoclaw.io)**, and company website is **[sipeed.com](https://sipeed.com)** +> - **Warning:** Many `.ai/.org/.com/.net/...` domains are registered by third parties. +> - **Warning:** picoclaw is in early development now and may have unresolved network security issues. Do not deploy to production environments before the v1.0 release. +> - **Note:** picoclaw has recently merged a lot of PRs, which may result in a larger memory footprint (10–20MB) in the latest versions. We plan to prioritize resource optimization as soon as the current feature set reaches a stable state. ## 📢 News @@ -234,7 +233,11 @@ picoclaw onboard "model_name": "gpt4", "max_tokens": 8192, "temperature": 0.7, - "max_tool_iterations": 20 + "max_tool_iterations": 20, + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "model_list": [ @@ -286,13 +289,13 @@ picoclaw onboard **3. Get API Keys** -* **LLM Provider**: [OpenRouter](https://openrouter.ai/keys) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) · [Anthropic](https://console.anthropic.com) · [OpenAI](https://platform.openai.com) · [Gemini](https://aistudio.google.com/api-keys) -* **Web Search** (optional): - * [Brave Search](https://brave.com/search/api) - Paid ($5/1000 queries, ~$5-6/month) - * [Perplexity](https://www.perplexity.ai) - AI-powered search with chat interface - * [SearXNG](https://github.com/searxng/searxng) - Self-hosted metasearch engine (free, no API key needed) - * [Tavily](https://tavily.com) - Optimized for AI Agents (1000 requests/month) - * DuckDuckGo - Built-in fallback (no API key required) +- **LLM Provider**: [OpenRouter](https://openrouter.ai/keys) · [Zhipu](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) · [Anthropic](https://console.anthropic.com) · [OpenAI](https://platform.openai.com) · [Gemini](https://aistudio.google.com/api-keys) +- **Web Search** (optional): + - [Brave Search](https://brave.com/search/api) - Paid ($5/1000 queries, ~$5-6/month) + - [Perplexity](https://www.perplexity.ai) - AI-powered search with chat interface + - [SearXNG](https://github.com/searxng/searxng) - Self-hosted metasearch engine (free, no API key needed) + - [Tavily](https://tavily.com) - Optimized for AI Agents (1000 requests/month) + - DuckDuckGo - Built-in fallback (no API key required) > **Note**: See `config.example.json` for a complete configuration template. @@ -312,24 +315,24 @@ Talk to your picoclaw through Telegram, Discord, WhatsApp, DingTalk, LINE, or We > **Note**: All webhook-based channels (LINE, WeCom, etc.) are served on a single shared Gateway HTTP server (`gateway.host`:`gateway.port`, default `127.0.0.1:18790`). There are no per-channel ports to configure. Note: Feishu uses WebSocket/SDK mode and does not use the shared HTTP webhook server. -| Channel | Setup | -| ------------ | ---------------------------------- | -| **Telegram** | Easy (just a token) | -| **Discord** | Easy (bot token + intents) | -| **WhatsApp** | Easy (native: QR scan; or bridge URL) | -| **QQ** | Easy (AppID + AppSecret) | -| **DingTalk** | Medium (app credentials) | -| **LINE** | Medium (credentials + webhook URL) | -| **WeCom AI Bot** | Medium (Token + AES key) | +| Channel | Setup | +| ---------------- | ------------------------------------- | +| **Telegram** | Easy (just a token) | +| **Discord** | Easy (bot token + intents) | +| **WhatsApp** | Easy (native: QR scan; or bridge URL) | +| **QQ** | Easy (AppID + AppSecret) | +| **DingTalk** | Medium (app credentials) | +| **LINE** | Medium (credentials + webhook URL) | +| **WeCom AI Bot** | Medium (Token + AES key) |
Telegram (Recommended) **1. Create a bot** -* Open Telegram, search `@BotFather` -* Send `/newbot`, follow prompts -* Copy the token +- Open Telegram, search `@BotFather` +- Send `/newbot`, follow prompts +- Copy the token **2. Configure** @@ -367,18 +370,19 @@ If command registration fails (network/API transient errors), the channel still **1. Create a bot** -* Go to -* Create an application → Bot → Add Bot -* Copy the bot token +- Go to +- Create an application → Bot → Add Bot +- Copy the bot token **2. Enable intents** -* In the Bot settings, enable **MESSAGE CONTENT INTENT** -* (Optional) Enable **SERVER MEMBERS INTENT** if you plan to use allow lists based on member data +- In the Bot settings, enable **MESSAGE CONTENT INTENT** +- (Optional) Enable **SERVER MEMBERS INTENT** if you plan to use allow lists based on member data **3. Get your User ID** -* Discord Settings → Advanced → enable **Developer Mode** -* Right-click your avatar → **Copy User ID** + +- Discord Settings → Advanced → enable **Developer Mode** +- Right-click your avatar → **Copy User ID** **4. Configure** @@ -396,10 +400,10 @@ If command registration fails (network/API transient errors), the channel still **5. Invite the bot** -* OAuth2 → URL Generator -* Scopes: `bot` -* Bot Permissions: `Send Messages`, `Read Message History` -* Open the generated invite URL and add the bot to your server +- OAuth2 → URL Generator +- Scopes: `bot` +- Bot Permissions: `Send Messages`, `Read Message History` +- Open the generated invite URL and add the bot to your server **Optional: Group trigger mode** @@ -500,9 +504,9 @@ picoclaw gateway **1. Create a bot** -* Go to [Open Platform](https://open.dingtalk.com/) -* Create an internal app -* Copy Client ID and Client Secret +- Go to [Open Platform](https://open.dingtalk.com/) +- Create an internal app +- Copy Client ID and Client Secret **2. Configure** @@ -526,6 +530,7 @@ picoclaw gateway ```bash picoclaw gateway ``` +
@@ -591,8 +596,8 @@ See [WeCom AI Bot Configuration Guide](docs/channels/wecom/wecom_aibot/README.zh **1. Create a bot** -* Go to WeCom Admin Console → Group Chat → Add Group Bot -* Copy the webhook URL (format: `https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=xxx`) +- Go to WeCom Admin Console → Group Chat → Add Group Bot +- Copy the webhook URL (format: `https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=xxx`) **2. Configure** @@ -617,15 +622,15 @@ See [WeCom AI Bot Configuration Guide](docs/channels/wecom/wecom_aibot/README.zh **1. Create an app** -* Go to WeCom Admin Console → App Management → Create App -* Copy **AgentId** and **Secret** -* Go to "My Company" page, copy **CorpID** +- Go to WeCom Admin Console → App Management → Create App +- Copy **AgentId** and **Secret** +- Go to "My Company" page, copy **CorpID** **2. Configure receive message** -* In App details, click "Receive Message" → "Set API" -* Set URL to `http://your-server:18790/webhook/wecom-app` -* Generate **Token** and **EncodingAESKey** +- In App details, click "Receive Message" → "Set API" +- Set URL to `http://your-server:18790/webhook/wecom-app` +- Generate **Token** and **EncodingAESKey** **3. Configure** @@ -658,9 +663,9 @@ picoclaw gateway **1. Create an AI Bot** -* Go to WeCom Admin Console → App Management → AI Bot -* In the AI Bot settings, configure callback URL: `http://your-server:18791/webhook/wecom-aibot` -* Copy **Token** and click "Random Generate" for **EncodingAESKey** +- Go to WeCom Admin Console → App Management → AI Bot +- In the AI Bot settings, configure callback URL: `http://your-server:18791/webhook/wecom-aibot` +- Copy **Token** and click "Random Generate" for **EncodingAESKey** **2. Configure** @@ -703,10 +708,10 @@ Config file: `~/.picoclaw/config.json` You can override default paths using environment variables. This is useful for portable installations, containerized deployments, or running picoclaw as a system service. These variables are independent and control different paths. -| Variable | Description | Default Path | -|-------------------|-----------------------------------------------------------------------------------------------------------------------------------------|---------------------------| +| Variable | Description | Default Path | +| ----------------- | ------------------------------------------------------------------------------------------------------------------------------------- | ------------------------- | | `PICOCLAW_CONFIG` | Overrides the path to the configuration file. This directly tells picoclaw which `config.json` to load, ignoring all other locations. | `~/.picoclaw/config.json` | -| `PICOCLAW_HOME` | Overrides the root directory for picoclaw data. This changes the default location of the `workspace` and other data directories. | `~/.picoclaw` | +| `PICOCLAW_HOME` | Overrides the root directory for picoclaw data. This changes the default location of the `workspace` and other data directories. | `~/.picoclaw` | **Examples:** @@ -802,12 +807,12 @@ When `restrict_to_workspace: true`, the following tools are sandboxed: Even with `restrict_to_workspace: false`, the `exec` tool blocks these dangerous commands: -* `rm -rf`, `del /f`, `rmdir /s` — Bulk deletion -* `format`, `mkfs`, `diskpart` — Disk formatting -* `dd if=` — Disk imaging -* Writing to `/dev/sd[a-z]` — Direct disk writes -* `shutdown`, `reboot`, `poweroff` — System shutdown -* Fork bomb `:(){ :|:& };:` +- `rm -rf`, `del /f`, `rmdir /s` — Bulk deletion +- `format`, `mkfs`, `diskpart` — Disk formatting +- `dd if=` — Disk imaging +- Writing to `/dev/sd[a-z]` — Direct disk writes +- `shutdown`, `reboot`, `poweroff` — System shutdown +- Fork bomb `:(){ :|:& };:` #### Error Examples @@ -933,8 +938,8 @@ The subagent has access to tools (message, web_search, etc.) and can communicate **Environment variables:** -* `PICOCLAW_HEARTBEAT_ENABLED=false` to disable -* `PICOCLAW_HEARTBEAT_INTERVAL=60` to change interval +- `PICOCLAW_HEARTBEAT_ENABLED=false` to disable +- `PICOCLAW_HEARTBEAT_INTERVAL=60` to change interval ### Providers @@ -967,7 +972,7 @@ This design also enables **multi-agent support** with flexible provider selectio #### 📋 All Supported Vendors | Vendor | `model` Prefix | Default API Base | Protocol | API Key | -| ------------------- | ----------------- |-----------------------------------------------------| --------- | ---------------------------------------------------------------- | +| ------------------- | ----------------- | --------------------------------------------------- | --------- | ---------------------------------------------------------------- | | **OpenAI** | `openai/` | `https://api.openai.com/v1` | OpenAI | [Get Key](https://platform.openai.com) | | **Anthropic** | `anthropic/` | `https://api.anthropic.com/v1` | Anthropic | [Get Key](https://console.anthropic.com) | | **智谱 AI (GLM)** | `zhipu/` | `https://open.bigmodel.cn/api/paas/v4` | OpenAI | [Get Key](https://open.bigmodel.cn/usercenter/proj-mgmt/apikeys) | @@ -979,7 +984,7 @@ This design also enables **multi-agent support** with flexible provider selectio | **NVIDIA** | `nvidia/` | `https://integrate.api.nvidia.com/v1` | OpenAI | [Get Key](https://build.nvidia.com) | | **Ollama** | `ollama/` | `http://localhost:11434/v1` | OpenAI | Local (no key needed) | | **OpenRouter** | `openrouter/` | `https://openrouter.ai/api/v1` | OpenAI | [Get Key](https://openrouter.ai/keys) | -| **LiteLLM Proxy** | `litellm/` | `http://localhost:4000/v1 | OpenAI | Your LiteLLM proxy key | +| **LiteLLM Proxy** | `litellm/` | `http://localhost:4000/v1 | OpenAI | Your LiteLLM proxy key | | **VLLM** | `vllm/` | `http://localhost:8000/v1` | OpenAI | Local | | **Cerebras** | `cerebras/` | `https://api.cerebras.ai/v1` | OpenAI | [Get Key](https://cerebras.ai) | | **火山引擎** | `volcengine/` | `https://ark.cn-beijing.volces.com/api/v3` | OpenAI | [Get Key](https://console.volcengine.com) | @@ -1176,7 +1181,7 @@ This keeps the runtime lightweight while making new OpenAI-compatible backends m **1. Get API key and base URL** -* Get [API key](https://bigmodel.cn/usercenter/proj-mgmt/apikeys) +- Get [API key](https://bigmodel.cn/usercenter/proj-mgmt/apikeys) **2. Configure** @@ -1215,7 +1220,11 @@ picoclaw agent -m "Hello" { "agents": { "defaults": { - "model": "anthropic/claude-opus-4-5" + "model": "anthropic/claude-opus-4-5", + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "session": { @@ -1314,9 +1323,9 @@ picoclaw agent -m "Hello" PicoClaw supports scheduled reminders and recurring tasks through the `cron` tool: -* **One-time reminders**: "Remind me in 10 minutes" → triggers once after 10min -* **Recurring tasks**: "Remind me every 2 hours" → triggers every 2 hours -* **Cron expressions**: "Remind me at 9am daily" → uses cron expression +- **One-time reminders**: "Remind me in 10 minutes" → triggers once after 10min +- **Recurring tasks**: "Remind me every 2 hours" → triggers every 2 hours +- **Cron expressions**: "Remind me at 9am daily" → uses cron expression Jobs are stored in `~/.picoclaw/workspace/cron/` and processed automatically. @@ -1343,6 +1352,7 @@ This is normal if you haven't configured a search API key yet. PicoClaw will pro #### Search Provider Priority PicoClaw automatically selects the best available search provider in this order: + 1. **Perplexity** (if enabled and API key configured) - AI-powered search with citations 2. **Brave Search** (if enabled and API key configured) - Privacy-focused paid API ($5/1000 queries) 3. **SearXNG** (if enabled and base_url configured) - Self-hosted metasearch aggregating 70+ engines (free) @@ -1351,6 +1361,7 @@ PicoClaw automatically selects the best available search provider in this order: #### Web Search Configuration Options **Option 1 (Best Results)**: Perplexity AI Search + ```json { "tools": { @@ -1366,6 +1377,7 @@ PicoClaw automatically selects the best available search provider in this order: ``` **Option 2 (Paid API)**: Get an API key at [https://brave.com/search/api](https://brave.com/search/api) ($5/1000 queries, ~$5-6/month) + ```json { "tools": { @@ -1381,6 +1393,7 @@ PicoClaw automatically selects the best available search provider in this order: ``` **Option 3 (Self-Hosted)**: Deploy your own [SearXNG](https://github.com/searxng/searxng) instance + ```json { "tools": { @@ -1396,6 +1409,7 @@ PicoClaw automatically selects the best available search provider in this order: ``` Benefits of SearXNG: + - **Zero cost**: No API fees or rate limits - **Privacy-focused**: Self-hosted, no tracking - **Aggregate results**: Queries 70+ search engines simultaneously @@ -1446,11 +1460,11 @@ This happens when another instance of the bot is running. Make sure only one `pi ## 📝 API Key Comparison -| Service | Free Tier | Use Case | -| ---------------- | ------------------------ | ------------------------------------- | -| **OpenRouter** | 200K tokens/month | Multiple models (Claude, GPT-4, etc.) | -| **Zhipu** | 200K tokens/month | Best for Chinese users | -| **Brave Search** | Paid ($5/1000 queries) | Web search functionality | -| **SearXNG** | Unlimited (self-hosted) | Privacy-focused metasearch (70+ engines) | -| **Groq** | Free tier available | Fast inference (Llama, Mixtral) | -| **Cerebras** | Free tier available | Fast inference (Llama, Qwen, etc.) | +| Service | Free Tier | Use Case | +| ---------------- | ----------------------- | ---------------------------------------- | +| **OpenRouter** | 200K tokens/month | Multiple models (Claude, GPT-4, etc.) | +| **Zhipu** | 200K tokens/month | Best for Chinese users | +| **Brave Search** | Paid ($5/1000 queries) | Web search functionality | +| **SearXNG** | Unlimited (self-hosted) | Privacy-focused metasearch (70+ engines) | +| **Groq** | Free tier available | Fast inference (Llama, Mixtral) | +| **Cerebras** | Free tier available | Fast inference (Llama, Qwen, etc.) | diff --git a/README.pt-br.md b/README.pt-br.md index 5f37ba457..392fcb70d 100644 --- a/README.pt-br.md +++ b/README.pt-br.md @@ -881,7 +881,11 @@ picoclaw agent -m "Ola, como vai?" { "agents": { "defaults": { - "model": "anthropic/claude-opus-4-5" + "model": "anthropic/claude-opus-4-5", + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "providers": { diff --git a/README.vi.md b/README.vi.md index 92c6ecbae..f44121661 100644 --- a/README.vi.md +++ b/README.vi.md @@ -853,7 +853,11 @@ picoclaw agent -m "Xin chào" { "agents": { "defaults": { - "model": "anthropic/claude-opus-4-5" + "model": "anthropic/claude-opus-4-5", + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "providers": { diff --git a/README.zh.md b/README.zh.md index d42b3cbb8..980db1057 100644 --- a/README.zh.md +++ b/README.zh.md @@ -229,7 +229,11 @@ picoclaw onboard "model_name": "gpt4", "max_tokens": 8192, "temperature": 0.7, - "max_tool_iterations": 20 + "max_tool_iterations": 20, + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "model_list": [ @@ -725,7 +729,11 @@ picoclaw agent -m "你好" { "agents": { "defaults": { - "model": "anthropic/claude-opus-4-5" + "model": "anthropic/claude-opus-4-5", + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "session": { diff --git a/config/config.example.json b/config/config.example.json index 2f643d41b..49a0b14eb 100644 --- a/config/config.example.json +++ b/config/config.example.json @@ -7,8 +7,10 @@ "max_tokens": 8192, "temperature": 0.7, "max_tool_iterations": 20, - "summarize_message_threshold": 20, - "summarize_token_percent": 75 + "summarization": { + "message_threshold": 20, + "token_percent": 75 + } } }, "model_list": [ diff --git a/docs/design/DDR-session-summarization.md b/docs/design/DDR-session-summarization.md new file mode 100644 index 000000000..4536885f2 --- /dev/null +++ b/docs/design/DDR-session-summarization.md @@ -0,0 +1,220 @@ +# DDR: Session Summarization Extraction & Configuration + +> Date: 2026-03-06 +> Status: In Review +> Packages: `pkg/session`, `pkg/config`, `pkg/agent` + +## Problem Frame + +Session context grows without bound during long conversations. The original implementation had several issues: + +- **Entangled logic**: Summarization was inlined in `pkg/agent/loop.go` — 4 functions (`maybeSummarize`, `forceCompression`, `summarizeSession`, `summarizeBatch`) mixed session I/O, LLM calls, and threshold logic in the agent loop. +- **Race condition**: Async summarization could silently discard messages appended by the user while the LLM was generating a summary. +- **Magic constants**: Token budget (`1024`), temperature (`0.3`), threshold (`20`), percent (`75`) were hardcoded with no user-facing knobs. +- **No injection point**: Tight coupling to `providers.LLMProvider` made unit testing impossible without a live LLM. +- **Flat config**: Two config fields (`summarize_message_threshold`, `summarize_token_percent`) lived as top-level scalars in `AgentDefaults`, not grouped or extensible. + +## Decision + +### Rule Changes + +1. **MUST** extract all summarization orchestration into `pkg/session.SessionManager`. + - `MaybeSummarize(key)` — threshold check + async dispatch. + - `ForceCompression(key)` — emergency drop of older half. + - `ApplySummarization(key, summary, snapshotLen, keepLast)` — atomic merge that compares `snapshotLen` to current message count, preserving messages appended during async work. + - `summarizeSession(ctx, key)` — unexported, handles single/multi-part batching. + +2. **MUST** define a `Summarizer` interface in `pkg/session`: + + ```go + type Summarizer interface { + Summarize(ctx context.Context, messages []providers.Message, existingSummary string) (string, error) + } + ``` + + A `SummarizeFunc` adapter MUST be provided for testing convenience. + +3. **MUST** inject the `Summarizer` at construction via functional options: + - `WithSummarizer(s Summarizer, cfg config.SummarizationConfig) Option` — generic. + - `WithLLMSummarizer(provider, model, agentID string, cfg config.SummarizationConfig) Option` — convenience, constructs `LLMSummarizer` internally. + - When `summarizer` is nil, `MaybeSummarize` is a no-op. + +4. **MUST** replace all magic constants with named defaults in `pkg/config`: + | Constant | Value | Field | + |---|---|---| + | `DefaultSummarizeMessageThreshold` | `20` | `MessageThreshold` | + | `DefaultSummarizeTokenPercent` | `75` | `TokenPercent` | + | `DefaultKeepLastMessages` | `4` | `KeepLastMessages` | + | `DefaultContextWindow` | `8192` | `ContextWindow` | + | `DefaultSummaryMaxTokens` | `1024` | `SummaryMaxTokens` | + | `DefaultSummaryTemperature` | `0.3` | `SummaryTemperature` | + | `DefaultMultiPartBatchThreshold` | `10` | `MultiPartBatchThreshold` | + | `DefaultTimeout` | `120s` | `Timeout` | + | `DefaultMaxSingleMsgTokenRatio` | `0.5` | `MaxSingleMsgTokenRatio` | + | `DefaultForceCompressionMinMsgs` | `4` | `ForceCompressionMinMessages` | + | `DefaultCharsPerToken` | `2.5` | `CharsPerToken` | + +5. **MUST** use a single `SummarizationConfig` struct in `pkg/config` as the sole config type for both user-facing JSON fields and internal tuning parameters: + + ```go + type SummarizationConfig struct { + // User-facing (documented in config.example.json) + MessageThreshold int `json:"message_threshold,omitempty"` + TokenPercent int `json:"token_percent,omitempty"` + KeepLastMessages int `json:"keep_last_messages,omitempty"` + ContextWindow int `json:"context_window,omitempty"` + SummaryMaxTokens int `json:"summary_max_tokens,omitempty"` + SummaryTemperature float64 `json:"summary_temperature,omitempty"` + + // Internal (undocumented, tunable for advanced use) + MultiPartBatchThreshold int `json:"multi_part_batch_threshold,omitempty"` + Timeout time.Duration `json:"timeout,omitempty"` + MaxSingleMsgTokenRatio float64 `json:"max_single_msg_token_ratio,omitempty"` + ForceCompressionMinMessages int `json:"force_compression_min_messages,omitempty"` + CharsPerToken float64 `json:"chars_per_token,omitempty"` + } + ``` + + Nested under `agents.defaults.summarization` in JSON config. The session package imports and uses this directly — no separate config struct. + +6. **MUST** preserve backward compatibility: `AgentDefaults.GetSummarization()` merges the new `Summarization *SummarizationConfig` with legacy flat fields `SummarizeMessageThreshold` / `SummarizeTokenPercent`. New struct takes priority; legacy fields apply only when the corresponding new field is zero. + +7. **SHOULD** deprecate `summarize_message_threshold` and `summarize_token_percent` in `AgentDefaults`. These fields MUST remain parseable with `omitempty` but SHOULD NOT appear in new documentation or `config.example.json`. + +8. **MAY** omit all `SummarizationConfig` fields — zero values are replaced by sensible defaults via `SummarizationConfig.WithDefaults()`. + +### Definitions + +- **Snapshot length**: The message count captured before dispatching async summarization. Used by `ApplySummarization` to detect concurrent writes. +- **Multi-part batching**: When the message batch exceeds `MultiPartBatchThreshold`, it is split in half; each half is summarized in parallel via goroutines; results are merged into a single summary. +- **Force compression**: Emergency fallback that drops the older half of messages and prepends a `[system]` note. No LLM call required. + +### Migration + +Old config (deprecated): + +```json +{ + "agents": { + "defaults": { + "summarize_message_threshold": 30, + "summarize_token_percent": 80 + } + } +} +``` + +New config (preferred): + +```json +{ + "agents": { + "defaults": { + "summarization": { + "message_threshold": 30, + "token_percent": 80 + } + } + } +} +``` + +Both forms are supported. If both are present, the nested struct wins. + +## Rationale + +### Core Reasons + +- **Separation of concerns**: Agent loop should orchestrate turn logic, not session housekeeping. Summarization is a session-level concern. +- **Testability**: Interface injection enables deterministic unit tests with `SummarizeFunc` stubs — no LLM, no network, no flakes. +- **Race safety**: `ApplySummarization` compares snapshot length to current state, so messages arriving during async summarization are never lost. +- **Configurability**: Users on constrained hardware (RISC-V, 10MB RAM) need different thresholds than users on GPT-4 with 128k context. +- **Extensibility**: New summarization strategies (e.g., embeddings-based, local model) plug in via the `Summarizer` interface. + +5. **Unified config** — Single `config.SummarizationConfig` used everywhere. No duplicate struct in session package, no field-by-field mapping. + +- **Config hygiene**: Grouping related fields under `summarization` prevents `AgentDefaults` from growing into a flat bag of 30+ fields. + +### Alternatives Considered + +1. **Do nothing** — Keep inlined in `loop.go`. + - Rejected: Untestable without LLM. Race condition on concurrent writes. Magic constants not user-configurable. + +2. **Separate `SessionSummarizer` struct** — Standalone struct with its own state. + - Rejected: Duplicates session access logic. Requires either passing sessions in/out or giving the struct direct storage access — both leak responsibilities. + +3. **Mutable closure via `ConfigureSummarizer(fn)`** — Post-construction injection. + - Rejected: Mutable after init violates immutability principle. Anonymous closures are harder to test and document. + +4. **Config as flat fields** — Add `keep_last_messages`, `summary_max_tokens`, etc. directly to `AgentDefaults`. + - Rejected: Would add 6 more top-level fields. Not cohesive. Harder to feature-flag or version. + +5. **Separate config struct per package** — `session.SummarizerConfig` with its own defaults, mapped field-by-field from `config.SummarizationConfig` in `instance.go`. + - Rejected: Duplicate types for the same data. Tedious mapping code. Easy to drift. + +### Trade-offs + +| Dimension | Cost | Benefit | +| ----------- | ------------------------------------- | ------------------------------------------------ | +| Complexity | New interface + options pattern | Clean test seam, swappable strategies | +| Migration | Must support two config shapes | Zero breaking changes for existing users | +| Performance | One extra goroutine per summarization | Non-blocking — user doesn't wait for LLM summary | +| Binary size | ~0 (no new deps) | N/A | + +### Risks & Mitigations + +| Risk | Mitigation | +| ------------------------------------------------------ | ----------------------------------------------------------------------------------------------------- | +| Snapshot stale (messages deleted during summarization) | `ApplySummarization` returns early if `snapshotLen > currentLen` | +| Duplicate in-flight summarizations | `sync.Map` inflight guard — second call is a no-op | +| LLM summarization fails | Logged and skipped; session continues unsummarized. `ForceCompression` as fallback at token overflow. | +| Legacy config silently ignored | `GetSummarization()` merges both; tested in `config_test.go` | + +## Consequences + +### Immediate Impacts + +| What | Change | +| ---------------------------- | ------------------------------------------------------------------------------------------------------------------ | +| `pkg/agent/loop.go` | Calls `agent.Sessions.MaybeSummarize(key)` / `ForceCompression(key)` — no direct summarization logic | +| `pkg/agent/instance.go` | Wires `WithLLMSummarizer(provider, model, agentID, cfg)` at construction | +| `pkg/session/manager.go` | Gains `summarizer`, `summarizerCfg` (`config.SummarizationConfig`), `inflight` fields + 4 new methods | +| `pkg/session/summarizer.go` | New file: interface, `SummarizeFunc` adapter, `LLMSummarizer`, prompt builder (no config types or constants) | +| `pkg/config/config.go` | `SummarizationConfig` struct (all fields + `WithDefaults()`), `GetSummarization()` method, named default constants | +| `config/config.example.json` | Uses nested `summarization` block | +| README config examples | Show `summarization` in quickstart snippets | + +### Follow-up Tasks + +| Task | Owner (role) | Priority | +| ---------------------------------------------------------------------------------- | ------------ | ------------------------ | +| Remove deprecated `summarize_message_threshold` / `summarize_token_percent` fields | Maintainer | Low — next major version | +| Add env var override for `summarization.*` fields | Maintainer | Medium | +| Evaluate embeddings-based summarization via `Summarizer` interface | Contributor | Low | + +### Test / Verification Plan + +| Acceptance Criterion | Test | +| -------------------------------------------- | ---------------------------------------------------- | +| Summarization triggers above threshold | `TestMaybeSummarize_AboveThreshold` | +| No-op below threshold | `TestMaybeSummarize_BelowThreshold` | +| Dedup prevents concurrent summarizations | `TestMaybeSummarize_Dedup` | +| Messages appended during async are preserved | `TestApplySummarization_PreservesNewMessages` | +| Stale snapshot rejected | `TestApplySummarization_StaleSnapshot` | +| Force compression drops older half | `TestForceCompression_DropsHalf` | +| Multi-part batching for large sessions | `TestSummarizeSession_MultiPart` | +| Concurrent read/write safety | `TestConcurrent_SummarizeAndWrite` | +| Token estimation (Latin + CJK) | `TestEstimateTokens_Latin`, `TestEstimateTokens_CJK` | +| Legacy config fallback works | `TestGetSummarization_LegacyFallback` | +| New config overrides legacy | `TestGetSummarization_NewOverridesLegacy` | +| Default config has correct thresholds | `TestDefaultConfig_SummarizationThresholds` | + +All tests pass as of 2026-03-06 (1455 passed, 0 failed). + +### Rollback Plan + +Revert the commits that introduce `pkg/session/summarizer.go`, the `Summarizer` interface, and the config changes. The legacy flat fields remain functional — no data migration needed. + +### Review Date + +2026-06-06 (3 months) — evaluate whether legacy flat config fields can be removed. diff --git a/pkg/agent/instance.go b/pkg/agent/instance.go index 97cf0fa05..ab97a1b9f 100644 --- a/pkg/agent/instance.go +++ b/pkg/agent/instance.go @@ -18,25 +18,23 @@ import ( // AgentInstance represents a fully configured agent with its own workspace, // session manager, context builder, and tool registry. type AgentInstance struct { - ID string - Name string - Model string - Fallbacks []string - Workspace string - MaxIterations int - MaxTokens int - Temperature float64 - ThinkingLevel ThinkingLevel - ContextWindow int - SummarizeMessageThreshold int - SummarizeTokenPercent int - Provider providers.LLMProvider - Sessions *session.SessionManager - ContextBuilder *ContextBuilder - Tools *tools.ToolRegistry - Subagents *config.SubagentsConfig - SkillsFilter []string - Candidates []providers.FallbackCandidate + ID string + Name string + Model string + Fallbacks []string + Workspace string + MaxIterations int + MaxTokens int + Temperature float64 + ThinkingLevel ThinkingLevel + ContextWindow int + Provider providers.LLMProvider + Sessions *session.SessionManager + ContextBuilder *ContextBuilder + Tools *tools.ToolRegistry + Subagents *config.SubagentsConfig + SkillsFilter []string + Candidates []providers.FallbackCandidate // Router is non-nil when model routing is configured and the light model // was successfully resolved. It scores each incoming message and decides @@ -94,7 +92,6 @@ func NewAgentInstance( } sessionsDir := filepath.Join(workspace, "sessions") - sessionsManager := session.NewSessionManager(sessionsDir) contextBuilder := NewContextBuilder(workspace) @@ -131,16 +128,6 @@ func NewAgentInstance( } thinkingLevel := parseThinkingLevel(thinkingLevelStr) - summarizeMessageThreshold := defaults.SummarizeMessageThreshold - if summarizeMessageThreshold == 0 { - summarizeMessageThreshold = 20 - } - - summarizeTokenPercent := defaults.SummarizeTokenPercent - if summarizeTokenPercent == 0 { - summarizeTokenPercent = 75 - } - // Resolve fallback candidates modelCfg := providers.ModelConfig{ Primary: model, @@ -207,28 +194,34 @@ func NewAgentInstance( } } + sc := defaults.GetSummarization() + if sc.ContextWindow == 0 { + sc.ContextWindow = maxTokens + } + sessionsManager := session.NewSessionManager(sessionsDir, session.WithLLMSummarizer( + provider, model, agentID, sc, + )) + return &AgentInstance{ - ID: agentID, - Name: agentName, - Model: model, - Fallbacks: fallbacks, - Workspace: workspace, - MaxIterations: maxIter, - MaxTokens: maxTokens, - Temperature: temperature, - ThinkingLevel: thinkingLevel, - ContextWindow: maxTokens, - SummarizeMessageThreshold: summarizeMessageThreshold, - SummarizeTokenPercent: summarizeTokenPercent, - Provider: provider, - Sessions: sessionsManager, - ContextBuilder: contextBuilder, - Tools: toolsRegistry, - Subagents: subagents, - SkillsFilter: skillsFilter, - Candidates: candidates, - Router: router, - LightCandidates: lightCandidates, + ID: agentID, + Name: agentName, + Model: model, + Fallbacks: fallbacks, + Workspace: workspace, + MaxIterations: maxIter, + MaxTokens: maxTokens, + Temperature: temperature, + ThinkingLevel: thinkingLevel, + ContextWindow: maxTokens, + Provider: provider, + Sessions: sessionsManager, + ContextBuilder: contextBuilder, + Tools: toolsRegistry, + Subagents: subagents, + SkillsFilter: skillsFilter, + Candidates: candidates, + Router: router, + LightCandidates: lightCandidates, } } diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index 19d13b2bb..846b22eab 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -17,7 +17,6 @@ import ( "sync" "sync/atomic" "time" - "unicode/utf8" "github.com/sipeed/picoclaw/pkg/bus" "github.com/sipeed/picoclaw/pkg/channels" @@ -42,7 +41,6 @@ type AgentLoop struct { registry *AgentRegistry state *state.Manager running atomic.Bool - summarizing sync.Map fallback *providers.FallbackChain channelManager *channels.Manager mediaStore media.MediaStore @@ -99,7 +97,6 @@ func NewAgentLoop( cfg: cfg, registry: registry, state: stateManager, - summarizing: sync.Map{}, fallback: fallbackChain, cmdRegistry: commands.NewRegistry(commands.BuiltinDefinitions()), } @@ -784,7 +781,7 @@ func (al *AgentLoop) runAgentLoop( // 6. Optional: summarization if opts.EnableSummary { - al.maybeSummarize(agent, opts.SessionKey, opts.Channel, opts.ChatID) + agent.Sessions.MaybeSummarize(opts.SessionKey) } // 7. Optional: send response via bus @@ -1017,7 +1014,7 @@ func (al *AgentLoop) runLLMIteration( }) } - al.forceCompression(agent, opts.SessionKey) + agent.Sessions.ForceCompression(opts.SessionKey) newHistory := agent.Sessions.GetHistory(opts.SessionKey) newSummary := agent.Sessions.GetSummary(opts.SessionKey) messages = agent.ContextBuilder.BuildMessages( @@ -1264,77 +1261,6 @@ func (al *AgentLoop) selectCandidates( return agent.LightCandidates, agent.Router.LightModel() } -// maybeSummarize triggers summarization if the session history exceeds thresholds. -func (al *AgentLoop) maybeSummarize(agent *AgentInstance, sessionKey, channel, chatID string) { - newHistory := agent.Sessions.GetHistory(sessionKey) - tokenEstimate := al.estimateTokens(newHistory) - threshold := agent.ContextWindow * agent.SummarizeTokenPercent / 100 - - if len(newHistory) > agent.SummarizeMessageThreshold || tokenEstimate > threshold { - summarizeKey := agent.ID + ":" + sessionKey - if _, loading := al.summarizing.LoadOrStore(summarizeKey, true); !loading { - go func() { - defer al.summarizing.Delete(summarizeKey) - logger.Debug("Memory threshold reached. Optimizing conversation history...") - al.summarizeSession(agent, sessionKey) - }() - } - } -} - -// forceCompression aggressively reduces context when the limit is hit. -// It drops the oldest 50% of messages (keeping system prompt and last user message). -func (al *AgentLoop) forceCompression(agent *AgentInstance, sessionKey string) { - history := agent.Sessions.GetHistory(sessionKey) - if len(history) <= 4 { - return - } - - // Keep system prompt (usually [0]) and the very last message (user's trigger) - // We want to drop the oldest half of the *conversation* - // Assuming [0] is system, [1:] is conversation - conversation := history[1 : len(history)-1] - if len(conversation) == 0 { - return - } - - // Helper to find the mid-point of the conversation - mid := len(conversation) / 2 - - // New history structure: - // 1. System Prompt (with compression note appended) - // 2. Second half of conversation - // 3. Last message - - droppedCount := mid - keptConversation := conversation[mid:] - - newHistory := make([]providers.Message, 0, 1+len(keptConversation)+1) - - // Append compression note to the original system prompt instead of adding a new system message - // This avoids having two consecutive system messages which some APIs (like Zhipu) reject - compressionNote := fmt.Sprintf( - "\n\n[System Note: Emergency compression dropped %d oldest messages due to context limit]", - droppedCount, - ) - enhancedSystemPrompt := history[0] - enhancedSystemPrompt.Content = enhancedSystemPrompt.Content + compressionNote - newHistory = append(newHistory, enhancedSystemPrompt) - - newHistory = append(newHistory, keptConversation...) - newHistory = append(newHistory, history[len(history)-1]) // Last message - - // Update session - agent.Sessions.SetHistory(sessionKey, newHistory) - agent.Sessions.Save(sessionKey) - - logger.WarnCF("agent", "Forced compression executed", map[string]any{ - "session_key": sessionKey, - "dropped_msgs": droppedCount, - "new_count": len(newHistory), - }) -} - // GetStartupInfo returns information about loaded tools and skills for logging. func (al *AgentLoop) GetStartupInfo() map[string]any { info := make(map[string]any) @@ -1422,139 +1348,6 @@ func formatToolsForLog(toolDefs []providers.ToolDefinition) string { return sb.String() } -// summarizeSession summarizes the conversation history for a session. -func (al *AgentLoop) summarizeSession(agent *AgentInstance, sessionKey string) { - ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) - defer cancel() - - history := agent.Sessions.GetHistory(sessionKey) - summary := agent.Sessions.GetSummary(sessionKey) - - // Keep last 4 messages for continuity - if len(history) <= 4 { - return - } - - toSummarize := history[:len(history)-4] - - // Oversized Message Guard - maxMessageTokens := agent.ContextWindow / 2 - validMessages := make([]providers.Message, 0) - omitted := false - - for _, m := range toSummarize { - if m.Role != "user" && m.Role != "assistant" { - continue - } - msgTokens := len(m.Content) / 2 - if msgTokens > maxMessageTokens { - omitted = true - continue - } - validMessages = append(validMessages, m) - } - - if len(validMessages) == 0 { - return - } - - // Multi-Part Summarization - var finalSummary string - if len(validMessages) > 10 { - mid := len(validMessages) / 2 - part1 := validMessages[:mid] - part2 := validMessages[mid:] - - s1, _ := al.summarizeBatch(ctx, agent, part1, "") - s2, _ := al.summarizeBatch(ctx, agent, part2, "") - - mergePrompt := fmt.Sprintf( - "Merge these two conversation summaries into one cohesive summary:\n\n1: %s\n\n2: %s", - s1, - s2, - ) - resp, err := agent.Provider.Chat( - ctx, - []providers.Message{{Role: "user", Content: mergePrompt}}, - nil, - agent.Model, - map[string]any{ - "max_tokens": 1024, - "temperature": 0.3, - "prompt_cache_key": agent.ID, - }, - ) - if err == nil { - finalSummary = resp.Content - } else { - finalSummary = s1 + " " + s2 - } - } else { - finalSummary, _ = al.summarizeBatch(ctx, agent, validMessages, summary) - } - - if omitted && finalSummary != "" { - finalSummary += "\n[Note: Some oversized messages were omitted from this summary for efficiency.]" - } - - if finalSummary != "" { - agent.Sessions.SetSummary(sessionKey, finalSummary) - agent.Sessions.TruncateHistory(sessionKey, 4) - agent.Sessions.Save(sessionKey) - } -} - -// summarizeBatch summarizes a batch of messages. -func (al *AgentLoop) summarizeBatch( - ctx context.Context, - agent *AgentInstance, - batch []providers.Message, - existingSummary string, -) (string, error) { - var sb strings.Builder - sb.WriteString( - "Provide a concise summary of this conversation segment, preserving core context and key points.\n", - ) - if existingSummary != "" { - sb.WriteString("Existing context: ") - sb.WriteString(existingSummary) - sb.WriteString("\n") - } - sb.WriteString("\nCONVERSATION:\n") - for _, m := range batch { - fmt.Fprintf(&sb, "%s: %s\n", m.Role, m.Content) - } - prompt := sb.String() - - response, err := agent.Provider.Chat( - ctx, - []providers.Message{{Role: "user", Content: prompt}}, - nil, - agent.Model, - map[string]any{ - "max_tokens": 1024, - "temperature": 0.3, - "prompt_cache_key": agent.ID, - }, - ) - if err != nil { - return "", err - } - return response.Content, nil -} - -// estimateTokens estimates the number of tokens in a message list. -// Uses a safe heuristic of 2.5 characters per token to account for CJK and other -// overheads better than the previous 3 chars/token. -func (al *AgentLoop) estimateTokens(messages []providers.Message) int { - totalChars := 0 - for _, m := range messages { - totalChars += utf8.RuneCountInString(m.Content) - } - // 2.5 chars per token = totalChars * 2 / 5 - return totalChars * 2 / 5 -} - func (al *AgentLoop) handleCommand( ctx context.Context, msg bus.InboundMessage, diff --git a/pkg/config/config.go b/pkg/config/config.go index 72af3e2fb..ca38741cd 100644 --- a/pkg/config/config.go +++ b/pkg/config/config.go @@ -5,6 +5,7 @@ import ( "fmt" "os" "sync/atomic" + "time" "github.com/caarlos0/env/v11" @@ -180,22 +181,135 @@ type RoutingConfig struct { } type AgentDefaults struct { - Workspace string `json:"workspace" env:"PICOCLAW_AGENTS_DEFAULTS_WORKSPACE"` - RestrictToWorkspace bool `json:"restrict_to_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE"` - AllowReadOutsideWorkspace bool `json:"allow_read_outside_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_ALLOW_READ_OUTSIDE_WORKSPACE"` - Provider string `json:"provider" env:"PICOCLAW_AGENTS_DEFAULTS_PROVIDER"` - ModelName string `json:"model_name,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL_NAME"` - Model string `json:"model" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` // Deprecated: use model_name instead - ModelFallbacks []string `json:"model_fallbacks,omitempty"` - ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"` - ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"` - MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` - Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` - 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"` - Routing *RoutingConfig `json:"routing,omitempty"` + Workspace string `json:"workspace" env:"PICOCLAW_AGENTS_DEFAULTS_WORKSPACE"` + RestrictToWorkspace bool `json:"restrict_to_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_RESTRICT_TO_WORKSPACE"` + AllowReadOutsideWorkspace bool `json:"allow_read_outside_workspace" env:"PICOCLAW_AGENTS_DEFAULTS_ALLOW_READ_OUTSIDE_WORKSPACE"` + Provider string `json:"provider" env:"PICOCLAW_AGENTS_DEFAULTS_PROVIDER"` + ModelName string `json:"model_name,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL_NAME"` + Model string `json:"model" env:"PICOCLAW_AGENTS_DEFAULTS_MODEL"` // Deprecated: use model_name instead + ModelFallbacks []string `json:"model_fallbacks,omitempty"` + ImageModel string `json:"image_model,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_IMAGE_MODEL"` + ImageModelFallbacks []string `json:"image_model_fallbacks,omitempty"` + MaxTokens int `json:"max_tokens" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOKENS"` + Temperature *float64 `json:"temperature,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_TEMPERATURE"` + MaxToolIterations int `json:"max_tool_iterations" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_TOOL_ITERATIONS"` + Summarization *SummarizationConfig `json:"summarization,omitempty"` + MaxMediaSize int `json:"max_media_size,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_MAX_MEDIA_SIZE"` + Routing *RoutingConfig `json:"routing,omitempty"` + + // Deprecated: use summarization.message_threshold instead. + SummarizeMessageThreshold int `json:"summarize_message_threshold,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_MESSAGE_THRESHOLD"` + // Deprecated: use summarization.token_percent instead. + SummarizeTokenPercent int `json:"summarize_token_percent,omitempty" env:"PICOCLAW_AGENTS_DEFAULTS_SUMMARIZE_TOKEN_PERCENT"` +} + +// SummarizationConfig controls session summarization behavior. +// All fields are optional — zero values are replaced by sensible defaults +// via WithDefaults(). +type SummarizationConfig struct { + // --- User-facing (documented in config.example.json) --- + + // MessageThreshold is the message count above which summarization is triggered. + MessageThreshold int `json:"message_threshold,omitempty"` + // TokenPercent is the percentage of context window above which summarization triggers. + TokenPercent int `json:"token_percent,omitempty"` + // KeepLastMessages is the number of recent messages to retain after summarization. + KeepLastMessages int `json:"keep_last_messages,omitempty"` + // ContextWindow overrides the token budget for threshold calculations (defaults to max_tokens). + ContextWindow int `json:"context_window,omitempty"` + // SummaryMaxTokens is the max_tokens limit for the LLM summarization call. + SummaryMaxTokens int `json:"summary_max_tokens,omitempty"` + // SummaryTemperature is the temperature for the LLM summarization call. + SummaryTemperature float64 `json:"summary_temperature,omitempty"` + + // --- Internal (undocumented, tunable for advanced use) --- + + // MultiPartBatchThreshold controls when summarization splits the + // message batch into two halves for separate summarization + merge. + MultiPartBatchThreshold int `json:"multi_part_batch_threshold,omitempty"` + // Timeout is the maximum duration for a single summarization operation. + Timeout time.Duration `json:"timeout,omitempty"` + // MaxSingleMsgTokenRatio is the fraction of ContextWindow above which + // a single message is considered oversized and skipped during summarization. + MaxSingleMsgTokenRatio float64 `json:"max_single_msg_token_ratio,omitempty"` + // ForceCompressionMinMessages is the minimum message count below which + // force-compression is a no-op (not enough to drop). + ForceCompressionMinMessages int `json:"force_compression_min_messages,omitempty"` + // CharsPerToken is the estimated character-to-token ratio used by + // EstimateTokens. Lower values are more conservative (estimate more tokens). + CharsPerToken float64 `json:"chars_per_token,omitempty"` +} + +// Defaults for SummarizationConfig fields. +const ( + DefaultSummarizeMessageThreshold = 20 + DefaultSummarizeTokenPercent = 75 + DefaultKeepLastMessages = 4 + DefaultContextWindow = 8192 + DefaultSummaryMaxTokens = 1024 + DefaultSummarizationTemperature float64 = 0.3 + DefaultMultiPartBatchThreshold = 10 + DefaultSummarizationTimeout = 120 * time.Second + DefaultMaxSingleMsgTokenRatio = 0.5 + DefaultForceCompressionMinMsgs = 4 + DefaultCharsPerToken = 2.5 +) + +// WithDefaults fills zero-valued fields with sensible defaults. +func (c SummarizationConfig) WithDefaults() SummarizationConfig { + if c.MessageThreshold == 0 { + c.MessageThreshold = DefaultSummarizeMessageThreshold + } + if c.TokenPercent == 0 { + c.TokenPercent = DefaultSummarizeTokenPercent + } + if c.KeepLastMessages == 0 { + c.KeepLastMessages = DefaultKeepLastMessages + } + if c.ContextWindow == 0 { + c.ContextWindow = DefaultContextWindow + } + if c.SummaryMaxTokens == 0 { + c.SummaryMaxTokens = DefaultSummaryMaxTokens + } + if c.SummaryTemperature == 0 { + c.SummaryTemperature = DefaultSummarizationTemperature + } + if c.MultiPartBatchThreshold == 0 { + c.MultiPartBatchThreshold = DefaultMultiPartBatchThreshold + } + if c.Timeout == 0 { + c.Timeout = DefaultSummarizationTimeout + } + if c.MaxSingleMsgTokenRatio == 0 { + c.MaxSingleMsgTokenRatio = DefaultMaxSingleMsgTokenRatio + } + if c.ForceCompressionMinMessages == 0 { + c.ForceCompressionMinMessages = DefaultForceCompressionMinMsgs + } + if c.CharsPerToken == 0 { + c.CharsPerToken = DefaultCharsPerToken + } + return c +} + +// GetSummarization returns the effective summarization config, merging the +// new nested field with the deprecated flat fields for backward compatibility. +// The new Summarization struct takes priority; legacy fields are used only +// when the corresponding new field is unset. +func (d *AgentDefaults) GetSummarization() SummarizationConfig { + var sc SummarizationConfig + if d.Summarization != nil { + sc = *d.Summarization + } + // Legacy fallback: only apply if the new field is zero. + if sc.MessageThreshold == 0 && d.SummarizeMessageThreshold != 0 { + sc.MessageThreshold = d.SummarizeMessageThreshold + } + if sc.TokenPercent == 0 && d.SummarizeTokenPercent != 0 { + sc.TokenPercent = d.SummarizeTokenPercent + } + return sc } const DefaultMaxMediaSize = 20 * 1024 * 1024 // 20 MB diff --git a/pkg/config/config_test.go b/pkg/config/config_test.go index 10ebc7c90..53f1d59f2 100644 --- a/pkg/config/config_test.go +++ b/pkg/config/config_test.go @@ -434,16 +434,48 @@ func TestLoadConfig_WebToolsProxy(t *testing.T) { } } -// 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) + sc := cfg.Agents.Defaults.GetSummarization() + if sc.MessageThreshold != 20 { + t.Errorf("MessageThreshold = %d, want 20", sc.MessageThreshold) } - if cfg.Agents.Defaults.SummarizeTokenPercent != 75 { - t.Errorf("SummarizeTokenPercent = %d, want 75", cfg.Agents.Defaults.SummarizeTokenPercent) + if sc.TokenPercent != 75 { + t.Errorf("TokenPercent = %d, want 75", sc.TokenPercent) + } +} + +func TestGetSummarization_LegacyFallback(t *testing.T) { + d := AgentDefaults{ + SummarizeMessageThreshold: 30, + SummarizeTokenPercent: 80, + } + sc := d.GetSummarization() + if sc.MessageThreshold != 30 { + t.Errorf("legacy MessageThreshold = %d, want 30", sc.MessageThreshold) + } + if sc.TokenPercent != 80 { + t.Errorf("legacy TokenPercent = %d, want 80", sc.TokenPercent) + } +} + +func TestGetSummarization_NewOverridesLegacy(t *testing.T) { + d := AgentDefaults{ + SummarizeMessageThreshold: 30, // legacy + SummarizeTokenPercent: 80, // legacy + Summarization: &SummarizationConfig{ + MessageThreshold: 50, + TokenPercent: 90, + }, + } + sc := d.GetSummarization() + if sc.MessageThreshold != 50 { + t.Errorf("new MessageThreshold = %d, want 50", sc.MessageThreshold) + } + if sc.TokenPercent != 90 { + t.Errorf("new TokenPercent = %d, want 90", sc.TokenPercent) } } diff --git a/pkg/config/defaults.go b/pkg/config/defaults.go index 1902480c5..ccd787a55 100644 --- a/pkg/config/defaults.go +++ b/pkg/config/defaults.go @@ -26,15 +26,17 @@ func DefaultConfig() *Config { return &Config{ Agents: AgentsConfig{ Defaults: AgentDefaults{ - Workspace: workspacePath, - RestrictToWorkspace: true, - Provider: "", - Model: "", - MaxTokens: 32768, - Temperature: nil, // nil means use provider default - MaxToolIterations: 50, - SummarizeMessageThreshold: 20, - SummarizeTokenPercent: 75, + Workspace: workspacePath, + RestrictToWorkspace: true, + Provider: "", + Model: "", + MaxTokens: 32768, + Temperature: nil, // nil means use provider default + MaxToolIterations: 50, + Summarization: &SummarizationConfig{ + MessageThreshold: 20, + TokenPercent: 75, + }, }, }, Bindings: []AgentBinding{}, diff --git a/pkg/session/manager.go b/pkg/session/manager.go index 08f0b0ad2..7cd00f9ff 100644 --- a/pkg/session/manager.go +++ b/pkg/session/manager.go @@ -1,13 +1,18 @@ package session import ( + "context" "encoding/json" + "fmt" "os" "path/filepath" "strings" "sync" "time" + "unicode/utf8" + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/logger" "github.com/sipeed/picoclaw/pkg/providers" ) @@ -20,17 +25,37 @@ type Session struct { } type SessionManager struct { - sessions map[string]*Session - mu sync.RWMutex - storage string + sessions map[string]*Session + mu sync.RWMutex + storage string + summarizer Summarizer // optional; nil = no summarization + summarizerCfg config.SummarizationConfig // filled via WithDefaults at construction + inflight sync.Map // sessionKey → true (dedup guard) } -func NewSessionManager(storage string) *SessionManager { +// Option configures a SessionManager at construction time. +type Option func(*SessionManager) + +// WithSummarizer enables background summarization and emergency compression. +// The Summarizer is called to produce summaries via the LLM. Zero-valued +// config fields are replaced with defaults. +func WithSummarizer(s Summarizer, cfg config.SummarizationConfig) Option { + return func(sm *SessionManager) { + sm.summarizer = s + sm.summarizerCfg = cfg.WithDefaults() + } +} + +func NewSessionManager(storage string, opts ...Option) *SessionManager { sm := &SessionManager{ sessions: make(map[string]*Session), storage: storage, } + for _, o := range opts { + o(sm) + } + if storage != "" { os.MkdirAll(storage, 0o755) sm.loadSessions() @@ -280,3 +305,250 @@ func (sm *SessionManager) SetHistory(key string, history []providers.Message) { session.Updated = time.Now() } } + +// ApplySummarization atomically sets a summary and trims old messages while +// preserving any messages that were appended after the caller's snapshot. +// +// snapshotLen is the len(session.Messages) at the time the caller took its +// snapshot. If the session has grown since then, the new tail is preserved. +// If messages were removed (e.g. another compression), the operation is +// skipped and returns false. +// +// keepLast is the number of messages from the *snapshot* to retain (counted +// from the end of the snapshot window, not the current session). +func (sm *SessionManager) ApplySummarization(key, summary string, snapshotLen, keepLast int) bool { + sm.mu.Lock() + defer sm.mu.Unlock() + + session, ok := sm.sessions[key] + if !ok { + return false + } + + currentLen := len(session.Messages) + + // Stale: session was truncated/replaced since our snapshot. + if currentLen < snapshotLen { + return false + } + + // Messages appended by the main loop after the snapshot was taken. + newTail := session.Messages[snapshotLen:] + + // Messages to keep from the original snapshot window. + var keptFromSnapshot []providers.Message + if keepLast > 0 && keepLast < snapshotLen { + keptFromSnapshot = session.Messages[snapshotLen-keepLast : snapshotLen] + } else if keepLast >= snapshotLen { + keptFromSnapshot = session.Messages[:snapshotLen] + } + + merged := make([]providers.Message, 0, len(keptFromSnapshot)+len(newTail)) + merged = append(merged, keptFromSnapshot...) + merged = append(merged, newTail...) + + session.Messages = merged + session.Summary = summary + session.Updated = time.Now() + return true +} + +// EstimateTokens estimates token count for a list of messages using the +// configured chars-per-token ratio and unicode rune counting (CJK-aware). +func (sm *SessionManager) EstimateTokens(messages []providers.Message) int { + totalRunes := 0 + for _, m := range messages { + totalRunes += utf8.RuneCountInString(m.Content) + } + return int(float64(totalRunes) / sm.summarizerCfg.CharsPerToken) +} + +// MaybeSummarize checks whether the session's history exceeds the configured +// thresholds and, if so, triggers background summarization. At most one +// summarization runs per session key at a time. +// No-op if no Summarizer was provided via WithSummarizer. +func (sm *SessionManager) MaybeSummarize(sessionKey string) { + if sm.summarizer == nil { + return + } + + history := sm.GetHistory(sessionKey) + tokenEstimate := sm.EstimateTokens(history) + cfg := sm.summarizerCfg + threshold := cfg.ContextWindow * cfg.TokenPercent / 100 + + if len(history) <= cfg.MessageThreshold && tokenEstimate <= threshold { + return + } + + if _, loaded := sm.inflight.LoadOrStore(sessionKey, true); loaded { + return // already running for this session + } + + go func() { + defer sm.inflight.Delete(sessionKey) + logger.Debug("Memory threshold reached. Optimizing conversation history...") + sm.summarizeSession(sessionKey) + }() +} + +// summarizeSession performs the actual summarization: snapshot history, +// generate summary via LLM, and atomically apply the result while +// preserving any messages added by the main loop during summarization. +func (sm *SessionManager) summarizeSession(sessionKey string) { + cfg := sm.summarizerCfg + ctx, cancel := context.WithTimeout(context.Background(), cfg.Timeout) + defer cancel() + + history := sm.GetHistory(sessionKey) + existingSummary := sm.GetSummary(sessionKey) + snapshotLen := len(history) + + keepLast := cfg.KeepLastMessages + if snapshotLen <= keepLast { + return + } + + toSummarize := history[:snapshotLen-keepLast] + + // Oversized message guard: skip individual messages that exceed + // MaxSingleMsgTokenRatio of the context window. + maxMessageTokens := int(float64(cfg.ContextWindow) * cfg.MaxSingleMsgTokenRatio) + var validMessages []providers.Message + omitted := false + + for _, m := range toSummarize { + if m.Role != "user" && m.Role != "assistant" { + continue + } + msgTokens := int(float64(utf8.RuneCountInString(m.Content)) / cfg.CharsPerToken) + if msgTokens > maxMessageTokens { + omitted = true + continue + } + validMessages = append(validMessages, m) + } + + if len(validMessages) == 0 { + return + } + + // Produce summary, splitting into two batches for large conversations. + var finalSummary string + if len(validMessages) > cfg.MultiPartBatchThreshold { + mid := len(validMessages) / 2 + part1 := validMessages[:mid] + part2 := validMessages[mid:] + + var s1, s2 string + var wg sync.WaitGroup + wg.Add(2) + go func() { + defer wg.Done() + s1, _ = sm.summarizer.Summarize(ctx, part1, "") + }() + go func() { + defer wg.Done() + s2, _ = sm.summarizer.Summarize(ctx, part2, "") + }() + wg.Wait() + + // Merge the two partial summaries. + mergeMessages := []providers.Message{ + { + Role: "user", + Content: fmt.Sprintf( + "Merge these two conversation summaries into one cohesive summary:\n\n1: %s\n\n2: %s", + s1, + s2, + ), + }, + } + merged, err := sm.summarizer.Summarize(ctx, mergeMessages, "") + if err == nil { + finalSummary = merged + } else { + finalSummary = s1 + " " + s2 + } + } else { + var err error + finalSummary, err = sm.summarizer.Summarize(ctx, validMessages, existingSummary) + if err != nil { + logger.WarnCF("session", "Summarization failed", map[string]any{"error": err.Error()}) + return + } + } + + if omitted && finalSummary != "" { + finalSummary += "\n[Note: Some oversized messages were omitted from this summary for efficiency.]" + } + + if finalSummary == "" { + return + } + + // Atomically apply: sets summary and trims old messages while preserving + // anything the main loop appended after our snapshot. + applied := sm.ApplySummarization(sessionKey, finalSummary, snapshotLen, keepLast) + if applied { + sm.Save(sessionKey) + logger.InfoCF("session", "Summarization applied", map[string]any{ + "session_key": sessionKey, + "snapshot_len": snapshotLen, + "keep_last": keepLast, + }) + } else { + logger.WarnCF("session", "Summarization skipped (stale snapshot)", map[string]any{ + "session_key": sessionKey, + "snapshot_len": snapshotLen, + }) + } +} + +// ForceCompression aggressively reduces context when the LLM returns a +// context-window error. It drops the oldest ~50% of conversation messages, +// preserving the system prompt (first message) and the last message. +// +// This is called synchronously from the LLM retry loop — the main loop is +// already blocked, so there is no concurrent-append concern here. +// No-op if no Summarizer was provided via WithSummarizer. +func (sm *SessionManager) ForceCompression(sessionKey string) { + cfg := sm.summarizerCfg + history := sm.GetHistory(sessionKey) + if len(history) <= cfg.ForceCompressionMinMessages { + return + } + + // history[0] is the system prompt, history[len-1] is the trigger message. + conversation := history[1 : len(history)-1] + if len(conversation) == 0 { + return + } + + mid := len(conversation) / 2 + droppedCount := mid + keptConversation := conversation[mid:] + + newHistory := make([]providers.Message, 0, 1+len(keptConversation)+1) + + // Append compression note to the system prompt to avoid consecutive + // system messages (rejected by some APIs like Zhipu). + compressionNote := fmt.Sprintf( + "\n\n[System Note: Emergency compression dropped %d oldest messages due to context limit]", + droppedCount, + ) + enhancedSystemPrompt := history[0] + enhancedSystemPrompt.Content += compressionNote + newHistory = append(newHistory, enhancedSystemPrompt) + newHistory = append(newHistory, keptConversation...) + newHistory = append(newHistory, history[len(history)-1]) + + sm.SetHistory(sessionKey, newHistory) + sm.Save(sessionKey) + + logger.WarnCF("session", "Forced compression executed", map[string]any{ + "session_key": sessionKey, + "dropped_msgs": droppedCount, + "new_count": len(newHistory), + }) +} diff --git a/pkg/session/summarizer.go b/pkg/session/summarizer.go new file mode 100644 index 000000000..36b22088c --- /dev/null +++ b/pkg/session/summarizer.go @@ -0,0 +1,102 @@ +package session + +import ( + "context" + "fmt" + "strings" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" +) + +// Summarizer is the LLM-calling strategy injected by the agent layer. +// The session package is agnostic to how the LLM is called — it only +// invokes this interface to produce textual summaries. +type Summarizer interface { + Summarize(ctx context.Context, messages []providers.Message, existingSummary string) (string, error) +} + +// Compile-time interface satisfaction checks. +var ( + _ Summarizer = (*LLMSummarizer)(nil) + _ Summarizer = SummarizeFunc(nil) +) + +// SummarizeFunc is a function adapter for Summarizer. +// It allows passing a plain function where a Summarizer is expected. +type SummarizeFunc func(ctx context.Context, messages []providers.Message, existingSummary string) (string, error) + +// Summarize implements the Summarizer interface. +func (f SummarizeFunc) Summarize( + ctx context.Context, + messages []providers.Message, + existingSummary string, +) (string, error) { + return f(ctx, messages, existingSummary) +} + +// BuildSummarizationPrompt formats messages into a summarization prompt string. +// This is intended to be used inside the SummarizeFunc implementation provided +// by the agent layer. +func BuildSummarizationPrompt(batch []providers.Message, existingSummary string) string { + var sb strings.Builder + sb.WriteString( + "Provide a concise summary of this conversation segment, preserving core context and key points.\n", + ) + if existingSummary != "" { + sb.WriteString("Existing context: ") + sb.WriteString(existingSummary) + sb.WriteString("\n") + } + sb.WriteString("\nCONVERSATION:\n") + for _, m := range batch { + fmt.Fprintf(&sb, "%s: %s\n", m.Role, m.Content) + } + return sb.String() +} + +// LLMSummarizer implements Summarizer by calling an LLM provider. +// This is the default implementation used in production. +type LLMSummarizer struct { + provider providers.LLMProvider + model string + agentID string + cfg config.SummarizationConfig +} + +// Summarize formats the messages into a summarization prompt and calls the +// LLM to produce a concise summary. +func (s *LLMSummarizer) Summarize( + ctx context.Context, + msgs []providers.Message, + existingSummary string, +) (string, error) { + prompt := BuildSummarizationPrompt(msgs, existingSummary) + resp, err := s.provider.Chat( + ctx, + []providers.Message{{Role: "user", Content: prompt}}, + nil, + s.model, + map[string]any{ + "max_tokens": s.cfg.SummaryMaxTokens, + "temperature": s.cfg.SummaryTemperature, + "prompt_cache_key": s.agentID, + }, + ) + if err != nil { + return "", err + } + return resp.Content, nil +} + +// WithLLMSummarizer is a convenience option that creates an LLMSummarizer +// and configures the SessionManager for summarization in one call. +func WithLLMSummarizer(provider providers.LLMProvider, model, agentID string, cfg config.SummarizationConfig) Option { + cfg = cfg.WithDefaults() + return WithSummarizer(&LLMSummarizer{ + provider: provider, + model: model, + agentID: agentID, + cfg: cfg, + }, cfg) +} diff --git a/pkg/session/summarizer_test.go b/pkg/session/summarizer_test.go new file mode 100644 index 000000000..ae9ec23e8 --- /dev/null +++ b/pkg/session/summarizer_test.go @@ -0,0 +1,454 @@ +package session + +import ( + "context" + "strings" + "sync" + "testing" + "time" + + "github.com/sipeed/picoclaw/pkg/config" + "github.com/sipeed/picoclaw/pkg/providers" +) + +func newTestManager(t *testing.T, fn SummarizeFunc, cfg config.SummarizationConfig) *SessionManager { + t.Helper() + return NewSessionManager(t.TempDir(), WithSummarizer(fn, cfg)) +} + +func noopSummarize(_ context.Context, _ []providers.Message, _ string) (string, error) { + return "summary", nil +} + +// --- EstimateTokens --- + +func TestEstimateTokens_Latin(t *testing.T) { + sm := newTestManager(t, noopSummarize, config.SummarizationConfig{}) + + msgs := []providers.Message{ + {Role: "user", Content: "hello world"}, // 11 runes + {Role: "assistant", Content: "hi there!"}, // 9 runes + } + // 20 runes / 2.5 = 8 tokens + got := sm.EstimateTokens(msgs) + if got != 8 { + t.Errorf("EstimateTokens = %d, want 8", got) + } +} + +func TestEstimateTokens_CJK(t *testing.T) { + sm := newTestManager(t, noopSummarize, config.SummarizationConfig{}) + + msgs := []providers.Message{ + {Role: "user", Content: "你好世界"}, //nolint:gosmopolitan // CJK text "你好世界" (4 runes) + } + // 4 / 2.5 = 1.6 → 1 + got := sm.EstimateTokens(msgs) + if got != 1 { + t.Errorf("EstimateTokens = %d, want 1", got) + } +} + +func TestEstimateTokens_CustomCharsPerToken(t *testing.T) { + cfg := config.SummarizationConfig{CharsPerToken: 5.0} + sm := newTestManager(t, noopSummarize, cfg) + + msgs := []providers.Message{ + {Role: "user", Content: "12345678901234567890"}, // 20 runes + } + // 20 / 5.0 = 4 + got := sm.EstimateTokens(msgs) + if got != 4 { + t.Errorf("EstimateTokens = %d, want 4", got) + } +} + +// --- MaybeSummarize --- + +func TestMaybeSummarize_BelowThreshold(t *testing.T) { + called := false + fn := func(_ context.Context, _ []providers.Message, _ string) (string, error) { + called = true + return "summary", nil + } + + cfg := config.SummarizationConfig{} + cfg.MessageThreshold = 100 + cfg.TokenPercent = 99 + sm := newTestManager(t, fn, cfg) + + // Add a few messages — well below thresholds. + for i := 0; i < 5; i++ { + sm.AddMessage("test", "user", "hello") + } + + sm.MaybeSummarize("test") + time.Sleep(50 * time.Millisecond) // give goroutine time to fire (it shouldn't) + + if called { + t.Error("SummarizeFunc was called despite being below thresholds") + } +} + +func TestMaybeSummarize_AboveMessageThreshold(t *testing.T) { + done := make(chan struct{}) + fn := func(_ context.Context, _ []providers.Message, _ string) (string, error) { + defer func() { + select { + case done <- struct{}{}: + default: + } + }() + return "summary of conversation", nil + } + + cfg := config.SummarizationConfig{} + cfg.MessageThreshold = 5 + cfg.KeepLastMessages = 2 + sm := newTestManager(t, fn, cfg) + + for i := 0; i < 10; i++ { + sm.AddMessage("test", "user", "msg") + sm.AddMessage("test", "assistant", "reply") + } + + sm.MaybeSummarize("test") + + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for summarization") + } + + // Wait a bit for ApplySummarization to complete. + time.Sleep(100 * time.Millisecond) + + summary := sm.GetSummary("test") + if summary == "" { + t.Error("expected non-empty summary after summarization") + } + + history := sm.GetHistory("test") + if len(history) > cfg.KeepLastMessages+2 { + // Some slack for messages added during the goroutine window. + t.Errorf("expected at most %d messages after summarization, got %d", + cfg.KeepLastMessages+2, len(history)) + } +} + +func TestMaybeSummarize_DeduplicatesConcurrent(t *testing.T) { + goroutineCount := 0 + var mu sync.Mutex + blocker := make(chan struct{}) + + fn := func(_ context.Context, _ []providers.Message, _ string) (string, error) { + mu.Lock() + goroutineCount++ + mu.Unlock() + <-blocker + return "summary", nil + } + + cfg := config.SummarizationConfig{} + cfg.MessageThreshold = 3 + cfg.KeepLastMessages = 2 + sm := newTestManager(t, fn, cfg) + + for i := 0; i < 10; i++ { + sm.AddMessage("test", "user", "msg") + sm.AddMessage("test", "assistant", "reply") + } + + // Trigger twice — only one goroutine should enter summarizeSession. + sm.MaybeSummarize("test") + sm.MaybeSummarize("test") + + // Let the blocker hold briefly so the second call sees inflight=true. + time.Sleep(50 * time.Millisecond) + + close(blocker) + time.Sleep(200 * time.Millisecond) + + mu.Lock() + got := goroutineCount + mu.Unlock() + + // summarizeSession may call Summarize multiple times internally + // (multi-part batching via parallel goroutines), but only ONE + // goroutine should have entered summarizeSession. + // With 18 valid messages and threshold=10: 3 calls (part1+part2+merge). + // If dedup failed, we'd see 6. + if got > 3 { + t.Errorf("SummarizeFunc called %d times, want at most 3 (single goroutine with multi-part)", got) + } +} + +// --- ForceCompression --- + +func TestForceCompression_DropsOldestHalf(t *testing.T) { + sm := newTestManager(t, noopSummarize, config.SummarizationConfig{}) + + // [system, msg1, msg2, msg3, msg4, msg5, trigger] + sm.AddFullMessage("test", providers.Message{Role: "system", Content: "system prompt"}) + for i := 1; i <= 5; i++ { + sm.AddMessage("test", "user", "old message") + } + sm.AddMessage("test", "user", "trigger") + + sm.ForceCompression("test") + + history := sm.GetHistory("test") + if len(history) == 0 { + t.Fatal("history is empty after compression") + } + + // System prompt should be first and contain the compression note. + if history[0].Role != "system" { + t.Errorf("first message role = %q, want system", history[0].Role) + } + if !strings.Contains(history[0].Content, "Emergency compression") { + t.Error("system prompt missing compression note") + } + + // Trigger message should be last. + last := history[len(history)-1] + if last.Content != "trigger" { + t.Errorf("last message = %q, want trigger", last.Content) + } +} + +func TestForceCompression_TooFewMessages(t *testing.T) { + sm := newTestManager(t, noopSummarize, config.SummarizationConfig{}) + + sm.AddMessage("test", "user", "one") + sm.AddMessage("test", "assistant", "two") + + sm.ForceCompression("test") + + history := sm.GetHistory("test") + if len(history) != 2 { + t.Errorf("expected 2 messages unchanged, got %d", len(history)) + } +} + +// --- ApplySummarization --- + +func TestApplySummarization_PreservesNewMessages(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + + // Seed 10 messages. + for i := 0; i < 10; i++ { + sm.AddMessage("test", "user", "original") + } + snapshotLen := 10 + + // Simulate the main loop appending 3 new messages after the snapshot. + sm.AddMessage("test", "user", "new-1") + sm.AddMessage("test", "assistant", "new-2") + sm.AddMessage("test", "user", "new-3") + + applied := sm.ApplySummarization("test", "the summary", snapshotLen, 4) + if !applied { + t.Fatal("ApplySummarization returned false") + } + + history := sm.GetHistory("test") + // Should have: 4 kept from snapshot + 3 new = 7 + if len(history) != 7 { + t.Errorf("len(history) = %d, want 7", len(history)) + } + + // Verify the new messages are at the end. + if history[4].Content != "new-1" { + t.Errorf("history[4] = %q, want new-1", history[4].Content) + } + if history[6].Content != "new-3" { + t.Errorf("history[6] = %q, want new-3", history[6].Content) + } + + // Summary should be set. + summary := sm.GetSummary("test") + if summary != "the summary" { + t.Errorf("summary = %q, want 'the summary'", summary) + } +} + +func TestApplySummarization_StaleSnapshot(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + + for i := 0; i < 10; i++ { + sm.AddMessage("test", "user", "original") + } + snapshotLen := 10 + + // Simulate another compression reducing messages below snapshot. + sm.SetHistory("test", []providers.Message{ + {Role: "user", Content: "compressed-1"}, + {Role: "assistant", Content: "compressed-2"}, + }) + + applied := sm.ApplySummarization("test", "stale summary", snapshotLen, 4) + if applied { + t.Error("ApplySummarization should return false for stale snapshot") + } + + // Summary should NOT have changed. + summary := sm.GetSummary("test") + if summary == "stale summary" { + t.Error("stale summary was applied despite stale snapshot") + } +} + +func TestApplySummarization_NoNewMessages(t *testing.T) { + sm := NewSessionManager(t.TempDir()) + + for i := 0; i < 10; i++ { + sm.AddMessage("test", "user", "original") + } + + applied := sm.ApplySummarization("test", "the summary", 10, 4) + if !applied { + t.Fatal("ApplySummarization returned false") + } + + history := sm.GetHistory("test") + // Should have exactly 4 kept. + if len(history) != 4 { + t.Errorf("len(history) = %d, want 4", len(history)) + } +} + +// --- SummarizeSession Multi-Part --- + +func TestSummarizeSession_MultiPart(t *testing.T) { + callCount := 0 + var mu sync.Mutex + fn := func(_ context.Context, msgs []providers.Message, existing string) (string, error) { + mu.Lock() + callCount++ + mu.Unlock() + return "partial summary", nil + } + + cfg := config.SummarizationConfig{} + cfg.MultiPartBatchThreshold = 5 + cfg.KeepLastMessages = 2 + sm := newTestManager(t, fn, cfg) + + // Add 12 user+assistant messages (well above threshold of 5). + for i := 0; i < 12; i++ { + sm.AddMessage("test", "user", "hello") + sm.AddMessage("test", "assistant", "world") + } + + sm.summarizeSession("test") + + mu.Lock() + got := callCount + mu.Unlock() + + // Should be 3 calls: part1 + part2 + merge. + if got != 3 { + t.Errorf("SummarizeFunc called %d times, want 3 (part1 + part2 + merge)", got) + } + + summary := sm.GetSummary("test") + if summary == "" { + t.Error("expected non-empty summary after multi-part summarization") + } +} + +// --- Concurrent safety: summarize + main loop writes --- + +func TestConcurrent_SummarizeAndWrite(t *testing.T) { + // Simulates the core race scenario: background summarization runs + // while the main loop adds new messages. + snapshotTaken := make(chan struct{}) + blocker := make(chan struct{}) + firstCall := true + var mu sync.Mutex + + fn := func(_ context.Context, _ []providers.Message, _ string) (string, error) { + mu.Lock() + if firstCall { + firstCall = false + mu.Unlock() + // Signal that summarizeSession has taken its snapshot and entered summarize. + close(snapshotTaken) + <-blocker // block until new messages are added + } else { + mu.Unlock() + } + return "concurrent summary", nil + } + + cfg := config.SummarizationConfig{} + cfg.KeepLastMessages = 2 + cfg.MultiPartBatchThreshold = 100 // force single-batch path + sm := newTestManager(t, fn, cfg) + + // Seed 20 messages. + for i := 0; i < 20; i++ { + sm.AddMessage("test", "user", "seed") + } + + // Start summarization in background. + done := make(chan struct{}) + go func() { + sm.summarizeSession("test") + close(done) + }() + + // Wait for the summarizer to take its snapshot and enter the LLM call. + <-snapshotTaken + + // While summarization is blocked, add 5 new messages. + for i := 0; i < 5; i++ { + sm.AddMessage("test", "user", "new-during-summarize") + } + + // Unblock the summarizer. + close(blocker) + <-done + + // All 5 new messages must survive. + history := sm.GetHistory("test") + newCount := 0 + for _, m := range history { + if m.Content == "new-during-summarize" { + newCount++ + } + } + if newCount != 5 { + t.Errorf("expected 5 new messages preserved, got %d (total history: %d)", + newCount, len(history)) + } + + summary := sm.GetSummary("test") + if summary == "" { + t.Error("expected non-empty summary") + } +} + +// --- BuildSummarizationPrompt --- + +func TestBuildSummarizationPrompt(t *testing.T) { + msgs := []providers.Message{ + {Role: "user", Content: "hello"}, + {Role: "assistant", Content: "hi"}, + } + + prompt := BuildSummarizationPrompt(msgs, "") + if !strings.Contains(prompt, "concise summary") { + t.Error("prompt missing summary instruction") + } + if !strings.Contains(prompt, "user: hello") { + t.Error("prompt missing user message") + } + + // With existing summary. + prompt2 := BuildSummarizationPrompt(msgs, "previous context") + if !strings.Contains(prompt2, "Existing context: previous context") { + t.Error("prompt missing existing summary") + } +}