Add ai-investigate shark-cli command for AI-driven leak investigation #2796
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,2 @@ | ||
| ./gradlew --quiet --no-configuration-cache :shark:shark-cli:installDist | ||
| ./shark/shark-cli/build/install/shark-cli/bin/shark-ai-investigate "$@" |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,95 @@ | ||
| #!/bin/sh | ||
| # | ||
| # Wrapper for 'shark-cli ... ai-investigate'. | ||
| # | ||
| # 1. Generates a unique session shortcode. | ||
| # 2. Optionally prints the investigation instructions (--no-help skips this). | ||
| # 3. Starts the ai-investigate daemon in the background with that shortcode. | ||
| # 4. Waits until the daemon has created its named pipes (happens after heap | ||
| # analysis, so may take a minute on large dumps). | ||
| # 5. Prints SHORT_CODE=<shortcode> — the AI agent uses this in every command. | ||
| # | ||
| # Usage: | ||
| # shark-ai-investigate --hprof <file> | ||
| # shark-ai-investigate --hprof <file> --obfuscation-mapping <map> | ||
| # shark-ai-investigate --hprof <file> --no-help # skip instructions | ||
| # | ||
|
|
||
| # Resolve the directory that contains this script, following symlinks. | ||
| app_path=$0 | ||
| while [ -h "$app_path" ]; do | ||
| ls=$(ls -ld "$app_path") | ||
| link=${ls#*' -> '} | ||
| case $link in | ||
| /*) app_path=$link ;; | ||
| *) app_path=$(dirname "$app_path")/$link ;; | ||
| esac | ||
| done | ||
| SCRIPT_DIR=$(cd "$(dirname "$app_path")" && pwd -P) | ||
| SHARK_CLI="$SCRIPT_DIR/shark-cli" | ||
|
|
||
| # If the first argument is "cmd", forward the rest to ai-investigate-cmd and exit. | ||
| # shark-ai-investigate cmd <shortcode> <command...> | ||
| if [ "$1" = "cmd" ]; then | ||
| shift | ||
| exec "$SHARK_CLI" ai-investigate-cmd "$@" | ||
| fi | ||
|
|
||
| # Generate a unique session shortcode. | ||
| shortcode=$(uuidgen | tr -d '-' | tr '[:upper:]' '[:lower:]' | cut -c1-8) | ||
|
|
||
| # Check whether --no-help was passed; if so skip the instructions step. | ||
| no_help=false | ||
| for arg in "$@"; do | ||
| case "$arg" in | ||
| --no-help) no_help=true ;; | ||
| esac | ||
| done | ||
|
|
||
| # --no-help is consumed by this wrapper and must not be forwarded to shark-cli. | ||
| # Rebuild $@ without it using a POSIX-safe eval loop. | ||
|
There was a problem hiding this comment. Do you really need POSIX compat? Ancient bash (like the version that ships with Mac) supports arrays, so you can write much simpler code: |
||
| orig_count=$# | ||
| i=1 | ||
| while [ $i -le $orig_count ]; do | ||
| eval "arg=\$$i" | ||
| case "$arg" in | ||
| --no-help) ;; | ||
| *) set -- "$@" "$arg" ;; | ||
| esac | ||
| i=$((i + 1)) | ||
| done | ||
| shift "$orig_count" | ||
| # $@ now contains only the shark-cli global options (--hprof, --device, etc.) | ||
|
|
||
| if ! "$no_help"; then | ||
| # Print instructions (blocking — agent reads them before the session opens). | ||
| "$SHARK_CLI" "$@" ai-investigate --instructions-only | ||
| fi | ||
|
|
||
| # Start the daemon in the background with the session shortcode. | ||
| # Redirect stdout/stderr to /dev/null so the daemon does not inherit this | ||
| # script's stdout pipe. Without this, any process that waits for stdout EOF | ||
| # (e.g. a task runner) would block until the daemon exits, because the daemon | ||
| # holds the write end of the pipe open. The daemon communicates via named | ||
| # pipes, not stdout, so nothing useful is lost. | ||
| "$SHARK_CLI" "$@" ai-investigate --session "$shortcode" >/dev/null 2>&1 & | ||
|
There was a problem hiding this comment. Note that the daemon will still be tied to the session - if you close the window within which you ran this, the daemon will die too. You can fix this with |
||
|
|
||
| # Wait until the daemon creates its named pipes. | ||
|
There was a problem hiding this comment. I don't see any code to capture the pid of the daemon - how to shut it down after? |
||
| # Pipes are created after heap analysis completes, so this may take a while. | ||
| printf 'Running heap analysis and creating pipes\n' | ||
| fifo="/tmp/shark-$shortcode.in" | ||
| i=0 | ||
| while [ ! -p "$fifo" ] && [ "$i" -lt 600 ]; do | ||
| sleep 0.5 | ||
| i=$((i + 1)) | ||
| done | ||
| if [ ! -p "$fifo" ]; then | ||
| printf '{"error":"Daemon did not start within 5 minutes."}\n' | ||
| exit 1 | ||
| fi | ||
|
|
||
| # Print a separator, then the shortcode. | ||
| # The agent passes the shortcode as the first argument to every command: | ||
| # shark-ai-investigate cmd <shortcode> <command> | ||
| printf '\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n' | ||
| printf 'SHORT_CODE=%s\n' "$shortcode" | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,96 @@ | ||
| package shark | ||
|
|
||
| import com.github.ajalt.clikt.core.CliktCommand | ||
| import com.github.ajalt.clikt.core.context | ||
| import com.github.ajalt.clikt.parameters.arguments.argument | ||
| import com.github.ajalt.clikt.parameters.arguments.multiple | ||
| import java.io.File | ||
| import java.io.FileOutputStream | ||
| import java.util.concurrent.CompletableFuture | ||
| import java.util.concurrent.TimeUnit | ||
| import java.util.concurrent.TimeoutException | ||
| import kotlinx.serialization.json.Json | ||
| import kotlinx.serialization.json.jsonObject | ||
| import kotlinx.serialization.json.jsonPrimitive | ||
| import shark.SharkCliCommand.Companion.echo | ||
|
|
||
| /** | ||
| * Thin client that sends a single command to a running [AiInvestigateCommand] daemon and prints | ||
| * the JSON response. The daemon and this command communicate via two named pipes identified by | ||
| * the session shortcode printed by shark-ai-investigate. | ||
| * | ||
| * Usage (start the session first with shark-ai-investigate, then send commands): | ||
| * shark-ai-investigate --hprof <file> # prints SHORT_CODE=<shortcode> | ||
| * shark-ai-investigate cmd <shortcode> trace | ||
| * shark-ai-investigate cmd <shortcode> fields 3 | ||
| * shark-ai-investigate cmd <shortcode> mark-leaking 3 mDestroyed is true | ||
| * shark-ai-investigate cmd <shortcode> close-session | ||
| */ | ||
| class AiInvestigateCmdCommand : CliktCommand( | ||
| name = "ai-investigate-cmd", | ||
| help = "Send a command to a running ai-investigate daemon. First argument is the session shortcode." | ||
| ) { | ||
|
|
||
| init { | ||
| context { | ||
| // Disable @-file expansion so `fields @<objectId>` isn't treated as "load args from file". | ||
| expandArgumentFiles = false | ||
| // Stop option-parsing after the first positional arg (session shortcode). | ||
| // Clikt 2.3.0 treats ANY token containing '=' as a long-option attempt, so a reason like | ||
| // "size=1" would otherwise trigger "no such option: ..." before even reaching the daemon. | ||
| allowInterspersedArgs = false | ||
| } | ||
| } | ||
|
|
||
| private val session by argument(help = "Session shortcode printed by shark-ai-investigate (SHORT_CODE=...)") | ||
|
|
||
| private val commandParts by argument(help = "Command and arguments to send to the daemon") | ||
| .multiple(required = true) | ||
|
|
||
| override fun run() { | ||
| val inPath = "/tmp/shark-$session.in" | ||
| val outPath = "/tmp/shark-$session.out" | ||
|
|
||
| if (!File(inPath).exists() || !File(outPath).exists()) { | ||
| echo("""{"error":"No active session '$session'. Run: shark-ai-investigate --hprof <file>"}""") | ||
| return | ||
| } | ||
|
|
||
| val cmd = commandParts.joinToString(" ") | ||
|
|
||
| // Opening a FIFO blocks until both ends are open. Start both opens concurrently so | ||
| // the client is never waiting sequentially while the daemon waits on the other FIFO. | ||
| // supplyAsync uses the common fork-join pool (daemon threads), matching the daemon side. | ||
| val writerFuture = CompletableFuture.supplyAsync { FileOutputStream(inPath).bufferedWriter() } | ||
| val readerFuture = CompletableFuture.supplyAsync { File(outPath).bufferedReader() } | ||
|
|
||
| try { | ||
| CompletableFuture.allOf(writerFuture, readerFuture).get(10_000L, TimeUnit.MILLISECONDS) | ||
| } catch (e: TimeoutException) { | ||
| echo("""{"error":"Session '$session' is not responding. Daemon may have exited."}""") | ||
| return | ||
| } | ||
|
|
||
| val writer = writerFuture.get() | ||
| val reader = readerFuture.get() | ||
|
|
||
| writer.write(cmd + "\n") | ||
| writer.flush() | ||
| writer.close() | ||
|
|
||
| val response = reader.readLine() | ||
| reader.close() | ||
|
|
||
| if (response != null) { | ||
| if (commandParts.firstOrNull() == "human-readable-trace") { | ||
| // Daemon sends a single JSON line {"trace":"..."}; extract and print plain text. | ||
| val trace = runCatching { | ||
| Json.parseToJsonElement(response).jsonObject["trace"]?.jsonPrimitive?.content | ||
| }.getOrNull() | ||
| echo(trace ?: response) | ||
| } else { | ||
| echo(response) | ||
| } | ||
| } | ||
| } | ||
| } |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I like this design, but I think the name
ai-investigateis misleading, even if it is how this tool is likely to be used. I'd suggestshark-daemon start --hprof <file>andshark-daemon stop --shortcode ...