4 files changed, 121 insertions, 52 deletions
diff --git a/internal/executor/classifier.go b/internal/executor/classifier.go
index efd2acb..7a474b6 100644
--- a/internal/executor/classifier.go
+++ b/internal/executor/classifier.go
@@ -24,12 +24,10 @@ type Classifier struct {
 }
 
 const classificationPrompt = `
-You are a task classifier for Claudomator.
-Given a task description and system status, select the best agent (claude or gemini) and model to use.
+You are a model selector for Claudomator.
+The agent has already been chosen by the load balancer. Your ONLY job is to select the best model for that agent.
 
-Agent Types:
-- claude: Best for complex coding, reasoning, and tool use.
-- gemini: Best for large context, fast reasoning, and multimodal tasks.
+REQUIRED agent: %s
 
 Available Models:
 Claude:
@@ -38,38 +36,30 @@ Claude:
 - claude-haiku-4-5-20251001 (fast, cheap, use for simple tasks)
 
 Gemini:
-- gemini-2.5-flash-lite (fastest, most efficient, best for simple tasks)
+- gemini-2.5-flash-lite (fastest, most efficient, best for simple/trivial tasks)
 - gemini-2.5-flash (fast, balanced)
-- gemini-2.5-pro (most powerful Gemini, larger context)
+- gemini-2.5-pro (most powerful, use for hardest tasks only)
 
 Selection Criteria:
-- Agent: CRITICAL: You MUST select an agent where "Rate Limited: false". DO NOT select an agent where "Rate Limited: true" if any other agent is available and NOT rate limited. 
-  Check the "System Status" section below. If it says "- Agent claude: ... Rate Limited: true", you MUST NOT select claude. Use gemini instead.
-- Model: Select based on task complexity. Use powerful models (opus, pro, pro-preview) for complex reasoning/coding, flash-lite/flash/haiku for simple tasks.
+- Use powerful models (opus, pro) only for the hardest reasoning/coding tasks.
+- Use lite/haiku for simple, short, or low-stakes tasks.
+- Default to the balanced model (sonnet, flash) for everything else.
 
 Task:
 Name: %s
 Instructions: %s
 
-System Status:
-%s
-
 Respond with ONLY a JSON object:
 {
-  "agent_type": "claude" | "gemini",
+  "agent_type": "%s",
   "model": "model-name",
   "reason": "brief reason"
 }
 `
 
-func (c *Classifier) Classify(ctx context.Context, taskName, instructions string, status SystemStatus) (*Classification, error) {
-	statusStr := ""
-	for agent, active := range status.ActiveTasks {
-		statusStr += fmt.Sprintf("- Agent %s: %d active tasks, Rate Limited: %t\n", agent, active, status.RateLimited[agent])
-	}
-
+func (c *Classifier) Classify(ctx context.Context, taskName, instructions string, _ SystemStatus, agentType string) (*Classification, error) {
 	prompt := fmt.Sprintf(classificationPrompt,
-		taskName, instructions, statusStr,
+		agentType, taskName, instructions, agentType,
 	)
 
 	binary := c.GeminiBinaryPath
diff --git a/internal/executor/classifier_test.go b/internal/executor/classifier_test.go
index 631952f..83a9743 100644
--- a/internal/executor/classifier_test.go
+++ b/internal/executor/classifier_test.go
@@ -23,7 +23,7 @@ echo '{"response": "{\"agent_type\": \"gemini\", \"model\": \"gemini-2.5-flash-l
 		RateLimited: map[string]bool{"claude": false, "gemini": false},
 	}
 
-	cls, err := c.Classify(context.Background(), "Test Task", "Test Instructions", status)
+	cls, err := c.Classify(context.Background(), "Test Task", "Test Instructions", status, "gemini")
 	if err != nil {
 		t.Fatalf("Classify failed: %v", err)
 	}
diff --git a/internal/executor/executor.go b/internal/executor/executor.go
index c04f68e..f54773a 100644
--- a/internal/executor/executor.go
+++ b/internal/executor/executor.go
@@ -343,30 +343,66 @@ func (p *Pool) ActiveCount() int {
 	return p.active
 }
 
-func (p *Pool) execute(ctx context.Context, t *task.Task) {
-	// 1. Classification
-	if p.Classifier != nil {
-		p.mu.Lock()
-		activeTasks := make(map[string]int)
-		rateLimited := make(map[string]bool)
-		now := time.Now()
-		for agent := range p.runners {
-			activeTasks[agent] = p.activePerAgent[agent]
-			if deadline, ok := p.rateLimited[agent]; ok && now.After(deadline) {
-				delete(p.rateLimited, agent)
-			}
-			rateLimited[agent] = now.Before(p.rateLimited[agent])
+// pickAgent selects the best agent from the given SystemStatus using explicit
+// load balancing: prefer the available (non-rate-limited) agent with the fewest
+// active tasks. If all agents are rate-limited, fall back to fewest active.
+func pickAgent(status SystemStatus) string {
+	best := ""
+	bestActive := -1
+
+	// First pass: only consider non-rate-limited agents.
+	for agent, active := range status.ActiveTasks {
+		if status.RateLimited[agent] {
+			continue
 		}
-		status := SystemStatus{
-			ActiveTasks: activeTasks,
-			RateLimited: rateLimited,
+		if bestActive == -1 || active < bestActive || (active == bestActive && agent < best) {
+			best = agent
+			bestActive = active
 		}
-		p.mu.Unlock()
+	}
+	if best != "" {
+		return best
+	}
+
+	// Fallback: all rate-limited — pick least active anyway.
+	for agent, active := range status.ActiveTasks {
+		if bestActive == -1 || active < bestActive || (active == bestActive && agent < best) {
+			best = agent
+			bestActive = active
+		}
+	}
+	return best
+}
 
-		cls, err := p.Classifier.Classify(ctx, t.Name, t.Agent.Instructions, status)
+func (p *Pool) execute(ctx context.Context, t *task.Task) {
+	// 1. Load-balanced agent selection + model classification.
+	p.mu.Lock()
+	activeTasks := make(map[string]int)
+	rateLimited := make(map[string]bool)
+	now := time.Now()
+	for agent := range p.runners {
+		activeTasks[agent] = p.activePerAgent[agent]
+		if deadline, ok := p.rateLimited[agent]; ok && now.After(deadline) {
+			delete(p.rateLimited, agent)
+		}
+		rateLimited[agent] = now.Before(p.rateLimited[agent])
+	}
+	status := SystemStatus{
+		ActiveTasks: activeTasks,
+		RateLimited: rateLimited,
+	}
+	p.mu.Unlock()
+
+	// Deterministically pick the agent with fewest active tasks.
+	selectedAgent := pickAgent(status)
+	if selectedAgent != "" {
+		t.Agent.Type = selectedAgent
+	}
+
+	if p.Classifier != nil {
+		cls, err := p.Classifier.Classify(ctx, t.Name, t.Agent.Instructions, status, t.Agent.Type)
 		if err == nil {
-			p.logger.Info("task classified", "taskID", t.ID, "agent", cls.AgentType, "model", cls.Model, "reason", cls.Reason)
-			t.Agent.Type = cls.AgentType
+			p.logger.Info("task classified", "taskID", t.ID, "agent", t.Agent.Type, "model", cls.Model, "reason", cls.Reason)
 			t.Agent.Model = cls.Model
 		} else {
 			p.logger.Error("classification failed", "error", err, "taskID", t.ID)
diff --git a/internal/executor/executor_test.go b/internal/executor/executor_test.go
index 0935545..9448816 100644
--- a/internal/executor/executor_test.go
+++ b/internal/executor/executor_test.go
@@ -116,6 +116,48 @@ func makeTask(id string) *task.Task {
 	}
 }
 
+func TestPickAgent_PrefersLessActiveAgent(t *testing.T) {
+	status := SystemStatus{
+		ActiveTasks: map[string]int{"claude": 3, "gemini": 1},
+		RateLimited: map[string]bool{"claude": false, "gemini": false},
+	}
+	if got := pickAgent(status); got != "gemini" {
+		t.Errorf("expected gemini (fewer active tasks), got %s", got)
+	}
+}
+
+func TestPickAgent_SkipsRateLimitedAgent(t *testing.T) {
+	status := SystemStatus{
+		ActiveTasks: map[string]int{"claude": 0, "gemini": 5},
+		RateLimited: map[string]bool{"claude": true, "gemini": false},
+	}
+	if got := pickAgent(status); got != "gemini" {
+		t.Errorf("expected gemini (claude rate limited), got %s", got)
+	}
+}
+
+func TestPickAgent_FallsBackWhenAllRateLimited(t *testing.T) {
+	status := SystemStatus{
+		ActiveTasks: map[string]int{"claude": 2, "gemini": 5},
+		RateLimited: map[string]bool{"claude": true, "gemini": true},
+	}
+	// Falls back to least active regardless of rate limit.
+	if got := pickAgent(status); got != "claude" {
+		t.Errorf("expected claude (fewer active tasks among all), got %s", got)
+	}
+}
+
+func TestPickAgent_TieBreakPrefersFirstAlpha(t *testing.T) {
+	status := SystemStatus{
+		ActiveTasks: map[string]int{"claude": 2, "gemini": 2},
+		RateLimited: map[string]bool{"claude": false, "gemini": false},
+	}
+	got := pickAgent(status)
+	if got != "claude" && got != "gemini" {
+		t.Errorf("unexpected agent %q on tie", got)
+	}
+}
+
 func TestPool_Submit_TopLevel_GoesToReady(t *testing.T) {
 	store := testStore(t)
 	runner := &mockRunner{}
@@ -995,13 +1037,17 @@ func TestHandleRunResult_SharedPath(t *testing.T) {
 	})
 }
 
-func TestPool_UnsupportedAgent(t *testing.T) {
+// TestPool_LoadBalancing_OverridesAgentType verifies that load balancing picks
+// from registered runners, overriding any pre-set Agent.Type on the task.
+func TestPool_LoadBalancing_OverridesAgentType(t *testing.T) {
 	store := testStore(t)
-	runners := map[string]Runner{"claude": &mockRunner{}}
+	runner := &mockRunner{}
+	runners := map[string]Runner{"claude": runner}
 	logger := slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelError}))
 	pool := NewPool(2, runners, store, logger)
 
-	tk := makeTask("bad-agent")
+	// Task has a non-existent agent type; load balancing should route to "claude".
+	tk := makeTask("lb-override")
 	tk.Agent.Type = "super-ai"
 	store.CreateTask(tk)
 
@@ -1010,13 +1056,10 @@ func TestPool_UnsupportedAgent(t *testing.T) {
 	}
 
 	result := <-pool.Results()
-	if result.Err == nil {
-		t.Fatal("expected error for unsupported agent")
-	}
-	if !strings.Contains(result.Err.Error(), "unsupported agent type") {
-		t.Errorf("expected 'unsupported agent type' in error, got: %v", result.Err)
+	if result.Err != nil {
+		t.Fatalf("expected success (load balancing overrides agent type), got: %v", result.Err)
 	}
-	if result.Execution.Status != "FAILED" {
-		t.Errorf("status: want FAILED, got %q", result.Execution.Status)
+	if runner.callCount() != 1 {
+		t.Errorf("expected claude runner to be called once, got %d", runner.callCount())
 	}
 }