From 0e98870d0cd00de50367abccb259f6328d82ef23 Mon Sep 17 00:00:00 2001 From: Jaron Swab Date: Sat, 21 Feb 2026 14:00:09 -0500 Subject: [PATCH] feat: start fresh sessions per chat --- pkg/agent/loop.go | 117 ++++++++++++++++++--- pkg/agent/loop_test.go | 224 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 327 insertions(+), 14 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index b36f4a0c4..f828673a9 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -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 diff --git a/pkg/agent/loop_test.go b/pkg/agent/loop_test.go index 4414398b1..c20b3a676 100644 --- a/pkg/agent/loop_test.go +++ b/pkg/agent/loop_test.go @@ -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 {