smartagent/ts/smartagent.classes.agent.ts

// Retry backoff and context overflow logic derived from opencode (MIT) — https://github.com/sst/opencode

import * as plugins from './plugins.js';
import type { IAgentRunOptions, IAgentRunResult, IAgentToolCallRecord } from './smartagent.interfaces.js';
import { ContextOverflowError } from './smartagent.interfaces.js';

// Retry constants
const RETRY_INITIAL_DELAY = 2000;
const RETRY_BACKOFF_FACTOR = 2;
const RETRY_MAX_DELAY = 30_000;
const MAX_RETRY_ATTEMPTS = 8;

function retryDelay(attempt: number, headers?: Record<string, string>): number {
  if (headers) {
    const ms = headers['retry-after-ms'];
    if (ms) {
      const n = parseFloat(ms);
      if (!isNaN(n)) return n;
    }
    const after = headers['retry-after'];
    if (after) {
      const secs = parseFloat(after);
      if (!isNaN(secs)) return Math.ceil(secs * 1000);
      const date = Date.parse(after) - Date.now();
      if (!isNaN(date) && date > 0) return Math.ceil(date);
    }
  }
  return Math.min(
    RETRY_INITIAL_DELAY * Math.pow(RETRY_BACKOFF_FACTOR, attempt - 1),
    RETRY_MAX_DELAY,
  );
}

async function sleep(ms: number, signal?: AbortSignal): Promise<void> {
  return new Promise((resolve, reject) => {
    if (signal?.aborted) {
      reject(new DOMException('Aborted', 'AbortError'));
      return;
    }
    const t = setTimeout(resolve, ms);
    signal?.addEventListener(
      'abort',
      () => {
        clearTimeout(t);
        reject(new DOMException('Aborted', 'AbortError'));
      },
      { once: true },
    );
  });
}

function isRetryableError(err: unknown): boolean {
  const status = (err as any)?.status ?? (err as any)?.statusCode;
  if (status === 429 || status === 529 || status === 503) return true;
  if (err instanceof Error) {
    const msg = err.message.toLowerCase();
    if (msg.includes('rate limit') || msg.includes('overloaded') || msg.includes('too many requests')) {
      return true;
    }
  }
  return false;
}

function isContextOverflow(err: unknown): boolean {
  if (err instanceof Error) {
    const msg = err.message.toLowerCase();
    return (
      msg.includes('context_length_exceeded') ||
      msg.includes('context window') ||
      msg.includes('maximum context length') ||
      msg.includes('too many tokens') ||
      msg.includes('input is too long') ||
      (err as any)?.name === 'AI_ContextWindowExceededError'
    );
  }
  return false;
}

function parseToolInput(input: unknown): unknown {
  if (typeof input !== 'string') return input;
  try {
    return JSON.parse(input);
  } catch {
    return input;
  }
}

function errorToString(error: unknown): string {
  if (error instanceof Error) return error.message;
  return String(error);
}

function recordToolCall(
  toolCalls: IAgentToolCallRecord[],
  toolCallIndexes: Map<string, number>,
  toolCall: unknown,
  update: { output?: unknown; error?: unknown } = {},
): void {
  const call = toolCall as any;
  const toolCallId = call?.toolCallId;
  const nextRecord: IAgentToolCallRecord = {
    toolName: String(call?.toolName ?? ''),
    input: parseToolInput(call?.input ?? call?.args),
  };
  const hasOutput = Object.prototype.hasOwnProperty.call(update, 'output');
  const hasError = Object.prototype.hasOwnProperty.call(update, 'error');

  if (hasOutput) nextRecord.output = update.output;
  if (hasError && update.error !== undefined) nextRecord.error = errorToString(update.error);

  const existingIndex = typeof toolCallId === 'string' ? toolCallIndexes.get(toolCallId) : undefined;
  if (existingIndex !== undefined) {
    const existingRecord = toolCalls[existingIndex];
    existingRecord.toolName = nextRecord.toolName || existingRecord.toolName;
    if (nextRecord.input !== undefined) existingRecord.input = nextRecord.input;
    if (hasOutput) existingRecord.output = nextRecord.output;
    if (nextRecord.error !== undefined) existingRecord.error = nextRecord.error;
    return;
  }

  toolCalls.push(nextRecord);
  if (typeof toolCallId === 'string') {
    toolCallIndexes.set(toolCallId, toolCalls.length - 1);
  }
}

export async function runAgent(options: IAgentRunOptions): Promise<IAgentRunResult> {
  let stepCount = 0;
  let attempt = 0;
  let totalInput = 0;
  let totalOutput = 0;
  let validationRetries = 0;
  const toolCalls: IAgentToolCallRecord[] = [];
  const toolCallIndexes = new Map<string, number>();

  const tools = options.tools ?? {};

  // Add a no-op sink for repaired-but-unrecognised tool calls
  const allTools: plugins.ToolSet = {
    ...tools,
    invalid: plugins.tool({
      description: 'Sink for unrecognised tool calls — returns an error message to the model',
      inputSchema: plugins.z.object({
        tool: plugins.z.string(),
        error: plugins.z.string(),
      }),
      execute: async ({ tool, error }: { tool: string; error: string }) =>
        `Unknown tool "${tool}": ${error}`,
    }),
  };

  // Build messages — streamText requires either prompt OR messages, not both
  let messages: plugins.ModelMessage[] = options.messages
    ? [...options.messages, { role: 'user' as const, content: options.prompt }]
    : [{ role: 'user' as const, content: options.prompt }];

  while (true) {
    try {
      const result = plugins.streamText({
        model: options.model,
        system: options.system,
        messages,
        tools: allTools,
        providerOptions: options.providerOptions,
        stopWhen: plugins.stepCountIs(options.maxSteps ?? 20),
        maxRetries: 0, // handled manually below
        abortSignal: options.abort,

        experimental_repairToolCall: async ({ toolCall, tools: availableTools, error }) => {
          const lower = toolCall.toolName.toLowerCase();
          if (lower !== toolCall.toolName && (availableTools as any)[lower]) {
            return { ...toolCall, toolName: lower };
          }
          return {
            ...toolCall,
            toolName: 'invalid',
            args: JSON.stringify({
              tool: toolCall.toolName,
              error: String(error),
            }),
          };
        },

        onChunk: ({ chunk }) => {
          if (chunk.type === 'text-delta' && options.onToken) {
            options.onToken((chunk as any).textDelta ?? (chunk as any).text ?? '');
          }
        },

        experimental_onToolCallStart: options.onToolCall
          ? ({ toolCall }) => {
              const input = parseToolInput((toolCall as any).input ?? (toolCall as any).args);
              recordToolCall(toolCalls, toolCallIndexes, toolCall);
              options.onToolCall!(toolCall.toolName, input);
            }
          : ({ toolCall }) => {
              recordToolCall(toolCalls, toolCallIndexes, toolCall);
            },

        experimental_onToolCallFinish: options.onToolResult
          ? (event) => {
              recordToolCall(
                toolCalls,
                toolCallIndexes,
                event.toolCall,
                event.success ? { output: event.output } : { error: event.error },
              );
              options.onToolResult!(event.toolCall.toolName, event.success ? event.output : undefined);
            }
          : (event) => {
              recordToolCall(
                toolCalls,
                toolCallIndexes,
                event.toolCall,
                event.success ? { output: event.output } : { error: event.error },
              );
            },

        onStepFinish: ({ usage, toolCalls: stepToolCalls, toolResults, content }) => {
          stepCount++;
          totalInput += usage?.inputTokens ?? 0;
          totalOutput += usage?.outputTokens ?? 0;
          for (const toolCall of stepToolCalls) {
            recordToolCall(toolCalls, toolCallIndexes, toolCall);
          }
          for (const toolResult of toolResults) {
            recordToolCall(toolCalls, toolCallIndexes, toolResult, { output: (toolResult as any).output });
          }
          for (const part of content) {
            if ((part as any).type === 'tool-error') {
              recordToolCall(toolCalls, toolCallIndexes, part, { error: (part as any).error });
            }
          }
        },
      });

      // Consume the stream and collect results
      const text = await result.text;
      const finishReason = await result.finishReason;
      const responseData = await result.response;
      const responseMessages = responseData.messages as plugins.ModelMessage[];

      attempt = 0; // reset on success

      const runResult: IAgentRunResult = {
        text,
        messages: responseMessages,
        steps: stepCount,
        finishReason,
        usage: {
          inputTokens: totalInput,
          outputTokens: totalOutput,
          totalTokens: totalInput + totalOutput,
        },
        toolCalls,
      };

      if (options.validateCompletion) {
        const validationPrompt = await options.validateCompletion(runResult);
        if (typeof validationPrompt === 'string') {
          if (validationRetries >= (options.maxValidationRetries ?? 0)) {
            throw new Error(`Agent completion validation failed: ${validationPrompt}`);
          }
          validationRetries++;
          messages = [
            ...messages,
            ...responseMessages,
            { role: 'user' as const, content: validationPrompt },
          ];
          continue;
        }
      }

      return runResult;
    } catch (err: unknown) {
      // Abort — don't retry
      if (err instanceof DOMException && err.name === 'AbortError') throw err;

      // Rate limit / overload — retry with backoff
      if (isRetryableError(err) && attempt < MAX_RETRY_ATTEMPTS) {
        attempt++;
        const headers = (err as any)?.responseHeaders ?? (err as any)?.headers;
        const delay = retryDelay(attempt, headers);
        await sleep(delay, options.abort);
        continue;
      }

      // Context overflow — compact and retry if handler provided
      if (isContextOverflow(err)) {
        if (!options.onContextOverflow) throw new ContextOverflowError();
        messages = await options.onContextOverflow(messages);
        continue;
      }

      throw err;
    }
  }
}