llmist 1.7.0 → 2.1.0

package/dist/testing/index.cjs

@@ -722,6 +722,44 @@ ${this.endPrefix}`
   }
 });
 
+// src/gadgets/exceptions.ts
+var BreakLoopException, HumanInputException, TimeoutException, AbortError;
+var init_exceptions = __esm({
+  "src/gadgets/exceptions.ts"() {
+    "use strict";
+    BreakLoopException = class extends Error {
+      constructor(message) {
+        super(message ?? "Agent loop terminated by gadget");
+        this.name = "BreakLoopException";
+      }
+    };
+    HumanInputException = class extends Error {
+      question;
+      constructor(question) {
+        super(`Human input required: ${question}`);
+        this.name = "HumanInputException";
+        this.question = question;
+      }
+    };
+    TimeoutException = class extends Error {
+      timeoutMs;
+      gadgetName;
+      constructor(gadgetName, timeoutMs) {
+        super(`Gadget '${gadgetName}' execution exceeded timeout of ${timeoutMs}ms`);
+        this.name = "TimeoutException";
+        this.gadgetName = gadgetName;
+        this.timeoutMs = timeoutMs;
+      }
+    };
+    AbortError = class extends Error {
+      constructor(message) {
+        super(message || "Gadget execution was aborted");
+        this.name = "AbortError";
+      }
+    };
+  }
+});
+
 // src/gadgets/schema-to-json.ts
 function schemaToJSONSchema(schema, options) {
   const jsonSchema = z2.toJSONSchema(schema, options ?? { target: "draft-7" });
@@ -915,6 +953,7 @@ var init_gadget = __esm({
   "src/gadgets/gadget.ts"() {
     "use strict";
     init_constants();
+    init_exceptions();
     init_schema_to_json();
     init_schema_validator();
     BaseGadget = class {
@@ -944,6 +983,136 @@ var init_gadget = __esm({
        * while maintaining runtime compatibility.
        */
       examples;
+      /**
+       * Throws an AbortError if the execution has been aborted.
+       *
+       * Call this at key checkpoints in long-running gadgets to allow early exit
+       * when the gadget has been cancelled (e.g., due to timeout). This enables
+       * resource cleanup and prevents unnecessary work after cancellation.
+       *
+       * @param ctx - The execution context containing the abort signal
+       * @throws AbortError if ctx.signal.aborted is true
+       *
+       * @example
+       * ```typescript
+       * class DataProcessor extends Gadget({
+       *   description: 'Processes data in multiple steps',
+       *   schema: z.object({ items: z.array(z.string()) }),
+       * }) {
+       *   async execute(params: this['params'], ctx?: ExecutionContext): Promise<string> {
+       *     const results: string[] = [];
+       *
+       *     for (const item of params.items) {
+       *       // Check before each expensive operation
+       *       this.throwIfAborted(ctx);
+       *
+       *       results.push(await this.processItem(item));
+       *     }
+       *
+       *     return results.join(', ');
+       *   }
+       * }
+       * ```
+       */
+      throwIfAborted(ctx) {
+        if (ctx?.signal?.aborted) {
+          throw new AbortError();
+        }
+      }
+      /**
+       * Register a cleanup function to run when execution is aborted (timeout or cancellation).
+       * The cleanup function is called immediately if the signal is already aborted.
+       * Errors thrown by the cleanup function are silently ignored.
+       *
+       * Use this to clean up resources like browser instances, database connections,
+       * or child processes when the gadget is cancelled due to timeout.
+       *
+       * @param ctx - The execution context containing the abort signal
+       * @param cleanup - Function to run on abort (can be sync or async)
+       *
+       * @example
+       * ```typescript
+       * class BrowserGadget extends Gadget({
+       *   description: 'Fetches web page content',
+       *   schema: z.object({ url: z.string() }),
+       * }) {
+       *   async execute(params: this['params'], ctx?: ExecutionContext): Promise<string> {
+       *     const browser = await chromium.launch();
+       *     this.onAbort(ctx, () => browser.close());
+       *
+       *     const page = await browser.newPage();
+       *     this.onAbort(ctx, () => page.close());
+       *
+       *     await page.goto(params.url);
+       *     const content = await page.content();
+       *
+       *     await browser.close();
+       *     return content;
+       *   }
+       * }
+       * ```
+       */
+      onAbort(ctx, cleanup) {
+        if (!ctx?.signal) return;
+        const safeCleanup = () => {
+          try {
+            const result = cleanup();
+            if (result && typeof result === "object" && "catch" in result) {
+              result.catch(() => {
+              });
+            }
+          } catch {
+          }
+        };
+        if (ctx.signal.aborted) {
+          safeCleanup();
+          return;
+        }
+        ctx.signal.addEventListener("abort", safeCleanup, { once: true });
+      }
+      /**
+       * Create an AbortController linked to the execution context's signal.
+       * When the parent signal aborts, the returned controller also aborts with the same reason.
+       *
+       * Useful for passing abort signals to child operations like fetch() while still
+       * being able to abort them independently if needed.
+       *
+       * @param ctx - The execution context containing the parent abort signal
+       * @returns A new AbortController linked to the parent signal
+       *
+       * @example
+       * ```typescript
+       * class FetchGadget extends Gadget({
+       *   description: 'Fetches data from URL',
+       *   schema: z.object({ url: z.string() }),
+       * }) {
+       *   async execute(params: this['params'], ctx?: ExecutionContext): Promise<string> {
+       *     const controller = this.createLinkedAbortController(ctx);
+       *
+       *     // fetch() will automatically abort when parent times out
+       *     const response = await fetch(params.url, { signal: controller.signal });
+       *     return response.text();
+       *   }
+       * }
+       * ```
+       */
+      createLinkedAbortController(ctx) {
+        const controller = new AbortController();
+        if (ctx?.signal) {
+          if (ctx.signal.aborted) {
+            controller.abort(ctx.signal.reason);
+          } else {
+            ctx.signal.addEventListener(
+              "abort",
+              () => {
+                controller.abort(ctx.signal.reason);
+              },
+              { once: true }
+            );
+          }
+        }
+        return controller;
+      }
       /**
        * Auto-generated instruction text for the LLM.
        * Combines name, description, and parameter schema into a formatted instruction.
@@ -1011,8 +1180,8 @@ function createGadget(config) {
     parameterSchema = config.schema;
     timeoutMs = config.timeoutMs;
     examples = config.examples;
-    execute(params) {
-      return config.execute(params);
+    execute(params, ctx) {
+      return config.execute(params, ctx);
     }
   }
   return new DynamicGadget();
@@ -2318,6 +2487,162 @@ var init_block_params = __esm({
   }
 });
 
+// src/gadgets/cost-reporting-client.ts
+var CostReportingLLMistWrapper;
+var init_cost_reporting_client = __esm({
+  "src/gadgets/cost-reporting-client.ts"() {
+    "use strict";
+    init_model_shortcuts();
+    CostReportingLLMistWrapper = class {
+      constructor(client, reportCost) {
+        this.client = client;
+        this.reportCost = reportCost;
+      }
+      /**
+       * Access to model registry for cost estimation.
+       */
+      get modelRegistry() {
+        return this.client.modelRegistry;
+      }
+      /**
+       * Quick completion with automatic cost reporting.
+       *
+       * Streams internally to track token usage, then reports the calculated cost.
+       *
+       * @param prompt - User prompt
+       * @param options - Optional configuration (model, temperature, etc.)
+       * @returns Complete text response
+       */
+      async complete(prompt, options) {
+        const model = resolveModel(options?.model ?? "haiku");
+        let result = "";
+        let inputTokens = 0;
+        let outputTokens = 0;
+        let cachedInputTokens = 0;
+        let cacheCreationInputTokens = 0;
+        const messages = [
+          ...options?.systemPrompt ? [{ role: "system", content: options.systemPrompt }] : [],
+          { role: "user", content: prompt }
+        ];
+        for await (const chunk of this.client.stream({
+          model,
+          messages,
+          temperature: options?.temperature,
+          maxTokens: options?.maxTokens
+        })) {
+          result += chunk.text ?? "";
+          if (chunk.usage) {
+            inputTokens = chunk.usage.inputTokens;
+            outputTokens = chunk.usage.outputTokens;
+            cachedInputTokens = chunk.usage.cachedInputTokens ?? 0;
+            cacheCreationInputTokens = chunk.usage.cacheCreationInputTokens ?? 0;
+          }
+        }
+        this.reportCostFromUsage(model, inputTokens, outputTokens, cachedInputTokens, cacheCreationInputTokens);
+        return result;
+      }
+      /**
+       * Quick streaming with automatic cost reporting when stream completes.
+       *
+       * Yields text chunks as they arrive, then reports cost in finally block.
+       *
+       * @param prompt - User prompt
+       * @param options - Optional configuration (model, temperature, etc.)
+       * @returns Async generator yielding text chunks
+       */
+      async *streamText(prompt, options) {
+        const model = resolveModel(options?.model ?? "haiku");
+        let inputTokens = 0;
+        let outputTokens = 0;
+        let cachedInputTokens = 0;
+        let cacheCreationInputTokens = 0;
+        const messages = [
+          ...options?.systemPrompt ? [{ role: "system", content: options.systemPrompt }] : [],
+          { role: "user", content: prompt }
+        ];
+        try {
+          for await (const chunk of this.client.stream({
+            model,
+            messages,
+            temperature: options?.temperature,
+            maxTokens: options?.maxTokens
+          })) {
+            if (chunk.text) {
+              yield chunk.text;
+            }
+            if (chunk.usage) {
+              inputTokens = chunk.usage.inputTokens;
+              outputTokens = chunk.usage.outputTokens;
+              cachedInputTokens = chunk.usage.cachedInputTokens ?? 0;
+              cacheCreationInputTokens = chunk.usage.cacheCreationInputTokens ?? 0;
+            }
+          }
+        } finally {
+          this.reportCostFromUsage(model, inputTokens, outputTokens, cachedInputTokens, cacheCreationInputTokens);
+        }
+      }
+      /**
+       * Low-level stream access with automatic cost reporting.
+       *
+       * Returns a wrapped stream that reports costs when iteration completes.
+       *
+       * @param options - Full LLM generation options
+       * @returns Wrapped LLM stream that auto-reports costs
+       */
+      stream(options) {
+        return this.createCostReportingStream(options);
+      }
+      /**
+       * Creates a wrapped stream that tracks usage and reports costs on completion.
+       */
+      createCostReportingStream(options) {
+        const innerStream = this.client.stream(options);
+        const reportCostFromUsage = this.reportCostFromUsage.bind(this);
+        const model = options.model;
+        async function* costReportingWrapper() {
+          let inputTokens = 0;
+          let outputTokens = 0;
+          let cachedInputTokens = 0;
+          let cacheCreationInputTokens = 0;
+          try {
+            for await (const chunk of innerStream) {
+              if (chunk.usage) {
+                inputTokens = chunk.usage.inputTokens;
+                outputTokens = chunk.usage.outputTokens;
+                cachedInputTokens = chunk.usage.cachedInputTokens ?? 0;
+                cacheCreationInputTokens = chunk.usage.cacheCreationInputTokens ?? 0;
+              }
+              yield chunk;
+            }
+          } finally {
+            if (inputTokens > 0 || outputTokens > 0) {
+              reportCostFromUsage(model, inputTokens, outputTokens, cachedInputTokens, cacheCreationInputTokens);
+            }
+          }
+        }
+        return costReportingWrapper();
+      }
+      /**
+       * Calculates and reports cost from token usage.
+       */
+      reportCostFromUsage(model, inputTokens, outputTokens, cachedInputTokens = 0, cacheCreationInputTokens = 0) {
+        if (inputTokens === 0 && outputTokens === 0) return;
+        const modelName = model.includes(":") ? model.split(":")[1] : model;
+        const estimate = this.client.modelRegistry.estimateCost(
+          modelName,
+          inputTokens,
+          outputTokens,
+          cachedInputTokens,
+          cacheCreationInputTokens
+        );
+        if (estimate && estimate.totalCost > 0) {
+          this.reportCost(estimate.totalCost);
+        }
+      }
+    };
+  }
+});
+
 // src/gadgets/error-formatter.ts
 var GadgetErrorFormatter;
 var init_error_formatter = __esm({
@@ -2401,38 +2726,6 @@ var init_error_formatter = __esm({
   }
 });
 
-// src/gadgets/exceptions.ts
-var BreakLoopException, HumanInputException, TimeoutException;
-var init_exceptions = __esm({
-  "src/gadgets/exceptions.ts"() {
-    "use strict";
-    BreakLoopException = class extends Error {
-      constructor(message) {
-        super(message ?? "Agent loop terminated by gadget");
-        this.name = "BreakLoopException";
-      }
-    };
-    HumanInputException = class extends Error {
-      question;
-      constructor(question) {
-        super(`Human input required: ${question}`);
-        this.name = "HumanInputException";
-        this.question = question;
-      }
-    };
-    TimeoutException = class extends Error {
-      timeoutMs;
-      gadgetName;
-      constructor(gadgetName, timeoutMs) {
-        super(`Gadget '${gadgetName}' execution exceeded timeout of ${timeoutMs}ms`);
-        this.name = "TimeoutException";
-        this.gadgetName = gadgetName;
-        this.timeoutMs = timeoutMs;
-      }
-    };
-  }
-});
-
 // src/gadgets/parser.ts
 function stripMarkdownFences(content) {
   let cleaned = content.trim();
@@ -2618,14 +2911,16 @@ var init_executor = __esm({
     init_constants();
     init_logger();
     init_block_params();
+    init_cost_reporting_client();
     init_error_formatter();
     init_exceptions();
     init_parser();
     GadgetExecutor = class {
-      constructor(registry, onHumanInputRequired, logger, defaultGadgetTimeoutMs, errorFormatterOptions) {
+      constructor(registry, onHumanInputRequired, logger, defaultGadgetTimeoutMs, errorFormatterOptions, client) {
         this.registry = registry;
         this.onHumanInputRequired = onHumanInputRequired;
         this.defaultGadgetTimeoutMs = defaultGadgetTimeoutMs;
+        this.client = client;
         this.logger = logger ?? createLogger({ name: "llmist:executor" });
         this.errorFormatter = new GadgetErrorFormatter(errorFormatterOptions);
         this.argPrefix = errorFormatterOptions?.argPrefix ?? GADGET_ARG_PREFIX;
@@ -2635,14 +2930,27 @@ var init_executor = __esm({
       argPrefix;
       /**
        * Creates a promise that rejects with a TimeoutException after the specified timeout.
+       * Aborts the provided AbortController before rejecting, allowing gadgets to clean up.
        */
-      createTimeoutPromise(gadgetName, timeoutMs) {
+      createTimeoutPromise(gadgetName, timeoutMs, abortController) {
         return new Promise((_, reject) => {
           setTimeout(() => {
-            reject(new TimeoutException(gadgetName, timeoutMs));
+            const timeoutError = new TimeoutException(gadgetName, timeoutMs);
+            abortController.abort(timeoutError.message);
+            reject(timeoutError);
           }, timeoutMs);
         });
       }
+      /**
+       * Normalizes gadget execute result to consistent format.
+       * Handles both string returns (backwards compat) and object returns with cost.
+       */
+      normalizeExecuteResult(raw) {
+        if (typeof raw === "string") {
+          return { result: raw, cost: 0 };
+        }
+        return { result: raw.result, cost: raw.cost ?? 0 };
+      }
       // Execute a gadget call asynchronously
       async execute(call) {
         const startTime = Date.now();
@@ -2737,30 +3045,53 @@ var init_executor = __esm({
             validatedParameters = schemaAwareParameters;
           }
           const timeoutMs = gadget.timeoutMs ?? this.defaultGadgetTimeoutMs;
-          let result;
+          const abortController = new AbortController();
+          let callbackCost = 0;
+          const reportCost = (amount) => {
+            if (amount > 0) {
+              callbackCost += amount;
+              this.logger.debug("Gadget reported cost via callback", {
+                gadgetName: call.gadgetName,
+                amount,
+                totalCallbackCost: callbackCost
+              });
+            }
+          };
+          const ctx = {
+            reportCost,
+            llmist: this.client ? new CostReportingLLMistWrapper(this.client, reportCost) : void 0,
+            signal: abortController.signal
+          };
+          let rawResult;
           if (timeoutMs && timeoutMs > 0) {
             this.logger.debug("Executing gadget with timeout", {
               gadgetName: call.gadgetName,
               timeoutMs
             });
-            result = await Promise.race([
-              Promise.resolve(gadget.execute(validatedParameters)),
-              this.createTimeoutPromise(call.gadgetName, timeoutMs)
+            rawResult = await Promise.race([
+              Promise.resolve(gadget.execute(validatedParameters, ctx)),
+              this.createTimeoutPromise(call.gadgetName, timeoutMs, abortController)
             ]);
           } else {
-            result = await Promise.resolve(gadget.execute(validatedParameters));
+            rawResult = await Promise.resolve(gadget.execute(validatedParameters, ctx));
           }
+          const { result, cost: returnCost } = this.normalizeExecuteResult(rawResult);
+          const totalCost = callbackCost + returnCost;
           const executionTimeMs = Date.now() - startTime;
           this.logger.info("Gadget executed successfully", {
             gadgetName: call.gadgetName,
             invocationId: call.invocationId,
-            executionTimeMs
+            executionTimeMs,
+            cost: totalCost > 0 ? totalCost : void 0,
+            callbackCost: callbackCost > 0 ? callbackCost : void 0,
+            returnCost: returnCost > 0 ? returnCost : void 0
           });
           this.logger.debug("Gadget result", {
             gadgetName: call.gadgetName,
             invocationId: call.invocationId,
             parameters: validatedParameters,
             result,
+            cost: totalCost,
             executionTimeMs
           });
           return {
@@ -2768,7 +3099,8 @@ var init_executor = __esm({
             invocationId: call.invocationId,
             parameters: validatedParameters,
             result,
-            executionTimeMs
+            executionTimeMs,
+            cost: totalCost
           };
         } catch (error) {
           if (error instanceof BreakLoopException) {
@@ -2799,6 +3131,19 @@ var init_executor = __esm({
               executionTimeMs: Date.now() - startTime
             };
           }
+          if (error instanceof AbortError) {
+            this.logger.info("Gadget execution was aborted", {
+              gadgetName: call.gadgetName,
+              executionTimeMs: Date.now() - startTime
+            });
+            return {
+              gadgetName: call.gadgetName,
+              invocationId: call.invocationId,
+              parameters: validatedParameters,
+              error: error.message,
+              executionTimeMs: Date.now() - startTime
+            };
+          }
           if (error instanceof HumanInputException) {
             this.logger.info("Gadget requested human input", {
               gadgetName: call.gadgetName,
@@ -2925,7 +3270,8 @@ var init_stream_processor = __esm({
           options.onHumanInputRequired,
           this.logger.getSubLogger({ name: "executor" }),
           options.defaultGadgetTimeoutMs,
-          { argPrefix: options.gadgetArgPrefix }
+          { argPrefix: options.gadgetArgPrefix },
+          options.client
         );
       }
       /**
@@ -3192,6 +3538,7 @@ var init_stream_processor = __esm({
               error: result.error,
               executionTimeMs: result.executionTimeMs,
               breaksLoop: result.breaksLoop,
+              cost: result.cost,
               logger: this.logger
             };
             await this.hooks.observers.onGadgetExecutionComplete(context);
@@ -3548,6 +3895,17 @@ var init_agent = __esm({
                 llmOptions = { ...llmOptions, ...action.modifiedOptions };
               }
             }
+            await this.safeObserve(async () => {
+              if (this.hooks.observers?.onLLMCallReady) {
+                const context = {
+                  iteration: currentIteration,
+                  maxIterations: this.maxIterations,
+                  options: llmOptions,
+                  logger: this.logger
+                };
+                await this.hooks.observers.onLLMCallReady(context);
+              }
+            });
             this.logger.info("Calling LLM", { model: this.model });
             this.logger.silly("LLM request details", {
               model: llmOptions.model,
@@ -3568,7 +3926,8 @@ var init_agent = __esm({
               onHumanInputRequired: this.onHumanInputRequired,
               stopOnGadgetError: this.stopOnGadgetError,
               shouldContinueAfterError: this.shouldContinueAfterError,
-              defaultGadgetTimeoutMs: this.defaultGadgetTimeoutMs
+              defaultGadgetTimeoutMs: this.defaultGadgetTimeoutMs,
+              client: this.client
             });
             const result = await processor.process(stream2);
             for (const output of result.outputs) {
@@ -6700,10 +7059,18 @@ async function testGadget(gadget, params, options) {
     validatedParams = validationResult.data;
   }
   try {
-    const result = await Promise.resolve(gadget.execute(validatedParams));
+    const rawResult = await Promise.resolve(gadget.execute(validatedParams));
+    if (typeof rawResult === "string") {
+      return {
+        result: rawResult,
+        validatedParams,
+        cost: 0
+      };
+    }
     return {
-      result,
-      validatedParams
+      result: rawResult.result,
+      validatedParams,
+      cost: rawResult.cost ?? 0
     };
   } catch (error) {
     return {