ituoke/internal/agent/memory_test.go

248 lines
9.6 KiB
Go
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

package agent
import (
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"net/http/httptest"
"os"
"strings"
"testing"
"unicode/utf8"
)
func memoryConversation(t *testing.T, s *Service, n int) Conversation {
t.Helper()
c, err := s.CreateConversation()
if err != nil {
t.Fatal(err)
}
for i := 1; i <= n; i++ {
c.Messages = append(c.Messages, ConversationMessage{ID: fmt.Sprintf("u%d", i), Role: "user", Content: fmt.Sprintf("原始需求-%d", i)}, ConversationMessage{ID: fmt.Sprintf("a%d", i), Role: "assistant", Content: fmt.Sprintf("确认结果-%d", i), Status: "complete"})
}
if _, err := s.SaveConversation(c); err != nil {
t.Fatal(err)
}
return c
}
func memoryReply(w http.ResponseWriter, text string) {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]any{"choices": []any{map[string]any{"message": map[string]any{"content": text}, "finish_reason": "stop"}}})
}
func TestAutomaticMemoryPersistsIncrementallyAndSurvivesStaleSave(t *testing.T) {
var summaryInputs []string
var answerInputs [][]modelMessage
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body struct {
Messages []modelMessage `json:"messages"`
Tools json.RawMessage `json:"tools"`
MaxTokens int `json:"max_tokens"`
}
json.NewDecoder(r.Body).Decode(&body)
if body.Messages[0].Content == memoryInstruction {
if len(body.Tools) > 0 {
t.Error("summary requested tools")
}
summaryInputs = append(summaryInputs, fmt.Sprint(body.Messages[1].Content))
memoryReply(w, "关键记忆:预算3500元,数量500,只起草不发送。")
} else {
answerInputs = append(answerInputs, body.Messages)
memoryReply(w, "继续执行需求分析。")
}
}))
defer server.Close()
s, secrets := testService(t)
cfg := testConfig(server.URL)
cfg.ContextTurns = 1
if _, err := s.SaveModelConfig(ModelConfigInput{Config: cfg, APIKey: "key"}); err != nil {
t.Fatal(err)
}
c := memoryConversation(t, s, 3)
c.Messages[0].Content = "预算3500元,数量500,只起草不发送。原始需求-1"
s.SaveConversation(c)
if _, err := s.Chat(c.ID, "继续"); err != nil {
t.Fatal(err)
}
stored, err := s.GetConversation(c.ID)
if err != nil || stored.Memory == nil || stored.Memory.MessageCount != 4 || len(stored.Messages) != 6 {
t.Fatalf("checkpoint: %+v %v", stored.Memory, err)
}
if len(summaryInputs) != 1 || !strings.Contains(summaryInputs[0], "原始需求-2") || strings.Contains(summaryInputs[0], "原始需求-3") {
t.Fatalf("wrong summary source: %v", summaryInputs)
}
serialized, _ := json.Marshal(answerInputs[0])
if !strings.Contains(string(serialized), "关键记忆") || !strings.Contains(string(serialized), "原始需求-3") || strings.Contains(string(serialized), "原始需求-2") {
t.Fatalf("context not compacted: %s", serialized)
}
// The frontend loaded c before the summary, then saves messages and a draft.
c.Draft = "新草稿"
s.SaveConversation(c)
stored, _ = s.GetConversation(c.ID)
if stored.Memory == nil || stored.Draft != "新草稿" {
t.Fatal("stale save erased server-owned memory")
}
restarted := newService(s.configPath, secrets)
if _, err := restarted.Chat(c.ID, "重启后继续"); err != nil {
t.Fatal(err)
}
if len(summaryInputs) != 1 {
t.Fatal("unchanged history was summarized again")
}
c.Messages = append(c.Messages, ConversationMessage{ID: "u4", Role: "user", Content: "第四轮"}, ConversationMessage{ID: "a4", Role: "assistant", Content: "第四轮确认", Status: "complete"})
s.SaveConversation(c)
if _, err := restarted.Chat(c.ID, "继续增量"); err != nil {
t.Fatal(err)
}
if len(summaryInputs) != 2 || !strings.Contains(summaryInputs[1], "关键记忆") || !strings.Contains(summaryInputs[1], "原始需求-3") || strings.Contains(summaryInputs[1], "原始需求-2") {
t.Fatalf("not incremental: %v", summaryInputs)
}
other := memoryConversation(t, s, 0)
if _, err := restarted.Chat(other.ID, "另一会话"); err != nil {
t.Fatal(err)
}
serialized, _ = json.Marshal(answerInputs[len(answerInputs)-1])
if strings.Contains(string(serialized), "关键记忆") {
t.Fatal("memory leaked across conversations")
}
c.Messages[0].Content = "用户修改了早期需求"
s.SaveConversation(c)
stored, _ = s.GetConversation(c.ID)
if stored.Memory != nil {
t.Fatal("edited prefix kept stale memory")
}
}
func TestLongHistoryChunksWithoutSplittingUnicode(t *testing.T) {
s, _ := testService(t)
c := memoryConversation(t, s, 2)
c.Messages[1].Content = strings.Repeat("重要信息🙂", 10000)
s.SaveConversation(c)
chunks := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body struct {
Messages []modelMessage `json:"messages"`
}
json.NewDecoder(r.Body).Decode(&body)
if body.Messages[0].Content != memoryInstruction {
t.Error("unexpected answer request")
}
text := fmt.Sprint(body.Messages[1].Content)
if !utf8.ValidString(text) || strings.ContainsRune(text, '\ufffd') {
t.Error("Unicode split")
}
chunks++
memoryReply(w, "记住重要信息,继续保留关键要求。")
}))
defer server.Close()
cfg := testConfig(server.URL)
cfg.ContextTurns = 10
turns, memory, found, err := s.prepareMemory(context.Background(), c.ID, cfg, "", nil)
if err != nil || !found || memory == nil || memory.MessageCount != 2 || len(turns) != 2 || chunks < 2 {
t.Fatalf("long history: %d %d %+v %v", chunks, len(turns), memory, err)
}
}
func TestSummaryFailureAndCancellationPreserveOriginal(t *testing.T) {
for _, mode := range []string{"provider-error", "truncated", "cancel", "changed", "deleted"} {
t.Run(mode, func(t *testing.T) {
s, _ := testService(t)
c := memoryConversation(t, s, 2)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
switch mode {
case "provider-error":
http.Error(w, "failed", 500)
return
case "truncated":
fmt.Fprint(w, `{"choices":[{"message":{"content":"半截摘要"},"finish_reason":"length"}]}`)
return
case "cancel":
cancel()
return
case "changed":
c.Messages[0].Content = "已修改"
s.SaveConversation(c)
case "deleted":
s.DeleteConversation(c.ID)
}
memoryReply(w, "摘要")
}))
defer server.Close()
cfg := testConfig(server.URL)
cfg.ContextTurns = 1
_, _, _, err := s.prepareMemory(ctx, c.ID, cfg, "", nil)
if err == nil {
t.Fatal("summary failure was hidden")
}
if mode == "cancel" && !errors.Is(err, context.Canceled) {
t.Fatalf("cancellation lost: %v", err)
}
path, _ := s.conversationPath(c.ID)
if mode == "deleted" {
if _, err := os.Stat(path); !os.IsNotExist(err) {
t.Fatal("summary recreated deleted conversation")
}
return
}
stored, err := s.GetConversation(c.ID)
if err != nil || stored.Memory != nil || len(stored.Messages) != 4 {
t.Fatalf("failed summary changed record: %+v %v", stored, err)
}
})
}
}
func TestZeroContextDoesNotReadOrGenerateMemory(t *testing.T) {
s, _ := testService(t)
c := memoryConversation(t, s, 3)
cfg := testConfig("http://127.0.0.1:1")
cfg.ContextTurns = 0
turns, memory, found, err := s.prepareMemory(context.Background(), c.ID, cfg, "", nil)
if err != nil || !found || len(turns) != 0 || memory != nil {
t.Fatalf("zero context: %v %+v %v", turns, memory, err)
}
}
func TestSummaryPreservesExistingCheckpointOnFailureAndIncludesToolStates(t *testing.T) {
s, _ := testService(t)
c := memoryConversation(t, s, 3)
c.Messages[1].Tools = []ToolExecution{{ID: "tool", Server: "Phone", Tool: "launch", Status: "complete", Result: "settings opened"}}
c.Messages = append(c.Messages[:2], append([]ConversationMessage{{ID: "failed-user", Role: "user", Content: "继续操作"}, {ID: "failed-answer", Role: "assistant", Content: "已中断", Status: "stopped", Tools: []ToolExecution{{ID: "tool-failed", Server: "Phone", Tool: "tap", Status: "error", Result: "device offline"}}}}, c.Messages[2:]...)...)
if _, err := s.SaveConversation(c); err != nil {
t.Fatal(err)
}
var source string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
var body struct {
Messages []modelMessage `json:"messages"`
}
json.NewDecoder(r.Body).Decode(&body)
source = fmt.Sprint(body.Messages[1].Content)
memoryReply(w, "已打开设置;后续操作中断,设备离线,不代表后续动作完成。")
}))
defer server.Close()
cfg := testConfig(server.URL)
cfg.ContextTurns = 1
_, memory, _, err := s.prepareMemory(context.Background(), c.ID, cfg, "", nil)
if err != nil || memory == nil || !strings.Contains(source, "状态=stopped") || !strings.Contains(source, "状态=error") || !strings.Contains(source, "settings opened") || !strings.Contains(source, "device offline") {
t.Fatalf("tool states lost: %s %v", source, err)
}
c.Messages = append(c.Messages, ConversationMessage{ID: "new-u", Role: "user", Content: "新的完成回合"}, ConversationMessage{ID: "new-a", Role: "assistant", Content: "新的结果", Status: "complete"})
s.SaveConversation(c)
failing := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "offline", 503) }))
defer failing.Close()
cfg.BaseURL = failing.URL
if _, _, _, err = s.prepareMemory(context.Background(), c.ID, cfg, "", nil); err == nil {
t.Fatal("failure hidden")
}
stored, _ := s.GetConversation(c.ID)
if stored.Memory == nil || stored.Memory.Text != memory.Text || stored.Memory.ThroughID != memory.ThroughID {
t.Fatal("failed update erased previous checkpoint")
}
}