stack/apps/gateway/src/agent/tools/web-tools.ts

import { Type } from '@sinclair/typebox';
import type { ToolDefinition } from '@mariozechner/pi-coding-agent';

const DEFAULT_TIMEOUT_MS = 15_000;
const MAX_RESPONSE_BYTES = 512 * 1024; // 512 KB

/**
 * Blocked URL patterns (private IP ranges, localhost, link-local).
 */
const BLOCKED_HOSTNAMES = [
  /^localhost$/i,
  /^127\./,
  /^10\./,
  /^172\.(1[6-9]|2\d|3[01])\./,
  /^192\.168\./,
  /^::1$/,
  /^fc[0-9a-f][0-9a-f]:/i,
  /^fe80:/i,
  /^0\.0\.0\.0$/,
  /^169\.254\./,
];

function isBlockedUrl(urlString: string): string | null {
  let parsed: URL;
  try {
    parsed = new URL(urlString);
  } catch {
    return `Invalid URL: ${urlString}`;
  }
  if (parsed.protocol !== 'http:' && parsed.protocol !== 'https:') {
    return `Unsupported protocol: ${parsed.protocol}. Only http and https are allowed.`;
  }
  const hostname = parsed.hostname;
  for (const pattern of BLOCKED_HOSTNAMES) {
    if (pattern.test(hostname)) {
      return `Blocked: requests to "${hostname}" are not allowed (private/local addresses).`;
    }
  }
  return null;
}

async function fetchWithLimit(
  url: string,
  options: RequestInit,
  timeoutMs: number,
): Promise<{ text: string; status: number; contentType: string }> {
  const controller = new AbortController();
  const timer = setTimeout(() => controller.abort(), timeoutMs);

  try {
    const response = await fetch(url, { ...options, signal: controller.signal });
    const contentType = response.headers.get('content-type') ?? '';

    // Stream response and enforce size limit
    const reader = response.body?.getReader();
    if (!reader) {
      return { text: '', status: response.status, contentType };
    }

    const chunks: Uint8Array[] = [];
    let totalBytes = 0;
    let truncated = false;

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      totalBytes += value.length;
      if (totalBytes > MAX_RESPONSE_BYTES) {
        const remaining = MAX_RESPONSE_BYTES - (totalBytes - value.length);
        chunks.push(value.subarray(0, remaining));
        truncated = true;
        reader.cancel();
        break;
      }
      chunks.push(value);
    }

    const combined = new Uint8Array(chunks.reduce((acc, c) => acc + c.length, 0));
    let offset = 0;
    for (const chunk of chunks) {
      combined.set(chunk, offset);
      offset += chunk.length;
    }

    let text = new TextDecoder().decode(combined);
    if (truncated) {
      text += '\n[response truncated at 512 KB limit]';
    }

    return { text, status: response.status, contentType };
  } finally {
    clearTimeout(timer);
  }
}

export function createWebTools(): ToolDefinition[] {
  const webGet: ToolDefinition = {
    name: 'web_get',
    label: 'HTTP GET',
    description:
      'Perform an HTTP GET request and return the response body. Private/local addresses are blocked.',
    parameters: Type.Object({
      url: Type.String({ description: 'URL to fetch (http/https only)' }),
      headers: Type.Optional(
        Type.Record(Type.String(), Type.String(), {
          description: 'Optional request headers as key-value pairs',
        }),
      ),
      timeout: Type.Optional(
        Type.Number({ description: 'Timeout in milliseconds (default 15000, max 30000)' }),
      ),
    }),
    async execute(_toolCallId, params) {
      const { url, headers, timeout } = params as {
        url: string;
        headers?: Record<string, string>;
        timeout?: number;
      };

      const blocked = isBlockedUrl(url);
      if (blocked) {
        return {
          content: [{ type: 'text' as const, text: `Error: ${blocked}` }],
          details: undefined,
        };
      }

      const timeoutMs = Math.min(timeout ?? DEFAULT_TIMEOUT_MS, 30_000);

      try {
        const result = await fetchWithLimit(
          url,
          { method: 'GET', headers: headers ?? {} },
          timeoutMs,
        );
        return {
          content: [
            {
              type: 'text' as const,
              text: `HTTP ${result.status} (${result.contentType})\n\n${result.text}`,
            },
          ],
          details: undefined,
        };
      } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        return {
          content: [{ type: 'text' as const, text: `Error fetching URL: ${msg}` }],
          details: undefined,
        };
      }
    },
  };

  const webPost: ToolDefinition = {
    name: 'web_post',
    label: 'HTTP POST',
    description:
      'Perform an HTTP POST request with a JSON or text body. Private/local addresses are blocked.',
    parameters: Type.Object({
      url: Type.String({ description: 'URL to POST to (http/https only)' }),
      body: Type.String({ description: 'Request body (JSON string or plain text)' }),
      contentType: Type.Optional(
        Type.String({ description: 'Content-Type header (default: application/json)' }),
      ),
      headers: Type.Optional(
        Type.Record(Type.String(), Type.String(), {
          description: 'Optional additional request headers',
        }),
      ),
      timeout: Type.Optional(
        Type.Number({ description: 'Timeout in milliseconds (default 15000, max 30000)' }),
      ),
    }),
    async execute(_toolCallId, params) {
      const { url, body, contentType, headers, timeout } = params as {
        url: string;
        body: string;
        contentType?: string;
        headers?: Record<string, string>;
        timeout?: number;
      };

      const blocked = isBlockedUrl(url);
      if (blocked) {
        return {
          content: [{ type: 'text' as const, text: `Error: ${blocked}` }],
          details: undefined,
        };
      }

      const timeoutMs = Math.min(timeout ?? DEFAULT_TIMEOUT_MS, 30_000);
      const ct = contentType ?? 'application/json';

      try {
        const result = await fetchWithLimit(
          url,
          {
            method: 'POST',
            headers: { 'Content-Type': ct, ...(headers ?? {}) },
            body,
          },
          timeoutMs,
        );
        return {
          content: [
            {
              type: 'text' as const,
              text: `HTTP ${result.status} (${result.contentType})\n\n${result.text}`,
            },
          ],
          details: undefined,
        };
      } catch (err) {
        const msg = err instanceof Error ? err.message : String(err);
        return {
          content: [{ type: 'text' as const, text: `Error posting to URL: ${msg}` }],
          details: undefined,
        };
      }
    },
  };

  return [webGet, webPost];
}