mirror of
https://github.com/tiennm99/goclaw.git
synced 2026-10-13 10:15:48 +00:00
Duy /zuey/
43049d3b3f
* feat(api): unblock trace follow and provider reconnect * test(cron): stabilize run log execution test * feat(http): add CLI P6 backend endpoints
108 lines
3.4 KiB
Go
108 lines
3.4 KiB
Go
//go:build sqlite || sqliteonly
|
|
|
|
package sqlitestore
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"maps"
|
|
"strconv"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
|
|
"github.com/nextlevelbuilder/goclaw/internal/providers"
|
|
"github.com/nextlevelbuilder/goclaw/internal/store"
|
|
)
|
|
|
|
func (s *SQLiteSessionStore) BranchSession(ctx context.Context, sourceKey string, opts store.SessionBranchOpts) (*store.SessionData, int, error) {
|
|
if opts.NewKey == "" || opts.UpToIndex < 0 {
|
|
return nil, 0, store.ErrInvalidSessionBranch
|
|
}
|
|
if s.Get(ctx, opts.NewKey) != nil {
|
|
return nil, 0, store.ErrSessionAlreadyExists
|
|
}
|
|
|
|
source := s.Get(ctx, sourceKey)
|
|
if source == nil {
|
|
return nil, 0, store.ErrSessionNotFound
|
|
}
|
|
|
|
s.mu.RLock()
|
|
sourceSnapshot := *source
|
|
sourceMsgs := make([]providers.Message, len(source.Messages))
|
|
copy(sourceMsgs, source.Messages)
|
|
sourceSnapshot.Messages = sourceMsgs
|
|
sourceMeta := make(map[string]string, len(source.Metadata))
|
|
maps.Copy(sourceMeta, source.Metadata)
|
|
sourceSnapshot.Metadata = sourceMeta
|
|
s.mu.RUnlock()
|
|
|
|
if opts.UpToIndex > len(sourceSnapshot.Messages) {
|
|
return nil, 0, store.ErrInvalidSessionBranch
|
|
}
|
|
|
|
now := time.Now().UTC()
|
|
copied := make([]providers.Message, opts.UpToIndex)
|
|
copy(copied, sourceSnapshot.Messages[:opts.UpToIndex])
|
|
meta := make(map[string]string, len(sourceSnapshot.Metadata)+len(opts.Metadata)+3)
|
|
maps.Copy(meta, sourceSnapshot.Metadata)
|
|
maps.Copy(meta, opts.Metadata)
|
|
meta["branched_from"] = sourceKey
|
|
meta["branched_from_index"] = strconv.Itoa(opts.UpToIndex)
|
|
meta["branched_at"] = now.Format(time.RFC3339Nano)
|
|
|
|
branch := &store.SessionData{
|
|
Key: opts.NewKey,
|
|
Messages: copied,
|
|
Summary: sourceSnapshot.Summary,
|
|
Created: now,
|
|
Updated: now,
|
|
AgentUUID: sourceSnapshot.AgentUUID,
|
|
UserID: sourceSnapshot.UserID,
|
|
Model: sourceSnapshot.Model,
|
|
Provider: sourceSnapshot.Provider,
|
|
Channel: "branch",
|
|
Label: opts.Label,
|
|
Metadata: meta,
|
|
ContextWindow: sourceSnapshot.ContextWindow,
|
|
}
|
|
if sourceSnapshot.TeamID != nil {
|
|
teamID := *sourceSnapshot.TeamID
|
|
branch.TeamID = &teamID
|
|
}
|
|
|
|
msgsJSON, _ := json.Marshal(branch.Messages)
|
|
metaJSON := []byte("{}")
|
|
if len(branch.Metadata) > 0 {
|
|
metaJSON, _ = json.Marshal(branch.Metadata)
|
|
}
|
|
|
|
res, err := s.db.ExecContext(ctx,
|
|
`INSERT INTO sessions (id, session_key, messages, summary, model, provider, channel,
|
|
input_tokens, output_tokens, compaction_count,
|
|
memory_flush_compaction_count, memory_flush_at,
|
|
label, spawned_by, spawn_depth, agent_id, user_id, metadata, updated_at, team_id, tenant_id, created_at)
|
|
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)
|
|
ON CONFLICT(session_key, tenant_id) DO NOTHING`,
|
|
uuid.Must(uuid.NewV7()), branch.Key, msgsJSON,
|
|
nilStr(branch.Summary), nilStr(branch.Model), nilStr(branch.Provider), nilStr(branch.Channel),
|
|
branch.InputTokens, branch.OutputTokens, branch.CompactionCount,
|
|
branch.MemoryFlushCompactionCount, branch.MemoryFlushAt,
|
|
nilStr(branch.Label), nilStr(branch.SpawnedBy), branch.SpawnDepth,
|
|
nilSessionUUID(branch.AgentUUID), nilStr(branch.UserID), metaJSON, branch.Updated,
|
|
branch.TeamID, tenantIDForInsert(ctx), branch.Created,
|
|
)
|
|
if err != nil {
|
|
return nil, 0, err
|
|
}
|
|
if n, _ := res.RowsAffected(); n == 0 {
|
|
return nil, 0, store.ErrSessionAlreadyExists
|
|
}
|
|
|
|
s.mu.Lock()
|
|
s.cache[sessionCacheKey(ctx, opts.NewKey)] = branch
|
|
s.mu.Unlock()
|
|
return branch, len(copied), nil
|
|
}
|