[codex] Fix list_devices MCP schema

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<string, unknown>,
+	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<string, z.ZodTypeAny>;
+	outputSchema: Record<string, z.ZodTypeAny>;
+};
+
+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,
+				}),
+			]),
+		);
+	});
+});

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 },