The wrapModelCall hook in createMiddleware currently only supports returning AIMessage instances. This issue requests adding support for returning state updates as well, enabling middleware to modify the agent’s state during model call interception.
Current Behavior
The WrapModelCallHook type is constrained to return only AIMessage: types.ts:130-136
The implementation validates that middleware returns an AIMessage and throws an error otherwise: AgentNode.ts:540-546
Desired Behavior
Allow wrapModelCall to return either:
-
An AIMessage (current behavior)
-
A state update object that can modify the agent’s state
-
A combination of both (e.g., { message: AIMessage, state: Partial<State> })
Could you modify the state as part of the beforeModel or are you missing a specific type of information in there?
Let me describe my current issue in more detail:
Currently I am trying to create a custom middleware called preHITLValidationMiddleware, which is designed to use in sequence with humanInTheLoopMiddleware, since I wish to have the model generated args to be validated first, before raising an interrupt to the client.
Sample code:
import { AIMessage, createMiddleware } from "langchain";
import { validateToolCalls } from "./validate-tool-calls";
const maxRetries = 3;
export const preHITLValidationMiddleware = createMiddleware({
name: "PreHITLValidationMiddleware",
wrapModelCall: async (request, handler) => {
/* Initial attempt + retries */
for (let attempt = 0; attempt <= maxRetries; attempt++) {
/* 1. Call the model */
const response = await handler(request);
/* 2. If no tool calls, return immediately */
if (!response.tool_calls?.length) return response;
/* 3. Assign a tool id to tool calls that don't have one */
const toolCallsWithId = response.tool_calls.map((call) => ({
...call,
id: call.id ?? crypto.randomUUID(),
}));
/* 4. Validate tool calls */
const result = await validateToolCalls(
toolCallsWithId,
request.tools,
);
if (result.success) {
/* 5. Success, replace the tool messages args with transformed args */
return new AIMessage({
...response,
tool_calls: toolCallsWithId.map((call) => ({
...call,
/* ! here should be correct */
args: result.argsByToolCallId[call.id]!,
})),
});
} else {
/* 4. Failed, run handler again */
if (attempt < maxRetries) {
request = {
...request,
messages: [
...request.messages,
...result.invalidToolCallMessages,
],
};
} else {
return new AIMessage({
content: "Failed to validate tool calls ",
});
// finalMessage = result.invalidToolCallMessages
// .map((m) => m.name)
// .join("\n\n");
}
}
}
return new AIMessage({
invalid_tool_calls: result.invalidToolCallMessages,
});
},
});
As you can see it is kind of like the toolRetryMIddleware, where provided a maxRetry, it will loop for certain of times, and if failed it will append a ToolMessage with status error to the request before passing to the model. Now, since wrap model call only allow returning AIMessage, the tool message containing the error status will be lost in the final output. If i use beforeModel or afterModel hooks instead, i will lose track of the retry state, and the graph might run in a n infinite loop potentially.