From 9b4eee673d6ab3bce79b3adebd58a6d2c11e874d Mon Sep 17 00:00:00 2001 From: Orion Date: Mon, 25 May 2026 16:36:58 +0800 Subject: [PATCH] feat: add subagent tool, undo system, history compaction, and TUI shortcuts - Add subagent_run tool for spawning stateless sub-agents in parallel - Implement UndoHistoryManager for multi-turn file edit rollback with /undo command - Add CompactActiveSession for session history compaction with /compact command - Wire CompactHistory hook across all LLM adapters (Anthropic, OpenAI, Genkit) - Add TUI prefix shortcuts: ! (shell), # (memory), & (background) - Add /skill slash command to list and inspect registered skills - Add integration and unit tests for compaction, subagent, and undo Co-Authored-By: Claude Opus 4.7 --- pkg/agent/auto_review.go | 1 + pkg/agent/compaction_integration_test.go | 172 ++++++++++++++++++++ pkg/agent/pool.go | 36 +++-- pkg/agent/runner.go | 83 +++++++++- pkg/agent/session_store.go | 99 ++++++++++++ pkg/agent/tools.go | 9 ++ pkg/agent/tools_subagent.go | 77 +++++++++ pkg/agent/tools_subagent_test.go | 91 +++++++++++ pkg/agent/undo_test.go | 80 ++++++++++ pkg/llm/adapter.go | 2 + pkg/llm/anthropic.go | 6 +- pkg/llm/genkit_adapter.go | 3 + pkg/llm/openai.go | 3 + pkg/tui/update_keys.go | 191 ++++++++++++++++++++++- 14 files changed, 833 insertions(+), 20 deletions(-) create mode 100644 pkg/agent/compaction_integration_test.go create mode 100644 pkg/agent/tools_subagent.go create mode 100644 pkg/agent/tools_subagent_test.go create mode 100644 pkg/agent/undo_test.go diff --git a/pkg/agent/auto_review.go b/pkg/agent/auto_review.go index 137f378..3da61f7 100644 --- a/pkg/agent/auto_review.go +++ b/pkg/agent/auto_review.go @@ -179,6 +179,7 @@ func ClassifyTool(toolName string, args any) (RiskTier, string) { case "todo", "task_create", "task_update", "task_list", "task_get", "check_background", "schedule_create", "schedule_list", "schedule_delete", + "subagent_run", // Teams "spawn_teammate", "list_teammates", "send_message", "read_inbox", "broadcast", // Protocols diff --git a/pkg/agent/compaction_integration_test.go b/pkg/agent/compaction_integration_test.go new file mode 100644 index 0000000..898b1ce --- /dev/null +++ b/pkg/agent/compaction_integration_test.go @@ -0,0 +1,172 @@ +package agent + +import ( + "context" + "fmt" + "iter" + "strings" + "testing" + + "google.golang.org/adk/model" + "google.golang.org/genai" +) + +type mockLLM struct { + generateFunc func(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] +} + +func (m *mockLLM) Name() string { + return "mock-llm" +} + +func (m *mockLLM) GenerateContent(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] { + if m.generateFunc != nil { + return m.generateFunc(ctx, req, stream) + } + return func(yield func(*model.LLMResponse, error) bool) { + yield(&model.LLMResponse{ + Content: &genai.Content{ + Role: "model", + Parts: []*genai.Part{ + {Text: "LLM-generated summary"}, + }, + }, + TurnComplete: true, + }, nil) + } +} + +func TestCompaction_LLMAndCircuitBreaker(t *testing.T) { + // Reset circuit breaker state before testing + compactionCircuitBreaker.mu.Lock() + compactionCircuitBreaker.failures = 0 + compactionCircuitBreaker.open = false + compactionCircuitBreaker.mu.Unlock() + + // 1. Create a history of 14 rounds (>12) + contents := make([]*genai.Content, 14) + for i := 0; i < 14; i++ { + role := "user" + if i%2 == 1 { + role = "model" + } + contents[i] = &genai.Content{ + Role: role, + Parts: []*genai.Part{ + {Text: fmt.Sprintf("Round message %d", i)}, + }, + } + } + + // 2. Test successful LLM summarization + mock := &mockLLM{} + compacted := CompactContents(contents, "session-test-llm", mock) + if len(compacted) != 6 { + t.Fatalf("Expected compacted length 6, got %d", len(compacted)) + } + // Verify that the second item is the system prompt with LLM summary + sysText := compacted[1].Parts[0].Text + if !strings.Contains(sysText, "LLM-generated summary") { + t.Errorf("Expected summary to contain LLM output, got: %q", sysText) + } + + // Verify circuit breaker did not trip + compactionCircuitBreaker.mu.Lock() + failures := compactionCircuitBreaker.failures + isOpen := compactionCircuitBreaker.open + compactionCircuitBreaker.mu.Unlock() + if failures != 0 || isOpen { + t.Errorf("Expected 0 failures and closed circuit breaker, got failures=%d, open=%v", failures, isOpen) + } + + // 3. Test failing LLM calls to trip circuit breaker + failMock := &mockLLM{ + generateFunc: func(ctx context.Context, req *model.LLMRequest, stream bool) iter.Seq2[*model.LLMResponse, error] { + panic("simulated LLM panic") + }, + } + + // First failure + _ = CompactContents(contents, "session-test-fail-1", failMock) + compactionCircuitBreaker.mu.Lock() + failures = compactionCircuitBreaker.failures + isOpen = compactionCircuitBreaker.open + compactionCircuitBreaker.mu.Unlock() + if failures != 1 || isOpen { + t.Errorf("Expected 1 failure, got %d, open=%v", failures, isOpen) + } + + // Second failure + _ = CompactContents(contents, "session-test-fail-2", failMock) + compactionCircuitBreaker.mu.Lock() + failures = compactionCircuitBreaker.failures + isOpen = compactionCircuitBreaker.open + compactionCircuitBreaker.mu.Unlock() + if failures != 2 || isOpen { + t.Errorf("Expected 2 failures, got %d, open=%v", failures, isOpen) + } + + // Third failure -> should trip circuit breaker + _ = CompactContents(contents, "session-test-fail-3", failMock) + compactionCircuitBreaker.mu.Lock() + failures = compactionCircuitBreaker.failures + isOpen = compactionCircuitBreaker.open + compactionCircuitBreaker.mu.Unlock() + if failures != 3 || !isOpen { + t.Errorf("Expected 3 failures and open circuit breaker, got failures=%d, open=%v", failures, isOpen) + } + + // 4. Verification that subsequent compaction bypasses LLM and falls back to truncation-only summary + compactedFallback := CompactContents(contents, "session-test-fallback", mock) + sysTextFallback := compactedFallback[1].Parts[0].Text + if !strings.Contains(sysTextFallback, "truncation-only mode") { + t.Errorf("Expected truncation-only fallback message, got: %q", sysTextFallback) + } +} + +func TestCompaction_StickyLatch(t *testing.T) { + // Reset circuit breaker state + compactionCircuitBreaker.mu.Lock() + compactionCircuitBreaker.failures = 0 + compactionCircuitBreaker.open = false + compactionCircuitBreaker.mu.Unlock() + + // 1. Create a history of 14 rounds with one sticky block + contents := make([]*genai.Content, 14) + for i := 0; i < 14; i++ { + role := "user" + if i%2 == 1 { + role = "model" + } + text := fmt.Sprintf("Round message %d", i) + if i == 5 { + text = "This is a [STICKY] instruction that must persist" + } + contents[i] = &genai.Content{ + Role: role, + Parts: []*genai.Part{ + {Text: text}, + }, + } + } + + // 2. Perform compaction + compacted := CompactContents(contents, "session-test-sticky") + // Expected: preserved prompt (1), summary (1), sticky (1), and last 4 rounds (4) = 7 + if len(compacted) != 7 { + t.Fatalf("Expected compacted length of 7, got %d", len(compacted)) + } + + // Check that the sticky block exists in the output + foundSticky := false + for _, c := range compacted { + for _, p := range c.Parts { + if strings.Contains(p.Text, "[STICKY]") { + foundSticky = true + } + } + } + if !foundSticky { + t.Error("Expected sticky content to be preserved, but it was not found") + } +} diff --git a/pkg/agent/pool.go b/pkg/agent/pool.go index 5cfb32b..d6db1b4 100644 --- a/pkg/agent/pool.go +++ b/pkg/agent/pool.go @@ -12,6 +12,7 @@ import ( "github.com/firebase/genkit/go/genkit" "google.golang.org/adk/agent" "google.golang.org/adk/agent/llmagent" + "google.golang.org/adk/model" "google.golang.org/adk/runner" "google.golang.org/adk/session" "google.golang.org/adk/tool" @@ -91,15 +92,23 @@ func TypePromptPrefix(typeName string) string { } func (ap *AgentPool) ExecuteMessage(teammate *Teammate, msg TeamMessage) (string, error) { + // Determine worktree directory + worktreePath := filepath.Join(GlobalWorktreeManager.worktreesDir, teammate.Name) + if _, err := os.Stat(worktreePath); os.IsNotExist(err) { + worktreePath, _ = os.Getwd() + } + return ap.ExecuteMessageInDir(teammate, msg, worktreePath) +} + +func (ap *AgentPool) ExecuteMessageInDir(teammate *Teammate, msg TeamMessage, dir string) (string, error) { ap.mu.Lock() subRunner, exists := ap.runners[teammate.Name] ap.mu.Unlock() if !exists { - // 1. Ensure worktree directory exists dynamically if it doesn't yet - worktreePath := filepath.Join(GlobalWorktreeManager.worktreesDir, teammate.Name) - if _, err := os.Stat(worktreePath); os.IsNotExist(err) { - _, _ = GlobalWorktreeManager.Create(teammate.Name, "") + // 1. Ensure directory exists dynamically if it doesn't yet + if dir != "" { + _ = os.MkdirAll(dir, 0755) } // 2. Setup subagent ADK Runner @@ -127,8 +136,14 @@ func (ap *AgentPool) ExecuteMessage(teammate *Teammate, msg TeamMessage) (string genkitReg := ap.GenkitRegistry ap.mu.RUnlock() - // Retrieve or build subagent adapter - subAdapter, err := llm.NewAdapter(genkitReg, prov, mod, key, base, systemPrompt, fmtFormat, runnerHooks{}) + // Retrieve or build subagent adapter with dynamic resolve hooks + var subAdapter model.LLM + hooks := runnerHooks{ + modelGetter: func() model.LLM { + return subAdapter + }, + } + subAdapter, err = llm.NewAdapter(genkitReg, prov, mod, key, base, systemPrompt, fmtFormat, hooks) if err != nil { return "", err } @@ -159,16 +174,11 @@ func (ap *AgentPool) ExecuteMessage(teammate *Teammate, msg TeamMessage) (string ap.mu.Unlock() } - // Determine worktree directory - worktreePath := filepath.Join(GlobalWorktreeManager.worktreesDir, teammate.Name) - if _, err := os.Stat(worktreePath); os.IsNotExist(err) { - worktreePath, _ = os.Getwd() - } - // 3. Execute prompt on the subRunner // Setup context with subagent name and workdir path - ctx := context.WithValue(context.Background(), WorkdirKey, worktreePath) + ctx := context.WithValue(context.Background(), WorkdirKey, dir) ctx = context.WithValue(ctx, AgentNameKey, teammate.Name) + ctx = context.WithValue(ctx, "session_id", teammate.Name+"-session") // Inject session id for history compaction userMsg := &genai.Content{ Role: "user", diff --git a/pkg/agent/runner.go b/pkg/agent/runner.go index 6c7da4e..9ab3c98 100644 --- a/pkg/agent/runner.go +++ b/pkg/agent/runner.go @@ -25,7 +25,9 @@ import ( ) // runnerHooks implements llm.AdapterHooks using the agent package's global managers. -type runnerHooks struct{} +type runnerHooks struct { + modelGetter func() model.LLM +} func (runnerHooks) NagReminder() string { if GlobalTodoManager.RoundsSinceUpdate() >= 3 { @@ -38,6 +40,18 @@ func (runnerHooks) NoteRound() { GlobalTodoManager.NoteRoundWithoutUpdate() } +func (r runnerHooks) CompactHistory(contents []*genai.Content) []*genai.Content { + sessID := GlobalLogger.sessionID + if sessID == "" { + sessID = "session-default" + } + var m model.LLM + if r.modelGetter != nil { + m = r.modelGetter() + } + return CompactContents(contents, sessID, m) +} + func buildSystemPrompt() string { builder := NewSystemPromptBuilder() return builder.Build() @@ -171,7 +185,14 @@ func NewCustomRunner(provider llm.ProviderType, modelName string, apiKey string, // 2. Create our abstract model adapter systemPrompt := buildSystemPrompt() - modelAdapter, err := llm.NewAdapter(g, provider, modelName, apiKey, baseURL, systemPrompt, apiFormat, runnerHooks{}) + var modelAdapter model.LLM + hooks := runnerHooks{ + modelGetter: func() model.LLM { + return modelAdapter + }, + } + var err error + modelAdapter, err = llm.NewAdapter(g, provider, modelName, apiKey, baseURL, systemPrompt, apiFormat, hooks) if err != nil { return nil, fmt.Errorf("failed to create model adapter: %w", err) } @@ -278,6 +299,10 @@ func (cr *CustomRunner) ModelName() string { return cr.llmModel.Name() } +func (cr *CustomRunner) GetModel() model.LLM { + return cr.llmModel +} + func (cr *CustomRunner) GetTokenUsage() int { if cr.llmModel == nil { return 0 @@ -383,8 +408,9 @@ func (cr *CustomRunner) Execute(ctx context.Context, userID, sessionID, prompt s }, } + runCtx := context.WithValue(ctx, "session_id", sessionID) runConfig := runner.WithStateDelta(nil) - events := cr.adkRunner.Run(ctx, userID, sessionID, userMsg, agent.RunConfig{ + events := cr.adkRunner.Run(runCtx, userID, sessionID, userMsg, agent.RunConfig{ StreamingMode: agent.StreamingModeSSE, }, runConfig) @@ -1018,9 +1044,58 @@ func rollbackPendingEdits() { pendingEditSnapshots.snapshots = make(map[string]string) } -// commitPendingEdits clears all snapshots after a successful turn. +// UndoGroup tracks pre-edit file states for a single conversational turn. +type UndoGroup struct { + Snapshots map[string]string +} + +// UndoHistoryManager maintains a session-wide history of UndoGroups. +type UndoHistoryManager struct { + mu sync.Mutex + history []UndoGroup +} + +func (u *UndoHistoryManager) Push(group UndoGroup) { + u.mu.Lock() + defer u.mu.Unlock() + u.history = append(u.history, group) +} + +func (u *UndoHistoryManager) PopAndUndo() (int, error) { + u.mu.Lock() + defer u.mu.Unlock() + + if len(u.history) == 0 { + return 0, fmt.Errorf("no changes to undo") + } + + last := u.history[len(u.history)-1] + u.history = u.history[:len(u.history)-1] + + count := 0 + for path, content := range last.Snapshots { + if content == "" { + _ = os.Remove(path) + } else { + _ = os.WriteFile(path, []byte(content), 0644) + } + count++ + } + + return count, nil +} + +var GlobalUndoManager = &UndoHistoryManager{} + +// commitPendingEdits pushes the current turn's snapshots onto the GlobalUndoManager and clears active pending snapshots. func commitPendingEdits() { pendingEditSnapshots.mu.Lock() defer pendingEditSnapshots.mu.Unlock() + + if len(pendingEditSnapshots.snapshots) > 0 { + GlobalUndoManager.Push(UndoGroup{ + Snapshots: pendingEditSnapshots.snapshots, + }) + } pendingEditSnapshots.snapshots = make(map[string]string) } diff --git a/pkg/agent/session_store.go b/pkg/agent/session_store.go index 41778bb..3efa5de 100644 --- a/pkg/agent/session_store.go +++ b/pkg/agent/session_store.go @@ -11,7 +11,9 @@ import ( "sync" "time" + "google.golang.org/adk/model" "google.golang.org/adk/session" + "google.golang.org/genai" ) // SerializedSession represents the full schema serialized to disk for a session. @@ -564,3 +566,100 @@ func (s *SerializedSession) ValidateResume() []string { // Ensure interface matching var _ session.Service = (*PersistentSessionService)(nil) + +// CompactActiveSession compacts the history events of the active session. +func (s *PersistentSessionService) CompactActiveSession(ctx context.Context, sessionID string, llm ...model.LLM) error { + s.mu.Lock() + defer s.mu.Unlock() + + // 1. Retrieve session from memory delegate + getResp, err := s.delegate.Get(ctx, &session.GetRequest{SessionID: sessionID}) + if err != nil { + return err + } + sess := getResp.Session + + // 2. Extract events + var events []*session.Event + if sess.Events() != nil { + for ev := range sess.Events().All() { + events = append(events, ev) + } + } + + if len(events) <= 12 { + return fmt.Errorf("session history too short to compact (has %d event(s), minimum 12 required)", len(events)) + } + + // 3. Convert events to genai.Content + var contents []*genai.Content + for _, ev := range events { + if ev == nil { + continue + } + if ev.Content != nil { + contents = append(contents, ev.Content) + } + if ev.LLMResponse.Content != nil { + contents = append(contents, ev.LLMResponse.Content) + } + } + + // 4. Run compaction + compacted := CompactContents(contents, sessionID, llm...) + + // 5. Convert compacted Contents back to session.Events + var compactedEvents []*session.Event + for _, c := range compacted { + if c == nil { + continue + } + ev := &session.Event{ + Timestamp: time.Now(), + } + if c.Role == "user" || c.Role == "tool" { + ev.Content = c + } else { + ev.LLMResponse = model.LLMResponse{ + Content: c, + TurnComplete: true, + } + } + compactedEvents = append(compactedEvents, ev) + } + + // 6. Delete old session and recreate it with compacted events in the delegate + _ = s.delegate.Delete(ctx, &session.DeleteRequest{SessionID: sessionID}) + + // Recreate + stateMap := make(map[string]any) + if sess.State() != nil { + for k, v := range sess.State().All() { + stateMap[k] = v + } + } + + res, err := s.delegate.Create(ctx, &session.CreateRequest{ + AppName: sess.AppName(), + UserID: sess.UserID(), + SessionID: sessionID, + State: stateMap, + }) + if err != nil { + return fmt.Errorf("failed to recreate session during compaction: %w", err) + } + + // Append compacted events + for _, ev := range compactedEvents { + if err := s.delegate.AppendEvent(ctx, res.Session, ev); err != nil { + return fmt.Errorf("failed to append compacted event during compaction: %w", err) + } + } + + // 7. Save session to disk + s.mu.Unlock() // unlock temporarily because SaveSession will lock it + err = s.SaveSession(ctx, res.Session) + s.mu.Lock() // relock for defer unlock + + return err +} diff --git a/pkg/agent/tools.go b/pkg/agent/tools.go index 86e8330..669b838 100644 --- a/pkg/agent/tools.go +++ b/pkg/agent/tools.go @@ -486,6 +486,14 @@ func GetSWETools() ([]tool.Tool, error) { return nil, err } + subagentRunTool, err := functiontool.New(functiontool.Config{ + Name: "subagent_run", + Description: "Spawn a specialized subagent to execute a specific sub-task or research question in parallel. The subagent operates within its own clean context window and returns a text summary of its findings.", + }, SubagentRunHandler) + if err != nil { + return nil, err + } + // s21 LSP Tools — load user LSP server config if available if cfg, err := config.LoadConfig(); err == nil && len(cfg.LSPServers) > 0 { servers := make([]LSPServerConfig, len(cfg.LSPServers)) @@ -546,6 +554,7 @@ func GetSWETools() ([]tool.Tool, error) { taskCreateTool, taskUpdateTool, taskListTool, taskGetTool, bgRunTool, bgCheckTool, schCreateTool, schListTool, schDeleteTool, + subagentRunTool, // s15 spawnTeammateTool, listTeammatesTool, sendMessageTool, readInboxTool, broadcastTool, diff --git a/pkg/agent/tools_subagent.go b/pkg/agent/tools_subagent.go new file mode 100644 index 0000000..06742c3 --- /dev/null +++ b/pkg/agent/tools_subagent.go @@ -0,0 +1,77 @@ +package agent + +import ( + "fmt" + + "github.com/google/uuid" + "google.golang.org/adk/tool" +) + +// SubagentRunArgs represents arguments for the subagent_run tool. +type SubagentRunArgs struct { + Prompt string `json:"prompt" description:"The specific instruction/task for the subagent to perform."` + ModelName string `json:"model_name,omitempty" description:"Optional. The model to use for the subagent (e.g. 'deepseek-chat', 'claude-3-5-haiku'). If omitted, defaults to the current session model."` +} + +// SubagentRunResult represents the output returned by the subagent. +type SubagentRunResult struct { + Summary string `json:"summary" description:"The summary of findings and results returned by the subagent."` +} + +// SubagentRunHandler dynamically spawns a stateless subagent in the current workspace directory. +func SubagentRunHandler(ctx tool.Context, args SubagentRunArgs) (SubagentRunResult, error) { + subagentName := "subagent-" + uuid.New().String()[:8] + + teammate := &Teammate{ + Name: subagentName, + Type: "executor", + SystemPrompt: "You are a specialized subagent. Complete the task as instructed. Return a clear and concise summary of your findings and results. Bypassing confirmation is NOT allowed; all edits and commands will prompt the user for permission.", + } + + GlobalAgentPool.mu.RLock() + originalModel := GlobalAgentPool.ModelName + GlobalAgentPool.mu.RUnlock() + + if args.ModelName != "" { + GlobalAgentPool.mu.Lock() + GlobalAgentPool.ModelName = args.ModelName + GlobalAgentPool.mu.Unlock() + + defer func() { + GlobalAgentPool.mu.Lock() + GlobalAgentPool.ModelName = originalModel + GlobalAgentPool.mu.Unlock() + }() + } + + parentWorkdir := getWorkdir(ctx) + msg := TeamMessage{ + Sender: "parent", + Content: args.Prompt, + } + + // Send status update to TUI + ToolBridge.Send(ToolStatus{ + Name: subagentName, + Args: args, + Running: true, + }) + + summary, err := GlobalAgentPool.ExecuteMessageInDir(teammate, msg, parentWorkdir) + + ToolBridge.Send(ToolStatus{ + Name: subagentName, + Args: args, + Running: false, + Success: err == nil, + Error: err, + }) + + if err != nil { + return SubagentRunResult{}, fmt.Errorf("subagent failed: %w", err) + } + + return SubagentRunResult{ + Summary: summary, + }, nil +} diff --git a/pkg/agent/tools_subagent_test.go b/pkg/agent/tools_subagent_test.go new file mode 100644 index 0000000..f60065f --- /dev/null +++ b/pkg/agent/tools_subagent_test.go @@ -0,0 +1,91 @@ +package agent + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "iroha/pkg/llm" +) + +func TestSubagentRunHandler_Success(t *testing.T) { + // 1. Setup mock SSE response server + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + + // Yield simple choices delta + _, _ = fmt.Fprint(w, "data: {\"choices\":[{\"delta\":{\"content\":\"mock subagent analysis completed successfully\"}}]}\n\n") + _, _ = fmt.Fprint(w, "data: [DONE]\n\n") + })) + defer server.Close() + + // 2. Configure GlobalAgentPool for testing + GlobalAgentPool.mu.Lock() + origProvider := GlobalAgentPool.Provider + origModelName := GlobalAgentPool.ModelName + origAPIKey := GlobalAgentPool.APIKey + origBaseURL := GlobalAgentPool.BaseURL + origAPIFormat := GlobalAgentPool.APIFormat + + GlobalAgentPool.Provider = llm.ProviderOpenAI + GlobalAgentPool.ModelName = "mock-gpt" + GlobalAgentPool.APIKey = "mock-key" + GlobalAgentPool.BaseURL = server.URL + GlobalAgentPool.APIFormat = llm.APIFormatOpenAI + GlobalAgentPool.mu.Unlock() + + defer func() { + GlobalAgentPool.mu.Lock() + GlobalAgentPool.Provider = origProvider + GlobalAgentPool.ModelName = origModelName + GlobalAgentPool.APIKey = origAPIKey + GlobalAgentPool.BaseURL = origBaseURL + GlobalAgentPool.APIFormat = origAPIFormat + GlobalAgentPool.mu.Unlock() + }() + + // Drain any leftover ToolBridge status updates + for len(ToolBridge.StatusChan) > 0 { + <-ToolBridge.StatusChan + } + + // 3. Build context & args + tempCwd := t.TempDir() + stdCtx := context.WithValue(context.Background(), WorkdirKey, tempCwd) + ctx := &mockToolContext{Context: stdCtx} + + args := SubagentRunArgs{ + Prompt: "Test run subagent please", + ModelName: "", // Use active default + } + + // 4. Run handler + res, err := SubagentRunHandler(ctx, args) + if err != nil { + t.Fatalf("SubagentRunHandler failed: %v", err) + } + + // 5. Verify results + if !strings.Contains(res.Summary, "mock subagent analysis completed successfully") { + t.Errorf("expected summary to contain mock subagent output, got %q", res.Summary) + } + + // Drain and verify status messages sent to TUI + s1 := <-ToolBridge.StatusChan + s2 := <-ToolBridge.StatusChan + + if !s1.Running { + t.Errorf("expected first status message to show running=true") + } + if s2.Running { + t.Errorf("expected second status message to show running=false") + } + if !s2.Success { + t.Errorf("expected second status message to show success=true") + } +} diff --git a/pkg/agent/undo_test.go b/pkg/agent/undo_test.go new file mode 100644 index 0000000..9c5b155 --- /dev/null +++ b/pkg/agent/undo_test.go @@ -0,0 +1,80 @@ +package agent + +import ( + "os" + "path/filepath" + "testing" +) + +func TestUndoHistoryManager_PushAndUndo(t *testing.T) { + tempDir := t.TempDir() + file1 := filepath.Join(tempDir, "test1.txt") + file2 := filepath.Join(tempDir, "test2.txt") + + // 1. Setup initial state: file1 exists, file2 does not exist + err := os.WriteFile(file1, []byte("original content 1"), 0644) + if err != nil { + t.Fatalf("Failed to write initial file: %v", err) + } + + // 2. Perform mock modifications: + // - file1 gets overwritten + // - file2 gets created + // Capture snapshots for undo (representing the original state before modification) + snapshots := map[string]string{ + file1: "original content 1", + file2: "", // Empty string means the file did not exist + } + + // Make changes + err = os.WriteFile(file1, []byte("modified content 1"), 0644) + if err != nil { + t.Fatalf("Failed to write modified file: %v", err) + } + err = os.WriteFile(file2, []byte("new file content 2"), 0644) + if err != nil { + t.Fatalf("Failed to write new file: %v", err) + } + + // 3. Register undo group + mgr := &UndoHistoryManager{} + mgr.Push(UndoGroup{ + Snapshots: snapshots, + }) + + // Verify manager history length + mgr.mu.Lock() + histLen := len(mgr.history) + mgr.mu.Unlock() + if histLen != 1 { + t.Fatalf("Expected history length 1, got %d", histLen) + } + + // 4. Perform Undo + count, err := mgr.PopAndUndo() + if err != nil { + t.Fatalf("PopAndUndo failed: %v", err) + } + if count != 2 { + t.Errorf("Expected 2 reverted files, got %d", count) + } + + // 5. Verify files are back to original state + content1, err := os.ReadFile(file1) + if err != nil { + t.Fatalf("Failed to read file1: %v", err) + } + if string(content1) != "original content 1" { + t.Errorf("Expected file1 to be 'original content 1', got %q", string(content1)) + } + + if _, err := os.Stat(file2); !os.IsNotExist(err) { + t.Errorf("Expected file2 to be deleted, but it still exists") + } + + // 6. Verify second undo returns error + _, err = mgr.PopAndUndo() + if err == nil { + t.Error("Expected error on empty undo history, got nil") + } +} diff --git a/pkg/llm/adapter.go b/pkg/llm/adapter.go index 77e160b..143dc4d 100644 --- a/pkg/llm/adapter.go +++ b/pkg/llm/adapter.go @@ -6,6 +6,7 @@ import ( "github.com/firebase/genkit/go/genkit" "google.golang.org/adk/model" + "google.golang.org/genai" ) // ProviderType represents the LLM provider @@ -33,6 +34,7 @@ const ( type AdapterHooks interface { NagReminder() string NoteRound() + CompactHistory(contents []*genai.Content) []*genai.Content } // TokenTracker tracks cumulative token usage across adapter calls. diff --git a/pkg/llm/anthropic.go b/pkg/llm/anthropic.go index 8e49938..8d69123 100644 --- a/pkg/llm/anthropic.go +++ b/pkg/llm/anthropic.go @@ -171,9 +171,13 @@ func (a *AnthropicAdapter) GenerateContent(ctx context.Context, req *model.LLMRe } } } + compactedContents := req.Contents + if a.hooks != nil { + compactedContents = a.hooks.CompactHistory(req.Contents) + } // Convert genai Contents to Anthropic messages - messages, err := convertToAnthropicMessages(req.Contents) + messages, err := convertToAnthropicMessages(compactedContents) if err != nil { yield(nil, err) return diff --git a/pkg/llm/genkit_adapter.go b/pkg/llm/genkit_adapter.go index 292b8eb..7409828 100644 --- a/pkg/llm/genkit_adapter.go +++ b/pkg/llm/genkit_adapter.go @@ -54,6 +54,9 @@ func (m *GenkitModelAdapter) GenerateContent(ctx context.Context, req *model.LLM } compactedContents := req.Contents + if m.hooks != nil { + compactedContents = m.hooks.CompactHistory(req.Contents) + } // Inject Nag Reminder if triggered if m.hooks != nil { diff --git a/pkg/llm/openai.go b/pkg/llm/openai.go index 8ad6337..e8e7c20 100644 --- a/pkg/llm/openai.go +++ b/pkg/llm/openai.go @@ -126,6 +126,9 @@ func (g *OpenAICompatibleAdapter) GenerateContent(ctx context.Context, req *mode } compactedContents := req.Contents + if g.hooks != nil { + compactedContents = g.hooks.CompactHistory(req.Contents) + } if g.hooks != nil { nagMsg := g.hooks.NagReminder() diff --git a/pkg/tui/update_keys.go b/pkg/tui/update_keys.go index 7d8f6ff..e22336a 100644 --- a/pkg/tui/update_keys.go +++ b/pkg/tui/update_keys.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "os" + "regexp" "strings" "time" @@ -322,9 +323,98 @@ func (m Model) handleKeyMsg(msg tea.KeyMsg) (Model, tea.Cmd, bool) { m.SlashMenuItems = nil // Fall through to execute the command } + inputVal := m.TextArea.Value() + + // Intercept Prefix Shortcuts + if strings.HasPrefix(inputVal, "!") { + m.HistoryManager.Add(inputVal) + userLog := StyleUserMsg.Render("> " + inputVal) + m.TextArea.SetValue("") + m.TextArea.SetHeight(2) + + cmdStr := strings.TrimSpace(inputVal[1:]) + var replyLog string + if cmdStr == "" { + replyLog = StyleToolError.Render("[error] Please specify a shell command: !") + } else { + res, err := agent.ShellRunHandler(nil, agent.ShellRunArgs{Command: cmdStr}) + if err != nil { + replyLog = StyleToolError.Render(fmt.Sprintf("[error] %v", err)) + } else { + if res.ExitCode != 0 { + replyLog = StyleToolError.Render(res.Output) + } else { + replyLog = res.Output + } + } + } + m.History = append(m.History, userLog, replyLog) + m.Viewport.SetContent(m.renderViewportContent()) + m.Viewport.GotoBottom() + return m, nil, true + } - inputVal := strings.TrimSpace(m.TextArea.Value()) - if inputVal == "" { + if strings.HasPrefix(inputVal, "#") { + m.HistoryManager.Add(inputVal) + userLog := StyleUserMsg.Render("> " + inputVal) + m.TextArea.SetValue("") + m.TextArea.SetHeight(2) + + memVal := strings.TrimSpace(inputVal[1:]) + var replyLog string + if memVal == "" { + replyLog = StyleToolError.Render("[error] Please specify memory content: #") + } else { + words := strings.Fields(memVal) + name := "user_preference" + if len(words) > 0 { + name = strings.ToLower(words[0]) + if len(words) > 1 { + name += "_" + strings.ToLower(words[1]) + } + } + reg := regexp.MustCompile(`[^a-z0-9_]`) + name = reg.ReplaceAllString(name, "") + + _, err := agent.MemorySaveHandler(nil, agent.MemorySaveArgs{ + Name: name, + Description: "User preference saved via shortcut", + Type: "user", + Content: memVal, + }) + if err != nil { + replyLog = StyleToolError.Render(fmt.Sprintf("[error] %v", err)) + } else { + replyLog = StyleToolSuccess.Render(fmt.Sprintf("Memory saved successfully as %q: %s", name, memVal)) + } + } + m.History = append(m.History, userLog, replyLog) + m.Viewport.SetContent(m.renderViewportContent()) + m.Viewport.GotoBottom() + return m, nil, true + } + + if strings.HasPrefix(inputVal, "&") { + m.HistoryManager.Add(inputVal) + userLog := StyleUserMsg.Render("> " + inputVal) + m.TextArea.SetValue("") + m.TextArea.SetHeight(2) + + cmdStr := strings.TrimSpace(inputVal[1:]) + var replyLog string + if cmdStr == "" { + replyLog = StyleToolError.Render("[error] Please specify a command: &") + } else { + res, err := agent.BackgroundRunHandler(nil, agent.BackgroundRunArgs{Command: cmdStr}) + if err != nil { + replyLog = StyleToolError.Render(fmt.Sprintf("[error] %v", err)) + } else { + replyLog = StyleToolSuccess.Render(res.Message) + } + } + m.History = append(m.History, userLog, replyLog) + m.Viewport.SetContent(m.renderViewportContent()) + m.Viewport.GotoBottom() return m, nil, true } @@ -855,5 +945,102 @@ func (m Model) handleSlashCommand(inputVal string) (Model, tea.Cmd, bool) { return m, nil, true } + if cmdName == "/compact" { + m.HistoryManager.Add(inputVal) + userLog := StyleUserMsg.Render("> " + inputVal) + m.TextArea.SetValue("") + m.TextArea.SetHeight(2) + + var replyLog string + if agent.GlobalSessionService == nil { + replyLog = StyleToolError.Render("[error] Session service not initialized") + } else { + var err error + if m.Runner != nil { + err = agent.GlobalSessionService.CompactActiveSession(context.Background(), m.SessionID, m.Runner.GetModel()) + } else { + err = agent.GlobalSessionService.CompactActiveSession(context.Background(), m.SessionID) + } + if err != nil { + replyLog = StyleToolError.Render(fmt.Sprintf("[error] Compaction failed: %v", err)) + } else { + replyLog = StyleToolSuccess.Render("Session history compacted successfully! Older rounds summarized to reclaim context window space.") + } + } + + m.History = append(m.History, userLog, replyLog) + m.Viewport.SetContent(m.renderViewportContent()) + m.Viewport.GotoBottom() + return m, nil, true + } + + if cmdName == "/undo" { + m.HistoryManager.Add(inputVal) + userLog := StyleUserMsg.Render("> " + inputVal) + m.TextArea.SetValue("") + m.TextArea.SetHeight(2) + + count, err := agent.GlobalUndoManager.PopAndUndo() + var replyLog string + if err != nil { + replyLog = StyleToolError.Render(fmt.Sprintf("[error] %v", err)) + } else { + replyLog = StyleToolSuccess.Render(fmt.Sprintf("Undo successful! Reverted %d file(s) changed in the last turn.", count)) + } + + m.History = append(m.History, userLog, replyLog) + m.Viewport.SetContent(m.renderViewportContent()) + m.Viewport.GotoBottom() + return m, nil, true + } + + if cmdName == "/skill" { + m.HistoryManager.Add(inputVal) + userLog := StyleUserMsg.Render("> " + inputVal) + m.TextArea.SetValue("") + m.TextArea.SetHeight(2) + + var replyLog string + if len(parts) < 2 { + skills := agent.GlobalSkillManager.AllSkills() + if len(skills) == 0 { + replyLog = StyleKeyHelp.Render("No skills registered. Place a skill under ~/.iroha/skills//") + } else { + var sb strings.Builder + sb.WriteString(StyleKeyActive.Render("Available Skills:") + "\n") + for i, s := range skills { + sb.WriteString(fmt.Sprintf(" %d. \x1b[1;36m%s\x1b[0m (%s) - %s\n", i+1, s.Name, s.ID, s.Description)) + if len(s.Triggers) > 0 { + sb.WriteString(fmt.Sprintf(" Triggers: %s\n", strings.Join(s.Triggers, ", "))) + } + } + replyLog = sb.String() + } + } else { + skillID := parts[1] + skill := agent.GlobalSkillManager.GetSkillByID(skillID) + if skill == nil { + replyLog = StyleToolError.Render(fmt.Sprintf("[error] Skill %q not found", skillID)) + } else { + instructions, err := agent.LoadInstructions(skill) + if err != nil { + replyLog = StyleToolError.Render(fmt.Sprintf("[error] Failed to load instructions: %v", err)) + } else { + var sb strings.Builder + sb.WriteString(StyleKeyActive.Render(fmt.Sprintf("Skill instructions for %s (%s):", skill.Name, skill.ID)) + "\n") + sb.WriteString(strings.Repeat("─", 72) + "\n") + sb.WriteString(instructions + "\n") + sb.WriteString(strings.Repeat("─", 72) + "\n") + replyLog = sb.String() + } + } + } + + m.History = append(m.History, userLog, replyLog) + m.Viewport.SetContent(m.renderViewportContent()) + m.Viewport.GotoBottom() + return m, nil, true + } + return m, nil, false }