Handling part.input for addToolResult in HITL Scenario with Enum Options

[Kinbaum]

Hi all,

I'm working on a HITL (Human-in-the-Loop) scenario where I need to specify the part.input for addToolResult. Unlike a confirm/deny flow, I have a set of options for the user to select from.

• I defined this in my input schema as an enum (e.g., a, b, c, etc).
• However, when users make a selection in the UI via addToolResult, the value of part.input is always the same and does not change based on the user’s selection.
• This behavior can be observed in the network tab and matches what I see in the HITL example in the ai repo.

This HITL workflow is confusing because:

• We are telling the LLM what the output should be,
• But when processing tool calls and running the execute function, we look at the input schema,
• Yet, there seems to be no way to control or update what part.input should be based on the user’s selection.

Could someone advise on how to properly handle this? Is there a way to ensure that part.input reflects the selected enum value during processToolCalls?

Thanks in advance for any guidance!
CC @nicoalbanese @lgrammel

[muniter]

The way I do editing of the input in human in the loop is instead of returning yes or no as the example shown in the docs I add the tool result like this:

type ToolApprovalOutput<TInput> =
  | { _brand: "tool_approved"; override?: { input: TInput } }
  | { _brand: "tool_rejected"; reason?: string };

You use the override.input and now the server when processing the messages to check for the tool approval you check if it was approved or not and what input should be used.

[Kinbaum]

Thanks for your reply @muniter. Are you able to provide a gist?

[muniter]

@Kinbaum It's not easy to provide a simple gist, here's one that shows kind of the general flow, I want to write a full article and example that shows the full flow, but this could be useful I have some types that allow the whole setup to be typesafe and tools declare dependencies and so on. (You're looking for the part that uses override.input and when sending the message how the ui modifies the original input so the ui shows the proper state.)

Feel free to reach out, happy to help.

On the server

import { z } from "zod/v4";
import {
  convertToModelMessages,
  createUIMessageStream,
  isToolUIPart,
  stepCountIs,
  streamText,
  tool,
  type InferUIMessageChunk,
  type Tool,
  type ToolCallOptions,
} from "ai";
import type { AppUIMessage, AppUIMessagePart } from "./agentTools.ts";
import type { IssueBuilder } from "#shared/utils/issueBuilder.ts";
import { openai } from "@ai-sdk/openai";

const ToolApprovalOutputSchema = z.discriminatedUnion("_brand", [
  z.object({
    _brand: z.literal("tool_approved"),
    override: z
      .record(z.string(), z.any())
      .describe("The override to apply to the tool call input if approved")
      .optional(),
  }),
  z.object({
    _brand: z.literal("tool_rejected"),
    reason: z.string().describe("The reason for rejecting the tool call").optional(),
  }),
]);

type ToolApprovalOutput<TInput> =
  | { _brand: "tool_approved"; override?: { input: TInput } }
  | { _brand: "tool_rejected"; reason?: string };

export function isToolApprovalOutput(output: unknown): output is ToolApprovalOutput<any> {
  return ToolApprovalOutputSchema.safeParse(output).success;
}

declare function getToolDefinitionFromPartType(
  partType: string,
): { execute: Tool<any, any>["execute"]; name: string } | undefined;

export async function processToolOutput(args: {
  messages: AppUIMessage[];
}): Promise<{ messages: AppUIMessage[]; parts: InferUIMessageChunk<AppUIMessage>[] }> {
  const result: {
    // How the UI messages end up looking
    messages: AppUIMessage[];
    // Parts are collected here to later use the writer to update the UI
    parts: InferUIMessageChunk<AppUIMessage>[];
  } = {
    messages: args.messages,
    parts: [],
  };
  const lastMessage = args.messages[args.messages.length - 1];
  if (!lastMessage) {
    return result;
  }

  const pendingToolCalls = lastMessage.parts.some(
    (part) => isToolUIPart(part) && isToolApprovalOutput(part.output),
  );

  if (!pendingToolCalls) {
    return result;
  }

  // Tool call receive modelMessages as part of the second arg options
  const modelMessages = convertToModelMessages(args.messages);
  const newParts: AppUIMessagePart[] = await Promise.all(
    lastMessage.parts.map(async (part) => {
      if (isToolUIPart(part) === false || isToolApprovalOutput(part.output) === false) {
        return part;
      }

      const lookup = getToolDefinitionFromPartType(part.type);
      if (!lookup) {
        return part;
      }

      const toolDefinition = lookup;
      if (!toolDefinition.execute) {
        throw new Error(`Could not find the tool execute function for tool ${toolDefinition.name}`);
      }

      const { execute, name } = toolDefinition;

      // The tool call was not approved by the user
      if (part.output._brand === "tool_rejected") {
        const message = `The user has rejected the attempt to use the tool: ${name}${part.output.reason ? `: ${part.output.reason}` : ""}`;
        const output: ToolApprovalOutput<unknown> = {
          _brand: "tool_rejected",
          reason: message,
        };
        result.parts.push({
          type: "tool-output-available",
          toolCallId: part.toolCallId,
          output,
        });
        return {
          type: part.type as any,
          toolCallId: part.toolCallId,
          state: "output-available",
          input: part.input,
          output,
        } satisfies AppUIMessagePart;
      }

      // Here if the approval result passed an override of the input, we use that
      // instead of what the LLM passed
      const input = part.output.override?.input ?? part.input;

      // Execute the tool
      const output = await execute(input, {
        toolCallId: part.toolCallId,
        messages: modelMessages,
      });

      // Write the output to the stream
      result.parts.push({
        type: "tool-output-available",
        toolCallId: part.toolCallId,
        output,
      });

      // Return the new part
      return {
        type: part.type as any,
        toolCallId: part.toolCallId,
        state: "output-available",
        input,
        output,
      } satisfies AppUIMessagePart;
    }),
  );

  // Check if the last and second last messages have the same id if they do we got to merge them
  // When passing a tool result from the client this is usually true, if we don't re arrange the messages
  // everything breaks
  if (args.messages[args.messages.length - 1]?.id === args.messages[args.messages.length - 2]?.id) {
    result.messages = [...args.messages.slice(0, -2), { ...lastMessage, parts: newParts }];
  }

  return result;
}

declare function loadOrCreateConversation(args: {
  conversationId: string;
}): Promise<{ id: string; messages: AppUIMessage[] }>;

const toolsDefinitions = {
  requires_approval: {
    tool: tool({
      description: "Tool that requires approval",
      inputSchema: z.object({
        input: z.string(),
      }),
      execute: async (input: { input: string }, options: ToolCallOptions) => {
        return "Tool 1 result";
      },
    }),
  },
} satisfies Record<
  string,
  {
    tool: Tool<any, any>;
    execute?: (input: any, options: ToolCallOptions) => Promise<any>;
  }
>;

// Object that only has the key: tool
const tools = Object.fromEntries(
  Object.entries(toolsDefinitions).map(([name, tool]) => [name, tool.tool]),
) as { [K in keyof typeof toolsDefinitions]: (typeof toolsDefinitions)[K]["tool"] };

export async function chatWithAgent(args: {
  message: AppUIMessage;
  conversationId: string;
  trigger: "submit-message" | "regenerate-message";
}): Promise<ReadableStream<InferUIMessageChunk<AppUIMessage>> | IssueBuilder> {
  const conversation = await loadOrCreateConversation({ conversationId: args.conversationId });

  let rawMessages: AppUIMessage[] = [];
  if (args.trigger === "submit-message") {
    rawMessages = [...conversation.messages, args.message];
  } else if (args.trigger === "regenerate-message") {
    const index = conversation.messages.findIndex((message) => message.id === args.message.id);
    if (index === -1) {
      throw new Error("Message not found");
    }
    rawMessages = conversation.messages.slice(0, index);
  }

  const processToolResult = await processToolOutput({
    messages: rawMessages,
  });

  const messages = processToolResult.messages;

  const uiMessageStream = createUIMessageStream<AppUIMessage>({
    originalMessages: messages,
    execute: async ({ writer }) => {
    // We write the parts to the stream that we previously collected
    // when running the tool after approval
      for (const part of processToolResult.parts) {
        writer.write(part);
      }
      const chatStream = streamText({
        model: openai.responses("gpt-4.1"),
        system: "Some prompt",
        messages: convertToModelMessages(messages),
        stopWhen: stepCountIs(5),
        tools: tools,
        onFinish: (stepResult) => {
          if (stepResult.finishReason === "tool-calls") {
            for (const entry of stepResult.content) {
              if (entry.type === "tool-call" && entry.input && !entry.providerExecuted) {
                  // The reason we finished is because there's a tool call that has not been executed,
                  // meaning in our context it needs approval.
                writer.write({
                  type: "tool-input-available",
                  toolCallId: entry.toolCallId,
                  toolName: entry.toolName,
                  input: entry.input,
                  providerMetadata: {
                    ...entry.providerMetadata,
                    app: {
                      waitingForApproval: true,
                    },
                  },
                });
              }
            }
          }
        },
      });

      chatStream.consumeStream();
      writer.merge(
        chatStream.toUIMessageStream({
          sendReasoning: true,
        }),
      );
    },
    onFinish: async (opts) => {
      // Validate and persist the messages
    },
  });

  return uiMessageStream;
}

On the Client

// Function to determine if a tool is waiting for approval
function isToolWaitingForApproval<T extends AppUIToolParts['type']>(
    part: AppUIToolParts,
  ): part is Extract<AppUIToolParts, { type: T; state: 'input-available' }> & {
    state: 'input-available';
    callProviderMetadata: { app: { waitingForApproval: true } };
  } {
    if (!isToolUIPart(part)) {
      return false;
    }
    return (
      part.state === 'input-available' &&
      (part.callProviderMetadata?.app?.waitingForApproval as boolean) === true
    );
  }

// Function to approve or reject a tool call
async function approveTool(
    args:
      | {
          approve: true;
          input?: { brand: 'tool_approved'; override?: { input: any } };
          partType: string;
          toolCallId: string;
        }
      | {
          approve: false;
          reason?: string;
          partType: string;
          toolCallId: string;
        },
  ) {
    if (args.approve) {
      if (args.input) {
        // Let's find the message and update the input
        let found = false;
        const messages = chat.messages;
        for (const message of messages) {
          for (const part of message.parts) {
            if (part.type === args.partType && part.toolCallId === args.toolCallId && part.input) {
              part.input = args.input;
              found = true;
              break;
            }
          }
        }
        if (!found) {
          throw new Error('Could not find the tool call to update the input');
        }
      }
      await chat.addToolResult({
        tool: args.partType.replace('tool-', '') as any,
        toolCallId: args.toolCallId,
        output: {
          _brand: 'tool_approved',
          override: { input: args.input },
        },
      });
    } else {
      await chat.addToolResult({
        tool: args.partType.replace('tool-', '') as any,
        toolCallId: args.toolCallId,
        output: {
          _brand: 'tool_rejected',
          reason: args.reason,
        },
      });
    }
    chat.sendMessage();
  }