222 lines
6.8 KiB
Go
222 lines
6.8 KiB
Go
package agent
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestConversationCRUDAndRestart(t *testing.T) {
|
|
s, secrets := testService(t)
|
|
a, err := s.CreateConversation()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
b, err := s.CreateConversation()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
a.Draft = "未发送的草稿"
|
|
a.Messages = []ConversationMessage{{ID: "u1", Role: "user", Content: "包装盒采购需求", CreatedAt: 1}, {ID: "a1", Role: "assistant", Content: "已保存的回复", Status: "complete", CreatedAt: 2}}
|
|
summary, err := s.SaveConversation(a)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if summary.Title != "包装盒采购需求" || summary.MessageCount != 2 {
|
|
t.Fatalf("summary: %+v", summary)
|
|
}
|
|
if _, err = s.RenameConversation(a.ID, "供应商咨询"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
// A queued autosave must preserve a more recent explicit rename.
|
|
if _, err = s.SaveConversation(a); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
restarted := newService(s.configPath, secrets)
|
|
loaded, err := restarted.GetConversation(a.ID)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if loaded.Title != "供应商咨询" || loaded.Draft != a.Draft || len(loaded.Messages) != 2 {
|
|
t.Fatalf("lost persisted data: %+v", loaded)
|
|
}
|
|
empty, err := restarted.GetConversation(b.ID)
|
|
if err != nil || len(empty.Messages) != 0 {
|
|
t.Fatal("conversation messages leaked")
|
|
}
|
|
list, err := restarted.ListConversations()
|
|
if err != nil || len(list) != 2 {
|
|
t.Fatalf("list: %v %v", list, err)
|
|
}
|
|
if err = restarted.DeleteConversation(a.ID); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if _, err = restarted.SaveConversation(a); err == nil {
|
|
t.Fatal("late update recreated a deleted conversation")
|
|
}
|
|
list, err = restarted.ListConversations()
|
|
if err != nil || len(list) != 1 || list[0].ID != b.ID {
|
|
t.Fatalf("delete affected wrong conversation: %+v %v", list, err)
|
|
}
|
|
}
|
|
|
|
func TestConversationRejectsTraversalAndBadMessages(t *testing.T) {
|
|
s, _ := testService(t)
|
|
for _, id := range []string{"../model", "a/b", "a\\b", "..", "C:target"} {
|
|
if _, err := s.GetConversation(id); err == nil {
|
|
t.Errorf("accepted invalid ID %q", id)
|
|
}
|
|
if err := s.DeleteConversation(id); err == nil {
|
|
t.Errorf("delete accepted invalid ID %q", id)
|
|
}
|
|
}
|
|
c, err := s.CreateConversation()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
c.Messages = []ConversationMessage{{ID: "x", Role: "user", Content: "a"}, {ID: "x", Role: "assistant", Content: "b"}}
|
|
if _, err = s.SaveConversation(c); err == nil {
|
|
t.Fatal("duplicate message IDs accepted")
|
|
}
|
|
c.Messages = []ConversationMessage{{ID: "x", Role: "developer", Content: "a"}}
|
|
if _, err = s.SaveConversation(c); err == nil {
|
|
t.Fatal("unsupported role accepted")
|
|
}
|
|
stored, err := s.GetConversation(c.ID)
|
|
if err != nil || len(stored.Messages) != 0 {
|
|
t.Fatal("invalid save overwrote original")
|
|
}
|
|
path, _ := s.conversationPath(c.ID)
|
|
if _, err = os.Stat(filepath.Dir(path)); err != nil {
|
|
t.Fatal("conversation directory not created")
|
|
}
|
|
}
|
|
|
|
func TestRestoredContextIsIsolatedAndExcludesIncompleteReplies(t *testing.T) {
|
|
requests := make(chan []modelMessage, 2)
|
|
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)
|
|
requests <- body.Messages
|
|
reply(w)
|
|
}))
|
|
defer server.Close()
|
|
s, secrets := testService(t)
|
|
config := testConfig(server.URL)
|
|
config.ContextTurns = 10
|
|
if _, err := s.SaveModelConfig(ModelConfigInput{Config: config, APIKey: "key"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
a, _ := s.CreateConversation()
|
|
b, _ := s.CreateConversation()
|
|
a.Messages = []ConversationMessage{
|
|
{ID: "u1", Role: "user", Content: "需求 A"}, {ID: "a1", Role: "assistant", Content: "回复 A", Status: "complete"},
|
|
{ID: "u2", Role: "user", Content: "已停止的问题"}, {ID: "a2", Role: "assistant", Content: "半截内容", Status: "stopped"},
|
|
{ID: "u3", Role: "user", Content: "本轮问题"}, {ID: "a3", Role: "assistant", Status: "streaming"},
|
|
}
|
|
if _, err := s.SaveConversation(a); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
restarted := newService(s.configPath, secrets)
|
|
if _, err := restarted.Chat(a.ID, "继续 A"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
first := <-requests
|
|
if len(first) != 4 || first[1].Content != "需求 A" || first[2].Content != "回复 A" || first[3].Content != "继续 A" {
|
|
t.Fatalf("wrong restored context: %+v", first)
|
|
}
|
|
if _, err := restarted.Chat(b.ID, "开始 B"); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
second := <-requests
|
|
if len(second) != 2 || second[1].Content != "开始 B" {
|
|
t.Fatalf("cross-conversation context leak: %+v", second)
|
|
}
|
|
}
|
|
|
|
func TestDifferentConversationsCanStreamConcurrently(t *testing.T) {
|
|
arrived := make(chan struct{}, 2)
|
|
release := make(chan struct{})
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
arrived <- struct{}{}
|
|
select {
|
|
case <-release:
|
|
case <-r.Context().Done():
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}]}\n\ndata: [DONE]\n\n")
|
|
}))
|
|
defer server.Close()
|
|
ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
|
|
defer cancel()
|
|
s, _ := testService(t)
|
|
if _, err := s.SaveModelConfig(ModelConfigInput{Config: testConfig(server.URL), APIKey: "key"}); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
errors := make(chan error, 2)
|
|
for _, id := range []string{"A", "B"} {
|
|
go func(id string) { _, err := s.ChatStream(ctx, id, id, "hello"); errors <- err }(id)
|
|
}
|
|
for i := 0; i < 2; i++ {
|
|
select {
|
|
case <-arrived:
|
|
case <-ctx.Done():
|
|
t.Fatal("one conversation blocked another")
|
|
}
|
|
}
|
|
close(release)
|
|
for i := 0; i < 2; i++ {
|
|
if err := <-errors; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
if len(s.chatLocks) != 0 {
|
|
t.Fatal("conversation locks leaked")
|
|
}
|
|
}
|
|
|
|
func TestCancelledWaiterDoesNotStartAnotherTurn(t *testing.T) {
|
|
s, _ := testService(t)
|
|
unlock, err := s.lockConversation(context.Background(), "same")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if release, err := s.lockConversation(ctx, "same"); err == nil {
|
|
release()
|
|
t.Fatal("cancelled waiter acquired locked chat")
|
|
}
|
|
unlock()
|
|
if len(s.chatLocks) != 0 {
|
|
t.Fatal("cancelled waiter leaked a lock")
|
|
}
|
|
}
|
|
|
|
func TestCorruptConversationIsReportedWithoutOverwriting(t *testing.T) {
|
|
s, _ := testService(t)
|
|
c, _ := s.CreateConversation()
|
|
path, _ := s.conversationPath(c.ID)
|
|
os.WriteFile(path, []byte("{corrupt"), 0600)
|
|
if _, err := s.ListConversations(); err == nil || !strings.Contains(err.Error(), "格式损坏") {
|
|
t.Fatalf("corruption hidden: %v", err)
|
|
}
|
|
if _, err := s.SaveConversation(c); err == nil {
|
|
t.Fatal("corrupt record overwritten silently")
|
|
}
|
|
raw, _ := os.ReadFile(path)
|
|
if string(raw) != "{corrupt" {
|
|
t.Fatal("corrupt original lost")
|
|
}
|
|
}
|