Agent + structured output + middleware error

Hi!

After calling an agent (via .invoke()) with a bound middleware and structured output, I get an error:

Invalid response from “wrapModelCall” in middleware “MiddlewareName”: expected AIMessage, got object.

Because of this, I have to do something like this at the end of each middleware.
Maybe I’m doing something wrong? Please share your advice on how to avoid this.

Example middleware
      ....
        if (AIMessage.isInstance(response)) {
          return response;
        }

        if (response && typeof response === 'object') {
          const messages = Array.isArray(
            (response as { messages?: unknown }).messages
          )
            ? (response as { messages: unknown[] }).messages
            : undefined;
          const lastMessage = messages?.at(-1);

          if (lastMessage && AIMessage.isInstance(lastMessage)) {
            if ('structuredResponse' in response) {
              Object.assign(lastMessage, {
                structuredResponse: (
                  response as { structuredResponse: unknown }
                ).structuredResponse,
                messages,
              });
            }

            return lastMessage;
          }
        }
      ....

If you’re modifying the response, is it possible for you to have this middleware in aftermodel instead?

I’ve created an middleware like this, which seems to be the same as the middleware in document. And it complains: Error: Invalid response from “wrapModelCall” in middleware “HandleModelErrors”: expected AIMessage, got object

const maxRetries = 3;

const handleModelErrors = createMiddleware({

name: "HandleModelErrors",

wrapModelCall: (request, handler) => {

for (let attempt = 0; attempt < maxRetries; attempt++) {

try {

return handler(request);

      } catch (e) {

if (attempt === maxRetries - 1) {

throw e;

        }

console.log(`Retry ${attempt + 1}/${maxRetries} after error: ${e}`);

      }

    }

throw new Error("Unreachable");

  },

})

Hi!
In my case, I can actually do this using beforeAgent, but I’m not sure if I can modify the systemPrompt and model at this stage (beforeAgent), and I need to do so.
In wrapModelCall, I did this by modifying the request that I passed to the handler().
How can I do this using the beforeAgent approach?