wrapTraced works with Promise and non-Promise callbacks

[nikitatome]

The core issue is unrelated to the ai-sdk integration. The customer is attempting to use the wrapTraced with a non-Promise and wrapTraced assumes that the callback is a Promise. There's existing code that tries to handle this, but the AsyncFlush args makes this hard to clean up.

---

Previously:

I'm using AI SDK v5 and am trying to use wrapTraced around my streamText function. However, I'm using the createUIMessageStream functionality as well, and it appears that the wrapTraced is dropping the toUIMessageStream function, so I'm unable to use it. Our code looks similar to the following simplified example:

  const stream = createUIMessageStream({
    originalMessages,
    execute: async ({ writer }) => {
      const result = tracer.wrap(
        'streamText',
        {
          tags: {
            threadId,
          },
        },
        wrapTraced(streamText, {
          name: 'streamText',
          event: {
            input: {
              metadata: {
                threadId,
                modelName,
              },
            },
          },
        })
      )({...});

      // Transform UI message stream to restore IDs for frontend
      const uiStream = result.toUIMessageStream();
    },
  });

  pipeUIMessageStreamToResponse({
    response: res,
    stream,
  });

[ibolmo]

@nikitatome on closer look I'd recommend to follow how we wrapped streamText in our ai-sdk integration (here).

A couple of notes:

I will see if we can some how make wrapTraced avoid always returning a Promise object. streamText returns a DefaultStreamTextResult class instance and that's why you're getting the error. We worked around this by avoiding the async syntax in our implementation.

I think for wrapTraced to do this well, we'll have to call the callback function, and check the returned value. I know this is highly used feature so I'm going to put in more time to get this right.