software-mansion__argent/packages/argent-tools-client/test/stream.test.ts

import { describe, it, expect, vi, afterEach } from "vitest";
import { createServer, type Server, type IncomingMessage, type ServerResponse } from "node:http";
import type { AddressInfo } from "node:net";
import { createToolsClient, errorBodyMessage, ToolInvocationError } from "../src/tools-client.js";

let server: Server | undefined;

/**
 * Stand up a stub tool-server: GET /tools advertises one tool ("streamy",
 * no fileInputs so the file boundary stays out of the way); POST behavior is
 * supplied per test. ARGENT_TOOLS_URL routes the client at it.
 */
async function startServer(
  onInvoke: (req: IncomingMessage, res: ServerResponse) => void
): Promise<void> {
  server = createServer((req, res) => {
    if (req.method === "GET" && req.url === "/tools") {
      res.setHeader("Content-Type", "application/json");
      res.end(JSON.stringify({ tools: [{ name: "streamy", description: "", inputSchema: {} }] }));
      return;
    }
    if (req.method === "POST" && req.url === "/tools/streamy") {
      onInvoke(req, res);
      return;
    }
    res.statusCode = 404;
    res.end(JSON.stringify({ error: "not found" }));
  });
  await new Promise<void>((resolve) => server!.listen(0, "127.0.0.1", resolve));
  const { port } = server!.address() as AddressInfo;
  vi.stubEnv("ARGENT_TOOLS_URL", `http://127.0.0.1:${port}`);
}

afterEach(async () => {
  vi.unstubAllEnvs();
  if (server) {
    await new Promise<void>((resolve) => server!.close(() => resolve()));
    server = undefined;
  }
});

describe("callTool progress streaming", () => {
  it("fires onProgress per NDJSON line and resolves with the terminal result", async () => {
    await startServer((_req, res) => {
      res.writeHead(200, { "Content-Type": "application/x-ndjson" });
      res.write(`${JSON.stringify({ event: "progress", data: { index: 0 } })}\n`);
      res.write(`${JSON.stringify({ event: "progress", data: { index: 1 } })}\n`);
      res.end(`${JSON.stringify({ event: "result", data: { ok: true }, note: "hi" })}\n`);
    });

    const events: unknown[] = [];
    const { callTool } = createToolsClient();
    const result = await callTool("streamy", {}, { onProgress: (e) => events.push(e) });

    expect(events).toEqual([{ index: 0 }, { index: 1 }]);
    expect(result.data).toEqual({ ok: true });
    expect(result.note).toBe("hi");
  });

  it("sends the Accept header only when a progress consumer is attached", async () => {
    const accepts: Array<string | undefined> = [];
    await startServer((req, res) => {
      accepts.push(req.headers.accept);
      res.setHeader("Content-Type", "application/json");
      res.end(JSON.stringify({ data: { ok: true } }));
    });

    const { callTool } = createToolsClient();
    await callTool("streamy", {});
    await callTool("streamy", {}, { onProgress: () => {} });

    expect(accepts[0] ?? "").not.toContain("application/x-ndjson");
    expect(accepts[1]).toContain("application/x-ndjson");
  });

  it("falls back to the buffered path against a pre-streaming server", async () => {
    // An old server ignores the Accept header and replies plain JSON.
    await startServer((_req, res) => {
      res.setHeader("Content-Type", "application/json");
      res.end(JSON.stringify({ data: { ok: true } }));
    });

    const events: unknown[] = [];
    const { callTool } = createToolsClient();
    const result = await callTool("streamy", {}, { onProgress: (e) => events.push(e) });

    expect(events).toEqual([]);
    expect(result.data).toEqual({ ok: true });
  });

  it("rejects with the in-band terminal error", async () => {
    await startServer((_req, res) => {
      res.writeHead(200, { "Content-Type": "application/x-ndjson" });
      res.write(`${JSON.stringify({ event: "progress", data: { index: 0 } })}\n`);
      res.end(`${JSON.stringify({ event: "error", error: "kaput" })}\n`);
    });

    const events: unknown[] = [];
    const { callTool } = createToolsClient();
    await expect(callTool("streamy", {}, { onProgress: (e) => events.push(e) })).rejects.toThrow(
      "kaput"
    );
    expect(events).toEqual([{ index: 0 }]);
  });

  it("surfaces the terminal error line's failure signal as a ToolInvocationError", async () => {
    await startServer((_req, res) => {
      res.writeHead(200, { "Content-Type": "application/x-ndjson" });
      res.end(
        `${JSON.stringify({
          event: "error",
          error: "flow file is not valid YAML",
          error_code: "FLOW_FILE_INVALID",
          error_kind: "validation",
        })}\n`
      );
    });

    const { callTool } = createToolsClient();
    const err = await callTool("streamy", {}, { onProgress: () => {} }).catch((e: unknown) => e);
    expect(err).toBeInstanceOf(ToolInvocationError);
    expect((err as ToolInvocationError).message).toBe("flow file is not valid YAML");
    expect((err as ToolInvocationError).errorCode).toBe("FLOW_FILE_INVALID");
    expect((err as ToolInvocationError).errorKind).toBe("validation");
  });

  it("surfaces a buffered JSON error response's failure signal the same way", async () => {
    await startServer((_req, res) => {
      res.statusCode = 500;
      res.setHeader("Content-Type", "application/json");
      res.end(
        JSON.stringify({ error: "boom", error_code: "FLOW_FILE_INVALID", error_kind: "validation" })
      );
    });

    const { callTool } = createToolsClient();
    const err = await callTool("streamy", {}).catch((e: unknown) => e);
    expect(err).toBeInstanceOf(ToolInvocationError);
    expect((err as ToolInvocationError).message).toBe("boom");
    expect((err as ToolInvocationError).errorKind).toBe("validation");
  });

  it("carries a 400's schema issue list beside its prose message", async () => {
    const issues = [{ code: "too_big", path: ["x"], message: "Too big: expected <=1" }];
    await startServer((_req, res) => {
      res.statusCode = 400;
      res.setHeader("Content-Type", "application/json");
      res.end(JSON.stringify({ error: "`x`: Too big: expected <=1. You sent: `x`.", issues }));
    });

    const { callTool } = createToolsClient();
    const err = await callTool("streamy", {}).catch((e: unknown) => e);
    expect((err as ToolInvocationError).issues).toEqual(issues);
  });

  it("reads the prose from `message`, leaving `error` to an older CLI", async () => {
    const issues = [{ code: "too_big", path: ["x"], message: "Too big: expected <=1" }];
    await startServer((_req, res) => {
      res.statusCode = 400;
      res.setHeader("Content-Type", "application/json");
      res.end(
        JSON.stringify({
          error: JSON.stringify(issues),
          message: "`x`: Too big: expected <=1. You sent: `x`.",
          issues,
        })
      );
    });

    const { callTool } = createToolsClient();
    const err = await callTool("streamy", {}).catch((e: unknown) => e);
    expect((err as ToolInvocationError).message).toBe("`x`: Too big: expected <=1. You sent: `x`.");
    expect((err as ToolInvocationError).issues).toEqual(issues);
  });

  it("keeps `error` as the message when the body carries no `issues`", async () => {
    await startServer((_req, res) => {
      res.statusCode = 500;
      res.setHeader("Content-Type", "application/json");
      res.end(JSON.stringify({ error: "boom", message: "something else entirely" }));
    });

    const { callTool } = createToolsClient();
    const err = await callTool("streamy", {}).catch((e: unknown) => e);
    expect((err as ToolInvocationError).message).toBe("boom");
  });

  it("leaves `issues` undefined when the body carries none, or a non-list", async () => {
    await startServer((_req, res) => {
      res.statusCode = 400;
      res.setHeader("Content-Type", "application/json");
      res.end(JSON.stringify({ error: "boom", issues: "not a list" }));
    });

    const { callTool } = createToolsClient();
    const err = await callTool("streamy", {}).catch((e: unknown) => e);
    expect((err as ToolInvocationError).issues).toBeUndefined();
  });

  it("rejects when the stream ends without a terminal line (connection lost)", async () => {
    await startServer((_req, res) => {
      res.writeHead(200, { "Content-Type": "application/x-ndjson" });
      res.write(`${JSON.stringify({ event: "progress", data: { index: 0 } })}\n`);
      res.end();
    });

    const { callTool } = createToolsClient();
    await expect(callTool("streamy", {}, { onProgress: () => {} })).rejects.toThrow(
      /without a result/
    );
  });
});

describe("errorBodyMessage", () => {
  it("takes the prose only when the body is the validation pair", () => {
    const issues = [{ code: "too_big", path: ["x"], message: "Too big" }];
    expect(errorBodyMessage({ error: "[...]", message: "prose", issues })).toBe("prose");
    expect(errorBodyMessage({ error: "boom", message: "prose" })).toBe("boom");
    expect(errorBodyMessage({ error: "boom" })).toBe("boom");
    expect(errorBodyMessage({ message: "prose" })).toBe("prose");
    expect(errorBodyMessage({})).toBeUndefined();
  });
});