Ling0925
/
buster
mirror of https://github.com/buster-so/buster.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor chat cancellation logic to improve type handling for tool calls and results. Introduced type guards for ToolCallPart and ToolResultPart, updated related functions to use these new types, and streamlined the message processing flow.

This commit is contained in:
dal 2025-08-18 23:52:26 -06:00
parent f0796b35af
commit e97b4d2ae1
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
1 changed files with 31 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

44
apps/server/src/api/v2/chats/cancel-chat.ts
View File

@ -1,10 +1,4 @@
import { canUserAccessChatCached } from '@buster/access-controls';
import {
type ToolCallContent,
type ToolResultContent,
isToolCallContent,
isToolResultContent,
} from '@buster/ai/utils/database/types';
import type { User } from '@buster/database';
import { and, eq, isNotNull, updateMessageFields } from '@buster/database';
import { db, messages } from '@buster/database';
@ -13,7 +7,7 @@ import type {
ChatMessageResponseMessage,
} from '@buster/server-shared/chats';
import { runs } from '@trigger.dev/sdk';
import type { ModelMessage, ToolResultPart } from 'ai';
import type { ModelMessage, ToolCallPart, ToolResultPart } from 'ai';
import { errorResponse } from '../../../utils/response';
/**
@ -95,24 +89,48 @@ export async function cancelChatHandler(chatId: string, user: User): Promise<voi
// Wait for all cleanups to complete
await Promise.allSettled(cleanupPromises);
}
/**
* Type guard for ToolCallPart
*/
function isToolCallPart(content: unknown): content is ToolCallPart {
return (
content !== null &&
typeof content === 'object' &&
'type' in content &&
content.type === 'tool-call'
);
}
/**
* Type guard for ToolResultPart
*/
function isToolResultPart(content: unknown): content is ToolResultPart {
return (
content !== null &&
typeof content === 'object' &&
'type' in content &&
content.type === 'tool-result'
);
}
/**
* Find tool calls without corresponding tool results
*/
function findIncompleteToolCalls(messages: ModelMessage[]): ToolCallContent[] {
const toolCalls = new Map<string, ToolCallContent>();
function findIncompleteToolCalls(messages: ModelMessage[]): ToolCallPart[] {
const toolCalls = new Map<string, ToolCallPart>();
const toolResults = new Set<string>();
// First pass: collect all tool calls and tool results
for (const message of messages) {
if (message.role === 'assistant' && Array.isArray(message.content)) {
for (const content of message.content) {
if (isToolCallContent(content)) {
if (isToolCallPart(content)) {
toolCalls.set(content.toolCallId, content);
}
}
} else if (message.role === 'tool' && Array.isArray(message.content)) {
for (const content of message.content) {
if (isToolResultContent(content)) {
if (isToolResultPart(content)) {
toolResults.add(content.toolCallId);
}
}
@ -120,7 +138,7 @@ function findIncompleteToolCalls(messages: ModelMessage[]): ToolCallContent[] {
}
// Second pass: find tool calls without results
const incompleteToolCalls: ToolCallContent[] = [];
const incompleteToolCalls: ToolCallPart[] = [];
for (const [toolCallId, toolCall] of toolCalls) {
if (!toolResults.has(toolCallId)) {
incompleteToolCalls.push(toolCall);
@ -133,7 +151,7 @@ function findIncompleteToolCalls(messages: ModelMessage[]): ToolCallContent[] {
/**
* Create tool result messages for incomplete tool calls
*/
function createCancellationToolResults(incompleteToolCalls: ToolCallContent[]): ModelMessage[] {
function createCancellationToolResults(incompleteToolCalls: ToolCallPart[]): ModelMessage[] {
if (incompleteToolCalls.length === 0) {
return [];
}