feat: start fresh sessions per chat
This commit is contained in:
Jaron Swab
parent
ab720f67cd
commit
0e98870d0c
2 changed files with 327 additions and 14 deletions
|
|
@ -30,14 +30,47 @@ import (
|
|||
)
|
||||
|
||||
type AgentLoop struct {
|
||||
bus *bus.MessageBus
|
||||
cfg *config.Config
|
||||
registry *AgentRegistry
|
||||
state *state.Manager
|
||||
running atomic.Bool
|
||||
summarizing sync.Map
|
||||
fallback *providers.FallbackChain
|
||||
channelManager *channels.Manager
|
||||
bus *bus.MessageBus
|
||||
cfg *config.Config
|
||||
registry *AgentRegistry
|
||||
state *state.Manager
|
||||
running atomic.Bool
|
||||
summarizing sync.Map
|
||||
fallback *providers.FallbackChain
|
||||
channelManager *channels.Manager
|
||||
sessionOverride *sessionOverrideStore
|
||||
}
|
||||
|
||||
type sessionOverrideStore struct {
|
||||
mu sync.RWMutex
|
||||
overrides map[string]string
|
||||
}
|
||||
|
||||
func newSessionOverrideStore() *sessionOverrideStore {
|
||||
return &sessionOverrideStore{overrides: make(map[string]string)}
|
||||
}
|
||||
|
||||
func (s *sessionOverrideStore) Get(key string) (string, bool) {
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
value, ok := s.overrides[key]
|
||||
return value, ok
|
||||
}
|
||||
|
||||
func (s *sessionOverrideStore) Set(key, value string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
s.overrides[key] = value
|
||||
}
|
||||
|
||||
func (s *sessionOverrideStore) Delete(key string) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
delete(s.overrides, key)
|
||||
}
|
||||
|
||||
func buildSessionOverrideKey(channel, chatID, agentID string) string {
|
||||
return channel + ":" + chatID + ":" + agentID
|
||||
}
|
||||
|
||||
// processOptions configures how a message is processed
|
||||
|
|
@ -70,12 +103,13 @@ func NewAgentLoop(cfg *config.Config, msgBus *bus.MessageBus, provider providers
|
|||
}
|
||||
|
||||
return &AgentLoop{
|
||||
bus: msgBus,
|
||||
cfg: cfg,
|
||||
registry: registry,
|
||||
state: stateManager,
|
||||
summarizing: sync.Map{},
|
||||
fallback: fallbackChain,
|
||||
bus: msgBus,
|
||||
cfg: cfg,
|
||||
registry: registry,
|
||||
state: stateManager,
|
||||
summarizing: sync.Map{},
|
||||
fallback: fallbackChain,
|
||||
sessionOverride: newSessionOverrideStore(),
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -312,6 +346,10 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage)
|
|||
sessionKey = msg.SessionKey
|
||||
}
|
||||
|
||||
if override, ok := al.sessionOverride.Get(buildSessionOverrideKey(msg.Channel, msg.ChatID, agent.ID)); ok {
|
||||
sessionKey = override
|
||||
}
|
||||
|
||||
logger.InfoCF("agent", "Routed message",
|
||||
map[string]any{
|
||||
"agent_id": agent.ID,
|
||||
|
|
@ -386,6 +424,47 @@ func (al *AgentLoop) processSystemMessage(ctx context.Context, msg bus.InboundMe
|
|||
})
|
||||
}
|
||||
|
||||
func (al *AgentLoop) startNewSessionForMessage(msg bus.InboundMessage) (string, error) {
|
||||
if constants.IsInternalChannel(msg.Channel) {
|
||||
return "", fmt.Errorf("new session commands are not supported in internal channels")
|
||||
}
|
||||
|
||||
route := al.registry.ResolveRoute(routing.RouteInput{
|
||||
Channel: msg.Channel,
|
||||
AccountID: msg.Metadata["account_id"],
|
||||
Peer: extractPeer(msg),
|
||||
ParentPeer: extractParentPeer(msg),
|
||||
GuildID: msg.Metadata["guild_id"],
|
||||
TeamID: msg.Metadata["team_id"],
|
||||
})
|
||||
|
||||
agent, ok := al.registry.GetAgent(route.AgentID)
|
||||
if !ok {
|
||||
agent = al.registry.GetDefaultAgent()
|
||||
}
|
||||
if agent == nil {
|
||||
return "", fmt.Errorf("no agent available for new session")
|
||||
}
|
||||
|
||||
sessionKey := route.SessionKey
|
||||
if msg.SessionKey != "" && strings.HasPrefix(msg.SessionKey, "agent:") {
|
||||
sessionKey = msg.SessionKey
|
||||
}
|
||||
|
||||
if sessionKey != "" {
|
||||
if err := agent.Sessions.Save(sessionKey); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
newSessionKey := fmt.Sprintf("agent:%s:new:%d", agent.ID, time.Now().UnixNano())
|
||||
agent.Sessions.GetOrCreate(newSessionKey)
|
||||
overrideKey := buildSessionOverrideKey(msg.Channel, msg.ChatID, agent.ID)
|
||||
al.sessionOverride.Set(overrideKey, newSessionKey)
|
||||
|
||||
return "Starting a new conversation...", nil
|
||||
}
|
||||
|
||||
// runAgentLoop is the core message processing logic.
|
||||
func (al *AgentLoop) runAgentLoop(ctx context.Context, agent *AgentInstance, opts processOptions) (string, error) {
|
||||
// 0. Record last channel for heartbeat notifications (skip internal channels)
|
||||
|
|
@ -1043,6 +1122,16 @@ func (al *AgentLoop) handleCommand(ctx context.Context, msg bus.InboundMessage)
|
|||
args := parts[1:]
|
||||
|
||||
switch cmd {
|
||||
case "/new", "/clear":
|
||||
if len(args) > 0 {
|
||||
return fmt.Sprintf("Usage: %s", cmd), true
|
||||
}
|
||||
response, err := al.startNewSessionForMessage(msg)
|
||||
if err != nil {
|
||||
return fmt.Sprintf("Error: %v", err), true
|
||||
}
|
||||
return response, true
|
||||
|
||||
case "/show":
|
||||
if len(args) < 1 {
|
||||
return "Usage: /show [model|channel|agents]", true
|
||||
|
|
|
|||
|
|
@ -5,15 +5,41 @@ import (
|
|||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/sipeed/picoclaw/pkg/bus"
|
||||
"github.com/sipeed/picoclaw/pkg/config"
|
||||
"github.com/sipeed/picoclaw/pkg/providers"
|
||||
"github.com/sipeed/picoclaw/pkg/routing"
|
||||
"github.com/sipeed/picoclaw/pkg/tools"
|
||||
)
|
||||
|
||||
func newTestAgentLoop(t *testing.T, provider providers.LLMProvider) (*AgentLoop, string) {
|
||||
t.Helper()
|
||||
|
||||
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
||||
if err != nil {
|
||||
t.Fatalf("Failed to create temp dir: %v", err)
|
||||
}
|
||||
|
||||
cfg := &config.Config{
|
||||
Agents: config.AgentsConfig{
|
||||
Defaults: config.AgentDefaults{
|
||||
Workspace: tmpDir,
|
||||
Model: "test-model",
|
||||
MaxTokens: 4096,
|
||||
MaxToolIterations: 10,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
msgBus := bus.NewMessageBus()
|
||||
al := NewAgentLoop(cfg, msgBus, provider)
|
||||
return al, tmpDir
|
||||
}
|
||||
|
||||
func TestRecordLastChannel(t *testing.T) {
|
||||
// Create temp workspace
|
||||
tmpDir, err := os.MkdirTemp("", "agent-test-*")
|
||||
|
|
@ -343,6 +369,204 @@ func TestAgentLoop_Stop(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestNewSessionCommand_SuccessAndIsolation(t *testing.T) {
|
||||
provider := &simpleMockProvider{response: "OK"}
|
||||
al, tmpDir := newTestAgentLoop(t, provider)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
t.Fatal("No default agent found")
|
||||
}
|
||||
|
||||
helper := testHelper{al: al}
|
||||
ctx := context.Background()
|
||||
msg := bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "hello",
|
||||
}
|
||||
_ = helper.executeAndGetResponse(t, ctx, msg)
|
||||
|
||||
oldKey := strings.ToLower(routing.BuildAgentMainSessionKey(agent.ID))
|
||||
oldHistory := agent.Sessions.GetHistory(oldKey)
|
||||
if len(oldHistory) == 0 {
|
||||
t.Fatalf("Expected old session history to be populated")
|
||||
}
|
||||
|
||||
response := helper.executeAndGetResponse(t, ctx, bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "/new",
|
||||
})
|
||||
if response != "Starting a new conversation..." {
|
||||
t.Fatalf("Expected confirmation response, got %q", response)
|
||||
}
|
||||
|
||||
overrideKey := buildSessionOverrideKey("test", "chat1", agent.ID)
|
||||
newKey, ok := al.sessionOverride.Get(overrideKey)
|
||||
if !ok {
|
||||
t.Fatalf("Expected session override to be set")
|
||||
}
|
||||
if newKey == oldKey {
|
||||
t.Fatalf("Expected new session key to differ from old session key")
|
||||
}
|
||||
newHistory := agent.Sessions.GetHistory(newKey)
|
||||
if len(newHistory) != 0 {
|
||||
t.Fatalf("Expected new session to start empty, got %d messages", len(newHistory))
|
||||
}
|
||||
|
||||
_ = helper.executeAndGetResponse(t, ctx, bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "next",
|
||||
})
|
||||
|
||||
newHistory = agent.Sessions.GetHistory(newKey)
|
||||
if len(newHistory) != 2 {
|
||||
t.Fatalf("Expected new session history to have 2 messages, got %d", len(newHistory))
|
||||
}
|
||||
oldHistoryAfter := agent.Sessions.GetHistory(oldKey)
|
||||
if len(oldHistoryAfter) != len(oldHistory) {
|
||||
t.Fatalf("Expected old session history to remain unchanged")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSessionCommand_RejectsArgs(t *testing.T) {
|
||||
provider := &simpleMockProvider{response: "OK"}
|
||||
al, tmpDir := newTestAgentLoop(t, provider)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
t.Fatal("No default agent found")
|
||||
}
|
||||
|
||||
helper := testHelper{al: al}
|
||||
ctx := context.Background()
|
||||
response := helper.executeAndGetResponse(t, ctx, bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "/new extra",
|
||||
})
|
||||
if response != "Usage: /new" {
|
||||
t.Fatalf("Expected usage response, got %q", response)
|
||||
}
|
||||
|
||||
overrideKey := buildSessionOverrideKey("test", "chat1", agent.ID)
|
||||
if _, ok := al.sessionOverride.Get(overrideKey); ok {
|
||||
t.Fatalf("Expected no session override to be set")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSessionCommand_PreservesOldSessionOnDisk(t *testing.T) {
|
||||
provider := &simpleMockProvider{response: "OK"}
|
||||
al, tmpDir := newTestAgentLoop(t, provider)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
t.Fatal("No default agent found")
|
||||
}
|
||||
|
||||
helper := testHelper{al: al}
|
||||
ctx := context.Background()
|
||||
_ = helper.executeAndGetResponse(t, ctx, bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "hello",
|
||||
})
|
||||
|
||||
oldKey := strings.ToLower(routing.BuildAgentMainSessionKey(agent.ID))
|
||||
filename := strings.ReplaceAll(oldKey, ":", "_") + ".json"
|
||||
sessionsDir := filepath.Join(agent.Workspace, "sessions")
|
||||
filePath := filepath.Join(sessionsDir, filename)
|
||||
_ = os.Remove(filePath)
|
||||
|
||||
_ = helper.executeAndGetResponse(t, ctx, bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "/new",
|
||||
})
|
||||
|
||||
data, err := os.ReadFile(filePath)
|
||||
if err != nil {
|
||||
t.Fatalf("Expected session file to exist after /new: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(data), "hello") {
|
||||
t.Fatalf("Expected session file to contain prior message")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewSessionCommand_NoCrossChatImpact(t *testing.T) {
|
||||
provider := &simpleMockProvider{response: "OK"}
|
||||
al, tmpDir := newTestAgentLoop(t, provider)
|
||||
defer os.RemoveAll(tmpDir)
|
||||
|
||||
agent := al.registry.GetDefaultAgent()
|
||||
if agent == nil {
|
||||
t.Fatal("No default agent found")
|
||||
}
|
||||
|
||||
helper := testHelper{al: al}
|
||||
ctx := context.Background()
|
||||
response := helper.executeAndGetResponse(t, ctx, bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user1",
|
||||
ChatID: "chat1",
|
||||
Content: "/new",
|
||||
Metadata: map[string]string{
|
||||
"peer_kind": "group",
|
||||
"peer_id": "group1",
|
||||
},
|
||||
})
|
||||
if response != "Starting a new conversation..." {
|
||||
t.Fatalf("Expected confirmation response, got %q", response)
|
||||
}
|
||||
|
||||
overrideKey := buildSessionOverrideKey("test", "chat1", agent.ID)
|
||||
newKey, ok := al.sessionOverride.Get(overrideKey)
|
||||
if !ok {
|
||||
t.Fatalf("Expected session override to be set")
|
||||
}
|
||||
|
||||
_ = helper.executeAndGetResponse(t, ctx, bus.InboundMessage{
|
||||
Channel: "test",
|
||||
SenderID: "user2",
|
||||
ChatID: "chat2",
|
||||
Content: "hello",
|
||||
Metadata: map[string]string{
|
||||
"peer_kind": "group",
|
||||
"peer_id": "group2",
|
||||
},
|
||||
})
|
||||
|
||||
chat2Key := strings.ToLower(routing.BuildAgentPeerSessionKey(routing.SessionKeyParams{
|
||||
AgentID: agent.ID,
|
||||
Channel: "test",
|
||||
Peer: &routing.RoutePeer{
|
||||
Kind: "group",
|
||||
ID: "group2",
|
||||
},
|
||||
DMScope: routing.DMScopeMain,
|
||||
}))
|
||||
chat2History := agent.Sessions.GetHistory(chat2Key)
|
||||
if len(chat2History) != 2 {
|
||||
t.Fatalf("Expected chat2 session history to have 2 messages, got %d", len(chat2History))
|
||||
}
|
||||
|
||||
newHistory := agent.Sessions.GetHistory(newKey)
|
||||
if len(newHistory) != 0 {
|
||||
t.Fatalf("Expected chat1 override session to remain empty, got %d messages", len(newHistory))
|
||||
}
|
||||
}
|
||||
|
||||
// Mock implementations for testing
|
||||
|
||||
type simpleMockProvider struct {
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue