Agent model/budget env tuning + self-correction retry

apps/api/.env.example

@@ -16,6 +16,20 @@ STACKFORGE_PROVIDER=auto
 # Codegen provider: default | nvidia
 STACKFORGE_CODEGEN_PROVIDER=default
 
+# ─── Per-agent model & budget tuning (all optional) ──────────────────────────
+# Override the model for every agent at once (e.g. a free OpenRouter model):
+#   STACKFORGE_MODEL=meta-llama/llama-3.3-70b-instruct:free
+STACKFORGE_MODEL=
+# Or override a single agent (planner|schema|api|frontend|devops|reviewer|codegen):
+#   STACKFORGE_MODEL_PLANNER=...
+#   STACKFORGE_MODEL_CODEGEN=...
+# Scale factor applied to every agent's output token budget (default 1).
+# Raise it to give weaker models more room, e.g. 1.5:
+STACKFORGE_OUTPUT_SCALE=
+# Global compression override: low | medium | high (default: per-agent setting).
+# "low" reduces context loss between agents at higher token cost:
+STACKFORGE_COMPRESSION=
+
 # OpenRouter API Keys (all 3 required for key rotation)
 OPENROUTER_API_KEY_1=
 OPENROUTER_API_KEY_2=

packages/agents/src/config/agent.configs.ts

@@ -1,94 +1,148 @@
 import type { AgentConfig, AgentName } from "@stackforge/shared";
 
+/**
+ * Resolve the model for an agent, allowing env overrides.
+ *
+ * Precedence (highest first):
+ *   1. STACKFORGE_MODEL_<AGENT>   e.g. STACKFORGE_MODEL_PLANNER
+ *   2. STACKFORGE_MODEL           global override for every agent
+ *   3. the hardcoded default passed in
+ *
+ * This lets a free OpenRouter model be swapped in without editing source,
+ * either per-agent or across the whole pipeline.
+ */
+function resolveModel(agent: AgentName, fallback: string): string {
+  const perAgent = process.env[`STACKFORGE_MODEL_${agent.toUpperCase()}`]?.trim();
+  if (perAgent !== undefined && perAgent.length > 0) {
+    return perAgent;
+  }
+
+  const global = process.env["STACKFORGE_MODEL"]?.trim();
+  if (global !== undefined && global.length > 0) {
+    return global;
+  }
+
+  return fallback;
+}
+
+/**
+ * Scale factor applied to every agent's output budget (maxOutputTokens and
+ * tokenBudget). Lets a demo give models more room without editing per-agent
+ * numbers. Defaults to 1 (no change), so existing behaviour and tests are
+ * unaffected. Set e.g. STACKFORGE_OUTPUT_SCALE=1.6 to loosen caps.
+ */
+const OUTPUT_SCALE = (() => {
+  const raw = Number.parseFloat(process.env["STACKFORGE_OUTPUT_SCALE"] ?? "");
+  return Number.isFinite(raw) && raw > 0 ? raw : 1;
+})();
+
+/**
+ * Global override for every agent's compression aggressiveness. When unset,
+ * each agent keeps its own default level. Set STACKFORGE_COMPRESSION=low to
+ * reduce context loss (better output on weak models, higher token cost).
+ */
+const COMPRESSION_OVERRIDE: AgentConfig["compressionLevel"] | undefined = (() => {
+  const raw = (process.env["STACKFORGE_COMPRESSION"] ?? "").trim().toLowerCase();
+  return raw === "low" || raw === "medium" || raw === "high" ? raw : undefined;
+})();
+
+function scaleOutput(value: number): number {
+  return Math.round(value * OUTPUT_SCALE);
+}
+
+function resolveCompression(fallback: AgentConfig["compressionLevel"]): AgentConfig["compressionLevel"] {
+  return COMPRESSION_OVERRIDE ?? fallback;
+}
+
 export const AGENT_CONFIGS: Record<AgentName, AgentConfig> = {
   planner: {
     name: "planner",
     description: "Resolves stack, project name, and top-level folder structure",
-    tokenBudget: 1000,
-    model: "openai/gpt-4o-mini",
+    tokenBudget: scaleOutput(1000),
+    model: resolveModel("planner", "openai/gpt-4o-mini"),
     maxInputTokens: 900,
-    maxOutputTokens: 700,
+    maxOutputTokens: scaleOutput(700),
     minOutputTokens: 220,
     temperature: 0.2,
-    compressionLevel: "high",
+    compressionLevel: resolveCompression("high"),
     budgetOverflowRetries: 2,
     retries: 1,
   },
   schema: {
     name: "schema",
     description: "Designs entities, fields, indexes, and relationships",
-    tokenBudget: 1500,
-    model: "openai/gpt-4o-mini",
+    tokenBudget: scaleOutput(1500),
+    model: resolveModel("schema", "openai/gpt-4o-mini"),
     maxInputTokens: 1200,
-    maxOutputTokens: 1100,
+    maxOutputTokens: scaleOutput(1100),
     minOutputTokens: 320,
     temperature: 0.1,
-    compressionLevel: "medium",
+    compressionLevel: resolveCompression("medium"),
     budgetOverflowRetries: 2,
     retries: 1,
   },
   api: {
     name: "api",
     description: "Plans REST routes, request/response shapes, and auth requirements",
-    tokenBudget: 1600,
-    model: "openai/gpt-4o-mini",
+    tokenBudget: scaleOutput(1600),
+    model: resolveModel("api", "openai/gpt-4o-mini"),
     maxInputTokens: 950,
-    maxOutputTokens: 1200,
+    maxOutputTokens: scaleOutput(1200),
     minOutputTokens: 280,
     temperature: 0.1,
-    compressionLevel: "medium",
+    compressionLevel: resolveCompression("medium"),
     budgetOverflowRetries: 2,
     retries: 1,
   },
   frontend: {
     name: "frontend",
     description: "Plans pages, React component trees, and routing",
-    tokenBudget: 1200,
-    model: "openai/gpt-4o-mini",
+    tokenBudget: scaleOutput(1200),
+    model: resolveModel("frontend", "openai/gpt-4o-mini"),
     maxInputTokens: 900,
-    maxOutputTokens: 900,
+    maxOutputTokens: scaleOutput(900),
     minOutputTokens: 250,
     temperature: 0.2,
-    compressionLevel: "high",
+    compressionLevel: resolveCompression("high"),
     budgetOverflowRetries: 2,
     retries: 1,
   },
   devops: {
     name: "devops",
     description: "Plans CI/CD pipelines, deployment targets, Docker, and env vars",
-    tokenBudget: 1000,
-    model: "openai/gpt-4o-mini",
+    tokenBudget: scaleOutput(1000),
+    model: resolveModel("devops", "openai/gpt-4o-mini"),
     maxInputTokens: 700,
-    maxOutputTokens: 850,
+    maxOutputTokens: scaleOutput(850),
     minOutputTokens: 180,
     temperature: 0.2,
-    compressionLevel: "high",
+    compressionLevel: resolveCompression("high"),
     budgetOverflowRetries: 3,
     retries: 1,
   },
   reviewer: {
     name: "reviewer",
     description: "Performs consistency checks and flags issues across the full blueprint",
-    tokenBudget: 1200,
-    model: "openai/gpt-4o-mini",
+    tokenBudget: scaleOutput(1200),
+    model: resolveModel("reviewer", "openai/gpt-4o-mini"),
     maxInputTokens: 850,
-    maxOutputTokens: 900,
+    maxOutputTokens: scaleOutput(900),
     minOutputTokens: 220,
     temperature: 0,
-    compressionLevel: "high",
+    compressionLevel: resolveCompression("high"),
     budgetOverflowRetries: 2,
     retries: 1,
   },
   codegen: {
     name: "codegen",
     description: "Generates concrete source files from the finalized blueprint",
-    tokenBudget: 2200,
-    model: "openai/gpt-4o-mini",
+    tokenBudget: scaleOutput(2200),
+    model: resolveModel("codegen", "openai/gpt-4o-mini"),
     maxInputTokens: 1100,
-    maxOutputTokens: 1200,
+    maxOutputTokens: scaleOutput(1200),
     minOutputTokens: 280,
     temperature: 0.1,
-    compressionLevel: "medium",
+    compressionLevel: resolveCompression("medium"),
     budgetOverflowRetries: 3,
     retries: 1,
   },

packages/agents/src/workflow/engine.ts

@@ -1,4 +1,4 @@
-import type { AgentName, SSEEvent } from "@stackforge/shared";
+import type { AgentName, AgentConfig, SSEEvent } from "@stackforge/shared";
 import { AGENT_CONFIGS } from "../config/agent.configs.js";
 import type { AgentCache } from "../cache/agent.cache.js";
 import type { LLMProvider, ProviderCallOutput } from "../provider/provider.interface.js";
@@ -357,7 +357,28 @@ export class WorkflowEngine {
         }
       }
 
-      validateStepOutput(step, primary.output);
+      try {
+        validateStepOutput(step, primary.output);
+      } catch (validationError) {
+        // One self-correction pass: weak/free models frequently emit malformed
+        // or incomplete JSON. Rather than fail the whole job, show the model its
+        // own output plus the validation error and ask for a corrected version.
+        const repaired = await this.attemptRepair(
+          workflow,
+          step,
+          config,
+          userPrompt,
+          primary.output,
+          validationError instanceof Error ? validationError.message : String(validationError),
+        );
+        primary = {
+          ...primary,
+          output: repaired.output,
+          tokensUsed: primary.tokensUsed + repaired.tokensUsed,
+        };
+        this.options.cache.set(primaryCacheKey, primary.output);
+        validateStepOutput(step, primary.output);
+      }
 
       // Convert raw code string to structured format for codegen
       let normalizedOutput = primary.output;
@@ -426,6 +447,52 @@ export class WorkflowEngine {
     }
   }
 
+  private async attemptRepair(
+    workflow: WorkflowDefinition,
+    step: WorkflowStep,
+    config: AgentConfig,
+    originalUserPrompt: string,
+    badOutput: unknown,
+    validationError: string,
+  ): Promise<{ output: unknown; tokensUsed: number }> {
+    const repairPrompt = [
+      "Your previous response for this workflow step was invalid and must be corrected.",
+      "",
+      `Validation error: ${validationError}`,
+      "",
+      "Your previous (invalid) output was:",
+      serializeOutput(badOutput),
+      "",
+      "Original step instructions:",
+      originalUserPrompt,
+      "",
+      "Return ONLY the corrected, strict JSON object for this step. No markdown, no code fences, no commentary.",
+    ].join("\n");
+
+    const response = await this.options.provider.call({
+      agentName: step.agent,
+      input: {
+        workflowId: workflow.id,
+        stepId: step.id,
+        repair: true,
+      },
+      options: {
+        systemPrompt:
+          "You repair invalid workflow-step output. Return strict JSON object output only that satisfies the stated schema.",
+        userPrompt: repairPrompt,
+        model: config.model,
+        maxInputTokens: config.maxInputTokens,
+        maxOutputTokens: config.maxOutputTokens,
+        temperature: 0,
+      },
+    });
+
+    return {
+      output: response.output,
+      tokensUsed: response.tokensUsed,
+    };
+  }
+
   private async compressStepOutput(
     workflow: WorkflowDefinition,
     step: WorkflowStep,

packages/agents/test/orchestrator-failure.test.ts

@@ -148,6 +148,38 @@ class CompressionFailureProvider extends FixedTokenProvider {
   }
 }
 
+/**
+ * Returns an invalid planner payload on the first attempt, then a valid payload
+ * once the engine issues a repair call (flagged via `repair: true` in input).
+ * Exercises the self-correction retry path in the workflow engine.
+ */
+class RepairablePlannerProvider extends FixedTokenProvider {
+  repairCalls = 0;
+
+  override async call(input: ProviderCallInput): Promise<ProviderCallOutput> {
+    const payload = input.input as Record<string, unknown>;
+
+    if (input.agentName === "planner" && "compressionTarget" in payload === false) {
+      if (payload["repair"] === true) {
+        this.repairCalls += 1;
+        return super.call(input);
+      }
+
+      // First (pre-repair) planner attempt: deliberately invalid shape.
+      return {
+        output: { bad: true },
+        tokensUsed: 10,
+        durationMs: 1,
+        inputTokens: 5,
+        outputTokens: 5,
+        model: input.options.model,
+      };
+    }
+
+    return super.call(input);
+  }
+}
+
 describe("Orchestrator failure propagation", () => {
   it("emits agent_failed and rejects when agent output schema validation fails", async () => {
     const events: SSEEvent[] = [];
@@ -228,4 +260,30 @@ describe("Orchestrator failure propagation", () => {
     expect(blueprint.reviewerNotes.length).toBeGreaterThan(0);
     expect((blueprint.generatedSourceFiles?.length ?? 0) > 0).toBe(true);
   });
+
+  it("recovers from invalid agent output via a single repair pass", async () => {
+    const events: SSEEvent[] = [];
+    const provider = new RepairablePlannerProvider();
+
+    const blueprint = await runOrchestrator({
+      jobId: crypto.randomUUID(),
+      prompt: "Build a task management platform",
+      projectName: "taskflow",
+      execution: {
+        enableCodeGeneration: false,
+      },
+      emit: (event) => events.push(event),
+      provider,
+      cache: new AgentCache(),
+    });
+
+    // Repair was triggered exactly once and the job still completed.
+    expect(provider.repairCalls).toBe(1);
+    expect(blueprint.projectName.length).toBeGreaterThan(0);
+    expect(blueprint.routePlan.length).toBeGreaterThan(0);
+
+    // No agent_failed should be emitted since the repair succeeded.
+    const failed = events.find((event) => event.type === "agent_failed");
+    expect(failed).toBeUndefined();
+  });
 });