3 compiler bugs related to closure variables

[pranaygp]

When working on #1362 I encountered a few compiler bugs with classes on closure variables. Tested and validated on the WDK compiler playground

Bug 1: Closer Variables not captured with class constructors

Input:

import { MockLanguageModelV3 } from 'ai/test';
import { xai as xaiProvider } from '@ai-sdk/xai';

export function mockModel(...args) {
  return async () => {
    'use step';
    return new MockLanguageModelV3(...args);
  };
}

export function xai(...args: Parameters<typeof xaiProvider>) {
  return async () => {
    'use step';
    return xaiProvider(...args);
  };
}

Notice that the args closure variable is not getting captured correctly in the workflow context or hydrated in the step context for the MockLanguageModelV3 but it is working for xaiProvider

Bug 2: Closure variables break steps from being reused as normal functions outside workflow

Sam input as above - notice that in both the functions in step output (i.e. even in the xai output), the ...args doesn't get propagated anymore correctly to the anonymous function that's extracted. This is fine when the step is being used from workflow context since we capture the the closure variables in workflow mode and rehydrate it inside the anonymous function - but this means that when using xai normal within the step, it wouldn't work since the args don't get captured and rehydrated inside the step.

The solution here is probably that in the step bundle itself, xai needs to wrap the replaced anonymous function with an AsyncLocalStorage that mimics the same way it would be called when used by a workflow - so that the step function continues to work as expected when executed from either, workflow or client code

Bug 3: Over-capturing of closure variables

I tried to sidestep the class constructor issue by wrapping it in a function like so:

import { MockLanguageModelV3 } from 'ai/test';

function mockProvider(...args) {
  return new MockLanguageModelV3(...args);
}

export function mockModel(...args) {
  return async () => {
    'use step';
    return mockProvider(...args);
  };
}

you can see in the step and workflow outputs that args is now being captured and rehydrated.

However, also notice that mockProvider is being captured and hydrated as if it was a closure variable. It shouldn't be since it's actually available inside the step bundle output, and the anonymous step should actually just be using the local function instead of the closure variable (the current output would fail at runtime since it would try and serialize a function as a closure variable).

furthermore, we need to make sure that DCE removes the mockProvider (and recursively removes the import) from the workflow output. otherwise the workflow bundle will unnecessarily try and bundle the entire import (which is likely incompatible) even though it's not being used in the workflow itself and is only being used inside the step

[pranaygp]

Found a 4th bug.

4th bug: Not properly capturing deeply nested use of closure variables

input:

import { mockProvider } from './mock-function-wrapper.js';

export function mockTextModel(text: string) {
  return async () => {
    'use step';
    // Bind closure var at step body level so SWC plugin detects it
    // const _text = text;
    return mockProvider({
      doStream: async () => ({
        stream: new ReadableStream({
          start(c) {
            for (const v of ([
              { type: 'stream-start', warnings: [] },
              { type: 'response-metadata', id: 'r', modelId: 'mock', timestamp: new Date() },
              { type: 'text-start', id: '1' },
              // the 'text' here should be detected and captured by
              // closure variables
              { type: 'text-delta', id: '1', delta: text },
              { type: 'text-end', id: '1' },
              { type: 'finish', finishReason: { unified: 'stop', raw: 'stop' }, usage: { inputTokens: { total: 5, noCache: 5 }, outputTokens: { total: 10, text: 10 } } },
            ]) as any[]) c.enqueue(v);
            c.close();
          },
        }),
      }),
    });
  };
}

notice that text is not getting tracked when it's used inside the readable stream directly. uncommenting the line that captures it as _text higher up allows us to properly capture the closure variable