# Plan for Implementing Concurrency

Below is a step-by-step plan for adding concurrency to the AI calls in your existing codebase. The main goal is to allow multiple AI calls to happen in parallel—especially useful for variations/iterations in Oneshot. The plan includes where to insert concurrency, how to structure it in a modular way, and sample code that fits into your existing structure with minimal impact.

### 1. Where Concurrency Matters

- Oneshot.process(): Currently, this method iterates over each variation (and each iteration per variation) in a straightforward nested loop, making calls one by one.
- Desired: Have multiple calls to the AI provider happen in parallel. For example, if you have 3 variations, each needing 2 iterations, that’s 6 calls. Instead of waiting for each call sequentially, you could run them concurrently (e.g., up to 2 or 3 at a time).

### 2. Minimal Codebase Changes

We will:
- Add a new concurrency option (e.g., maxConcurrency) to OneshotOptions.
- Build a small concurrency utility that can run tasks with a cap on how many run at once.
- Replace the sequential loops in Oneshot.process() with this concurrency approach.

You won’t need to rewrite your entire project. The rest of the code (e.g., PromptScript, Oneshotcat, ps, etc.) stays intact.

### 3. Concurrency Utility

Create a reusable module (e.g., src/concurrency.ts) to manage concurrency:

```typescript
/**
 * concurrency.ts
 * A lightweight concurrency queue/limiter.
 */

/**
 * Runs an array of async tasks with a given concurrency limit.
 * 
 * @param tasks - Array of no-arg functions that return a Promise.
 * @param maxConcurrency - Maximum number of tasks to run at once.
 * @returns Array of the resolved values from each task, in the same order.
 */
export async function runConcurrent<T>(
  tasks: Array<() => Promise<T>>,
  maxConcurrency: number
): Promise<T[]> {
  if (maxConcurrency <= 0) {
    throw new Error('maxConcurrency must be > 0');
  }

  const results: T[] = new Array(tasks.length);
  let currentIndex = 0;
  let activeCount = 0;

  return new Promise((resolve, reject) => {
    // Start worker function
    function runNext() {
      if (currentIndex >= tasks.length) {
        // All tasks have been scheduled
        if (activeCount === 0) {
          // Nothing is running, so we’re done
          resolve(results);
        }
        return;
      }

      const idx = currentIndex;
      currentIndex++;
      activeCount++;

      const task = tasks[idx];
      task()
        .then((res) => {
          results[idx] = res;
        })
        .catch((err) => {
          reject(err);
        })
        .finally(() => {
          activeCount--;
          runNext(); // Schedule the next task
        });
    }

    // Start up to maxConcurrency tasks initially
    for (let i = 0; i < Math.min(maxConcurrency, tasks.length); i++) {
      runNext();
    }
  });
}
```

Key Points:
- Order Preservation: We keep an array of results and place each result by index. This ensures the final array aligns with the original task order, even though they finish at various times.
- Max Concurrency: We never exceed maxConcurrency active tasks.


### 4. Extend OneshotOptions

In src/oneshot.ts, update the interface:

```typescript
export interface OneshotOptions {
  model: string;
  promptFile?: string;
  promptText?: string;  // if you later decide to allow direct text
  system?: string;
  systemFile?: string;
  variations?: string[];
  iterations?: number;
  /**
   * Maximum number of AI calls at once. Default: 1 (sequential).
   */
  maxConcurrency?: number;
}
```

Rationale:
- We introduce maxConcurrency?: number.
- If not provided, default to 1 (the old sequential behavior).

### 5. Modify Oneshot.process()

Now we’ll collect tasks for each variation and iteration, and then use runConcurrent():

```typescript
import { runConcurrent } from './concurrency';

export class Oneshot {
  // ...

  public async process(): Promise<ResponseEnvelope[]> {
    const {
      promptFile,
      systemFile,
      system,
      variations,
      iterations,
      model,
      maxConcurrency = 1
    } = this.options;

    // 1. Read userPrompt from file or text
    let userPrompt = '';
    if (promptFile && existsSync(promptFile)) {
      userPrompt = readFileSync(promptFile, 'utf-8');
    } 
    // else if (this.options.promptText) {
    //   userPrompt = this.options.promptText;
    // }
    // else throw new ToolError(...);

    // 2. Read systemPrompt from file if provided
    let systemPrompt = system || '';
    if (systemFile) {
      // ...
    }

    // 3. Prepare variations array
    const finalVariations = (variations && variations.length > 0) 
      ? variations 
      : [undefined];
    
    const totalTasks: Array<() => Promise<ResponseEnvelope>> = [];

    // 4. Build tasks array
    finalVariations.forEach((variation) => {
      for (let i = 0; i < (iterations || 1); i++) {
        totalTasks.push(async () => {
          // Combine user prompt + variation
          const combinedUserPrompt = variation 
            ? `${userPrompt}\n\n${variation}`
            : userPrompt;
          
          const response = await this.provider.sendPrompt({
            model,
            systemPrompt,
            userPrompt: combinedUserPrompt
          });

          return {
            variation,
            systemPrompt,
            response
          };
        });
      }
    });

    // 5. Execute tasks with concurrency
    const results = await runConcurrent(totalTasks, maxConcurrency);
    return results;
  }
}
```

Notes:
- totalTasks: We gather all calls we plan to make into a single array of “task” functions.
- runConcurrent(...): We pass in the array, along with maxConcurrency. This will run them in parallel (up to the limit) and return the results in order.

### 6. Update CLI

oneshot.ts in bin/oneshot.ts

Allow the user to specify --max-concurrency <number>:

```typescript
#!/usr/bin/env node
import { Command } from 'commander';
import { Oneshot, AnthropicProvider, OpenAIProvider } from '../src/oneshot';
import { loadConfig } from '../src/config';

const program = new Command()
  .argument('<model>', 'Model identifier')
  .argument('<prompt>', 'Prompt file')
  .option('--system <prompt>', 'System prompt')
  .option('--system-file <path>', 'System prompt file')
  .option('--variations <json>', 'JSON array of variation prompts')
  .option('--iterations <number>', 'Number of iterations', '1')
  .option('--max-concurrency <number>', 'Maximum concurrent calls', '1')
  .action(async (model, prompt, options) => {
    const config = loadConfig();
    const provider = model.includes('gpt')
      ? new OpenAIProvider(config.openaiApiKey)
      : new AnthropicProvider(config.anthropicApiKey);

    let parsedVariations: string[] | undefined;
    if (options.variations) {
      parsedVariations = JSON.parse(options.variations);
    }

    const oneshot = new Oneshot({
      model,
      promptFile: prompt,
      system: options.system,
      systemFile: options['systemFile'],
      variations: parsedVariations,
      iterations: parseInt(options.iterations),
      maxConcurrency: parseInt(options.maxConcurrency)
    }, provider);

    const results = await oneshot.process();
    
    // Print or save results
    console.log(JSON.stringify(results, null, 2));
  });

program.parse();
```

oneshotcat.ts in bin/oneshotcat.ts

Similarly, just pass maxConcurrency along:

```typescript
#!/usr/bin/env node
import { Command } from 'commander';
import { Oneshotcat } from '../src/oneshotcat';
import { loadConfig } from '../src/config';

const program = new Command()
  .argument('<model>', 'Model identifier')
  .argument('<input>', 'Input .ps.md file')
  // ... other existing options
  .option('--max-concurrency <number>', 'Maximum concurrent calls', '1')
  .action(async (model, input, options) => {
    // Merge options into the Oneshotcat
    const finalOptions = {
      model,
      promptFile: input,
      // ... other options
      maxConcurrency: parseInt(options.maxConcurrency)
    };
    
    const cat = new Oneshotcat(finalOptions);
    const results = await cat.process();
    console.log(JSON.stringify(results, null, 2));
  });

program.parse();
```

### 7. Oneshotcat Changes

Internally, Oneshotcat just sets up a Oneshot instance with whatever maxConcurrency is provided. Because you’ll likely store the concurrency setting in the combined options, no additional major changes are required—just ensure Oneshotcat passes it along:

```typescript
// oneshotcat.ts
// ...
public async process(): Promise<ResponseEnvelope[]> {
  // Step 1: ps.process() → expanded
  const expanded = await ps.process();

  // Step 2: create temp file or store in memory
  writeFileSync(tempFilePath, expanded, 'utf-8');

  // Step 3: Oneshot
  const config = loadConfig();
  const provider = this.oneshotOptions.model.includes('gpt')
    ? new OpenAIProvider(config.openaiApiKey ?? '')
    : new AnthropicProvider(config.anthropicApiKey ?? '');

  // The key here is to include `maxConcurrency`.
  const oneshot = new Oneshot({
    model: this.oneshotOptions.model,
    promptFile: tempFilePath,
    system: this.oneshotOptions.system,
    systemFile: this.oneshotOptions.systemFile,
    variations: this.oneshotOptions.variations,
    iterations: this.oneshotOptions.iterations,
    maxConcurrency: this.oneshotOptions.maxConcurrency
  }, provider);

  return oneshot.process();
}
``` 

### 7. Testing Strategy

1. Unit Test runConcurrent() in isolation:
- Scenario A: 5 tasks, maxConcurrency = 2. Ensure tasks run in the correct order and the results are in the correct array positions.
- Scenario B: Test a failing task. Confirm the promise rejects immediately (the queue stops scheduling new tasks).
- Scenario C: Check boundary conditions (0 tasks, maxConcurrency=1, etc.).
2. Integration Test:
- Use a mock AIProvider that simulates different response times.
- Provide a series of variations/iterations.
- Confirm parallel requests are happening by checking timestamps or internal logging.
- Confirm the final aggregated results are in the right order.
3. CLI Test:
- For oneshot with --max-concurrency 3 and 6 total requests (2 variations × 3 iterations each).
- Mock provider to see all 6 requests start but only 3 at a time.

### 8. Summary of Changes

1. Create src/concurrency.ts with a runConcurrent function.
2. Add maxConcurrency to OneshotOptions.
3. Refactor loops in Oneshot.process() to collect tasks and run them via runConcurrent(...).
4. Update CLI commands (oneshot.ts, oneshotcat.ts) to accept --max-concurrency.
5. Test thoroughly with a mock AI provider and concurrency-limited runs.