diff --git a/README.md b/README.md index 22fabd0b..ba8c9769 100644 --- a/README.md +++ b/README.md @@ -291,6 +291,10 @@ Features adopted from open-source agent projects. All are off by default unless | Atomic install transactions | `internal/installtxn` | Cross-process staged install/remove with rollback. Wired into skill install (atomic `SKILL.md` publish) | | Stale-lock reclaim | `internal/lockutil` | Race-correct atomic reclaim of O_EXCL lock files with live-restore (ready for O_EXCL lock sites) | | Test command discovery | `internal/testrunner` | Auto-detect test/verify commands (Go/npm/bun/pnpm/yarn/pytest/cargo) and parse runner output into structured results. Wired into `hawk verify` | +| Circuit breaker | `internal/circuitbreaker` | Closed/open/half-open retry-storm protection with cooldown — generalized primitive (auto-compact style: 3 failures → cooldown → half-open probe) | +| Smart turn routing | `internal/smartrouting` | Deterministic simple/strong turn classifier with fail-toward-strong safety (cheap model for trivial turns) | +| Conversation arc | `internal/conversationarc` | Durable sidecar memory of goals/decisions/milestones/phase with a byte-stable model-visible summary | +| Relevance pruning | `internal/relevanceprune` | Token-budgeted context pruning that keeps relevant history while preserving recent turns/tool calls/errors | ## Usage diff --git a/internal/circuitbreaker/circuitbreaker.go b/internal/circuitbreaker/circuitbreaker.go new file mode 100644 index 00000000..cd342e05 --- /dev/null +++ b/internal/circuitbreaker/circuitbreaker.go @@ -0,0 +1,123 @@ +// Package circuitbreaker provides a small, self-contained circuit breaker used +// to protect repeated operations from retry storms. It is the Go port of +// OpenClaude's auto-compact circuit breaker: after a threshold of consecutive +// failures, the breaker "opens" and skips the operation for a cooldown window, +// then re-arms in a half-open state so a single success closes it and a single +// failure re-opens it. Hawk uses it to stop runaway auto-compaction (an +// irrecoverable prompt_too_long would otherwise retry thousands of times). +package circuitbreaker + +import ( + "sync" + "time" +) + +// State describes the breaker's lifecycle. +type State int + +const ( + // Closed allows calls; a success keeps it closed. + Closed State = iota + // Open skips calls until the cooldown elapses. + Open + // HalfOpen allows a probe call after the cooldown; one failure re-opens. + HalfOpen +) + +// Breaker is a concurrency-safe circuit breaker. The zero value is not usable; +// use New. +type Breaker struct { + maxConsecutiveFailures int + cooldown time.Duration + + mu sync.Mutex + consecutiveFailures int + nextRetryAt time.Time + lastFailureAt time.Time +} + +// New returns a breaker that opens after maxConsecutiveFailures (must be >= 1) +// consecutive failures, skipping calls for cooldown. +func New(maxConsecutiveFailures int, cooldown time.Duration) *Breaker { + if maxConsecutiveFailures < 1 { + maxConsecutiveFailures = 1 + } + if cooldown < 0 { + cooldown = 0 + } + return &Breaker{ + maxConsecutiveFailures: maxConsecutiveFailures, + cooldown: cooldown, + } +} + +// Decision is the outcome of ShouldAllow. +type Decision struct { + // Allow reports whether the caller may proceed with the operation. + Allow bool + // WasHalfOpen is true when the cooldown had elapsed and a probe was admitted + // (so the caller knows one more failure will re-open the breaker). + WasHalfOpen bool + // EffectiveConsecutiveFailures is the count to use for a subsequent failure + // record (the half-open probe is seeded one below the threshold). + EffectiveConsecutiveFailures int +} + +// ShouldAllow consults the breaker at time now. When the breaker is open and +// the cooldown has not elapsed, it returns Allow=false (skip). When the +// cooldown has elapsed it admits a probe with WasHalfOpen=true, seeding the +// failure count one below the threshold so a single failure re-opens. +func (b *Breaker) ShouldAllow(now time.Time) Decision { + b.mu.Lock() + defer b.mu.Unlock() + if b.consecutiveFailures < b.maxConsecutiveFailures { + return Decision{Allow: true, EffectiveConsecutiveFailures: b.consecutiveFailures} + } + if b.nextRetryAt.IsZero() && !b.lastFailureAt.IsZero() { + b.nextRetryAt = b.lastFailureAt.Add(b.cooldown) + } + if !b.nextRetryAt.IsZero() && now.Before(b.nextRetryAt) { + return Decision{Allow: false, EffectiveConsecutiveFailures: b.consecutiveFailures} + } + b.consecutiveFailures = b.maxConsecutiveFailures - 1 + return Decision{Allow: true, WasHalfOpen: true, EffectiveConsecutiveFailures: b.consecutiveFailures} +} + +// RecordSuccess resets the consecutive-failure count, closing the breaker. +func (b *Breaker) RecordSuccess() { + b.mu.Lock() + defer b.mu.Unlock() + b.consecutiveFailures = 0 + b.nextRetryAt = time.Time{} + b.lastFailureAt = time.Time{} +} + +// RecordFailure increments the consecutive-failure count and records the time +// the cooldown window should start from. +func (b *Breaker) RecordFailure(now time.Time) { + b.mu.Lock() + defer b.mu.Unlock() + b.consecutiveFailures++ + b.lastFailureAt = now + b.nextRetryAt = now.Add(b.cooldown) +} + +// State reports the current logical state at time now. +func (b *Breaker) State(now time.Time) State { + b.mu.Lock() + defer b.mu.Unlock() + if b.consecutiveFailures < b.maxConsecutiveFailures { + return Closed + } + if !b.nextRetryAt.IsZero() && now.Before(b.nextRetryAt) { + return Open + } + return HalfOpen +} + +// ConsecutiveFailures returns the current consecutive-failure count. +func (b *Breaker) ConsecutiveFailures() int { + b.mu.Lock() + defer b.mu.Unlock() + return b.consecutiveFailures +} diff --git a/internal/circuitbreaker/circuitbreaker_test.go b/internal/circuitbreaker/circuitbreaker_test.go new file mode 100644 index 00000000..b0ca45a5 --- /dev/null +++ b/internal/circuitbreaker/circuitbreaker_test.go @@ -0,0 +1,88 @@ +package circuitbreaker + +import ( + "testing" + "time" +) + +func TestClosedAllowsUntilThreshold(t *testing.T) { + b := New(3, 5*time.Minute) + base := time.Now() + for i := 0; i < 3; i++ { + d := b.ShouldAllow(base) + if !d.Allow || d.WasHalfOpen { + t.Fatalf("call %d: allow=%v halfOpen=%v, want allow", i+1, d.Allow, d.WasHalfOpen) + } + b.RecordFailure(base) + } + if b.State(base) != Open { + t.Fatal("expected Open after 3 failures") + } +} + +func TestOpenSkipsUntilCooldownElapses(t *testing.T) { + b := New(3, time.Minute) + base := time.Now() + for i := 0; i < 3; i++ { + b.RecordFailure(base) + } + // Within cooldown: skipped. + if d := b.ShouldAllow(base.Add(30 * time.Second)); d.Allow { + t.Fatal("expected skip while open") + } + // After cooldown: half-open probe admitted. + d := b.ShouldAllow(base.Add(2 * time.Minute)) + if !d.Allow || !d.WasHalfOpen { + t.Fatalf("expected half-open probe, got allow=%v halfOpen=%v", d.Allow, d.WasHalfOpen) + } + if d.EffectiveConsecutiveFailures != 2 { + t.Fatalf("effective failures = %d, want 2", d.EffectiveConsecutiveFailures) + } +} + +func TestHalfOpenSingleFailureReopens(t *testing.T) { + b := New(3, time.Minute) + base := time.Now() + for i := 0; i < 3; i++ { + b.RecordFailure(base) + } + // Half-open probe admitted after cooldown, then fails once. + b.ShouldAllow(base.Add(2 * time.Minute)) + b.RecordFailure(base.Add(2 * time.Minute)) + if b.State(base.Add(2*time.Minute+1)) != Open { + t.Fatal("expected Open again after half-open failure") + } +} + +func TestHalfOpenSuccessCloses(t *testing.T) { + b := New(3, time.Minute) + base := time.Now() + for i := 0; i < 3; i++ { + b.RecordFailure(base) + } + b.ShouldAllow(base.Add(2 * time.Minute)) + b.RecordSuccess() + if b.State(base.Add(3*time.Minute)) != Closed { + t.Fatal("expected Closed after half-open success") + } + if b.ConsecutiveFailures() != 0 { + t.Fatalf("consecutive failures = %d, want 0", b.ConsecutiveFailures()) + } +} + +func TestMinThresholdFloor(t *testing.T) { + if New(0, time.Minute).maxConsecutiveFailures != 1 { + t.Fatal("threshold should floor to 1") + } +} + +func TestNegativeCooldownTreatedAsZero(t *testing.T) { + b := New(1, -time.Second) + base := time.Now() + b.RecordFailure(base) + // Zero cooldown: next call is a half-open probe immediately. + d := b.ShouldAllow(base) + if !d.Allow || !d.WasHalfOpen { + t.Fatalf("expected immediate half-open probe, got %+v", d) + } +} diff --git a/internal/conversationarc/arc.go b/internal/conversationarc/arc.go new file mode 100644 index 00000000..690fb304 --- /dev/null +++ b/internal/conversationarc/arc.go @@ -0,0 +1,315 @@ +// Package conversationarc tracks a session's durable "arc": goals, decisions, +// milestones, and a current phase, persisted to a sidecar JSON file so it can +// be summarized and injected into later turns. It is the Go port of +// OpenClaude's conversation arc (.arc.json) memory. +// +// The phase is advanced deterministically from message keywords +// (init→exploring→implementing→reviewing→completed) and goals/decisions are +// extracted from user text by lightweight heuristics. The Summary() string is +// byte-stable: volatile per-request timestamps are stripped so an unchanged +// arc does not rewrite the model-visible summary every turn. +package conversationarc + +import ( + "encoding/json" + "fmt" + "os" + "path/filepath" + "strings" + "time" +) + +// Phase is the conversation's current lifecycle stage. +type Phase string + +const ( + PhaseInit Phase = "init" + PhaseExploring Phase = "exploring" + PhaseImplementing Phase = "implementing" + PhaseReviewing Phase = "reviewing" + PhaseCompleted Phase = "completed" +) + +// GoalStatus tracks a goal's lifecycle. +type GoalStatus string + +const ( + GoalPending GoalStatus = "pending" + GoalActive GoalStatus = "active" + GoalCompleted GoalStatus = "completed" + GoalAbandoned GoalStatus = "abandoned" +) + +// Goal is a tracked objective. +type Goal struct { + ID string `json:"id"` + Description string `json:"description"` + Status GoalStatus `json:"status"` + CreatedAt int64 `json:"createdAt"` + CompletedAt *int64 `json:"completedAt,omitempty"` +} + +// Decision is a recorded decision with optional rationale. +type Decision struct { + ID string `json:"id"` + Description string `json:"description"` + Rationale string `json:"rationale,omitempty"` + Timestamp int64 `json:"timestamp"` +} + +// Milestone is a recorded achievement. +type Milestone struct { + ID string `json:"id"` + Description string `json:"description"` + AchievedAt int64 `json:"achievedAt"` +} + +// Arc is the durable conversation summary. +type Arc struct { + ID string `json:"id"` + Goals []Goal `json:"goals"` + Decisions []Decision `json:"decisions"` + Milestones []Milestone `json:"milestones"` + CurrentPhase Phase `json:"currentPhase"` + StartTime int64 `json:"startTime"` + LastUpdate int64 `json:"lastUpdateTime"` +} + +const ( + defaultCap = 50 + arcFileName = ".arc.json" + byteStableMarker = "" +) + +// phaseOrder is the monotonic phase ladder; a detected phase only advances. +var phaseOrder = []Phase{PhaseInit, PhaseExploring, PhaseImplementing, PhaseReviewing, PhaseCompleted} + +// phaseKeywords drive deterministic phase detection from a message. +var phaseKeywords = map[Phase][]string{ + PhaseInit: {"start", "begin", "help", "please"}, + PhaseExploring: {"check", "find", "look", "what", "how", "where", "show"}, + PhaseImplementing: {"write", "create", "add", "fix", "update", "modify", "implement"}, + PhaseReviewing: {"test", "review", "verify", "ensure"}, + PhaseCompleted: {"done", "complete", "finished", "ready", "good"}, +} + +// New returns an empty arc with a fresh id and the init phase. +func New() *Arc { + now := time.Now().UnixMilli() + return &Arc{ + ID: "arc", + Goals: []Goal{}, + Decisions: []Decision{}, + Milestones: []Milestone{}, + CurrentPhase: PhaseInit, + StartTime: now, + LastUpdate: now, + } +} + +// DetectPhase returns the highest phase whose keywords appear in text, or the +// current phase if none do. Callers may then advance via AdvancePhase. +func DetectPhase(text string) Phase { + lower := strings.ToLower(text) + best := -1 + for _, ph := range phaseOrder { + for _, kw := range phaseKeywords[ph] { + if strings.Contains(lower, kw) { + if idx := indexOfPhase(ph); idx > best { + best = idx + } + break + } + } + } + if best < 0 { + return "" + } + return phaseOrder[best] +} + +// AdvancePhase moves the arc to detected if it is later on the phase ladder. +// Returns true when the phase changed. +func (a *Arc) AdvancePhase(detected Phase) bool { + if detected == "" { + return false + } + cur := indexOfPhase(a.CurrentPhase) + next := indexOfPhase(detected) + if next > cur { + a.CurrentPhase = detected + a.touch() + return true + } + return false +} + +// AddGoal appends a pending goal (capped) and, if the arc is still in init, +// advances it to exploring. +func (a *Arc) AddGoal(description string) Goal { + g := Goal{ + ID: fmt.Sprintf("goal_%d", len(a.Goals)+1), + Description: description, + Status: GoalPending, + CreatedAt: time.Now().UnixMilli(), + } + a.Goals = append(a.Goals, g) + a.Goals = capGoals(a.Goals) + if a.CurrentPhase == PhaseInit { + a.CurrentPhase = PhaseExploring + } + a.touch() + return g +} + +// UpdateGoalStatus sets a goal's status by id, stamping completion time. +func (a *Arc) UpdateGoalStatus(id string, status GoalStatus) bool { + for i := range a.Goals { + if a.Goals[i].ID == id { + a.Goals[i].Status = status + if status == GoalCompleted { + now := time.Now().UnixMilli() + a.Goals[i].CompletedAt = &now + } + a.touch() + return true + } + } + return false +} + +// AddDecision appends a decision (capped). +func (a *Arc) AddDecision(description, rationale string) Decision { + d := Decision{ + ID: fmt.Sprintf("decision_%d", len(a.Decisions)+1), + Description: description, + Rationale: rationale, + Timestamp: time.Now().UnixMilli(), + } + a.Decisions = append(a.Decisions, d) + a.Decisions = capDecisions(a.Decisions) + a.touch() + return d +} + +// AddMilestone appends a milestone (capped). +func (a *Arc) AddMilestone(description string) Milestone { + m := Milestone{ + ID: fmt.Sprintf("milestone_%d", len(a.Milestones)+1), + Description: description, + AchievedAt: time.Now().UnixMilli(), + } + a.Milestones = append(a.Milestones, m) + a.Milestones = capMilestones(a.Milestones) + a.touch() + return m +} + +// Summary renders a model-visible, byte-stable arc summary. The volatile +// timestamp line is normalized so an unchanged arc produces identical output +// across turns (no prompt-cache churn). +func (a *Arc) Summary() string { + var b strings.Builder + fmt.Fprintf(&b, "conversation phase: %s\n", a.CurrentPhase) + if len(a.Goals) > 0 { + b.WriteString("goals:\n") + for _, g := range a.Goals { + fmt.Fprintf(&b, " - [%s] %s\n", g.Status, g.Description) + } + } + if len(a.Decisions) > 0 { + b.WriteString("decisions:\n") + for _, d := range a.Decisions { + line := fmt.Sprintf(" - %s", d.Description) + if d.Rationale != "" { + line += fmt.Sprintf(" (%s)", d.Rationale) + } + b.WriteString(line) + b.WriteByte('\n') + } + } + if len(a.Milestones) > 0 { + b.WriteString("milestones:\n") + for _, m := range a.Milestones { + fmt.Fprintf(&b, " - %s\n", m.Description) + } + } + // Byte-stable marker in place of a live timestamp. + fmt.Fprintf(&b, "detectedAt: %s\n", byteStableMarker) + return b.String() +} + +// Load reads an arc from dir/.arc.json; returns nil if absent or malformed. +func Load(dir string) (*Arc, error) { + path := filepath.Join(dir, arcFileName) + data, err := os.ReadFile(path) + if err != nil { + if os.IsNotExist(err) { + return nil, nil + } + return nil, err + } + var a Arc + if err := json.Unmarshal(data, &a); err != nil { + return nil, nil + } + if a.CurrentPhase == "" || a.Goals == nil { + return nil, nil + } + return &a, nil +} + +// Save persists the arc to dir/.arc.json (creating dir). Writes are +// best-effort and non-fatal, like the source. +func (a *Arc) Save(dir string) error { + a.touch() + data, err := json.MarshalIndent(a, "", " ") + if err != nil { + return err + } + if err := os.MkdirAll(dir, 0o755); err != nil { + return err + } + return os.WriteFile(filepath.Join(dir, arcFileName), data, 0o600) +} + +// Reset clears the arc's tracked state back to init. +func (a *Arc) Reset() { + a.Goals = []Goal{} + a.Decisions = []Decision{} + a.Milestones = []Milestone{} + a.CurrentPhase = PhaseInit + a.touch() +} + +func (a *Arc) touch() { a.LastUpdate = time.Now().UnixMilli() } + +func indexOfPhase(p Phase) int { + for i, ph := range phaseOrder { + if ph == p { + return i + } + } + return -1 +} + +func capGoals(goals []Goal) []Goal { + if len(goals) > defaultCap { + return goals[len(goals)-defaultCap:] + } + return goals +} + +func capDecisions(ds []Decision) []Decision { + if len(ds) > defaultCap { + return ds[len(ds)-defaultCap:] + } + return ds +} + +func capMilestones(ms []Milestone) []Milestone { + if len(ms) > defaultCap { + return ms[len(ms)-defaultCap:] + } + return ms +} diff --git a/internal/conversationarc/arc_test.go b/internal/conversationarc/arc_test.go new file mode 100644 index 00000000..4c6da431 --- /dev/null +++ b/internal/conversationarc/arc_test.go @@ -0,0 +1,130 @@ +package conversationarc + +import ( + "os" + "path/filepath" + "strings" + "testing" +) + +func TestNewArcInit(t *testing.T) { + a := New() + if a.CurrentPhase != PhaseInit { + t.Fatalf("phase = %s", a.CurrentPhase) + } + if len(a.Goals) != 0 { + t.Fatal("expected no goals") + } +} + +func TestDetectAndAdvancePhase(t *testing.T) { + a := New() + if !a.AdvancePhase(DetectPhase("find the bug and write a fix")) { + t.Fatal("expected phase advance for exploring") + } + if a.CurrentPhase != PhaseImplementing { + t.Fatalf("phase = %s, want implementing", a.CurrentPhase) + } + // Non-monotonic: a completed keyword should advance past implementing. + if !a.AdvancePhase(DetectPhase("done now")) { + t.Fatal("expected advance to completed") + } + if a.CurrentPhase != PhaseCompleted { + t.Fatalf("phase = %s, want completed", a.CurrentPhase) + } +} + +func TestPhaseDoesNotRegress(t *testing.T) { + a := New() + a.CurrentPhase = PhaseImplementing + if a.AdvancePhase(PhaseExploring) { + t.Fatal("phase should not regress") + } +} + +func TestAddGoalAdvancesFromInit(t *testing.T) { + a := New() + g := a.AddGoal("implement X") + if a.CurrentPhase != PhaseExploring { + t.Fatalf("phase = %s, want exploring", a.CurrentPhase) + } + if g.Status != GoalPending { + t.Fatalf("status = %s", g.Status) + } +} + +func TestUpdateGoalStatusCompletion(t *testing.T) { + a := New() + g := a.AddGoal("implement X") + if !a.UpdateGoalStatus(g.ID, GoalCompleted) { + t.Fatal("update failed") + } + if a.Goals[0].Status != GoalCompleted || a.Goals[0].CompletedAt == nil { + t.Fatalf("goal not completed: %+v", a.Goals[0]) + } + if a.UpdateGoalStatus("nope", GoalCompleted) { + t.Fatal("update of unknown id should fail") + } +} + +func TestSummaryByteStable(t *testing.T) { + a := New() + a.AddGoal("implement X") + a.AddDecision("use go", "stdlib") + a.AddMilestone("prototype done") + s1 := a.Summary() + s2 := a.Summary() + if s1 != s2 { + t.Fatalf("summary not byte-stable:\n%q\n%q", s1, s2) + } + if !strings.Contains(s1, "implement X") { + t.Fatalf("summary missing goal:\n%s", s1) + } +} + +func TestSaveAndLoadRoundTrip(t *testing.T) { + dir := t.TempDir() + a := New() + a.AddGoal("implement X") + a.CurrentPhase = PhaseImplementing + if err := a.Save(dir); err != nil { + t.Fatal(err) + } + loaded, err := Load(dir) + if err != nil { + t.Fatal(err) + } + if loaded == nil { + t.Fatal("loaded nil") + } + if len(loaded.Goals) != 1 || loaded.Goals[0].Description != "implement X" { + t.Fatalf("loaded goals wrong: %+v", loaded.Goals) + } + if loaded.CurrentPhase != PhaseImplementing { + t.Fatalf("loaded phase = %s", loaded.CurrentPhase) + } +} + +func TestLoadMissingReturnsNil(t *testing.T) { + loaded, err := Load(filepath.Join(t.TempDir(), "nope")) + if err != nil { + t.Fatal(err) + } + if loaded != nil { + t.Fatal("expected nil for missing arc") + } +} + +func TestLoadMalformedReturnsNil(t *testing.T) { + dir := t.TempDir() + if err := os.WriteFile(filepath.Join(dir, arcFileName), []byte("{bad json"), 0o644); err != nil { + t.Fatal(err) + } + loaded, err := Load(dir) + if err != nil { + t.Fatal(err) + } + if loaded != nil { + t.Fatal("expected nil for malformed arc") + } +} diff --git a/internal/relevanceprune/relevanceprune.go b/internal/relevanceprune/relevanceprune.go new file mode 100644 index 00000000..9b0a9694 --- /dev/null +++ b/internal/relevanceprune/relevanceprune.go @@ -0,0 +1,217 @@ +// Package relevanceprune provides deterministic, token-budgeted context +// pruning that keeps messages relevant to the current task while preserving +// recent turns, tool calls, and error messages verbatim. It is the Go port of +// OpenClaude's relevance-based context pruning — a cheap, predictable +// alternative to pure LLM summarization: score older messages by keyword +// overlap against the task context, keep the highest-relevance groups up to a +// token budget, and always retain the most recent tail. +package relevanceprune + +import ( + "sort" + "strings" + "time" +) + +// DefaultCompactTailTurns is the number of recent messages preserved verbatim. +const DefaultCompactTailTurns = 3 + +// NormalizeCompactTailTurns floors any finite value >= 1 to an integer and +// falls back to the default for everything else (0, negatives, fractions below +// 1, NaN). UI and runtime must share this single rule. +func NormalizeCompactTailTurns(value int) int { + if value >= 1 { + return value + } + return DefaultCompactTailTurns +} + +// Message is a minimal conversation message the pruner understands. +type Message struct { + Role string // "user", "assistant", or "system" + Content string + Timestamp time.Time + HasToolCall bool + IsError bool +} + +// Options configures a pruning pass. +type Options struct { + TargetTokens int + TaskContext string + PreserveRecent int + PreserveTools bool + PreserveErrors bool +} + +var stopWords = map[string]bool{ + "the": true, "and": true, "for": true, "are": true, "but": true, + "not": true, "you": true, "all": true, "can": true, "had": true, + "her": true, "was": true, "one": true, "our": true, "out": true, + "has": true, "have": true, "they": true, "will": true, "would": true, +} + +func extractKeywords(text string) map[string]bool { + words := strings.Fields(strings.ToLower(text)) + keywords := map[string]bool{} + for _, w := range words { + cleaned := strings.Map(func(r rune) rune { + if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') { + return r + } + return -1 + }, w) + if len(cleaned) > 3 && !stopWords[cleaned] { + keywords[cleaned] = true + } + } + return keywords +} + +func keywordOverlap(a, b string) float64 { + ka := extractKeywords(a) + kb := extractKeywords(b) + overlap := 0 + for k := range ka { + if kb[k] { + overlap++ + } + } + total := len(ka) + len(kb) + if total == 0 { + return 0 + } + return float64(2*overlap) / float64(total) +} + +// CalculateRelevance scores a message in [0,1]. Base 0.5, boosted by keyword +// overlap with the task context, preserved tool calls / errors, recency, and +// user role. +func CalculateRelevance(m Message, o Options) float64 { + score := 0.5 + if o.TaskContext != "" { + score += keywordOverlap(m.Content, o.TaskContext) * 0.3 + } + if m.HasToolCall && o.PreserveTools { + score += 0.25 + } + if m.IsError && o.PreserveErrors { + score += 0.3 + } + if !m.Timestamp.IsZero() && time.Since(m.Timestamp) < time.Hour { + score += 0.15 + } + if m.Role == "user" { + score += 0.1 + } + if score > 1 { + return 1 + } + return score +} + +// PruneByRelevance drops low-relevance history to fit targetTokens while +// preserving the most recent preserveRecent messages verbatim. Older messages +// are grouped by API round and scored by average relevance; the highest-scoring +// groups are kept until the token budget is exhausted, then everything is +// re-sorted chronologically. +func PruneByRelevance(messages []Message, o Options) []Message { + target := o.TargetTokens + if target <= 0 { + target = 5000 + } + preserve := NormalizeCompactTailTurns(o.PreserveRecent) + if len(messages) <= preserve { + return messages + } + recent := messages[len(messages)-preserve:] + older := messages[:len(messages)-preserve] + + groups := groupByRound(older) + scored := make([]struct { + group []Message + score float64 + }, 0, len(groups)) + for _, g := range groups { + var sum float64 + for _, m := range g { + sum += CalculateRelevance(m, o) + } + scored = append(scored, struct { + group []Message + score float64 + }{group: g, score: sum / float64(len(g))}) + } + sort.SliceStable(scored, func(i, j int) bool { + if scored[i].score != scored[j].score { + return scored[i].score > scored[j].score + } + return scored[i].group[0].Timestamp.After(scored[j].group[0].Timestamp) + }) + + result := make([]Message, 0, len(recent)) + result = append(result, recent...) + totalTokens := 0 + for _, s := range scored { + var content strings.Builder + for _, m := range s.group { + content.WriteString(m.Content) + } + tokens := RoughTokenCount(content.String()) + if totalTokens+tokens > target { + continue + } + result = append(result, s.group...) + totalTokens += tokens + } + sort.SliceStable(result, func(i, j int) bool { return result[i].Timestamp.Before(result[j].Timestamp) }) + return result +} + +// RoughTokenCount estimates tokens as chars/4, matching the source's cheap +// heuristic. +func RoughTokenCount(text string) int { + return len(text) / 4 +} + +// GetRelevanceStats summarizes a set of messages. +func GetRelevanceStats(messages []Message, o Options) (average float64, highCount, toolCalls, errors int) { + var sum float64 + for _, m := range messages { + score := CalculateRelevance(m, o) + sum += score + if score > 0.7 { + highCount++ + } + if m.HasToolCall { + toolCalls++ + } + if m.IsError { + errors++ + } + } + if len(messages) > 0 { + average = sum / float64(len(messages)) + } + return +} + +// groupByRound groups consecutive messages that share an assistant response +// round (an assistant message starts a new group unless it repeats the previous +// assistant's id — here we approximate by starting a new group at each assistant +// message). +func groupByRound(messages []Message) [][]Message { + var groups [][]Message + var current []Message + for _, m := range messages { + if m.Role == "assistant" && len(current) > 0 { + groups = append(groups, current) + current = []Message{} + } + current = append(current, m) + } + if len(current) > 0 { + groups = append(groups, current) + } + return groups +} diff --git a/internal/relevanceprune/relevanceprune_test.go b/internal/relevanceprune/relevanceprune_test.go new file mode 100644 index 00000000..5d2bd12b --- /dev/null +++ b/internal/relevanceprune/relevanceprune_test.go @@ -0,0 +1,92 @@ +package relevanceprune + +import ( + "testing" + "time" +) + +func msg(role, content string, hasTool, isErr bool, t time.Time) Message { + return Message{Role: role, Content: content, HasToolCall: hasTool, IsError: isErr, Timestamp: t} +} + +func TestNormalizeCompactTailTurns(t *testing.T) { + if NormalizeCompactTailTurns(5) != 5 { + t.Fatal("5 should stay 5") + } + if NormalizeCompactTailTurns(0) != DefaultCompactTailTurns { + t.Fatal("0 should fall back to default") + } + if NormalizeCompactTailTurns(-2) != DefaultCompactTailTurns { + t.Fatal("negative should fall back to default") + } +} + +func TestCalculateRelevanceBasics(t *testing.T) { + now := time.Now() + base := msg("user", "hello there", false, false, now) + if got := CalculateRelevance(base, Options{}); got < 0.5 || got > 1 { + t.Fatalf("base score = %v", got) + } + // Error preserved should score higher. + errMsg := msg("assistant", "boom", false, true, now) + withErr := CalculateRelevance(errMsg, Options{PreserveErrors: true}) + withoutErr := CalculateRelevance(errMsg, Options{}) + if withErr <= withoutErr { + t.Fatalf("preserved error should boost: %v vs %v", withErr, withoutErr) + } + // Task-context overlap should boost. + rel := msg("user", "refactor the payment module and update tests", false, false, now) + if CalculateRelevance(rel, Options{TaskContext: "refactor payment module tests"}) <= CalculateRelevance(rel, Options{}) { + t.Fatal("task overlap should boost relevance") + } +} + +func TestPruneByRelevancePreservesRecent(t *testing.T) { + base := time.Now() + var msgs []Message + for i := 0; i < 8; i++ { + msgs = append(msgs, msg("user", "unrelated chatter filler content here", false, false, base.Add(time.Duration(i)*time.Minute))) + } + // The last message strongly matches the task. + msgs[len(msgs)-1] = msg("user", "refactor payment module", false, false, base.Add(7*time.Minute)) + + pruned := PruneByRelevance(msgs, Options{TargetTokens: 100, TaskContext: "refactor payment module", PreserveRecent: 3, PreserveTools: true, PreserveErrors: true}) + if len(pruned) < 3 { + t.Fatalf("pruned below preserve-recent floor: %d", len(pruned)) + } + // The most recent message must survive verbatim. + last := pruned[len(pruned)-1] + if last.Content != "refactor payment module" { + t.Fatalf("most recent message not preserved: %q", last.Content) + } +} + +func TestPruneShortListUnchanged(t *testing.T) { + msgs := []Message{ + msg("user", "a", false, false, time.Now()), + msg("assistant", "b", false, false, time.Now()), + } + out := PruneByRelevance(msgs, Options{TargetTokens: 5000, PreserveRecent: 3}) + if len(out) != 2 { + t.Fatalf("short list should be unchanged, got %d", len(out)) + } +} + +func TestGetRelevanceStats(t *testing.T) { + now := time.Now() + msgs := []Message{ + msg("user", "refactor payment module deeply", false, false, now), + msg("assistant", "done", true, false, now), + msg("assistant", "failed", false, true, now), + } + avg, high, tools, errs := GetRelevanceStats(msgs, Options{TaskContext: "refactor payment module", PreserveTools: true, PreserveErrors: true}) + if tools != 1 || errs != 1 { + t.Fatalf("tools=%d errors=%d", tools, errs) + } + if high == 0 { + t.Fatal("expected at least one high-relevance message") + } + if avg <= 0 || avg > 1 { + t.Fatalf("avg = %v", avg) + } +} diff --git a/internal/smartrouting/smartrouting.go b/internal/smartrouting/smartrouting.go new file mode 100644 index 00000000..1d66cded --- /dev/null +++ b/internal/smartrouting/smartrouting.go @@ -0,0 +1,165 @@ +// Package smartrouting provides a deterministic, pure classifier that decides +// whether a user turn should go to a cheap "simple" model or a strong model. +// It is the Go port of OpenClaude's smart model routing: trivial turns ("ok", +// "rename this", "what does this do?") route to a cheaper model while the +// strong model handles anything non-trivial. When in doubt it routes to the +// strong model, so the failure mode is "no savings on a cheap turn," never a +// silently degraded answer on a turn you cared about. The classifier is a pure +// function — the caller supplies config and input. +package smartrouting + +import "strings" + +// Config selects the simple and strong models and the size thresholds that +// qualify a turn as "simple". Enabled is opt-in. +type Config struct { + Enabled bool + SimpleModel string + StrongModel string + SimpleMax int // max characters to qualify as simple (default 160) + SimpleMaxWords int // max whitespace-separated words (default 28) +} + +// Complexity is the classifier's verdict. +type Complexity string + +const ( + Simple Complexity = "simple" + Strong Complexity = "strong" +) + +// Input describes the user turn to classify. +type Input struct { + UserText string + HasNonText bool // image/document or other non-text blocks + TurnNumber int // 1-indexed turn in the session +} + +// Decision is the classifier's output. +type Decision struct { + Model string + Complexity Complexity + Reason string +} + +const ( + defaultSimpleMax = 160 + defaultSimpleMaxWords = 28 +) + +// strongKeywords strongly suggest reasoning/planning/design work. +var strongKeywords = []string{ + "plan", "design", "architect", "architecture", "refactor", "debug", + "investigate", "analyze", "analyse", "implement", "optimize", "optimise", + "review", "audit", "diagnose", "root cause", "root-cause", "why does", + "why is", "how should", "why did", "propose", "trace", "reproduce", +} + +// Route decides which model to use for the turn. +func Route(input Input, cfg Config) Decision { + if !cfg.Enabled { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "smart-routing disabled"} + } + if cfg.SimpleModel == "" || cfg.StrongModel == "" { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "simpleModel or strongModel missing"} + } + if cfg.SimpleModel == cfg.StrongModel { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "simpleModel equals strongModel"} + } + + text := strings.TrimSpace(input.UserText) + + if input.HasNonText { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "contains non-text content"} + } + if text == "" { + return Decision{Model: cfg.SimpleModel, Complexity: Simple, Reason: "empty user text"} + } + if input.TurnNumber == 1 { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "first turn of session"} + } + + maxChars := cfg.SimpleMax + if maxChars <= 0 { + maxChars = defaultSimpleMax + } + maxWords := cfg.SimpleMaxWords + if maxWords <= 0 { + maxWords = defaultSimpleMaxWords + } + + if hasCode(text) { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "contains code block or inline code"} + } + if hasStrongKeyword(text) { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "contains reasoning/planning keyword"} + } + if hasMultiParagraph(text) { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "multi-paragraph input"} + } + if len(text) > maxChars { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "input exceeds char threshold"} + } + if countWords(text) > maxWords { + return Decision{Model: cfg.StrongModel, Complexity: Strong, Reason: "input exceeds word threshold"} + } + return Decision{ + Model: cfg.SimpleModel, + Complexity: Simple, + Reason: "short input", + } +} + +// hasCode detects a fenced code block or an inline backtick span. +func hasCode(text string) bool { + if strings.Contains(text, "```") { + return true + } + return strings.Contains(text, "`") +} + +// hasStrongKeyword matches reasoning/planning keywords with word boundaries, +// case-insensitive. +func hasStrongKeyword(text string) bool { + lower := strings.ToLower(text) + for _, kw := range strongKeywords { + if containsWord(lower, kw) { + return true + } + } + return false +} + +func containsWord(haystack, needle string) bool { + start := 0 + for { + idx := strings.Index(haystack[start:], needle) + if idx < 0 { + return false + } + pos := start + idx + beforeOK := pos == 0 || !isWordByte(haystack[pos-1]) + end := pos + len(needle) + afterOK := end >= len(haystack) || !isWordByte(haystack[end]) + if beforeOK && afterOK { + return true + } + start = pos + 1 + } +} + +func isWordByte(b byte) bool { + return b == '_' || (b >= 'a' && b <= 'z') || (b >= '0' && b <= '9') +} + +func hasMultiParagraph(text string) bool { + return strings.Contains(text, "\n\n") || strings.Contains(text, "\n \n") +} + +func countWords(text string) int { + trimmed := strings.TrimSpace(text) + if trimmed == "" { + return 0 + } + return len(strings.Fields(trimmed)) +} diff --git a/internal/smartrouting/smartrouting_test.go b/internal/smartrouting/smartrouting_test.go new file mode 100644 index 00000000..8af18872 --- /dev/null +++ b/internal/smartrouting/smartrouting_test.go @@ -0,0 +1,77 @@ +package smartrouting + +import "testing" + +func on() Config { + return Config{Enabled: true, SimpleModel: "mini", StrongModel: "main"} +} + +func TestDisabledRoutesStrong(t *testing.T) { + d := Route(Input{UserText: "ok"}, Config{Enabled: false, StrongModel: "main"}) + if d.Complexity != Strong || d.Model != "main" { + t.Fatalf("got %+v", d) + } +} + +func TestMissingOrEqualModelsRouteStrong(t *testing.T) { + if Route(Input{}, Config{Enabled: true, SimpleModel: "", StrongModel: "main"}).Complexity != Strong { + t.Fatal("missing simple should route strong") + } + if Route(Input{}, Config{Enabled: true, SimpleModel: "same", StrongModel: "same"}).Complexity != Strong { + t.Fatal("equal models should route strong") + } +} + +func TestShortChatRoutesSimple(t *testing.T) { + d := Route(Input{UserText: "ok", TurnNumber: 2}, on()) + if d.Complexity != Simple || d.Model != "mini" { + t.Fatalf("got %+v", d) + } +} + +func TestFirstTurnRoutesStrong(t *testing.T) { + d := Route(Input{UserText: "ok", TurnNumber: 1}, on()) + if d.Complexity != Strong { + t.Fatalf("first turn should route strong, got %+v", d) + } +} + +func TestNonTextRoutesStrong(t *testing.T) { + d := Route(Input{UserText: "ok", TurnNumber: 2, HasNonText: true}, on()) + if d.Complexity != Strong { + t.Fatalf("got %+v", d) + } +} + +func TestCodeRoutesStrong(t *testing.T) { + for _, s := range []string{"```go\nx\n```", "use `foo`"} { + if d := Route(Input{UserText: s, TurnNumber: 2}, on()); d.Complexity != Strong { + t.Fatalf("code %q should route strong, got %+v", s, d) + } + } +} + +func TestStrongKeywordRoutesStrong(t *testing.T) { + for _, s := range []string{"plan the refactor", "why does this fail", "root cause analysis"} { + if d := Route(Input{UserText: s, TurnNumber: 2}, on()); d.Complexity != Strong { + t.Fatalf("keyword %q should route strong, got %+v", s, d) + } + } +} + +func TestLongInputRoutesStrong(t *testing.T) { + long := make([]byte, 161) + for i := range long { + long[i] = 'a' + } + if d := Route(Input{UserText: string(long), TurnNumber: 2}, on()); d.Complexity != Strong { + t.Fatalf("long input should route strong, got %+v", d) + } +} + +func TestEmptyRoutesSimple(t *testing.T) { + d := Route(Input{TurnNumber: 2}, on()) + if d.Complexity != Simple { + t.Fatalf("empty should route simple, got %+v", d) + } +}