fix(agent): route media turns to image models and embed onboard workspace

cmd/picoclaw/internal/onboard/command.go

@@ -1,14 +1,12 @@
 package onboard
 
 import (
-	"embed"
-
 	"github.com/spf13/cobra"
+
+	picoclaw "github.com/sipeed/picoclaw"
 )
 
-//go:generate go run ../../../../scripts/copydir.go ../../../../workspace ./workspace
-//go:embed workspace
-var embeddedFiles embed.FS
+var embeddedFiles = picoclaw.OnboardWorkspace
 
 func NewOnboardCommand() *cobra.Command {
 	var encrypt bool

onboard_workspace_embed.go

@@ -0,0 +1,12 @@
+package picoclaw
+
+import "embed"
+
+// OnboardWorkspace embeds the default onboarding workspace template.
+//
+// Keeping this embed at the module root lets us source files directly from the
+// tracked `workspace/` tree, instead of relying on a generated copy inside
+// `cmd/...` that may be absent in clean checkouts and CI lint runs.
+//
+//go:embed workspace
+var OnboardWorkspace embed.FS

pkg/agent/agent_test.go

@@ -4401,7 +4401,90 @@ func (p *visionUnsupportedMediaProvider) GetDefaultModel() string {
 	return "mock-fail-model"
 }
 
-func TestAgentLoop_VisionUnsupportedErrorStripsSessionMedia(t *testing.T) {
+type loadImagePlanningProvider struct {
+	path        string
+	followUpErr error
+	calls       int
+	models      []string
+}
+
+func (p *loadImagePlanningProvider) Chat(
+	ctx context.Context,
+	messages []providers.Message,
+	tools []providers.ToolDefinition,
+	model string,
+	opts map[string]any,
+) (*providers.LLMResponse, error) {
+	p.calls++
+	p.models = append(p.models, model)
+
+	if p.calls == 1 {
+		return &providers.LLMResponse{
+			Content: "Let me inspect the image.",
+			ToolCalls: []providers.ToolCall{{
+				ID:        "call_load_image_test",
+				Type:      "function",
+				Name:      "load_image",
+				Arguments: map[string]any{"path": p.path},
+			}},
+		}, nil
+	}
+
+	if p.followUpErr != nil {
+		return nil, p.followUpErr
+	}
+
+	return nil, fmt.Errorf("load_image follow-up should not be handled by the text model")
+}
+
+func (p *loadImagePlanningProvider) GetDefaultModel() string {
+	return "load-image-planner"
+}
+
+type visionAnswerProvider struct {
+	calls     int
+	models    []string
+	mediaSeen []bool
+}
+
+func (p *visionAnswerProvider) Chat(
+	ctx context.Context,
+	messages []providers.Message,
+	tools []providers.ToolDefinition,
+	model string,
+	opts map[string]any,
+) (*providers.LLMResponse, error) {
+	p.calls++
+	p.models = append(p.models, model)
+
+	hasMedia := false
+	for _, msg := range messages {
+		for _, ref := range msg.Media {
+			if strings.TrimSpace(ref) != "" {
+				hasMedia = true
+				break
+			}
+		}
+		if hasMedia {
+			break
+		}
+	}
+	p.mediaSeen = append(p.mediaSeen, hasMedia)
+	if !hasMedia {
+		return nil, fmt.Errorf("vision provider expected image media in follow-up request")
+	}
+
+	return &providers.LLMResponse{
+		Content:   "vision answer",
+		ToolCalls: []providers.ToolCall{},
+	}, nil
+}
+
+func (p *visionAnswerProvider) GetDefaultModel() string {
+	return "vision-answer-model"
+}
+
+func TestAgentLoop_VisionUnsupportedErrorReturnsClearFailure(t *testing.T) {
 	workspace := t.TempDir()
 
 	cfg := &config.Config{
@@ -4436,17 +4519,20 @@ func TestAgentLoop_VisionUnsupportedErrorStripsSessionMedia(t *testing.T) {
 		Media:      []string{"data:image/png;base64,abc123"},
 		SessionKey: sessionKey,
 	}))
-	if err != nil {
-		t.Fatalf("processMessage() error = %v", err)
+	if err == nil {
+		t.Fatal("processMessage() error = nil, want vision unsupported failure")
 	}
-	if resp != "ok" {
-		t.Fatalf("response = %q, want %q", resp, "ok")
+	if resp != "" {
+		t.Fatalf("response = %q, want empty response on error", resp)
 	}
-	if provider.calls != 2 {
-		t.Fatalf("calls = %d, want %d (fail with media, then retry without media)", provider.calls, 2)
+	if !strings.Contains(err.Error(), `active model "test-model" does not support image input`) {
+		t.Fatalf("error = %q, want clear vision unsupported guidance", err.Error())
 	}
-	if !slices.Equal(provider.mediaSeen, []bool{true, false}) {
-		t.Fatalf("mediaSeen = %v, want %v", provider.mediaSeen, []bool{true, false})
+	if provider.calls != 1 {
+		t.Fatalf("calls = %d, want %d (no retry without media)", provider.calls, 1)
+	}
+	if !slices.Equal(provider.mediaSeen, []bool{true}) {
+		t.Fatalf("mediaSeen = %v, want %v", provider.mediaSeen, []bool{true})
 	}
 
 	agent := al.registry.GetDefaultAgent()
@@ -4454,37 +4540,161 @@ func TestAgentLoop_VisionUnsupportedErrorStripsSessionMedia(t *testing.T) {
 		t.Fatal("expected default agent")
 	}
 	history := agent.Sessions.GetHistory(sessionKey)
-	for i, msg := range history {
-		if len(msg.Media) > 0 {
-			t.Fatalf("history[%d].Media = %v, want no media after stripping", i, msg.Media)
+	if len(history) == 0 {
+		t.Fatal("expected user message to remain in session history")
 	}
+	if len(history[0].Media) == 0 {
+		t.Fatalf("history[0].Media = %v, want original media preserved", history[0].Media)
+	}
+}
+
+func TestAgentLoop_LoadImageFollowUpRoutesToImageModel(t *testing.T) {
+	workspace := t.TempDir()
+	pngPath := filepath.Join(workspace, "sample.png")
+	pngBytes := []byte{
+		0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A,
+		0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52,
+		0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02,
+		0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xDE,
+	}
+	if err := os.WriteFile(pngPath, pngBytes, 0o644); err != nil {
+		t.Fatalf("WriteFile() error = %v", err)
 	}
 
-	timeoutCtx2, cancel2 := context.WithTimeout(context.Background(), responseTimeout)
-	defer cancel2()
+	cfg := &config.Config{
+		Agents: config.AgentsConfig{
+			Defaults: config.AgentDefaults{
+				Workspace:         workspace,
+				ModelName:         "text-model",
+				ImageModel:        "vision-model",
+				MaxTokens:         4096,
+				MaxToolIterations: 3,
+			},
+		},
+		Tools: config.ToolsConfig{
+			LoadImage: config.ToolConfig{Enabled: true},
+		},
+		ModelList: []*config.ModelConfig{
+			{ModelName: "text-model", Model: "openai/text-model"},
+			{ModelName: "vision-model", Model: "openai/vision-model"},
+		},
+	}
 
-	resp2, err := al.processMessage(timeoutCtx2, testInboundMessage(bus.InboundMessage{
+	msgBus := bus.NewMessageBus()
+	planner := &loadImagePlanningProvider{path: pngPath}
+	al := NewAgentLoop(cfg, msgBus, planner)
+	al.SetMediaStore(media.NewFileMediaStore())
+
+	agent := al.registry.GetDefaultAgent()
+	if agent == nil {
+		t.Fatal("expected default agent")
+	}
+	if len(agent.ImageCandidates) != 1 {
+		t.Fatalf("len(ImageCandidates) = %d, want 1", len(agent.ImageCandidates))
+	}
+
+	visionProvider := &visionAnswerProvider{}
+	agent.CandidateProviders[providers.ModelKey("openai", "vision-model")] = visionProvider
+
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), responseTimeout)
+	defer cancel()
+
+	resp, err := al.processMessage(timeoutCtx, testInboundMessage(bus.InboundMessage{
 		Context: bus.InboundContext{
 			Channel:   "telegram",
 			ChatID:    "chat1",
 			ChatType:  "direct",
 			SenderID:  "user1",
-			MessageID: "m2",
+			MessageID: "m1",
 		},
-		Content:    "hello again",
-		SessionKey: sessionKey,
+		Content:    "describe the image you load",
+		SessionKey: "agent:main:telegram:direct:user1",
 	}))
 	if err != nil {
-		t.Fatalf("processMessage() second call error = %v", err)
+		t.Fatalf("processMessage() error = %v", err)
 	}
-	if resp2 != "ok" {
-		t.Fatalf("second response = %q, want %q", resp2, "ok")
+	if resp != "vision answer" {
+		t.Fatalf("response = %q, want %q", resp, "vision answer")
 	}
-	if provider.calls != 3 {
-		t.Fatalf("calls after second turn = %d, want %d", provider.calls, 3)
+	if planner.calls != 1 {
+		t.Fatalf("planner calls = %d, want %d", planner.calls, 1)
 	}
-	if !slices.Equal(provider.mediaSeen, []bool{true, false, false}) {
-		t.Fatalf("mediaSeen = %v, want %v", provider.mediaSeen, []bool{true, false, false})
+	if visionProvider.calls != 1 {
+		t.Fatalf("visionProvider calls = %d, want %d", visionProvider.calls, 1)
+	}
+	if !slices.Equal(visionProvider.models, []string{"vision-model"}) {
+		t.Fatalf("visionProvider models = %v, want %v", visionProvider.models, []string{"vision-model"})
+	}
+	if !slices.Equal(visionProvider.mediaSeen, []bool{true}) {
+		t.Fatalf("visionProvider mediaSeen = %v, want %v", visionProvider.mediaSeen, []bool{true})
+	}
+}
+
+func TestAgentLoop_LoadImageFollowUpWithoutImageModelFailsClearly(t *testing.T) {
+	workspace := t.TempDir()
+	pngPath := filepath.Join(workspace, "sample.png")
+	pngBytes := []byte{
+		0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A,
+		0x00, 0x00, 0x00, 0x0D, 0x49, 0x48, 0x44, 0x52,
+		0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x08, 0x02,
+		0x00, 0x00, 0x00, 0x90, 0x77, 0x53, 0xDE,
+	}
+	if err := os.WriteFile(pngPath, pngBytes, 0o644); err != nil {
+		t.Fatalf("WriteFile() error = %v", err)
+	}
+
+	cfg := &config.Config{
+		Agents: config.AgentsConfig{
+			Defaults: config.AgentDefaults{
+				Workspace:         workspace,
+				ModelName:         "text-model",
+				MaxTokens:         4096,
+				MaxToolIterations: 3,
+			},
+		},
+		Tools: config.ToolsConfig{
+			LoadImage: config.ToolConfig{Enabled: true},
+		},
+		ModelList: []*config.ModelConfig{
+			{ModelName: "text-model", Model: "openai/text-model"},
+		},
+	}
+
+	msgBus := bus.NewMessageBus()
+	planner := &loadImagePlanningProvider{
+		path: pngPath,
+		followUpErr: fmt.Errorf(
+			`API request failed: Status: 404 Body: {"error":{"message":"No endpoints found that support image input"}}`,
+		),
+	}
+	al := NewAgentLoop(cfg, msgBus, planner)
+	al.SetMediaStore(media.NewFileMediaStore())
+
+	timeoutCtx, cancel := context.WithTimeout(context.Background(), responseTimeout)
+	defer cancel()
+
+	resp, err := al.processMessage(timeoutCtx, testInboundMessage(bus.InboundMessage{
+		Context: bus.InboundContext{
+			Channel:   "telegram",
+			ChatID:    "chat1",
+			ChatType:  "direct",
+			SenderID:  "user1",
+			MessageID: "m1",
+		},
+		Content:    "describe the image you load",
+		SessionKey: "agent:main:telegram:direct:user1",
+	}))
+	if err == nil {
+		t.Fatal("processMessage() error = nil, want vision unsupported failure")
+	}
+	if resp != "" {
+		t.Fatalf("response = %q, want empty response on error", resp)
+	}
+	if !strings.Contains(err.Error(), `active model "text-model" does not support image input`) {
+		t.Fatalf("error = %q, want clear vision unsupported guidance", err.Error())
+	}
+	if planner.calls != 2 {
+		t.Fatalf("planner calls = %d, want %d", planner.calls, 2)
 	}
 }
 

pkg/agent/instance.go

@@ -43,6 +43,7 @@ type AgentInstance struct {
 	SkillsFilter              []string
 	MCPServerAllowlist        map[string]struct{}
 	Candidates                []providers.FallbackCandidate
+	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
@@ -198,9 +199,19 @@ func NewAgentInstance(
 
 	// Resolve fallback candidates
 	candidates := resolveModelCandidates(cfg, defaults.Provider, model, fallbacks)
+	imageCandidates := resolveModelCandidates(
+		cfg,
+		defaults.Provider,
+		defaults.ImageModel,
+		defaults.ImageModelFallbacks,
+	)
 
 	candidateProviders := make(map[string]providers.LLMProvider)
 	populateCandidateProvidersFromNames(cfg, workspace, fallbacks, candidateProviders)
+	if strings.TrimSpace(defaults.ImageModel) != "" {
+		imageNames := append([]string{defaults.ImageModel}, defaults.ImageModelFallbacks...)
+		populateCandidateProvidersFromNames(cfg, workspace, imageNames, candidateProviders)
+	}
 
 	// Model routing setup: pre-resolve light model candidates at creation time
 	// to avoid repeated model_list lookups on every incoming message.
@@ -265,6 +276,7 @@ func NewAgentInstance(
 		SkillsFilter:              skillsFilter,
 		MCPServerAllowlist:        agentMCPServerAllowlist,
 		Candidates:                candidates,
+		ImageCandidates:           imageCandidates,
 		Router:                    router,
 		LightCandidates:           lightCandidates,
 		LightProvider:             lightProvider,

pkg/agent/llm_media.go

@@ -1,8 +1,10 @@
 package agent
 
 import (
+	"fmt"
 	"strings"
 
+	"github.com/sipeed/picoclaw/pkg/logger"
 	"github.com/sipeed/picoclaw/pkg/providers"
 )
 
@@ -65,3 +67,112 @@ func isVisionUnsupportedError(err error) bool {
 
 	return false
 }
+
+func visionUnsupportedModelError(modelName string, imageModelConfigured bool) error {
+	modelName = strings.TrimSpace(modelName)
+	if imageModelConfigured {
+		if modelName != "" {
+			return fmt.Errorf(
+				"selected vision model %q does not support image input; update agents.defaults.image_model to a multimodal model",
+				modelName,
+			)
+		}
+		return fmt.Errorf(
+			"selected vision model does not support image input; update agents.defaults.image_model to a multimodal model",
+		)
+	}
+	if modelName != "" {
+		return fmt.Errorf(
+			"active model %q does not support image input; configure agents.defaults.image_model with a multimodal model",
+			modelName,
+		)
+	}
+	return fmt.Errorf(
+		"the active model does not support image input; configure agents.defaults.image_model with a multimodal model",
+	)
+}
+
+func sameCandidateSet(a, b []providers.FallbackCandidate) bool {
+	if len(a) != len(b) {
+		return false
+	}
+	for i := range a {
+		if a[i].StableKey() != b[i].StableKey() {
+			return false
+		}
+	}
+	return true
+}
+
+func (p *Pipeline) routeMediaTurn(ts *turnState, exec *turnExecution) error {
+	if p == nil || ts == nil || ts.agent == nil || exec == nil || !hasMediaRefs(exec.callMessages) {
+		return nil
+	}
+
+	var targetCandidates []providers.FallbackCandidate
+	var targetModelName string
+	var routeReason string
+
+	switch {
+	case len(ts.agent.ImageCandidates) > 0:
+		targetCandidates = append([]providers.FallbackCandidate(nil), ts.agent.ImageCandidates...)
+		targetModelName = strings.TrimSpace(p.Cfg.Agents.Defaults.ImageModel)
+		routeReason = "configured_image_model"
+	case exec.usedLight && len(ts.agent.Candidates) > 0:
+		targetCandidates = append([]providers.FallbackCandidate(nil), ts.agent.Candidates...)
+		targetModelName = strings.TrimSpace(ts.agent.Model)
+		routeReason = "bypass_light_model_for_media"
+	default:
+		return nil
+	}
+
+	if len(targetCandidates) == 0 {
+		return nil
+	}
+
+	targetModel := resolvedCandidateModel(targetCandidates, targetModelName)
+	targetProvider := exec.activeProvider
+	firstCandidate := targetCandidates[0]
+	if provider, err := providerForFallbackCandidate(
+		ts.agent,
+		ts.agent.Provider,
+		targetCandidates,
+		firstCandidate.Provider,
+		firstCandidate.Model,
+	); err != nil {
+		return err
+	} else if provider != nil {
+		targetProvider = provider
+	}
+
+	resolvedModelName := resolvedCandidateModelName(targetCandidates, targetModelName)
+	if sameCandidateSet(exec.activeCandidates, targetCandidates) &&
+		exec.activeModel == targetModel &&
+		exec.llmModelName == resolvedModelName {
+		return nil
+	}
+
+	exec.activeCandidates = targetCandidates
+	exec.activeModel = targetModel
+	exec.activeProvider = targetProvider
+	exec.activeModelConfig = resolveActiveModelConfig(
+		p.Cfg,
+		ts.agent.Workspace,
+		targetCandidates,
+		targetModel,
+		p.Cfg.Agents.Defaults.Provider,
+	)
+	exec.llmModelName = resolvedModelName
+	exec.usedLight = false
+
+	logger.InfoCF("agent", "Media turn routing selected model", map[string]any{
+		"agent_id":       ts.agent.ID,
+		"reason":         routeReason,
+		"model":          exec.activeModel,
+		"model_name":     exec.llmModelName,
+		"candidates":     len(exec.activeCandidates),
+		"messages_count": len(exec.callMessages),
+	})
+
+	return nil
+}

pkg/agent/pipeline_llm.go

@@ -64,6 +64,9 @@ func (p *Pipeline) CallLLM(
 		exec.providerToolDefs = nil
 		ts.markGracefulTerminalUsed()
 	}
+	if err := p.routeMediaTurn(ts, exec); err != nil {
+		return ControlBreak, err
+	}
 
 	exec.llmOpts = map[string]any{
 		"max_tokens":       ts.agent.MaxTokens,
@@ -170,11 +173,10 @@ func (p *Pipeline) CallLLM(
 			return response, streamErr
 		}
 
-		if len(exec.activeCandidates) > 1 && p.Fallback != nil {
-			fbResult, fbErr := p.Fallback.ExecuteCandidate(
-				providerCtx,
-				exec.activeCandidates,
-				func(ctx context.Context, candidate providers.FallbackCandidate) (*providers.LLMResponse, error) {
+		runCandidate := func(
+			ctx context.Context,
+			candidate providers.FallbackCandidate,
+		) (*providers.LLMResponse, error) {
 			candidateProvider, err := providerForFallbackCandidate(
 				ts.agent,
 				exec.activeProvider,
@@ -198,8 +200,35 @@ func (p *Pipeline) CallLLM(
 			applyThinkingOption(callOpts, candidateProvider, candidateThinking, true, ts.agent.ID)
 			exec.suppressReasoning = shouldSuppressReasoningFor(candidateThinking)
 			return candidateProvider.Chat(ctx, messagesForCall, toolDefsForCall, candidate.Model, callOpts)
+		}
+
+		if len(exec.activeCandidates) > 1 && p.Fallback != nil {
+			var (
+				fbResult *providers.FallbackResult
+				fbErr    error
+			)
+			if hasMediaRefs(messagesForCall) {
+				fbResult, fbErr = p.Fallback.ExecuteImage(
+					providerCtx,
+					exec.activeCandidates,
+					func(ctx context.Context, provider, model string) (*providers.LLMResponse, error) {
+						candidate := providers.FallbackCandidate{Provider: provider, Model: model}
+						for _, configured := range exec.activeCandidates {
+							if configured.Provider == provider && configured.Model == model {
+								candidate = configured
+								break
+							}
+						}
+						return runCandidate(ctx, candidate)
 					},
 				)
+			} else {
+				fbResult, fbErr = p.Fallback.ExecuteCandidate(
+					providerCtx,
+					exec.activeCandidates,
+					runCandidate,
+				)
+			}
 			if fbErr != nil {
 				return nil, fbErr
 			}
@@ -250,33 +279,11 @@ func (p *Pipeline) CallLLM(
 			break
 		}
 
-		// Retry without media if vision is unsupported
-		if hasMediaRefs(exec.callMessages) && isVisionUnsupportedError(err) && retry < maxRetries {
-			al.emitEvent(
-				runtimeevents.KindAgentLLMRetry,
-				ts.eventMeta("runTurn", "turn.llm.retry"),
-				LLMRetryPayload{
-					Attempt:    retry + 1,
-					MaxRetries: maxRetries,
-					Reason:     "vision_unsupported",
-					Error:      err.Error(),
-					Backoff:    0,
-				},
+		if hasMediaRefs(exec.callMessages) && isVisionUnsupportedError(err) {
+			return ControlBreak, visionUnsupportedModelError(
+				exec.llmModelName,
+				len(ts.agent.ImageCandidates) > 0,
 			)
-			logger.WarnCF("agent", "Vision unsupported, retrying without media", map[string]any{
-				"error": err.Error(),
-				"retry": retry,
-			})
-			exec.callMessages = stripMessageMedia(exec.callMessages)
-			if !ts.opts.NoHistory {
-				exec.history = stripMessageMedia(exec.history)
-				ts.agent.Sessions.SetHistory(ts.sessionKey, exec.history)
-				for i := range ts.persistedMessages {
-					ts.persistedMessages[i].Media = nil
-				}
-				ts.refreshRestorePointFromSession(ts.agent)
-			}
-			continue
 		}
 
 		errMsg := strings.ToLower(err.Error())