fixes

Author: jonit-dev

src/commands/run.ts

@@ -141,7 +141,6 @@ export class Run extends Command {
 
       const options: CrackedAgentOptions = {
         ...config,
-        model: this.modelManager.getCurrentModel(),
         options: this.parseOptions(config.options || ""),
         provider: config.provider as LLMProviderType,
       };
@@ -152,7 +151,7 @@ export class Run extends Command {
       }
 
       console.log(
-        `Using ${options.provider} provider and model: ${options.model}`,
+        `Using ${options.provider} provider and model: ${this.modelManager.getCurrentModel()}`,
       );
 
       const agent = container.resolve(CrackedAgent);

src/constants/modelScaling.ts

@@ -1,84 +1,5 @@
 import "reflect-metadata";
 
-// Mock modelScaling module
-jest.mock("@constants/modelScaling", () => ({
-  __esModule: true,
-  autoScaleAvailableModels: [
-    {
-      id: "qwen/qwen-2.5-coder-32b-instruct",
-      description: "Cheap, fast, slightly better than GPT4o-mini",
-      maxWriteTries: 2,
-      maxGlobalTries: 5,
-    },
-    {
-      id: "anthropic/claude-3.5-sonnet:beta",
-      description: "Scaled model for retry attempts",
-      maxWriteTries: 3,
-      maxGlobalTries: 10,
-    },
-    {
-      id: "openai/gpt-4o-2024-11-20",
-      description: "Scaled model for retry attempts",
-      maxWriteTries: 5,
-      maxGlobalTries: 15,
-    },
-    {
-      id: "openai/o1-mini",
-      description: "Final model for complex cases (currently inactive)",
-      maxWriteTries: 2,
-      maxGlobalTries: 20,
-    },
-  ],
-  getModelForTryCount: (
-    tryCount: string | null,
-    globalTries: number,
-  ): string => {
-    const models = [
-      {
-        id: "qwen/qwen-2.5-coder-32b-instruct",
-        maxWriteTries: 2,
-        maxGlobalTries: 5,
-      },
-      {
-        id: "anthropic/claude-3.5-sonnet:beta",
-        maxWriteTries: 3,
-        maxGlobalTries: 10,
-      },
-      {
-        id: "openai/gpt-4o-2024-11-20",
-        maxWriteTries: 5,
-        maxGlobalTries: 15,
-      },
-      {
-        id: "openai/o1-mini",
-        maxWriteTries: 2,
-        maxGlobalTries: 20,
-      },
-    ];
-
-    if (!tryCount) return models[0].id;
-
-    const tries = parseInt(tryCount, 10);
-
-    for (let i = 0; i < models.length; i++) {
-      const previousTriesSum = models
-        .slice(0, i)
-        .reduce((sum, model) => sum + model.maxWriteTries, 0);
-
-      if (
-        tries >= previousTriesSum + models[i].maxWriteTries ||
-        globalTries >= models[i].maxGlobalTries
-      ) {
-        continue;
-      }
-
-      return models[i].id;
-    }
-
-    return models[models.length - 1].id;
-  },
-}));
-
 // Mock chalk with a default export that matches how it's used
 jest.mock("chalk", () => ({
   __esModule: true,
@@ -102,3 +23,41 @@ jest.mock("chalk", () => ({
     strikethrough: jest.fn((text) => text),
   },
 }));
+
+// Global mock for ModelInfo
+jest.mock("@services/LLM/ModelInfo", () => {
+  const originalModule = jest.requireActual("@services/LLM/ModelInfo");
+
+  return {
+    ...originalModule,
+    ModelInfo: jest.fn().mockImplementation(() => ({
+      initialize: jest.fn().mockResolvedValue(undefined),
+      ensureInitialized: jest.fn().mockResolvedValue(undefined),
+      getCurrentModel: jest.fn().mockReturnValue("gpt-4"),
+      getCurrentModelInfo: jest.fn().mockReturnValue({
+        id: "gpt-4",
+        context_length: 8192,
+        top_provider: {
+          max_completion_tokens: 4096,
+        },
+      }),
+      setCurrentModel: jest.fn().mockResolvedValue(undefined),
+      getModelInfo: jest.fn().mockResolvedValue({
+        id: "gpt-4",
+        context_length: 8192,
+        top_provider: {
+          max_completion_tokens: 4096,
+        },
+      }),
+      isModelAvailable: jest.fn().mockResolvedValue(true),
+      getAllModels: jest.fn().mockResolvedValue(["gpt-4"]),
+      getCurrentModelContextLength: jest.fn().mockResolvedValue(8192),
+      getModelContextLength: jest.fn().mockResolvedValue(8192),
+      getCurrentModelMaxCompletionTokens: jest.fn().mockResolvedValue(4096),
+      getModelMaxCompletionTokens: jest.fn().mockResolvedValue(4096),
+      logCurrentModelUsage: jest.fn().mockResolvedValue(undefined),
+    })),
+  };
+});
+
+process.env.IS_UNIT_TEST = "true";

src/jest/jestSetupFilesAfterEnv.ts

@@ -24,16 +24,14 @@ const configSchema = z.object({
     .string()
     .optional()
     .default("anthropic/claude-3.5-sonnet:beta"),
-  autoScaleAvailableModels: z
-    .array(
-      z.object({
-        id: z.string(),
-        description: z.string(),
-        maxWriteTries: z.number(),
-        maxGlobalTries: z.number(),
-      }),
-    )
-    .optional(),
+  autoScaleAvailableModels: z.array(
+    z.object({
+      id: z.string(),
+      description: z.string(),
+      maxWriteTries: z.number(),
+      maxGlobalTries: z.number(),
+    }),
+  ),
   runAllTestsCmd: z.string().optional(),
   runOneTestCmd: z.string().optional(),
   runTypeCheckCmd: z.string().optional(),

src/services/ConfigService.ts

@@ -1,20 +1,19 @@
+import { DEFAULT_INSTRUCTIONS } from "@constants/defaultInstructions";
 import { FileReader } from "@services/FileManagement/FileReader";
 import { ActionsParser } from "@services/LLM/actions/ActionsParser";
 import { ILLMProvider } from "@services/LLM/ILLMProvider";
 import { LLMContextCreator } from "@services/LLM/LLMContextCreator";
 import { LLMProvider, LLMProviderType } from "@services/LLM/LLMProvider";
-import { ModelScaler } from "@services/LLM/ModelScaler";
 import { DebugLogger } from "@services/logging/DebugLogger";
 import { StreamHandler } from "@services/streaming/StreamHandler";
-import { HtmlEntityDecoder } from "@services/text/HTMLEntityDecoder";
 import { autoInjectable, singleton } from "tsyringe";
-import { DEFAULT_INSTRUCTIONS } from "@constants/defaultInstructions";
+import { ModelManager } from "./LLM/ModelManager";
+import { PhaseManager } from "./LLM/PhaseManager";
 
 export interface CrackedAgentOptions {
   root?: string;
   instructionsPath?: string;
   instructions?: string;
-  model: string;
   provider?: LLMProviderType;
   stream?: boolean;
   debug?: boolean;
@@ -41,16 +40,18 @@ export class CrackedAgent {
     private debugLogger: DebugLogger,
     private actionsParser: ActionsParser,
     private streamHandler: StreamHandler,
-    private htmlEntityDecoder: HtmlEntityDecoder,
-    private modelScaler: ModelScaler,
+    private phaseManager: PhaseManager,
+    private modelManager: ModelManager,
   ) {}
 
   async execute(
     message: string,
     options: CrackedAgentOptions,
   ): Promise<ExecutionResult> {
+    this.phaseManager.initializePhaseConfigs();
+
     const finalOptions = await this.setupExecution(options);
-    this.currentModel = finalOptions.model;
+    this.currentModel = this.modelManager.getCurrentModel();
 
     const formattedMessage = await this.contextCreator.create(
       message,
@@ -108,14 +109,7 @@ export class CrackedAgent {
       this.clearConversationHistory();
     }
 
-    // Always explicitly set auto-scaler state with the model
-    // If autoScaler is not specified, it defaults to false and uses the provided model
-    this.modelScaler.setAutoScaler(
-      finalOptions.autoScaler || false,
-      finalOptions.model,
-    );
-
-    await this.validateModel(finalOptions.model);
+    await this.validateModel(this.modelManager.getCurrentModel());
     await this.setupInstructions(finalOptions);
 
     return finalOptions;
@@ -286,4 +280,4 @@ export class CrackedAgent {
     this.llm.clearConversationContext();
     this.isFirstInteraction = true;
   }
-}
+}

src/services/CrackedAgent.ts

@@ -96,7 +96,7 @@ Run Type Check: ${runTypeCheckCmd}`;
     const phaseConfig = this.phaseManager.getCurrentPhaseConfig();
 
     const envDetails = config.includeAllFilesOnEnvToContext
-      ? `\n${context.environmentDetails}`
+      ? context.environmentDetails
       : "";
 
     const promptArgs: IPhasePromptArgs = {
@@ -119,6 +119,8 @@ ${context.message}
 - Break tasks into prioritized steps.
 - Use available actions sequentially.
 
+${envDetails ? `\n${envDetails}` : ""}
+${context.projectInfo ? `\n${context.projectInfo}` : ""}
 
 ## Instructions
 ${phaseConfig.generatePrompt(promptArgs)}

src/services/LLM/LLMContextCreator.ts

@@ -47,25 +47,25 @@ export class ModelInfo {
       return; // Already tracking this model
     }
 
-    const modelInfo = this.modelInfoMap.get(modelId);
-    if (modelInfo) {
-      this.currentModel = modelId;
-      this.currentModelInfo = modelInfo;
-      this.debugLogger.log("ModelInfo", "Current model info", {
-        model: modelId,
-        contextLength: modelInfo.context_length,
-        maxCompletionTokens: modelInfo.top_provider.max_completion_tokens,
-      });
-    } else {
-      const availableModels = Array.from(this.modelInfoMap.keys());
-      this.debugLogger.log("ModelInfo", "Model not found in available models", {
-        modelId,
-        availableModels,
-      });
+    if (!(await this.isModelAvailable(modelId))) {
       throw new Error(
-        `Invalid model: ${modelId}. Available models: ${availableModels.join(", ")}`,
+        `Invalid model: ${modelId}. Available models: ${Array.from(this.modelInfoMap.keys()).join(", ")}`,
       );
     }
+
+    const modelInfo = this.modelInfoMap.get(modelId)!; // Safe to use ! here since we checked availability
+    this.currentModel = modelId;
+    this.currentModelInfo = modelInfo;
+
+    if (!modelInfo) {
+      return;
+    }
+
+    this.debugLogger.log("ModelInfo", "Current model info", {
+      model: modelId,
+      contextLength: modelInfo.context_length,
+      maxCompletionTokens: modelInfo.top_provider.max_completion_tokens,
+    });
   }
 
   getCurrentModel(): string | null {
@@ -99,7 +99,17 @@ export class ModelInfo {
 
   async isModelAvailable(modelId: string): Promise<boolean> {
     await this.ensureInitialized();
-    return this.modelInfoMap.has(modelId);
+    const available = this.modelInfoMap.has(modelId);
+
+    if (!available) {
+      const availableModels = Array.from(this.modelInfoMap.keys());
+      this.debugLogger.log("ModelInfo", "Model not found in available models", {
+        modelId,
+        availableModels,
+      });
+    }
+
+    return available;
   }
 
   async getModelMaxCompletionTokens(modelId: string): Promise<number> {

src/services/LLM/ModelInfo.ts

@@ -1,4 +1,3 @@
-import { getModelForTryCount } from "@constants/modelScaling";
 import { DebugLogger } from "@services/logging/DebugLogger";
 import { singleton } from "tsyringe";
 import { ConfigService } from "../ConfigService";
@@ -15,9 +14,9 @@ export class ModelManager {
     private messageContextManager: MessageContextManager,
     private configService: ConfigService,
   ) {
-    // Initialize with discovery phase model from config or fallback to default
+    // Initialize with discovery phase model from config
     const config = this.configService.getConfig();
-    this.currentModel = config.discoveryModel || getModelForTryCount(null, 0);
+    this.currentModel = config.discoveryModel;
     this.modelInfo.setCurrentModel(this.currentModel);
     this.debugLogger.log("Model", "Initialized model manager", {
       model: this.currentModel,