Merge pull request #1313 from wgqqqqq/fix/multimodal-model-capabilities

fix: preserve multimodal model capabilities

Author: GCWing

src/crates/assembly/core/src/agentic/execution/execution_engine.rs

@@ -655,6 +655,49 @@ impl ExecutionEngine {
             .unwrap_or_else(|| "auto".to_string())
     }
 
+    async fn resolve_primary_model_context(
+        model_id: &str,
+        ai_client_model: &str,
+        ai_client_provider: &str,
+        unavailable_log_message: &str,
+    ) -> (String, bool) {
+        let config_service = get_global_config_service().await.ok();
+        if let Some(service) = config_service {
+            let ai_config: crate::service::config::types::AIConfig =
+                service.get_config(Some("ai")).await.unwrap_or_default();
+
+            let resolved_id = Self::resolve_configured_model_id(&ai_config, model_id);
+            let model_cfg = ai_config
+                .models
+                .iter()
+                .find(|m| m.id == resolved_id)
+                .or_else(|| ai_config.models.iter().find(|m| m.name == resolved_id))
+                .or_else(|| {
+                    ai_config
+                        .models
+                        .iter()
+                        .find(|m| m.model_name == resolved_id)
+                })
+                .or_else(|| {
+                    ai_config.models.iter().find(|m| {
+                        m.model_name == ai_client_model && m.provider == ai_client_provider
+                    })
+                });
+
+            let supports = model_cfg.is_some_and(|m| {
+                m.capabilities
+                    .iter()
+                    .any(|cap| matches!(cap, ModelCapability::ImageUnderstanding))
+                    || matches!(m.category, ModelCategory::Multimodal)
+            });
+
+            (resolved_id, supports)
+        } else {
+            warn!("{}", unavailable_log_message);
+            (model_id.to_string(), false)
+        }
+    }
+
     async fn build_tool_listing_sections(
         manifest: &ResolvedToolManifest,
         tool_context: &crate::agentic::tools::framework::ToolUseContext,
@@ -1527,46 +1570,14 @@ impl ExecutionEngine {
                 ))
             })?;
 
-        let (resolved_primary_model_id, primary_supports_image_understanding) = {
-            let config_service = get_global_config_service().await.ok();
-            if let Some(service) = config_service {
-                let ai_config: crate::service::config::types::AIConfig =
-                    service.get_config(Some("ai")).await.unwrap_or_default();
-
-                let resolved_id = Self::resolve_configured_model_id(&ai_config, &model_id);
-                let model_cfg = ai_config
-                    .models
-                    .iter()
-                    .find(|m| m.id == resolved_id)
-                    .or_else(|| ai_config.models.iter().find(|m| m.name == resolved_id))
-                    .or_else(|| {
-                        ai_config
-                            .models
-                            .iter()
-                            .find(|m| m.model_name == resolved_id)
-                    })
-                    .or_else(|| {
-                        ai_config.models.iter().find(|m| {
-                            m.model_name == ai_client.config.model
-                                && m.provider == ai_client.config.format
-                        })
-                    });
-
-                let supports = model_cfg.is_some_and(|m| {
-                    m.capabilities
-                        .iter()
-                        .any(|cap| matches!(cap, ModelCapability::ImageUnderstanding))
-                        || matches!(m.category, ModelCategory::Multimodal)
-                });
-
-                (resolved_id, supports)
-            } else {
-                warn!(
-                    "Config service unavailable, assuming compression model is text-only for image input gating"
-                );
-                (model_id.clone(), false)
-            }
-        };
+        let (resolved_primary_model_id, primary_supports_image_understanding) =
+            Self::resolve_primary_model_context(
+                &model_id,
+                &ai_client.config.model,
+                &ai_client.config.format,
+                "Config service unavailable, assuming compression model is text-only for image input gating",
+            )
+            .await;
 
         let model_capability_profile = ModelCapabilityProfile::from_resolved_model(
             &resolved_primary_model_id,
@@ -1602,6 +1613,9 @@ impl ExecutionEngine {
             &context.agent_type,
             context.workspace.as_ref(),
             context.workspace_services.as_ref(),
+            Some(&resolved_primary_model_id),
+            Some(&ai_client.config.model),
+            Some(&ai_client.config.format),
             primary_supports_image_understanding,
             &tool_manifest_context_vars,
         );
@@ -2198,47 +2212,14 @@ impl ExecutionEngine {
             })?;
 
         // Primary model vision capability (tools + system prompt appendix; also used below for API message stripping).
-        let (resolved_primary_model_id, primary_supports_image_understanding) = {
-            let config_service = get_global_config_service().await.ok();
-            if let Some(service) = config_service {
-                let ai_config: crate::service::config::types::AIConfig =
-                    service.get_config(Some("ai")).await.unwrap_or_default();
-
-                let resolved_id = Self::resolve_configured_model_id(&ai_config, &model_id);
-
-                let model_cfg = ai_config
-                    .models
-                    .iter()
-                    .find(|m| m.id == resolved_id)
-                    .or_else(|| ai_config.models.iter().find(|m| m.name == resolved_id))
-                    .or_else(|| {
-                        ai_config
-                            .models
-                            .iter()
-                            .find(|m| m.model_name == resolved_id)
-                    })
-                    .or_else(|| {
-                        ai_config.models.iter().find(|m| {
-                            m.model_name == ai_client.config.model
-                                && m.provider == ai_client.config.format
-                        })
-                    });
-
-                let supports = model_cfg.is_some_and(|m| {
-                    m.capabilities
-                        .iter()
-                        .any(|cap| matches!(cap, ModelCapability::ImageUnderstanding))
-                        || matches!(m.category, ModelCategory::Multimodal)
-                });
-
-                (resolved_id, supports)
-            } else {
-                warn!(
-                    "Config service unavailable, assuming primary model is text-only for image input gating"
-                );
-                (model_id.clone(), false)
-            }
-        };
+        let (resolved_primary_model_id, primary_supports_image_understanding) =
+            Self::resolve_primary_model_context(
+                &model_id,
+                &ai_client.config.model,
+                &ai_client.config.format,
+                "Config service unavailable, assuming primary model is text-only for image input gating",
+            )
+            .await;
 
         let model_context_window = ai_client.config.context_window as usize;
         let session_max_tokens = session.config.max_context_tokens;
@@ -2296,6 +2277,9 @@ impl ExecutionEngine {
             &agent_type,
             context.workspace.as_ref(),
             context.workspace_services.as_ref(),
+            Some(&resolved_primary_model_id),
+            Some(&ai_client.config.model),
+            Some(&ai_client.config.format),
             primary_supports_image_understanding,
             &tool_manifest_context_vars,
         );

src/crates/assembly/core/src/agentic/skill_agent_snapshot.rs

@@ -51,6 +51,9 @@ pub async fn resolve_skill_agent_snapshot(
         agent_type,
         workspace,
         workspace_services,
+        None,
+        None,
+        None,
         true,
         context_vars,
     );

src/crates/assembly/core/src/agentic/tools/product_runtime/catalog.rs

@@ -239,6 +239,18 @@ mod tests {
         }
     }
 
+    fn multimodal_anthropic_tool_context(agent_type: Option<&str>) -> ToolUseContext {
+        let mut context = tool_context(agent_type);
+        context.custom_data.insert(
+            "primary_model_supports_image_understanding".to_string(),
+            json!(true),
+        );
+        context
+            .custom_data
+            .insert("primary_model_provider".to_string(), json!("anthropic"));
+        context
+    }
+
     fn context_without_agent_type() -> ToolUseContext {
         tool_context(None)
     }
@@ -483,6 +495,24 @@ mod tests {
             .any(|tool| tool.name == GET_TOOL_SPEC_TOOL_NAME));
     }
 
+    #[tokio::test]
+    async fn product_manifest_keeps_view_image_for_multimodal_anthropic_context() {
+        let allowed_tools = vec!["Read".to_string(), "view_image".to_string()];
+
+        let manifest = resolve_product_resolved_tool_manifest(
+            &allowed_tools,
+            &AgentToolPolicyOverrides::default(),
+            &multimodal_anthropic_tool_context(Some("test-agent")),
+        )
+        .await;
+
+        assert_eq!(manifest.allowed_tool_names, allowed_tools);
+        assert!(manifest
+            .tool_definitions
+            .iter()
+            .any(|tool| tool.name == "view_image"));
+    }
+
     #[tokio::test]
     async fn product_manifest_snapshot_preserves_collapsed_tool_discovery_contract() {
         let allowed_tools = vec![

src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs

@@ -224,10 +224,31 @@ pub(crate) fn build_tool_description_context(
     agent_type: &str,
     workspace: Option<&WorkspaceBinding>,
     workspace_services: Option<&WorkspaceServices>,
+    primary_model_id: Option<&str>,
+    primary_model_name: Option<&str>,
+    primary_model_provider: Option<&str>,
     primary_supports_image_understanding: bool,
     context_vars: &HashMap<String, String>,
 ) -> ToolUseContext {
     let mut custom_data = HashMap::new();
+    if let Some(primary_model_id) = primary_model_id {
+        custom_data.insert(
+            "primary_model_id".to_string(),
+            Value::String(primary_model_id.to_string()),
+        );
+    }
+    if let Some(primary_model_name) = primary_model_name {
+        custom_data.insert(
+            "primary_model_name".to_string(),
+            Value::String(primary_model_name.to_string()),
+        );
+    }
+    if let Some(primary_model_provider) = primary_model_provider {
+        custom_data.insert(
+            "primary_model_provider".to_string(),
+            Value::String(primary_model_provider.to_string()),
+        );
+    }
     custom_data.insert(
         "primary_model_supports_image_understanding".to_string(),
         Value::Bool(primary_supports_image_understanding),
@@ -1197,7 +1218,16 @@ mod context_builder_tests {
             "false".to_string(),
         );
 
-        let context = build_tool_description_context("coding", None, None, true, &context_vars);
+        let context = build_tool_description_context(
+            "coding",
+            None,
+            None,
+            Some("model_1"),
+            Some("vision-model"),
+            Some("anthropic"),
+            true,
+            &context_vars,
+        );
 
         assert_eq!(context.agent_type.as_deref(), Some("coding"));
         assert!(context.tool_call_id.is_none());
@@ -1212,6 +1242,15 @@ mod context_builder_tests {
             context.custom_data["primary_model_supports_image_understanding"],
             json!("false")
         );
+        assert_eq!(context.custom_data["primary_model_id"], json!("model_1"));
+        assert_eq!(
+            context.custom_data["primary_model_name"],
+            json!("vision-model")
+        );
+        assert_eq!(
+            context.custom_data["primary_model_provider"],
+            json!("anthropic")
+        );
     }
 }
 

src/web-ui/src/infrastructure/config/components/AIModelConfig.tsx

@@ -6,10 +6,10 @@ import {
   AIModelConfig as AIModelConfigType, 
   ProxyConfig, 
   ModelCategory,
-  ModelCapability,
   ReasoningMode
 } from '../types';
 import { configManager } from '../services/ConfigManager';
+import { getCapabilitiesByCategory, resolveModelCategory } from '../services/modelCategory';
 import { PROVIDER_TEMPLATES, getModelDisplayName, getProviderDisplayName, getProviderTemplateId } from '../services/modelConfigs';
 import { DEFAULT_REASONING_MODE, getEffectiveReasoningMode, supportsAnthropicAdaptive, supportsAnthropicReasoning, supportsAnthropicThinkingBudget, supportsDeepSeekReasoningEffort, supportsResponsesReasoning } from '../utils/reasoning';
 import { aiApi, systemAPI } from '@/infrastructure/api';
@@ -233,15 +233,6 @@ function dedupeSelectedModelDraftsByModelName(drafts: SelectedModelDraft[]): Sel
   return out;
 }
 
-function getCapabilitiesByCategory(category: ModelCategory): ModelCapability[] {
-  switch (category) {
-    case 'general_chat':
-    case 'multimodal':
-    default:
-      return ['text_chat', 'function_calling'];
-  }
-}
-
 /**
  * Compute the stored request URL from a base URL and provider format.
  * For gemini, stores the bare base (no /v1beta/models/... suffix) —
@@ -1047,8 +1038,18 @@ const AIModelConfig: React.FC = () => {
           enabled: editingConfig.enabled ?? true,
           context_window: draft.contextWindow,
           max_tokens: draft.maxTokens,
-          category: draft.category,
-          capabilities: getCapabilitiesByCategory(draft.category),
+          category: resolveModelCategory(
+            draft.modelName,
+            draft.category,
+            editingConfig.provider || 'openai'
+          ),
+          capabilities: getCapabilitiesByCategory(
+            resolveModelCategory(
+              draft.modelName,
+              draft.category,
+              editingConfig.provider || 'openai'
+            )
+          ),
           recommended_for: editingConfig.recommended_for || [],
           metadata: {
             ...(editingConfig.metadata || {}),