From f8196c111c526e19016e58556b2627bbbbe052ab Mon Sep 17 00:00:00 2001 From: Aska Date: Thu, 20 Aug 2026 20:34:08 +0800 Subject: [PATCH] feat(gateway): persist chat sessions and messages in PostgreSQL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - store: new chat_sessions/chat_messages tables with CRUD, ILIKE message search, and auto-titling from the first user message - context: optional Persister on the sliding-window store — messages are saved asynchronously on Add, sessions not in memory are restored from the database on Get - orchestrator: ProcessInput/FinishRound now carry userID so persisted messages have an owner - api: real session handlers (list/create/delete), message history and clear endpoints, working message search and three-format export (markdown/txt/json), fixing the short-ID panic in export - main: migration, store wiring, ctxStore persister adapter Tests: chat store handler lifecycle (fake store), context persister save/restore, stub auth login flow updated. --- backend/cmd/cyrene/main.go | 41 ++- backend/services/ai-core/context/context.go | 48 +++- .../services/ai-core/context/context_test.go | 98 +++++++ .../ai-core/orchestrator/orchestrator.go | 6 +- .../orchestrator/orchestrator_tools_test.go | 2 +- backend/services/gateway/api/handlers_more.go | 60 ++++- backend/services/gateway/api/router.go | 90 ++++++- backend/services/gateway/api/router_test.go | 2 +- backend/services/gateway/api/sessions_test.go | 239 ++++++++++++++++++ backend/services/gateway/store/chat_store.go | 224 ++++++++++++++++ backend/services/gateway/ws/hub.go | 4 +- 11 files changed, 787 insertions(+), 27 deletions(-) create mode 100644 backend/services/ai-core/context/context_test.go create mode 100644 backend/services/gateway/api/sessions_test.go create mode 100644 backend/services/gateway/store/chat_store.go diff --git a/backend/cmd/cyrene/main.go b/backend/cmd/cyrene/main.go index f429556..fe8c87d 100644 --- a/backend/cmd/cyrene/main.go +++ b/backend/cmd/cyrene/main.go @@ -37,10 +37,12 @@ import ( tracepkg "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/trace" perfpkg "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/perf" aiCoreCtx "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/context" + "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model" "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/orchestrator" "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/persona" "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/subsession" "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/api" + "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/store" "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/ws" "git.yeij.top/AskaEth/Cyrene_v2/services/memory" bridge "git.yeij.top/AskaEth/Cyrene_v2/services/platform-bridge" @@ -79,6 +81,7 @@ func main() { } else { defer database.Close() db.RunMigration(context.Background(), database, "memories", db.SplitSQL(memory.CreateTableSQL())) + db.RunMigration(context.Background(), database, "chat", db.SplitSQL(store.CreateTableSQL())) } // === 3. NATS (in-process) === @@ -161,6 +164,14 @@ func main() { ctxStore := aiCoreCtx.NewStore(4000, 50) logger.Infof("[context] sliding window ready (4000 tokens, 50 turns)") + // === 9.5 Chat persistence (sessions + messages) === + var chatStore *store.ChatStore + if database != nil { + chatStore = store.NewChatStore(database) + ctxStore.SetPersister(&chatPersisterAdapter{s: chatStore}) + logger.Infof("[chat] session/message persistence ready") + } + // === 10. AI-Core: Sub-session Manager === subMgr := subsession.NewManager( &subsession.GeneralProvider{}, @@ -238,6 +249,9 @@ func main() { if memSvc != nil { gwDeps.MemorySvc = memory.NewAPIAdapter(memSvc) } + if chatStore != nil { + gwDeps.ChatStore = chatStore + } gwDeps.VoiceSvc = voiceSvc wsHub := ws.NewHub(orch) @@ -331,6 +345,31 @@ func (a *perfAdapter) History() []map[string]any { return result } +// chatPersisterAdapter adapts store.ChatStore to aiCoreCtx.Persister. +type chatPersisterAdapter struct{ s *store.ChatStore } + +func (a *chatPersisterAdapter) SaveMessage(sessionID, userID string, msg model.Message) { + if userID == "" { + userID = "unknown" + } + if err := a.s.AddMessage(context.Background(), sessionID, userID, msg.Role, msg.Content); err != nil { + logger.Infof("[chat] persist message failed: %v", err) + } +} + +func (a *chatPersisterAdapter) LoadRecent(sessionID string, limit int) []model.Message { + msgs, err := a.s.ListMessages(context.Background(), sessionID, limit) + if err != nil { + logger.Infof("[chat] load history failed: %v", err) + return nil + } + out := make([]model.Message, len(msgs)) + for i, m := range msgs { + out[i] = model.Message{SessionID: m.SessionID, Role: m.Role, Content: m.Content, CreatedAt: m.CreatedAt} + } + return out +} + // toolExecAdapter adapts the plugin registry to orchestrator.ToolExecutor. type toolExecAdapter struct{ reg *plugin.Registry } @@ -389,7 +428,7 @@ func loadAdaptersFromConfig(router *bridge.Router, cfg *bridge.BridgeConfig, orc router.SetMessageHandler(func(msg *bridge.UnifiedMessage) { sessionID := "bridge-" + msg.Source - tokenCh, _, err := orch.ProcessInput(context.Background(), sessionID, msg.Content) + tokenCh, _, err := orch.ProcessInput(context.Background(), sessionID, msg.SenderID, msg.Content) if err != nil { return } var fc string for t := range tokenCh { fc += t.Content } diff --git a/backend/services/ai-core/context/context.go b/backend/services/ai-core/context/context.go index 3b3c513..ff73b81 100644 --- a/backend/services/ai-core/context/context.go +++ b/backend/services/ai-core/context/context.go @@ -10,12 +10,23 @@ import ( "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model" ) +// Persister saves conversation messages durably and can restore them +// (e.g. from PostgreSQL) when a session is not in memory. +type Persister interface { + // SaveMessage persists one message. Called synchronously by Add — + // implementations should be fast or dispatch async work themselves. + SaveMessage(sessionID, userID string, msg model.Message) + // LoadRecent returns the most recent messages of a session, ascending. + LoadRecent(sessionID string, limit int) []model.Message +} + // Store holds conversation history with a sliding window. type Store struct { mu sync.RWMutex sessions map[string]*SessionHistory maxTokens int // Maximum tokens to include in context maxTurns int // Maximum message turns to keep + persister Persister } // SessionHistory holds messages for a single session. @@ -39,14 +50,31 @@ func NewStore(maxTokens, maxTurns int) *Store { } } +// SetPersister injects durable persistence for conversation messages. +func (s *Store) SetPersister(p Persister) { + s.mu.Lock() + s.persister = p + s.mu.Unlock() +} + // Add appends a message to the session history. func (s *Store) Add(sessionID string, msg model.Message) { + s.add(sessionID, "", msg) +} + +// AddWithUser appends a message and records the owning user for persistence. +func (s *Store) AddWithUser(sessionID, userID string, msg model.Message) { + s.add(sessionID, userID, msg) +} + +func (s *Store) add(sessionID, userID string, msg model.Message) { s.mu.Lock() sh, ok := s.sessions[sessionID] if !ok { sh = &SessionHistory{} s.sessions[sessionID] = sh } + persister := s.persister s.mu.Unlock() sh.mu.Lock() @@ -56,15 +84,33 @@ func (s *Store) Add(sessionID string, msg model.Message) { sh.Messages = sh.Messages[len(sh.Messages)-s.maxTurns:] } sh.mu.Unlock() + + // Persist asynchronously — never block the conversation on the database. + if persister != nil && (msg.Role == "user" || msg.Role == "assistant") { + go persister.SaveMessage(sessionID, userID, msg) + } } // Get retrieves recent messages for a session, trimmed to fit token budget. +// Sessions missing from memory are restored from the persister if available. func (s *Store) Get(sessionID string) []model.Message { s.mu.RLock() sh, ok := s.sessions[sessionID] + persister := s.persister s.mu.RUnlock() + if !ok { - return nil + if persister == nil { + return nil + } + loaded := persister.LoadRecent(sessionID, s.maxTurns) + if len(loaded) == 0 { + return nil + } + sh = &SessionHistory{Messages: loaded} + s.mu.Lock() + s.sessions[sessionID] = sh + s.mu.Unlock() } sh.mu.RLock() diff --git a/backend/services/ai-core/context/context_test.go b/backend/services/ai-core/context/context_test.go new file mode 100644 index 0000000..837abad --- /dev/null +++ b/backend/services/ai-core/context/context_test.go @@ -0,0 +1,98 @@ +package context + +import ( + "sync" + "testing" + "time" + + "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model" +) + +// fakePersister records saved messages and can restore them. +type fakePersister struct { + mu sync.Mutex + saved map[string][]model.Message +} + +func newFakePersister() *fakePersister { + return &fakePersister{saved: map[string][]model.Message{}} +} + +func (f *fakePersister) SaveMessage(sessionID, userID string, msg model.Message) { + f.mu.Lock() + defer f.mu.Unlock() + f.saved[sessionID] = append(f.saved[sessionID], msg) +} + +func (f *fakePersister) LoadRecent(sessionID string, limit int) []model.Message { + f.mu.Lock() + defer f.mu.Unlock() + return f.saved[sessionID] +} + +func (f *fakePersister) count(sessionID string) int { + f.mu.Lock() + defer f.mu.Unlock() + return len(f.saved[sessionID]) +} + +func TestPersisterSaveOnAdd(t *testing.T) { + s := NewStore(4000, 50) + p := newFakePersister() + s.SetPersister(p) + + s.AddWithUser("s1", "u1", model.Message{Role: "user", Content: "你好"}) + + // SaveMessage is dispatched asynchronously — wait briefly + for i := 0; i < 50 && p.count("s1") == 0; i++ { + time.Sleep(10 * time.Millisecond) + } + if p.count("s1") != 1 { + t.Fatalf("expected 1 persisted message, got %d", p.count("s1")) + } +} + +func TestPersisterSkipsNonChatRoles(t *testing.T) { + s := NewStore(4000, 50) + p := newFakePersister() + s.SetPersister(p) + + s.Add("s1", model.Message{Role: "system", Content: "上下文块"}) + time.Sleep(50 * time.Millisecond) + if p.count("s1") != 0 { + t.Fatalf("system messages should not be persisted, got %d", p.count("s1")) + } +} + +func TestRestoreFromPersisterOnGet(t *testing.T) { + s := NewStore(4000, 50) + p := newFakePersister() + p.saved["s2"] = []model.Message{ + {Role: "user", Content: "以前说过的话"}, + {Role: "assistant", Content: "以前的回复"}, + } + s.SetPersister(p) + + // Session not in memory — should restore from persister + msgs := s.Get("s2") + if len(msgs) != 2 { + t.Fatalf("expected 2 restored messages, got %d", len(msgs)) + } + if msgs[0].Content != "以前说过的话" { + t.Errorf("unexpected restored content: %q", msgs[0].Content) + } + + // Second Get should hit memory, not reload + p.saved["s2"] = nil + msgs = s.Get("s2") + if len(msgs) != 2 { + t.Errorf("expected cached messages after restore, got %d", len(msgs)) + } +} + +func TestGetWithoutPersister(t *testing.T) { + s := NewStore(4000, 50) + if msgs := s.Get("missing"); msgs != nil { + t.Errorf("expected nil for unknown session, got %v", msgs) + } +} diff --git a/backend/services/ai-core/orchestrator/orchestrator.go b/backend/services/ai-core/orchestrator/orchestrator.go index 7a09d21..bfcbbdd 100644 --- a/backend/services/ai-core/orchestrator/orchestrator.go +++ b/backend/services/ai-core/orchestrator/orchestrator.go @@ -79,7 +79,7 @@ func (p *Pipeline) SetToolExecutor(t ToolExecutor) { // ProcessInput processes a user message through the full pipeline. // Returns a channel of StreamToken for the LLM response. -func (p *Pipeline) ProcessInput(ctx context.Context, sessionID, userMessage string) (<-chan model.StreamToken, *model.Intent, error) { +func (p *Pipeline) ProcessInput(ctx context.Context, sessionID, userID, userMessage string) (<-chan model.StreamToken, *model.Intent, error) { // 1. Cache check if cached := p.cache.Get(userMessage); cached != nil { ch := make(chan model.StreamToken, 1) @@ -122,7 +122,7 @@ func (p *Pipeline) ProcessInput(ctx context.Context, sessionID, userMessage stri } // 7. Store user message - p.ctxStore.Add(sessionID, model.Message{Role: "user", Content: userMessage}) + p.ctxStore.AddWithUser(sessionID, userID, model.Message{Role: "user", Content: userMessage}) return tokenCh, intent, nil } @@ -193,7 +193,7 @@ func (p *Pipeline) streamWithTools(ctx context.Context, adapter *llm.Adapter, me // after the response stream completes. func (p *Pipeline) FinishRound(sessionID, userID, assistantReply string) { if assistantReply != "" { - p.ctxStore.Add(sessionID, model.Message{Role: "assistant", Content: assistantReply}) + p.ctxStore.AddWithUser(sessionID, userID, model.Message{Role: "assistant", Content: assistantReply}) } p.PostProcess(sessionID, userID) } diff --git a/backend/services/ai-core/orchestrator/orchestrator_tools_test.go b/backend/services/ai-core/orchestrator/orchestrator_tools_test.go index 595537c..6f4562d 100644 --- a/backend/services/ai-core/orchestrator/orchestrator_tools_test.go +++ b/backend/services/ai-core/orchestrator/orchestrator_tools_test.go @@ -87,7 +87,7 @@ func TestToolCallLoop(t *testing.T) { exec := &fakeToolExec{} p.SetToolExecutor(exec) - tokenCh, _, err := p.ProcessInput(context.Background(), "s1", "打开客厅灯") + tokenCh, _, err := p.ProcessInput(context.Background(), "s1", "u1", "打开客厅灯") if err != nil { t.Fatalf("ProcessInput: %v", err) } diff --git a/backend/services/gateway/api/handlers_more.go b/backend/services/gateway/api/handlers_more.go index c48d4fa..7327f74 100644 --- a/backend/services/gateway/api/handlers_more.go +++ b/backend/services/gateway/api/handlers_more.go @@ -166,20 +166,60 @@ func (r *Router) notificationHandler(w http.ResponseWriter, req *http.Request) { func (r *Router) messageSearchHandler(w http.ResponseWriter, req *http.Request) { q := req.URL.Query().Get("q") - uid := req.Header.Get("X-User-ID") + uid := r.userID(req) if q == "" { writeJSON(w, 400, M{"error": "missing query"}); return } - // Uses session store DB — returns matching messages - writeJSON(w, 200, M{"messages": []any{}, "query": q, "user_id": uid, "total": 0}) + if r.deps.ChatStore == nil { writeJSON(w, 503, M{"error": "chat store unavailable"}); return } + + messages, err := r.deps.ChatStore.SearchMessages(req.Context(), uid, q, 20) + if err != nil { + writeJSON(w, 500, M{"error": err.Error()}) + return + } + writeJSON(w, 200, M{"messages": messages, "query": q, "total": len(messages)}) } func (r *Router) sessionExportHandler(w http.ResponseWriter, req *http.Request) { - id := strings.TrimPrefix(req.URL.Path, "/api/v1/sessions/") - id = strings.TrimSuffix(id, "/export") - // Export session as markdown - w.Header().Set("Content-Type", "text/markdown; charset=utf-8") - w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=session-%s.md", id[:8])) - w.WriteHeader(200) - fmt.Fprintf(w, "# Cyrene Chat - Session %s\n\n*Exported at %s*\n\n", id[:8], time.Now().Format(time.RFC3339)) + id := strings.TrimSuffix(strings.TrimPrefix(req.URL.Path, "/api/v1/sessions/"), "/export") + if r.deps.ChatStore == nil { writeJSON(w, 503, M{"error": "chat store unavailable"}); return } + + messages, err := r.deps.ChatStore.ListMessages(req.Context(), id, 200) + if err != nil { + writeJSON(w, 500, M{"error": err.Error()}) + return + } + + short := id + if len(short) > 8 { short = short[:8] } + format := req.URL.Query().Get("format") + exportedAt := time.Now().Format(time.RFC3339) + + switch format { + case "json": + w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=session-%s.json", short)) + writeJSON(w, 200, M{"session_id": id, "exported_at": exportedAt, "messages": messages}) + case "txt": + var sb strings.Builder + fmt.Fprintf(&sb, "Cyrene Chat - Session %s (exported %s)\n\n", short, exportedAt) + for _, m := range messages { + fmt.Fprintf(&sb, "[%s] %s: %s\n\n", m.CreatedAt.Format("2006-01-02 15:04"), m.Role, m.Content) + } + w.Header().Set("Content-Type", "text/plain; charset=utf-8") + w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=session-%s.txt", short)) + w.WriteHeader(200) + w.Write([]byte(sb.String())) + default: // markdown + var sb strings.Builder + fmt.Fprintf(&sb, "# Cyrene Chat - Session %s\n\n*Exported at %s*\n\n", short, exportedAt) + for _, m := range messages { + who := "用户" + if m.Role == "assistant" { who = "昔涟" } + fmt.Fprintf(&sb, "**%s**（%s）：\n\n%s\n\n---\n\n", who, m.CreatedAt.Format("2006-01-02 15:04"), m.Content) + } + w.Header().Set("Content-Type", "text/markdown; charset=utf-8") + w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=session-%s.md", short)) + w.WriteHeader(200) + w.Write([]byte(sb.String())) + } } func (r *Router) imageAnalyzeHandler(w http.ResponseWriter, req *http.Request) { diff --git a/backend/services/gateway/api/router.go b/backend/services/gateway/api/router.go index 6799156..9480758 100644 --- a/backend/services/gateway/api/router.go +++ b/backend/services/gateway/api/router.go @@ -12,6 +12,7 @@ import ( "time" "git.yeij.top/AskaEth/Cyrene_v2/services/ai-core/model" + "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/store" ) // --- Interfaces for injected services --- @@ -51,7 +52,7 @@ type VoiceService interface { // OrchestratorPipeline handles the main dialogue flow. type OrchestratorPipeline interface { - ProcessInput(ctx context.Context, sessionID, userMessage string) (<-chan model.StreamToken, *model.Intent, error) + ProcessInput(ctx context.Context, sessionID, userID, userMessage string) (<-chan model.StreamToken, *model.Intent, error) } // WebSocketHub manages WS connections. @@ -85,6 +86,16 @@ type BridgeConfigManager interface { DeleteAdapter(name string) error } +// ChatStorer persists chat sessions and messages. +type ChatStorer interface { + CreateSession(ctx context.Context, id, userID, title string) (*store.Session, error) + ListSessions(ctx context.Context, userID string) ([]store.Session, error) + DeleteSession(ctx context.Context, id, userID string) error + ListMessages(ctx context.Context, sessionID string, limit int) ([]store.Message, error) + ClearMessages(ctx context.Context, sessionID string) error + SearchMessages(ctx context.Context, userID, query string, limit int) ([]store.Message, error) +} + // Dependencies holds all injectable services. type Dependencies struct { AuthMgr AuthManager @@ -96,6 +107,7 @@ type Dependencies struct { TraceLog TraceLogger PerfMon PerfMonitor BridgeCfg BridgeConfigManager + ChatStore ChatStorer ModelConfig *ModelConfigStore Automation *AutomationStore Reminders *ReminderStore @@ -257,29 +269,91 @@ func (r *Router) registerHandler(w http.ResponseWriter, req *http.Request) { // --- Sessions --- func (r *Router) sessionsHandler(w http.ResponseWriter, req *http.Request) { + if r.deps.ChatStore == nil { + writeJSON(w, 503, M{"error": "chat store unavailable"}) + return + } + uid := r.userID(req) switch req.Method { case "GET": - writeJSON(w, 200, M{"sessions": []any{}}) + sessions, err := r.deps.ChatStore.ListSessions(req.Context(), uid) + if err != nil { + writeJSON(w, 500, M{"error": err.Error()}) + return + } + writeJSON(w, 200, M{"sessions": sessions}) case "POST": - writeJSON(w, 201, M{"session_id": randID()}) + var body struct { + Title string `json:"title"` + } + json.NewDecoder(req.Body).Decode(&body) + sess, err := r.deps.ChatStore.CreateSession(req.Context(), randID(), uid, body.Title) + if err != nil { + writeJSON(w, 500, M{"error": err.Error()}) + return + } + writeJSON(w, 201, M{"session_id": sess.ID, "title": sess.Title}) default: writeJSON(w, 405, M{"error": "method not allowed"}) } } func (r *Router) sessionByIDHandler(w http.ResponseWriter, req *http.Request) { - id := strings.TrimPrefix(req.URL.Path, "/api/v1/sessions/") - // Handle export - if strings.HasSuffix(id, "/export") { - id = strings.TrimSuffix(id, "/export") + if r.deps.ChatStore == nil { + writeJSON(w, 503, M{"error": "chat store unavailable"}) + return + } + rest := strings.TrimPrefix(req.URL.Path, "/api/v1/sessions/") + + // Sub-resources: /sessions/{id}/export and /sessions/{id}/messages + if strings.HasSuffix(rest, "/export") { r.sessionExportHandler(w, req) return } + if strings.HasSuffix(rest, "/messages") { + r.sessionMessagesHandler(w, req) + return + } + + id := rest switch req.Method { case "DELETE": + if err := r.deps.ChatStore.DeleteSession(req.Context(), id, r.userID(req)); err != nil { + writeJSON(w, 404, M{"error": err.Error()}) + return + } writeJSON(w, 204, nil) case "GET": - writeJSON(w, 200, M{"session_id": id}) + messages, err := r.deps.ChatStore.ListMessages(req.Context(), id, 100) + if err != nil { + writeJSON(w, 500, M{"error": err.Error()}) + return + } + writeJSON(w, 200, M{"session_id": id, "messages": messages}) + default: + writeJSON(w, 405, M{"error": "method not allowed"}) + } +} + +// sessionMessagesHandler serves GET (recent history) and DELETE (clear). +func (r *Router) sessionMessagesHandler(w http.ResponseWriter, req *http.Request) { + id := strings.TrimSuffix(strings.TrimPrefix(req.URL.Path, "/api/v1/sessions/"), "/messages") + switch req.Method { + case "GET": + limit := 0 + fmt.Sscanf(req.URL.Query().Get("limit"), "%d", &limit) + messages, err := r.deps.ChatStore.ListMessages(req.Context(), id, limit) + if err != nil { + writeJSON(w, 500, M{"error": err.Error()}) + return + } + writeJSON(w, 200, M{"messages": messages}) + case "DELETE": + if err := r.deps.ChatStore.ClearMessages(req.Context(), id); err != nil { + writeJSON(w, 500, M{"error": err.Error()}) + return + } + writeJSON(w, 204, nil) default: writeJSON(w, 405, M{"error": "method not allowed"}) } diff --git a/backend/services/gateway/api/router_test.go b/backend/services/gateway/api/router_test.go index fa87b96..9165574 100644 --- a/backend/services/gateway/api/router_test.go +++ b/backend/services/gateway/api/router_test.go @@ -25,7 +25,7 @@ func TestHealth(t *testing.T) { func TestLoginAndProtectedEndpoint(t *testing.T) { auth := NewStubAuthManager() - r := New(Dependencies{AuthMgr: auth}) + r := New(Dependencies{AuthMgr: auth, ChatStore: newFakeChatStore()}) srv := httptest.NewServer(r) defer srv.Close() diff --git a/backend/services/gateway/api/sessions_test.go b/backend/services/gateway/api/sessions_test.go new file mode 100644 index 0000000..604e889 --- /dev/null +++ b/backend/services/gateway/api/sessions_test.go @@ -0,0 +1,239 @@ +package api + +import ( + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "git.yeij.top/AskaEth/Cyrene_v2/services/gateway/store" +) + +// fakeChatStore is an in-memory ChatStorer for tests. +type fakeChatStore struct { + mu sync.Mutex + sessions map[string]*store.Session + messages map[string][]store.Message + nextID int64 +} + +func newFakeChatStore() *fakeChatStore { + return &fakeChatStore{sessions: map[string]*store.Session{}, messages: map[string][]store.Message{}} +} + +func (f *fakeChatStore) CreateSession(_ context.Context, id, userID, title string) (*store.Session, error) { + f.mu.Lock() + defer f.mu.Unlock() + s := &store.Session{ID: id, UserID: userID, Title: title, CreatedAt: time.Now(), UpdatedAt: time.Now()} + f.sessions[id] = s + return s, nil +} + +func (f *fakeChatStore) ListSessions(_ context.Context, userID string) ([]store.Session, error) { + f.mu.Lock() + defer f.mu.Unlock() + out := []store.Session{} + for _, s := range f.sessions { + if s.UserID == userID { + out = append(out, *s) + } + } + return out, nil +} + +func (f *fakeChatStore) DeleteSession(_ context.Context, id, userID string) error { + f.mu.Lock() + defer f.mu.Unlock() + if _, ok := f.sessions[id]; !ok { + return fmt.Errorf("session not found") + } + delete(f.sessions, id) + delete(f.messages, id) + return nil +} + +func (f *fakeChatStore) AddMessage(sessionID, role, content string) { + f.mu.Lock() + defer f.mu.Unlock() + f.nextID++ + f.messages[sessionID] = append(f.messages[sessionID], store.Message{ + ID: f.nextID, SessionID: sessionID, Role: role, Content: content, CreatedAt: time.Now(), + }) +} + +func (f *fakeChatStore) ListMessages(_ context.Context, sessionID string, limit int) ([]store.Message, error) { + f.mu.Lock() + defer f.mu.Unlock() + msgs := f.messages[sessionID] + if msgs == nil { + msgs = []store.Message{} + } + return msgs, nil +} + +func (f *fakeChatStore) ClearMessages(_ context.Context, sessionID string) error { + f.mu.Lock() + defer f.mu.Unlock() + f.messages[sessionID] = nil + return nil +} + +func (f *fakeChatStore) SearchMessages(_ context.Context, userID, query string, limit int) ([]store.Message, error) { + f.mu.Lock() + defer f.mu.Unlock() + out := []store.Message{} + for _, msgs := range f.messages { + for _, m := range msgs { + if strings.Contains(m.Content, query) { + out = append(out, m) + } + } + } + return out, nil +} + +func authReq(method, url, body string) *http.Request { + var r *http.Request + if body != "" { + r, _ = http.NewRequest(method, url, strings.NewReader(body)) + r.Header.Set("Content-Type", "application/json") + } else { + r, _ = http.NewRequest(method, url, nil) + } + r.Header.Set("Authorization", "Bearer stub-token-admin") + return r +} + +func TestSessionLifecycle(t *testing.T) { + auth := NewStubAuthManager() + auth.Login("admin", "admin") // register stub-token-admin + chat := newFakeChatStore() + r := New(Dependencies{AuthMgr: auth, ChatStore: chat}) + srv := httptest.NewServer(r) + defer srv.Close() + + // Create session + resp, err := http.DefaultClient.Do(authReq("POST", srv.URL+"/api/v1/sessions", `{"title":""}`)) + if err != nil { + t.Fatalf("POST sessions: %v", err) + } + if resp.StatusCode != http.StatusCreated { + t.Fatalf("create: expected 201, got %d", resp.StatusCode) + } + var created struct { + SessionID string `json:"session_id"` + } + readJSON(t, resp, &created) + if created.SessionID == "" { + t.Fatal("empty session_id") + } + + // Seed messages directly + chat.AddMessage(created.SessionID, "user", "今晚想吃火锅") + chat.AddMessage(created.SessionID, "assistant", "好呀，要不要叫上朋友") + + // List sessions + resp, _ = http.DefaultClient.Do(authReq("GET", srv.URL+"/api/v1/sessions", "")) + var listed struct { + Sessions []store.Session `json:"sessions"` + } + readJSON(t, resp, &listed) + if len(listed.Sessions) != 1 { + t.Fatalf("list: expected 1 session, got %d", len(listed.Sessions)) + } + + // Get messages + resp, _ = http.DefaultClient.Do(authReq("GET", srv.URL+"/api/v1/sessions/"+created.SessionID+"/messages", "")) + var msgs struct { + Messages []store.Message `json:"messages"` + } + readJSON(t, resp, &msgs) + if len(msgs.Messages) != 2 { + t.Fatalf("messages: expected 2, got %d", len(msgs.Messages)) + } + + // Search + resp, _ = http.DefaultClient.Do(authReq("GET", srv.URL+"/api/v1/messages/search?q=火锅", "")) + var found struct { + Total int `json:"total"` + } + readJSON(t, resp, &found) + if found.Total != 1 { + t.Errorf("search: expected 1 hit, got %d", found.Total) + } + + // Export markdown + resp, _ = http.DefaultClient.Do(authReq("GET", srv.URL+"/api/v1/sessions/"+created.SessionID+"/export", "")) + body := readBody(t, resp) + if !strings.Contains(body, "火锅") || !strings.Contains(body, "# Cyrene Chat") { + t.Errorf("export markdown missing content: %q", body) + } + + // Export json + resp, _ = http.DefaultClient.Do(authReq("GET", srv.URL+"/api/v1/sessions/"+created.SessionID+"/export?format=json", "")) + var exported struct { + Messages []store.Message `json:"messages"` + } + readJSON(t, resp, &exported) + if len(exported.Messages) != 2 { + t.Errorf("export json: expected 2 messages, got %d", len(exported.Messages)) + } + + // Clear messages + resp, _ = http.DefaultClient.Do(authReq("DELETE", srv.URL+"/api/v1/sessions/"+created.SessionID+"/messages", "")) + if resp.StatusCode != http.StatusNoContent { + t.Errorf("clear: expected 204, got %d", resp.StatusCode) + } + resp.Body.Close() + + // Delete session + resp, _ = http.DefaultClient.Do(authReq("DELETE", srv.URL+"/api/v1/sessions/"+created.SessionID, "")) + if resp.StatusCode != http.StatusNoContent { + t.Errorf("delete: expected 204, got %d", resp.StatusCode) + } + resp.Body.Close() + + resp, _ = http.DefaultClient.Do(authReq("DELETE", srv.URL+"/api/v1/sessions/"+created.SessionID, "")) + if resp.StatusCode != http.StatusNotFound { + t.Errorf("delete again: expected 404, got %d", resp.StatusCode) + } + resp.Body.Close() +} + +func TestSessionsUnavailableWithoutStore(t *testing.T) { + auth := NewStubAuthManager() + auth.Login("admin", "admin") + r := New(Dependencies{AuthMgr: auth}) + srv := httptest.NewServer(r) + defer srv.Close() + + resp, _ := http.DefaultClient.Do(authReq("GET", srv.URL+"/api/v1/sessions", "")) + if resp.StatusCode != http.StatusServiceUnavailable { + t.Errorf("expected 503 without store, got %d", resp.StatusCode) + } + resp.Body.Close() +} + +func readJSON(t *testing.T, resp *http.Response, v any) { + t.Helper() + defer resp.Body.Close() + if err := json.NewDecoder(resp.Body).Decode(v); err != nil { + t.Fatalf("decode response: %v", err) + } +} + +func readBody(t *testing.T, resp *http.Response) string { + t.Helper() + defer resp.Body.Close() + b, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatalf("read body: %v", err) + } + return string(b) +} diff --git a/backend/services/gateway/store/chat_store.go b/backend/services/gateway/store/chat_store.go new file mode 100644 index 0000000..2905ac6 --- /dev/null +++ b/backend/services/gateway/store/chat_store.go @@ -0,0 +1,224 @@ +// Package store provides PostgreSQL persistence for chat sessions and messages. +package store + +import ( + "context" + "database/sql" + "fmt" + "strings" + "time" +) + +// Session represents a chat session. +type Session struct { + ID string `json:"id"` + UserID string `json:"user_id"` + Title string `json:"title"` + CreatedAt time.Time `json:"created_at"` + UpdatedAt time.Time `json:"updated_at"` +} + +// Message represents a single chat message. +type Message struct { + ID int64 `json:"id"` + SessionID string `json:"session_id"` + Role string `json:"role"` // "user" | "assistant" + Content string `json:"content"` + CreatedAt time.Time `json:"created_at"` +} + +// ChatStore persists sessions and messages in PostgreSQL. +type ChatStore struct { + db *sql.DB +} + +// NewChatStore creates a store backed by the given database. +func NewChatStore(db *sql.DB) *ChatStore { + return &ChatStore{db: db} +} + +// CreateTableSQL returns the schema for session/message persistence. +func CreateTableSQL() string { + return ` +CREATE TABLE IF NOT EXISTS chat_sessions ( + id TEXT PRIMARY KEY, + user_id TEXT NOT NULL, + title TEXT NOT NULL DEFAULT '', + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_chat_sessions_user ON chat_sessions(user_id, updated_at DESC); + +CREATE TABLE IF NOT EXISTS chat_messages ( + id BIGSERIAL PRIMARY KEY, + session_id TEXT NOT NULL REFERENCES chat_sessions(id) ON DELETE CASCADE, + user_id TEXT NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() +); +CREATE INDEX IF NOT EXISTS idx_chat_messages_session ON chat_messages(session_id, id); +CREATE INDEX IF NOT EXISTS idx_chat_messages_user ON chat_messages(user_id, id); +` +} + +// --- Sessions --- + +// CreateSession creates a new session for the user. +func (s *ChatStore) CreateSession(ctx context.Context, id, userID, title string) (*Session, error) { + sess := &Session{ID: id, UserID: userID, Title: title} + err := s.db.QueryRowContext(ctx, ` + INSERT INTO chat_sessions (id, user_id, title) VALUES ($1, $2, $3) + RETURNING created_at, updated_at + `, id, userID, title).Scan(&sess.CreatedAt, &sess.UpdatedAt) + if err != nil { + return nil, fmt.Errorf("create session: %w", err) + } + return sess, nil +} + +// ListSessions returns the user's sessions, most recently active first. +func (s *ChatStore) ListSessions(ctx context.Context, userID string) ([]Session, error) { + rows, err := s.db.QueryContext(ctx, ` + SELECT id, user_id, title, created_at, updated_at + FROM chat_sessions WHERE user_id = $1 + ORDER BY updated_at DESC LIMIT 100 + `, userID) + if err != nil { + return nil, fmt.Errorf("list sessions: %w", err) + } + defer rows.Close() + + sessions := []Session{} + for rows.Next() { + var sess Session + if err := rows.Scan(&sess.ID, &sess.UserID, &sess.Title, &sess.CreatedAt, &sess.UpdatedAt); err != nil { + return nil, err + } + sessions = append(sessions, sess) + } + return sessions, rows.Err() +} + +// DeleteSession removes a session and its messages (CASCADE). +func (s *ChatStore) DeleteSession(ctx context.Context, id, userID string) error { + res, err := s.db.ExecContext(ctx, `DELETE FROM chat_sessions WHERE id = $1 AND user_id = $2`, id, userID) + if err != nil { + return fmt.Errorf("delete session: %w", err) + } + if n, _ := res.RowsAffected(); n == 0 { + return fmt.Errorf("session not found") + } + return nil +} + +// --- Messages --- + +// AddMessage appends a message, creating the session row if needed and +// bumping updated_at. The first user message becomes the session title. +func (s *ChatStore) AddMessage(ctx context.Context, sessionID, userID, role, content string) error { + tx, err := s.db.BeginTx(ctx, nil) + if err != nil { + return fmt.Errorf("begin tx: %w", err) + } + defer tx.Rollback() + + if _, err := tx.ExecContext(ctx, ` + INSERT INTO chat_sessions (id, user_id) VALUES ($1, $2) ON CONFLICT (id) DO NOTHING + `, sessionID, userID); err != nil { + return fmt.Errorf("ensure session: %w", err) + } + + if _, err := tx.ExecContext(ctx, ` + INSERT INTO chat_messages (session_id, user_id, role, content) VALUES ($1, $2, $3, $4) + `, sessionID, userID, role, content); err != nil { + return fmt.Errorf("insert message: %w", err) + } + + title := "" + if role == "user" { + title = truncateRunes(strings.TrimSpace(content), 24) + } + if _, err := tx.ExecContext(ctx, ` + UPDATE chat_sessions SET + updated_at = now(), + title = CASE WHEN title = '' AND $3 <> '' THEN $3 ELSE title END + WHERE id = $1 AND user_id = $2 + `, sessionID, userID, title); err != nil { + return fmt.Errorf("touch session: %w", err) + } + + return tx.Commit() +} + +// ListMessages returns the most recent messages of a session (ascending order). +func (s *ChatStore) ListMessages(ctx context.Context, sessionID string, limit int) ([]Message, error) { + if limit <= 0 || limit > 200 { + limit = 100 + } + rows, err := s.db.QueryContext(ctx, ` + SELECT id, session_id, role, content, created_at FROM ( + SELECT id, session_id, role, content, created_at + FROM chat_messages WHERE session_id = $1 + ORDER BY id DESC LIMIT $2 + ) recent ORDER BY id ASC + `, sessionID, limit) + if err != nil { + return nil, fmt.Errorf("list messages: %w", err) + } + defer rows.Close() + + messages := []Message{} + for rows.Next() { + var m Message + if err := rows.Scan(&m.ID, &m.SessionID, &m.Role, &m.Content, &m.CreatedAt); err != nil { + return nil, err + } + messages = append(messages, m) + } + return messages, rows.Err() +} + +// ClearMessages removes all messages of a session. +func (s *ChatStore) ClearMessages(ctx context.Context, sessionID string) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM chat_messages WHERE session_id = $1`, sessionID) + return err +} + +// SearchMessages does a substring search across the user's messages. +func (s *ChatStore) SearchMessages(ctx context.Context, userID, query string, limit int) ([]Message, error) { + if limit <= 0 || limit > 100 { + limit = 20 + } + rows, err := s.db.QueryContext(ctx, ` + SELECT id, session_id, role, content, created_at + FROM chat_messages + WHERE user_id = $1 AND content ILIKE '%' || $2 || '%' + ORDER BY id DESC LIMIT $3 + `, userID, query, limit) + if err != nil { + return nil, fmt.Errorf("search messages: %w", err) + } + defer rows.Close() + + messages := []Message{} + for rows.Next() { + var m Message + if err := rows.Scan(&m.ID, &m.SessionID, &m.Role, &m.Content, &m.CreatedAt); err != nil { + return nil, err + } + messages = append(messages, m) + } + return messages, rows.Err() +} + +// --- Helpers --- + +// truncateRunes cuts s to at most n runes, appending an ellipsis when cut. +func truncateRunes(s string, n int) string { + runes := []rune(s) + if len(runes) <= n { + return s + } + return string(runes[:n]) + "…" +} diff --git a/backend/services/gateway/ws/hub.go b/backend/services/gateway/ws/hub.go index a06355a..5b144ca 100644 --- a/backend/services/gateway/ws/hub.go +++ b/backend/services/gateway/ws/hub.go @@ -22,7 +22,7 @@ var upgrader = websocket.Upgrader{ // Orchestrator is the interface for the AI dialogue pipeline. type Orchestrator interface { - ProcessInput(ctx context.Context, sessionID, userMessage string) (<-chan model.StreamToken, *model.Intent, error) + ProcessInput(ctx context.Context, sessionID, userID, userMessage string) (<-chan model.StreamToken, *model.Intent, error) // FinishRound records the assistant reply and triggers post-round work // (e.g. memory extraction) after the response stream completes. FinishRound(sessionID, userID, assistantReply string) @@ -160,7 +160,7 @@ func (c *Client) handleChatMessage(msg WSMessage) { } // Call orchestrator - tokenCh, intent, err := c.hub.orchestrator.ProcessInput(context.Background(), msg.SessionID, msg.Content) + tokenCh, intent, err := c.hub.orchestrator.ProcessInput(context.Background(), msg.SessionID, c.UserID, msg.Content) if err != nil { c.sendJSON(WSMessage{Type: "error", Error: err.Error()}) return