feat(cli): add --batch-size parameter to prevent context leaking

.changeset/add-batch-size.md

@@ -0,0 +1,5 @@
+---
+"lingo.dev": minor
+---
+
+feat: add `--batch-size` parameter to `run` and `i18n` commands to prevent context leaking

packages/cli/src/cli/cmd/i18n.ts

@@ -90,6 +90,11 @@ export default new Command()
     "--strict",
     "Stop immediately on first error instead of continuing to process remaining buckets and locales (fail-fast mode)",
   )
+  .option(
+    "--batch-size <number>",
+    "Number of translations to process in a single batch",
+    parseInt,
+  )
   .action(async function (options) {
     updateGitignore();
 
@@ -440,6 +445,7 @@ export default new Command()
                   apiKey: settings.auth.apiKey,
                   apiUrl: settings.auth.apiUrl,
                   engineId: i18nConfig!.engineId,
+                  batchSize: flags.batchSize,
                 });
                 processPayload = withExponentialBackoff(
                   processPayload,
@@ -662,6 +668,7 @@ function parseFlags(options: any) {
     file: Z.array(Z.string()).optional(),
     interactive: Z.boolean().prefault(false),
     debug: Z.boolean().prefault(false),
+    batchSize: Z.number().min(1).max(250).optional(),
   }).parse(options);
 }
 

packages/cli/src/cli/cmd/run/_types.ts

@@ -56,5 +56,6 @@ export const flagsSchema = z.object({
   debounce: z.number().positive().prefault(5000), // 5 seconds default
   sound: z.boolean().optional(),
   pseudo: z.boolean().optional(),
+  batchSize: z.number().min(1).max(250).optional(),
 });
 export type CmdRunFlags = z.infer<typeof flagsSchema>;

packages/cli/src/cli/cmd/run/index.ts

@@ -123,6 +123,11 @@ export default new Command()
     "--pseudo",
     "Enable pseudo-localization mode: automatically pseudo-translates all extracted strings with accented characters and visual markers without calling any external API. Useful for testing UI internationalization readiness",
   )
+  .option(
+    "--batch-size <number>",
+    "Number of translations to process in a single batch (not applicable when using lingo.dev provider)",
+    (val: string) => parseInt(val),
+  )
   .action(async (args) => {
     let userIdentity: UserIdentity = null;
     try {

packages/cli/src/cli/cmd/run/setup.ts

@@ -54,7 +54,12 @@ export default async function setup(input: CmdRunContext) {
             ctx.flags.pseudo || ctx.config?.dev?.usePseudotranslator;
           const provider = isPseudo ? "pseudo" : ctx.config?.provider;
           const engineId = ctx.config?.engineId;
-          ctx.localizer = createLocalizer(provider, engineId, ctx.flags.apiKey);
+          ctx.localizer = createLocalizer(
+            provider,
+            engineId,
+            ctx.flags.apiKey,
+            ctx.flags.batchSize,
+          );
           if (!ctx.localizer) {
             throw new Error(
               "Could not create localization provider. Please check your i18n.json configuration.",

packages/cli/src/cli/localizer/explicit.ts

@@ -6,14 +6,16 @@ import { createMistral } from "@ai-sdk/mistral";
 import { I18nConfig } from "@lingo.dev/_spec";
 import chalk from "chalk";
 import dedent from "dedent";
-import { ILocalizer, LocalizerData } from "./_types";
+import { ILocalizer, LocalizerData, LocalizerProgressFn } from "./_types";
 import { LanguageModel, ModelMessage, generateText } from "ai";
 import { colors } from "../constants";
 import { jsonrepair } from "jsonrepair";
 import { createOllama } from "ollama-ai-provider-v2";
-
+import _ from "lodash";
+import { extractPayloadChunks } from "../utils/chunk";
 export default function createExplicitLocalizer(
   provider: NonNullable<I18nConfig["provider"]>,
+  batchSize?: number,
 ): ILocalizer {
   const settings = provider.settings || {};
 
@@ -26,10 +28,10 @@ export default function createExplicitLocalizer(
           To fix this issue:
           1. Switch to one of the supported providers, or
           2. Remove the ${chalk.italic(
-          "provider",
-        )} node from your i18n.json configuration to switch to ${chalk.hex(
-          colors.green,
-        )("Lingo.dev")}
+            "provider",
+          )} node from your i18n.json configuration to switch to ${chalk.hex(
+            colors.green,
+          )("Lingo.dev")}
 
           ${chalk.hex(colors.blue)("Docs: https://lingo.dev/go/docs")}
         `,
@@ -42,6 +44,7 @@ export default function createExplicitLocalizer(
         apiKeyName: "OPENAI_API_KEY",
         baseUrl: provider.baseUrl,
         settings,
+        batchSize,
       });
     case "anthropic":
       return createAiSdkLocalizer({
@@ -52,6 +55,7 @@ export default function createExplicitLocalizer(
         apiKeyName: "ANTHROPIC_API_KEY",
         baseUrl: provider.baseUrl,
         settings,
+        batchSize,
       });
     case "google":
       return createAiSdkLocalizer({
@@ -62,6 +66,7 @@ export default function createExplicitLocalizer(
         apiKeyName: "GOOGLE_API_KEY",
         baseUrl: provider.baseUrl,
         settings,
+        batchSize,
       });
     case "openrouter":
       return createAiSdkLocalizer({
@@ -72,6 +77,7 @@ export default function createExplicitLocalizer(
         apiKeyName: "OPENROUTER_API_KEY",
         baseUrl: provider.baseUrl,
         settings,
+        batchSize,
       });
     case "ollama":
       return createAiSdkLocalizer({
@@ -80,6 +86,7 @@ export default function createExplicitLocalizer(
         prompt: provider.prompt,
         skipAuth: true,
         settings,
+        batchSize,
       });
     case "mistral":
       return createAiSdkLocalizer({
@@ -90,6 +97,7 @@ export default function createExplicitLocalizer(
         apiKeyName: "MISTRAL_API_KEY",
         baseUrl: provider.baseUrl,
         settings,
+        batchSize,
       });
   }
 }
@@ -120,26 +128,29 @@ function createAiSdkLocalizer(params: {
   baseUrl?: string;
   skipAuth?: boolean;
   settings?: { temperature?: number };
+  batchSize?: number;
 }): ILocalizer {
   const skipAuth = params.skipAuth === true;
 
   const apiKey = process.env[params?.apiKeyName ?? ""];
   if (!skipAuth && (!apiKey || !params.apiKeyName)) {
     throw new Error(
       dedent`
-        You're trying to use raw ${chalk.dim(params.id)} API for translation. ${params.apiKeyName
-          ? `However, ${chalk.dim(
-            params.apiKeyName,
-          )} environment variable is not set.`
-          : "However, that provider is unavailable."
+        You're trying to use raw ${chalk.dim(params.id)} API for translation. ${
+          params.apiKeyName
+            ? `However, ${chalk.dim(
+                params.apiKeyName,
+              )} environment variable is not set.`
+            : "However, that provider is unavailable."
         }
 
         To fix this issue:
-        1. ${params.apiKeyName
-          ? `Set ${chalk.dim(
-            params.apiKeyName,
-          )} in your environment variables`
-          : "Set the environment variable for your provider (if required)"
+        1. ${
+          params.apiKeyName
+            ? `Set ${chalk.dim(
+                params.apiKeyName,
+              )} in your environment variables`
+            : "Set the environment variable for your provider (if required)"
         }, or
         2. Remove the ${chalk.italic(
           "provider",
@@ -183,85 +194,132 @@ function createAiSdkLocalizer(params: {
         return { valid: false, error: errorMessage };
       }
     },
-    localize: async (input: LocalizerData) => {
-      const systemPrompt = params.prompt
-        .replaceAll("{source}", input.sourceLocale)
-        .replaceAll("{target}", input.targetLocale);
-      const shots = [
-        [
-          {
-            sourceLocale: "en",
-            targetLocale: "es",
-            data: {
-              message: "Hello, world!",
-            },
-          },
-          {
-            sourceLocale: "en",
-            targetLocale: "es",
-            data: {
-              message: "Hola, mundo!",
+    localize: async (
+      input: LocalizerData,
+      onProgress?: LocalizerProgressFn,
+    ) => {
+      const chunks = extractPayloadChunks(
+        input.processableData,
+        params.batchSize,
+      );
+      const subResults: Record<string, any>[] = [];
+
+      for (let i = 0; i < chunks.length; i++) {
+        const chunk = chunks[i];
+
+        const systemPrompt = params.prompt
+          .replaceAll("{source}", input.sourceLocale)
+          .replaceAll("{target}", input.targetLocale);
+
+        const shots = [
+          [
+            {
+              sourceLocale: "en",
+              targetLocale: "es",
+              data: {
+                message: "Hello, world!",
+              },
             },
-          },
-        ],
-        [
-          {
-            sourceLocale: "en",
-            targetLocale: "es",
-            data: {
-              spring: "Spring",
+            {
+              sourceLocale: "en",
+              targetLocale: "es",
+              data: {
+                message: "Hola, mundo!",
+              },
             },
-            hints: {
-              spring: ["A source of water"],
+          ],
+          [
+            {
+              sourceLocale: "en",
+              targetLocale: "es",
+              data: {
+                spring: "Spring",
+              },
+              hints: {
+                spring: ["A source of water"],
+              },
             },
-          },
-          {
-            sourceLocale: "en",
-            targetLocale: "es",
-            data: {
-              spring: "Manantial",
+            {
+              sourceLocale: "en",
+              targetLocale: "es",
+              data: {
+                spring: "Manantial",
+              },
             },
-          },
-        ],
-      ];
+          ],
+        ];
 
-      const hasHints = input.hints && Object.keys(input.hints).length > 0;
+        const chunkHints = input.hints
+          ? _.pick(input.hints, Object.keys(chunk))
+          : undefined;
+        const hasHints = chunkHints && Object.keys(chunkHints).length > 0;
 
-      const payload = {
-        sourceLocale: input.sourceLocale,
-        targetLocale: input.targetLocale,
-        data: input.processableData,
-        ...(hasHints && { hints: input.hints }),
-      };
+        const payload = {
+          sourceLocale: input.sourceLocale,
+          targetLocale: input.targetLocale,
+          data: chunk,
+          ...(hasHints && { hints: chunkHints }),
+        };
 
-      const response = await generateText({
-        model,
-        ...params.settings,
-        messages: [
-          { role: "system", content: systemPrompt },
-          ...shots.flatMap(
-            ([userShot, assistantShot]) =>
-              [
-                { role: "user", content: JSON.stringify(userShot) },
-                { role: "assistant", content: JSON.stringify(assistantShot) },
-              ] as ModelMessage[],
-          ),
-          { role: "user", content: JSON.stringify(payload) },
-        ],
-      });
+        const response = await generateText({
+          model,
+          ...params.settings,
+          messages: [
+            { role: "system", content: systemPrompt },
+            ...shots.flatMap(
+              ([userShot, assistantShot]) =>
+                [
+                  { role: "user", content: JSON.stringify(userShot) },
+                  { role: "assistant", content: JSON.stringify(assistantShot) },
+                ] as ModelMessage[],
+            ),
+            { role: "user", content: JSON.stringify(payload) },
+          ],
+        });
 
-      const result = parseModelResponse(response.text);
+        let result: any;
+        try {
+          result = parseModelResponse(response.text);
+        } catch (e2) {
+          const snippet =
+            response.text.length > 500
+              ? `${response.text.slice(0, 500)}…`
+              : response.text;
+          console.error(
+            `Failed to parse response from ${params.id}. Response snippet: ${snippet}`,
+          );
+          throw new Error(
+            `Failed to parse response from ${params.id}: ${e2} (Snippet: ${snippet})`,
+          );
+        }
+        let finalResult: Record<string, any> = {};
 
-      // Handle both object and string responses
-      if (typeof result.data === "object" && result.data !== null) {
-        return result.data;
+        // Handle both object and string responses
+        if (typeof result?.data === "object" && result.data !== null) {
+          finalResult = result.data;
+        } else if (typeof result?.data === "string") {
+          // Handle string responses where the model double-stringified the JSON
+          try {
+            const parsed = parseModelResponse(result.data);
+            finalResult = parsed.data || parsed || {};
+          } catch (e) {
+            console.error(
+              `Failed to parse nested JSON response. Snippet: ${result.data.slice(0, 100)}...`,
+            );
+            throw new Error(
+              `Failed to parse nested JSON response: ${e} (Snippet: ${result.data.slice(0, 100)}...)`,
+            );
+          }
+        }
+
+        subResults.push(finalResult);
+        if (onProgress) {
+          onProgress(((i + 1) / chunks.length) * 100, chunk, finalResult);
+        }
       }
 
-      // Handle string responses - extract and repair JSON
-      const index = result.data.indexOf("{");
-      const lastIndex = result.data.lastIndexOf("}");
-      const trimmed = result.data.slice(index, lastIndex + 1);
-      return JSON.parse(jsonrepair(trimmed)).data;
+      const finalMergedResult = _.merge({}, ...subResults);
+      return finalMergedResult;
     },
   };
 }