126 lines
3.1 KiB
TypeScript
126 lines
3.1 KiB
TypeScript
import type { ToolCall } from "@earendil-works/pi-ai";
|
|
import type {
|
|
ExtensionAPI,
|
|
ExtensionContext,
|
|
SessionEntry,
|
|
} from "@earendil-works/pi-coding-agent";
|
|
import { findPayloadForSourceEntry } from "./ask-payload-store.ts";
|
|
import { validateParams } from "./ask-tool-helpers.ts";
|
|
import type { AskParams } from "./types.ts";
|
|
|
|
export const ASK_PENDING_DISMISSED_ENTRY_TYPE = "ask:pending-dismissed";
|
|
const ASK_TOOL_NAME = "ask_user";
|
|
|
|
export interface PendingAskToolCall {
|
|
params: AskParams;
|
|
toolCallId: string;
|
|
}
|
|
|
|
type AskToolCall = ToolCall;
|
|
|
|
export function appendPendingAskDismissal(
|
|
pi: Pick<ExtensionAPI, "appendEntry">,
|
|
toolCallId: string
|
|
): void {
|
|
pi.appendEntry(ASK_PENDING_DISMISSED_ENTRY_TYPE, { toolCallId });
|
|
}
|
|
|
|
export function findPendingAskToolCall(
|
|
ctx: Pick<ExtensionContext, "sessionManager">
|
|
): PendingAskToolCall | undefined {
|
|
const branch = ctx.sessionManager.getBranch();
|
|
const resolvedToolCallIds = collectResolvedToolCallIds(branch);
|
|
|
|
for (let entryIndex = branch.length - 1; entryIndex >= 0; entryIndex--) {
|
|
const toolCall = findUnresolvedAskToolCall(
|
|
branch[entryIndex],
|
|
resolvedToolCallIds
|
|
);
|
|
if (!toolCall) {
|
|
continue;
|
|
}
|
|
|
|
const params = resolvePendingAskParams(ctx, toolCall);
|
|
if (params) {
|
|
return { params, toolCallId: toolCall.id };
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
function collectResolvedToolCallIds(
|
|
branch: readonly SessionEntry[]
|
|
): Set<string> {
|
|
const resolved = new Set<string>();
|
|
for (const entry of branch) {
|
|
const dismissedToolCallId = getDismissedToolCallId(entry);
|
|
if (dismissedToolCallId) {
|
|
resolved.add(dismissedToolCallId);
|
|
continue;
|
|
}
|
|
if (entry.type === "message" && entry.message.role === "toolResult") {
|
|
resolved.add(entry.message.toolCallId);
|
|
}
|
|
}
|
|
return resolved;
|
|
}
|
|
|
|
function findUnresolvedAskToolCall(
|
|
entry: SessionEntry,
|
|
resolvedToolCallIds: ReadonlySet<string>
|
|
): AskToolCall | undefined {
|
|
if (
|
|
entry.type !== "message" ||
|
|
entry.message.role !== "assistant" ||
|
|
entry.message.stopReason !== "toolUse"
|
|
) {
|
|
return;
|
|
}
|
|
|
|
for (
|
|
let partIndex = entry.message.content.length - 1;
|
|
partIndex >= 0;
|
|
partIndex--
|
|
) {
|
|
const part = entry.message.content[partIndex];
|
|
if (
|
|
part.type === "toolCall" &&
|
|
part.name === ASK_TOOL_NAME &&
|
|
!resolvedToolCallIds.has(part.id)
|
|
) {
|
|
return part;
|
|
}
|
|
}
|
|
return;
|
|
}
|
|
|
|
function resolvePendingAskParams(
|
|
ctx: Pick<ExtensionContext, "sessionManager">,
|
|
toolCall: AskToolCall
|
|
): AskParams | undefined {
|
|
const persistedPayload = findPayloadForSourceEntry(ctx, toolCall.id, "tool");
|
|
if (persistedPayload) {
|
|
return persistedPayload.params;
|
|
}
|
|
|
|
const argumentsFallback = toolCall.arguments as AskParams | undefined;
|
|
if (argumentsFallback && validateParams(argumentsFallback).ok) {
|
|
return argumentsFallback;
|
|
}
|
|
return;
|
|
}
|
|
|
|
function getDismissedToolCallId(entry: SessionEntry): string | undefined {
|
|
if (
|
|
entry.type !== "custom" ||
|
|
entry.customType !== ASK_PENDING_DISMISSED_ENTRY_TYPE ||
|
|
!entry.data ||
|
|
typeof entry.data !== "object"
|
|
) {
|
|
return;
|
|
}
|
|
|
|
const toolCallId = (entry.data as { toolCallId?: unknown }).toolCallId;
|
|
return typeof toolCallId === "string" ? toolCallId : undefined;
|
|
}
|