Merge pull request #368 from abhijit360/akamath/assign-model-to-role

Different roles to different models
2026-01-12 01:46:28 -08:00
parent a9080ed04f 833724a7ed
commit 39a51c0d14
4 changed files with 338 additions and 6 deletions
--- a/internal/config/loader.go
+++ b/internal/config/loader.go
@@ -898,6 +898,102 @@ func ResolveAgentConfigWithOverride(townRoot, rigPath, agentOverride string) (*R
 	return lookupAgentConfig(agentName, townSettings, rigSettings), agentName, nil
 }

+// ResolveRoleAgentConfig resolves the agent configuration for a specific role.
+// It checks role-specific agent assignments before falling back to the default agent.
+//
+// Resolution order:
+//  1. Rig's RoleAgents[role] - if set, look up that agent
+//  2. Town's RoleAgents[role] - if set, look up that agent
+//  3. Fall back to ResolveAgentConfig (rig's Agent → town's DefaultAgent → "claude")
+//
+// role is one of: "mayor", "deacon", "witness", "refinery", "polecat", "crew".
+// townRoot is the path to the town directory (e.g., ~/gt).
+// rigPath is the path to the rig directory (e.g., ~/gt/gastown), or empty for town-level roles.
+func ResolveRoleAgentConfig(role, townRoot, rigPath string) *RuntimeConfig {
+	// Load rig settings (may be nil for town-level roles like mayor/deacon)
+	var rigSettings *RigSettings
+	if rigPath != "" {
+		var err error
+		rigSettings, err = LoadRigSettings(RigSettingsPath(rigPath))
+		if err != nil {
+			rigSettings = nil
+		}
+	}
+
+	// Load town settings
+	townSettings, err := LoadOrCreateTownSettings(TownSettingsPath(townRoot))
+	if err != nil {
+		townSettings = NewTownSettings()
+	}
+
+	// Load custom agent registries
+	_ = LoadAgentRegistry(DefaultAgentRegistryPath(townRoot))
+	if rigPath != "" {
+		_ = LoadRigAgentRegistry(RigAgentRegistryPath(rigPath))
+	}
+
+	// Check rig's RoleAgents first
+	if rigSettings != nil && rigSettings.RoleAgents != nil {
+		if agentName, ok := rigSettings.RoleAgents[role]; ok && agentName != "" {
+			return lookupAgentConfig(agentName, townSettings, rigSettings)
+		}
+	}
+
+	// Check town's RoleAgents
+	if townSettings.RoleAgents != nil {
+		if agentName, ok := townSettings.RoleAgents[role]; ok && agentName != "" {
+			return lookupAgentConfig(agentName, townSettings, rigSettings)
+		}
+	}
+
+	// Fall back to existing resolution (rig's Agent → town's DefaultAgent → "claude")
+	return ResolveAgentConfig(townRoot, rigPath)
+}
+
+// ResolveRoleAgentName returns the agent name that would be used for a specific role.
+// This is useful for logging and diagnostics.
+// Returns the agent name and whether it came from role-specific configuration.
+func ResolveRoleAgentName(role, townRoot, rigPath string) (agentName string, isRoleSpecific bool) {
+	// Load rig settings
+	var rigSettings *RigSettings
+	if rigPath != "" {
+		var err error
+		rigSettings, err = LoadRigSettings(RigSettingsPath(rigPath))
+		if err != nil {
+			rigSettings = nil
+		}
+	}
+
+	// Load town settings
+	townSettings, err := LoadOrCreateTownSettings(TownSettingsPath(townRoot))
+	if err != nil {
+		townSettings = NewTownSettings()
+	}
+
+	// Check rig's RoleAgents first
+	if rigSettings != nil && rigSettings.RoleAgents != nil {
+		if name, ok := rigSettings.RoleAgents[role]; ok && name != "" {
+			return name, true
+		}
+	}
+
+	// Check town's RoleAgents
+	if townSettings.RoleAgents != nil {
+		if name, ok := townSettings.RoleAgents[role]; ok && name != "" {
+			return name, true
+		}
+	}
+
+	// Fall back to existing resolution
+	if rigSettings != nil && rigSettings.Agent != "" {
+		return rigSettings.Agent, false
+	}
+	if townSettings.DefaultAgent != "" {
+		return townSettings.DefaultAgent, false
+	}
+	return "claude", false
+}
+
 // lookupAgentConfig looks up an agent by name.
 // Checks rig-level custom agents first, then town's custom agents, then built-in presets from agents.go.
 func lookupAgentConfig(name string, townSettings *TownSettings, rigSettings *RigSettings) *RuntimeConfig {
--- a/internal/config/loader_test.go
+++ b/internal/config/loader_test.go
@@ -1750,3 +1750,207 @@ func TestLookupAgentConfigWithRigSettings(t *testing.T) {
 		})
 	}
 }
+
+func TestResolveRoleAgentConfig(t *testing.T) {
+	t.Parallel()
+	townRoot := t.TempDir()
+	rigPath := filepath.Join(townRoot, "testrig")
+
+	// Create town settings with role-specific agents
+	townSettings := NewTownSettings()
+	townSettings.DefaultAgent = "claude"
+	townSettings.RoleAgents = map[string]string{
+		"mayor":   "claude", // mayor uses default claude
+		"witness": "gemini", // witness uses gemini
+		"polecat": "codex",  // polecats use codex
+	}
+	townSettings.Agents = map[string]*RuntimeConfig{
+		"claude-haiku": {
+			Command: "claude",
+			Args:    []string{"--model", "haiku", "--dangerously-skip-permissions"},
+		},
+	}
+	if err := SaveTownSettings(TownSettingsPath(townRoot), townSettings); err != nil {
+		t.Fatalf("SaveTownSettings: %v", err)
+	}
+
+	// Create rig settings that override some roles
+	rigSettings := NewRigSettings()
+	rigSettings.Agent = "gemini" // default for this rig
+	rigSettings.RoleAgents = map[string]string{
+		"witness": "claude-haiku", // override witness to use haiku
+	}
+	if err := SaveRigSettings(RigSettingsPath(rigPath), rigSettings); err != nil {
+		t.Fatalf("SaveRigSettings: %v", err)
+	}
+
+	t.Run("rig RoleAgents overrides town RoleAgents", func(t *testing.T) {
+		rc := ResolveRoleAgentConfig("witness", townRoot, rigPath)
+		// Should get claude-haiku from rig's RoleAgents
+		if rc.Command != "claude" {
+			t.Errorf("Command = %q, want %q", rc.Command, "claude")
+		}
+		cmd := rc.BuildCommand()
+		if !strings.Contains(cmd, "--model haiku") {
+			t.Errorf("BuildCommand() = %q, should contain --model haiku", cmd)
+		}
+	})
+
+	t.Run("town RoleAgents used when rig has no override", func(t *testing.T) {
+		rc := ResolveRoleAgentConfig("polecat", townRoot, rigPath)
+		// Should get codex from town's RoleAgents (rig doesn't override polecat)
+		if rc.Command != "codex" {
+			t.Errorf("Command = %q, want %q", rc.Command, "codex")
+		}
+	})
+
+	t.Run("falls back to default agent when role not in RoleAgents", func(t *testing.T) {
+		rc := ResolveRoleAgentConfig("crew", townRoot, rigPath)
+		// crew is not in any RoleAgents, should use rig's default agent (gemini)
+		if rc.Command != "gemini" {
+			t.Errorf("Command = %q, want %q", rc.Command, "gemini")
+		}
+	})
+
+	t.Run("town-level role (no rigPath) uses town RoleAgents", func(t *testing.T) {
+		rc := ResolveRoleAgentConfig("mayor", townRoot, "")
+		// mayor is in town's RoleAgents
+		if rc.Command != "claude" {
+			t.Errorf("Command = %q, want %q", rc.Command, "claude")
+		}
+	})
+}
+
+func TestResolveRoleAgentName(t *testing.T) {
+	t.Parallel()
+	townRoot := t.TempDir()
+	rigPath := filepath.Join(townRoot, "testrig")
+
+	// Create town settings with role-specific agents
+	townSettings := NewTownSettings()
+	townSettings.DefaultAgent = "claude"
+	townSettings.RoleAgents = map[string]string{
+		"witness": "gemini",
+		"polecat": "codex",
+	}
+	if err := SaveTownSettings(TownSettingsPath(townRoot), townSettings); err != nil {
+		t.Fatalf("SaveTownSettings: %v", err)
+	}
+
+	// Create rig settings
+	rigSettings := NewRigSettings()
+	rigSettings.Agent = "amp"
+	rigSettings.RoleAgents = map[string]string{
+		"witness": "cursor", // override witness
+	}
+	if err := SaveRigSettings(RigSettingsPath(rigPath), rigSettings); err != nil {
+		t.Fatalf("SaveRigSettings: %v", err)
+	}
+
+	t.Run("rig role-specific agent", func(t *testing.T) {
+		name, isRoleSpecific := ResolveRoleAgentName("witness", townRoot, rigPath)
+		if name != "cursor" {
+			t.Errorf("name = %q, want %q", name, "cursor")
+		}
+		if !isRoleSpecific {
+			t.Error("isRoleSpecific = false, want true")
+		}
+	})
+
+	t.Run("town role-specific agent", func(t *testing.T) {
+		name, isRoleSpecific := ResolveRoleAgentName("polecat", townRoot, rigPath)
+		if name != "codex" {
+			t.Errorf("name = %q, want %q", name, "codex")
+		}
+		if !isRoleSpecific {
+			t.Error("isRoleSpecific = false, want true")
+		}
+	})
+
+	t.Run("falls back to rig default agent", func(t *testing.T) {
+		name, isRoleSpecific := ResolveRoleAgentName("crew", townRoot, rigPath)
+		if name != "amp" {
+			t.Errorf("name = %q, want %q", name, "amp")
+		}
+		if isRoleSpecific {
+			t.Error("isRoleSpecific = true, want false")
+		}
+	})
+
+	t.Run("falls back to town default agent when no rig path", func(t *testing.T) {
+		name, isRoleSpecific := ResolveRoleAgentName("refinery", townRoot, "")
+		if name != "claude" {
+			t.Errorf("name = %q, want %q", name, "claude")
+		}
+		if isRoleSpecific {
+			t.Error("isRoleSpecific = true, want false")
+		}
+	})
+}
+
+func TestRoleAgentsRoundTrip(t *testing.T) {
+	t.Parallel()
+	dir := t.TempDir()
+	townSettingsPath := filepath.Join(dir, "settings", "config.json")
+	rigSettingsPath := filepath.Join(dir, "rig", "settings", "config.json")
+
+	// Test TownSettings with RoleAgents
+	t.Run("town settings with role_agents", func(t *testing.T) {
+		original := NewTownSettings()
+		original.RoleAgents = map[string]string{
+			"mayor":   "claude-opus",
+			"witness": "claude-haiku",
+			"polecat": "claude-sonnet",
+		}
+
+		if err := SaveTownSettings(townSettingsPath, original); err != nil {
+			t.Fatalf("SaveTownSettings: %v", err)
+		}
+
+		loaded, err := LoadOrCreateTownSettings(townSettingsPath)
+		if err != nil {
+			t.Fatalf("LoadOrCreateTownSettings: %v", err)
+		}
+
+		if len(loaded.RoleAgents) != 3 {
+			t.Errorf("RoleAgents count = %d, want 3", len(loaded.RoleAgents))
+		}
+		if loaded.RoleAgents["mayor"] != "claude-opus" {
+			t.Errorf("RoleAgents[mayor] = %q, want %q", loaded.RoleAgents["mayor"], "claude-opus")
+		}
+		if loaded.RoleAgents["witness"] != "claude-haiku" {
+			t.Errorf("RoleAgents[witness] = %q, want %q", loaded.RoleAgents["witness"], "claude-haiku")
+		}
+		if loaded.RoleAgents["polecat"] != "claude-sonnet" {
+			t.Errorf("RoleAgents[polecat] = %q, want %q", loaded.RoleAgents["polecat"], "claude-sonnet")
+		}
+	})
+
+	// Test RigSettings with RoleAgents
+	t.Run("rig settings with role_agents", func(t *testing.T) {
+		original := NewRigSettings()
+		original.RoleAgents = map[string]string{
+			"witness": "gemini",
+			"crew":    "codex",
+		}
+
+		if err := SaveRigSettings(rigSettingsPath, original); err != nil {
+			t.Fatalf("SaveRigSettings: %v", err)
+		}
+
+		loaded, err := LoadRigSettings(rigSettingsPath)
+		if err != nil {
+			t.Fatalf("LoadRigSettings: %v", err)
+		}
+
+		if len(loaded.RoleAgents) != 2 {
+			t.Errorf("RoleAgents count = %d, want 2", len(loaded.RoleAgents))
+		}
+		if loaded.RoleAgents["witness"] != "gemini" {
+			t.Errorf("RoleAgents[witness] = %q, want %q", loaded.RoleAgents["witness"], "gemini")
+		}
+		if loaded.RoleAgents["crew"] != "codex" {
+			t.Errorf("RoleAgents[crew] = %q, want %q", loaded.RoleAgents["crew"], "codex")
+		}
+	})
+}
--- a/internal/config/types.go
+++ b/internal/config/types.go
@@ -49,6 +49,13 @@ type TownSettings struct {
 	// Values override or extend the built-in presets.
 	// Example: {"gemini": {"command": "/custom/path/to/gemini"}}
 	Agents map[string]*RuntimeConfig `json:"agents,omitempty"`
+
+	// RoleAgents maps role names to agent aliases for per-role model selection.
+	// Keys are role names: "mayor", "deacon", "witness", "refinery", "polecat", "crew".
+	// Values are agent names (built-in presets or custom agents defined in Agents).
+	// This allows cost optimization by using different models for different roles.
+	// Example: {"mayor": "claude-opus", "witness": "claude-haiku", "polecat": "claude-sonnet"}
+	RoleAgents map[string]string `json:"role_agents,omitempty"`
 }

 // NewTownSettings creates a new TownSettings with defaults.
@@ -58,6 +65,7 @@ func NewTownSettings() *TownSettings {
 		Version:      CurrentTownSettingsVersion,
 		DefaultAgent: "claude",
 		Agents:       make(map[string]*RuntimeConfig),
+		RoleAgents:   make(map[string]string),
 	}
 }

@@ -209,6 +217,13 @@ type RigSettings struct {
 	// Similar to TownSettings.Agents but applies to this rig only.
 	// Allows per-rig custom agents for polecats and crew members.
 	Agents map[string]*RuntimeConfig `json:"agents,omitempty"`
+
+	// RoleAgents maps role names to agent aliases for per-role model selection.
+	// Keys are role names: "witness", "refinery", "polecat", "crew".
+	// Values are agent names (built-in presets or custom agents).
+	// Overrides TownSettings.RoleAgents for this specific rig.
+	// Example: {"witness": "claude-haiku", "polecat": "claude-sonnet"}
+	RoleAgents map[string]string `json:"role_agents,omitempty"`
 }

 // CrewConfig represents crew workspace settings for a rig.
--- a/internal/daemon/lifecycle.go
+++ b/internal/daemon/lifecycle.go
@@ -458,7 +458,7 @@ func (d *Daemon) getNeedsPreSync(config *beads.RoleConfig, parsed *ParsedIdentit
 }

 // getStartCommand determines the startup command for an agent.
-// Uses role bead config if available, falls back to hardcoded defaults.
+// Uses role bead config if available, then role-based agent selection, then hardcoded defaults.
 func (d *Daemon) getStartCommand(roleConfig *beads.RoleConfig, parsed *ParsedIdentity) string {
 	// If role bead has explicit config, use it
 	if roleConfig != nil && roleConfig.StartCommand != "" {
@@ -471,16 +471,33 @@ func (d *Daemon) getStartCommand(roleConfig *beads.RoleConfig, parsed *ParsedIde
 		rigPath = filepath.Join(d.config.TownRoot, parsed.RigName)
 	}

-	// Default command for all agents - use runtime config
-	defaultCmd := "exec " + config.GetRuntimeCommand(rigPath)
-	runtimeConfig := config.LoadRuntimeConfig(rigPath)
+	// Use role-based agent resolution for per-role model selection
+	runtimeConfig := config.ResolveRoleAgentConfig(parsed.RoleType, d.config.TownRoot, rigPath)
+
+	// Build default command using the role-resolved runtime config
+	defaultCmd := "exec " + runtimeConfig.BuildCommand()
 	if runtimeConfig.Session != nil && runtimeConfig.Session.SessionIDEnv != "" {
 		defaultCmd = config.PrependEnv(defaultCmd, map[string]string{"GT_SESSION_ID_ENV": runtimeConfig.Session.SessionIDEnv})
 	}

-	// Polecats need environment variables set in the command
+	// Polecats and crew need environment variables set in the command
 	if parsed.RoleType == "polecat" {
-		return config.BuildPolecatStartupCommand(parsed.RigName, parsed.AgentName, rigPath, "")
+		envVars := config.AgentEnvSimple("polecat", parsed.RigName, parsed.AgentName)
+		// Add GT_ROOT and session ID env if available
+		envVars["GT_ROOT"] = d.config.TownRoot
+		if runtimeConfig.Session != nil && runtimeConfig.Session.SessionIDEnv != "" {
+			envVars["GT_SESSION_ID_ENV"] = runtimeConfig.Session.SessionIDEnv
+		}
+		return config.PrependEnv("exec "+runtimeConfig.BuildCommand(), envVars)
+	}
+
+	if parsed.RoleType == "crew" {
+		envVars := config.AgentEnvSimple("crew", parsed.RigName, parsed.AgentName)
+		envVars["GT_ROOT"] = d.config.TownRoot
+		if runtimeConfig.Session != nil && runtimeConfig.Session.SessionIDEnv != "" {
+			envVars["GT_SESSION_ID_ENV"] = runtimeConfig.Session.SessionIDEnv
+		}
+		return config.PrependEnv("exec "+runtimeConfig.BuildCommand(), envVars)
 	}

 	return defaultCmd