mcpSamplingLog.test.ts

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/

import * as assert from "assert";
import * as sinon from "sinon";
import { ensureNoDisposablesAreLeakedInTestSuite } from "../../../../../base/test/common/utils.js";
import {
	StorageScope
} from "../../../../../platform/storage/common/storage.js";
import { TestStorageService } from "../../../../test/common/workbenchTestServices.js";
import { McpSamplingLog } from "../../common/mcpSamplingLog.js";
import { IMcpServer } from "../../common/mcpTypes.js";

suite("MCP - Sampling Log", () => {
	const ds = ensureNoDisposablesAreLeakedInTestSuite();
	const fakeServer: IMcpServer = {
		definition: { id: "testServer" },
		readDefinitions: () => ({
			get: () => ({ collection: { scope: StorageScope.APPLICATION } }),
		}),
	} as any;

	let log: McpSamplingLog;
	let storage: TestStorageService;
	let clock: sinon.SinonFakeTimers;

	setup(() => {
		storage = ds.add(new TestStorageService());
		log = ds.add(new McpSamplingLog(storage));
		clock = sinon.useFakeTimers();
		clock.setSystemTime(new Date("2023-10-01T00:00:00Z").getTime());
	});

	teardown(() => {
		clock.restore();
	});

	test("logs a single request", async () => {
		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "test request" } }],
			"test response here",
			"foobar9000",
		);

		// storage.testEmitWillSaveState(WillSaveStateReason.NONE);
		await storage.flush();
		assert.deepStrictEqual(
			(storage.getObject("mcp.sampling.logs", StorageScope.APPLICATION) as any),
			[
				[
					"testServer",
					{
						head: 19631,
						bins: [1, 0, 0, 0, 0, 0, 0],
						lastReqs: [
							{
								request: [{ role: "user", content: { type: "text", text: "test request" } }],
								response: "test response here",
								at: 1696118400000,
								model: "foobar9000",
							},
						],
					},
				],
			],
		);
	});

	test("logs multiple requests on the same day", async () => {
		// First request
		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "first request" } }],
			"first response",
			"foobar9000",
		);

		// Advance time by a few hours but stay on the same day
		clock.tick(5 * 60 * 60 * 1000); // 5 hours

		// Second request
		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "second request" } }],
			"second response",
			"foobar9000",
		);

		await storage.flush();
		const data = (storage.getObject("mcp.sampling.logs", StorageScope.APPLICATION) as any)[0][1];

		// Verify the bin for the current day has 2 requests
		assert.strictEqual(data.bins[0], 2);

		// Verify both requests are in the lastReqs array, with the most recent first
		assert.strictEqual(data.lastReqs.length, 2);
		assert.strictEqual(data.lastReqs[0].request[0].content.text, "second request");
		assert.strictEqual(data.lastReqs[1].request[0].content.text, "first request");
	});

	test("shifts bins when adding requests on different days", async () => {
		// First request on day 1
		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "day 1 request" } }],
			"day 1 response",
			"foobar9000",
		);

		// Advance time to the next day
		clock.tick(24 * 60 * 60 * 1000);

		// Second request on day 2
		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "day 2 request" } }],
			"day 2 response",
			"foobar9000",
		);

		await storage.flush();
		const data = (storage.getObject("mcp.sampling.logs", StorageScope.APPLICATION) as any)[0][1];

		// Verify the bins: day 2 should have 1 request, day 1 should have 1 request
		assert.strictEqual(data.bins[0], 1); // day 2
		assert.strictEqual(data.bins[1], 1); // day 1

		// Advance time by 5 more days
		clock.tick(5 * 24 * 60 * 60 * 1000);

		// Request on day 7
		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "day 7 request" } }],
			"day 7 response",
			"foobar9000",
		);

		await storage.flush();
		const updatedData = (storage.getObject("mcp.sampling.logs", StorageScope.APPLICATION) as any)[0][1];

		// Verify the bins have shifted correctly
		assert.strictEqual(updatedData.bins[0], 1); // day 7
		assert.strictEqual(updatedData.bins[5], 1); // day 2
		assert.strictEqual(updatedData.bins[6], 1); // day 1
	});

	test("limits the number of stored requests", async () => {
		// Add more than the maximum number of requests (Constants.SamplingLastNMessage = 30)
		for (let i = 0; i < 35; i++) {
			log.add(
				fakeServer,
				[{ role: "user", content: { type: "text", text: `request ${i}` } }],
				`response ${i}`,
				"foobar9000",
			);
		}

		await storage.flush();
		const data = (storage.getObject("mcp.sampling.logs", StorageScope.APPLICATION) as any)[0][1];

		// Verify only the last 30 requests are kept
		assert.strictEqual(data.lastReqs.length, 30);
		assert.strictEqual(data.lastReqs[0].request[0].content.text, "request 34");
		assert.strictEqual(data.lastReqs[29].request[0].content.text, "request 5");
	});

	test("handles different content types", async () => {
		// Add a request with text content
		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "text request" } }],
			"text response",
			"foobar9000",
		);

		// Add a request with image content
		log.add(
			fakeServer,
			[{
				role: "user",
				content: {
					type: "image",
					data: "base64data",
					mimeType: "image/png"
				}
			}],
			"image response",
			"foobar9000",
		);

		// Add a request with mixed content
		log.add(
			fakeServer,
			[
				{ role: "user", content: { type: "text", text: "text and image" } },
				{
					role: "assistant",
					content: {
						type: "image",
						data: "base64data",
						mimeType: "image/jpeg"
					}
				}
			],
			"mixed response",
			"foobar9000",
		);

		await storage.flush();
		const data = (storage.getObject("mcp.sampling.logs", StorageScope.APPLICATION) as any)[0][1];

		// Verify all requests are stored correctly
		assert.strictEqual(data.lastReqs.length, 3);
		assert.strictEqual(data.lastReqs[0].request.length, 2); // Mixed content request has 2 messages
		assert.strictEqual(data.lastReqs[1].request[0].content.type, "image");
		assert.strictEqual(data.lastReqs[2].request[0].content.type, "text");
	});

	test("handles multiple servers", async () => {
		const fakeServer2: IMcpServer = {
			definition: { id: "testServer2" },
			readDefinitions: () => ({
				get: () => ({ collection: { scope: StorageScope.APPLICATION } }),
			}),
		} as any;

		log.add(
			fakeServer,
			[{ role: "user", content: { type: "text", text: "server1 request" } }],
			"server1 response",
			"foobar9000",
		);

		log.add(
			fakeServer2,
			[{ role: "user", content: { type: "text", text: "server2 request" } }],
			"server2 response",
			"foobar9000",
		);

		await storage.flush();
		const storageData = (storage.getObject("mcp.sampling.logs", StorageScope.APPLICATION) as any);

		// Verify both servers have their data stored
		assert.strictEqual(storageData.length, 2);
		assert.strictEqual(storageData[0][0], "testServer");
		assert.strictEqual(storageData[1][0], "testServer2");
	});
});