Skip to content

fix(ai): collect trailing local tool approvals #17036

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
iroiro147 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(ai): collect trailing local tool approvals #17036

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
173 changes: 173 additions & 0 deletions packages/ai/src/generate-text/collect-tool-approvals.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,47 @@ describe('collectToolApprovals', () => {
`);
});

it('should return approved approval with approved response followed by user context', () => {
const result = collectToolApprovals({
messages: [
{
role: 'assistant',
content: [
{
type: 'tool-call',
toolCallId: 'call-1',
toolName: 'tool1',
input: { value: 'test-input' },
},
{
type: 'tool-approval-request',
approvalId: 'approval-id-1',
toolCallId: 'call-1',
},
],
},
{
role: 'tool',
content: [
{
type: 'tool-approval-response',
approvalId: 'approval-id-1',
approved: true,
},
],
},
{
role: 'user',
content: 'extra context',
},
],
});

expect(result.approvedToolApprovals).toHaveLength(1);
expect(result.approvedToolApprovals[0].toolCall.toolCallId).toBe('call-1');
expect(result.deniedToolApprovals).toHaveLength(0);
});

it('should return processed approval with approved response and tool result', () => {
const result = collectToolApprovals({
messages: [
Expand Down Expand Up @@ -156,6 +197,54 @@ describe('collectToolApprovals', () => {
`);
});

it('should return processed approval with approved response and tool result followed by user context', () => {
const result = collectToolApprovals({
messages: [
{
role: 'assistant',
content: [
{
type: 'tool-call',
toolCallId: 'call-1',
toolName: 'tool1',
input: { value: 'test-input' },
},
{
type: 'tool-approval-request',
approvalId: 'approval-id-1',
toolCallId: 'call-1',
},
],
},
{
role: 'tool',
content: [
{
type: 'tool-approval-response',
approvalId: 'approval-id-1',
approved: true,
},
{
type: 'tool-result',
toolCallId: 'call-1',
toolName: 'tool1',
output: { type: 'text', value: 'test-output' },
},
],
},
{
role: 'user',
content: 'extra context',
},
],
});

expect(result).toEqual({
approvedToolApprovals: [],
deniedToolApprovals: [],
});
});

it('should return denied approval with denied response', () => {
const result = collectToolApprovals({
messages: [
Expand Down Expand Up @@ -219,6 +308,48 @@ describe('collectToolApprovals', () => {
`);
});

it('should return denied approval with denied response followed by system context', () => {
const result = collectToolApprovals({
messages: [
{
role: 'assistant',
content: [
{
type: 'tool-call',
toolCallId: 'call-1',
toolName: 'tool1',
input: { value: 'test-input' },
},
{
type: 'tool-approval-request',
approvalId: 'approval-id-1',
toolCallId: 'call-1',
},
],
},
{
role: 'tool',
content: [
{
type: 'tool-approval-response',
approvalId: 'approval-id-1',
approved: false,
reason: 'test-reason',
},
],
},
{
role: 'system',
content: 'extra context',
},
],
});

expect(result.approvedToolApprovals).toHaveLength(0);
expect(result.deniedToolApprovals).toHaveLength(1);
expect(result.deniedToolApprovals[0].toolCall.toolCallId).toBe('call-1');
});

it('should return processed approval with denied response and tool result', () => {
const result = collectToolApprovals({
messages: [
Expand Down Expand Up @@ -266,6 +397,48 @@ describe('collectToolApprovals', () => {
`);
});

it('should ignore approval responses before a later assistant response', () => {
const result = collectToolApprovals({
messages: [
{
role: 'assistant',
content: [
{
type: 'tool-call',
toolCallId: 'call-1',
toolName: 'tool1',
input: { value: 'test-input' },
},
{
type: 'tool-approval-request',
approvalId: 'approval-id-1',
toolCallId: 'call-1',
},
],
},
{
role: 'tool',
content: [
{
type: 'tool-approval-response',
approvalId: 'approval-id-1',
approved: true,
},
],
},
{
role: 'assistant',
content: 'later response',
},
],
});

expect(result).toEqual({
approvedToolApprovals: [],
deniedToolApprovals: [],
});
});

it('should throw when approval response has unknown approvalId', () => {
expect(() =>
collectToolApprovals({
Expand Down
75 changes: 43 additions & 32 deletions packages/ai/src/generate-text/collect-tool-approvals.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import type {
} from '@ai-sdk/provider-utils';
import { InvalidToolApprovalError } from '../error/invalid-tool-approval-error';
import { ToolCallNotFoundForApprovalError } from '../error/tool-call-not-found-for-approval-error';
import { getLatestAssistantApprovalTurn } from '../util/get-latest-assistant-approval-turn';
import type { TypedToolCall } from './tool-call';
import type { TypedToolResult } from './tool-result';

Expand All @@ -15,9 +16,20 @@ export type CollectedToolApprovals<TOOLS extends ToolSet> = {
toolCall: TypedToolCall<TOOLS>;
};

function createEmptyToolApprovals<TOOLS extends ToolSet>(): {
approvedToolApprovals: Array<CollectedToolApprovals<TOOLS>>;
deniedToolApprovals: Array<CollectedToolApprovals<TOOLS>>;
} {
return {
approvedToolApprovals: [],
deniedToolApprovals: [],
};
}

/**
* If the last message is a tool message, this function collects all tool approvals
* from that message.
* Collects unresolved tool approvals for the latest assistant turn. Approval
* responses may be followed by non-assistant context messages, but are not
* replayed once a later assistant response exists.
*/
export function collectToolApprovals<TOOLS extends ToolSet>({
messages,
Expand All @@ -27,15 +39,14 @@ export function collectToolApprovals<TOOLS extends ToolSet>({
approvedToolApprovals: Array<CollectedToolApprovals<TOOLS>>;
deniedToolApprovals: Array<CollectedToolApprovals<TOOLS>>;
} {
const lastMessage = messages.at(-1);
const approvalTurn = getLatestAssistantApprovalTurn(messages);

if (lastMessage?.role != 'tool') {
return {
approvedToolApprovals: [],
deniedToolApprovals: [],
};
if (approvalTurn == null) {
return createEmptyToolApprovals();
}

const { latestAssistantContent, suffixMessages } = approvalTurn;

// gather tool calls and prepare lookup.
//
// These maps are keyed by client-supplied ids (`toolCallId`, `approvalId`)
Expand All @@ -49,47 +60,47 @@ export function collectToolApprovals<TOOLS extends ToolSet>({
string,
TypedToolCall<TOOLS>
> = Object.create(null);
for (const message of messages) {
if (message.role === 'assistant' && typeof message.content !== 'string') {
const content = message.content;
for (const part of content) {
if (part.type === 'tool-call') {
toolCallsByToolCallId[part.toolCallId] = part as TypedToolCall<TOOLS>;
}
}
for (const part of latestAssistantContent) {
if (part.type === 'tool-call') {
toolCallsByToolCallId[part.toolCallId] = part as TypedToolCall<TOOLS>;
}
}

// gather approval responses and prepare lookup
// gather approval requests from the latest assistant turn only. If a later
// assistant message exists, older approval responses belong to already
// continued history and must not trigger side effects again.
const toolApprovalRequestsByApprovalId: Record<string, ToolApprovalRequest> =
Object.create(null);
for (const message of messages) {
if (message.role === 'assistant' && typeof message.content !== 'string') {
const content = message.content;
for (const part of content) {
if (part.type === 'tool-approval-request') {
toolApprovalRequestsByApprovalId[part.approvalId] = part;
}
}
for (const part of latestAssistantContent) {
if (part.type === 'tool-approval-request') {
toolApprovalRequestsByApprovalId[part.approvalId] = part;
}
}

// gather tool results from the last tool message
const approvalResponses: ToolApprovalResponse[] = [];

// gather tool results from the unresolved suffix after the latest assistant
// turn, allowing user/system context to trail approval responses.
const toolResults: Record<string, TypedToolResult<TOOLS>> = Object.create(
null,
);
for (const part of lastMessage.content) {
if (part.type === 'tool-result') {
toolResults[part.toolCallId] = part as TypedToolResult<TOOLS>;
for (const message of suffixMessages) {
if (message.role !== 'tool') {
continue;
}

for (const part of message.content) {
if (part.type === 'tool-approval-response') {
approvalResponses.push(part);
} else if (part.type === 'tool-result') {
toolResults[part.toolCallId] = part as TypedToolResult<TOOLS>;
}
}
}

const approvedToolApprovals: Array<CollectedToolApprovals<TOOLS>> = [];
const deniedToolApprovals: Array<CollectedToolApprovals<TOOLS>> = [];

const approvalResponses = lastMessage.content.filter(
part => part.type === 'tool-approval-response',
);
for (const approvalResponse of approvalResponses) {
const approvalRequest =
toolApprovalRequestsByApprovalId[approvalResponse.approvalId];
Expand Down
Loading