diff --git a/packages/mcp/src/tools/devices/list-devices/list-devices.test.ts b/packages/mcp/src/tools/devices/list-devices/list-devices.test.ts new file mode 100644 index 00000000000..562b323d9d1 --- /dev/null +++ b/packages/mcp/src/tools/devices/list-devices/list-devices.test.ts @@ -0,0 +1,182 @@ +import { beforeEach, describe, expect, it, mock } from "bun:test"; +import { z } from "zod"; + +const getMcpContextMock = mock(() => ({ organizationId: "org-1" })); + +let fetchedDevices = [ + { + deviceId: "device-online", + deviceName: "Ada's MacBook", + deviceType: "desktop", + lastSeenAt: new Date(Date.now() - 30_000), + ownerId: "user-1", + ownerName: "Ada", + ownerEmail: "ada@example.com", + }, + { + deviceId: "device-offline", + deviceName: "Grace's iPhone", + deviceType: "mobile", + lastSeenAt: new Date(Date.now() - 120_000), + ownerId: "user-2", + ownerName: "Grace", + ownerEmail: "grace@example.com", + }, +]; + +const selectMock = mock(() => ({ + from: () => ({ + innerJoin: () => ({ + where: () => ({ + orderBy: async () => fetchedDevices, + }), + }), + }), +})); + +mock.module("@superset/db/client", () => ({ + db: { + select: selectMock, + }, +})); + +mock.module("../../utils", () => ({ + getMcpContext: getMcpContextMock, +})); + +const { register } = await import("./index"); + +type RegisteredToolHandler = ( + args: Record, + extra: unknown, +) => Promise<{ + content?: Array<{ text?: string }>; + isError?: boolean; + structuredContent?: { + devices: Array<{ + deviceId: string; + deviceName: string | null; + deviceType: string; + lastSeenAt: string; + ownerId: string; + ownerName: string | null; + ownerEmail: string; + isOnline: boolean; + }>; + }; +}>; + +type RegisteredToolConfig = { + inputSchema: Record; + outputSchema: Record; +}; + +function createTool() { + let config: RegisteredToolConfig | null = null; + let handler: RegisteredToolHandler | null = null; + + register({ + registerTool: ( + name: string, + nextConfig: RegisteredToolConfig, + nextHandler: RegisteredToolHandler, + ) => { + if (name === "list_devices") { + config = nextConfig; + handler = nextHandler; + } + }, + } as never); + + if (!config || !handler) { + throw new Error("list_devices was not registered"); + } + + return { + config: config as RegisteredToolConfig, + handler: handler as RegisteredToolHandler, + }; +} + +describe("list_devices MCP tool", () => { + beforeEach(() => { + fetchedDevices = [ + { + deviceId: "device-online", + deviceName: "Ada's MacBook", + deviceType: "desktop", + lastSeenAt: new Date(Date.now() - 30_000), + ownerId: "user-1", + ownerName: "Ada", + ownerEmail: "ada@example.com", + }, + { + deviceId: "device-offline", + deviceName: "Grace's iPhone", + deviceType: "mobile", + lastSeenAt: new Date(Date.now() - 120_000), + ownerId: "user-2", + ownerName: "Grace", + ownerEmail: "grace@example.com", + }, + ]; + getMcpContextMock.mockClear(); + selectMock.mockClear(); + }); + + it("registers input and output schemas that validate includeOffline and isOnline", async () => { + const { config, handler } = createTool(); + const inputSchema = z.object(config.inputSchema); + const outputSchema = z.object(config.outputSchema); + + expect(inputSchema.parse({})).toEqual({ includeOffline: false }); + expect(inputSchema.parse({ includeOffline: true })).toEqual({ + includeOffline: true, + }); + + const result = await handler({ includeOffline: true }, {}); + + expect(() => outputSchema.parse(result.structuredContent)).not.toThrow(); + }); + + it("returns only online devices by default", async () => { + const { handler } = createTool(); + + const result = await handler({}, {}); + + expect(getMcpContextMock).toHaveBeenCalledTimes(1); + expect(selectMock).toHaveBeenCalledTimes(1); + expect(result.structuredContent?.devices).toEqual([ + { + deviceId: "device-online", + deviceName: "Ada's MacBook", + deviceType: "desktop", + lastSeenAt: expect.any(String), + ownerId: "user-1", + ownerName: "Ada", + ownerEmail: "ada@example.com", + isOnline: true, + }, + ]); + }); + + it("includes offline devices when requested and marks them offline", async () => { + const { handler } = createTool(); + + const result = await handler({ includeOffline: true }, {}); + + expect(result.structuredContent?.devices).toHaveLength(2); + expect(result.structuredContent?.devices).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + deviceId: "device-online", + isOnline: true, + }), + expect.objectContaining({ + deviceId: "device-offline", + isOnline: false, + }), + ]), + ); + }); +}); diff --git a/packages/mcp/src/tools/devices/list-devices/list-devices.ts b/packages/mcp/src/tools/devices/list-devices/list-devices.ts index e3153ce24cc..dc06856f4c9 100644 --- a/packages/mcp/src/tools/devices/list-devices/list-devices.ts +++ b/packages/mcp/src/tools/devices/list-devices/list-devices.ts @@ -1,32 +1,43 @@ import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; import { db } from "@superset/db/client"; -import { devicePresence, users } from "@superset/db/schema"; +import { devicePresence, deviceTypeValues, users } from "@superset/db/schema"; import { desc, eq } from "drizzle-orm"; import { z } from "zod"; import { getMcpContext } from "../../utils"; +const DEVICE_ONLINE_WINDOW_MS = 60_000; + export function register(server: McpServer) { server.registerTool( "list_devices", { - description: "List registered devices in the organization", - inputSchema: {}, + description: + "List devices in the organization. By default, only devices seen within the last 60 seconds are returned.", + inputSchema: { + includeOffline: z + .boolean() + .default(false) + .describe("Include devices that have not checked in recently"), + }, outputSchema: { devices: z.array( z.object({ deviceId: z.string(), deviceName: z.string().nullable(), - deviceType: z.string(), - lastSeenAt: z.string(), + deviceType: z.enum(deviceTypeValues), + lastSeenAt: z.string().datetime(), ownerId: z.string(), ownerName: z.string().nullable(), ownerEmail: z.string(), + isOnline: z.boolean(), }), ), }, }, - async (_args, extra) => { + async (args, extra) => { const ctx = getMcpContext(extra); + const includeOffline = args.includeOffline === true; + const onlineCutoff = Date.now() - DEVICE_ONLINE_WINDOW_MS; const devices = await db .select({ @@ -43,10 +54,17 @@ export function register(server: McpServer) { .where(eq(devicePresence.organizationId, ctx.organizationId)) .orderBy(desc(devicePresence.lastSeenAt)); - const result = devices.map((d) => ({ - ...d, - lastSeenAt: d.lastSeenAt.toISOString(), - })); + const result = devices + .map((d) => { + const isOnline = d.lastSeenAt.getTime() >= onlineCutoff; + + return { + ...d, + lastSeenAt: d.lastSeenAt.toISOString(), + isOnline, + }; + }) + .filter((device) => includeOffline || device.isOnline); return { structuredContent: { devices: result },