pi tool call filter

pi-tool-call-filter.md

/**
 * Tool Call Filter Extension
 *
 * Filters out invalid tool calls before they're sent back to the provider.
 * This addresses issues where some models construct malformed tool call blocks.
 *
 * Install: Place in ~/.pi/extensions/tool-call-filter.ts or project .pi/extensions/
 */

import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { AssistantMessage, ToolCall } from "@mariozechner/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";

// ============================================================================
// Configuration (customize these filters)
// ============================================================================

/**
 * Filter function: returns true if the tool call should be kept.
 * Return false to filter it out.
 */
async function shouldKeepToolCall(
    toolName: string,
    arguments_: Record<string, unknown>,
    ctx: ExtensionContext,
): Promise<boolean> {
    // Filter out tool calls with empty/whitespace-only names
    if (!toolName || toolName.trim().length === 0) {
        ctx.ui.notify(`Filtered tool call with empty name`, "info");
        return false;
    }

    // Filter out tool calls with empty or null arguments
    if (arguments_ === null || typeof arguments_ !== "object") {
        ctx.ui.notify(`Filtered tool call with invalid arguments type`, "warning");
        return false;
    }

    // Filter out tool calls with empty arguments object (no keys)
    const keys = Object.keys(arguments_);
    if (keys.length === 0) {
        ctx.ui.notify(`Filtered tool call with empty arguments: ${toolName}`, "info");
        return false;
    }

    // Filter out tool calls where all argument values are empty/undefined/null
    const hasNonEmptyValue = keys.some((key) => {
        const value = arguments_[key];
        // Consider undefined, null, empty string, empty array, empty object as empty
        if (value === undefined || value === null) return false;
        if (typeof value === "string" && value.trim().length === 0) return false;
        if (Array.isArray(value) && value.length === 0) return false;
        if (typeof value === "object" && !Array.isArray(value) && Object.keys(value).length === 0) return false;
        return true;
    });

    if (!hasNonEmptyValue) {
        ctx.ui.notify(`Filtered tool call with all empty argument values: ${toolName}`, "info");
        return false;
    }

    return true; // Keep the tool call
}

// Simple in-memory tracker for duplicate detection (per turn)
let lastToolCalls: Array<{ name: string; args: string }> = [];

function resetDuplicateTracker() {
    lastToolCalls = [];
}

function isDuplicate(toolName: string, args: Record<string, unknown>): boolean {
    const argsString = JSON.stringify(args);
    return lastToolCalls.some((call) => call.name === toolName && call.args === argsString);
}

function recordToolCall(toolName: string, args: Record<string, unknown>) {
    lastToolCalls.push({ name: toolName, args: JSON.stringify(args) });
}

// ============================================================================
// Extension Factory
// ============================================================================

export default function (api: ExtensionAPI): Promise<void> | void {
    // Register the context event handler
    api.on("context", async (event, ctx) => {
        // Reset duplicate tracker at the start of each context processing
        resetDuplicateTracker();

        const messages = event.messages;
        const filteredMessages: AgentMessage[] = [];

        for (const message of messages) {
            // Only process assistant messages
            if (message.role !== "assistant") {
                filteredMessages.push(message);
                continue;
            }

            const assistantMsg = message as AssistantMessage;
            const filteredContent: AssistantMessage["content"] = [];

            for (const content of assistantMsg.content) {
                if (content.type === "toolCall") {
                    const toolCall = content as ToolCall;

                    // Validate the tool call
                    const isValid = await shouldKeepToolCall(
                        toolCall.name,
                        toolCall.arguments,
                        ctx,
                    );

                    if (isValid) {
                        // Check for duplicates
                        if (isDuplicate(toolCall.name, toolCall.arguments)) {
                            ctx.ui.notify(`Filtered duplicate tool call: ${toolCall.name}`, "info");
                        } else {
                            filteredContent.push(content);
                            recordToolCall(toolCall.name, toolCall.arguments);
                        }
                    }
                    // If filtered, skip adding to filteredContent
                } else {
                    // Non-tool-call content (text, thinking, etc.) - keep as-is
                    filteredContent.push(content);
                }
            }

            // Only add the message if it still has content after filtering
            if (filteredContent.length > 0) {
                filteredMessages.push({
                    ...assistantMsg,
                    content: filteredContent,
                });
            } else {
                ctx.ui.notify(`Removed assistant message with no valid content after filtering`, "info");
            }
        }

        return { messages: filteredMessages };
    });

    // Optional: Register a command to toggle filtering or adjust settings
    api.registerCommand("tool-filter-toggle", {
        description: "Toggle tool call filtering (for debugging)",
        handler: async (args, ctx) => {
            // This is a placeholder - you could add runtime configuration
            ctx.ui.notify("Tool call filter is active", "info");
        },
    });
}