fix(agent): route image messages through image_model

pkg/agent/instance.go

@@ -39,6 +39,8 @@ type AgentInstance struct {
 	Subagents                 *config.SubagentsConfig
 	SkillsFilter              []string
 	Candidates                []providers.FallbackCandidate
+	ImageModel                string
+	ImageCandidates           []providers.FallbackCandidate
 
 	// Router is non-nil when model routing is configured and the light model
 	// was successfully resolved. It scores each incoming message and decides
@@ -195,6 +197,20 @@ func NewAgentInstance(
 
 	candidates := providers.ResolveCandidatesWithLookup(modelCfg, defaults.Provider, resolveFromModelList)
 
+	imageModel := strings.TrimSpace(defaults.ImageModel)
+	var imageCandidates []providers.FallbackCandidate
+	if imageModel != "" {
+		imageModelCfg := providers.ModelConfig{
+			Primary:   imageModel,
+			Fallbacks: defaults.ImageModelFallbacks,
+		}
+		imageCandidates = providers.ResolveCandidatesWithLookup(
+			imageModelCfg,
+			defaults.Provider,
+			resolveFromModelList,
+		)
+	}
+
 	// Model routing setup: pre-resolve light model candidates at creation time
 	// to avoid repeated model_list lookups on every incoming message.
 	var router *routing.Router
@@ -234,6 +250,8 @@ func NewAgentInstance(
 		Subagents:                 subagents,
 		SkillsFilter:              skillsFilter,
 		Candidates:                candidates,
+		ImageModel:                imageModel,
+		ImageCandidates:           imageCandidates,
 		Router:                    router,
 		LightCandidates:           lightCandidates,
 	}

pkg/agent/instance_test.go

@@ -1,10 +1,14 @@
 package agent
 
 import (
+	"context"
 	"os"
+	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/sipeed/picoclaw/pkg/config"
+	"github.com/sipeed/picoclaw/pkg/media"
 )
 
 func TestNewAgentInstance_UsesDefaultsTemperatureAndMaxTokens(t *testing.T) {
@@ -160,3 +164,128 @@ func TestNewAgentInstance_ResolveCandidatesFromModelListAlias(t *testing.T) {
 		})
 	}
 }
+func TestNewAgentInstance_ResolveImageCandidatesFromModelListAlias(t *testing.T) {
+	tmpDir, err := os.MkdirTemp("", "agent-instance-test-*")
+	if err != nil {
+		t.Fatalf("Failed to create temp dir: %v", err)
+	}
+	defer os.RemoveAll(tmpDir)
+
+	cfg := &config.Config{
+		Agents: config.AgentsConfig{
+			Defaults: config.AgentDefaults{
+				Workspace:           tmpDir,
+				Model:               "text-main",
+				ImageModel:          "vision-main",
+				ImageModelFallbacks: []string{"vision-backup"},
+			},
+		},
+		ModelList: []config.ModelConfig{
+			{
+				ModelName: "vision-main",
+				Model:     "gemini/gemini-2.5-flash-lite",
+			},
+			{
+				ModelName: "vision-backup",
+				Model:     "anthropic/claude-3-7-sonnet",
+			},
+		},
+	}
+
+	agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
+
+	if agent.ImageModel != "vision-main" {
+		t.Fatalf("ImageModel = %q, want %q", agent.ImageModel, "vision-main")
+	}
+	if len(agent.ImageCandidates) != 2 {
+		t.Fatalf("len(ImageCandidates) = %d, want 2", len(agent.ImageCandidates))
+	}
+	if agent.ImageCandidates[0].Provider != "gemini" || agent.ImageCandidates[0].Model != "gemini-2.5-flash-lite" {
+		t.Fatalf("first image candidate = %+v, want gemini/gemini-2.5-flash-lite", agent.ImageCandidates[0])
+	}
+	if agent.ImageCandidates[1].Provider != "anthropic" || agent.ImageCandidates[1].Model != "claude-3-7-sonnet" {
+		t.Fatalf("second image candidate = %+v, want anthropic/claude-3-7-sonnet", agent.ImageCandidates[1])
+	}
+}
+
+func TestNewAgentInstance_AllowsMediaTempDirForReadListAndExec(t *testing.T) {
+	workspace := t.TempDir()
+	mediaDir := media.TempDir()
+	if err := os.MkdirAll(mediaDir, 0o700); err != nil {
+		t.Fatalf("MkdirAll(mediaDir) error = %v", err)
+	}
+
+	mediaFile, err := os.CreateTemp(mediaDir, "instance-tool-*.txt")
+	if err != nil {
+		t.Fatalf("CreateTemp(mediaDir) error = %v", err)
+	}
+	mediaPath := mediaFile.Name()
+	if _, err := mediaFile.WriteString("attachment content"); err != nil {
+		mediaFile.Close()
+		t.Fatalf("WriteString(mediaFile) error = %v", err)
+	}
+	if err := mediaFile.Close(); err != nil {
+		t.Fatalf("Close(mediaFile) error = %v", err)
+	}
+	t.Cleanup(func() { _ = os.Remove(mediaPath) })
+
+	cfg := &config.Config{
+		Agents: config.AgentsConfig{
+			Defaults: config.AgentDefaults{
+				Workspace:           workspace,
+				Model:               "test-model",
+				RestrictToWorkspace: true,
+			},
+		},
+		Tools: config.ToolsConfig{
+			ReadFile: config.ReadFileToolConfig{Enabled: true},
+			ListDir:  config.ToolConfig{Enabled: true},
+			Exec: config.ExecConfig{
+				ToolConfig:         config.ToolConfig{Enabled: true},
+				EnableDenyPatterns: true,
+				AllowRemote:        true,
+			},
+		},
+	}
+
+	agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
+
+	readTool, ok := agent.Tools.Get("read_file")
+	if !ok {
+		t.Fatal("read_file tool not registered")
+	}
+	readResult := readTool.Execute(context.Background(), map[string]any{"path": mediaPath})
+	if readResult.IsError {
+		t.Fatalf("read_file should allow media temp dir, got: %s", readResult.ForLLM)
+	}
+	if !strings.Contains(readResult.ForLLM, "attachment content") {
+		t.Fatalf("read_file output missing media content: %s", readResult.ForLLM)
+	}
+
+	listTool, ok := agent.Tools.Get("list_dir")
+	if !ok {
+		t.Fatal("list_dir tool not registered")
+	}
+	listResult := listTool.Execute(context.Background(), map[string]any{"path": mediaDir})
+	if listResult.IsError {
+		t.Fatalf("list_dir should allow media temp dir, got: %s", listResult.ForLLM)
+	}
+	if !strings.Contains(listResult.ForLLM, filepath.Base(mediaPath)) {
+		t.Fatalf("list_dir output missing media file: %s", listResult.ForLLM)
+	}
+
+	execTool, ok := agent.Tools.Get("exec")
+	if !ok {
+		t.Fatal("exec tool not registered")
+	}
+	execResult := execTool.Execute(context.Background(), map[string]any{
+		"command":     "cat " + filepath.Base(mediaPath),
+		"working_dir": mediaDir,
+	})
+	if execResult.IsError {
+		t.Fatalf("exec should allow media temp dir, got: %s", execResult.ForLLM)
+	}
+	if !strings.Contains(execResult.ForLLM, "attachment content") {
+		t.Fatalf("exec output missing media content: %s", execResult.ForLLM)
+	}
+}

pkg/agent/loop.go

@@ -982,7 +982,7 @@ func (al *AgentLoop) runLLMIteration(
 	// selectCandidates evaluates routing once and the decision is sticky for
 	// all tool-follow-up iterations within the same turn so that a multi-step
 	// tool chain doesn't switch models mid-way through.
-	activeCandidates, activeModel := al.selectCandidates(agent, opts.UserMessage, messages)
+	activeCandidates, activeModel, useImageFallback := al.selectCandidates(agent, opts.UserMessage, messages)
 
 	for iteration < agent.MaxIterations {
 		iteration++
@@ -1040,13 +1040,19 @@ func (al *AgentLoop) runLLMIteration(
 
 		callLLM := func() (*providers.LLMResponse, error) {
 			if len(activeCandidates) > 1 && al.fallback != nil {
-				fbResult, fbErr := al.fallback.Execute(
+				runCandidate := func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
-					ctx,
+					return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
-					activeCandidates,
+				}
-					func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
+
-						return agent.Provider.Chat(ctx, messages, providerToolDefs, model, llmOpts)
+				var (
-					},
+					fbResult *providers.FallbackResult
+					fbErr    error
 				)
+				if useImageFallback {
+					fbResult, fbErr = al.fallback.ExecuteImage(ctx, activeCandidates, runCandidate)
+				} else {
+					fbResult, fbErr = al.fallback.Execute(ctx, activeCandidates, runCandidate)
+				}
 				if fbErr != nil {
 					return nil, fbErr
 				}
@@ -1387,9 +1393,18 @@ func (al *AgentLoop) selectCandidates(
 	agent *AgentInstance,
 	userMsg string,
 	history []providers.Message,
-) (candidates []providers.FallbackCandidate, model string) {
+) (candidates []providers.FallbackCandidate, model string, useImageFallback bool) {
+	if hasImageMedia(history) && len(agent.ImageCandidates) > 0 {
+		logger.InfoCF("agent", "Image model selected",
+			map[string]any{
+				"agent_id":    agent.ID,
+				"image_model": agent.ImageModel,
+			})
+		return agent.ImageCandidates, agent.ImageModel, true
+	}
+
 	if agent.Router == nil || len(agent.LightCandidates) == 0 {
-		return agent.Candidates, agent.Model
+		return agent.Candidates, agent.Model, false
 	}
 
 	_, usedLight, score := agent.Router.SelectModel(userMsg, history, agent.Model)
@@ -1400,7 +1415,7 @@ func (al *AgentLoop) selectCandidates(
 				"score":     score,
 				"threshold": agent.Router.Threshold(),
 			})
-		return agent.Candidates, agent.Model
+		return agent.Candidates, agent.Model, false
 	}
 
 	logger.InfoCF("agent", "Model routing: light model selected",
@@ -1410,7 +1425,18 @@ func (al *AgentLoop) selectCandidates(
 			"score":       score,
 			"threshold":   agent.Router.Threshold(),
 		})
-	return agent.LightCandidates, agent.Router.LightModel()
+	return agent.LightCandidates, agent.Router.LightModel(), false
+}
+
+func hasImageMedia(messages []providers.Message) bool {
+	for _, msg := range messages {
+		for _, ref := range msg.Media {
+			if strings.HasPrefix(strings.ToLower(ref), "data:image/") {
+				return true
+			}
+		}
+	}
+	return false
 }
 
 // maybeSummarize triggers summarization if the session history exceeds thresholds.

pkg/agent/loop_test.go

@@ -88,6 +88,85 @@ func TestRecordLastChatID(t *testing.T) {
 	}
 }
 
+func TestSelectCandidates_UsesImageModelWhenImageMediaPresent(t *testing.T) {
+	tmpDir := t.TempDir()
+	cfg := &config.Config{
+		Agents: config.AgentsConfig{
+			Defaults: config.AgentDefaults{
+				Workspace:  tmpDir,
+				ModelName:  "text-main",
+				ImageModel: "vision-main",
+			},
+		},
+		ModelList: []config.ModelConfig{
+			{ModelName: "text-main", Model: "openai/gpt-5.4"},
+			{ModelName: "vision-main", Model: "gemini/gemini-2.5-flash-lite"},
+		},
+	}
+
+	agent := NewAgentInstance(nil, &cfg.Agents.Defaults, cfg, &mockProvider{})
+	candidates, model, useImageFallback := (&AgentLoop{}).selectCandidates(
+		agent,
+		"describe this image",
+		[]providers.Message{
+			{Role: "user", Content: "describe this image", Media: []string{"data:image/png;base64,AAAA"}},
+		},
+	)
+
+	if !useImageFallback {
+		t.Fatal("expected image fallback to be selected")
+	}
+	if model != "vision-main" {
+		t.Fatalf("model = %q, want %q", model, "vision-main")
+	}
+	if len(candidates) != 1 {
+		t.Fatalf("len(candidates) = %d, want 1", len(candidates))
+	}
+	if candidates[0].Provider != "gemini" || candidates[0].Model != "gemini-2.5-flash-lite" {
+		t.Fatalf("candidate = %+v, want gemini/gemini-2.5-flash-lite", candidates[0])
+	}
+}
+
+func TestRunAgentLoop_UsesImageModelForImageMessages(t *testing.T) {
+	tmpDir := t.TempDir()
+	cfg := &config.Config{
+		Agents: config.AgentsConfig{
+			Defaults: config.AgentDefaults{
+				Workspace:         tmpDir,
+				ModelName:         "text-main",
+				ImageModel:        "vision-main",
+				MaxTokens:         4096,
+				MaxToolIterations: 5,
+			},
+		},
+		ModelList: []config.ModelConfig{
+			{ModelName: "text-main", Model: "openai/gpt-5.4"},
+			{ModelName: "vision-main", Model: "gemini/gemini-2.5-flash-lite"},
+		},
+	}
+
+	msgBus := bus.NewMessageBus()
+	provider := &mockProvider{}
+	al := NewAgentLoop(cfg, msgBus, provider)
+	agent := al.registry.GetDefaultAgent()
+
+	_, err := al.runAgentLoop(context.Background(), agent, processOptions{
+		SessionKey:      "image-session",
+		Channel:         "telegram",
+		ChatID:          "chat-1",
+		UserMessage:     "describe this image",
+		Media:           []string{"data:image/png;base64,AAAA"},
+		DefaultResponse: "fallback",
+		SendResponse:    false,
+	})
+	if err != nil {
+		t.Fatalf("runAgentLoop returned error: %v", err)
+	}
+	if provider.lastModel != "vision-main" {
+		t.Fatalf("provider lastModel = %q, want %q", provider.lastModel, "vision-main")
+	}
+}
+
 func TestNewAgentLoop_StateInitialized(t *testing.T) {
 	// Create temp workspace
 	tmpDir, err := os.MkdirTemp("", "agent-test-*")

pkg/agent/mock_provider_test.go

@@ -6,7 +6,9 @@ import (
 	"github.com/sipeed/picoclaw/pkg/providers"
 )
 
-type mockProvider struct{}
+type mockProvider struct {
+	lastModel string
+}
 
 func (m *mockProvider) Chat(
 	ctx context.Context,
@@ -15,6 +17,7 @@ func (m *mockProvider) Chat(
 	model string,
 	opts map[string]any,
 ) (*providers.LLMResponse, error) {
+	m.lastModel = model
 	return &providers.LLMResponse{
 		Content:   "Mock response",
 		ToolCalls: []providers.ToolCall{},