@worktango/ai-assistant 0.0.22 → 0.0.24

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@worktango/ai-assistant",
-  "version": "0.0.22",
+  "version": "0.0.24",
   "types": "./types.d.ts",
   "files": [
     "dist",

package/src/backend/express.test.ts

@@ -1,41 +1,208 @@
+import http from "http";
+
 import { Express } from "express";
 
 import { setupAiAssistantRoutes } from "./express";
 
+jest.mock("http", () => ({
+  request: jest.fn().mockImplementation((options, callback) => {
+    const mockProxyRes = {
+      statusCode: 200,
+      headers: { "content-type": "application/json" },
+      pipe: jest.fn(),
+    };
+    callback(mockProxyRes);
+    return {
+      on: jest.fn(),
+      end: jest.fn(),
+    };
+  }),
+}));
+
 describe("setupAiAssistantRoutes", () => {
   let mockApp: jest.Mocked<Express>;
   let mockPost: jest.Mock;
-  let mockSend: jest.Mock;
+  let mockGetBearerToken: jest.Mock;
+  let mockReq: any;
+  let mockRes: any;
+  let handlerPromise: Promise<void>;
+
+  const mockLogger = {
+    log: jest.fn(),
+    error: jest.fn(),
+    debug: jest.fn(),
+  };
 
   beforeEach(() => {
-    mockSend = jest.fn();
+    mockReq = {
+      url: "http://localhost/ai-assistant",
+      method: "POST",
+      headers: { "content-type": "application/json" },
+      readable: false,
+    };
+    mockRes = {
+      statusCode: 200,
+      setHeader: jest.fn(),
+      status: jest.fn().mockReturnThis(),
+      send: jest.fn(),
+    };
     mockPost = jest.fn().mockImplementation((path, handler) => {
-      // Store the handler so we can call it in our tests
-      handler({}, { send: mockSend });
+      handlerPromise = handler(mockReq, mockRes);
     });
+    mockGetBearerToken = jest.fn().mockResolvedValue("mock-token");
 
     mockApp = {
       post: mockPost,
+      all: mockPost,
     } as unknown as jest.Mocked<Express>;
+
+    (http.request as jest.Mock).mockClear();
   });
 
-  it("sets up the default route correctly", () => {
-    setupAiAssistantRoutes(mockApp);
+  it("sets up the default route correctly", async () => {
+    setupAiAssistantRoutes({
+      app: mockApp,
+      getBearerToken: mockGetBearerToken,
+      targetHost: "localhost",
+      logger: mockLogger,
+    });
 
     expect(mockPost).toHaveBeenCalledWith(
       "/ai-assistant",
       expect.any(Function)
     );
-    expect(mockSend).toHaveBeenCalledWith("Hello world!");
+
+    await handlerPromise;
+
+    expect(http.request).toHaveBeenCalledWith(
+      expect.objectContaining({
+        hostname: "localhost",
+        path: "/ai-assistant",
+        headers: expect.objectContaining({
+          Authorization: "Bearer mock-token",
+        }),
+      }),
+      expect.any(Function)
+    );
   });
 
-  it("sets up a custom route when provided", () => {
-    setupAiAssistantRoutes(mockApp, { route: "/custom-route" });
+  it("sets up a custom route when provided", async () => {
+    mockReq.url = "http://localhost/custom-route";
+
+    setupAiAssistantRoutes({
+      app: mockApp,
+      getBearerToken: mockGetBearerToken,
+      targetHost: "localhost",
+      route: "/custom-route",
+      logger: mockLogger,
+    });
 
     expect(mockPost).toHaveBeenCalledWith(
       "/custom-route",
       expect.any(Function)
     );
-    expect(mockSend).toHaveBeenCalledWith("Hello world!");
+
+    await handlerPromise;
+
+    expect(http.request).toHaveBeenCalledWith(
+      expect.objectContaining({
+        hostname: "localhost",
+        path: "/custom-route",
+        headers: expect.objectContaining({
+          Authorization: "Bearer mock-token",
+        }),
+      }),
+      expect.any(Function)
+    );
+  });
+
+  it("handles errors when getting the bearer token", async () => {
+    mockGetBearerToken.mockRejectedValueOnce(
+      new Error("Error getting bearer token")
+    );
+
+    setupAiAssistantRoutes({
+      app: mockApp,
+      getBearerToken: mockGetBearerToken,
+      targetHost: "localhost",
+      logger: mockLogger,
+    });
+
+    await handlerPromise;
+
+    expect(mockRes.status).toHaveBeenCalledWith(500);
+    expect(mockRes.send).toHaveBeenCalledWith("Error getting bearer token");
+  });
+
+  it("handles proxy errors correctly", async () => {
+    const mockError = new Error("Proxy error");
+    (http.request as jest.Mock).mockImplementationOnce(() => ({
+      on: jest.fn().mockImplementation((event, handler) => {
+        if (event === "error") {
+          handler(mockError);
+        }
+      }),
+      end: jest.fn(),
+    }));
+
+    setupAiAssistantRoutes({
+      app: mockApp,
+      getBearerToken: mockGetBearerToken,
+      targetHost: "localhost",
+      logger: mockLogger,
+    });
+
+    await handlerPromise;
+
+    expect(mockRes.status).toHaveBeenCalledWith(500);
+    expect(mockRes.send).toHaveBeenCalledWith("Proxy Error: Proxy error");
+  });
+
+  it("handles readable request bodies", async () => {
+    mockReq.readable = true;
+    mockReq.pipe = jest.fn();
+
+    setupAiAssistantRoutes({
+      app: mockApp,
+      getBearerToken: mockGetBearerToken,
+      targetHost: "localhost",
+      logger: mockLogger,
+    });
+
+    await handlerPromise;
+
+    expect(mockReq.pipe).toHaveBeenCalled();
+  });
+
+  it("handles response headers and status code correctly", async () => {
+    const mockProxyRes = {
+      statusCode: undefined,
+      headers: {
+        "content-type": undefined,
+        "x-custom-header": "value",
+      },
+      pipe: jest.fn(),
+    };
+
+    (http.request as jest.Mock).mockImplementationOnce((options, callback) => {
+      callback(mockProxyRes);
+      return {
+        on: jest.fn(),
+        end: jest.fn(),
+      };
+    });
+
+    setupAiAssistantRoutes({
+      app: mockApp,
+      getBearerToken: mockGetBearerToken,
+      targetHost: "localhost",
+      logger: mockLogger,
+    });
+
+    await handlerPromise;
+
+    expect(mockRes.statusCode).toBe(500);
+    expect(mockRes.setHeader).toHaveBeenCalledWith("x-custom-header", "value");
+    expect(mockRes.setHeader).toHaveBeenCalledWith("content-type", "");
   });
 });

package/src/backend/express.ts

@@ -1,17 +1,98 @@
-import { Express } from "express";
+import http from "http";
+
+import { Express, RequestHandler } from "express";
+
+import { coalesce } from "@kazoohr/helpers";
 
 /**
  * Sets up the AI assistant routes on the given Express app.
  */
-export function setupAiAssistantRoutes(
-  app: Express,
-  options?: {
-    route?: string;
-  }
-) {
-  const route = options?.route ?? "/ai-assistant";
-  app.post(route, (req, res) => {
-    // TODO: Generate the server-to-server request token to the KW AI assistant and proxy the response back to the client.
-    res.send("Hello world!");
-  });
+export function setupAiAssistantRoutes(args: {
+  app: Express;
+  /**
+   * A function that returns the bearer token for the AI assistant to use for
+   * back-channel request proxying. This should be a JWT token that can be
+   * decoded on the `kazoo-web` application to authenticate the request.
+   */
+  getBearerToken: () => Promise<string>;
+  /**
+   * The host of the application to proxy requests to.
+   */
+  targetHost: string;
+  /**
+   * The port of the application to proxy requests to. Defaults to 80.
+   */
+  targetPort?: number;
+  /**
+   * The route path to the local service's AI assistant. Defaults to
+   * "/ai-assistant".
+   */
+  route?: string;
+  /**
+   * The logger to use for logging.
+   */
+  logger?: Pick<typeof console, "log" | "error" | "debug">;
+  /**
+   * Middleware functions to run before the proxy request.
+   */
+  preflightMiddlewares?: RequestHandler[];
+}) {
+  const logger = coalesce(args.logger, console);
+  const routePath = args.route ?? "/ai-assistant";
+
+  args.app.all(
+    routePath,
+    ...coalesce(args.preflightMiddlewares, []),
+    async (req, res) => {
+      const bearerToken = await args.getBearerToken().catch((error) => {
+        logger.error("Error getting bearer token:", error);
+        res.status(500).send("Error getting bearer token");
+        return null;
+      });
+
+      if (!bearerToken) {
+        return;
+      }
+
+      const parsedUrl = new URL(req.url);
+
+      const options = {
+        hostname: args.targetHost,
+        port: args.targetPort,
+        path: parsedUrl.pathname,
+        method: req.method,
+        headers: { ...req.headers },
+      };
+
+      options.headers["Authorization"] = `Bearer ${bearerToken}`;
+      delete options.headers.host;
+
+      logger.debug("Proxying request:", {
+        method: req.method,
+        headers: options.headers,
+        proxyOptions: options,
+      });
+
+      const proxyReq = http.request(options, (proxyRes) => {
+        res.statusCode = proxyRes.statusCode || 500;
+
+        Object.keys(proxyRes.headers).forEach((key) => {
+          res.setHeader(key, proxyRes.headers[key] || "");
+        });
+
+        proxyRes.pipe(res);
+      });
+
+      proxyReq.on("error", (error) => {
+        logger.error("Proxy request error:", error);
+        res.status(500).send("Proxy Error: " + error.message);
+      });
+
+      if (req.readable) {
+        req.pipe(proxyReq);
+      } else {
+        proxyReq.end();
+      }
+    }
+  );
 }

package/types.d.ts

@@ -32,3 +32,38 @@ declare module "@worktango/ai-assistant/react" {
     tools?: ToolDeclaration[];
   }): JSX.Element;
 }
+
+declare module "@worktango/ai-assistant/express" {
+  export function setupAiAssistantRoutes(args: {
+    app: {
+      all: (path: string, ...args: any[]) => void;
+    };
+    /**
+     * A function that returns the bearer token for the AI assistant to use for
+     * back-channel request proxying. This should be a JWT token that can be
+     * decoded on the `kazoo-web` application to authenticate the request.
+     */
+    getBearerToken: () => Promise<string>;
+    /**
+     * The host of the application to proxy requests to.
+     */
+    targetHost: string;
+    /**
+     * The port of the application to proxy requests to. Defaults to 80.
+     */
+    targetPort?: number;
+    /**
+     * The route path to the local service's AI assistant. Defaults to
+     * "/ai-assistant".
+     */
+    route?: string;
+    /**
+     * The logger to use for logging.
+     */
+    logger?: Pick<typeof console, "log" | "error" | "debug">;
+    /**
+     * Middleware functions to run before the proxy request.
+     */
+    preflightMiddlewares?: RequestHandler[];
+  }): void;
+}