339 lines
7.7 KiB
Go
339 lines
7.7 KiB
Go
package agent
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
type SessionStore struct {
|
|
mu sync.RWMutex
|
|
path string
|
|
sessions map[string]*Session
|
|
order []string
|
|
}
|
|
|
|
func NewSessionStore(dataDir string) (*SessionStore, error) {
|
|
if err := os.MkdirAll(dataDir, 0o755); err != nil {
|
|
return nil, err
|
|
}
|
|
path := filepath.Join(dataDir, "sessions.json")
|
|
store := &SessionStore{
|
|
path: path,
|
|
sessions: make(map[string]*Session),
|
|
order: make([]string, 0),
|
|
}
|
|
_ = store.load()
|
|
return store, nil
|
|
}
|
|
|
|
func (s *SessionStore) load() error {
|
|
data, err := os.ReadFile(s.path)
|
|
if err != nil {
|
|
if errors.Is(err, os.ErrNotExist) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|
|
if len(strings.TrimSpace(string(data))) == 0 {
|
|
return nil
|
|
}
|
|
var sessions map[string]*Session
|
|
if err := json.Unmarshal(data, &sessions); err != nil {
|
|
return err
|
|
}
|
|
for id, session := range sessions {
|
|
if session == nil || session.ID == "" {
|
|
continue
|
|
}
|
|
s.sessions[id] = session
|
|
s.order = append(s.order, id)
|
|
}
|
|
s.sortOrder()
|
|
return nil
|
|
}
|
|
|
|
func (s *SessionStore) saveLocked() error {
|
|
data, err := json.MarshalIndent(s.sessions, "", " ")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tmp := s.path + ".tmp"
|
|
if err := os.WriteFile(tmp, data, 0o600); err != nil {
|
|
return err
|
|
}
|
|
return os.Rename(tmp, s.path)
|
|
}
|
|
|
|
func (s *SessionStore) List() []SessionSummary {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
return s.listLocked("", 0, 0)
|
|
}
|
|
|
|
// Search 按关键词过滤会话（标题或任意消息内容，忽略大小写）。
|
|
func (s *SessionStore) Search(query string) []SessionSummary {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
return s.listLocked(query, 0, 0)
|
|
}
|
|
|
|
// listLocked 返回会话摘要列表；query 非空时按关键词过滤，limit>0 时截断。
|
|
// 必须持读锁调用。
|
|
func (s *SessionStore) listLocked(query string, limit, offset int) []SessionSummary {
|
|
query = strings.ToLower(strings.TrimSpace(query))
|
|
out := make([]SessionSummary, 0, len(s.order))
|
|
skip := 0
|
|
for _, id := range s.order {
|
|
session, ok := s.sessions[id]
|
|
if !ok {
|
|
continue
|
|
}
|
|
if query != "" && !sessionMatches(session, query) {
|
|
continue
|
|
}
|
|
if offset > 0 && skip < offset {
|
|
skip++
|
|
continue
|
|
}
|
|
if limit > 0 && len(out) >= limit {
|
|
break
|
|
}
|
|
out = append(out, summarize(session))
|
|
}
|
|
return out
|
|
}
|
|
|
|
func sessionMatches(session *Session, query string) bool {
|
|
if strings.Contains(strings.ToLower(session.Title), query) {
|
|
return true
|
|
}
|
|
for _, message := range session.Messages {
|
|
if message.Content != nil && strings.Contains(strings.ToLower(*message.Content), query) {
|
|
return true
|
|
}
|
|
for _, call := range message.ToolCalls {
|
|
if strings.Contains(strings.ToLower(call.Function.Name), query) ||
|
|
strings.Contains(strings.ToLower(call.Function.Arguments), query) {
|
|
return true
|
|
}
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
// Messages 分页返回会话消息（按时间正序），offset 从最新一条往前数，
|
|
// 即 offset=0 返回最新 limit 条；返回 has_more 表示还有更早的消息。
|
|
func (s *SessionStore) Messages(id string, limit, offset int) ([]Message, bool, bool) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
session, ok := s.sessions[id]
|
|
if !ok {
|
|
return nil, false, false
|
|
}
|
|
total := len(session.Messages)
|
|
if limit <= 0 {
|
|
limit = 50
|
|
}
|
|
if offset < 0 {
|
|
offset = 0
|
|
}
|
|
start := total - offset - limit
|
|
if start < 0 {
|
|
start = 0
|
|
}
|
|
messages := make([]Message, 0, limit)
|
|
for i := start; i < total-offset; i++ {
|
|
messages = append(messages, cloneMessage(session.Messages[i]))
|
|
}
|
|
hasMore := start > 0
|
|
return messages, true, hasMore
|
|
}
|
|
|
|
// Clear 清空会话消息（保留会话本身）。
|
|
func (s *SessionStore) Clear(id string) bool {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
session, ok := s.sessions[id]
|
|
if !ok {
|
|
return false
|
|
}
|
|
session.Messages = make([]Message, 0)
|
|
session.UpdatedAt = time.Now()
|
|
_ = s.saveLocked()
|
|
return true
|
|
}
|
|
|
|
// TotalMessages 返回会话消息总数（用于分页计算）。
|
|
func (s *SessionStore) TotalMessages(id string) (int, bool) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
session, ok := s.sessions[id]
|
|
if !ok {
|
|
return 0, false
|
|
}
|
|
return len(session.Messages), true
|
|
}
|
|
|
|
func cloneMessage(message Message) Message {
|
|
copy := message
|
|
if message.Content != nil {
|
|
content := *message.Content
|
|
copy.Content = &content
|
|
}
|
|
copy.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
|
|
return copy
|
|
}
|
|
|
|
func (s *SessionStore) Get(id string) (*Session, bool) {
|
|
s.mu.RLock()
|
|
defer s.mu.RUnlock()
|
|
session, ok := s.sessions[id]
|
|
return session, ok
|
|
}
|
|
|
|
func (s *SessionStore) Create() *Session {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
session := &Session{
|
|
ID: newID(),
|
|
Title: "新对话",
|
|
CreatedAt: time.Now(),
|
|
UpdatedAt: time.Now(),
|
|
Messages: make([]Message, 0),
|
|
}
|
|
s.sessions[session.ID] = session
|
|
s.order = append([]string{session.ID}, s.order...)
|
|
_ = s.saveLocked()
|
|
return session
|
|
}
|
|
|
|
func (s *SessionStore) Delete(id string) bool {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if _, ok := s.sessions[id]; !ok {
|
|
return false
|
|
}
|
|
delete(s.sessions, id)
|
|
for i, item := range s.order {
|
|
if item == id {
|
|
s.order = append(s.order[:i], s.order[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
_ = s.saveLocked()
|
|
return true
|
|
}
|
|
|
|
func (s *SessionStore) Save(session *Session) {
|
|
session.UpdatedAt = time.Now()
|
|
s.mu.Lock()
|
|
// 存储克隆而非原对象：会话对象可能仍在 Agent 循环中被修改，
|
|
// 直接存指针会导致 saveLocked 序列化时与其他会话的写入产生数据竞争。
|
|
s.sessions[session.ID] = cloneSession(session)
|
|
s.order = append([]string{session.ID}, removeString(s.order, session.ID)...)
|
|
err := s.saveLocked()
|
|
s.mu.Unlock()
|
|
if err != nil {
|
|
// A failed save should not break an in-memory conversation.
|
|
return
|
|
}
|
|
}
|
|
|
|
func (s *SessionStore) sortOrder() {
|
|
sort.SliceStable(s.order, func(i, j int) bool {
|
|
a, aOK := s.sessions[s.order[i]]
|
|
b, bOK := s.sessions[s.order[j]]
|
|
if !aOK || !bOK {
|
|
return aOK && !bOK
|
|
}
|
|
return a.UpdatedAt.After(b.UpdatedAt)
|
|
})
|
|
}
|
|
|
|
func summarize(session *Session) SessionSummary {
|
|
summary := SessionSummary{
|
|
ID: session.ID,
|
|
Title: session.Title,
|
|
CreatedAt: session.CreatedAt,
|
|
UpdatedAt: session.UpdatedAt,
|
|
MessageCount: len(session.Messages),
|
|
}
|
|
for i := len(session.Messages) - 1; i >= 0; i-- {
|
|
message := session.Messages[i]
|
|
if message.Role != "user" && message.Role != "assistant" {
|
|
continue
|
|
}
|
|
if message.Content != nil {
|
|
summary.Preview = truncateRunes(*message.Content, 120)
|
|
break
|
|
}
|
|
}
|
|
return summary
|
|
}
|
|
|
|
func newID() string {
|
|
buf := make([]byte, 12)
|
|
_, _ = rand.Read(buf)
|
|
return hex.EncodeToString(buf)
|
|
}
|
|
|
|
// cloneSession 深拷贝会话对象，保证持久化对象在存储后不再被外部修改，
|
|
// 从而消除 Agent 循环写入与 saveLocked 序列化之间的数据竞争。
|
|
func cloneSession(session *Session) *Session {
|
|
if session == nil {
|
|
return nil
|
|
}
|
|
clone := *session
|
|
clone.Messages = make([]Message, len(session.Messages))
|
|
for i, message := range session.Messages {
|
|
copy := message
|
|
if message.Content != nil {
|
|
content := *message.Content
|
|
copy.Content = &content
|
|
}
|
|
copy.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
|
|
clone.Messages[i] = copy
|
|
}
|
|
return &clone
|
|
}
|
|
|
|
func removeString(items []string, target string) []string {
|
|
result := make([]string, 0, len(items))
|
|
for _, item := range items {
|
|
if item != target {
|
|
result = append(result, item)
|
|
}
|
|
}
|
|
return result
|
|
}
|
|
|
|
func truncateRunes(value string, max int) string {
|
|
runes := []rune(value)
|
|
if len(runes) <= max {
|
|
return value
|
|
}
|
|
return string(runes[:max]) + "..."
|
|
}
|
|
|
|
func firstRunes(value string, max int) string {
|
|
runes := []rune(strings.TrimSpace(value))
|
|
if len(runes) <= max {
|
|
return string(runes)
|
|
}
|
|
return string(runes[:max])
|
|
}
|
|
|
|
func countRunes(value string) int {
|
|
return utf8.RuneCountInString(value)
|
|
}
|