fix: tests, lint, workflow model name logic

Author: RomneyDa

extensions/cli/src/args.test.ts

@@ -1,4 +1,5 @@
 import { vi } from "vitest";
+
 import { processRule as processPromptOrRule } from "./hubLoader.js";
 describe("processPromptOrRule (loadRuleFromHub integration)", () => {
   // Mock fetch for hub tests

extensions/cli/src/configEnhancer.test.ts

@@ -13,6 +13,9 @@ vi.mock("./hubLoader.js", () => ({
     // Return as-is for direct content
     return Promise.resolve(rule);
   }),
+  loadPackageFromHub: vi.fn(() =>
+    Promise.resolve({ name: "test", provider: "test" }),
+  ),
   loadPackagesFromHub: vi.fn(() => Promise.resolve([])),
   mcpProcessor: {},
   modelProcessor: {},
@@ -23,6 +26,8 @@ vi.mock("./services/ServiceContainer.js", () => ({
   serviceContainer: {
     get: vi.fn(() =>
       Promise.resolve({
+        workflowService: null,
+        workflowModelName: null,
         workflowFile: null,
         workflow: null,
       }),

extensions/cli/src/configEnhancer.ts

@@ -6,6 +6,7 @@ import {
 
 import { BaseCommandOptions } from "./commands/BaseCommandOptions.js";
 import {
+  loadPackageFromHub,
   loadPackagesFromHub,
   mcpProcessor,
   modelProcessor,
@@ -27,12 +28,13 @@ export class ConfigEnhancer {
     _options: BaseCommandOptions,
   ): Promise<AssistantUnrolled> {
     let enhancedConfig = { ...config };
+    const options = { ..._options };
 
     // Add workflow rules/mcp servers if present
-    const options = { ..._options };
-    const { workflowFile } = await serviceContainer.get<WorkflowServiceState>(
+    const workflowState = await serviceContainer.get<WorkflowServiceState>(
       SERVICE_NAMES.WORKFLOW,
     );
+    const { workflowFile, workflowService } = workflowState;
 
     if (workflowFile) {
       const { rules, model, tools, prompt } = workflowFile;
@@ -50,12 +52,17 @@ export class ConfigEnhancer {
         }
       }
 
-      // Note that --model takes precedence over workflow model
+      // --model takes precedence over workflow model
       if (model) {
-        options.model = [...(options.model ?? []), model];
+        const workflowModel = await loadPackageFromHub(model, modelProcessor);
+        enhancedConfig.models = [
+          workflowModel,
+          ...(enhancedConfig.models ?? []),
+        ];
+        workflowService?.setWorkflowModelName(workflowModel.name);
       }
 
-      // Add workflow prompt as prefix
+      // Add workflow prompt as prefix (see processAndCombinePrompts)
       if (prompt) {
         options.prompt = [prompt, ...(options.prompt || [])];
       }
@@ -141,11 +148,8 @@ export class ConfigEnhancer {
     const modifiedConfig = { ...config };
 
     // Prepend processed MCPs to existing mcpServers array for consistency
-    const existingMcpServers = (modifiedConfig as any).mcpServers || [];
-    (modifiedConfig as any).mcpServers = [
-      ...processedMcps,
-      ...existingMcpServers,
-    ];
+    const existingMcpServers = modifiedConfig.mcpServers || [];
+    modifiedConfig.mcpServers = [...processedMcps, ...existingMcpServers];
 
     return modifiedConfig;
   }
@@ -159,12 +163,11 @@ export class ConfigEnhancer {
   ): Promise<AssistantUnrolled> {
     const processedModels = await loadPackagesFromHub(models, modelProcessor);
 
-    // Clone the config to avoid mutating the original
     const modifiedConfig = { ...config };
 
     // Prepend processed models to existing models array so they become the default
-    const existingModels = (modifiedConfig as any).models || [];
-    (modifiedConfig as any).models = [...processedModels, ...existingModels];
+    const existingModels = modifiedConfig.models || [];
+    modifiedConfig.models = [...processedModels, ...existingModels];
 
     return modifiedConfig;
   }

extensions/cli/src/hubLoader.ts

@@ -1,6 +1,6 @@
+import { parseWorkflowFile, WorkflowFile } from "@continuedev/config-yaml";
 import JSZip from "jszip";
 
-import { parseWorkflowFile, WorkflowFile } from "@continuedev/config-yaml";
 import { env } from "./env.js";
 import { logger } from "./util/logger.js";
 

extensions/cli/src/services/ModelService.ts

@@ -56,15 +56,14 @@ export class ModelService
         model && (model.roles?.includes("chat") || model.roles === undefined),
     ) || []) as ModelConfig[];
 
-    // Check if workflow has a model specified (highest priority)
-    let preferredModelName: string | null = null;
+    let preferredModelName: string | null | undefined = null;
     let modelSource = "default";
 
-    if (workflowServiceState?.workflowFile?.model) {
-      preferredModelName = workflowServiceState.workflowFile.model;
+    // Priority = workflow -> last selected model
+    if (workflowServiceState?.workflowModelName) {
+      preferredModelName = workflowServiceState.workflowModelName;
       modelSource = "workflow";
     } else {
-      // Fall back to persisted model name if no workflow model
       preferredModelName = getModelName(authConfig);
       if (preferredModelName) {
         modelSource = "persisted";

extensions/cli/src/services/ModelService.workflow-priority.test.ts

@@ -60,6 +60,8 @@ describe("ModelService workflow model prioritization", () => {
         prompt: "Test workflow",
       },
       slug: "test/workflow",
+      workflowModelName: "gpt-4",
+      workflowService: null,
     };
 
     // Mock createLlmApi to return different models based on the selected model
@@ -100,6 +102,8 @@ describe("ModelService workflow model prioritization", () => {
         // No model specified
       },
       slug: "test/workflow",
+      workflowModelName: null,
+      workflowService: null,
     };
 
     // Mock getModelName to return a persisted model
@@ -142,6 +146,8 @@ describe("ModelService workflow model prioritization", () => {
         prompt: "Test workflow",
       },
       slug: "test/workflow",
+      workflowModelName: "non-existent-model",
+      workflowService: null,
     };
 
     const getLlmApiMock = vi.mocked(config.getLlmApi);
@@ -171,6 +177,8 @@ describe("ModelService workflow model prioritization", () => {
     const workflowServiceState: WorkflowServiceState = {
       workflowFile: null,
       slug: null,
+      workflowModelName: null,
+      workflowService: null,
     };
 
     // Make sure getModelName returns null (no persisted model)
@@ -207,6 +215,8 @@ describe("ModelService workflow model prioritization", () => {
         prompt: "Test workflow",
       },
       slug: "test/workflow",
+      workflowModelName: "gpt-4",
+      workflowService: null,
     };
 
     // Mock getModelName to return a different persisted model

extensions/cli/src/services/WorkflowService.test.ts

@@ -1,6 +1,7 @@
 import { Mock, vi } from "vitest";
 
 import { workflowProcessor } from "../hubLoader.js";
+
 import { WorkflowService } from "./WorkflowService.js";
 
 // Mock the hubLoader module
@@ -44,6 +45,8 @@ describe("WorkflowService", () => {
       expect(state).toEqual({
         workflowFile: null,
         slug: null,
+        workflowModelName: null,
+        workflowService: service,
       });
     });
 
@@ -53,6 +56,8 @@ describe("WorkflowService", () => {
       expect(state).toEqual({
         workflowFile: null,
         slug: null,
+        workflowModelName: null,
+        workflowService: service,
       });
     });
 
@@ -62,6 +67,8 @@ describe("WorkflowService", () => {
       expect(state).toEqual({
         workflowFile: null,
         slug: null,
+        workflowModelName: null,
+        workflowService: service,
       });
 
       // Should not call loadPackageFromHub with invalid slug
@@ -74,6 +81,8 @@ describe("WorkflowService", () => {
       expect(state).toEqual({
         workflowFile: null,
         slug: null,
+        workflowModelName: null,
+        workflowService: service,
       });
 
       expect(mockLoadPackageFromHub).not.toHaveBeenCalled();
@@ -96,6 +105,8 @@ describe("WorkflowService", () => {
       expect(state).toEqual({
         workflowFile: mockWorkflowFile,
         slug: "owner/package",
+        workflowModelName: null,
+        workflowService: service,
       });
 
       expect(mockLoadPackageFromHub).toHaveBeenCalledWith(
@@ -112,6 +123,8 @@ describe("WorkflowService", () => {
       expect(state).toEqual({
         workflowFile: null,
         slug: null,
+        workflowModelName: null,
+        workflowService: service,
       });
 
       expect(mockLoadPackageFromHub).toHaveBeenCalledWith(
@@ -133,6 +146,8 @@ describe("WorkflowService", () => {
       expect(state).toEqual({
         workflowFile: mockWorkflowFile,
         slug: "owner/minimal",
+        workflowModelName: null,
+        workflowService: service,
       });
     });
   });

extensions/cli/src/services/WorkflowService.ts

@@ -9,10 +9,21 @@ import { WorkflowServiceState } from "./types.js";
  * Loads workflows from the hub and extracts model, tools, and prompt information
  */
 export class WorkflowService extends BaseService<WorkflowServiceState> {
+  /**
+   * Set the resolved workflow model name after it's been processed
+   * Called by ConfigEnhancer after resolving the model slug
+   */
+  setWorkflowModelName(modelName: string): void {
+    this.setState({
+      workflowModelName: modelName,
+    });
+  }
   constructor() {
     super("WorkflowService", {
+      workflowService: null,
       workflowFile: null,
       slug: null,
+      workflowModelName: null,
     });
   }
 
@@ -22,8 +33,10 @@ export class WorkflowService extends BaseService<WorkflowServiceState> {
   async doInitialize(workflowSlug?: string): Promise<WorkflowServiceState> {
     if (!workflowSlug) {
       return {
+        workflowService: this,
         workflowFile: null,
         slug: null,
+        workflowModelName: null,
       };
     }
 
@@ -41,14 +54,18 @@ export class WorkflowService extends BaseService<WorkflowServiceState> {
       );
 
       return {
+        workflowService: this,
         workflowFile,
         slug: workflowSlug,
+        workflowModelName: null, // Will be set by ConfigEnhancer after model resolution
       };
     } catch (error: any) {
       logger.error("Failed to initialize WorkflowService:", error);
       return {
+        workflowService: this,
         workflowFile: null,
         slug: null,
+        workflowModelName: null,
       };
     }
   }

extensions/cli/src/services/types.ts

@@ -13,6 +13,7 @@ import { BaseCommandOptions } from "../commands/BaseCommandOptions.js";
 import { PermissionMode, ToolPermissions } from "../permissions/types.js";
 
 import { MCPService } from "./MCPService.js";
+import { WorkflowService } from "./WorkflowService.js";
 
 /**
  * Service lifecycle states
@@ -125,6 +126,8 @@ export interface StorageSyncServiceState {
 export interface WorkflowServiceState {
   workflowFile: WorkflowFile | null;
   slug: string | null;
+  workflowModelName: string | null;
+  workflowService: WorkflowService | null;
 }
 
 export type { ChatHistoryState } from "./ChatHistoryService.js";

extensions/cli/src/services/workflow-integration.test.ts

@@ -136,6 +136,12 @@ describe("Workflow Integration Tests", () => {
       const workflowState = workflowService.getState();
       expect(workflowState.workflowFile?.model).toBe("gpt-4-workflow");
 
+      // Mock loadPackageFromHub to return a model for the workflow model
+      mockLoadPackageFromHub.mockResolvedValueOnce({
+        name: "gpt-4-workflow",
+        provider: "openai",
+      });
+
       // Test that ConfigEnhancer adds the workflow model to options
       const baseConfig = { models: [] };
       const baseOptions = {}; // No --model flag
@@ -145,11 +151,18 @@ describe("Workflow Integration Tests", () => {
         baseOptions,
       );
 
-      // Should have loaded the workflow model via injectModels
-      expect(mockLoadPackagesFromHub).toHaveBeenCalledWith(
-        ["gpt-4-workflow"],
-        expect.anything(),
+      // Should have loaded the workflow model directly via loadPackageFromHub
+      expect(mockLoadPackageFromHub).toHaveBeenCalledWith(
+        "gpt-4-workflow",
+        mockModelProcessor,
       );
+
+      // The workflow model should be prepended to the models array
+      expect(enhancedConfig.models).toHaveLength(1);
+      expect(enhancedConfig.models?.[0]).toEqual({
+        name: "gpt-4-workflow",
+        provider: "openai",
+      });
     });
 
     it("should not add workflow model when no workflow is active", async () => {
@@ -177,6 +190,18 @@ describe("Workflow Integration Tests", () => {
       mockLoadPackageFromHub.mockResolvedValue(mockWorkflowFile);
       await workflowService.initialize("owner/workflow");
 
+      // Mock loadPackageFromHub for workflow model and loadPackagesFromHub for user models
+      mockLoadPackageFromHub.mockResolvedValueOnce({
+        name: "gpt-4-workflow",
+        provider: "openai",
+      });
+      mockLoadPackagesFromHub.mockResolvedValueOnce([
+        {
+          name: "user-specified-model",
+          provider: "anthropic",
+        },
+      ]);
+
       // Test that --model flag takes precedence
       const baseConfig = { models: [] };
       const baseOptions = { model: ["user-specified-model"] }; // User specified --model
@@ -186,12 +211,28 @@ describe("Workflow Integration Tests", () => {
         baseOptions,
       );
 
-      // Should process both the user model and workflow model
-      // The ConfigEnhancer adds workflow model to the existing model options
+      // Should process the user model via loadPackagesFromHub
       expect(mockLoadPackagesFromHub).toHaveBeenCalledWith(
-        ["user-specified-model", "gpt-4-workflow"],
-        expect.anything(),
+        ["user-specified-model"],
+        mockModelProcessor,
+      );
+
+      // Should also load the workflow model
+      expect(mockLoadPackageFromHub).toHaveBeenCalledWith(
+        "gpt-4-workflow",
+        mockModelProcessor,
       );
+
+      // Both models should be in the config, with user model first (takes precedence)
+      expect(enhancedConfig.models).toHaveLength(2);
+      expect(enhancedConfig.models?.[0]).toEqual({
+        name: "user-specified-model",
+        provider: "anthropic",
+      });
+      expect(enhancedConfig.models?.[1]).toEqual({
+        name: "gpt-4-workflow",
+        provider: "openai",
+      });
     });
   });