@skyramp/mcp
Version:
Skyramp MCP (Model Context Protocol) Server - AI-powered test generation and execution
198 lines (193 loc) • 9.07 kB
JavaScript
import { z } from "zod";
import { SkyrampClient } from "@skyramp/skyramp";
import openProxyTerminalTracked from "../../utils/proxy-terminal.js";
import { TELEMETRY_entrypoint_FIELD_NAME } from "../../utils/utils.js";
import { logger } from "../../utils/logger.js";
import { basePlaywrightSchema, baseSchema, } from "../../types/TestTypes.js";
import { AnalyticsService } from "../../services/AnalyticsService.js";
import path from "path";
const TOOL_NAME = "skyramp_start_trace_collection";
export function registerTraceTool(server) {
server.registerTool(TOOL_NAME, {
description: `Start trace collection using Skyramp's comprehensive tracing capabilities.
Trace collection monitors your application in real-time to capture user interactions, API calls, service communications, and data flows. This captured data is then used to generate comprehensive integration, load, and E2E tests that reflect real-world usage patterns.
COLLECTION TYPES:
• UI-Only Tracing: Captures frontend user interactions for UI test generation
• Backend-Only Tracing: Records API calls and service communications for integration tests
• Full-Stack Tracing: Combines UI and backend tracing for comprehensive E2E test generation
WORKFLOW:
1. Start trace collection with desired configuration
2. Interact with your application (UI clicks, API calls, workflows)
3. Stop trace collection to save captured data
4. Use traces to generate realistic test scenarios
For detailed documentation visit: https://www.skyramp.dev/docs/load-test/advanced-generation#start-trace-collection`,
inputSchema: {
playwright: z
.boolean()
.describe("Whether to enable Playwright for trace collection. Set to true for UI interactions, false for API-only tracing")
.default(true),
playwrightStoragePath: basePlaywrightSchema.shape.playwrightStoragePath,
playwrightSaveStoragePath: basePlaywrightSchema.shape.playwrightSaveStoragePath,
playwrightViewportSize: basePlaywrightSchema.shape.playwrightViewportSize,
runtime: z
.string()
.default("docker")
.describe("Runtime environment for trace collection. Currently only 'docker' is supported and is used as the default."),
include: z
.array(z.string())
.default([])
.describe("List of service names or patterns to include in trace collection"),
exclude: z
.array(z.string())
.default([])
.describe("List of service names or patterns to exclude from trace collection"),
noProxy: z
.array(z.string())
.default([])
.describe("List of hosts or patterns that should bypass proxy during tracing"),
dockerNetwork: z
.string()
.default("skyramp")
.describe("Docker network name to use for containerized trace collection"),
dockerWorkerPort: z
.number()
.default(35142)
.describe("Port number for the Docker worker service during trace collection"),
outputDir: baseSchema.shape.outputDir,
prompt: z
.string()
.describe("The prompt user provided to start trace collection"),
},
_meta: {
keywords: ["start trace", "trace collection", "trace generation"],
},
}, async (params, extra) => {
let errorResult;
// Helper to send progress notifications to the MCP client
const sendProgress = async (progress, total, message) => {
const progressToken = extra._meta?.progressToken;
if (progressToken !== undefined) {
const notification = {
method: "notifications/progress",
params: {
progressToken,
progress,
total,
message,
},
};
await extra.sendNotification(notification);
}
};
logger.info("Starting trace collection", {
playwright: params.playwright,
runtime: params.runtime,
include: params.include,
exclude: params.exclude,
noProxy: params.noProxy,
dockerNetwork: params.dockerNetwork,
dockerWorkerPort: params.dockerWorkerPort,
outputDir: params.outputDir,
prompt: params.prompt,
});
let saveStoragePath = params.playwrightSaveStoragePath;
if (saveStoragePath) {
// If saveStoragePath is just a filename (no directory separators), use outputDir
if (params.outputDir && !saveStoragePath.includes(path.sep) && !path.isAbsolute(saveStoragePath)) {
saveStoragePath = path.join(params.outputDir, saveStoragePath);
}
logger.info("Session storage will be saved to:", { saveStoragePath });
}
try {
// Send initial progress
await sendProgress(0, 100, "Initializing trace collection...");
const client = new SkyrampClient();
const generateOptions = {
testType: "trace",
runtime: params.runtime,
dockerNetwork: params.dockerNetwork,
dockerWorkerPort: params.dockerWorkerPort.toString(),
generateInclude: params.include,
generateExclude: params.exclude,
generateNoProxy: params.noProxy,
unblock: true,
playwright: params.playwright,
playwrightStoragePath: params.playwrightStoragePath,
playwrightViewportSize: params.playwrightViewportSize,
entrypoint: TELEMETRY_entrypoint_FIELD_NAME,
};
if (saveStoragePath) {
generateOptions.playwrightSaveStoragePath = saveStoragePath;
}
// Send progress for configuration
const traceMode = params.playwright ? "UI + Backend" : "Backend-only";
await sendProgress(30, 100, `Configuring ${traceMode} trace collection...`);
// Start periodic progress updates during the long-running operation
const startTime = Date.now();
let progressTick = 0;
const progressInterval = setInterval(() => {
progressTick++;
const elapsed = Math.floor((Date.now() - startTime) / 1000);
// Progress moves from 30% to 65% during service startup
const percent = Math.min(30 + progressTick * 5, 65);
sendProgress(percent, 100, `Starting trace collection service... (${elapsed}s elapsed)`).catch(() => {
// Ignore progress notification errors
});
// Also log to server for visibility
logger.info(`Trace collection progress: ${percent}% (${elapsed}s elapsed)`);
}, 10000); // 10 seconds
let result;
try {
result = await client.generateRestTest(generateOptions);
}
finally {
clearInterval(progressInterval);
}
if (result.toLowerCase().includes("failed")) {
errorResult = {
content: [
{
type: "text",
text: `Trace collection failed: ${result}`,
},
],
isError: true,
};
return errorResult;
}
// Send progress for opening proxy terminal
await sendProgress(70, 100, "Opening proxy terminal...");
await openProxyTerminalTracked();
// Send completion progress
await sendProgress(100, 100, "Trace collection started successfully");
errorResult = {
content: [
{
type: "text",
text: `Trace collection started: ${result}. Please let me know when you are ready to stop the trace collection.`,
},
],
};
return errorResult;
}
catch (error) {
errorResult = {
content: [
{
type: "text",
text: `Trace generation failed: ${error.message}`,
},
],
isError: true,
};
return errorResult;
}
finally {
const recordParams = {
prompt: params.prompt,
playwright: params.playwright.toString(),
};
AnalyticsService.pushMCPToolEvent(TOOL_NAME, errorResult, recordParams);
}
});
}