diff --git a/.github/workflows/mcp-e2e.yml b/.github/workflows/mcp-e2e.yml new file mode 100644 index 0000000..037b8b6 --- /dev/null +++ b/.github/workflows/mcp-e2e.yml @@ -0,0 +1,90 @@ +name: MCP e2e + +# Runs the BTrace MCP server end-to-end tests against a BTrace distribution built from source, so the +# job never depends on a published or locally installed io.btrace:btrace artifact. No LLM is involved. + +on: + push: + branches: [main] + paths: + - 'plugins/btrace-observability/mcp/**' + - '.github/workflows/mcp-e2e.yml' + pull_request: + paths: + - 'plugins/btrace-observability/mcp/**' + - '.github/workflows/mcp-e2e.yml' + workflow_dispatch: + inputs: + btrace-ref: + description: 'BTrace branch, tag or SHA to build (default: the repository default branch)' + required: false + default: '' + +permissions: + contents: read + +concurrency: + group: mcp-e2e-${{ github.ref }} + cancel-in-progress: true + +jobs: + e2e: + runs-on: ubuntu-latest + timeout-minutes: 30 + steps: + # setup-java registers every JDK in ~/.m2/toolchains.xml, where Gradle discovers toolchains. + # BTrace builds with toolchains 8 to 26; the MCP module uses 11. The last entry is JAVA_HOME. + - name: Set up JDKs + uses: actions/setup-java@03ad4de0992f5dab5e18fcb136590ce7c4a0ac95 # v5 + with: + distribution: temurin + java-version: | + 8 + 11 + 17 + 24 + 26 + 21 + + - name: Check out agent plugins + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + + - name: Check out BTrace + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7 + with: + repository: btraceio/btrace + ref: ${{ inputs.btrace-ref }} + path: btrace + + - name: Set up Gradle + uses: gradle/actions/setup-gradle@3f131e8634966bd73d06cc69884922b02e6faf92 # v6 + + - name: Build the BTrace distribution JAR + working-directory: btrace + run: ./gradlew --no-daemon :btrace-dist:btraceJar + + - name: Locate the BTrace JAR + run: | + shopt -s nullglob + jars=("$GITHUB_WORKSPACE"/btrace/btrace-dist/build/resources/main/v*/libs/btrace.jar) + if [ "${#jars[@]}" -ne 1 ]; then + echo "::error::expected exactly one btrace.jar, found ${#jars[@]}: ${jars[*]}" + exit 1 + fi + echo "BTRACE_JAR=${jars[0]}" >> "$GITHUB_ENV" + echo "Using ${jars[0]} from BTrace $(git -C btrace rev-parse --short HEAD)" + + - name: Run MCP end-to-end tests + working-directory: plugins/btrace-observability/mcp + run: ./gradlew --no-daemon e2eTest + + - name: Upload e2e logs and reports + if: failure() + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 + with: + name: mcp-e2e-logs + retention-days: 7 + path: | + plugins/btrace-observability/mcp/build/e2e + plugins/btrace-observability/mcp/build/reports/tests/e2eTest + plugins/btrace-observability/mcp/build/test-results/e2eTest diff --git a/plugins/btrace-observability/.codex-plugin/mcp.json b/plugins/btrace-observability/.codex-plugin/mcp.json new file mode 100644 index 0000000..dcda961 --- /dev/null +++ b/plugins/btrace-observability/.codex-plugin/mcp.json @@ -0,0 +1,9 @@ +{ + "mcpServers": { + "btrace": { + "command": "jbang", + "args": ["mcp/src/main/java/io/btrace/mcp/BTraceMcpServer.java"], + "cwd": "." + } + } +} diff --git a/plugins/btrace-observability/.codex-plugin/plugin.json b/plugins/btrace-observability/.codex-plugin/plugin.json index 7709c81..3d34c1f 100644 --- a/plugins/btrace-observability/.codex-plugin/plugin.json +++ b/plugins/btrace-observability/.codex-plugin/plugin.json @@ -15,6 +15,7 @@ "diagnostics" ], "skills": "./skills/", + "mcpServers": "./.codex-plugin/mcp.json", "interface": { "displayName": "BTrace Observability", "shortDescription": "Diagnose Java behavior with focused probes.", diff --git a/plugins/btrace-observability/mcp/README.md b/plugins/btrace-observability/mcp/README.md index 15e5fa8..bed24aa 100644 --- a/plugins/btrace-observability/mcp/README.md +++ b/plugins/btrace-observability/mcp/README.md @@ -20,3 +20,21 @@ Run its unit tests with: ``` `./gradlew check` enforces at least 80% line coverage for the MCP Java code. + +Run the end-to-end tests with: + +```sh +./gradlew e2eTest +``` + +They exercise every tool and prompt against a sample JVM, observing probe effects through `@Export` +counters read with `jcmd`. No LLM is involved. By default they launch the server through JBang, as +the plugin manifests do, which needs `jbang` on `PATH` and a resolvable `io.btrace:btrace` artifact, +so they are not part of `check`. To test against a BTrace build instead, such as in CI, point them +at its masked JAR: + +```sh +BTRACE_JAR=/path/to/btrace/btrace-dist/build/resources/main/v/libs/btrace.jar ./gradlew e2eTest +``` + +The `list_probes` case is disabled because the BTrace client exits the JVM after listing probes. diff --git a/plugins/btrace-observability/mcp/build.gradle b/plugins/btrace-observability/mcp/build.gradle index 90a48e6..724f8b7 100644 --- a/plugins/btrace-observability/mcp/build.gradle +++ b/plugins/btrace-observability/mcp/build.gradle @@ -34,6 +34,43 @@ tasks.withType(Test).configureEach { jvmArgs '--add-modules', 'jdk.attach' } +// End-to-end tests run the real server against a live sample JVM, so they are not part of `check`. +// By default the server is launched through JBang, as the plugin manifests do, which needs jbang on +// PATH and a resolvable io.btrace:btrace artifact. Set BTRACE_JAR (or -PbtraceJar) to a masked +// btrace.jar, e.g. from a CI build of BTrace, to run the compiled server against that JAR instead. +def btraceJar = providers.gradleProperty('btraceJar') + .orElse(providers.environmentVariable('BTRACE_JAR')) + .map { it.trim() ? it.trim() : null } + +sourceSets { + e2eTest +} + +configurations { + e2eTestImplementation.extendsFrom testImplementation + e2eTestRuntimeOnly.extendsFrom testRuntimeOnly +} + +dependencies { + e2eTestImplementation sourceSets.main.output +} + +tasks.register('e2eTest', Test) { + description = 'Runs the MCP server end-to-end against a sample JVM.' + group = 'verification' + testClassesDirs = sourceSets.e2eTest.output.classesDirs + classpath = sourceSets.e2eTest.runtimeClasspath + shouldRunAfter tasks.named('test') + outputs.upToDateWhen { false } + systemProperty 'e2e.serverScript', file('src/main/java/io/btrace/mcp/BTraceMcpServer.java').absolutePath + systemProperty 'e2e.sampleClasspath', sourceSets.e2eTest.output.classesDirs.asPath + systemProperty 'e2e.logDir', layout.buildDirectory.dir('e2e').get().asFile.absolutePath + if (btraceJar.isPresent()) { + systemProperty 'e2e.btraceJar', file(btraceJar.get()).absolutePath + systemProperty 'e2e.serverClasspath', sourceSets.main.runtimeClasspath.asPath + } +} + jacoco { toolVersion = '0.8.13' } diff --git a/plugins/btrace-observability/mcp/src/e2eTest/java/io/btrace/mcp/McpServerE2ETest.java b/plugins/btrace-observability/mcp/src/e2eTest/java/io/btrace/mcp/McpServerE2ETest.java new file mode 100644 index 0000000..8045d58 --- /dev/null +++ b/plugins/btrace-observability/mcp/src/e2eTest/java/io/btrace/mcp/McpServerE2ETest.java @@ -0,0 +1,346 @@ +package io.btrace.mcp; + +import static io.btrace.mcp.McpStdioClient.isError; +import static io.btrace.mcp.McpStdioClient.names; +import static io.btrace.mcp.McpStdioClient.text; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; + +import sample.orders.SampleOrderApp; +import java.io.BufferedReader; +import java.io.File; +import java.io.InputStreamReader; +import java.net.ServerSocket; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.function.LongPredicate; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Disabled; +import org.junit.jupiter.api.MethodOrderer; +import org.junit.jupiter.api.Order; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; +import org.junit.jupiter.api.TestMethodOrder; + +/** + * Drives the real MCP server against a live sample JVM. The server is launched through JBang, as the + * plugin manifests launch it, unless the build supplies a BTrace JAR (see build.gradle). Probe + * effects are observed through {@code @Export} jvmstat counters read with {@code jcmd}, because no + * MCP tool returns probe output after deployment. + * + *

The steps share one server and one target JVM and must run in order. + */ +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +@TestMethodOrder(MethodOrderer.OrderAnnotation.class) +class McpServerE2ETest { + private static final Duration STARTUP = Duration.ofMinutes(5); + private static final Duration CALL = Duration.ofSeconds(90); + private static final Duration PROBE_EFFECT = Duration.ofSeconds(15); + + private static final String TARGET = SampleOrderApp.class.getName(); + private static final String RUNNING_PROBE = "E2ERunningProbe"; + private static final String EXIT_PROBE = "E2EExitProbe"; + + private Process sampleApp; + private String pid; + private int port; + private McpStdioClient server; + private Map initializeResult; + + @BeforeAll + void startSampleAppAndServer() throws Exception { + File logDir = new File(System.getProperty("e2e.logDir", "build/e2e")); + logDir.mkdirs(); + + sampleApp = + new ProcessBuilder(javaTool("java"), "-cp", requiredProperty("e2e.sampleClasspath"), TARGET) + .redirectError(new File(logDir, "sample-app.err")) + .start(); + pid = String.valueOf(sampleApp.pid()); + awaitReady(sampleApp); + + // A dedicated agent port keeps the run independent of other BTrace sessions on this host. + try (ServerSocket socket = new ServerSocket(0)) { + port = socket.getLocalPort(); + } + + server = new McpStdioClient(serverCommand(), new File(logDir, "mcp-server.err")); + initializeResult = server.initialize(STARTUP); + } + + @AfterAll + void stopProcesses() throws Exception { + if (server != null) { + server.close(); + } + if (sampleApp != null) { + sampleApp.destroy(); + if (!sampleApp.waitFor(10, TimeUnit.SECONDS)) { + sampleApp.destroyForcibly(); + } + } + } + + @Test + @Order(1) + @SuppressWarnings("unchecked") + void handshakeAdvertisesEveryTool() throws Exception { + Map serverInfo = (Map) initializeResult.get("serverInfo"); + assertEquals("btrace-mcp-server", serverInfo.get("name")); + + List tools = names(server.request("tools/list", new LinkedHashMap<>(), CALL).get("tools")); + assertEquals( + List.of( + "list_jvms", + "deploy_oneliner", + "deploy_script", + "list_probes", + "send_event", + "detach_probe", + "exit_probe"), + tools); + } + + @Test + @Order(2) + @SuppressWarnings("unchecked") + void servesEveryPromptForTheTarget() throws Exception { + List prompts = + names(server.request("prompts/list", new LinkedHashMap<>(), CALL).get("prompts")); + assertEquals( + List.of("diagnose_slow_endpoint", "find_exception_source", "profile_method"), prompts); + + Map> arguments = new LinkedHashMap<>(); + arguments.put( + "diagnose_slow_endpoint", args("endpoint_class", TARGET, "endpoint_method", "process")); + arguments.put( + "find_exception_source", args("exception_class", "java.lang.IllegalStateException")); + arguments.put("profile_method", args("class_name", TARGET, "method_name", "process")); + for (Map.Entry> prompt : arguments.entrySet()) { + prompt.getValue().put("pid", pid); + Map params = args("name", prompt.getKey(), "arguments", prompt.getValue()); + Map result = server.request("prompts/get", params, CALL); + List messages = (List) result.get("messages"); + assertFalse(messages.isEmpty(), prompt.getKey() + " returned no messages"); + String rendered = McpProtocol.toJson(messages); + assertTrue(rendered.contains(pid), prompt.getKey() + " does not mention the target PID"); + } + } + + @Test + @Order(3) + void listJvmsFindsTheSampleApp() throws Exception { + Map result = server.callTool("list_jvms", new LinkedHashMap<>(), CALL); + assertSuccess(result); + assertTrue( + text(result).contains("PID: " + pid + " | Main Class: " + TARGET), + "sample app missing from: " + text(result)); + } + + @Test + @Order(4) + void deployScriptInstrumentsTheTarget() throws Exception { + Map result = tool("deploy_script", "script", runningProbeScript()); + assertSuccess(result); + assertTrue(text(result).contains("Script deployed successfully to PID " + pid)); + awaitCounter(RUNNING_PROBE, "processCalls", calls -> calls > 0); + } + + @Test + @Order(5) + void sendEventRunsTheNamedHandler() throws Exception { + assertSuccess(tool("send_event", "event_name", "bump")); + awaitCounter(RUNNING_PROBE, "events", events -> events == 1); + } + + @Test + @Order(6) + void detachProbeLeavesTheProbeRunning() throws Exception { + assertSuccess(tool("detach_probe")); + + Map afterDetach = tool("send_event", "event_name", "bump"); + assertTrue(isError(afterDetach), "session should be gone after detach"); + assertTrue(text(afterDetach).contains("No active BTrace session")); + + long calls = counter(RUNNING_PROBE, "processCalls"); + awaitCounter(RUNNING_PROBE, "processCalls", now -> now > calls); + } + + @Test + @Order(7) + void deployOnelinerAndExitProbe() throws Exception { + Map deployed = + tool( + "deploy_oneliner", + "oneliner", + TARGET + "::process @return if duration>50ms { print method, duration }"); + assertSuccess(deployed); + assertTrue(text(deployed).contains("Probe deployed successfully to PID " + pid)); + + assertSuccess(tool("exit_probe")); + assertTrue(isError(tool("send_event")), "session should be gone after exit"); + } + + @Test + @Order(8) + void exitProbeRemovesTheInstrumentation() throws Exception { + assertSuccess(tool("deploy_script", "script", exitProbeScript())); + awaitCounter(EXIT_PROBE, "calls", calls -> calls > 0); + + assertSuccess(tool("exit_probe")); + // Give in-flight handler invocations time to drain, then require the counter to stay frozen. + Thread.sleep(1_000); + long frozen = counter(EXIT_PROBE, "calls"); + Thread.sleep(2_000); + assertEquals(frozen, counter(EXIT_PROBE, "calls"), "probe still firing after exit_probe"); + } + + @Test + @Order(9) + @Disabled( + "BTrace Client.connectAndListProbes calls System.exit(0) after the reply, which terminates" + + " the MCP server process; needs a BTrace client change") + void listProbesReportsTheRunningProbe() throws Exception { + Map result = tool("list_probes"); + assertSuccess(result); + assertTrue(text(result).contains(RUNNING_PROBE)); + assertTrue(server.isAlive(), "list_probes terminated the MCP server"); + } + + @Test + @Order(10) + void targetSurvivesTheWholeSession() { + assertTrue(sampleApp.isAlive(), "sample app died during the session"); + assertTrue(server.isAlive(), "MCP server died during the session"); + } + + private Map tool(String name, Object... keyValues) throws Exception { + Map arguments = args(keyValues); + arguments.put("pid", pid); + arguments.put("port", port); + return server.callTool(name, arguments, CALL); + } + + private static String runningProbeScript() { + return "import io.btrace.core.annotations.*;\n" + + "@BTrace public class " + + RUNNING_PROBE + + " {\n" + + " @io.btrace.core.annotations.Export static long processCalls;\n" + + " @io.btrace.core.annotations.Export static long events;\n" + + " @OnMethod(clazz = \"" + + TARGET + + "\", method = \"process\", location = @Location(Kind.RETURN))\n" + + " public static void onProcess() { processCalls++; }\n" + + " @OnEvent(\"bump\") public static void onBump() { events++; }\n" + + "}\n"; + } + + private static String exitProbeScript() { + return "import io.btrace.core.annotations.*;\n" + + "@BTrace public class " + + EXIT_PROBE + + " {\n" + + " @io.btrace.core.annotations.Export static long calls;\n" + + " @OnMethod(clazz = \"" + + TARGET + + "\", method = \"process\")\n" + + " public static void onProcess() { calls++; }\n" + + "}\n"; + } + + private void awaitCounter(String probe, String field, LongPredicate condition) throws Exception { + long deadline = System.nanoTime() + PROBE_EFFECT.toNanos(); + long value = -1; + while (System.nanoTime() < deadline) { + value = counter(probe, field); + if (condition.test(value)) { + return; + } + Thread.sleep(200); + } + fail(probe + "." + field + " did not reach the expected value; last=" + value); + } + + /** Reads an {@code @Export} field of a probe from the target's jvmstat counters; -1 if absent. */ + private long counter(String probe, String field) throws Exception { + Process jcmd = + new ProcessBuilder(javaTool("jcmd"), pid, "PerfCounter.print") + .redirectErrorStream(true) + .start(); + String suffix = "/" + probe + "." + field + "="; + long value = -1; + try (BufferedReader out = + new BufferedReader(new InputStreamReader(jcmd.getInputStream(), StandardCharsets.UTF_8))) { + String line; + while ((line = out.readLine()) != null) { + int at = line.indexOf(suffix); + if (at >= 0) { + value = Long.parseLong(line.substring(at + suffix.length()).trim()); + } + } + } + assertTrue(jcmd.waitFor(30, TimeUnit.SECONDS), "jcmd did not finish"); + return value; + } + + private static void assertSuccess(Map result) { + assertFalse(isError(result), "tool failed: " + text(result)); + } + + private static void awaitReady(Process app) throws Exception { + BufferedReader out = + new BufferedReader(new InputStreamReader(app.getInputStream(), StandardCharsets.UTF_8)); + String line = out.readLine(); + if (!SampleOrderApp.READY.equals(line)) { + fail("sample app did not start: " + line); + } + } + + private static Map args(Object... keyValues) { + Map map = new LinkedHashMap<>(); + for (int i = 0; i < keyValues.length; i += 2) { + map.put((String) keyValues[i], keyValues[i + 1]); + } + return map; + } + + /** + * JBang launch as in the plugin manifests, or, when the build supplies a BTrace JAR, the compiled + * server on a plain classpath so the run does not depend on the published BTrace artifact. + */ + private static List serverCommand() { + String btraceJar = System.getProperty("e2e.btraceJar"); + if (btraceJar == null || btraceJar.isEmpty()) { + return List.of("jbang", requiredProperty("e2e.serverScript")); + } + if (!new File(btraceJar).isFile()) { + throw new IllegalStateException("BTrace JAR not found: " + btraceJar); + } + return List.of( + javaTool("java"), + "--add-modules=jdk.attach", + "-cp", + requiredProperty("e2e.serverClasspath") + File.pathSeparator + btraceJar, + BTraceMcpServer.class.getName()); + } + + private static String javaTool(String name) { + return new File(System.getProperty("java.home"), "bin/" + name).getAbsolutePath(); + } + + private static String requiredProperty(String name) { + String value = System.getProperty(name); + if (value == null || value.isEmpty()) { + throw new IllegalStateException("Run through `./gradlew e2eTest`; missing -D" + name); + } + return value; + } +} diff --git a/plugins/btrace-observability/mcp/src/e2eTest/java/io/btrace/mcp/McpStdioClient.java b/plugins/btrace-observability/mcp/src/e2eTest/java/io/btrace/mcp/McpStdioClient.java new file mode 100644 index 0000000..3f17166 --- /dev/null +++ b/plugins/btrace-observability/mcp/src/e2eTest/java/io/btrace/mcp/McpStdioClient.java @@ -0,0 +1,171 @@ +package io.btrace.mcp; + +import java.io.BufferedReader; +import java.io.File; +import java.io.IOException; +import java.io.InputStreamReader; +import java.io.OutputStream; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; + +/** Minimal MCP client that drives a server process over JSON-RPC on stdin/stdout. */ +final class McpStdioClient implements AutoCloseable { + private static final String EOF = "\u0000EOF"; + + private final Process process; + private final OutputStream stdin; + private final BlockingQueue lines = new LinkedBlockingQueue<>(); + private int nextId = 1; + + McpStdioClient(List command, File stderrLog) throws IOException { + process = + new ProcessBuilder(command) + .redirectError(ProcessBuilder.Redirect.to(stderrLog)) + .start(); + stdin = process.getOutputStream(); + Thread reader = + new Thread( + () -> { + try (BufferedReader in = + new BufferedReader( + new InputStreamReader(process.getInputStream(), StandardCharsets.UTF_8))) { + String line; + while ((line = in.readLine()) != null) { + lines.add(line); + } + } catch (IOException ignored) { + // process went away; the EOF marker below reports it + } + lines.add(EOF); + }, + "mcp-stdout"); + reader.setDaemon(true); + reader.start(); + } + + Map initialize(Duration timeout) throws Exception { + Map clientInfo = new LinkedHashMap<>(); + clientInfo.put("name", "btrace-mcp-e2e"); + clientInfo.put("version", "0"); + Map params = new LinkedHashMap<>(); + params.put("protocolVersion", "2024-11-05"); + params.put("capabilities", new LinkedHashMap<>()); + params.put("clientInfo", clientInfo); + Map result = request("initialize", params, timeout); + Map initialized = new LinkedHashMap<>(); + initialized.put("jsonrpc", "2.0"); + initialized.put("method", "notifications/initialized"); + write(initialized); + return result; + } + + /** Sends a request and returns its {@code result}; fails on a JSON-RPC error or bad stdout. */ + @SuppressWarnings("unchecked") + Map request(String method, Map params, Duration timeout) + throws Exception { + int id = nextId++; + Map message = new LinkedHashMap<>(); + message.put("jsonrpc", "2.0"); + message.put("id", id); + message.put("method", method); + message.put("params", params); + write(message); + + long deadline = System.nanoTime() + timeout.toNanos(); + while (true) { + long remaining = deadline - System.nanoTime(); + String line = remaining > 0 ? lines.poll(remaining, TimeUnit.NANOSECONDS) : null; + if (line == null) { + throw new AssertionError(method + " got no response within " + timeout); + } + if (line.equals(EOF)) { + throw new AssertionError( + method + ": server closed stdout (exit " + exitCodeIfDone() + ")"); + } + Map response; + try { + response = McpProtocol.parseJson(line); + } catch (RuntimeException e) { + throw new AssertionError("Non-JSON line on the MCP stdout transport: " + line, e); + } + if (!Integer.valueOf(id).equals(asInt(response.get("id")))) { + continue; + } + if (response.containsKey("error")) { + throw new AssertionError(method + " failed: " + McpProtocol.toJson(response.get("error"))); + } + return (Map) response.get("result"); + } + } + + Map callTool(String name, Map arguments, Duration timeout) + throws Exception { + Map params = new LinkedHashMap<>(); + params.put("name", name); + params.put("arguments", arguments); + return request("tools/call", params, timeout); + } + + @SuppressWarnings("unchecked") + static String text(Map toolResult) { + StringBuilder text = new StringBuilder(); + for (Object item : (List) toolResult.get("content")) { + text.append(((Map) item).get("text")); + } + return text.toString(); + } + + static boolean isError(Map toolResult) { + return Boolean.TRUE.equals(toolResult.get("isError")); + } + + @SuppressWarnings("unchecked") + static List names(Object items) { + List names = new ArrayList<>(); + for (Object item : (List) items) { + names.add((String) ((Map) item).get("name")); + } + return names; + } + + boolean isAlive() { + return process.isAlive(); + } + + private synchronized void write(Map message) throws IOException { + stdin.write((McpProtocol.toJson(message) + "\n").getBytes(StandardCharsets.UTF_8)); + stdin.flush(); + } + + private String exitCodeIfDone() { + try { + return process.waitFor(2, TimeUnit.SECONDS) ? String.valueOf(process.exitValue()) : "n/a"; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + return "n/a"; + } + } + + private static Integer asInt(Object value) { + return value instanceof Number ? ((Number) value).intValue() : null; + } + + @Override + public void close() throws Exception { + try { + stdin.close(); + } catch (IOException ignored) { + // already gone + } + if (!process.waitFor(10, TimeUnit.SECONDS)) { + process.destroyForcibly().waitFor(10, TimeUnit.SECONDS); + } + } +} diff --git a/plugins/btrace-observability/mcp/src/e2eTest/java/sample/orders/SampleOrderApp.java b/plugins/btrace-observability/mcp/src/e2eTest/java/sample/orders/SampleOrderApp.java new file mode 100644 index 0000000..b868b69 --- /dev/null +++ b/plugins/btrace-observability/mcp/src/e2eTest/java/sample/orders/SampleOrderApp.java @@ -0,0 +1,32 @@ +package sample.orders; + +/** + * Target JVM for the MCP end-to-end tests: calls {@link #process(int)} in a loop, with every fifth + * call slow and every seventh call throwing, so latency and failure probes have something to see. + * It lives outside {@code io.btrace}, whose classes BTrace refuses to instrument. + */ +public final class SampleOrderApp { + public static final String READY = "SAMPLE_ORDER_APP_READY"; + + private SampleOrderApp() {} + + static int process(int id) throws InterruptedException { + Thread.sleep(id % 5 == 0 ? 60 : 5); + if (id % 7 == 0) { + throw new IllegalStateException("order " + id + " rejected"); + } + return id * 2; + } + + public static void main(String[] args) throws Exception { + System.out.println(READY); + System.out.flush(); + for (int id = 1; ; id++) { + try { + process(id); + } catch (IllegalStateException expected) { + // failures are part of the workload + } + } + } +} diff --git a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceClient.java b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceClient.java index 543dc0a..9d4684a 100644 --- a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceClient.java +++ b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceClient.java @@ -24,6 +24,8 @@ import java.net.URL; import java.security.CodeSource; import java.security.ProtectionDomain; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * Reflective adapter for the client section of BTrace's masked distribution JAR. @@ -33,6 +35,9 @@ * {@code MaskedClassLoader} boundary used by {@code java -jar btrace.jar}. */ public final class BTraceClient { + private static final Logger log = LoggerFactory.getLogger(BTraceClient.class); + private static final String JAR_PATH_PROPERTY = "btrace.jar.path"; + @FunctionalInterface public interface CommandListener { void onCommand(Object command) throws Exception; @@ -60,6 +65,11 @@ private BTraceClient( public static BTraceClient create(int port) throws Exception { Class loaderClass = Class.forName("io.btrace.boot.MaskedClassLoader"); File jar = jarFile(loaderClass); + // The client locates the agent JAR via this property or a classpath entry named btrace.jar; + // JBang puts the versioned Maven artifact (btrace-.jar) on the classpath instead. + if (jar.isFile() && System.getProperty(JAR_PATH_PROPERTY) == null) { + System.setProperty(JAR_PATH_PROPERTY, jar.getAbsolutePath()); + } Constructor constructor = loaderClass.getConstructor(File.class, String.class, ClassLoader.class); ClassLoader loader = (ClassLoader) constructor.newInstance(jar, "client", BTraceClient.class.getClassLoader()); @@ -115,6 +125,27 @@ public void submit(String host, String fileName, byte[] code, String[] args, Com .invoke(delegate, host, fileName, code, args, listenerProxy(listener)); } + /** + * Submits a probe on a daemon thread. The client's submit runs the probe command loop until the + * probe exits or the client is closed, so it must not block the MCP request thread. + */ + public Thread submitInBackground( + String host, String fileName, byte[] code, String[] args, CommandListener listener) { + Thread thread = + new Thread( + () -> { + try { + submit(host, fileName, code, args, listener); + } catch (Exception e) { + log.debug("Probe command loop for {} ended", fileName, e); + } + }, + "btrace-probe-" + fileName); + thread.setDaemon(true); + thread.start(); + return thread; + } + public void listProbes(String host, CommandListener listener) throws Exception { clientClass .getMethod("connectAndListProbes", String.class, commandListenerClass) diff --git a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployOnelinerHandler.java b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployOnelinerHandler.java index c41dfe1..3f47286 100644 --- a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployOnelinerHandler.java +++ b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployOnelinerHandler.java @@ -114,12 +114,16 @@ public static Map execute(Map arguments) { AtomicBoolean success = new AtomicBoolean(false); AtomicBoolean exited = new AtomicBoolean(false); - client.submit( + client.submitInBackground( "localhost", fileName, code, new String[0], cmd -> { + if (statusLatch.getCount() == 0) { + // Later probe output is not reported; do not buffer it for the probe's lifetime. + return; + } output.append(client.printableText(cmd)); int type = client.commandType(cmd); if (type == client.commandConstant("STATUS")) { @@ -136,10 +140,12 @@ public static Map execute(Map arguments) { boolean started = statusLatch.await(PROBE_TIMEOUT_SECONDS, TimeUnit.SECONDS); if (!started) { + closeQuietly(client); return toolResult("Probe deployment timed out after " + PROBE_TIMEOUT_SECONDS + "s", true); } if (exited.get() && !success.get()) { + closeQuietly(client); return toolResult("Probe exited with error:\n" + output.toString(), true); } @@ -167,6 +173,14 @@ public static Map execute(Map arguments) { } } + private static void closeQuietly(BTraceClient client) { + try { + client.close(); + } catch (Exception e) { + log.debug("Failed to close BTrace client", e); + } + } + private static String getStringArg(Map args, String key) { Object val = args == null ? null : args.get(key); return val == null ? null : val.toString(); diff --git a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployScriptHandler.java b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployScriptHandler.java index 61ee8e5..302181e 100644 --- a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployScriptHandler.java +++ b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/DeployScriptHandler.java @@ -27,6 +27,8 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.regex.Matcher; +import java.util.regex.Pattern; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -35,6 +37,9 @@ public final class DeployScriptHandler { private static final Logger log = LoggerFactory.getLogger(DeployScriptHandler.class); private static final int DEFAULT_PORT = 2020; private static final int PROBE_TIMEOUT_SECONDS = 30; + private static final Pattern COMMENTS = Pattern.compile("//[^\\n]*|/\\*.*?\\*/", Pattern.DOTALL); + private static final Pattern CLASS_NAME = + Pattern.compile("\\bclass\\s+([A-Za-z_$][A-Za-z0-9_$]*)"); private DeployScriptHandler() {} @@ -107,7 +112,7 @@ public static Map execute(Map arguments) { } try { - String fileName = "BTraceScript_" + System.currentTimeMillis() + ".java"; + String fileName = scriptFileName(script); // Compile the script BTraceClient client = ClientManager.getClient(port); @@ -129,12 +134,16 @@ public static Map execute(Map arguments) { AtomicBoolean exited = new AtomicBoolean(false); final String[] finalArgs = btraceArgs; - client.submit( + client.submitInBackground( "localhost", fileName, code, finalArgs, cmd -> { + if (statusLatch.getCount() == 0) { + // Later probe output is not reported; do not buffer it for the probe's lifetime. + return; + } output.append(client.printableText(cmd)); int type = client.commandType(cmd); if (type == client.commandConstant("STATUS")) { @@ -150,10 +159,12 @@ public static Map execute(Map arguments) { boolean started = statusLatch.await(PROBE_TIMEOUT_SECONDS, TimeUnit.SECONDS); if (!started) { + closeQuietly(client); return toolResult("Probe deployment timed out after " + PROBE_TIMEOUT_SECONDS + "s", true); } if (exited.get() && !success.get()) { + closeQuietly(client); return toolResult("Probe exited with error:\n" + output.toString(), true); } @@ -172,6 +183,25 @@ public static Map execute(Map arguments) { } } + /** + * Names the source file after the script's first declared class, as javac requires for public + * classes; a non-public class compiles as BTrace short syntax, which rejects fields and statics. + */ + private static String scriptFileName(String script) { + String code = COMMENTS.matcher(script).replaceAll(" "); + Matcher matcher = CLASS_NAME.matcher(code); + String name = matcher.find() ? matcher.group(1) : "BTraceScript_" + System.currentTimeMillis(); + return name + ".java"; + } + + private static void closeQuietly(BTraceClient client) { + try { + client.close(); + } catch (Exception e) { + log.debug("Failed to close BTrace client", e); + } + } + private static String getStringArg(Map args, String key) { Object val = args == null ? null : args.get(key); return val == null ? null : val.toString(); diff --git a/plugins/btrace-observability/mcp/src/test/java/io/btrace/client/Client.java b/plugins/btrace-observability/mcp/src/test/java/io/btrace/client/Client.java index 437c198..0c4ce35 100644 --- a/plugins/btrace-observability/mcp/src/test/java/io/btrace/client/Client.java +++ b/plugins/btrace-observability/mcp/src/test/java/io/btrace/client/Client.java @@ -4,6 +4,7 @@ import io.btrace.core.comm.CommandListener; import io.btrace.core.comm.FakePrintableCommand; import java.io.PrintWriter; +import java.util.concurrent.CountDownLatch; public class Client { public static Client last; @@ -13,6 +14,8 @@ public class Client { public String event; public boolean disconnected; public boolean closed; + public String submittedFileName; + private final CountDownLatch closeLatch = new CountDownLatch(1); public Client(int port) { this.port = port; @@ -29,7 +32,10 @@ public void attach(String pid, String systemClassPath, String bootstrapClassPath public void submit( String host, String fileName, byte[] code, String[] args, CommandListener listener) throws Exception { + submittedFileName = fileName; listener.onCommand(new FakePrintableCommand(Command.STATUS, host + ':' + fileName)); + // Like the real client, keep running the probe command loop until the client is closed. + closeLatch.await(); } public void connectAndListProbes(String host, CommandListener listener) throws Exception { @@ -54,5 +60,6 @@ public void sendEvent(String name) { public void close() { closed = true; + closeLatch.countDown(); } } diff --git a/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceClientTest.java b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceClientTest.java index 1f96e17..cddf0c8 100644 --- a/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceClientTest.java +++ b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceClientTest.java @@ -9,6 +9,8 @@ import io.btrace.core.comm.Command; import java.io.PrintWriter; import java.io.StringWriter; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; @@ -29,7 +31,19 @@ void invokesTheMaskedClientContractReflectively() throws Exception { assertEquals("42", delegate.attachedPid); AtomicReference submitted = new AtomicReference<>(); - client.submit("localhost", "Trace.java", new byte[] {1}, new String[] {"arg"}, submitted::set); + CountDownLatch status = new CountDownLatch(1); + Thread probe = + client.submitInBackground( + "localhost", + "Trace.java", + new byte[] {1}, + new String[] {"arg"}, + command -> { + submitted.set(command); + status.countDown(); + }); + assertTrue(status.await(5, TimeUnit.SECONDS)); + assertTrue(probe.isDaemon()); assertEquals("localhost:Trace.java", client.printableText(submitted.get())); assertEquals(Command.STATUS, client.commandType(submitted.get())); assertEquals(Command.STATUS, client.commandConstant("STATUS")); @@ -47,6 +61,8 @@ void invokesTheMaskedClientContractReflectively() throws Exception { client.sendDisconnect(); client.sendExit(7); client.close(); + probe.join(5000); + assertFalse(probe.isAlive()); assertTrue(delegate.disconnected); assertEquals(7, delegate.exitCode); assertTrue(delegate.closed); diff --git a/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/ToolHandlersTest.java b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/ToolHandlersTest.java index 3aad883..35adcde 100644 --- a/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/ToolHandlersTest.java +++ b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/ToolHandlersTest.java @@ -4,6 +4,7 @@ import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertTrue; +import io.btrace.client.Client; import io.btrace.mcp.tools.DeployOnelinerHandler; import io.btrace.mcp.tools.DeployScriptHandler; import io.btrace.mcp.tools.DetachProbeHandler; @@ -73,6 +74,24 @@ void runsTheMaskedClientToolWorkflowWithFixtures() { assertTrue(text(probes).contains("Active probes on PID 53")); } + @Test + void namesScriptSourceAfterItsDeclaredClass() { + String script = "// this class traces orders\n/* class Ignored */\n@BTrace public class OrderTrace {}"; + assertFalse( + (Boolean) + DeployScriptHandler.execute(Map.of("pid", "54", "script", script, "port", 2054)) + .get("isError")); + assertEquals("OrderTrace.java", Client.last.submittedFileName); + ExitProbeHandler.execute(Map.of("pid", "54", "port", 2054)); + + assertFalse( + (Boolean) + DeployScriptHandler.execute(Map.of("pid", "55", "script", "@BTrace", "port", 2055)) + .get("isError")); + assertTrue(Client.last.submittedFileName.matches("BTraceScript_\\d+\\.java")); + ExitProbeHandler.execute(Map.of("pid", "55", "port", 2055)); + } + @SuppressWarnings("unchecked") private static void assertRequired(Map schema, String... expected) { assertEquals(List.of(expected), schema.get("inputSchema") instanceof Map diff --git a/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md b/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md index 7dfda19..57ea461 100644 --- a/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md +++ b/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md @@ -1,13 +1,22 @@ --- name: btrace-mcp-operations -description: Use when an AI client should operate BTrace through the BTrace MCP server to list local JVMs, deploy probes, inspect output, or clean up diagnostic sessions. +description: Use when an AI client should operate BTrace through the BTrace MCP server to list local JVMs, deploy probes, trigger probe events, or clean up diagnostic sessions. --- # MCP Operations The BTrace MCP server offers structured local-JVM operations such as `list_jvms`, `deploy_oneliner`, -`deploy_script`, `list_probes`, and `exit_probe`. It is suitable when the AI client and target JVM -are on the same host and the operator wants an auditable conversational workflow. +`deploy_script`, `send_event`, `detach_probe`, and `exit_probe`. It is suitable when the AI client +and target JVM are on the same host and the operator wants an auditable conversational workflow. + +Do not call `list_probes`. The current BTrace client exits the process after returning the probe +list, which terminates the MCP server and drops every open probe session, so `send_event` and +`exit_probe` can no longer reach probes deployed earlier. Track each deployment yourself instead: +target PID, agent port, probe class or oneliner, and observation window. + +Deployment results include only the probe's initial output; no tool returns what the probe prints +afterwards. When the operator needs ongoing results, have the probe write them somewhere they can +read, such as exported counters (`@Export` fields, read with `jcmd PerfCounter.print`). - The MCP server uses the local JVM Attach API; it cannot attach across SSH, Docker, or Kubernetes boundaries. For those cases, first use `btrace-runtime-access` to place the client/server in the