From af7bc26a582f2ec7918a633b917d4e53d8c423ba Mon Sep 17 00:00:00 2001 From: Anton Bogdanovich <27antonb@gmail.com> Date: Thu, 7 May 2026 14:27:43 -0700 Subject: [PATCH] Add non-destructive session reset command --- pkg/agent/agent_command.go | 22 +++++ pkg/agent/agent_message.go | 19 ++-- pkg/agent/agent_steering.go | 2 +- pkg/agent/agent_test.go | 172 ++++++++++++++++++++++++++++++++++ pkg/agent/dispatch_request.go | 18 ++-- pkg/agent/session_reset.go | 60 ++++++++++++ pkg/commands/builtin.go | 1 + pkg/commands/cmd_reset.go | 38 ++++++++ pkg/commands/runtime.go | 1 + pkg/state/state.go | 65 +++++++++++++ pkg/state/state_test.go | 30 ++++++ 11 files changed, 411 insertions(+), 17 deletions(-) create mode 100644 pkg/agent/session_reset.go create mode 100644 pkg/commands/cmd_reset.go diff --git a/pkg/agent/agent_command.go b/pkg/agent/agent_command.go index ae0293d71..3f9e3cb99 100644 --- a/pkg/agent/agent_command.go +++ b/pkg/agent/agent_command.go @@ -335,6 +335,28 @@ func (al *AgentLoop) buildCommandsRuntime( return al.contextManager.Clear(ctx, opts.SessionKey) } + rt.ResetSession = func(clear bool) (string, error) { + if opts == nil { + return "", fmt.Errorf("process options not available") + } + routeSessionKey := strings.TrimSpace(opts.Dispatch.RouteSessionKey) + if routeSessionKey == "" { + return "", fmt.Errorf("route session key not available") + } + if clear { + return "", al.clearSessionOverride(routeSessionKey) + } + + nextSessionKey := buildResetSessionKey(agent.ID, routeSessionKey) + if nextSessionKey == "" { + return "", fmt.Errorf("failed to allocate reset session key") + } + if err := al.setSessionOverride(routeSessionKey, nextSessionKey); err != nil { + return "", err + } + return nextSessionKey, nil + } + rt.AskSideQuestion = func(ctx context.Context, question string) (string, error) { return al.askSideQuestion(ctx, agent, opts, question) } diff --git a/pkg/agent/agent_message.go b/pkg/agent/agent_message.go index 96b0b0817..5b82b25ff 100644 --- a/pkg/agent/agent_message.go +++ b/pkg/agent/agent_message.go @@ -27,7 +27,7 @@ func (al *AgentLoop) buildContinuationTarget(msg bus.InboundMessage) (*continuat allocation := al.allocateRouteSession(route, msg) return &continuationTarget{ - SessionKey: resolveScopeKey(allocation.SessionKey, msg.SessionKey), + SessionKey: al.resolveEffectiveSessionKey(allocation.SessionKey, msg.SessionKey), Channel: msg.Channel, ChatID: msg.ChatID, }, nil @@ -146,7 +146,7 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) // Resolve session key from the route allocation, while preserving explicit // agent-scoped keys supplied by the caller. - scopeKey := resolveScopeKey(allocation.SessionKey, msg.SessionKey) + scopeKey := al.resolveEffectiveSessionKey(allocation.SessionKey, msg.SessionKey) sessionKey := scopeKey // Reset message-tool state for this round so we don't skip publishing due to a previous round. @@ -169,13 +169,14 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) opts := processOptions{ Dispatch: DispatchRequest{ - SessionKey: sessionKey, - SessionAliases: buildSessionAliases(sessionKey, append(allocation.SessionAliases, msg.SessionKey)...), - InboundContext: cloneInboundContext(&msg.Context), - RouteResult: cloneResolvedRoute(&route), - SessionScope: session.CloneScope(&allocation.Scope), - UserMessage: msg.Content, - Media: append([]string(nil), msg.Media...), + RouteSessionKey: allocation.SessionKey, + SessionKey: sessionKey, + SessionAliases: buildSessionAliases(sessionKey, sessionAliasCandidates(allocation.SessionKey, sessionKey, allocation.SessionAliases, msg.SessionKey)...), + InboundContext: cloneInboundContext(&msg.Context), + RouteResult: cloneResolvedRoute(&route), + SessionScope: session.CloneScope(&allocation.Scope), + UserMessage: msg.Content, + Media: append([]string(nil), msg.Media...), }, SenderID: msg.SenderID, SenderDisplayName: msg.Sender.DisplayName, diff --git a/pkg/agent/agent_steering.go b/pkg/agent/agent_steering.go index 9b136e7cd..700b0a165 100644 --- a/pkg/agent/agent_steering.go +++ b/pkg/agent/agent_steering.go @@ -108,5 +108,5 @@ func (al *AgentLoop) resolveSteeringTarget(msg bus.InboundMessage) (string, stri } allocation := al.allocateRouteSession(route, msg) - return resolveScopeKey(allocation.SessionKey, msg.SessionKey), agent.ID, true + return al.resolveEffectiveSessionKey(allocation.SessionKey, msg.SessionKey), agent.ID, true } diff --git a/pkg/agent/agent_test.go b/pkg/agent/agent_test.go index a75919912..3343042b2 100644 --- a/pkg/agent/agent_test.go +++ b/pkg/agent/agent_test.go @@ -681,6 +681,178 @@ func TestHandleCommand_UseCommandRejectsUnknownSkill(t *testing.T) { } } +func TestProcessMessage_ResetCommandStartsFreshSession(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + msgBus := bus.NewMessageBus() + provider := &recordingProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + + initialMsg := testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "first message", + }) + if _, err := al.processMessage(context.Background(), initialMsg); err != nil { + t.Fatalf("initial processMessage() error = %v", err) + } + + route, _, err := al.resolveMessageRoute(initialMsg) + if err != nil { + t.Fatalf("resolveMessageRoute() error = %v", err) + } + allocation := al.allocateRouteSession(route, initialMsg) + routeSessionKey := allocation.SessionKey + defaultAgent := al.GetRegistry().GetDefaultAgent() + originalHistory := defaultAgent.Sessions.GetHistory(routeSessionKey) + if len(originalHistory) == 0 { + t.Fatal("expected initial history in routed session") + } + + resetReply, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "/reset", + })) + if err != nil { + t.Fatalf("reset processMessage() error = %v", err) + } + if !strings.Contains(resetReply, "Started a fresh session") { + t.Fatalf("reset reply = %q, want fresh-session confirmation", resetReply) + } + + overrideSessionKey := al.getSessionOverride(routeSessionKey) + if overrideSessionKey == "" || overrideSessionKey == routeSessionKey { + t.Fatalf("override session key = %q, want distinct session key", overrideSessionKey) + } + + if _, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "second message", + })); err != nil { + t.Fatalf("second processMessage() error = %v", err) + } + + gotOriginalHistory := defaultAgent.Sessions.GetHistory(routeSessionKey) + if len(gotOriginalHistory) != len(originalHistory) { + t.Fatalf("original history len = %d, want preserved len %d", len(gotOriginalHistory), len(originalHistory)) + } + + resetHistory := defaultAgent.Sessions.GetHistory(overrideSessionKey) + if len(resetHistory) == 0 { + t.Fatal("expected fresh session history after /reset") + } +} + +func TestProcessMessage_ResetClearRestoresDefaultSession(t *testing.T) { + tmpDir := t.TempDir() + cfg := &config.Config{ + Agents: config.AgentsConfig{ + Defaults: config.AgentDefaults{ + Workspace: tmpDir, + ModelName: "test-model", + MaxTokens: 4096, + MaxToolIterations: 10, + }, + }, + } + msgBus := bus.NewMessageBus() + provider := &recordingProvider{} + al := NewAgentLoop(cfg, msgBus, provider) + + initialMsg := testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "before reset", + }) + if _, err := al.processMessage(context.Background(), initialMsg); err != nil { + t.Fatalf("initial processMessage() error = %v", err) + } + + route, _, err := al.resolveMessageRoute(initialMsg) + if err != nil { + t.Fatalf("resolveMessageRoute() error = %v", err) + } + allocation := al.allocateRouteSession(route, initialMsg) + routeSessionKey := allocation.SessionKey + defaultAgent := al.GetRegistry().GetDefaultAgent() + originalHistoryLen := len(defaultAgent.Sessions.GetHistory(routeSessionKey)) + + if _, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "/reset", + })); err != nil { + t.Fatalf("reset processMessage() error = %v", err) + } + overrideSessionKey := al.getSessionOverride(routeSessionKey) + if overrideSessionKey == "" { + t.Fatal("expected override session key after /reset") + } + + if _, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "during reset", + })); err != nil { + t.Fatalf("during-reset processMessage() error = %v", err) + } + overrideHistoryLen := len(defaultAgent.Sessions.GetHistory(overrideSessionKey)) + if overrideHistoryLen == 0 { + t.Fatal("expected override history after reset message") + } + + clearReply, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "/reset clear", + })) + if err != nil { + t.Fatalf("reset clear processMessage() error = %v", err) + } + if !strings.Contains(clearReply, "Soft reset cleared") { + t.Fatalf("reset clear reply = %q, want clear confirmation", clearReply) + } + if got := al.getSessionOverride(routeSessionKey); got != "" { + t.Fatalf("override session key after clear = %q, want empty", got) + } + + if _, err := al.processMessage(context.Background(), testInboundMessage(bus.InboundMessage{ + Channel: "telegram", + SenderID: "telegram:123", + ChatID: "chat-1", + Content: "after clear", + })); err != nil { + t.Fatalf("after-clear processMessage() error = %v", err) + } + + gotOriginalHistoryLen := len(defaultAgent.Sessions.GetHistory(routeSessionKey)) + if gotOriginalHistoryLen <= originalHistoryLen { + t.Fatalf("original history len = %d, want > %d after reset clear", gotOriginalHistoryLen, originalHistoryLen) + } + gotOverrideHistoryLen := len(defaultAgent.Sessions.GetHistory(overrideSessionKey)) + if gotOverrideHistoryLen != overrideHistoryLen { + t.Fatalf("override history len = %d, want preserved len %d after reset clear", gotOverrideHistoryLen, overrideHistoryLen) + } +} + func TestProcessMessage_UseCommandArmsSkillForNextMessage(t *testing.T) { tmpDir := t.TempDir() skillDir := filepath.Join(tmpDir, "skills", "shell") diff --git a/pkg/agent/dispatch_request.go b/pkg/agent/dispatch_request.go index cb54264d6..e71ca7cb4 100644 --- a/pkg/agent/dispatch_request.go +++ b/pkg/agent/dispatch_request.go @@ -11,13 +11,14 @@ import ( // DispatchRequest is the normalized runtime input passed into the agent loop // after routing and session allocation have completed. type DispatchRequest struct { - SessionKey string - SessionAliases []string - InboundContext *bus.InboundContext - RouteResult *routing.ResolvedRoute - SessionScope *session.SessionScope - UserMessage string - Media []string + RouteSessionKey string + SessionKey string + SessionAliases []string + InboundContext *bus.InboundContext + RouteResult *routing.ResolvedRoute + SessionScope *session.SessionScope + UserMessage string + Media []string } func (r DispatchRequest) Channel() string { @@ -63,6 +64,9 @@ func normalizeProcessOptionsInPlace(opts *processOptions) { } func normalizeProcessOptions(opts processOptions) processOptions { + if opts.Dispatch.RouteSessionKey == "" { + opts.Dispatch.RouteSessionKey = strings.TrimSpace(opts.SessionKey) + } if opts.Dispatch.SessionKey == "" { opts.Dispatch.SessionKey = strings.TrimSpace(opts.SessionKey) } diff --git a/pkg/agent/session_reset.go b/pkg/agent/session_reset.go new file mode 100644 index 000000000..3880b7c6a --- /dev/null +++ b/pkg/agent/session_reset.go @@ -0,0 +1,60 @@ +package agent + +import ( + "fmt" + "strings" + "time" + + "github.com/sipeed/picoclaw/pkg/session" +) + +func buildResetSessionKey(agentID, routeSessionKey string) string { + alias := fmt.Sprintf( + "agent:%s:reset:%s:%d", + strings.ToLower(strings.TrimSpace(agentID)), + strings.ToLower(strings.TrimSpace(routeSessionKey)), + time.Now().UnixNano(), + ) + return session.BuildOpaqueSessionKey(alias) +} + +func (al *AgentLoop) getSessionOverride(routeSessionKey string) string { + if al == nil || al.state == nil { + return "" + } + return al.state.GetSessionOverride(routeSessionKey) +} + +func (al *AgentLoop) setSessionOverride(routeSessionKey, sessionKey string) error { + if al == nil || al.state == nil { + return fmt.Errorf("state manager not initialized") + } + return al.state.SetSessionOverride(routeSessionKey, sessionKey) +} + +func (al *AgentLoop) clearSessionOverride(routeSessionKey string) error { + if al == nil || al.state == nil { + return fmt.Errorf("state manager not initialized") + } + return al.state.ClearSessionOverride(routeSessionKey) +} + +func (al *AgentLoop) resolveEffectiveSessionKey(routeSessionKey, msgSessionKey string) string { + if isExplicitSessionKey(msgSessionKey) { + return msgSessionKey + } + if override := al.getSessionOverride(routeSessionKey); override != "" { + return override + } + return routeSessionKey +} + +func sessionAliasCandidates(routeSessionKey, effectiveSessionKey string, routeAliases []string, msgSessionKey string) []string { + if isExplicitSessionKey(msgSessionKey) { + return []string{msgSessionKey} + } + if strings.TrimSpace(routeSessionKey) == strings.TrimSpace(effectiveSessionKey) { + return routeAliases + } + return nil +} diff --git a/pkg/commands/builtin.go b/pkg/commands/builtin.go index e268812a0..ee67196dd 100644 --- a/pkg/commands/builtin.go +++ b/pkg/commands/builtin.go @@ -15,6 +15,7 @@ func BuiltinDefinitions() []Definition { btwCommand(), switchCommand(), checkCommand(), + resetCommand(), clearCommand(), contextCommand(), subagentsCommand(), diff --git a/pkg/commands/cmd_reset.go b/pkg/commands/cmd_reset.go new file mode 100644 index 000000000..8a3bee532 --- /dev/null +++ b/pkg/commands/cmd_reset.go @@ -0,0 +1,38 @@ +package commands + +import ( + "context" + "fmt" + "strings" +) + +func resetCommand() Definition { + return Definition{ + Name: "reset", + Description: "Start a fresh session without deleting stored history", + Usage: "/reset [clear|off]", + Handler: func(_ context.Context, req Request, rt *Runtime) error { + if rt == nil || rt.ResetSession == nil { + return req.Reply(unavailableMsg) + } + + arg := strings.ToLower(strings.TrimSpace(nthToken(req.Text, 1))) + clearOverride := arg == "clear" || arg == "off" + if arg != "" && !clearOverride { + return req.Reply("Usage: /reset [clear|off]") + } + + sessionKey, err := rt.ResetSession(clearOverride) + if err != nil { + return req.Reply("Failed to reset session: " + err.Error()) + } + if clearOverride { + return req.Reply("Soft reset cleared. Future messages will use the default routed session again.") + } + return req.Reply(fmt.Sprintf( + "Started a fresh session. Previous history was preserved. New session key: %s", + sessionKey, + )) + }, + } +} diff --git a/pkg/commands/runtime.go b/pkg/commands/runtime.go index b0327c863..f8596c2e3 100644 --- a/pkg/commands/runtime.go +++ b/pkg/commands/runtime.go @@ -59,6 +59,7 @@ type Runtime struct { GetContextStats func() *ContextStats SwitchModel func(value string) (oldModel string, err error) SwitchChannel func(value string) error + ResetSession func(clear bool) (sessionKey string, err error) ClearHistory func() error ReloadConfig func() error StopActiveTurn func() (StopResult, error) diff --git a/pkg/state/state.go b/pkg/state/state.go index 5da7bbde1..2f1489b2f 100644 --- a/pkg/state/state.go +++ b/pkg/state/state.go @@ -6,6 +6,7 @@ import ( "log" "os" "path/filepath" + "strings" "sync" "time" @@ -21,6 +22,10 @@ type State struct { // LastChatID is the last chat ID used for communication LastChatID string `json:"last_chat_id,omitempty"` + // SessionOverrides maps the default routed session key for a conversation + // onto an explicit replacement session key created by a soft reset. + SessionOverrides map[string]string `json:"session_overrides,omitempty"` + // Timestamp is the last time this state was updated Timestamp time.Time `json:"timestamp"` } @@ -108,6 +113,66 @@ func (sm *Manager) SetLastChatID(chatID string) error { return nil } +// SetSessionOverride persists a replacement session key for a routed session. +func (sm *Manager) SetSessionOverride(routeSessionKey, sessionKey string) error { + sm.mu.Lock() + defer sm.mu.Unlock() + + routeSessionKey = strings.TrimSpace(routeSessionKey) + sessionKey = strings.TrimSpace(sessionKey) + if routeSessionKey == "" || sessionKey == "" { + return fmt.Errorf("route session key and session key are required") + } + + if sm.state.SessionOverrides == nil { + sm.state.SessionOverrides = make(map[string]string) + } + sm.state.SessionOverrides[routeSessionKey] = sessionKey + sm.state.Timestamp = time.Now() + + if err := sm.saveAtomic(); err != nil { + return fmt.Errorf("failed to save state atomically: %w", err) + } + + return nil +} + +// ClearSessionOverride removes a previously persisted replacement session key. +func (sm *Manager) ClearSessionOverride(routeSessionKey string) error { + sm.mu.Lock() + defer sm.mu.Unlock() + + routeSessionKey = strings.TrimSpace(routeSessionKey) + if routeSessionKey == "" { + return fmt.Errorf("route session key is required") + } + if len(sm.state.SessionOverrides) == 0 { + return nil + } + + delete(sm.state.SessionOverrides, routeSessionKey) + if len(sm.state.SessionOverrides) == 0 { + sm.state.SessionOverrides = nil + } + sm.state.Timestamp = time.Now() + + if err := sm.saveAtomic(); err != nil { + return fmt.Errorf("failed to save state atomically: %w", err) + } + + return nil +} + +// GetSessionOverride returns the replacement session key for a routed session. +func (sm *Manager) GetSessionOverride(routeSessionKey string) string { + sm.mu.RLock() + defer sm.mu.RUnlock() + if len(sm.state.SessionOverrides) == 0 { + return "" + } + return sm.state.SessionOverrides[strings.TrimSpace(routeSessionKey)] +} + // GetLastChannel returns the last channel from the state. func (sm *Manager) GetLastChannel() string { sm.mu.RLock() diff --git a/pkg/state/state_test.go b/pkg/state/state_test.go index 3924e5533..da21db7d4 100644 --- a/pkg/state/state_test.go +++ b/pkg/state/state_test.go @@ -216,6 +216,36 @@ func TestNewManager_EmptyWorkspace(t *testing.T) { } } +func TestSessionOverridesPersist(t *testing.T) { + tmpDir := t.TempDir() + + sm := NewManager(tmpDir) + if err := sm.SetSessionOverride("route-session", "reset-session"); err != nil { + t.Fatalf("SetSessionOverride failed: %v", err) + } + + if got := sm.GetSessionOverride("route-session"); got != "reset-session" { + t.Fatalf("GetSessionOverride() = %q, want %q", got, "reset-session") + } + + sm2 := NewManager(tmpDir) + if got := sm2.GetSessionOverride("route-session"); got != "reset-session" { + t.Fatalf("persisted GetSessionOverride() = %q, want %q", got, "reset-session") + } + + if err := sm2.ClearSessionOverride("route-session"); err != nil { + t.Fatalf("ClearSessionOverride failed: %v", err) + } + if got := sm2.GetSessionOverride("route-session"); got != "" { + t.Fatalf("GetSessionOverride() after clear = %q, want empty", got) + } + + sm3 := NewManager(tmpDir) + if got := sm3.GetSessionOverride("route-session"); got != "" { + t.Fatalf("persisted GetSessionOverride() after clear = %q, want empty", got) + } +} + func TestNewManager_MkdirFailureDoesNotCrash(t *testing.T) { if os.Getenv("BE_CRASHER") == "1" { tmpDir := os.Getenv("CRASH_DIR")