diff --git a/CHANGELOG.md b/CHANGELOG.md index 9920e8f..5a8caaa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,39 @@ from this file's releases; see CONTRACT.md for the stability policy. ## v0.3.0 — unreleased +### Added +- Permission engine (`permission`): modes `plan|build|edit|yolo|auto`, + per-tool allow/deny/ask rules from `~/.rakit/settings.json` and + `/.rakit/settings.json`, per-call decisions, project/session + rules, and typed literal rules like `Bash(git push:*)` / `Edit(src/**)`. +- Hooks (`hooks`): Claude-Code-compatible `hooks.json` (global + workspace) + for `PreToolUse`, `PostToolUse`, `UserPromptSubmit`, `Stop`, `SessionStart`, + `SessionEnd`, `PreCompact`. Exit code 2 blocks with a reason; JSON output + can modify tool input or append context. +- Checkpoints (`checkpoint`): files touched by a mutating batch are snapshotted + (plus a `git stash create` tree) and restorable via `/undo` and + `Engine.UndoLast`. +- Resume journal: interrupted approval/question turns are recovered on startup + and resumable via `/resume` / `Engine.ResumeTurn`; tools are never + re-executed automatically. +- Real plan mode tools `enter_plan_mode` / `exit_plan_mode`; plans are written + to `/.rakit/plans/`. +- `ToolSpec` classification metadata (read-only / destructive / approval / + side-effect scope / timeout / max output) is the single source for the + scheduler, permissions and the agent loop. +- LLM retry with jittered exponential backoff for 429/5xx/network errors, + `Retry-After` support, and byte-identical retry bodies. +- New events: `agent:updated`, `agent:usage`, `agent:tool_progress`, + `agent:permission_request`, `agent:subagent_update`, `agent:compaction`, + `agent:turn_resumed`, `agent:plan_update`. +- TUI: `/mode`, `/undo`, `/resume`; approval prompts can be denied with + `esc`/`ctrl+c`; submitting while busy now steers the turn instead of + dropping the draft. + ### Changed +- Tool failures are fed back to the model instead of ending the turn; only + repeated consecutive failures stop it. +- Denied tool calls get an error result immediately; asks pause per call. - Repository consolidation: the legacy crush-architecture port under `internal/` and `cmd/rakitkode` has been removed. The top-level packages (`agent`, `harness`, `llm`, `tools`, …) are the single runtime and the diff --git a/README.md b/README.md index 9a24e01..9c4aa53 100644 --- a/README.md +++ b/README.md @@ -76,7 +76,10 @@ exactly what they need. |---|---| | `engine` | facade: one constructor, `Run`/`RunIn` one-shot turns | | `agent` | session lifecycle, turn loop, approvals, subagents | -| `harness` | provider-agnostic turn runtime, permissions, scheduling | +| `permission` | collaboration modes, allow/deny/ask rules, per-tool decisions | +| `hooks` | lifecycle hooks (`PreToolUse`, `PostToolUse`, `Stop`, …) | +| `checkpoint` | file snapshots, undo/rewind | +| `harness` | provider-agnostic tool scheduler and metadata | | `llm` | OpenAI-compatible multi-provider client with streaming | | `tools` | canonical tool set + registry (read/write/edit/bash/search/…) | | `mcp` | MCP client manager (stdio + HTTP), tool bridging | diff --git a/agent/agent.go b/agent/agent.go index 968f24c..3332395 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -17,11 +17,14 @@ import ( "sync/atomic" "time" + "github.com/Haslab-dev/RakitKode/checkpoint" "github.com/Haslab-dev/RakitKode/decisions" "github.com/Haslab-dev/RakitKode/events" + "github.com/Haslab-dev/RakitKode/hooks" "github.com/Haslab-dev/RakitKode/llm" "github.com/Haslab-dev/RakitKode/mcp" "github.com/Haslab-dev/RakitKode/memory" + "github.com/Haslab-dev/RakitKode/permission" "github.com/Haslab-dev/RakitKode/plugins" "github.com/Haslab-dev/RakitKode/sandbox" "github.com/Haslab-dev/RakitKode/skills" @@ -74,25 +77,32 @@ const ( ) type Session struct { - ID string `json:"id"` - Name string `json:"name"` - RoleFilter RoleFilter `json:"role_filter"` - State SessionState `json:"state"` - Folder string `json:"folder"` - ProjectName string `json:"project_name,omitempty"` - Model string `json:"model,omitempty"` - AllowBash bool `json:"allow_bash,omitempty"` - ReadOnly bool `json:"read_only,omitempty"` - Tools []string `json:"tools,omitempty"` - MCPServers []string `json:"mcp_servers,omitempty"` - Skills []string `json:"skills,omitempty"` - ParentSessionID string `json:"parent_session_id,omitempty"` - SubagentIDs []string `json:"subagent_ids,omitempty"` - Messages []AgentMessage `json:"messages"` - Tasks []TaskItem `json:"tasks"` - Progress SessionProgress `json:"progress,omitempty"` - TokenUsage llm.TokenStats `json:"token_usage"` - AutoApprove bool `json:"auto_approve"` + ID string `json:"id"` + Name string `json:"name"` + RoleFilter RoleFilter `json:"role_filter"` + State SessionState `json:"state"` + Folder string `json:"folder"` + ProjectName string `json:"project_name,omitempty"` + Model string `json:"model,omitempty"` + AllowBash bool `json:"allow_bash,omitempty"` + ReadOnly bool `json:"read_only,omitempty"` + Tools []string `json:"tools,omitempty"` + MCPServers []string `json:"mcp_servers,omitempty"` + Skills []string `json:"skills,omitempty"` + ParentSessionID string `json:"parent_session_id,omitempty"` + SubagentIDs []string `json:"subagent_ids,omitempty"` + Messages []AgentMessage `json:"messages"` + Tasks []TaskItem `json:"tasks"` + Progress SessionProgress `json:"progress,omitempty"` + TokenUsage llm.TokenStats `json:"token_usage"` + AutoApprove bool `json:"auto_approve"` + // PermissionMode overrides the workspace mode for this session + // (plan|build|edit|yolo|auto). Empty falls back to the settings file. + PermissionMode string `json:"permission_mode,omitempty"` + // Interrupted marks a turn recovered from the journal after a crash/quit. + Interrupted bool `json:"interrupted,omitempty"` + // SessionAllow holds "always allow" rules added during this session. + SessionAllow []permission.Rule `json:"permission_allow,omitempty"` PendingTools []ContentBlock `json:"pending_tools,omitempty"` PendingQuestions []tools.AskQuestion `json:"pending_questions,omitempty"` Dialect string `json:"dialect,omitempty"` // "" = native tool calling, "xml" = in-band @@ -138,6 +148,8 @@ type Manager struct { schedulesMu sync.Mutex schedules map[string]*AgentSchedule storePath string + permSettings permission.Settings + hookRunner *hooks.Runner // defaultFolder is the workspace root the app is currently serving. // Sessions created without an explicit folder are scoped to it — never to @@ -223,6 +235,13 @@ type ManagerOptions struct { // Memory optionally overrides the default memory manager (which is // rooted at DataDir). Memory *memory.Manager + // Permissions optionally seeds the permission pipeline (mode + rules). + Permissions permission.Settings + // Hooks optionally supplies a lifecycle hook runner. nil disables hooks. + Hooks *hooks.Runner + // Workspace is the project root used for hook execution and checkpoints. + // Empty falls back to SetDefaultFolder. + Workspace string } // NewManagerWithOptions builds the Manager from explicit options. See @@ -266,6 +285,11 @@ func NewManagerWithOptions(opts ManagerOptions) *Manager { subagentNotes: make(map[string][]SubagentNotification), schedules: make(map[string]*AgentSchedule), storePath: storePath, + permSettings: opts.Permissions, + hookRunner: opts.Hooks, + } + if opts.Workspace != "" { + m.defaultFolder.Store(opts.Workspace) } m.loadSessions() m.loadDefinitions() @@ -427,6 +451,124 @@ func (m *Manager) MemoryManager() *memory.Manager { return m.memoryMgr } +// SetPermissionSettings replaces the workspace permission settings. +func (m *Manager) SetPermissionSettings(s permission.Settings) { + m.mu.Lock() + defer m.mu.Unlock() + m.permSettings = s +} + +// SetHookRunner installs the lifecycle hook runner. +func (m *Manager) SetHookRunner(r *hooks.Runner) { + m.mu.Lock() + defer m.mu.Unlock() + m.hookRunner = r +} + +// hooksSnapshot returns the current hook runner without holding the lock. +func (m *Manager) hooksSnapshot() *hooks.Runner { + m.mu.RLock() + defer m.mu.RUnlock() + return m.hookRunner +} + +// SetPermissionMode overrides the collaboration mode for one session. +func (m *Manager) SetPermissionMode(sessionID, mode string) error { + m.mu.Lock() + defer m.mu.Unlock() + sess, ok := m.sessions[sessionID] + if !ok { + return fmt.Errorf("session %s not found", sessionID) + } + sess.PermissionMode = mode + m.saveSessionLocked(sessionID) + return nil +} + +// PermissionModeFor returns the effective mode for a session. +func (m *Manager) PermissionModeFor(sessionID string) string { + m.mu.RLock() + defer m.mu.RUnlock() + sess, ok := m.sessions[sessionID] + if !ok { + return string(m.effectiveMode(nil)) + } + return string(m.effectiveMode(sess)) +} + +func (m *Manager) effectiveMode(sess *Session) permission.Mode { + if sess != nil && sess.PermissionMode != "" { + return permission.Mode(sess.PermissionMode) + } + if sess != nil && sess.AutoApprove { + return permission.ModeYolo + } + if m.permSettings.Mode != "" { + return m.permSettings.Mode + } + return permission.ModeBuild +} + +func (m *Manager) permissionContext(sess *Session) permission.Context { + mode := m.effectiveMode(sess) + ctx := permission.Context{ + Mode: mode, + PlanEnabled: mode == permission.ModePlan, + Rules: m.permSettings.Permissions, + Disallowed: m.permSettings.DisallowedTools, + } + if sess != nil { + ctx.SessionAllow = sess.SessionAllow + } + return ctx +} + +// UndoLast restores the most recent checkpoint for a session. +func (m *Manager) UndoLast(sessionID string) error { + folder := m.getFolder(sessionID) + if folder == "" { + return fmt.Errorf("session has no workspace folder") + } + home := filepath.Dir(m.storePath) + cp := checkpoint.New(home, folder, sessionID) + _, err := cp.RestoreLast() + if err == nil { + m.emitSessionUpdate(sessionID) + } + return err +} + +// ResumeTurn re-enters an interrupted turn recovered from the journal. +func (m *Manager) ResumeTurn(sessionID string) error { + m.mu.Lock() + sess, ok := m.sessions[sessionID] + if !ok { + m.mu.Unlock() + return fmt.Errorf("session %s not found", sessionID) + } + if !sess.Interrupted && sess.State == StateIdle { + m.mu.Unlock() + return fmt.Errorf("session %s has no interrupted turn", sessionID) + } + sess.Interrupted = false + if len(sess.PendingTools) > 0 { + sess.State = StateAwaiting + } else if len(sess.PendingQuestions) > 0 { + sess.State = StateAwaitingInput + } else { + sess.State = StateIdle + } + m.mu.Unlock() + + if sess.State == StateAwaiting || sess.State == StateAwaitingInput { + m.emitSessionUpdate(sessionID) + return nil + } + go m.runAgentTurn(context.Background(), sessionID) + m.emitAgentEvent(events.AgentTurnResumed, sessionID, map[string]interface{}{}) + return nil +} + func (m *Manager) loadSessions() { m.mu.Lock() defer m.mu.Unlock() @@ -452,6 +594,11 @@ func (m *Manager) loadSessions() { if !strings.Contains(s.SystemPrompt, "Response style") { s.SystemPrompt = buildSystemPrompt(s.RoleFilter, s.CustomPrompt, s.CustomRules) } + // A journal means the previous process died mid-turn. Restore the + // interrupted state; tools are never re-executed automatically. + if m.applyJournal(&s) { + stateChanged = true + } m.sessions[s.ID] = &s if stateChanged { m.saveSessionLocked(s.ID) @@ -1014,6 +1161,21 @@ func (m *Manager) SteerMessage(sessionID string, instruction string) error { } func (m *Manager) SendMessage(ctx context.Context, sessionID string, userContent string, mentionedPaths []string) error { + // UserPromptSubmit hooks may block or augment the prompt before it enters + // the session. They run outside the session lock. + if hr := m.hooksSnapshot(); hr != nil && hr.Has(hooks.UserPromptSubmit) { + res := hr.Run(ctx, hooks.UserPromptSubmit, "", map[string]any{ + "session_id": sessionID, + "prompt": userContent, + }) + if res.Blocked { + return fmt.Errorf("prompt blocked by hook: %s", res.Reason) + } + if res.Output != "" { + userContent += "\n\n" + res.Output + } + } + m.mu.Lock() sess, ok := m.sessions[sessionID] if !ok { @@ -1438,6 +1600,10 @@ func parseSkillInvocation(text string) (string, string, bool) { // is generous enough for real coding tasks. const maxReasoningSteps = 120 +// maxConsecutiveToolFailures stops a turn after this many back-to-back tool +// errors when the model has not recovered. +const maxConsecutiveToolFailures = 3 + // maxToolBudget is the per-turn tool budget in points (cheap=1, medium=3, // high=10), surfaced to the model each iteration as guidance. const maxToolBudget = 300 @@ -1544,6 +1710,7 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { <-routeDone budgetSpent := 0 + consecutiveToolFailures := 0 transcriptWindowSize := 40 isFinalSummaryStep := false for iter := 0; iter < maxReasoningSteps; iter++ { @@ -1998,41 +2165,83 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { } } - // Approval gate: pause the whole batch if any mutating tool needs - // approval (the old code gated only the first tool and dropped the - // rest). Scratch writes inside the project sandbox are exempt — - // that is what the sandbox is for — but bash is never exempt - // because a command cannot be path-scoped. - if len(toolCalls) > 0 && !sess.AutoApprove { - needApproval := false + // Permission gate: evaluate every call individually. Denied calls get + // an error result immediately; asks pause the batch; allowed calls run. + // Scratch writes inside the project sandbox are exempt. + if len(toolCalls) > 0 { + m.mu.RLock() + permCtx := m.permissionContext(sess) + m.mu.RUnlock() + + var allowed []ContentBlock + var asks []ContentBlock for _, tc := range toolCalls { - if isMutatingTool(tc.Name) && !isSandboxScopedCall(tc, folder) { - needApproval = true - break + md := tools.MetadataFor(tc.Name) + meta := permission.ToolMeta{ + Name: tc.Name, + ReadOnly: md.ReadOnly, + Destructive: md.Destructive, + NeedsApproval: md.NeedsApproval, + RequiresUserInteraction: md.RequiresUserInteraction, + SideEffectScope: md.SideEffectScope, + IsMCP: strings.Contains(tc.Name, "/") || strings.HasPrefix(tc.Name, "mcp__"), + } + dec := permission.Check(tc.Name, meta, tc.Arguments, permCtx) + if dec.Action == permission.ActionAsk && isSandboxScopedCall(tc, folder) { + dec.Action = permission.ActionAllow + } + switch dec.Action { + case permission.ActionDeny: + text := fmt.Sprintf("Tool %q denied: %s", tc.Name, dec.Message) + _ = m.appendToolResult(sessionID, tc, text, true) + case permission.ActionAsk: + asks = append(asks, tc) + default: + allowed = append(allowed, tc) } } - if needApproval { + toolCalls = allowed + + if len(asks) > 0 { cleanupStep() m.mu.Lock() sess.State = StateAwaiting - sess.PendingTools = toolCalls + sess.PendingTools = asks m.saveSessionsLocked() + if s, ok := m.sessions[sessionID]; ok { + m.writeJournalLocked(s, turnNum) + } m.mu.Unlock() m.emitSessionUpdate(sessionID) + for _, tc := range asks { + m.emitAgentEvent(events.AgentPermission, sessionID, map[string]interface{}{ + "tool_call_id": tc.ToolCallID, + "name": tc.Name, + "args": tc.Arguments, + "mode": string(permCtx.Mode), + "turn": turnNum, + }) + } return } + if len(toolCalls) == 0 { + // Every call was denied; let the model see the errors. + cleanupStep() + continue + } } // Execute the batch under harness scheduler semantics (read-only // calls in parallel groups, mutating calls sequential), then loop. if len(toolCalls) > 0 { - // Phase 6 snapshot: capture git tree hash before mutating tool executions - for _, tc := range toolCalls { - if isMutatingTool(tc.Name) { - if snap := m.captureTurnGitSnapshot(folder); snap != "" { - log.Printf("[agent:%s] Created pre-mutation git snapshot: %s", sessionID, snap) - } - break + // Checkpoint touched files before a mutating batch so /undo and + // /rewind can restore them. + if files := mutatedPaths(toolCalls); len(files) > 0 && folder != "" { + cp := checkpoint.New(filepath.Dir(m.storePath), folder, sessionID) + if snap, err := cp.Capture(turnNum, files); err == nil { + log.Printf("[agent:%s] checkpoint %d: %d file(s)", sessionID, snap.Turn, len(snap.Files)) + } else { + log.Printf("[agent:%s] checkpoint failed: %v", sessionID, err) } } @@ -2042,8 +2251,13 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { m.mu.Unlock() m.emitSessionUpdate(sessionID) - spent, fatal := m.executeScheduledBatch(stepCtx, sessionID, toolCalls, turnNum) + spent, fatal, failures := m.executeScheduledBatch(stepCtx, sessionID, toolCalls, turnNum) budgetSpent += spent + if failures > 0 { + consecutiveToolFailures += failures + } else { + consecutiveToolFailures = 0 + } // Capture cancellation before cleanupStep cancels stepCtx (see // the LLM error path above): otherwise fatal tool failures are // misread as steering preemption and the turn never ends. @@ -2092,9 +2306,17 @@ func (m *Manager) runAgentTurn(ctx context.Context, sessionID string) { continue } if fatal { - // A tool failed fatally; surface it and stop rather than - // feeding the error back into an infinite loop. - m.finishTurn(sessionID, "One or more tool calls failed.") + // The step context was cancelled; surface it and stop. + m.finishTurn(sessionID, "One or more tool calls were interrupted.") + return + } + // Tool failures are fed back to the model so it can adapt. Only a + // long run of consecutive failures stops the turn, with a notice + // instead of a silent dead end. + if consecutiveToolFailures >= maxConsecutiveToolFailures { + m.finishTurn(sessionID, fmt.Sprintf( + "Stopped after %d consecutive tool failures. Review the errors above and retry with a different approach.", + consecutiveToolFailures)) return } @@ -2150,6 +2372,7 @@ func (m *Manager) RespondApproval(ctx context.Context, sessionID string, approve pending := sess.PendingTools sess.PendingTools = nil + m.clearJournal(sessionID) if autoApproveAll { sess.AutoApprove = true @@ -2203,6 +2426,62 @@ func (m *Manager) executeToolCall(ctx context.Context, sessionID string, block C return content, err } +// handlePlanModeCall applies enter_plan_mode / exit_plan_mode to the session +// and records the plan artifact. It returns handled=false for other tools. +func (m *Manager) handlePlanModeCall(sessionID string, block ContentBlock, turn int) (bool, string) { + name := strings.ToLower(block.Name) + if name != "enter_plan_mode" && name != "exit_plan_mode" { + return false, "" + } + m.emitToolStart(sessionID, block, turn) + + var content, planPath, mode string + switch name { + case "enter_plan_mode": + m.mu.Lock() + if s, ok := m.sessions[sessionID]; ok { + s.PermissionMode = string(permission.ModePlan) + s.UpdatedAt = time.Now() + m.saveSessionLocked(sessionID) + } + m.mu.Unlock() + mode = "plan" + content = "Plan mode enabled. Only read-only tools are available. Investigate the codebase, then call exit_plan_mode with the plan." + case "exit_plan_mode": + plan, _ := block.Arguments["plan"].(string) + if plan != "" { + if folder := m.getFolder(sessionID); folder != "" { + dir := filepath.Join(folder, ".rakit", "plans") + if err := os.MkdirAll(dir, 0o755); err == nil { + planPath = filepath.Join(dir, time.Now().Format("20060102-150405")+"-plan.md") + _ = os.WriteFile(planPath, []byte(plan), 0o644) + } + } + } + m.mu.Lock() + if s, ok := m.sessions[sessionID]; ok { + s.PermissionMode = string(permission.ModeBuild) + s.UpdatedAt = time.Now() + m.saveSessionLocked(sessionID) + } + m.mu.Unlock() + mode = "build" + content = "Plan mode disabled; build mode enabled." + if planPath != "" { + content += " Plan saved to " + planPath + "." + } + } + + m.emitToolEnd(sessionID, block, content, nil, false, turn) + _ = m.appendToolResult(sessionID, block, content, false) + m.emitAgentEvent(events.AgentPlanUpdate, sessionID, map[string]interface{}{ + "mode": mode, + "plan_path": planPath, + "approved": name == "exit_plan_mode", + }) + return true, content +} + // executeToolCallDetailed additionally returns the structured display payload // (raw tool result map) the UI renders from the tool_end event. func (m *Manager) executeToolCallDetailed(ctx context.Context, sessionID string, block ContentBlock, turn ...int) (string, map[string]any, error) { @@ -2211,6 +2490,10 @@ func (m *Manager) executeToolCallDetailed(ctx context.Context, sessionID string, currentTurn = turn[0] } + if handled, content := m.handlePlanModeCall(sessionID, block, currentTurn); handled { + return content, nil, nil + } + // Role-gating safety check for subagents and permissions var sess *Session m.mu.RLock() @@ -2228,6 +2511,33 @@ func (m *Manager) executeToolCallDetailed(ctx context.Context, sessionID string, return contentStr, nil, err } + // PreToolUse hooks may modify arguments or block the call outright. + m.mu.RLock() + hookRunner := m.hookRunner + m.mu.RUnlock() + if hookRunner != nil { + res := hookRunner.Run(ctx, hooks.PreToolUse, block.Name, map[string]any{ + "session_id": sessionID, + "tool_name": block.Name, + "tool_input": block.Arguments, + }) + if res.ModifiedInput != nil { + if block.Arguments == nil { + block.Arguments = map[string]any{} + } + for k, v := range res.ModifiedInput { + block.Arguments[k] = v + } + } + if res.Blocked { + m.emitToolStart(sessionID, block, currentTurn) + contentStr := fmt.Sprintf("Blocked by %s hook: %s", hooks.PreToolUse, res.Reason) + m.emitToolEnd(sessionID, block, contentStr, nil, true, currentTurn) + _ = m.appendToolResult(sessionID, block, contentStr, true) + return contentStr, nil, fmt.Errorf("blocked by hook: %s", res.Reason) + } + } + // Attach the session bridge so `todo` / `ask` tools can touch session state, // and the per-session read-state tracker so edit/write can enforce // read-before-edit. @@ -2287,6 +2597,19 @@ func (m *Manager) executeToolCallDetailed(ctx context.Context, sessionID string, } } + // PostToolUse hooks can append context (formatters, linters, policy notes). + if err == nil && hookRunner != nil { + hr := hookRunner.Run(ctx, hooks.PostToolUse, block.Name, map[string]any{ + "session_id": sessionID, + "tool_name": block.Name, + "tool_input": block.Arguments, + "tool_output": contentStr, + }) + if hr.Output != "" { + contentStr += "\n\n" + hr.Output + } + } + m.emitToolEnd(sessionID, block, contentStr, display, isErr, currentTurn) if err := m.appendToolResult(sessionID, block, contentStr, isErr); err != nil { return contentStr, display, err @@ -2372,6 +2695,10 @@ func (m *Manager) finishTurn(sessionID string, errMsg string) { detector.Reset() } sess.State = StateIdle + sess.Interrupted = false + sess.PendingTools = nil + sess.PendingQuestions = nil + m.clearJournal(sessionID) sess.UpdatedAt = time.Now() isCancel := errMsg != "" && strings.Contains(strings.ToLower(errMsg), "canceled") if errMsg != "" && !isCancel { @@ -2495,6 +2822,12 @@ func (m *Manager) finishTurn(sessionID string, errMsg string) { m.emitAgentEvent(events.AgentTurnEnd, sessionID, map[string]interface{}{ "error": errMsg, }) + if hr := m.hooksSnapshot(); hr != nil && hr.Has(hooks.Stop) { + go hr.Run(context.Background(), hooks.Stop, "", map[string]any{ + "session_id": sessionID, + "error": errMsg, + }) + } if errMsg != "" && !isCancel { m.emitMessageEnd(sessionID, &AgentMessage{ ID: uuid.New().String(), @@ -2745,6 +3078,7 @@ func (b *sessionBridge) Ask(questions []tools.AskQuestion) error { sess.PendingQuestions = questions sess.UpdatedAt = time.Now() b.m.saveSessionsLocked() + b.m.writeJournalLocked(sess, 0) b.m.emitSessionUpdate(b.sessionID) b.m.emitAgentEvent(events.AgentAsk, b.sessionID, map[string]interface{}{ "questions": questions, @@ -3616,6 +3950,7 @@ func (m *Manager) RespondAsk(sessionID string, answers map[string]any) error { return fmt.Errorf("no pending questions for session %s", sessionID) } sess.PendingQuestions = nil + m.clearJournal(sessionID) b, _ := json.MarshalIndent(answers, "", " ") var foundPendingResult bool @@ -4216,18 +4551,8 @@ func (m *Manager) GetUsageSummary() map[string]interface{} { // message with the matching tool_call_id. // isReadOnlyTool reports whether a tool only reads state (safe to dedup // identical calls within one batch without changing semantics). -var readOnlyTools = map[string]bool{ - "read": true, - "read_multiple": true, - "read_directory_files": true, - "search": true, - "glob": true, - "find": true, - "git_status": true, -} - func isReadOnlyTool(name string) bool { - return readOnlyTools[name] + return tools.MetadataFor(name).ReadOnly } // canonicalArgs renders tool args deterministically for dedup keys. @@ -4691,15 +5016,32 @@ func isBashTool(name string) bool { } func isFileMutationTool(name string) bool { - switch name { - case "write", "write_file", "create_file", "edit", "edit_file", "replace_file_content": - return true + md := tools.MetadataFor(name) + return md.Category == "file" && !md.ReadOnly +} + +// mutatedPaths extracts the file paths a batch may modify, for checkpointing. +func mutatedPaths(toolCalls []ContentBlock) []string { + var paths []string + seen := map[string]bool{} + for _, tc := range toolCalls { + md := tools.MetadataFor(tc.Name) + if md.Category != "file" || md.ReadOnly { + continue + } + for _, key := range []string{"path", "file_path", "filepath", "filePath", "target_file", "file"} { + if v, ok := tc.Arguments[key].(string); ok && v != "" && !seen[v] { + seen[v] = true + paths = append(paths, v) + } + } } - return false + return paths } func isMutatingTool(name string) bool { - return isBashTool(name) || isFileMutationTool(name) || isComputerInputTool(name) + md := tools.MetadataFor(name) + return md.Destructive || md.NeedsApproval || md.Category == "file" || md.Category == "shell" || isComputerInputTool(name) } // isComputerInputTool reports whether a tool takes over the user's desktop. diff --git a/agent/harness_engine.go b/agent/harness_engine.go index 83d0a8c..39b3efb 100644 --- a/agent/harness_engine.go +++ b/agent/harness_engine.go @@ -3,8 +3,10 @@ package agent import ( "context" "sync" + "time" "github.com/Haslab-dev/RakitKode/harness" + "github.com/Haslab-dev/RakitKode/tools" ) // Harness-scheduled tool batching: this wires the harness tool scheduler @@ -30,17 +32,16 @@ func (m *Manager) schedRegistry() *harness.ToolRegistry { if name == "" { continue } - meta := harness.ToolMetadata{Name: name} - switch { - case isReadOnlyTool(name): - meta.ReadOnly = true - meta.SideEffectScope = "none" - case isMutatingTool(name): - meta.Destructive = true - meta.SideEffectScope = "workspace" - default: - // Unknown classification: never schedule in parallel. - meta.SideEffectScope = "workspace" + md := tools.MetadataFor(name) + meta := harness.ToolMetadata{ + Name: name, + ReadOnly: md.ReadOnly, + Destructive: md.Destructive, + ConcurrentSafe: md.ConcurrencySafe, + NeedsApproval: md.NeedsApproval, + SideEffectScope: md.SideEffectScope, + TimeoutMs: int(md.Timeout / time.Millisecond), + MaxOutputBytes: md.MaxOutputBytes, } _ = reg.Register(&harness.ToolEntry{Metadata: meta}) } @@ -49,10 +50,11 @@ func (m *Manager) schedRegistry() *harness.ToolRegistry { // batchState carries dedup + budget accumulation across a scheduled batch. type batchState struct { - mu sync.Mutex - dedup map[string]dedupEntry - spent int - fatal bool + mu sync.Mutex + dedup map[string]dedupEntry + spent int + failures int + fatal bool } type dedupEntry struct { @@ -62,11 +64,12 @@ type dedupEntry struct { } // executeScheduledBatch runs one batch of tool calls under harness scheduler -// semantics and returns the points spent. fatal reports a tool failure the -// same way the previous sequential loop did (turn ends). -func (m *Manager) executeScheduledBatch(ctx context.Context, sessionID string, toolCalls []ContentBlock, turn int) (spent int, fatal bool) { +// semantics and returns the points spent, whether the turn was cancelled, and +// how many individual calls failed. Tool failures are reported to the model as +// error results instead of ending the turn. +func (m *Manager) executeScheduledBatch(ctx context.Context, sessionID string, toolCalls []ContentBlock, turn int) (spent int, fatal bool, failures int) { if ctx != nil && ctx.Err() != nil { - return 0, true + return 0, true, 0 } reg := m.schedRegistry() calls := make([]harness.ToolCall, 0, len(toolCalls)) @@ -88,17 +91,16 @@ func (m *Manager) executeScheduledBatch(ctx context.Context, sessionID string, t // Fall back to the sequential path on scheduling failure. return m.executeBatchSequential(ctx, sessionID, toolCalls, turn) } - st := &batchState{dedup: make(map[string]dedupEntry)} for _, group := range sched.ParallelGroups { if ctx != nil && ctx.Err() != nil { - return st.spent, true + return st.spent, true, st.failures } if len(group) == 1 { tc := byID[calls[group[0]].ID] m.executeOneBatched(ctx, sessionID, tc, st, turn) if st.fatal { - return st.spent, true + return st.spent, true, st.failures } continue } @@ -136,7 +138,7 @@ func (m *Manager) executeScheduledBatch(ctx context.Context, sessionID string, t } wg.Wait() if st.fatal { - return st.spent, true + return st.spent, true, st.failures } // Mirror dedup'd results to duplicate caller ids in this group. for key, callers := range seen { @@ -152,7 +154,7 @@ func (m *Manager) executeScheduledBatch(ctx context.Context, sessionID string, t } } } - return st.spent, st.fatal + return st.spent, st.fatal, st.failures } // executeOneBatched executes one call with read-only dedup + bookkeeping. @@ -180,23 +182,23 @@ func (m *Manager) executeOneBatched(ctx context.Context, sessionID string, tc Co st.mu.Lock() st.spent += m.toolCost(tc.Name) if execErr != nil { - st.fatal = true + st.failures++ } st.mu.Unlock() } // executeBatchSequential is the pre-harness sequential path, kept as the // fallback and for parity testing. -func (m *Manager) executeBatchSequential(ctx context.Context, sessionID string, toolCalls []ContentBlock, turn int) (spent int, fatal bool) { +func (m *Manager) executeBatchSequential(ctx context.Context, sessionID string, toolCalls []ContentBlock, turn int) (spent int, fatal bool, failures int) { st := &batchState{dedup: make(map[string]dedupEntry)} for _, tc := range toolCalls { if ctx != nil && ctx.Err() != nil { - return st.spent, true + return st.spent, true, st.failures } m.executeOneBatched(ctx, sessionID, tc, st, turn) if st.fatal { - return st.spent, true + return st.spent, true, st.failures } } - return st.spent, st.fatal + return st.spent, st.fatal, st.failures } diff --git a/agent/journal.go b/agent/journal.go new file mode 100644 index 0000000..a31cc2c --- /dev/null +++ b/agent/journal.go @@ -0,0 +1,101 @@ +package agent + +import ( + "encoding/json" + "os" + "path/filepath" + "time" + + "github.com/Haslab-dev/RakitKode/tools" +) + +// turnJournal persists enough state to resume an interrupted turn after a +// crash, a quit, or a killed process. It never stores model output; resuming +// always asks for approval before executing any pending tool call. +type turnJournal struct { + SessionID string `json:"session_id"` + State SessionState `json:"state"` + Step int `json:"step"` + Model string `json:"model,omitempty"` + PendingTools []ContentBlock `json:"pending_tools,omitempty"` + PendingQuestions []tools.AskQuestion `json:"pending_questions,omitempty"` + UpdatedAt time.Time `json:"updated_at"` +} + +func (m *Manager) journalPath(sessionID string) string { + if m.storePath == "" { + return "" + } + return filepath.Join(filepath.Dir(m.storePath), "sessions", sessionID, "turn.json") +} + +// writeJournalLocked persists the current turn state. Caller holds m.mu. +func (m *Manager) writeJournalLocked(sess *Session, step int) { + path := m.journalPath(sess.ID) + if path == "" { + return + } + if sess.State == StateIdle { + _ = os.Remove(path) + return + } + j := turnJournal{ + SessionID: sess.ID, + State: sess.State, + Step: step, + Model: sess.Model, + PendingTools: sess.PendingTools, + PendingQuestions: sess.PendingQuestions, + UpdatedAt: time.Now(), + } + data, err := json.MarshalIndent(j, "", " ") + if err != nil { + return + } + _ = os.MkdirAll(filepath.Dir(path), 0o755) + _ = os.WriteFile(path, data, 0o644) +} + +// clearJournal removes the journal for a session. +func (m *Manager) clearJournal(sessionID string) { + if path := m.journalPath(sessionID); path != "" { + _ = os.Remove(path) + } +} + +// readJournal loads a journal without locking. +func (m *Manager) readJournal(sessionID string) *turnJournal { + path := m.journalPath(sessionID) + if path == "" { + return nil + } + data, err := os.ReadFile(path) + if err != nil { + return nil + } + var j turnJournal + if err := json.Unmarshal(data, &j); err != nil { + return nil + } + if j.State == "" || j.State == StateIdle { + return nil + } + return &j +} + +// applyJournal restores an interrupted turn onto a loaded session. +func (m *Manager) applyJournal(s *Session) bool { + j := m.readJournal(s.ID) + if j == nil { + return false + } + s.State = j.State + if len(j.PendingTools) > 0 { + s.PendingTools = j.PendingTools + } + if len(j.PendingQuestions) > 0 { + s.PendingQuestions = j.PendingQuestions + } + s.Interrupted = true + return true +} diff --git a/checkpoint/checkpoint.go b/checkpoint/checkpoint.go new file mode 100644 index 0000000..afe79ae --- /dev/null +++ b/checkpoint/checkpoint.go @@ -0,0 +1,214 @@ +// Package checkpoint captures a restorable snapshot of the files a mutating +// tool batch is about to touch. It powers /undo and /rewind. +package checkpoint + +import ( + "encoding/json" + "fmt" + "io" + "os" + "os/exec" + "path/filepath" + "sort" + "strconv" + "strings" + "time" +) + +const maxFileBytes = 10 << 20 // 10 MiB per file + +// FileRecord remembers whether a file existed before the batch. +type FileRecord struct { + Path string `json:"path"` // relative to workspace + Existed bool `json:"existed"` +} + +// Snapshot is one restorable checkpoint. +type Snapshot struct { + Turn int `json:"turn"` + CreatedAt time.Time `json:"created_at"` + GitCommit string `json:"git_commit,omitempty"` + Files []FileRecord `json:"files"` +} + +// Manager stores snapshots for one session. +type Manager struct { + home string + workspace string + sessionID string +} + +func New(home, workspace, sessionID string) *Manager { + return &Manager{home: home, workspace: workspace, sessionID: sessionID} +} + +func (m *Manager) root() string { + return filepath.Join(m.home, "sessions", m.sessionID, "checkpoints") +} + +func (m *Manager) dir(turn int) string { + return filepath.Join(m.root(), fmt.Sprintf("%04d", turn)) +} + +func (m *Manager) filesDir(turn int) string { + return filepath.Join(m.dir(turn), "files") +} + +// Capture snapshots paths (absolute or workspace-relative). Missing paths are +// recorded as non-existent so undo can delete files the batch creates. +func (m *Manager) Capture(turn int, paths []string) (*Snapshot, error) { + snap := &Snapshot{Turn: turn, CreatedAt: time.Now()} + if err := os.MkdirAll(m.filesDir(turn), 0o755); err != nil { + return nil, err + } + + seen := map[string]bool{} + for _, p := range paths { + rel, err := m.rel(p) + if err != nil || rel == "" || seen[rel] { + continue + } + seen[rel] = true + rec := FileRecord{Path: rel} + src := filepath.Join(m.workspace, rel) + info, err := os.Stat(src) + if err == nil && info.Mode().IsRegular() && info.Size() <= maxFileBytes { + rec.Existed = true + dst := filepath.Join(m.filesDir(turn), rel) + if err := copyFile(src, dst); err != nil { + return nil, err + } + } + snap.Files = append(snap.Files, rec) + } + + snap.GitCommit = m.gitStashCreate() + data, err := json.MarshalIndent(snap, "", " ") + if err != nil { + return nil, err + } + if err := os.WriteFile(filepath.Join(m.dir(turn), "snapshot.json"), data, 0o644); err != nil { + return nil, err + } + return snap, nil +} + +func (m *Manager) rel(p string) (string, error) { + if p == "" { + return "", fmt.Errorf("empty path") + } + abs := p + if !filepath.IsAbs(abs) { + abs = filepath.Join(m.workspace, abs) + } + rel, err := filepath.Rel(m.workspace, abs) + if err != nil { + return "", err + } + if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return "", fmt.Errorf("path outside workspace") + } + return rel, nil +} + +func (m *Manager) gitStashCreate() string { + if _, err := os.Stat(filepath.Join(m.workspace, ".git")); err != nil { + return "" + } + cmd := exec.Command("git", "stash", "create", "rakit checkpoint") + cmd.Dir = m.workspace + out, err := cmd.Output() + if err != nil { + return "" + } + return strings.TrimSpace(string(out)) +} + +// Load reads a snapshot for a turn. +func (m *Manager) Load(turn int) (*Snapshot, error) { + data, err := os.ReadFile(filepath.Join(m.dir(turn), "snapshot.json")) + if err != nil { + return nil, err + } + var snap Snapshot + if err := json.Unmarshal(data, &snap); err != nil { + return nil, err + } + return &snap, nil +} + +// Last returns the highest-numbered snapshot. +func (m *Manager) Last() (*Snapshot, error) { + entries, err := os.ReadDir(m.root()) + if err != nil { + return nil, err + } + var turns []int + for _, e := range entries { + if !e.IsDir() { + continue + } + n, err := strconv.Atoi(e.Name()) + if err == nil { + turns = append(turns, n) + } + } + if len(turns) == 0 { + return nil, fmt.Errorf("no checkpoints") + } + sort.Ints(turns) + return m.Load(turns[len(turns)-1]) +} + +// Restore writes snapshot contents back into the workspace. Files that did not +// exist at capture time are removed. +func (m *Manager) Restore(snap *Snapshot) error { + if snap == nil { + return fmt.Errorf("nil snapshot") + } + for _, rec := range snap.Files { + dst := filepath.Join(m.workspace, rec.Path) + if !rec.Existed { + _ = os.Remove(dst) + continue + } + src := filepath.Join(m.filesDir(snap.Turn), rec.Path) + if err := copyFile(src, dst); err != nil { + return err + } + } + return nil +} + +// RestoreLast restores the most recent checkpoint and returns it. +func (m *Manager) RestoreLast() (*Snapshot, error) { + snap, err := m.Last() + if err != nil { + return nil, err + } + return snap, m.Restore(snap) +} + +func copyFile(src, dst string) error { + if err := os.MkdirAll(filepath.Dir(dst), 0o755); err != nil { + return err + } + in, err := os.Open(src) + if err != nil { + return err + } + defer in.Close() + info, err := in.Stat() + if err != nil { + return err + } + out, err := os.OpenFile(dst, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, info.Mode().Perm()) + if err != nil { + return err + } + if _, err := io.Copy(out, in); err != nil { + out.Close() + return err + } + return out.Close() +} diff --git a/checkpoint/checkpoint_test.go b/checkpoint/checkpoint_test.go new file mode 100644 index 0000000..3cf9938 --- /dev/null +++ b/checkpoint/checkpoint_test.go @@ -0,0 +1,79 @@ +package checkpoint + +import ( + "os" + "path/filepath" + "testing" +) + +func TestCaptureAndRestore(t *testing.T) { + home := t.TempDir() + ws := t.TempDir() + target := filepath.Join(ws, "a.txt") + if err := os.WriteFile(target, []byte("original"), 0o644); err != nil { + t.Fatal(err) + } + + m := New(home, ws, "s1") + snap, err := m.Capture(1, []string{"a.txt", "new.txt"}) + if err != nil { + t.Fatal(err) + } + if len(snap.Files) != 2 { + t.Fatalf("files = %d", len(snap.Files)) + } + + if err := os.WriteFile(target, []byte("changed"), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(ws, "new.txt"), []byte("created"), 0o644); err != nil { + t.Fatal(err) + } + + if err := m.Restore(snap); err != nil { + t.Fatal(err) + } + got, err := os.ReadFile(target) + if err != nil { + t.Fatal(err) + } + if string(got) != "original" { + t.Fatalf("a.txt = %q", got) + } + if _, err := os.Stat(filepath.Join(ws, "new.txt")); !os.IsNotExist(err) { + t.Fatalf("new.txt should have been removed, err=%v", err) + } +} + +func TestLastRestoresNewest(t *testing.T) { + home := t.TempDir() + ws := t.TempDir() + if err := os.WriteFile(filepath.Join(ws, "a.txt"), []byte("v1"), 0o644); err != nil { + t.Fatal(err) + } + m := New(home, ws, "s1") + if _, err := m.Capture(1, []string{"a.txt"}); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(ws, "a.txt"), []byte("v2"), 0o644); err != nil { + t.Fatal(err) + } + if _, err := m.Capture(2, []string{"a.txt"}); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(ws, "a.txt"), []byte("v3"), 0o644); err != nil { + t.Fatal(err) + } + + snap, err := m.RestoreLast() + if err != nil { + t.Fatal(err) + } + if snap.Turn != 2 { + t.Fatalf("restored turn %d", snap.Turn) + } + got, _ := os.ReadFile(filepath.Join(ws, "a.txt")) + if string(got) != "v2" { + t.Fatalf("a.txt = %q", got) + } +} diff --git a/cmd/rakit/tui/slash.go b/cmd/rakit/tui/slash.go index d6d635c..c5bec20 100644 --- a/cmd/rakit/tui/slash.go +++ b/cmd/rakit/tui/slash.go @@ -66,6 +66,21 @@ var slashCommands = []command{ description: "MCP servers: /mcp [on|off|reconnect] [name]", action: actionMCP, }, + { + name: "/mode", + description: "permission mode: /mode plan|build|edit|yolo|auto", + action: actionMode, + }, + { + name: "/resume", + description: "resume an interrupted turn", + action: actionResume, + }, + { + name: "/undo", + description: "restore files from the last mutating batch", + action: actionUndo, + }, { name: "/model", description: "pick the model for this session (interactive)", @@ -243,6 +258,48 @@ func actionHelp(m *Model) ([]string, tea.Cmd) { return out, nil } +// actionMode shows or sets the session's permission mode. +func actionMode(m *Model) ([]string, tea.Cmd) { + if m.eng == nil { + return []string{errorStyle.Render("no engine attached")}, nil + } + if len(m.pendingSlashArgs) == 0 { + return []string{statusStyle.Render("permission mode: " + m.eng.PermissionModeFor(m.sessionID))}, nil + } + mode := strings.ToLower(m.pendingSlashArgs[0]) + switch mode { + case "plan", "build", "edit", "yolo", "auto": + default: + return []string{errorStyle.Render("mode must be plan|build|edit|yolo|auto")}, nil + } + if err := m.eng.SetPermissionMode(m.sessionID, mode); err != nil { + return []string{errorStyle.Render("failed: " + err.Error())}, nil + } + return []string{toolDotOKStyle.Render("✓ permission mode: " + mode)}, nil +} + +// actionUndo restores files from the last mutating batch. +func actionUndo(m *Model) ([]string, tea.Cmd) { + if m.eng == nil { + return []string{errorStyle.Render("no engine attached")}, nil + } + if err := m.eng.UndoLast(m.sessionID); err != nil { + return []string{errorStyle.Render("undo failed: " + err.Error())}, nil + } + return []string{toolDotOKStyle.Render("✓ restored the last checkpoint")}, nil +} + +// actionResume resumes an interrupted turn recovered from the journal. +func actionResume(m *Model) ([]string, tea.Cmd) { + if m.eng == nil { + return []string{errorStyle.Render("no engine attached")}, nil + } + if err := m.eng.ResumeTurn(m.sessionID); err != nil { + return []string{errorStyle.Render("resume failed: " + err.Error())}, nil + } + return []string{toolDotOKStyle.Render("✓ resumed turn")}, nil +} + // actionToggleYolo flips the session's auto-approval flag. func actionToggleYolo(m *Model) ([]string, tea.Cmd) { if m.mgr == nil { @@ -324,6 +381,9 @@ func actionStatus(m *Model) ([]string, tea.Cmd) { yolo = "on" } out = append(out, fmt.Sprintf(" %-9s %s", "yolo", yolo)) + if m.eng != nil { + out = append(out, fmt.Sprintf(" %-9s %s", "mode", m.eng.PermissionModeFor(m.sessionID))) + } return out, nil } diff --git a/cmd/rakit/tui/tui.go b/cmd/rakit/tui/tui.go index fbd68bf..5daaa11 100644 --- a/cmd/rakit/tui/tui.go +++ b/cmd/rakit/tui/tui.go @@ -306,12 +306,21 @@ func (m *Model) interceptKey(msg tea.KeyPressMsg) (bool, tea.Cmd) { case "n": m.respondApproval(false, false) m.appendBlock(toolDotErrStyle.Render(" ✗ denied")) + case "esc", "ctrl+c": + m.respondApproval(false, false) + m.appendBlock(toolDotErrStyle.Render(" ✗ denied (esc)")) + case "ctrl+d": + return true, tea.Quit } return true, nil } // Structured ask: number keys pick options for the current question. if m.pendingAsk != nil { + switch key { + case "ctrl+c", "ctrl+d": + return true, tea.Quit + } return true, m.interceptAskKey(key) } @@ -434,10 +443,23 @@ func (m *Model) interceptKey(msg tea.KeyPressMsg) (bool, tea.Cmd) { // starts one agent turn via the manager. func (m *Model) submit() tea.Cmd { value := strings.TrimSpace(m.textarea.Value()) - m.textarea.Reset() - if value == "" || m.busy { + if value == "" { return nil } + if m.busy { + // Keep the turn alive and steer it instead of discarding the draft. + if m.mgr != nil { + if err := m.mgr.SteerMessage(m.sessionID, value); err == nil { + m.appendBlock(userStyle.Render("↳ steer ") + value) + m.textarea.Reset() + m.status = "steering current turn" + return nil + } + } + m.status = "turn running — ctrl+c interrupts" + return nil + } + m.textarea.Reset() m.menuOpen = false m.appendBlock(userStyle.Render("❯ ") + value) if strings.HasPrefix(value, "/") { diff --git a/engine/engine.go b/engine/engine.go index e3b950f..4d867eb 100644 --- a/engine/engine.go +++ b/engine/engine.go @@ -15,9 +15,11 @@ import ( "github.com/Haslab-dev/RakitKode/agent" "github.com/Haslab-dev/RakitKode/events" + "github.com/Haslab-dev/RakitKode/hooks" "github.com/Haslab-dev/RakitKode/llm" "github.com/Haslab-dev/RakitKode/mcp" "github.com/Haslab-dev/RakitKode/memory" + "github.com/Haslab-dev/RakitKode/permission" "github.com/Haslab-dev/RakitKode/plugins" "github.com/Haslab-dev/RakitKode/search" "github.com/Haslab-dev/RakitKode/skills" @@ -138,6 +140,16 @@ func New(opts Options) (*Engine, error) { mcpMgr = mcp.NewManager(home) } + permSettings, err := permission.LoadSettings(permission.DefaultSettingsPaths(home, workspace)...) + if err != nil { + log.Printf("[engine] permissions: %v", err) + } + hookCfg, err := hooks.Load(hooks.DefaultPaths(home, workspace)...) + if err != nil { + log.Printf("[engine] hooks: %v", err) + } + hookRunner := hooks.NewRunner(hookCfg, workspace) + mgr := agent.NewManagerWithOptions(agent.ManagerOptions{ LLMClient: llmClient, ToolRegistry: toolReg, @@ -148,6 +160,9 @@ func New(opts Options) (*Engine, error) { Bus: bus, DataDir: home, Memory: opts.Memory, + Permissions: permSettings, + Hooks: hookRunner, + Workspace: workspace, }) mgr.SetDefaultFolder(workspace) @@ -179,6 +194,26 @@ func New(opts Options) (*Engine, error) { return e, nil } +// SetPermissionMode overrides the collaboration mode for one session. +func (e *Engine) SetPermissionMode(sessionID, mode string) error { + return e.mgr.SetPermissionMode(sessionID, mode) +} + +// PermissionModeFor returns the effective permission mode for a session. +func (e *Engine) PermissionModeFor(sessionID string) string { + return e.mgr.PermissionModeFor(sessionID) +} + +// UndoLast restores the most recent file checkpoint for a session. +func (e *Engine) UndoLast(sessionID string) error { + return e.mgr.UndoLast(sessionID) +} + +// ResumeTurn resumes an interrupted turn recovered from the journal. +func (e *Engine) ResumeTurn(sessionID string) error { + return e.mgr.ResumeTurn(sessionID) +} + // RegisterMCPTools connects nothing but registers whatever tools the MCP // manager reports as connected into the engine's tool registry. func (e *Engine) RegisterMCPTools() { diff --git a/events/bus.go b/events/bus.go index 61d7731..dd3df1a 100644 --- a/events/bus.go +++ b/events/bus.go @@ -68,6 +68,14 @@ const ( AgentAsk EventType = "agent:ask" AgentTaskUpdate EventType = "agent:task_update" AgentScheduleUpdate EventType = "agent:schedule_update" + AgentUpdated EventType = "agent:updated" + AgentUsage EventType = "agent:usage" + AgentToolProgress EventType = "agent:tool_progress" + AgentPermission EventType = "agent:permission_request" + AgentSubagentUpdate EventType = "agent:subagent_update" + AgentCompaction EventType = "agent:compaction" + AgentTurnResumed EventType = "agent:turn_resumed" + AgentPlanUpdate EventType = "agent:plan_update" // Search events SearchIndexUpdated EventType = "search:index:updated" diff --git a/events/doc.go b/events/doc.go index 252009c..01fb5ab 100644 --- a/events/doc.go +++ b/events/doc.go @@ -28,6 +28,14 @@ // when finished) // agent:schedule_update — schedule snapshot / change record // agent:updated (session-snapshot change signal; no payload contract) +// agent:usage prompt_tokens, completion_tokens, cached_tokens, +// cache_hit_rate, context_tokens, turn +// agent:tool_progress tool_call_id, name, elapsed_ms, turn +// agent:permission_request tool_call_id, name, args, reason, turn +// agent:subagent_update subagent_id, phase, summary, usage +// agent:compaction reason, before_tokens, after_tokens +// agent:turn_resumed (no additional keys) +// agent:plan_update mode, plan_path, approved // // Non-agent events (workspace:*, file:*, terminal:*, git:*, ssh:*, plugin:*) // are bus-internal to hosts that also use those features and follow the same diff --git a/hooks/hooks.go b/hooks/hooks.go new file mode 100644 index 0000000..de95c8f --- /dev/null +++ b/hooks/hooks.go @@ -0,0 +1,262 @@ +// Package hooks executes user-configured lifecycle hooks (PreToolUse, +// PostToolUse, UserPromptSubmit, Stop, SessionStart, SessionEnd, PreCompact). +// +// Hook configuration is Claude-Code compatible and merged global → workspace: +// +// {"hooks": { +// "PreToolUse": [{"matcher": "Bash|Edit", "hooks": [{"type": "command", +// "command": ".rakit/hooks/guard.sh", "timeout": 30}]}] +// }} +// +// A hook receives a JSON payload on stdin plus RAKIT_* environment variables. +// Exit code 2 blocks the action with stderr as the reason. JSON on stdout can +// return {"decision":"block","reason":"…","modified_input":{…}, +// "system_message":"…","output":"…"}. +package hooks + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "regexp" + "strings" + "time" +) + +// Event names. +type Event string + +const ( + PreToolUse Event = "PreToolUse" + PostToolUse Event = "PostToolUse" + UserPromptSubmit Event = "UserPromptSubmit" + Stop Event = "Stop" + SessionStart Event = "SessionStart" + SessionEnd Event = "SessionEnd" + PreCompact Event = "PreCompact" +) + +// Command is one hook entry. +type Command struct { + Type string `json:"type"` + Command string `json:"command"` + Timeout int `json:"timeout,omitempty"` // seconds, default 30 +} + +// Matcher groups hooks under a tool-name regex. +type Matcher struct { + Matcher string `json:"matcher,omitempty"` + Hooks []Command `json:"hooks"` +} + +// Config maps an event to its matchers. +type Config map[string][]Matcher + +// FileConfig is the on-disk shape. +type FileConfig struct { + Hooks Config `json:"hooks"` +} + +// Result is the aggregated outcome of a hook event. +type Result struct { + Blocked bool + Reason string + ModifiedInput map[string]any + Output string + SystemMessage string +} + +type hookResponse struct { + Decision string `json:"decision"` + Reason string `json:"reason"` + ModifiedInput map[string]any `json:"modified_input"` + SystemMessage string `json:"system_message"` + Output string `json:"output"` +} + +// Runner executes hooks for one workspace. +type Runner struct { + config Config + workspace string + // Timeout is the fallback when a hook does not set one. + Timeout time.Duration +} + +// NewRunner builds a runner. A nil config is valid and disables hooks. +func NewRunner(cfg Config, workspace string) *Runner { + return &Runner{config: cfg, workspace: workspace, Timeout: 30 * time.Second} +} + +// Load merges hook files in order. Missing files are ignored. +func Load(paths ...string) (Config, error) { + out := Config{} + for _, p := range paths { + if p == "" { + continue + } + data, err := os.ReadFile(p) + if err != nil { + if os.IsNotExist(err) { + continue + } + return out, err + } + var fc FileConfig + if err := json.Unmarshal(data, &fc); err != nil { + return out, fmt.Errorf("parse %s: %w", p, err) + } + for event, matchers := range fc.Hooks { + out[event] = append(out[event], matchers...) + } + } + return out, nil +} + +// DefaultPaths returns global and workspace hook config paths. +func DefaultPaths(home, workspace string) []string { + var paths []string + if home != "" { + paths = append(paths, filepath.Join(home, "hooks.json")) + } + if workspace != "" { + paths = append(paths, filepath.Join(workspace, ".rakit", "hooks.json")) + } + return paths +} + +// Has reports whether any hooks are configured for the event. +func (r *Runner) Has(event Event) bool { + if r == nil { + return false + } + return len(r.config[string(event)]) > 0 +} + +// Run executes all matching hooks for the event. toolName is used for the +// matcher (empty for non-tool events). payload is sent on stdin. +func (r *Runner) Run(ctx context.Context, event Event, toolName string, payload map[string]any) Result { + var result Result + if r == nil || len(r.config[string(event)]) == 0 { + return result + } + if payload == nil { + payload = map[string]any{} + } + payload["hook_event"] = string(event) + payload["workspace"] = r.workspace + if toolName != "" { + payload["tool_name"] = toolName + } + + for _, m := range r.config[string(event)] { + if !matchTool(m.Matcher, toolName) { + continue + } + for _, h := range m.Hooks { + if h.Type != "" && h.Type != "command" { + continue + } + res, blocked, reason := r.runOne(ctx, h, payload) + if res.ModifiedInput != nil { + result.ModifiedInput = res.ModifiedInput + for k, v := range res.ModifiedInput { + payload["tool_input_modified_"+k] = v + } + } + if res.Output != "" { + if result.Output != "" { + result.Output += "\n" + } + result.Output += res.Output + } + if res.SystemMessage != "" { + if result.SystemMessage != "" { + result.SystemMessage += "\n" + } + result.SystemMessage += res.SystemMessage + } + if blocked { + result.Blocked = true + result.Reason = reason + return result + } + } + } + return result +} + +func (r *Runner) runOne(ctx context.Context, h Command, payload map[string]any) (hookResponse, bool, string) { + timeout := time.Duration(h.Timeout) * time.Second + if timeout <= 0 { + timeout = r.Timeout + } + runCtx, cancel := context.WithTimeout(ctx, timeout) + defer cancel() + + cmd := exec.CommandContext(runCtx, "sh", "-c", h.Command) + cmd.Dir = r.workspace + input, _ := json.Marshal(payload) + cmd.Stdin = bytes.NewReader(input) + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + cmd.Env = append(os.Environ(), + "RAKIT_EVENT="+fmt.Sprint(payload["hook_event"]), + "RAKIT_TOOL="+fmt.Sprint(payload["tool_name"]), + "RAKIT_WORKSPACE="+r.workspace, + "RAKIT_SESSION="+fmt.Sprint(payload["session_id"]), + ) + + err := cmd.Run() + out := strings.TrimSpace(stdout.String()) + errOut := strings.TrimSpace(stderr.String()) + + var resp hookResponse + if out != "" && strings.HasPrefix(out, "{") { + _ = json.Unmarshal([]byte(out), &resp) + } + if resp.Output == "" && out != "" && !strings.HasPrefix(out, "{") { + resp.Output = out + } + if resp.Reason == "" && resp.Decision == "block" { + resp.Reason = errOut + } + + if runCtx.Err() == context.DeadlineExceeded { + return resp, true, fmt.Sprintf("hook timed out after %s", timeout) + } + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok && exitErr.ExitCode() == 2 { + reason := resp.Reason + if reason == "" { + reason = errOut + } + if reason == "" { + reason = "blocked by hook" + } + return resp, true, reason + } + return resp, false, "" + } + if resp.Decision == "block" { + reason := resp.Reason + if reason == "" { + reason = "blocked by hook" + } + return resp, true, reason + } + return resp, false, "" +} + +func matchTool(pattern, tool string) bool { + pattern = strings.TrimSpace(pattern) + if pattern == "" || pattern == "*" { + return true + } + ok, err := regexp.MatchString(pattern, tool) + return err == nil && ok +} diff --git a/hooks/hooks_test.go b/hooks/hooks_test.go new file mode 100644 index 0000000..04bb112 --- /dev/null +++ b/hooks/hooks_test.go @@ -0,0 +1,90 @@ +package hooks + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestRunBlocksOnJSONDecision(t *testing.T) { + dir := t.TempDir() + cfg := Config{ + string(PreToolUse): []Matcher{{ + Matcher: "bash|Bash", + Hooks: []Command{{Type: "command", Command: `cat >/dev/null; echo '{"decision":"block","reason":"no rm"}'`}}, + }}, + } + r := NewRunner(cfg, dir) + res := r.Run(context.Background(), PreToolUse, "bash", map[string]any{"tool_input": map[string]any{"command": "rm -rf /"}}) + if !res.Blocked || res.Reason != "no rm" { + t.Fatalf("got blocked=%v reason=%q", res.Blocked, res.Reason) + } +} + +func TestRunBlocksOnExitCode2(t *testing.T) { + dir := t.TempDir() + cfg := Config{ + string(PreToolUse): []Matcher{{ + Matcher: "", + Hooks: []Command{{Type: "command", Command: `echo "denied by policy" >&2; exit 2`}}, + }}, + } + r := NewRunner(cfg, dir) + res := r.Run(context.Background(), PreToolUse, "edit", nil) + if !res.Blocked || res.Reason != "denied by policy" { + t.Fatalf("got blocked=%v reason=%q", res.Blocked, res.Reason) + } +} + +func TestMatcherSkipsNonMatchingTool(t *testing.T) { + dir := t.TempDir() + cfg := Config{ + string(PreToolUse): []Matcher{{ + Matcher: "bash", + Hooks: []Command{{Command: `exit 2`}}, + }}, + } + r := NewRunner(cfg, dir) + res := r.Run(context.Background(), PreToolUse, "read", nil) + if res.Blocked { + t.Fatal("read should not match bash matcher") + } +} + +func TestLoadMergesFiles(t *testing.T) { + dir := t.TempDir() + a := filepath.Join(dir, "a.json") + b := filepath.Join(dir, "b.json") + if err := os.WriteFile(a, []byte(`{"hooks":{"Stop":[{"hooks":[{"command":"echo a"}]}]}}`), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(b, []byte(`{"hooks":{"Stop":[{"hooks":[{"command":"echo b"}]}]}}`), 0o644); err != nil { + t.Fatal(err) + } + cfg, err := Load(a, b) + if err != nil { + t.Fatal(err) + } + if len(cfg[string(Stop)]) != 2 { + t.Fatalf("expected 2 matchers, got %d", len(cfg[string(Stop)])) + } +} + +func TestPostToolUseOutput(t *testing.T) { + dir := t.TempDir() + cfg := Config{ + string(PostToolUse): []Matcher{{ + Matcher: "edit", + Hooks: []Command{{Command: `cat >/dev/null; echo formatted`}}, + }}, + } + r := NewRunner(cfg, dir) + res := r.Run(context.Background(), PostToolUse, "edit", nil) + if res.Blocked { + t.Fatal("post hook should not block") + } + if res.Output != "formatted" { + t.Fatalf("output = %q", res.Output) + } +} diff --git a/llm/provider.go b/llm/provider.go index 830333d..8789a63 100644 --- a/llm/provider.go +++ b/llm/provider.go @@ -8,10 +8,12 @@ import ( "fmt" "io" "log" + "math/rand" "net" "net/http" "os" "path/filepath" + "strconv" "strings" "sync" "time" @@ -963,23 +965,58 @@ func (c *LLMClient) ChatWithProviderStream( if err != nil { return nil, fmt.Errorf("marshal request: %w", err) } - req, err := http.NewRequestWithContext(reqCtx, "POST", url, bytes.NewBuffer(bodyBytes)) - if err != nil { - return nil, fmt.Errorf("create http request: %w", err) - } - req.Header.Set("Content-Type", "application/json") - if apiKey != "" { - req.Header.Set("Authorization", "Bearer "+apiKey) - } - if sessionID := SessionIDFrom(ctx); sessionID != "" { - if targetProviderID == "xai" || strings.Contains(strings.ToLower(model), "grok") { - req.Header.Set("x-grok-conv-id", sessionID) + var lastErr error + for attempt := 0; attempt < maxRetryAttempts; attempt++ { + req, err := http.NewRequestWithContext(reqCtx, "POST", url, bytes.NewBuffer(bodyBytes)) + if err != nil { + return nil, fmt.Errorf("create http request: %w", err) + } + req.Header.Set("Content-Type", "application/json") + if apiKey != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } + if sessionID := SessionIDFrom(ctx); sessionID != "" { + if targetProviderID == "xai" || strings.Contains(strings.ToLower(model), "grok") { + req.Header.Set("x-grok-conv-id", sessionID) + } + if targetProviderID == "openrouter" { + req.Header.Set("x-session-affinity", sessionID) + } + } + + resp, err := c.httpClient.Do(req) + if err != nil { + if reqCtx.Err() != nil { + return nil, err + } + lastErr = err + if attempt == maxRetryAttempts-1 { + break + } + log.Printf("[llm] request error (attempt %d/%d): %v — retrying", attempt+1, maxRetryAttempts, err) + if err := sleepContext(reqCtx, retryDelay(attempt+1, 0)); err != nil { + return nil, err + } + continue + } + + if !isRetryableStatus(resp.StatusCode) { + return resp, nil } - if targetProviderID == "openrouter" { - req.Header.Set("x-session-affinity", sessionID) + + wait := retryAfterDelay(resp.Header.Get("Retry-After")) + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + lastErr = fmt.Errorf("api error (status %d)", resp.StatusCode) + if attempt == maxRetryAttempts-1 { + break + } + log.Printf("[llm] status %d (attempt %d/%d) — backing off", resp.StatusCode, attempt+1, maxRetryAttempts) + if err := sleepContext(reqCtx, retryDelay(attempt+1, wait)); err != nil { + return nil, err } } - return c.httpClient.Do(req) + return nil, lastErr } resp, err := doRequest(reqBody) @@ -1500,3 +1537,69 @@ func MergeToolCallArgs(existing, incoming string) string { } return existing + incoming } + +// Retry policy for transient provider failures. The variables are mutable so +// tests can shorten the backoff. +var ( + maxRetryAttempts = 5 + retryBaseDelay = 250 * time.Millisecond + retryMaxDelay = 8 * time.Second +) + +func isRetryableStatus(code int) bool { + switch code { + case http.StatusRequestTimeout, http.StatusTooEarly, http.StatusTooManyRequests: + return true + } + return code >= 500 && code <= 599 +} + +func retryAfterDelay(header string) time.Duration { + header = strings.TrimSpace(header) + if header == "" { + return 0 + } + if secs, err := strconv.Atoi(header); err == nil && secs > 0 { + return time.Duration(secs) * time.Second + } + if t, err := http.ParseTime(header); err == nil { + if d := time.Until(t); d > 0 { + return d + } + } + return 0 +} + +// retryDelay returns a jittered exponential backoff. A server-provided wait +// (Retry-After) always wins, capped at 30s. +func retryDelay(attempt int, serverWait time.Duration) time.Duration { + if serverWait > 0 { + if serverWait > 30*time.Second { + return 30 * time.Second + } + return serverWait + } + if attempt < 1 { + attempt = 1 + } + d := retryBaseDelay << (attempt - 1) + if d > retryMaxDelay { + d = retryMaxDelay + } + jitter := time.Duration(rand.Int63n(int64(retryBaseDelay))) + return d + jitter +} + +func sleepContext(ctx context.Context, d time.Duration) error { + if d <= 0 { + return nil + } + t := time.NewTimer(d) + defer t.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-t.C: + return nil + } +} diff --git a/llm/retry_test.go b/llm/retry_test.go new file mode 100644 index 0000000..09876bc --- /dev/null +++ b/llm/retry_test.go @@ -0,0 +1,93 @@ +package llm + +import ( + "context" + "net/http" + "net/http/httptest" + "sync/atomic" + "testing" + "time" +) + +func TestRetryOnRateLimitThenSucceed(t *testing.T) { + oldBase, oldMax := retryBaseDelay, maxRetryAttempts + retryBaseDelay = time.Millisecond + maxRetryAttempts = 5 + defer func() { retryBaseDelay, maxRetryAttempts = oldBase, oldMax }() + + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + n := atomic.AddInt32(&calls, 1) + if n <= 2 { + w.Header().Set("Retry-After", "0") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":"rate limited"}`)) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"ok"}}],"usage":{"prompt_tokens":10,"completion_tokens":2,"total_tokens":12}}`)) + })) + defer srv.Close() + + c := &LLMClient{ + baseURL: srv.URL, + model: "test-model", + httpClient: srv.Client(), + } + resp, err := c.ChatWithProviderStream(context.Background(), "", "test-model", []LLMMessage{{Role: RoleUser, Content: "hi"}}, nil, nil, nil) + if err != nil { + t.Fatal(err) + } + if resp.Content != "ok" { + t.Fatalf("content = %q", resp.Content) + } + if got := atomic.LoadInt32(&calls); got != 3 { + t.Fatalf("calls = %d, want 3", got) + } +} + +func TestRetryExhaustedReturnsError(t *testing.T) { + oldBase, oldMax := retryBaseDelay, maxRetryAttempts + retryBaseDelay = time.Millisecond + maxRetryAttempts = 3 + defer func() { retryBaseDelay, maxRetryAttempts = oldBase, oldMax }() + + var calls int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&calls, 1) + w.WriteHeader(http.StatusServiceUnavailable) + _, _ = w.Write([]byte(`{"error":"unavailable"}`)) + })) + defer srv.Close() + + c := &LLMClient{baseURL: srv.URL, model: "test-model", httpClient: srv.Client()} + _, err := c.ChatWithProviderStream(context.Background(), "", "test-model", []LLMMessage{{Role: RoleUser, Content: "hi"}}, nil, nil, nil) + if err == nil { + t.Fatal("expected error after exhausting retries") + } + if got := atomic.LoadInt32(&calls); got != 3 { + t.Fatalf("calls = %d, want 3", got) + } +} + +func TestIsRetryableStatus(t *testing.T) { + for _, code := range []int{408, 425, 429, 500, 502, 503, 504} { + if !isRetryableStatus(code) { + t.Fatalf("%d should be retryable", code) + } + } + for _, code := range []int{200, 400, 401, 403, 404, 422} { + if isRetryableStatus(code) { + t.Fatalf("%d should not be retryable", code) + } + } +} + +func TestRetryAfterDelay(t *testing.T) { + if d := retryAfterDelay("2"); d != 2*time.Second { + t.Fatalf("got %s", d) + } + if d := retryAfterDelay(""); d != 0 { + t.Fatalf("got %s", d) + } +} diff --git a/permission/permission.go b/permission/permission.go new file mode 100644 index 0000000..6ed6cf8 --- /dev/null +++ b/permission/permission.go @@ -0,0 +1,313 @@ +// Package permission implements the tool permission pipeline: collaboration +// modes, allow/deny/ask rules and per-tool decisions. It is the live engine +// counterpart of the ported ruleset in harness/permission.go. +package permission + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "regexp" + "strings" +) + +// Mode is the collaboration mode for a session. +type Mode string + +const ( + ModePlan Mode = "plan" + ModeBuild Mode = "build" + ModeEdit Mode = "edit" + ModeYolo Mode = "yolo" + ModeAuto Mode = "auto" +) + +// Action is the result of a permission check. +type Action string + +const ( + ActionAllow Action = "allow" + ActionDeny Action = "deny" + ActionAsk Action = "ask" +) + +// Rule matches a tool and optionally a content subject (command, path, url…). +type Rule struct { + Tool string `json:"tool"` + Pattern string `json:"pattern,omitempty"` +} + +// UnmarshalJSON accepts either "Bash(git push:*)" or +// {"tool":"Bash","pattern":"git push:*}. +func (r *Rule) UnmarshalJSON(data []byte) error { + var s string + if err := json.Unmarshal(data, &s); err == nil { + *r = ParseRule(s) + return nil + } + type alias Rule + var a alias + if err := json.Unmarshal(data, &a); err != nil { + return err + } + *r = Rule(a) + return nil +} + +// MarshalJSON renders a rule back to the compact string form. +func (r Rule) MarshalJSON() ([]byte, error) { + return json.Marshal(r.String()) +} + +func (r Rule) String() string { + if r.Pattern == "" { + return r.Tool + } + return fmt.Sprintf("%s(%s)", r.Tool, r.Pattern) +} + +// ParseRule parses "Bash(git push:*)" / "Edit" into a Rule. +func ParseRule(s string) Rule { + s = strings.TrimSpace(s) + open := strings.IndexByte(s, '(') + if open <= 0 || !strings.HasSuffix(s, ")") { + return Rule{Tool: s} + } + return Rule{Tool: strings.TrimSpace(s[:open]), Pattern: s[open+1 : len(s)-1]} +} + +// Ruleset groups allow/deny/ask rules. +type Ruleset struct { + Allow []Rule `json:"allow,omitempty"` + Deny []Rule `json:"deny,omitempty"` + Ask []Rule `json:"ask,omitempty"` +} + +// Settings is the on-disk permission configuration. +type Settings struct { + Mode Mode `json:"mode,omitempty"` + Permissions Ruleset `json:"permissions"` + DisallowedTools []string `json:"disallowedTools,omitempty"` +} + +// ToolMeta is the subset of tool metadata the permission pipeline needs. +type ToolMeta struct { + Name string + ReadOnly bool + Destructive bool + NeedsApproval bool + RequiresUserInteraction bool + SideEffectScope string + IsMCP bool +} + +// Context is the per-session permission state. +type Context struct { + Mode Mode + PlanEnabled bool + Rules Ruleset + SessionAllow []Rule + Disallowed []string +} + +// Decision is the pipeline result for one tool call. +type Decision struct { + Action Action + Message string + AlwaysAsk bool +} + +// Check evaluates one tool call. Ordering mirrors the reference pipeline: +// plan transitions, user interaction, approval capability, yolo, auto, +// disallowed, project deny/ask, plan restriction, project allow, mode rules. +func Check(tool string, meta ToolMeta, input map[string]any, s Context) Decision { + if tool == "EnterPlanMode" || tool == "enter_plan_mode" { + return Decision{Action: ActionAllow} + } + if (tool == "ExitPlanMode" || tool == "exit_plan_mode") && !s.PlanEnabled && s.Mode != ModePlan { + return Decision{Action: ActionDeny, Message: "plan mode not active"} + } + + // Hard denials always win, in every mode. + if contains(s.Disallowed, tool) { + return Decision{Action: ActionDeny, Message: "tool is disallowed"} + } + if _, ok := MatchRules(s.Rules.Deny, tool, input); ok { + return Decision{Action: ActionDeny, Message: "denied by project rule"} + } + + if meta.RequiresUserInteraction { + // The `ask` tool is itself the user-interaction surface; it must run + // so it can pause the session and present the question. + if tool == "ask" { + return Decision{Action: ActionAllow} + } + if s.Mode == ModeAuto { + return Decision{Action: ActionDeny, Message: "mode.auto requires explicit approval"} + } + return Decision{Action: ActionAsk} + } + + // Plan mode: read-only (and non-destructive MCP / session-scoped) tools only. + if s.PlanEnabled || s.Mode == ModePlan { + if meta.ReadOnly && !meta.Destructive { + return Decision{Action: ActionAllow} + } + if meta.SideEffectScope == "session" && !meta.Destructive { + return Decision{Action: ActionAllow} + } + if meta.IsMCP && !meta.Destructive { + return Decision{Action: ActionAllow} + } + return Decision{Action: ActionDeny, Message: "plan mode allows read-only tools only"} + } + + if s.Mode == ModeYolo { + return Decision{Action: ActionAllow} + } + if s.Mode == ModeAuto { + if meta.ReadOnly && !meta.Destructive { + return Decision{Action: ActionAllow} + } + return Decision{Action: ActionDeny, Message: "mode.auto requires explicit approval"} + } + + if _, ok := MatchRules(s.SessionAllow, tool, input); ok { + return Decision{Action: ActionAllow} + } + if _, ok := MatchRules(s.Rules.Ask, tool, input); ok { + return Decision{Action: ActionAsk} + } + if _, ok := MatchRules(s.Rules.Allow, tool, input); ok { + return Decision{Action: ActionAllow} + } + + if meta.ReadOnly && !meta.Destructive { + return Decision{Action: ActionAllow} + } + + if s.Mode == ModeEdit { + switch tool { + case "Write", "Edit", "MultiEdit", "NotebookEdit", "write", "edit", "multiedit": + return Decision{Action: ActionAllow} + } + } + if meta.NeedsApproval || meta.Destructive || meta.SideEffectScope == "workspace" { + return Decision{Action: ActionAsk, AlwaysAsk: meta.NeedsApproval} + } + return Decision{Action: ActionAllow} +} + +// MatchRules returns the first matching rule. +func MatchRules(rules []Rule, tool string, input map[string]any) (Rule, bool) { + for _, r := range rules { + if !toolMatches(r.Tool, tool) { + continue + } + if r.Pattern == "" { + return r, true + } + subject := subjectOf(input) + if subject == "" { + continue + } + if contentMatches(r.Pattern, subject) { + return r, true + } + } + return Rule{}, false +} + +func toolMatches(ruleTool, tool string) bool { + if strings.EqualFold(ruleTool, tool) { + return true + } + // An Edit rule also covers Write, matching the reference behaviour. + if strings.EqualFold(ruleTool, "Edit") && strings.EqualFold(tool, "Write") { + return true + } + return false +} + +func subjectOf(input map[string]any) string { + for _, k := range []string{"command", "url", "file_path", "path", "pattern", "patch_text", "cmd"} { + if v, ok := input[k].(string); ok && v != "" { + return v + } + } + return "" +} + +func contentMatches(rule, subject string) bool { + if strings.HasSuffix(rule, ":*") { + prefix := strings.TrimSuffix(rule, ":*") + return subject == prefix || strings.HasPrefix(subject, prefix+" ") || strings.HasPrefix(subject, prefix+"\t") + } + if strings.HasSuffix(rule, "**") { + prefix := strings.TrimSuffix(rule, "**") + if strings.HasSuffix(prefix, "/") { + prefix = strings.TrimSuffix(prefix, "/") + } + return subject == prefix || strings.HasPrefix(subject, prefix+"/") || strings.HasPrefix(subject, prefix) + } + if strings.Contains(rule, "*") { + pattern := regexp.QuoteMeta(rule) + pattern = strings.ReplaceAll(pattern, `\*`, ".*") + ok, err := regexp.MatchString("^"+pattern+"$", subject) + return err == nil && ok + } + return rule == subject +} + +func contains(list []string, s string) bool { + for _, x := range list { + if strings.EqualFold(x, s) { + return true + } + } + return false +} + +// LoadSettings reads settings files in order and merges them. Missing files +// are ignored; later files override mode and extend the rule lists. +func LoadSettings(paths ...string) (Settings, error) { + var out Settings + for _, p := range paths { + if p == "" { + continue + } + data, err := os.ReadFile(p) + if err != nil { + if os.IsNotExist(err) { + continue + } + return out, err + } + var s Settings + if err := json.Unmarshal(data, &s); err != nil { + return out, fmt.Errorf("parse %s: %w", p, err) + } + if s.Mode != "" { + out.Mode = s.Mode + } + out.Permissions.Allow = append(out.Permissions.Allow, s.Permissions.Allow...) + out.Permissions.Deny = append(out.Permissions.Deny, s.Permissions.Deny...) + out.Permissions.Ask = append(out.Permissions.Ask, s.Permissions.Ask...) + out.DisallowedTools = append(out.DisallowedTools, s.DisallowedTools...) + } + return out, nil +} + +// DefaultSettingsPaths returns the global and workspace settings paths in +// merge order. +func DefaultSettingsPaths(home, workspace string) []string { + var paths []string + if home != "" { + paths = append(paths, filepath.Join(home, "settings.json")) + } + if workspace != "" { + paths = append(paths, filepath.Join(workspace, ".rakit", "settings.json")) + } + return paths +} diff --git a/permission/permission_test.go b/permission/permission_test.go new file mode 100644 index 0000000..da4e4c9 --- /dev/null +++ b/permission/permission_test.go @@ -0,0 +1,122 @@ +package permission + +import ( + "os" + "path/filepath" + "testing" +) + +func meta(readOnly, needsApproval, destructive bool) ToolMeta { + return ToolMeta{ReadOnly: readOnly, NeedsApproval: needsApproval, Destructive: destructive, SideEffectScope: "workspace"} +} + +func TestCheckModes(t *testing.T) { + read := meta(true, false, false) + write := meta(false, true, false) + write.SideEffectScope = "workspace" + + tests := []struct { + name string + mode Mode + tool string + m ToolMeta + want Action + }{ + {"build allows read", ModeBuild, "read", read, ActionAllow}, + {"build asks write", ModeBuild, "write", write, ActionAsk}, + {"yolo allows write", ModeYolo, "write", write, ActionAllow}, + {"plan denies write", ModePlan, "write", write, ActionDeny}, + {"plan allows read", ModePlan, "read", read, ActionAllow}, + {"auto denies write", ModeAuto, "write", write, ActionDeny}, + {"edit allows edit", ModeEdit, "edit", write, ActionAllow}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := Check(tt.tool, tt.m, map[string]any{"file_path": "a.go"}, Context{Mode: tt.mode, PlanEnabled: tt.mode == ModePlan}) + if got.Action != tt.want { + t.Fatalf("got %s, want %s (%s)", got.Action, tt.want, got.Message) + } + }) + } +} + +func TestRulesPrecedence(t *testing.T) { + write := meta(false, true, false) + ctx := Context{ + Mode: ModeBuild, + Rules: Ruleset{ + Deny: []Rule{ParseRule("Bash(rm -rf:*)")}, + Allow: []Rule{ParseRule("Bash(go test:*)")}, + }, + } + tests := []struct { + cmd string + want Action + }{ + {"rm -rf /", ActionDeny}, + {"go test ./...", ActionAllow}, + {"git push", ActionAsk}, + } + for _, tt := range tests { + got := Check("bash", ToolMeta{NeedsApproval: true, SideEffectScope: "workspace"}, map[string]any{"command": tt.cmd}, ctx) + if got.Action != tt.want { + t.Fatalf("%q: got %s, want %s", tt.cmd, got.Action, tt.want) + } + } + _ = write +} + +func TestParseRule(t *testing.T) { + tests := map[string]Rule{ + "Bash": {Tool: "Bash"}, + "Bash(git push:*)": {Tool: "Bash", Pattern: "git push:*"}, + "Edit(src/**)": {Tool: "Edit", Pattern: "src/**"}, + " Bash(npm test) ": {Tool: "Bash", Pattern: "npm test"}, + "not a rule (broken)": {Tool: "not a rule", Pattern: "broken"}, + } + for in, want := range tests { + if got := ParseRule(in); got != want { + t.Fatalf("%q: got %+v, want %+v", in, got, want) + } + } +} + +func TestLoadSettingsMerge(t *testing.T) { + dir := t.TempDir() + global := filepath.Join(dir, "settings.json") + project := filepath.Join(dir, "project.json") + if err := os.WriteFile(global, []byte(`{"mode":"build","permissions":{"allow":["Read"]}}`), 0o644); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(project, []byte(`{"permissions":{"deny":["Bash(rm -rf:*)"]},"disallowedTools":["web_fetch"]}`), 0o644); err != nil { + t.Fatal(err) + } + s, err := LoadSettings(global, project) + if err != nil { + t.Fatal(err) + } + if s.Mode != ModeBuild { + t.Fatalf("mode = %q", s.Mode) + } + if len(s.Permissions.Allow) != 1 || len(s.Permissions.Deny) != 1 { + t.Fatalf("rules not merged: %+v", s.Permissions) + } + if len(s.DisallowedTools) != 1 || s.DisallowedTools[0] != "web_fetch" { + t.Fatalf("disallowed not merged: %v", s.DisallowedTools) + } +} + +func TestWildcardAndPrefix(t *testing.T) { + if !contentMatches("src/**", "src/a/b.go") { + t.Fatal("** should match nested path") + } + if !contentMatches("go test:*", "go test ./...") { + t.Fatal("prefix should match") + } + if contentMatches("go test:*", "gotest") { + t.Fatal("prefix must be a word boundary") + } + if !contentMatches("*.go", "main.go") { + t.Fatal("* wildcard should match") + } +} diff --git a/tools/core_tools.go b/tools/core_tools.go index d3a7272..bf42530 100644 --- a/tools/core_tools.go +++ b/tools/core_tools.go @@ -439,6 +439,9 @@ func (r *Registry) registerCoreTools(searchMgr searchAPI) { // Fast web fetch and search r.Register(webFetchTool()) r.Register(webSearchTool()) + // Plan mode transitions (handled by the agent loop, tool schemas expose them). + r.Register(EnterPlanModeTool()) + r.Register(ExitPlanModeTool()) } // readMultipleTool batches file reads into one tool call so the agent doesn't diff --git a/tools/metadata.go b/tools/metadata.go new file mode 100644 index 0000000..1db324d --- /dev/null +++ b/tools/metadata.go @@ -0,0 +1,112 @@ +package tools + +import ( + "strings" + "time" +) + +// Default timeouts for tool categories. +const ( + DefaultToolTimeout = 120 * time.Second + ShellToolTimeout = 300 * time.Second + MCPToolTimeout = 300 * time.Second +) + +// Tool classification metadata. It lives in one table so the scheduler, the +// permission pipeline and the agent loop cannot drift apart. +type ToolMetadata struct { + ReadOnly bool + Destructive bool + ConcurrencySafe *bool + NeedsApproval bool + RequiresUserInteraction bool + SideEffectScope string // none|session|workspace + Timeout time.Duration + MaxOutputBytes int + Category string // file|shell|search|web|skill|plugin|task|mcp|plan +} + +// MetadataFor returns the classification for a tool name. Unknown tools are +// treated as workspace-mutating and requiring approval (fail closed). +func MetadataFor(name string) ToolMetadata { + n := strings.ToLower(name) + switch n { + case "read", "read_file", "view_file", "cat", "ls", "list_dir", "read_multiple", + "read_directory_files", "search", "search_workspace", "rg", "grep", "ripgrep", + "find", "glob", "index", "git_status", "diff", "list_skills", "list_plugins", + "list_subagents", "get_subagent_output", "list_schedules", "task_output", + "task_list", "recall_memory", "load_skill", "skill", "call_skill", "run_skill", + "invoke_skill": + return ToolMetadata{ReadOnly: true, SideEffectScope: "none", Timeout: DefaultToolTimeout, Category: "search"} + case "web_fetch", "fetch", "http_get", "read_url", "web_search", "search_web", "google", "ddg": + return ToolMetadata{ReadOnly: true, SideEffectScope: "none", Timeout: DefaultToolTimeout, Category: "web"} + case "ask": + return ToolMetadata{ReadOnly: true, RequiresUserInteraction: true, SideEffectScope: "none", Category: "task"} + case "write", "write_file", "create_file", "edit", "edit_file", "replace_file_content", + "multiedit", "multi_edit": + return ToolMetadata{NeedsApproval: true, SideEffectScope: "workspace", Timeout: DefaultToolTimeout, Category: "file"} + case "bash", "run_shell", "exec", "run_command", "background_shell": + return ToolMetadata{Destructive: true, NeedsApproval: true, SideEffectScope: "workspace", Timeout: ShellToolTimeout, MaxOutputBytes: 100 << 10, Category: "shell"} + case "python_eval", "python", "eval", "run_code": + return ToolMetadata{Destructive: true, NeedsApproval: true, SideEffectScope: "workspace", Timeout: ShellToolTimeout, Category: "shell"} + case "todo", "update_tasks", "task", "invoke_subagent", "select_subagent", + "schedule", "cancel_schedule", "enter_plan_mode", "exit_plan_mode": + return ToolMetadata{SideEffectScope: "session", Timeout: DefaultToolTimeout, Category: "task"} + case "create_skill", "create_plugin", "register_plugin", "toggle_plugin": + return ToolMetadata{NeedsApproval: true, SideEffectScope: "workspace", Timeout: DefaultToolTimeout, Category: "plugin"} + case "learn_memory": + return ToolMetadata{SideEffectScope: "session", Timeout: DefaultToolTimeout, Category: "task"} + } + // MCP tools are registered as server/tool or mcp__server__tool. + if strings.Contains(name, "/") || strings.HasPrefix(n, "mcp__") { + return ToolMetadata{NeedsApproval: true, SideEffectScope: "workspace", Timeout: MCPToolTimeout, Category: "mcp"} + } + return ToolMetadata{NeedsApproval: true, SideEffectScope: "workspace", Timeout: DefaultToolTimeout} +} + +// applyMetadata fills unset classification fields on a spec. +func applyMetadata(spec *ToolSpec) { + md := MetadataFor(spec.Name) + if !spec.ReadOnly { + spec.ReadOnly = md.ReadOnly + } + if !spec.Destructive { + spec.Destructive = md.Destructive + } + if spec.ConcurrencySafe == nil { + spec.ConcurrencySafe = md.ConcurrencySafe + } + if !spec.NeedsApproval { + spec.NeedsApproval = md.NeedsApproval + } + if !spec.RequiresUserInteraction { + spec.RequiresUserInteraction = md.RequiresUserInteraction + } + if spec.SideEffectScope == "" { + spec.SideEffectScope = md.SideEffectScope + } + if spec.Timeout == 0 { + spec.Timeout = md.Timeout + } + if spec.MaxOutputBytes == 0 { + spec.MaxOutputBytes = md.MaxOutputBytes + } + if spec.Category == "" { + spec.Category = md.Category + } +} + +// MetaForSpec converts a ToolSpec to metadata (used by the permission gate). +func MetaForSpec(spec ToolSpec) ToolMetadata { + return ToolMetadata{ + ReadOnly: spec.ReadOnly, + Destructive: spec.Destructive, + ConcurrencySafe: spec.ConcurrencySafe, + NeedsApproval: spec.NeedsApproval, + RequiresUserInteraction: spec.RequiresUserInteraction, + SideEffectScope: spec.SideEffectScope, + Timeout: spec.Timeout, + MaxOutputBytes: spec.MaxOutputBytes, + Category: spec.Category, + } +} diff --git a/tools/plan_mode.go b/tools/plan_mode.go new file mode 100644 index 0000000..3696847 --- /dev/null +++ b/tools/plan_mode.go @@ -0,0 +1,35 @@ +package tools + +import "context" + +// EnterPlanModeTool switches the session into plan mode: read-only tools only, +// and the model produces a plan artifact before exiting. +func EnterPlanModeTool() ToolSpec { + return ToolSpec{ + Name: "enter_plan_mode", + Cost: "cheap", + Description: "Switch to plan mode. In plan mode only read-only tools are allowed; investigate and write a plan, then call exit_plan_mode for approval.", + Parameters: map[string]any{"type": "object", "properties": map[string]any{}}, + Handler: func(ctx context.Context, args map[string]any) (any, error) { + return toolResult(map[string]any{"mode": "plan", "status": "enabled"}), nil + }, + } +} + +// ExitPlanModeTool requests approval of the plan and returns to build mode. +func ExitPlanModeTool() ToolSpec { + return ToolSpec{ + Name: "exit_plan_mode", + Cost: "cheap", + Description: "Present the current plan for approval and return to build mode. Call only after the plan is complete.", + Parameters: map[string]any{ + "type": "object", + "properties": map[string]any{ + "plan": map[string]any{"type": "string", "description": "The plan summary to present to the user"}, + }, + }, + Handler: func(ctx context.Context, args map[string]any) (any, error) { + return toolResult(map[string]any{"mode": "build", "status": "exit_requested"}), nil + }, + } +} diff --git a/tools/registry.go b/tools/registry.go index 618db06..3bbbd80 100644 --- a/tools/registry.go +++ b/tools/registry.go @@ -9,6 +9,7 @@ import ( "sort" "strings" "sync" + "time" "github.com/Haslab-dev/RakitKode/llm" "github.com/Haslab-dev/RakitKode/search" @@ -26,7 +27,19 @@ type ToolSpec struct { Order int // Cost is one of "cheap" (1pt), "medium" (3pt), "high" (10pt). Used for // the per-turn tool budget the agent sees in its system prompt. - Cost string + Cost string + // Classification used by the scheduler, the permission pipeline and the + // agent loop. Populated by applyMetadata when not set explicitly. + ReadOnly bool + Destructive bool + ConcurrencySafe *bool + NeedsApproval bool + RequiresUserInteraction bool + SideEffectScope string + Timeout time.Duration + MaxOutputBytes int + Category string + Handler ToolHandler IsMCP bool ServerName string @@ -147,6 +160,7 @@ func findExecutable(name string, candidatePaths []string) string { func (r *Registry) Register(spec ToolSpec) { r.mu.Lock() defer r.mu.Unlock() + applyMetadata(&spec) if spec.Order == 0 { r.nextOrder++ spec.Order = r.nextOrder