diff --git a/evals/btrace-observability.json b/evals/btrace-observability.json index 644b8d6..9b79c67 100644 --- a/evals/btrace-observability.json +++ b/evals/btrace-observability.json @@ -161,12 +161,12 @@ "list_jvms", "list_probes", "pid", - "btrace -r" + "stop_detached_probe" ], "mustAvoid": [ "kill -9" ], - "rubric": "Find the target PID, list detached probes with list_probes, explain that the MCP server cannot stop a probe it did not deploy in this session, and hand the operator the CLI reconnect command to stop it without killing the target JVM." + "rubric": "Find the target PID, list detached probes with list_probes, confirm with the operator that the probe is no longer needed, and stop it with stop_detached_probe by its id without killing the target JVM." }, { "id": "legacy-libs-migration", 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 index d22b307..c7bf877 100644 --- 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 @@ -103,7 +103,8 @@ void handshakeAdvertisesEveryTool() throws Exception { "list_probes", "send_event", "detach_probe", - "exit_probe"), + "exit_probe", + "stop_detached_probe"), tools); } @@ -212,6 +213,29 @@ void listProbesReportsTheRunningProbe() throws Exception { @Test @Order(10) + void stopDetachedProbeRemovesTheListedProbe() throws Exception { + String entry = + text(tool("list_probes")) + .lines() + .filter(line -> line.contains(RUNNING_PROBE)) + .findFirst() + .orElseThrow(() -> new AssertionError("detached probe is not listed")); + + Map stopped = tool("stop_detached_probe", "probe_id", entry); + assertSuccess(stopped); + // Give in-flight handler invocations time to drain, then require the counter to stay frozen. + Thread.sleep(1_000); + long frozen = counter(RUNNING_PROBE, "processCalls"); + Thread.sleep(2_000); + assertEquals( + frozen, counter(RUNNING_PROBE, "processCalls"), "probe still firing after stop_detached_probe"); + + assertFalse(text(tool("list_probes")).contains(RUNNING_PROBE), "stopped probe is still listed"); + assertTrue(isError(tool("stop_detached_probe", "probe_id", entry)), "probe stopped twice"); + } + + @Test + @Order(11) void targetSurvivesTheWholeSession() { assertTrue(sampleApp.isAlive(), "sample app died during the session"); assertTrue(server.isAlive(), "MCP server died during the session"); 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 9d4684a..a3fce05 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 @@ -152,6 +152,13 @@ public void listProbes(String host, CommandListener listener) throws Exception { .invoke(delegate, host, listenerProxy(listener)); } + /** Reconnects to a detached probe and stops it; returns once the agent confirms the exit. */ + public void exitProbe(String host, String probeId) throws Exception { + clientClass + .getMethod("connectAndExitProbe", String.class, String.class) + .invoke(delegate, host, probeId); + } + public void sendExit(int code) throws Exception { clientClass.getMethod("sendExit", int.class).invoke(delegate, code); } diff --git a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceMcpServer.java b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceMcpServer.java index 50a600b..2da640e 100644 --- a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceMcpServer.java +++ b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/BTraceMcpServer.java @@ -14,6 +14,7 @@ //SOURCES ./tools/ListJvmsHandler.java //SOURCES ./tools/ListProbesHandler.java //SOURCES ./tools/SendEventHandler.java +//SOURCES ./tools/StopDetachedProbeHandler.java /* * Copyright (c) 2008, 2024, Jaroslav Bachorik . @@ -41,6 +42,7 @@ import io.btrace.mcp.tools.ListJvmsHandler; import io.btrace.mcp.tools.ListProbesHandler; import io.btrace.mcp.tools.SendEventHandler; +import io.btrace.mcp.tools.StopDetachedProbeHandler; import java.io.IOException; import java.util.ArrayList; import java.util.LinkedHashMap; @@ -167,6 +169,7 @@ private void handleToolsList(Object id) throws IOException { tools.add(SendEventHandler.schema()); tools.add(DetachProbeHandler.schema()); tools.add(ExitProbeHandler.schema()); + tools.add(StopDetachedProbeHandler.schema()); Map result = new LinkedHashMap<>(); result.put("tools", tools); @@ -208,6 +211,9 @@ private void handleToolsCall(Object id, Map params) throws IOExc case "exit_probe": result = ExitProbeHandler.execute(arguments); break; + case "stop_detached_probe": + result = StopDetachedProbeHandler.execute(arguments); + break; default: protocol.sendError(id, -32602, "Unknown tool: " + toolName); return; diff --git a/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/StopDetachedProbeHandler.java b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/StopDetachedProbeHandler.java new file mode 100644 index 0000000..503b007 --- /dev/null +++ b/plugins/btrace-observability/mcp/src/main/java/io/btrace/mcp/tools/StopDetachedProbeHandler.java @@ -0,0 +1,179 @@ +/* + * Copyright (c) 2008, 2026, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.mcp.tools; + +import io.btrace.mcp.BTraceClient; +import io.btrace.mcp.ClientManager; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.FutureTask; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Handles the stop_detached_probe MCP tool - reconnects to a probe that a client detached from and + * stops it. + */ +public final class StopDetachedProbeHandler { + private static final Logger log = LoggerFactory.getLogger(StopDetachedProbeHandler.class); + private static final int DEFAULT_PORT = 2020; + private static final int TIMEOUT_SECONDS = 30; + + private StopDetachedProbeHandler() {} + + /** Returns tool schema for MCP tools/list. */ + public static Map schema() { + Map tool = new LinkedHashMap<>(); + tool.put("name", "stop_detached_probe"); + tool.put("description", "Stop and remove a detached probe by its list_probes id."); + + Map properties = new LinkedHashMap<>(); + + Map pidProp = new LinkedHashMap<>(); + pidProp.put("type", "string"); + pidProp.put("description", "Target JVM PID."); + properties.put("pid", pidProp); + + Map probeIdProp = new LinkedHashMap<>(); + probeIdProp.put("type", "string"); + probeIdProp.put("description", "Probe id from list_probes."); + properties.put("probe_id", probeIdProp); + + Map portProp = new LinkedHashMap<>(); + portProp.put("type", "integer"); + portProp.put("description", "Agent port (2020)."); + properties.put("port", portProp); + + List required = new ArrayList<>(); + required.add("pid"); + required.add("probe_id"); + + Map inputSchema = new LinkedHashMap<>(); + inputSchema.put("type", "object"); + inputSchema.put("properties", properties); + inputSchema.put("required", required); + tool.put("inputSchema", inputSchema); + return tool; + } + + /** Executes the stop_detached_probe tool. */ + public static Map execute(Map arguments) { + return execute(arguments, TimeUnit.SECONDS.toMillis(TIMEOUT_SECONDS)); + } + + static Map execute(Map arguments, long timeoutMillis) { + String pid = getStringArg(arguments, "pid"); + String probeId = probeId(getStringArg(arguments, "probe_id")); + int port = getIntArg(arguments, "port", DEFAULT_PORT); + + if (pid == null || pid.isEmpty()) { + return toolResult("Error: 'pid' parameter is required", true); + } + if (probeId == null) { + return toolResult("Error: 'probe_id' parameter is required", true); + } + + try { + BTraceClient client = ClientManager.getClient(port); + client.attach(pid, null, "."); + // The client waits for the agent to confirm the exit; never let that block the request. + FutureTask stop = + new FutureTask<>( + () -> { + client.exitProbe("localhost", probeId); + return null; + }); + Thread stopper = new Thread(stop, "btrace-stop-" + probeId); + stopper.setDaemon(true); + stopper.start(); + try { + stop.get(timeoutMillis, TimeUnit.MILLISECONDS); + } catch (TimeoutException e) { + try { + client.close(); + } catch (Exception ignore) { + // best effort; closing the socket ends the stopper's wait + } + return toolResult( + "Timed out waiting for PID " + pid + " to confirm that probe " + probeId + " exited", + true); + } catch (ExecutionException e) { + throw e.getCause() instanceof Exception ? (Exception) e.getCause() : e; + } + return toolResult("Probe " + probeId + " stopped and removed from PID " + pid, false); + } catch (Exception e) { + log.error("Failed to stop detached probe", e); + return toolResult( + "Error stopping probe " + probeId + " on PID " + pid + ": " + rootMessage(e), true); + } + } + + /** Accepts a bare id or a whole list_probes line ({@code : []}). */ + private static String probeId(String value) { + if (value == null) { + return null; + } + String entry = value.trim().replaceFirst("^\\d+:\\s+", ""); + return entry.isEmpty() ? null : entry.split("\\s+", 2)[0]; + } + + /** Reflective calls wrap the client's exception; report the client's own message. */ + private static String rootMessage(Throwable e) { + Throwable cause = e; + while (cause.getCause() != null && cause.getMessage() == null) { + cause = cause.getCause(); + } + return cause.getMessage(); + } + + private static String getStringArg(Map args, String key) { + Object val = args == null ? null : args.get(key); + return val == null ? null : val.toString(); + } + + private static int getIntArg(Map args, String key, int defaultVal) { + Object val = args == null ? null : args.get(key); + if (val == null) { + return defaultVal; + } + if (val instanceof Number) { + return ((Number) val).intValue(); + } + try { + return Integer.parseInt(val.toString()); + } catch (NumberFormatException e) { + return defaultVal; + } + } + + private static Map toolResult(String text, boolean isError) { + Map content = new LinkedHashMap<>(); + content.put("type", "text"); + content.put("text", text); + List contentList = new ArrayList<>(); + contentList.add(content); + Map result = new LinkedHashMap<>(); + result.put("content", contentList); + result.put("isError", isError); + return result; + } +} 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 0c4ce35..6919bc7 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 @@ -15,6 +15,7 @@ public class Client { public boolean disconnected; public boolean closed; public String submittedFileName; + public String exitedProbe; private final CountDownLatch closeLatch = new CountDownLatch(1); public Client(int port) { @@ -42,6 +43,14 @@ public void connectAndListProbes(String host, CommandListener listener) throws E listener.onCommand(new FakePrintableCommand(Command.LIST_PROBES, host)); } + public void connectAndExitProbe(String host, String probeId) throws Exception { + exitedProbe = host + ':' + probeId; + if (probeId.equals("hang")) { + // An agent that never confirms the exit: only closing the client ends the wait. + closeLatch.await(); + } + } + public void sendExit(int code) { exitCode = code; } 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 cddf0c8..17f2639 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 @@ -53,6 +53,8 @@ void invokesTheMaskedClientContractReflectively() throws Exception { assertEquals("localhost", client.printableText(probes.get())); assertEquals(Command.LIST_PROBES, client.commandType(probes.get())); assertEquals("", client.printableText(new Command(Command.EXIT))); + client.exitProbe("localhost", "probe-1"); + assertEquals("localhost:probe-1", delegate.exitedProbe); client.sendEvent(); assertEquals("", delegate.event); diff --git a/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceMcpServerTest.java b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceMcpServerTest.java index adc5a59..84989d4 100644 --- a/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceMcpServerTest.java +++ b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/BTraceMcpServerTest.java @@ -27,6 +27,7 @@ void servesInitializeAndCompactToolSchemas() throws Exception { assertTrue(responses[1].contains("\"description\":\"Deploy an oneliner probe.\"")); assertTrue(responses[1].contains("\"description\":\"List local attachable JVMs.\"")); assertTrue(responses[1].contains("\"description\":\"Stop and remove an active probe.\"")); + assertTrue(responses[1].contains("\"name\":\"stop_detached_probe\"")); } @Test @@ -49,7 +50,8 @@ void reportsUnknownMethodsAndMissingToolParameters() throws Exception { String input = "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"missing\"}\n" + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"tools/call\"}\n" - + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"tools/call\",\"params\":{\"name\":\"missing\"}}\n"; + + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"tools/call\",\"params\":{\"name\":\"missing\"}}\n" + + "{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"tools/call\",\"params\":{\"name\":\"stop_detached_probe\",\"arguments\":{\"pid\":\"42\"}}}\n"; ByteArrayOutputStream output = new ByteArrayOutputStream(); new BTraceMcpServer( new McpProtocol( @@ -60,5 +62,6 @@ void reportsUnknownMethodsAndMissingToolParameters() throws Exception { assertTrue(response.contains("Method not found: missing")); assertTrue(response.contains("Missing params")); assertTrue(response.contains("Unknown tool: missing")); + assertTrue(response.contains("Error: 'probe_id' parameter is required")); } } diff --git a/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/tools/StopDetachedProbeHandlerTest.java b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/tools/StopDetachedProbeHandlerTest.java new file mode 100644 index 0000000..ac076f0 --- /dev/null +++ b/plugins/btrace-observability/mcp/src/test/java/io/btrace/mcp/tools/StopDetachedProbeHandlerTest.java @@ -0,0 +1,43 @@ +package io.btrace.mcp.tools; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.btrace.client.Client; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class StopDetachedProbeHandlerTest { + @Test + void stopsTheProbeNamedByAListProbesLine() { + Map result = + StopDetachedProbeHandler.execute(args("1: 2f0c-probe [io.example.Trace]")); + + assertEquals(false, result.get("isError")); + assertEquals("42", Client.last.attachedPid); + assertEquals("localhost:2f0c-probe", Client.last.exitedProbe); + } + + @Test + void givesUpAndClosesTheClientWhenTheAgentNeverConfirms() { + Map result = StopDetachedProbeHandler.execute(args("hang"), 200); + + assertEquals(true, result.get("isError")); + assertTrue(text(result).contains("Timed out"), text(result)); + assertTrue(Client.last.closed, "client must be closed to release the connection"); + } + + private static Map args(String probeId) { + Map args = new LinkedHashMap<>(); + args.put("pid", "42"); + args.put("probe_id", probeId); + return args; + } + + @SuppressWarnings("unchecked") + private static String text(Map result) { + return (String) ((List>) result.get("content")).get(0).get("text"); + } +} diff --git a/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md b/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md index f64edc8..2ae32ff 100644 --- a/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md +++ b/plugins/btrace-observability/skills/btrace-mcp-operations/SKILL.md @@ -6,17 +6,17 @@ description: Use when an AI client should operate BTrace through the BTrace MCP # MCP Operations The BTrace MCP server offers structured local-JVM operations such as `list_jvms`, `deploy_oneliner`, -`deploy_script`, `list_probes`, `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. +`deploy_script`, `list_probes`, `send_event`, `detach_probe`, `exit_probe`, and +`stop_detached_probe`. It is suitable when the AI client and target JVM are on the same host and the +operator wants an auditable conversational workflow. `list_probes` reports only probes that were detached with `detach_probe` and are still running on the target; probes deployed in the current session that are still attached do not appear. Call it before deploying, so probes left behind by an earlier session are found rather than stacked on, and -after any `detach_probe` to confirm what is still running. The MCP server cannot reconnect to or -stop a listed probe: report it to the operator, who can reconnect with `btrace -r ` -and stop it. Track your own deployments yourself: target PID, agent port, probe class or oneliner, -and observation window. +after any `detach_probe` to confirm what is still running. Stop a listed probe with +`stop_detached_probe`, passing its id or the whole `list_probes` line, only after the operator +confirms it is no longer needed; another person or session may own it. Track your own deployments +yourself: 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 @@ -27,7 +27,7 @@ read, such as exported counters (`@Export` fields, read with `jcmd PerfCou target environment. - Treat deployment and cleanup as one operation: record the target PID and call `exit_probe` after the observation window. Prefer `exit_probe` over `detach_probe`; a detached probe keeps running - and is visible afterwards only through `list_probes`. + and can afterwards only be found with `list_probes` and stopped with `stop_detached_probe`. - Use the server's `diagnose_slow_endpoint`, `find_exception_source`, and `profile_method` prompts as starting workflows, then apply the other plugin skills for scope and data safety. diff --git a/plugins/btrace-observability/skills/btrace-probe-lifecycle/SKILL.md b/plugins/btrace-observability/skills/btrace-probe-lifecycle/SKILL.md index c331e03..51cf60b 100644 --- a/plugins/btrace-observability/skills/btrace-probe-lifecycle/SKILL.md +++ b/plugins/btrace-observability/skills/btrace-probe-lifecycle/SKILL.md @@ -14,7 +14,8 @@ condition. 3. Exercise the scenario or observe only for the agreed window. 4. List probes that a client detached from and that are still running with `btrace -lp ` (the `list_probes` tool when working through the BTrace MCP server); reconnect only when - intentionally continuing an existing probe with `btrace -r `. + intentionally continuing an existing probe with `btrace -r `, and stop one with + `btrace -r exit ` (`stop_detached_probe` through the MCP server). 5. Stop/remove the probe at the end of the window, archive or delete the output according to policy, and remove temporary probe files from the target host/container/pod.